agent_patch_distr.py

import sys
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors
import json
import math
from abm.replay.replay import ExperimentReplay

########################### DOCUMENTATION ######################################
#
# What agent_patch_distr.py does:
# When doing an experiment with different values for RADIUS_RESOURCE the script creates
# for each value a 2D heatmap plot of the exploitable area (area covered by patches)
# and a plot of the locations of the centers of the patches for testing reasons. The
# plots contain the parameter values in the respective filename and the plot title itself.
# Furthermore normalized_pixel_array.npz is created for doing quick adaptations in the plots.
# The scipt can also handles the case if N=0 and further allows plot the experiment data
# generated by nb_exp13_small_patch_only.py

#
# How to use the scripts_
# - run nb_exp12.py or nb_exp13_small_patch_only.py (define EXPERIMENT_NAME, e.g. patch_place_distr)
# - adapt nb_exp12.py by varying `num_patches`. Define at least two values, e.g. `num_patches = [1, 5]`
#   agent_patch_distr.py fails if only one value is given!
#   adapt nb_exp13_small_patch_only.py by varying `N` and/or `num_batches`
# - adapt filepath in agent_patch_distr.py (e.g. filepath = 'home/ABM/abm/data/simulation_data/patch_place_distr')
# - run agent_patch_distr.py (the directory where to patch_exploitable_area.py is arbitrary)
# - a folder patch_place_distr/summary/agent_patch_distr will be created containing the plots named after
#   the the different parameters used and the .npz-files for redoing the plots
#
################################################################################


# for loading data from /summary
filepath = '/ABM/abm/data/simulation_data/agent_patch_place_distr/'

# for storing plots of patch distribution
folderpath = os.path.join(filepath, 'summary/agent_patch_distr')
window_pad = 30  # hardcoded

# see "Equation of a circle" in  https://en.wikipedia.org/wiki/Circle#Equations
def modifiy_pixel_array_circle(pixel_array, m_x, m_y, radius):
    """ Modifies a 2D pixel array such that in the area of a circle with given
    position (center of the circle) and radius a one is added.

    Parameters
    ----------
    pixel_array : 2d array
    m_x : x-location of circle.
    m_y : y-location of circle.
    radius : radius of circle.

    Returns: modified pixel_array
    """
    # presuming that arena consists of pixel id = 0 until pixel id = ENV_WIDTH-1
    # instead of id = 1 until pixel id = ENV_WIDTH (and respectively for ENV_HEIGHT)
    ENV_HEIGHT, ENV_WIDTH = np.shape(pixel_array)

    #  create arrays of indices
    x_ind = np.arange(0, ENV_WIDTH)
    y_ind = np.arange(0, ENV_HEIGHT)

    # create numpy mask
    mask_circle = (x_ind[np.newaxis, :] - m_x)**2 + (y_ind[:, np.newaxis] - m_y)**2 <= radius**2
    pixel_array[mask_circle] += 1

    return pixel_array


def save_pixel_array(normalized_pixel_array, locations_or_circles, agent_or_patch,
                    R, N, radii_resources, N_R, num_batches, folderpath):
    """ Saves a pixel array for adapting plots."""
    filename = (f'{agent_or_patch}_normalized_pixel_array_{locations_or_circles}_R_{R}'
                f'_N_{N}_R_R_{radii_resources}'
                f'_N_R_{N_R}_batches_{num_batches}.npz')
    final_path = os.path.join(folderpath, filename)
    np.savez(final_path, normalized_pixel_array=normalized_pixel_array)
    return

def loop_params_batches(agent_or_patch, folderpath, ENV_HEIGHT, ENV_WIDTH, posx, posy,
                        RADIUS_AGENT, N, values_N_RESOURCES, radii_resources, nr_of_radii,
                        num_batches, window_pad):
    """ Loops through the different radii values of the patches and
        saves plots of either the patches or the agents and the respective pixel arrays
        from which the plots are genereated.

    Parameters:

    agent_or_patch : (string) Set to 'patch' for plotting patches or 'agent' for plotting agents.
    folderpath : Path where plots and pixel arrays are saved.
    ENV_HEIGHT : Height of arena in pixels.
    ENV_WIDTH : Width of arena in pixels.
    posx : numpy.ndarray containing x-positions of either patches or agents depending agent_or_patch
    posy : numpy.ndarray containing y-positions of either patches or agents depending agent_or_patch
    RADIUS_AGENT : (float) Radius of the agents.
    N : (int) Number of agents.
    values_N_RESOURCES : (list) Different values of the number of resources
    radii_resources : (list) Different values of the radii of the resources
    nr_of_radii : (int) Number of different radii.
    num_batches : (int) number of batches
    window_pad : (int) parameter for padding of arena in pygame window.
    """
    for i in range(0, nr_of_radii):
        if agent_or_patch == 'patch':
            number_of_entities = values_N_RESOURCES[i] # "entity" as abstraction for agent and patch
            radius_of_entity = radii_resources[nr_of_radii - 1 - i]
        else:
            number_of_entities = N
            radius_of_entity = RADIUS_AGENT

        # See x = np.random.randint(self.window_pad - radius, self.WIDTH + self.window_pad - radius) in sims.py.
        # For x = np.random.randint(float, float) it rounds down so here we might need to ceil.
        x_values = posx[:, i, nr_of_radii - 1 - i, 0:int(number_of_entities), 0].flatten() - window_pad + int(radius_of_entity)
        y_values = posy[:, i, nr_of_radii - 1 - i, 0:int(number_of_entities), 0].flatten() - window_pad + int(radius_of_entity)

        # making sure that no drawn location is outside the arena borders
        max_x_values = np.max(x_values)
        max_y_values = np.max(y_values)
        if (max_x_values > ENV_WIDTH) | (max_y_values > ENV_HEIGHT):
            sys.exit('At least one patch location is outside the arena borders.')

        # loop through agent/patch locations and "create" a filled CIRCLE for each location in pixel array
        pixel_array = np.zeros([ENV_HEIGHT, ENV_WIDTH])
        for k in range(0, len(x_values)):
            pixel_array = modifiy_pixel_array_circle(pixel_array, x_values[k], y_values[k], radius_of_entity)

        # probability density for num_batches initializations
        normalized_pixel_array = pixel_array / pixel_array.sum()

        save_pixel_array(normalized_pixel_array, 'circles', agent_or_patch, round(RADIUS_AGENT, 2),
                        int(N), round(radii_resources[nr_of_radii - 1 - i], 2),
                        int(values_N_RESOURCES[i]), num_batches, folderpath)

        plot_circles(agent_or_patch, normalized_pixel_array, RADIUS_AGENT, N,
                    radii_resources[nr_of_radii - 1 - i], values_N_RESOURCES[i], num_batches, folderpath)

        # for all agent/patch locations add + 1 PIXEL in pixel array
        pixel_array = np.zeros([ENV_HEIGHT, ENV_WIDTH])
        pixel_array[[int(i) for i in y_values], [int(i) for i in x_values]] += 1

        # probability density for num_batches initializations
        normalized_pixel_array = pixel_array / pixel_array.sum()

        save_pixel_array(normalized_pixel_array, 'locations', agent_or_patch, round(RADIUS_AGENT, 2),
                        int(N), round(radii_resources[nr_of_radii - 1 - i], 2),
                        int(values_N_RESOURCES[i]), num_batches, folderpath)

        plot_centers(agent_or_patch,normalized_pixel_array, RADIUS_AGENT,
                    N, radii_resources[nr_of_radii - 1 - i], values_N_RESOURCES[i], num_batches, folderpath)
    return

def loop_batches_small_patch(agent_or_patch, folderpath, ENV_HEIGHT, ENV_WIDTH,
                            posx, posy, RADIUS_AGENT, N, N_RESOURCES, RADIUS_RESOURCE,
                            num_batches, window_pad):
    """ Saves plots of the agents for a fixed set of parameters and
        and the respective pixel arrays from which the plots are genereated. It is made
        for plotting the agents for the case of only one resource patch af radius 1.

    Parameters:

    agent_or_patch : (string) Set to 'patch' for plotting patches or 'agent' for plotting agents.
    folderpath : Path where plots and pixel arrays are saved.
    ENV_HEIGHT : Height of arena in pixels.
    ENV_WIDTH : Width of arena in pixels.
    posx : numpy.ndarray containing x-positions of agents.
    posy : numpy.ndarray containing y-positions of agents.
    RADIUS_AGENT : (float) Radius of the agents.
    N : (int) Number of agents.
    N_RESOURCES : (int) Number of resources
    RADIUS_RESOURCE : (int) Radius of the resources.
    num_batches : (int) number of batches
    window_pad : (int) parameter for padding of arena in pygame window.
    """
    # See x = np.random.randint(self.window_pad - radius, self.WIDTH + self.window_pad - radius) in sims.py.
    # For x = np.random.randint(float, float) it rounds down so here we might need to ceil.
    x_values = posx[:, 0:int(N), 0].flatten() - window_pad + int(RADIUS_AGENT)
    y_values = posy[:, 0:int(N), 0].flatten() - window_pad + int(RADIUS_AGENT)

    # making sure that no drawn location is outside the arena borders
    max_x_values = np.max(x_values)
    max_y_values = np.max(y_values)
    if (max_x_values > ENV_WIDTH) | (max_y_values > ENV_HEIGHT):
        sys.exit('At least one patch location is outside the arena borders.')

    # loop through agent/patch locations and "create" a filled CIRCLE for each location in pixel array
    pixel_array = np.zeros([ENV_HEIGHT, ENV_WIDTH])
    for k in range(0, len(x_values)):
        pixel_array = modifiy_pixel_array_circle(pixel_array, x_values[k], y_values[k], RADIUS_AGENT)

    # probability density for num_batches initializations
    normalized_pixel_array = pixel_array / pixel_array.sum()

    save_pixel_array(normalized_pixel_array, 'circles', agent_or_patch, round(RADIUS_AGENT, 2),
                    int(N), round(RADIUS_RESOURCE, 2), int(N_RESOURCES), num_batches, folderpath)

    plot_circles(agent_or_patch, normalized_pixel_array, RADIUS_AGENT, N,
                RADIUS_RESOURCE, N_RESOURCES, num_batches, folderpath)

    # for all agent/patch locations add + 1 PIXEL in pixel array
    pixel_array = np.zeros([ENV_HEIGHT, ENV_WIDTH])
    pixel_array[[int(i) for i in y_values], [int(i) for i in x_values]] += 1

    # probability density for num_batches initializations
    normalized_pixel_array = pixel_array / pixel_array.sum()

    save_pixel_array(normalized_pixel_array, 'locations', agent_or_patch, round(RADIUS_AGENT, 2),
                    int(N), round(RADIUS_RESOURCE, 2), int(N_RESOURCES), num_batches, folderpath)

    plot_centers(agent_or_patch, normalized_pixel_array, RADIUS_AGENT, N, RADIUS_RESOURCE,
                N_RESOURCES, num_batches, folderpath)
    return


def plot_circles(agent_or_patch, pixel_array, RADIUS_AGENT, N,
                RADIUS_RESOURCE, N_RESOURCES, num_batches, folderpath):
    """ Plots filled circles for agents and patches."""
    plt.figure()
    plt.xlabel('location x')
    plt.ylabel('location y')
    plt.title(f'{agent_or_patch} R={round(RADIUS_AGENT, 2)} N={int(N)} '
              f'R_R={round(RADIUS_RESOURCE, 2)} N_R={int(N_RESOURCES)} batches={num_batches}')
    color_map = plt.imshow(pixel_array)
    color_map.set_cmap("viridis")
    plt.colorbar(label='Probability density')
    plt.tight_layout()
    filename = (f'{agent_or_patch}_covered_area_R_{round(RADIUS_AGENT, 2)}_N_{int(N)}'
                f'_R_R_{round(RADIUS_RESOURCE, 2)}_N_R_{int(N_RESOURCES)}_batches_{num_batches}.png')
    final_path = os.path.join(folderpath, filename)
    plt.savefig(final_path)
    return

def plot_centers(agent_or_patch, pixel_array, RADIUS_AGENT, N, RADIUS_RESOURCE,
                N_RESOURCES, num_batches, folderpath):
    """ Plots locations (centers) of agents and patches."""
    plt.figure()
    plt.xlabel('location x')
    plt.ylabel('location y')
    plt.title(f'{agent_or_patch} R={round(RADIUS_AGENT, 2)} N={int(N)} R_R={round(RADIUS_RESOURCE, 2)} '
              f'N_R={int(N_RESOURCES)} batches={num_batches}')
    color_map = plt.imshow(pixel_array)
    color_map.set_cmap("viridis")
    plt.colorbar(label='Probability density')
    plt.tight_layout()
    filename = (f'{agent_or_patch}_locations_R_{round(RADIUS_AGENT, 2)}_N_{int(N)}'
                f'_R_R_{round(RADIUS_RESOURCE, 2)}_N_R_{int(N_RESOURCES)}_batches_{num_batches}.png')
    final_path = os.path.join(folderpath, filename)
    plt.savefig(final_path)

    return


# load data from /summary
data = ExperimentReplay(filepath, undersample=1)

# create file for storing plots
if os.path.exists(folderpath) == False:
    os.mkdir(folderpath)

# read out parameters from .json-files for printing/showing them in plots and filesnames
# fixed_env
# abm/data/simulation_data/experiment_name/summary/fixed_env.json
with open(os.path.join(filepath, 'summary/fixed_env.json'), 'r') as k:
    fixed_env = json.loads(k.read())

ENV_WIDTH = fixed_env['ENV_WIDTH']
ENV_HEIGHT = fixed_env['ENV_HEIGHT']

# get agent parameters
RADIUS_AGENT = fixed_env['RADIUS_AGENT']
N = fixed_env['N']
posx = data.posx_z
posy = data.posy_z

# for experiments with N_RESOURCES=1 RADIUS_RESOURCE=1
if len(posx.shape) == 3:
    agent_or_patch = 'agent'
    RADIUS_RESOURCE = fixed_env['RADIUS_RESOURCE']
    N_RESOURCES = fixed_env['N_RESOURCES']
    num_batches = posx.shape[0]
    loop_batches_small_patch(agent_or_patch, folderpath, ENV_HEIGHT, ENV_WIDTH,
                            posx, posy, RADIUS_AGENT, N, N_RESOURCES, RADIUS_RESOURCE,
                            num_batches, window_pad)
    sys.exit()

# tuned_env
# abm/data/simulation_data/experiment_name/summary/tuned_env.json
with open(os.path.join(filepath, 'summary/tuned_env.json'), 'r') as k:
    tuned_env = json.loads(k.read())

# get resource parameters
radii_resources = tuned_env['RADIUS_RESOURCE'] # different radii used in experiment
values_N_RESOURCES = tuned_env['N_RESOURCES']
res_posx = data.res_pos_x_z
res_posy = data.res_pos_y_z

num_batches, nr_of_radii, length_num_patches, max_N_RESOURCES, T = res_posx.shape

agent_or_patch = 'patch'
loop_params_batches(agent_or_patch, folderpath, ENV_HEIGHT, ENV_WIDTH,
                    res_posx, res_posy, RADIUS_AGENT, N, values_N_RESOURCES, radii_resources,
                    nr_of_radii, num_batches, window_pad)

# in case of N=0 no agents are created, so plotting in following lines is not executed
if N == 0:
    sys.exit('No plots for agents created as number of agents N=0. Plots for patches created successfully.')
agent_or_patch = 'agent'
loop_params_batches(agent_or_patch, folderpath, ENV_HEIGHT, ENV_WIDTH, posx, posy,
                    RADIUS_AGENT, N, values_N_RESOURCES, radii_resources,
                    nr_of_radii, num_batches, window_pad)