olfactory_navigation

Agent

A generic agent class.

It is meant to define the general structure for an agent meant to evolve in a environment of olfactory cues. To define such agent, a set of methods need to be implemented. This methods can be seperated into 3 categories:

1. Training methods
2. Simulation methods
3. General methods

The training methods are meant to train the agent before testing their performance in a simulation. A single method is needed for this:

• train()

The simulation methods are meant for the agent to make choices and receiving observations during a simulation. The following methods are required for this:

• initialize_state(): This method is meant for the state of the agent(s) to be initialized before the simulation loop starts. The state of the agent can be an internal clock, a belief or something else arbitrary.
• choose_action(): Here the agent(s) is asked to choose an action to play based on its internal state.
• update_state(): Then, after the agent(s) has taken an action, the observation it makes along with whether he reached the source or not is returned to him using this method. This allows the agent to update its internal state.
• kill(): Finally, the method asks for a set of agents to be terminated. The basic case happens when the agent reaches the source but it can also be asked to terminate if it has reached the end of the simulation without success.

The general methods are methods to perform general actions with the agent. These methods are:

• save(): To save the agent to long term storage.
• load(): To load the agent from long term storage.
• modify_environment(): To provide an equivalent agent with a different environment linked to it. If the agent has previously been trained, the trained components needs to be adapted to this new environment.
• to_gpu(): To create an alternative version of the agent whether the array instances are stored on the GPU memory instead of the CPU memory.
• to_cpu(): To create an alternative version of the agent whether the array instances are stored on the CPU memory instead of the GPU memory.

For a user to implement an agent, the main methods to define are the Simulation methods! The training method is, as stated, optional, as some agent definitions do not require it. And the General methods all have some default behavior and are therefore only needed to be overwritten in specific cases.

Parameters:

Name	Type	Description	Default
environment	Environment	The olfactory environment the agent is meant to evolve in.	required
thresholds	float or list[float] or dict[str, float] or dict[str, list[float]]	The olfactory thresholds. If an odor cue above this threshold is detected, the agent detects it, else it does not. If a list of thresholds is provided, he agent should be able to detect |thresholds|+1 levels of odor. A dictionary of (list of) thresholds can also be provided when the environment is layered. In such case, the number of layers provided must match the environment's layers and their labels must match. The thresholds provided will be converted to an array where the levels start with -inf and end with +inf.	3e-6
space_aware	bool	Whether the agent is aware of it's own position in space. This is to be used in scenarios where, for example, the agent is an enclosed container and the source is the variable. Note: The observation array will have a different shape when returned to the update_state function!	False
spacial_subdivisions	ndarray	How many spacial compartments the agent has to internally represent the space it lives in. By default, it will be as many as there are grid points in the environment.	None
actions	dict or ndarray	The set of action available to the agent. It should match the type of environment (ie: if the environment has layers, it should contain a layer component to the action vector, and similarly for a third dimension). Else, a dict of strings and action vectors where the strings represent the action labels. If none is provided, by default, all unit movement vectors are included and shuch for all layers (if the environment has layers.)	None
name	str	A custom name for the agent. If it is not provided it will be named like "-thresh_".	None
seed	int	For reproducible randomness.	12131415

Attributes:

Name	Type	Description
environment	Environment
thresholds	ndarray	An array of the thresholds of detection, starting with -inf and ending with +inf. In the case of a 2D array of thresholds, the rows of thresholds apply to the different layers of the environment.
space_aware	bool
spacial_subdivisions	ndarray
name	str
action_set	ndarray	The actions allowed of the agent. Formulated as movement vectors as [(layer,) (dz,) dy, dx].
action_labels	list[str]	The labels associated to the action vectors present in the action set.
saved_at	str	If the agent has been saved, the path at which it is saved is recorded in this variable.
on_gpu	bool	Whether the arrays are on the GPU memory or not. For this, the support for Cupy needs to be enabled and the agent needs to have been moved to the GPU using the to_gpu() function.
class_name	str	The name of the class of the agent.
seed	int	The seed used for the random operations (to allow for reproducability).
rnd_state	RandomState	The random state variable used to generate random values.
cpu_version	Agent	An instance of the agent on the CPU. If it already is, it returns itself.
gpu_version	Agent	An instance of the agent on the CPU. If it already is, it returns itself.

Source code in olfactory_navigation/agent.py

class Agent:
    '''
    A generic agent class.

    It is meant to define the general structure for an agent meant to evolve in a environment of olfactory cues.
    To define such agent, a set of methods need to be implemented. This methods can be seperated into 3 categories:

    1. Training methods
    2. Simulation methods
    3. General methods

    The training methods are meant to train the agent before testing their performance in a simulation. A single method is needed for this:

    - train()

    The simulation methods are meant for the agent to make choices and receiving observations during a simulation. The following methods are required for this:

    - initialize_state(): This method is meant for the state of the agent(s) to be initialized before the simulation loop starts. The state of the agent can be an internal clock, a belief or something else arbitrary.
    - choose_action(): Here the agent(s) is asked to choose an action to play based on its internal state.
    - update_state(): Then, after the agent(s) has taken an action, the observation it makes along with whether he reached the source or not is returned to him using this method. This allows the agent to update its internal state.
    - kill(): Finally, the method asks for a set of agents to be terminated. The basic case happens when the agent reaches the source but it can also be asked to terminate if it has reached the end of the simulation without success.

    The general methods are methods to perform general actions with the agent. These methods are:

    - save(): To save the agent to long term storage.
    - load(): To load the agent from long term storage.
    - modify_environment(): To provide an equivalent agent with a different environment linked to it. If the agent has previously been trained, the trained components needs to be adapted to this new environment.
    - to_gpu(): To create an alternative version of the agent whether the array instances are stored on the GPU memory instead of the CPU memory.
    - to_cpu(): To create an alternative version of the agent whether the array instances are stored on the CPU memory instead of the GPU memory.

    For a user to implement an agent, the main methods to define are the Simulation methods! The training method is, as stated, optional, as some agent definitions do not require it.
    And the General methods all have some default behavior and are therefore only needed to be overwritten in specific cases.


    Parameters
    ----------
    environment : Environment
        The olfactory environment the agent is meant to evolve in.

    thresholds : float or list[float] or dict[str, float] or dict[str, list[float]], default=3e-6
        The olfactory thresholds. If an odor cue above this threshold is detected, the agent detects it, else it does not.
        If a list of thresholds is provided, he agent should be able to detect |thresholds|+1 levels of odor.
        A dictionary of (list of) thresholds can also be provided when the environment is layered.
        In such case, the number of layers provided must match the environment's layers and their labels must match.
        The thresholds provided will be converted to an array where the levels start with -inf and end with +inf.
    space_aware : bool, default=False
        Whether the agent is aware of it's own position in space.
        This is to be used in scenarios where, for example, the agent is an enclosed container and the source is the variable.
        Note: The observation array will have a different shape when returned to the update_state function!
    spacial_subdivisions : np.ndarray, optional
        How many spacial compartments the agent has to internally represent the space it lives in.
        By default, it will be as many as there are grid points in the environment.
    actions : dict or np.ndarray, optional
        The set of action available to the agent. It should match the type of environment (ie: if the environment has layers, it should contain a layer component to the action vector, and similarly for a third dimension).
        Else, a dict of strings and action vectors where the strings represent the action labels.
        If none is provided, by default, all unit movement vectors are included and shuch for all layers (if the environment has layers.)
    name : str, optional
        A custom name for the agent. If it is not provided it will be named like "<class_name>-thresh_<threshold>".
    seed : int, default=12131415
        For reproducible randomness.

    Attributes
    ----------
    environment : Environment
    thresholds : np.ndarray
        An array of the thresholds of detection, starting with -inf and ending with +inf.
        In the case of a 2D array of thresholds, the rows of thresholds apply to the different layers of the environment.
    space_aware : bool
    spacial_subdivisions : np.ndarray
    name : str
    action_set : np.ndarray
        The actions allowed of the agent. Formulated as movement vectors as [(layer,) (dz,) dy, dx].
    action_labels : list[str]
        The labels associated to the action vectors present in the action set.
    saved_at : str
        If the agent has been saved, the path at which it is saved is recorded in this variable.
    on_gpu : bool
        Whether the arrays are on the GPU memory or not. For this, the support for Cupy needs to be enabled and the agent needs to have been moved to the GPU using the to_gpu() function.
    class_name : str
        The name of the class of the agent.
    seed : int
        The seed used for the random operations (to allow for reproducability).
    rnd_state : np.random.RandomState
        The random state variable used to generate random values.
    cpu_version : Agent
        An instance of the agent on the CPU. If it already is, it returns itself.
    gpu_version : Agent
        An instance of the agent on the CPU. If it already is, it returns itself.
    '''
    def __init__(self,
                 environment: Environment,
                 thresholds: float | list[float] | dict[str, float] | dict[str, list[float]] = 3e-6,
                 space_aware: bool = False,
                 spacial_subdivisions: np.ndarray | None = None,
                 actions: dict[str, np.ndarray] | np.ndarray | None = None,
                 name: str | None = None,
                 seed: int = 12131415
                 ) -> None:
        self.environment = environment
        self.space_aware = space_aware

        # Handle thresholds
        if isinstance(thresholds, float):
            thresholds = [thresholds]

        if isinstance(thresholds, list):
            # Ensure -inf and +inf begin and end the thresholds list
            if -np.inf not in thresholds:
                thresholds = [-np.inf] + thresholds

            if np.inf not in thresholds:
                thresholds = thresholds + [np.inf]

            self.thresholds = np.array(thresholds)

        # Handle the case where a dict of thresholds is provided in the case of a layered environment and ensuring they are ordered.
        elif isinstance(thresholds, dict):
            assert self.environment.has_layers, "Thresholds can only be provided as a dict if the environment is layered."
            assert len(thresholds) == len(self.environment.layer_labels), "Thresholds must be given for each layers."

            if all(isinstance(layer_thresholds, list) for layer_thresholds in thresholds.values()):
                layer_thresholds_count = len(thresholds.values()[0])
                assert all(len(layer_thresholds) == layer_thresholds_count for layer_thresholds in thresholds.values()), "If provided as lists, the threshold lists must all be of the same length."
            else:
                assert all(isinstance(layer_threshold, float) for layer_threshold in thresholds.values()), "The thresholds provided in a dictionary must all be list or float, not a combination of both."

            # Processing layers of thresholds
            layer_thresholds_list = []
            for layer_label in enumerate(self.environment.layer_labels):
                assert layer_label in thresholds, f"Environment's layer [{layer_label}] was not matched to any layers of the environment."
                layer_thresholds = thresholds[layer_label]

                if not isinstance(layer_thresholds, list):
                    layer_thresholds = [layer_thresholds]

                # Ensure -inf and +inf begin and end the thresholds list
                if -np.inf not in layer_thresholds:
                    layer_thresholds = [-np.inf] + layer_thresholds

                if np.inf not in layer_thresholds:
                    layer_thresholds = layer_thresholds + [np.inf]

                layer_thresholds_list.append(layer_thresholds)

            # Building an array with the layers of thresholds
            self.thresholds = np.array(layer_thresholds_list)

        # Ensuring the thresholds are in ascending order
        self.thresholds.sort()

        # Spacial subdivisions
        if spacial_subdivisions is None:
            self.spacial_subdivisions = np.array(self.environment.shape)
        else:
            self.spacial_subdivisions = np.array(spacial_subdivisions)
        assert len(self.spacial_subdivisions) == self.environment.dimensions, "The amount of spacial divisions must match the amount of dimensions of the environment."

        # Mapping environment to spacial subdivisions
        env_shape = np.array(self.environment.shape)

        std_size = (env_shape / self.spacial_subdivisions).astype(int)
        overflows = (env_shape % self.spacial_subdivisions)

        cell_sizes = [(np.repeat(size, partitions) + np.array([np.floor(overflow / 2), *np.zeros(partitions-2), np.ceil(overflow / 2)])).astype(int)
                    for partitions, size, overflow in zip(self.spacial_subdivisions, std_size, overflows)]

        # Finding the edges of the cells and filling a grid with ids
        cell_edges = [np.concatenate(([0], np.cumsum(ax_sizes))) for ax_sizes in cell_sizes]

        lower_bounds = np.array([ax_arr.ravel() for ax_arr in np.meshgrid(*[bounds_arr[:-1] for bounds_arr in cell_edges], indexing='ij')]).T
        upper_bounds = np.array([ax_arr.ravel() for ax_arr in np.meshgrid(*[bounds_arr[1 :] for bounds_arr in cell_edges], indexing='ij')]).T

        self._environment_to_subdivision_mapping = np.full(self.environment.shape, -1)
        for i, (lower_b, upper_b) in enumerate(zip(lower_bounds, upper_bounds)):
            slices = [slice(ax_lower, ax_upper) for ax_lower, ax_upper in zip(lower_b, upper_b)]

            # Grid to cell mapping
            self._environment_to_subdivision_mapping[*slices] = i

        # Allowed actions
        self.action_labels = None
        if actions is None:
            if environment.dimensions == 2:
                self.action_set = np.array([
                    [-1,  0], # North
                    [ 0,  1], # East
                    [ 1,  0], # South
                    [ 0, -1]  # West
                ])
                self.action_labels = [
                    'North',
                    'East',
                    'South',
                    'West'
                ]
            elif environment.dimensions == 3:
                self.action_set = np.array([
                    [ 0, -1,  0], # North
                    [ 0,  0,  1], # East
                    [ 0,  1,  0], # South
                    [ 0,  0, -1], # West
                    [ 1,  0,  0], # Up
                    [-1,  0,  0]  # Down
                ])
                self.action_labels = [
                    'North',
                    'East',
                    'South',
                    'West',
                    'Up',
                    'Down'
                ]
            else: # ND
                self.action_set = np.zeros((2*environment.dimensions, environment.dimensions))
                self.action_labels = []
                for dim in range(environment.dimensions):
                    # Increase in dimension 'dim'
                    self.action_set[dim*2, -dim-1] = 1
                    self.action_labels.append(f'd{dim}+1')

                    # Decrease in dimension 'dim'
                    self.action_set[(dim*2) + 1, -dim-1] = -1
                    self.action_labels.append(f'd{dim}-1')

            # Layered
            if environment.has_layers:
                self.action_set = np.array([[layer, *action_vector] for layer in environment.layers for action_vector in self.action_set])
                self.action_labels = [f'l_{layer}_{action}' for  layer in environment.layer_labels for action in self.action_labels]

        # Actions provided as numpy array
        elif isinstance(actions, np.ndarray):
            self.action_set = actions
            self.action_labels = ['a_' + '_'.join([str(dim_a) for dim_a in action_vector]) for action_vector in self.action_set]

        # Actions provided as dict
        else:
            self.action_set = np.ndarray(list(actions.values()))
            self.action_labels = list(actions.keys())

        # Asertion that the shape of the actions set if right
        layered = 0 if not environment.has_layers else 1
        assert self.action_set.shape[1] == (layered + environment.dimensions), f"The shape of the action_set provided is not right. (Found {self.action_set.shape}; expected (., {layered + environment.dimensions}))"

        # setup name
        if name is None:
            self.name = self.class_name
            if len(self.thresholds.shape) == 2:
                thresh_string = '-'.join(['_'.join([f'l{row_i}'] + [str(item) for item in row]) for row_i, row in enumerate(self.thresholds[:,1:-1])])
            else:
                thresh_string = '_'.join([str(item) for item in self.thresholds])
            self.name += f'-tresholds-' + thresh_string
            self.name += '-space_aware' if self.space_aware else ''
        else:
            self.name = name

        # Other variables
        self.saved_at = None

        self.on_gpu = False
        self._alternate_version = None

        # random state
        self.seed = seed
        self.rnd_state = np.random.RandomState(seed = seed)


    @property
    def class_name(self):
        '''
        The name of the class of the agent.
        '''
        return self.__class__.__name__


    # ----------------
    # Training methods
    # ----------------
    def train(self) -> None:
        '''
        Optional function to train the agent in the olfactory environment it is in.
        This function is optional as some agents have some fixed behavior and therefore dont require training.
        '''
        raise NotImplementedError('The train function is not implemented, make an agent subclass to implement the method')


    # ------------------
    # Simulation methods
    # ------------------
    def initialize_state(self,
                         n: int = 1
                         ) -> None:
        '''
        Function to initialize the internal state of the agent(s) for the simulation process. The internal state can be concepts such as the "memory" or "belief" of the agent.
        The n parameter corresponds to how many "instances" need to instanciated. This is meant so that we work with a "group" of agents instead of individual instances.

        This is done with the purpose that the state of the group of agents be stored in (Numpy) arrays to allow vectorization instead of sequential loops.

        Parameters
        ----------
        n : int, default=1
            How many agents to initialize.
        '''
        raise NotImplementedError('The initialize_state function is not implemented, make an agent subclass to implement the method')


    def choose_action(self) -> np.ndarray:
        '''
        Function to allow for the agent(s) to choose an action to take based on its current state.

        It should return a 2D array of shape n by 2 (or 3, or 4 depending of whether the environment has layers and/or a 3rd dimension),
        where n is how many agents are to choose an action. It should be n 2D vectors of (the layer and) the change in the (z,) y, and x positions.

        Returns
        -------
        movement_vector : np.ndarray
            An array of n vectors in 2D space of the movement(s) the agent(s) will take.
        '''
        raise NotImplementedError('The choose_action function is not implemented, make an agent subclass to implement the method')


    def discretize_observations(self,
                                observation: np.ndarray,
                                action: np.ndarray,
                                source_reached: np.ndarray
                                ) -> np.ndarray:
        '''
        Function to convert a set of observations to discrete observation ids.
        It uses the olfaction thresholds of the agent to discretize the odor concentrations.

        In the case where the agent is also aware of it's own position in space, in which case the observation matrix is of size n by 1 + environment.dimensions,
        the agent first converts the points on the grid to position ids and then multiply the id by olfactory observation id.

        Parameters
        ----------
        observation : np.ndarray
            The observations the agent receives, in the case the agent is space_aware, the position point is also included.
        action: np.ndarray
            The action the agent did. This parameter is used in the particular case where the environment has layers and the odor thresholds are layer dependent.
        source_reached : np.ndarray
            A 1D array of boolean values signifying whether each agent reached or not the source.

        Returns
        -------
        discrete_observations: np.ndarray
            An integer array of discrete observations
        '''
        # GPU support
        xp = np if not self.on_gpu else cp

        n = len(observation)
        discrete_observations = xp.ones(n, dtype=int) * -1

        # Gather olfactory observation
        olfactory_observation = observation if not self.space_aware else observation[:,0]

        # Compute the amount of observations available
        observation_count = self.thresholds.shape[-1] - 1 # |thresholds| - 1 observation buckets

        # Handle the special case where we have a layered environment and different thresholds for these layers
        if self.environment.has_layers and len(self.thresholds.shape) == 2:
            layer_ids = action[:,0] # First column of actions is the layer
            action_layer_thresholds = self.thresholds[layer_ids]
            observation_ids = xp.argwhere((olfactory_observation[:,None] >= action_layer_thresholds[:,:-1]) & (olfactory_observation[:,None] < action_layer_thresholds[:,1:]))[:,1]
        else:
            # Setting observation ids
            observation_ids = xp.argwhere((olfactory_observation[:,None] >= self.thresholds[:-1][None,:]) & (olfactory_observation[:,None] < self.thresholds[1:][None,:]))[:,1]

        # If needed, multiply observation indices by position indices
        if not self.space_aware:
            discrete_observations = observation_ids
        else:
            position_clipped = xp.clip(observation[:,1:], a_min=0, a_max=(xp.array(self.environment.shape)-1))
            position_count = int(xp.prod(self.spacial_subdivisions))
            position_ids = self._environment_to_subdivision_mapping[*position_clipped.astype(int).T]

            # Add the amount of possible positions to the observation count
            observation_count *= position_count

            # Combine with observation ids
            discrete_observations = (observation_ids * position_count) + position_ids

        # Adding the goal observation
        discrete_observations[source_reached] = observation_count

        return discrete_observations


    def update_state(self,
                     action: np.ndarray,
                     observation: np.ndarray,
                     source_reached: np.ndarray
                     ) -> None | np.ndarray:
        '''
        Function to update the internal state(s) of the agent(s) based on the action(s) taken and the observation(s) received.
        The observations are then compared with the thresholds to decide whether something was sensed or not or to what level.

        Parameters
        ----------
        action : np.ndarray
            A 2D array of n movement vectors. If the environment is layered, the 1st component should be the layer.
        observation : np.ndarray
            A n by 1 (or 1 + environment.dimensions if space_aware) array of odor cues (float values) retrieved from the environment.
        source_reached : np.array
            A 1D array of boolean values signifying whether each agent reached or not the source.

        Returns
        -------
        update_successfull : np.ndarray, optional
            If nothing is returned, it means all the agent's state updates have been successfull.
            Else, a boolean np.ndarray of size n can be returned confirming for each agent whether the update has been successful or not.
        '''
        raise NotImplementedError('The update_state function is not implemented, make an agent subclass to implement the method')


    def kill(self,
             simulations_to_kill: np.ndarray
             ) -> None:
        '''
        Function to kill any agents that either reached the source or failed by not reaching the source before the horizon or failing to update its own state.
        The agents where the simulations_to_kill paramater is True have to removed from the list of agents.
        It is necessary because their reference will also be removed from the simulation loop. Therefore, if they are not removed, the array sizes will not match anymore.

        Parameters
        ----------
        simulations_to_kill : np.ndarray
            An array of size n containing boolean values of whether or not agent's simulations are terminated and therefore should be removed.
        '''
        raise NotImplementedError('The kill function is not implemented, make an agent subclass to implement the method')


    # ---------------
    # General methods
    # ---------------
    def save(self,
             folder: str | None = None,
             force: bool = False,
             save_environment: bool = False
             ) -> None:
        '''
        Function to save a trained agent to long term storage.
        By default, the agent is saved in its entirety using pickle.

        However, it is strongly advised to overwrite this method to only save save the necessary components of the agents in order to be able to load it and reproduce its behavior.
        For instance, if the agent is saved after the simulation is run, the state would also be saved within the pickle which is not wanted.

        Parameters
        ----------
        folder : str, optional
            The folder in which the agent's data should be saved.
        force : bool, default=False
            If the agent is already saved at the folder provided, the saving should fail.
            If the already saved agent should be overwritten, this parameter should be toggled to True.
        save_environment : bool, default=False
            Whether to save the agent's linked environment alongside the agent itself.
        '''
        if self.on_gpu:
            cpu_agent = self.to_cpu()
            cpu_agent.save(folder=folder, force=force, save_environment=save_environment)
            return

        # Adding env name to folder path
        if folder is None:
            folder = f'./Agent-{self.name}'
        else:
            folder += f'/Agent-{self.name}'

        # Checking the folder exists or creates it
        if not os.path.exists(folder):
            os.mkdir(folder)
        elif len(os.listdir(folder)) > 0:
            if force:
                shutil.rmtree(folder)
                os.mkdir(folder)
            else:
                raise Exception(f'{folder} is not empty. If you want to overwrite the saved agent, enable "force".')

        # Send self to pickle
        with open(folder + '/binary.pkl', 'wb') as f:
            pickle.dump(self, f)

        # Save environment in folder too if requested
        if save_environment:
            self.environment.save(folder=(folder + f'/Env-{self.environment.name}'))


    @classmethod
    def load(cls,
             folder: str
             ) -> 'Agent':
        '''
        Function to load a trained agent from long term storage.
        By default, as for the save function, it will load the agent from the folder assuming it is a pickle file.

        Parameters
        ----------
        folder : str
            The folder in which the agent was saved.

        Returns
        -------
        loaded_agent : Agent
            The agent loaded from the folder.
        '''
        from olfactory_navigation import agents

        for name, obj in inspect.getmembers(agents):
            if inspect.isclass(obj) and (name in folder) and issubclass(obj, cls) and (obj != cls):
                return obj.load(folder)

        # Default loading with pickle
        with open(folder + '/binary.pkl', 'rb') as f:
            return pickle.load(f)


    def modify_environment(self,
                           new_environment: Environment
                           ) -> 'Agent':
        '''
        Function to modify the environment of the agent.

        Note: By default, a new agent is created with the same thresholds and name but with a this new environment!
        If there are any trained elements to the agent, they are to be modified in this method to be adapted to this new environment.

        Parameters
        ----------
        new_environment : Environment
            The new environment to replace the agent in an equivalent agent.

        Returns
        -------
        modified_agent : Agent
            A new Agent whose environment has been replaced.
        '''
        # TODO: Fix this to account for other init parameters
        modified_agent = self.__class__(environment=new_environment,
                                        thresholds=self.thresholds,
                                        name=self.name)
        return modified_agent


    def to_gpu(self) -> 'Agent':
        '''
        Function to send the numpy arrays of the agent to the gpu.
        It returns a new instance of the Agent class with the arrays on the gpu.

        Returns
        -------
        gpu_agent : Agent
            A new environment instance where the arrays are on the gpu memory.
        '''
        # Check whether the agent is already on the gpu or not
        if self.on_gpu:
            return self

        # Warn and overwrite alternate_version in case it already exists
        if self._alternate_version is not None:
            print('[warning] A GPU instance already existed and is being recreated.')
            self._alternate_version = None

        assert gpu_support, "GPU support is not enabled, Cupy might need to be installed..."

        # Generating a new instance
        cls = self.__class__
        gpu_agent = cls.__new__(cls)

        # Copying arguments to gpu
        for arg, val in self.__dict__.items():
            if isinstance(val, np.ndarray):
                setattr(gpu_agent, arg, cp.array(val))
            elif arg == 'rnd_state':
                setattr(gpu_agent, arg, cp.random.RandomState(self.seed))
            else:
                setattr(gpu_agent, arg, val)

        # Self reference instances
        self._alternate_version = gpu_agent
        gpu_agent._alternate_version = self

        gpu_agent.on_gpu = True
        return gpu_agent


    def to_cpu(self) -> 'Agent':
        '''
        Function to send the numpy arrays of the agent to the cpu.
        It returns a new instance of the Agent class with the arrays on the cpu.

        Returns
        -------
        cpu_agent : Agent
            A new environment instance where the arrays are on the cpu memory.
        '''
        # Check whether the agent is already on the cpu or not
        if not self.on_gpu:
            return self

        if self._alternate_version is not None:
            print('[warning] A CPU instance already existed and is being recreated.')
            self._alternate_version = None

        # Generating a new instance
        cls = self.__class__
        cpu_agent = cls.__new__(cls)

        # Copying arguments to gpu
        for arg, val in self.__dict__.items():
            if isinstance(val, cp.ndarray):
                setattr(cpu_agent, arg, cp.asnumpy(val))
            elif arg == 'rnd_state':
                setattr(cpu_agent, arg, np.random.RandomState(self.seed))
            else:
                setattr(cpu_agent, arg, val)

        # Self reference instances
        self._alternate_version = cpu_agent
        cpu_agent._alternate_version = self

        cpu_agent.on_gpu = True
        return cpu_agent


    @property
    def gpu_version(self):
        '''
        A version of the Agent on the GPU.
        If the agent is already on the GPU it returns itself, otherwise the to_gpu function is called to generate a new one.
        '''
        if self.on_gpu:
            return self
        else:
            if self._alternate_version is not None: # Check if an alternate version already exists
                return self._alternate_version
            else: # Generate an alternate version on the gpu
                return self.to_gpu()


    @property
    def cpu_version(self):
        '''
        A version of the Agent on the CPU.
        If the agent is already on the CPU it returns itself, otherwise the to_cpu function is called to generate a new one.
        '''
        if not self.on_gpu:
            return self
        else:
            if self._alternate_version is not None: # Check if an alternate version already exists
                return self._alternate_version
            else: # Generate an alternate version on the cpu
                return self.to_cpu()

class_name property

The name of the class of the agent.

cpu_version property

A version of the Agent on the CPU. If the agent is already on the CPU it returns itself, otherwise the to_cpu function is called to generate a new one.

gpu_version property

A version of the Agent on the GPU. If the agent is already on the GPU it returns itself, otherwise the to_gpu function is called to generate a new one.

choose_action()

Function to allow for the agent(s) to choose an action to take based on its current state.

It should return a 2D array of shape n by 2 (or 3, or 4 depending of whether the environment has layers and/or a 3rd dimension), where n is how many agents are to choose an action. It should be n 2D vectors of (the layer and) the change in the (z,) y, and x positions.

Returns:

Name	Type	Description
movement_vector	ndarray	An array of n vectors in 2D space of the movement(s) the agent(s) will take.

Source code in olfactory_navigation/agent.py

def choose_action(self) -> np.ndarray:
    '''
    Function to allow for the agent(s) to choose an action to take based on its current state.

    It should return a 2D array of shape n by 2 (or 3, or 4 depending of whether the environment has layers and/or a 3rd dimension),
    where n is how many agents are to choose an action. It should be n 2D vectors of (the layer and) the change in the (z,) y, and x positions.

    Returns
    -------
    movement_vector : np.ndarray
        An array of n vectors in 2D space of the movement(s) the agent(s) will take.
    '''
    raise NotImplementedError('The choose_action function is not implemented, make an agent subclass to implement the method')

discretize_observations(observation, action, source_reached)

Function to convert a set of observations to discrete observation ids. It uses the olfaction thresholds of the agent to discretize the odor concentrations.

In the case where the agent is also aware of it's own position in space, in which case the observation matrix is of size n by 1 + environment.dimensions, the agent first converts the points on the grid to position ids and then multiply the id by olfactory observation id.

Parameters:

Name	Type	Description	Default
observation	ndarray	The observations the agent receives, in the case the agent is space_aware, the position point is also included.	required
action	ndarray	The action the agent did. This parameter is used in the particular case where the environment has layers and the odor thresholds are layer dependent.	required
source_reached	ndarray	A 1D array of boolean values signifying whether each agent reached or not the source.	required

Returns:

Name	Type	Description
discrete_observations	ndarray	An integer array of discrete observations

Source code in olfactory_navigation/agent.py

def discretize_observations(self,
                            observation: np.ndarray,
                            action: np.ndarray,
                            source_reached: np.ndarray
                            ) -> np.ndarray:
    '''
    Function to convert a set of observations to discrete observation ids.
    It uses the olfaction thresholds of the agent to discretize the odor concentrations.

    In the case where the agent is also aware of it's own position in space, in which case the observation matrix is of size n by 1 + environment.dimensions,
    the agent first converts the points on the grid to position ids and then multiply the id by olfactory observation id.

    Parameters
    ----------
    observation : np.ndarray
        The observations the agent receives, in the case the agent is space_aware, the position point is also included.
    action: np.ndarray
        The action the agent did. This parameter is used in the particular case where the environment has layers and the odor thresholds are layer dependent.
    source_reached : np.ndarray
        A 1D array of boolean values signifying whether each agent reached or not the source.

    Returns
    -------
    discrete_observations: np.ndarray
        An integer array of discrete observations
    '''
    # GPU support
    xp = np if not self.on_gpu else cp

    n = len(observation)
    discrete_observations = xp.ones(n, dtype=int) * -1

    # Gather olfactory observation
    olfactory_observation = observation if not self.space_aware else observation[:,0]

    # Compute the amount of observations available
    observation_count = self.thresholds.shape[-1] - 1 # |thresholds| - 1 observation buckets

    # Handle the special case where we have a layered environment and different thresholds for these layers
    if self.environment.has_layers and len(self.thresholds.shape) == 2:
        layer_ids = action[:,0] # First column of actions is the layer
        action_layer_thresholds = self.thresholds[layer_ids]
        observation_ids = xp.argwhere((olfactory_observation[:,None] >= action_layer_thresholds[:,:-1]) & (olfactory_observation[:,None] < action_layer_thresholds[:,1:]))[:,1]
    else:
        # Setting observation ids
        observation_ids = xp.argwhere((olfactory_observation[:,None] >= self.thresholds[:-1][None,:]) & (olfactory_observation[:,None] < self.thresholds[1:][None,:]))[:,1]

    # If needed, multiply observation indices by position indices
    if not self.space_aware:
        discrete_observations = observation_ids
    else:
        position_clipped = xp.clip(observation[:,1:], a_min=0, a_max=(xp.array(self.environment.shape)-1))
        position_count = int(xp.prod(self.spacial_subdivisions))
        position_ids = self._environment_to_subdivision_mapping[*position_clipped.astype(int).T]

        # Add the amount of possible positions to the observation count
        observation_count *= position_count

        # Combine with observation ids
        discrete_observations = (observation_ids * position_count) + position_ids

    # Adding the goal observation
    discrete_observations[source_reached] = observation_count

    return discrete_observations

initialize_state(n=1)

Function to initialize the internal state of the agent(s) for the simulation process. The internal state can be concepts such as the "memory" or "belief" of the agent. The n parameter corresponds to how many "instances" need to instanciated. This is meant so that we work with a "group" of agents instead of individual instances.

This is done with the purpose that the state of the group of agents be stored in (Numpy) arrays to allow vectorization instead of sequential loops.

Parameters:

Name	Type	Description	Default
n	int	How many agents to initialize.	1

Source code in olfactory_navigation/agent.py

def initialize_state(self,
                     n: int = 1
                     ) -> None:
    '''
    Function to initialize the internal state of the agent(s) for the simulation process. The internal state can be concepts such as the "memory" or "belief" of the agent.
    The n parameter corresponds to how many "instances" need to instanciated. This is meant so that we work with a "group" of agents instead of individual instances.

    This is done with the purpose that the state of the group of agents be stored in (Numpy) arrays to allow vectorization instead of sequential loops.

    Parameters
    ----------
    n : int, default=1
        How many agents to initialize.
    '''
    raise NotImplementedError('The initialize_state function is not implemented, make an agent subclass to implement the method')

kill(simulations_to_kill)

Function to kill any agents that either reached the source or failed by not reaching the source before the horizon or failing to update its own state. The agents where the simulations_to_kill paramater is True have to removed from the list of agents. It is necessary because their reference will also be removed from the simulation loop. Therefore, if they are not removed, the array sizes will not match anymore.

Parameters:

Name	Type	Description	Default
simulations_to_kill	ndarray	An array of size n containing boolean values of whether or not agent's simulations are terminated and therefore should be removed.	required

Source code in olfactory_navigation/agent.py

def kill(self,
         simulations_to_kill: np.ndarray
         ) -> None:
    '''
    Function to kill any agents that either reached the source or failed by not reaching the source before the horizon or failing to update its own state.
    The agents where the simulations_to_kill paramater is True have to removed from the list of agents.
    It is necessary because their reference will also be removed from the simulation loop. Therefore, if they are not removed, the array sizes will not match anymore.

    Parameters
    ----------
    simulations_to_kill : np.ndarray
        An array of size n containing boolean values of whether or not agent's simulations are terminated and therefore should be removed.
    '''
    raise NotImplementedError('The kill function is not implemented, make an agent subclass to implement the method')

load(folder) classmethod

Function to load a trained agent from long term storage. By default, as for the save function, it will load the agent from the folder assuming it is a pickle file.

Parameters:

Name	Type	Description	Default
folder	str	The folder in which the agent was saved.	required

Returns:

Name	Type	Description
loaded_agent	Agent	The agent loaded from the folder.

Source code in olfactory_navigation/agent.py

@classmethod
def load(cls,
         folder: str
         ) -> 'Agent':
    '''
    Function to load a trained agent from long term storage.
    By default, as for the save function, it will load the agent from the folder assuming it is a pickle file.

    Parameters
    ----------
    folder : str
        The folder in which the agent was saved.

    Returns
    -------
    loaded_agent : Agent
        The agent loaded from the folder.
    '''
    from olfactory_navigation import agents

    for name, obj in inspect.getmembers(agents):
        if inspect.isclass(obj) and (name in folder) and issubclass(obj, cls) and (obj != cls):
            return obj.load(folder)

    # Default loading with pickle
    with open(folder + '/binary.pkl', 'rb') as f:
        return pickle.load(f)

modify_environment(new_environment)

Function to modify the environment of the agent.

Note: By default, a new agent is created with the same thresholds and name but with a this new environment! If there are any trained elements to the agent, they are to be modified in this method to be adapted to this new environment.

Parameters:

Name	Type	Description	Default
new_environment	Environment	The new environment to replace the agent in an equivalent agent.	required

Returns:

Name	Type	Description
modified_agent	Agent	A new Agent whose environment has been replaced.

Source code in olfactory_navigation/agent.py

def modify_environment(self,
                       new_environment: Environment
                       ) -> 'Agent':
    '''
    Function to modify the environment of the agent.

    Note: By default, a new agent is created with the same thresholds and name but with a this new environment!
    If there are any trained elements to the agent, they are to be modified in this method to be adapted to this new environment.

    Parameters
    ----------
    new_environment : Environment
        The new environment to replace the agent in an equivalent agent.

    Returns
    -------
    modified_agent : Agent
        A new Agent whose environment has been replaced.
    '''
    # TODO: Fix this to account for other init parameters
    modified_agent = self.__class__(environment=new_environment,
                                    thresholds=self.thresholds,
                                    name=self.name)
    return modified_agent

save(folder=None, force=False, save_environment=False)

Function to save a trained agent to long term storage. By default, the agent is saved in its entirety using pickle.

However, it is strongly advised to overwrite this method to only save save the necessary components of the agents in order to be able to load it and reproduce its behavior. For instance, if the agent is saved after the simulation is run, the state would also be saved within the pickle which is not wanted.

Parameters:

Name	Type	Description	Default
folder	str	The folder in which the agent's data should be saved.	None
force	bool	If the agent is already saved at the folder provided, the saving should fail. If the already saved agent should be overwritten, this parameter should be toggled to True.	False
save_environment	bool	Whether to save the agent's linked environment alongside the agent itself.	False

Source code in olfactory_navigation/agent.py

def save(self,
         folder: str | None = None,
         force: bool = False,
         save_environment: bool = False
         ) -> None:
    '''
    Function to save a trained agent to long term storage.
    By default, the agent is saved in its entirety using pickle.

    However, it is strongly advised to overwrite this method to only save save the necessary components of the agents in order to be able to load it and reproduce its behavior.
    For instance, if the agent is saved after the simulation is run, the state would also be saved within the pickle which is not wanted.

    Parameters
    ----------
    folder : str, optional
        The folder in which the agent's data should be saved.
    force : bool, default=False
        If the agent is already saved at the folder provided, the saving should fail.
        If the already saved agent should be overwritten, this parameter should be toggled to True.
    save_environment : bool, default=False
        Whether to save the agent's linked environment alongside the agent itself.
    '''
    if self.on_gpu:
        cpu_agent = self.to_cpu()
        cpu_agent.save(folder=folder, force=force, save_environment=save_environment)
        return

    # Adding env name to folder path
    if folder is None:
        folder = f'./Agent-{self.name}'
    else:
        folder += f'/Agent-{self.name}'

    # Checking the folder exists or creates it
    if not os.path.exists(folder):
        os.mkdir(folder)
    elif len(os.listdir(folder)) > 0:
        if force:
            shutil.rmtree(folder)
            os.mkdir(folder)
        else:
            raise Exception(f'{folder} is not empty. If you want to overwrite the saved agent, enable "force".')

    # Send self to pickle
    with open(folder + '/binary.pkl', 'wb') as f:
        pickle.dump(self, f)

    # Save environment in folder too if requested
    if save_environment:
        self.environment.save(folder=(folder + f'/Env-{self.environment.name}'))

to_cpu()

Function to send the numpy arrays of the agent to the cpu. It returns a new instance of the Agent class with the arrays on the cpu.

Returns:

Name	Type	Description
cpu_agent	Agent	A new environment instance where the arrays are on the cpu memory.

Source code in olfactory_navigation/agent.py

def to_cpu(self) -> 'Agent':
    '''
    Function to send the numpy arrays of the agent to the cpu.
    It returns a new instance of the Agent class with the arrays on the cpu.

    Returns
    -------
    cpu_agent : Agent
        A new environment instance where the arrays are on the cpu memory.
    '''
    # Check whether the agent is already on the cpu or not
    if not self.on_gpu:
        return self

    if self._alternate_version is not None:
        print('[warning] A CPU instance already existed and is being recreated.')
        self._alternate_version = None

    # Generating a new instance
    cls = self.__class__
    cpu_agent = cls.__new__(cls)

    # Copying arguments to gpu
    for arg, val in self.__dict__.items():
        if isinstance(val, cp.ndarray):
            setattr(cpu_agent, arg, cp.asnumpy(val))
        elif arg == 'rnd_state':
            setattr(cpu_agent, arg, np.random.RandomState(self.seed))
        else:
            setattr(cpu_agent, arg, val)

    # Self reference instances
    self._alternate_version = cpu_agent
    cpu_agent._alternate_version = self

    cpu_agent.on_gpu = True
    return cpu_agent

to_gpu()

Function to send the numpy arrays of the agent to the gpu. It returns a new instance of the Agent class with the arrays on the gpu.

Returns:

Name	Type	Description
gpu_agent	Agent	A new environment instance where the arrays are on the gpu memory.

Source code in olfactory_navigation/agent.py

def to_gpu(self) -> 'Agent':
    '''
    Function to send the numpy arrays of the agent to the gpu.
    It returns a new instance of the Agent class with the arrays on the gpu.

    Returns
    -------
    gpu_agent : Agent
        A new environment instance where the arrays are on the gpu memory.
    '''
    # Check whether the agent is already on the gpu or not
    if self.on_gpu:
        return self

    # Warn and overwrite alternate_version in case it already exists
    if self._alternate_version is not None:
        print('[warning] A GPU instance already existed and is being recreated.')
        self._alternate_version = None

    assert gpu_support, "GPU support is not enabled, Cupy might need to be installed..."

    # Generating a new instance
    cls = self.__class__
    gpu_agent = cls.__new__(cls)

    # Copying arguments to gpu
    for arg, val in self.__dict__.items():
        if isinstance(val, np.ndarray):
            setattr(gpu_agent, arg, cp.array(val))
        elif arg == 'rnd_state':
            setattr(gpu_agent, arg, cp.random.RandomState(self.seed))
        else:
            setattr(gpu_agent, arg, val)

    # Self reference instances
    self._alternate_version = gpu_agent
    gpu_agent._alternate_version = self

    gpu_agent.on_gpu = True
    return gpu_agent

train()

Optional function to train the agent in the olfactory environment it is in. This function is optional as some agents have some fixed behavior and therefore dont require training.

Source code in olfactory_navigation/agent.py

def train(self) -> None:
    '''
    Optional function to train the agent in the olfactory environment it is in.
    This function is optional as some agents have some fixed behavior and therefore dont require training.
    '''
    raise NotImplementedError('The train function is not implemented, make an agent subclass to implement the method')

update_state(action, observation, source_reached)

Function to update the internal state(s) of the agent(s) based on the action(s) taken and the observation(s) received. The observations are then compared with the thresholds to decide whether something was sensed or not or to what level.

Parameters:

Name	Type	Description	Default
action	ndarray	A 2D array of n movement vectors. If the environment is layered, the 1st component should be the layer.	required
observation	ndarray	A n by 1 (or 1 + environment.dimensions if space_aware) array of odor cues (float values) retrieved from the environment.	required
source_reached	array	A 1D array of boolean values signifying whether each agent reached or not the source.	required

Returns:

Name	Type	Description
update_successfull	(ndarray, optional)	If nothing is returned, it means all the agent's state updates have been successfull. Else, a boolean np.ndarray of size n can be returned confirming for each agent whether the update has been successful or not.

Source code in olfactory_navigation/agent.py

def update_state(self,
                 action: np.ndarray,
                 observation: np.ndarray,
                 source_reached: np.ndarray
                 ) -> None | np.ndarray:
    '''
    Function to update the internal state(s) of the agent(s) based on the action(s) taken and the observation(s) received.
    The observations are then compared with the thresholds to decide whether something was sensed or not or to what level.

    Parameters
    ----------
    action : np.ndarray
        A 2D array of n movement vectors. If the environment is layered, the 1st component should be the layer.
    observation : np.ndarray
        A n by 1 (or 1 + environment.dimensions if space_aware) array of odor cues (float values) retrieved from the environment.
    source_reached : np.array
        A 1D array of boolean values signifying whether each agent reached or not the source.

    Returns
    -------
    update_successfull : np.ndarray, optional
        If nothing is returned, it means all the agent's state updates have been successfull.
        Else, a boolean np.ndarray of size n can be returned confirming for each agent whether the update has been successful or not.
    '''
    raise NotImplementedError('The update_state function is not implemented, make an agent subclass to implement the method')

Environment

Class to represent an olfactory environment.

It is defined based on an olfactory data set provided as either a numpy file or an array directly with shape time, y, x. From this environment, the various parameters are applied in the following order:

1. The source position is set
2. The margins are added and the shape (total size) of the environment are set.
3. The data file's x and y components are squished and streched the to fit the inter-marginal shape of the environment.
4. The source's position is also moved to stay at the same position within the data.
5. The multiplier is finally applied to modify the data file's x and y components a final time by growing or shrinking the margins to account for the multiplier. (The multiplication applies with the source position as a center point)

Note: to modify the shape of the data file's x and y components the OpenCV library's resize function is used. And the interpolation method is controlled by the interpolation_method parameter.

Then, the starting probability map is built. Either an array can be provided directly or preset option can be chosen:

• 'data_zone': The agent can start at any point in the data_zone (after all the modification parameters have been applied)
• 'odor_present': The agent can start at any point where an odor cue above the odor_present_threshold can be found at any timestep during the simulation

Parameters:

Name	Type	Description	Default
data_file	str or ndarray	The dataset containing the olfactory data. It can be provided as a path to a file containing said array.	required
data_source_position	list or ndarray	The center point of the source provided as a list or a 1D array with the components being x,y. This position is computed in the olfactory data zone (so excluding the margins).	required
source_radius	float	The radius from the center point of the source in which we consider the agent has reached the source.	1.0
layers	bool or list[int] or list[str]	Whether or not the data provided contains layers or not. If a list of strings is provided, it will be either used to name the layers found (if numpy data), or it is used to querry the datasets of the h5 file.	False
shape	list or ndarray	A 2-element array or list of how many units should be kept in the final array (including the margins). As it should include the margins, the shape should be strictly larger than the sum of the margins in each direction. By default, the shape of the olfactory data will be maintained.	None
margins	int or list or ndarray	How many units have to be added to the data as margins. (Before the multiplier is applied) If a unique element is provided, the margin will be this same value on each side. If a list or array of 2 elements is provided, the first number will be vertical margins (y-axis), while the other will be on the x-axis (horizontal).	0
multiplier	list or ndarray	A 2-element array or list of how much the odor field should be streched in each direction. If a value larger than 1 is provided, the margins will be reduced to accomodate for the larger size of the olfactory data size. And inversly, less than 1 will increase the margins. By default, the multipliers will be set to 1.0.	[1.0,1.0]
interpolation_method	Nearest or Linear or Cubic	The interpolation method to be used in the case the data needs to be reshaped to fit the shape, margins and multiplier parameters. By default, it uses Bi-linear interpolation. The interpolation is performed using the OpenCV library.	'Linear'
preprocess_data	bool	Applicable only for data_file being a path to a h5 file. Whether to reshape of the data at the creation of the environment. Reshaping the data ahead of time will require more processing at the creation and more memory overall. While if this is disabled, when gathering observations, more time will be required but less memory will need to be used at once.	False
boundary_condition	stop or wrap or wrap_vertical or wrap_horizontal or clip	How the agent should behave at the boundary. Stop means for the agent to stop at the boundary, if the agent tries to move north while being on the top edge, it will stay in the same state. Wrap means for the borders to be like portals, when entering on one side, it reappears on the other side. Wrap can be specified to be only vertically or horizontally	'stop'
start_zone	odor_present or data_zone or ndarray	Either an array or a string representing how the starting probabilities should be constructed. - odor_present: The start probabilities will be uniform where odor cues can be found above 0 (or a given odor_present_threshold) - data_zone: Uniform over the data zone, so without the margins. Note that the points within the source radius will be excluded from this probability grid.	'data_zone'
odor_present_threshold	float	An olfactory threshold, under which the odor is considered too low to be noticed. It is used only to build the starting zone if the 'odor_present' option is selected.	None
name	str	A custom name to be given to the agent. If it is not provided, by default it will have the format: -marg_-edge_-start_-source__radius	None
seed	int	For reproducible randomness.	12131415

Attributes:

Name	Type	Description
data	ndarray	An array containing the olfactory data after the modification parameters have been applied.
data_file_path	str	If the data is loaded from a path, the path will be recorded here.
data_source_position	ndarray	The position of the source in the original data file (after modifications have been applied).
layers	ndarray	A numbered list of the IDs of the layers.
layer_labels	list[str]	A list of how the layers are named.
has_layers	bool	Whether or not the environment is made up of layers.
margins	ndarray	An array of the margins vertically and horizontally (after multiplier is applied).
timestamps	int	The amount of timeslices available in the environment.
data_shape	tuple[int]	The shape of the data's odor field (after modifications have been applied).
dimensions	int	The amount of dimensions of the physical space of the olfactory environment.
shape	tuple[int]	The shape of the environment. It is a tuple of the size in each axis of the environment.
data_bounds	ndarray	The bounds between which the original olfactory data stands in the coordinate system of the environment (after modifications have been applied).
source_position	ndarray	The position of the source in the padded grid (after modifications have been applied).
source_radius	float	The radius of the source.
interpolation_method	str	The interpolation used to modify the shape of the original data.
data_processed	bool	Whether the data was processed (ie the shape is at it should be) or not.
boundary_condition	str	How the agent should behave when reaching the boundary.
start_probabilities	ndarray	A probability map of where the agent is likely to start within the environment. Note: Zero within the source radius.
start_type	str	The type of the start probability map building. For instance: 'data_zone', 'odor_present', or 'custom' (if an array is provided).
odor_present_threshold	float	The threshold used to uild the start probabilities if the option 'odor_present' is used.
name	str	The name set to the agent as defined in the parameters.
saved_at	str	If the environment is saved, the path at which it is saved will be recorded here.
on_gpu	bool	Whether the environment's arrays are on the gpu's memory or not.
seed	int	The seed used for the random operations (to allow for reproducability).
rnd_state	RandomState	The random state variable used to generate random values.
cpu_version	Environment	An instance of the environment on the CPU. If it already is, it returns itself.
gpu_version	Environment	An instance of the environment on the CPU. If it already is, it returns itself.

Source code in olfactory_navigation/environment.py

class Environment:
    '''
    Class to represent an olfactory environment.

    It is defined based on an olfactory data set provided as either a numpy file or an array directly with shape time, y, x.
    From this environment, the various parameters are applied in the following order:

    0. The source position is set
    1. The margins are added and the shape (total size) of the environment are set.
    2. The data file's x and y components are squished and streched the to fit the inter-marginal shape of the environment.
    3. The source's position is also moved to stay at the same position within the data.
    4. The multiplier is finally applied to modify the data file's x and y components a final time by growing or shrinking the margins to account for the multiplier. (The multiplication applies with the source position as a center point)

    Note: to modify the shape of the data file's x and y components the OpenCV library's resize function is used. And the interpolation method is controlled by the interpolation_method parameter.


    Then, the starting probability map is built. Either an array can be provided directly or preset option can be chosen:

    - 'data_zone': The agent can start at any point in the data_zone (after all the modification parameters have been applied)
    - 'odor_present': The agent can start at any point where an odor cue above the odor_present_threshold can be found at any timestep during the simulation

    Parameters
    ----------
    data_file : str or np.ndarray
        The dataset containing the olfactory data. It can be provided as a path to a file containing said array.
    data_source_position : list or np.ndarray
        The center point of the source provided as a list or a 1D array with the components being x,y.
        This position is computed in the olfactory data zone (so excluding the margins).
    source_radius : float, default=1.0
        The radius from the center point of the source in which we consider the agent has reached the source.
    layers : bool or list[int] or list[str], default=False
        Whether or not the data provided contains layers or not.
        If a list of strings is provided, it will be either used to name the layers found (if numpy data), or it is used to querry the datasets of the h5 file.
    shape : list or np.ndarray, optional
        A 2-element array or list of how many units should be kept in the final array (including the margins).
        As it should include the margins, the shape should be strictly larger than the sum of the margins in each direction.
        By default, the shape of the olfactory data will be maintained.
    margins : int or list or np.ndarray, default=0
        How many units have to be added to the data as margins. (Before the multiplier is applied)
        If a unique element is provided, the margin will be this same value on each side.
        If a list or array of 2 elements is provided, the first number will be vertical margins (y-axis), while the other will be on the x-axis (horizontal).
    multiplier : list or np.ndarray, default=[1.0,1.0]
        A 2-element array or list of how much the odor field should be streched in each direction.
        If a value larger than 1 is provided, the margins will be reduced to accomodate for the larger size of the olfactory data size.
        And inversly, less than 1 will increase the margins.
        By default, the multipliers will be set to 1.0.
    interpolation_method : 'Nearest' or 'Linear' or 'Cubic', default='Linear'
        The interpolation method to be used in the case the data needs to be reshaped to fit the shape, margins and multiplier parameters.
        By default, it uses Bi-linear interpolation. The interpolation is performed using the OpenCV library.
    preprocess_data : bool, default=False
        Applicable only for data_file being a path to a h5 file.
        Whether to reshape of the data at the creation of the environment.
        Reshaping the data ahead of time will require more processing at the creation and more memory overall.
        While if this is disabled, when gathering observations, more time will be required but less memory will need to be used at once.
    boundary_condition : 'stop' or 'wrap' or 'wrap_vertical' or 'wrap_horizontal' or 'clip', default='stop'
        How the agent should behave at the boundary.
        Stop means for the agent to stop at the boundary, if the agent tries to move north while being on the top edge, it will stay in the same state.
        Wrap means for the borders to be like portals, when entering on one side, it reappears on the other side.
        Wrap can be specified to be only vertically or horizontally
    start_zone : 'odor_present' or 'data_zone' or np.ndarray, default='data_zone'
        Either an array or a string representing how the starting probabilities should be constructed.
        - odor_present: The start probabilities will be uniform where odor cues can be found above 0 (or a given odor_present_threshold)
        - data_zone: Uniform over the data zone, so without the margins.
        Note that the points within the source radius will be excluded from this probability grid.
    odor_present_threshold : float, optional
        An olfactory threshold, under which the odor is considered too low to be noticed.
        It is used only to build the starting zone if the 'odor_present' option is selected.
    name : str, optional
        A custom name to be given to the agent.
        If it is not provided, by default it will have the format:
        <shape>-marg_<margins>-edge_<boundary_condition>-start_<start_zone>-source_<source_point>_radius<source_radius>
    seed : int, default=12131415
        For reproducible randomness.

    Attributes
    ----------
    data : np.ndarray
        An array containing the olfactory data after the modification parameters have been applied.
    data_file_path : str
        If the data is loaded from a path, the path will be recorded here.
    data_source_position : np.ndarray
        The position of the source in the original data file (after modifications have been applied).
    layers : np.ndarray
        A numbered list of the IDs of the layers.
    layer_labels : list[str]
        A list of how the layers are named.
    has_layers : bool
        Whether or not the environment is made up of layers.
    margins : np.ndarray
        An array of the margins vertically and horizontally (after multiplier is applied).
    timestamps : int
        The amount of timeslices available in the environment.
    data_shape : tuple[int]
        The shape of the data's odor field (after modifications have been applied).
    dimensions : int
        The amount of dimensions of the physical space of the olfactory environment.
    shape : tuple[int]
        The shape of the environment. It is a tuple of the size in each axis of the environment.
    data_bounds : np.ndarray
        The bounds between which the original olfactory data stands in the coordinate system of the environment (after modifications have been applied).
    source_position : np.ndarray
        The position of the source in the padded grid (after modifications have been applied).
    source_radius : float
        The radius of the source.
    interpolation_method : str
        The interpolation used to modify the shape of the original data.
    data_processed : bool
        Whether the data was processed (ie the shape is at it should be) or not.
    boundary_condition : str
        How the agent should behave when reaching the boundary.
    start_probabilities : np.ndarray
        A probability map of where the agent is likely to start within the environment.
        Note: Zero within the source radius.
    start_type : str
        The type of the start probability map building. For instance: 'data_zone', 'odor_present', or 'custom' (if an array is provided).
    odor_present_threshold : float
        The threshold used to uild the start probabilities if the option 'odor_present' is used.
    name : str
        The name set to the agent as defined in the parameters.
    saved_at : str
        If the environment is saved, the path at which it is saved will be recorded here.
    on_gpu : bool
        Whether the environment's arrays are on the gpu's memory or not.
    seed : int
        The seed used for the random operations (to allow for reproducability).
    rnd_state : np.random.RandomState
        The random state variable used to generate random values.
    cpu_version : Environment
        An instance of the environment on the CPU. If it already is, it returns itself.
    gpu_version : Environment
        An instance of the environment on the CPU. If it already is, it returns itself.
    '''
    def __init__(self,
                 data_file: str | np.ndarray,
                 data_source_position: list | np.ndarray,
                 source_radius: float = 1.0,
                 layers: bool | list[str] = False,
                 shape: list | np.ndarray | None = None,
                 margins: int | list | np.ndarray = 0,
                 multiplier: list| np.ndarray = [1.0, 1.0],
                 interpolation_method: Literal['Nearest', 'Linear', 'Cubic'] = 'Linear',
                 preprocess_data: bool = False,
                 boundary_condition: Literal['stop', 'wrap', 'wrap_vertical', 'wrap_horizontal', 'clip', 'no'] = 'stop',
                 start_zone: Literal['odor_present', 'data_zone'] | np.ndarray = 'data_zone',
                 odor_present_threshold: float | None = None,
                 name: str | None = None,
                 seed: int = 12131415,
                 ) -> None:
        self.saved_at: str = None

        # Layer properties
        self.layers = None
        self.layer_labels = None
        self.has_layers = False

        if isinstance(layers, list):
            self.has_layers = True
            self.layers = np.arange(len(layers))
            self.layer_labels = [layer for layer in layers]
        elif isinstance(layers, bool):
            self.has_layers = layers

        # Load from file if string provided
        self.data_file_path = None
        self._preprocess_data: bool = preprocess_data

        loaded_data = None
        if isinstance(data_file, str):
            self.data_file_path = data_file

            # NUMPY
            if data_file.endswith('.npy'):
                loaded_data = np.load(data_file)

                # Layered data
                if self.has_layers:
                    if self.layers is None:
                        self.layers = np.arange(len(loaded_data))
                        self.layer_labels = [str(layer) for layer in range(len(loaded_data))]
                    else:
                        assert (len(self.layers) == len(loaded_data)), "The amount of layers provided dont match the amount in the dataset."

                        # Re-ordering the layers
                        loaded_data = loaded_data[self.layers]

            # H5
            elif data_file.endswith('.h5'):
                loaded_data = h5py.File(data_file,'r')

                # Layered data
                if self.has_layers:

                    # Converting layers to strings
                    data_layer_labels = list(loaded_data.keys())
                    if self.layers is None:
                        self.layers = np.arange(len(data_layer_labels))
                        self.layer_labels = data_layer_labels

                    # Getting the labels based on the list of integers provided
                    elif all(isinstance(layer, int) for layer in layers):
                        self.layer_labels = [data_layer_labels[layer_id] for layer_id in self.layers]

                    # Loading the list of slices from the data
                    loaded_data = [[loaded_data[layer][f"{t}"] for t in range(len(loaded_data[layer]))] for layer in self.layer_labels]

                else:
                    loaded_data = [loaded_data[f"{t}"] for t in range(len(loaded_data))]

            # Not supported
            else:
                raise NotImplementedError('File format loading not implemented')

        elif not isinstance(data_file, np.ndarray):
            raise NotImplementedError("Data file should be either a path or an object that is either an h5 object or a numpy array")

        self._data: np.ndarray = loaded_data if loaded_data is not None else data_file

        # Unmodified sizes
        self.timesteps = len(self._data if not self.has_layers else self._data[0])
        self.data_shape = (self._data[0] if not self.has_layers else self._data[0][0]).shape
        self.dimensions = len(self.data_shape)
        self.data_source_position = np.array(data_source_position)
        self.original_data_source_position = self.data_source_position

        original_data_shape = self.data_shape

        # Making margins a |dims|x2 array
        if isinstance(margins, int):
            self.margins = np.ones((self.dimensions, 2), dtype=int) * margins
        elif isinstance(margins, list) or isinstance(margins, np.ndarray):
            margins = np.array(margins)
            if margins.shape == (self.dimensions,): # Symmetric min and max margins
                self.margins = np.hstack((margins[:,None], margins[:,None]))
            elif margins.shape == (self.dimensions,2):
                self.margins = margins
            else:
                raise ValueError('The array or lists of Margins provided have a shape not supported. (Supported formats (2,) or (2,2))')
        else:
            raise ValueError('margins argument should be either an integer or a 1D or 2D array with either shape (2) or (2,2)')
        assert (self.margins.dtype == int), 'margins should be integers'

        # Process shape parameter
        new_data_shape = None
        if shape is not None:
            shape = np.array(shape)

            assert np.all(shape > np.sum(self.margins, axis=1)), "The shape of the environment must be strictly larger than the sum of margins."

            # Computing the new shape of the data
            new_data_shape: np.ndarray = (shape - np.sum(self.margins, axis=1)).astype(int)

            # New source position
            self.data_source_position = (self.data_source_position * (new_data_shape / self.data_shape)).astype(int)
        else:
            shape = self.data_shape + np.sum(self.margins, axis=1)

        if new_data_shape is not None:
            self.data_shape = (*new_data_shape,)

        # Process multiplier
        multiplier = np.array(multiplier)

        # Assert multiplier value is correct
        with np.errstate(divide='ignore'):
            low_max_mult = ((self.margins[:,0] / self.data_source_position) + 1)
            high_max_mult = (1 + (self.margins[:,1] / (self.data_shape - self.data_source_position)))
            max_mult = np.min(np.vstack([low_max_mult, high_max_mult]), axis=0)

            assert np.all(multiplier <= max_mult), f"The multiplier given is larger than allowed (the values should be lower than {max_mult})"

        # Compute new data shape with the multiplier
        if new_data_shape is None:
            new_data_shape = self.data_shape
        new_data_shape = (new_data_shape * multiplier).astype(int)

        # New source position based on multiplier
        new_source_position = (self.data_source_position * multiplier).astype(int)

        # Recomputing margins with new source position
        self.margins[:,0] -= (new_source_position - self.data_source_position)
        self.margins[:,1] = (shape - (self.margins[:,0] + new_data_shape))

        # Re-Setting new source position
        self.data_source_position = new_source_position

        # Interpolation method choice
        self.interpolation_method = interpolation_method

        # Input the new shape of the data if set by custom shape or multiplier
        if new_data_shape is not None:
            self.data_shape: tuple[int] = (*new_data_shape,)

        # Check if data is already processed by default
        self.data_processed = (self.data_shape == original_data_shape)

        # If requested process all the slices of data into a single
        if preprocess_data and not self.data_processed:
            assert self.dimensions == 2, "Higher dimensional data doesnt support reshaping yet, ensure it is done beforehand.."
            if self.has_layers:
                new_data = np.zeros((len(self.layers), self.timesteps, *self.data_shape))
                for layer in self.layers:
                    for i in range(self.timesteps):
                        new_data[layer, i] = _resize_array(np.array(self._data[layer][i]),
                                                           new_shape=self.data_shape,
                                                           interpolation=self.interpolation_method.lower())
            else:
                new_data = np.zeros((self.timesteps, *self.data_shape))
                for i in range(self.timesteps):
                    new_data[i] = _resize_array(np.array(self._data[i]),
                                                new_shape=self.data_shape,
                                                interpolation=self.interpolation_method.lower())
            self._data = new_data
            self.data_processed = True

        # Reading shape of data array
        self.shape = (*(self.data_shape + np.sum(self.margins, axis=1)),)

        # Converting the shape tuple to integer sets
        self.shape: tuple[int] = tuple([int(el) for el in self.shape])
        self.data_shape: tuple[int] = tuple([int(el) for el in self.data_shape])

        # Building a data bounds
        self.data_bounds = np.array([self.margins[:,0], self.margins[:,0] + np.array(self.data_shape)]).T

        # Saving arguments
        self.source_position = self.data_source_position + self.margins[:,0]
        self.source_radius = source_radius

        # Boundary conditions
        assert not ((self.dimensions > 2) and (boundary_condition in ['wrap_vertical', 'wrap_horizontal'])), "There are more than 2 dimensions, the options of 'wrap_horizontal' and 'wrap_vertical' are disabled."
        self.boundary_condition = boundary_condition

        # Starting zone
        self.start_probabilities = np.zeros(self.shape)
        self.start_type = start_zone if isinstance(start_zone, str) else 'custom'

        if isinstance(start_zone, np.ndarray):
            if start_zone.shape == (self.dimensions,2):
                slices = tuple(slice(low, high) for low, high in start_zone)
                self.start_probabilities[slices] = 1.0
                self.start_type += '_' + '_'.join([str(el) for el in start_zone.ravel()])
            elif start_zone.shape == self.shape:
                self.start_probabilities = start_zone
            else:
                raise ValueError('If an np.ndarray is provided for the start_zone it has to be |dim| x 2...')

        elif start_zone == 'data_zone':
            slices = tuple(slice(low, high) for low, high in self.data_bounds)
            self.start_probabilities[slices] = 1.0

        elif start_zone == 'odor_present':
            if self.data_processed and isinstance(self._data, np.ndarray):
                odor_present_map = (np.mean((self._data > (odor_present_threshold if odor_present_threshold is not None else 0)).astype(int), axis=0) > 0).astype(float)
                self.start_probabilities[tuple(slice(low, high) for low, high in self.data_bounds)] = odor_present_map
            else:
                odor_sum = np.zeros(self.data_shape, dtype=float)
                for i in range(self.timesteps):
                    data_slice = np.array(self._data[i]) if not self.has_layers else np.array(self._data[0][i])
                    reshaped_data_slice = _resize_array(data_slice,
                                                        new_shape=self.data_shape,
                                                        interpolation=self.interpolation_method.lower())
                    odor_sum += (reshaped_data_slice > (odor_present_threshold if odor_present_threshold is not None else 0))
                self.start_probabilities[tuple(slice(low, high) for low, high in self.data_bounds)] = (odor_sum / self.timesteps)
        else:
            raise ValueError('start_zone value is wrong')

        # Odor present tresh
        self.odor_present_threshold = odor_present_threshold

        # Removing the source area from the starting zone
        source_mask = np.fromfunction((lambda *points: np.sum((np.array(points).transpose([i+1 for i in range(len(self.shape))] + [0]) - self.source_position[None,:])**2, axis=-1) <= self.source_radius**2), shape=self.shape)
        self.start_probabilities[source_mask] = 0
        self.start_probabilities /= np.sum(self.start_probabilities) # Normalization

        # Name
        self.name = name
        if self.name is None:
            self.name =  '_'.join([str(axis_size) for axis_size in self.shape]) # Size of env
            self.name += f'-marg_' + '_'.join(['_'.join([str(marg) for marg in dim_margins]) for dim_margins in self.margins]) # margins
            self.name += f'-edge_{self.boundary_condition}' # Boundary condition
            self.name += f'-start_{self.start_type}' # Start zone
            self.name += f'-source_' + '_'.join([str(pos) for pos in self.source_position]) + f'_radius{self.source_radius}' # Source

        # gpu support
        self._alternate_version = None
        self.on_gpu = False

        # random state
        self.seed = seed
        self.rnd_state = np.random.RandomState(seed = seed)


    @property
    def data(self) -> np.ndarray:
        '''
        The whole dataset with the right shape. If not preprocessed to modify its shape the data will be processed when querrying this object.
        '''
        if not self._data_is_numpy or not self.data_processed:
            xp = cp if self.on_gpu else np
            print('[Warning] The whole dataset is being querried, it will be reshaped at this time. To avoid this, avoid querrying environment.data directly.')

            # Reshaping
            if self.has_layers:
                new_data = np.zeros((len(self.layers), self.timesteps, *self.data_shape))
                for layer in self.layers:
                    for i in range(self.timesteps):
                        new_data[layer, i] = _resize_array(np.array(self._data[layer][i]),
                                                           new_shape=self.data_shape,
                                                           interpolation=self.interpolation_method.lower())
            else:
                new_data = np.zeros((self.timesteps, *self.data_shape))
                for i in range(self.timesteps):
                    new_data[i] = _resize_array(np.array(self._data[i]),
                                                new_shape=self.data_shape,
                                                interpolation=self.interpolation_method.lower())

            self._data = xp.array(new_data)
            self.data_processed = True

        return self._data


    @property
    def _data_is_numpy(self) -> bool:
        '''
        Wheter the data is a numpy array or not.
        '''
        xp = cp if self.on_gpu else np
        return isinstance(self._data, xp.ndarray)


    def plot(self,
             frame: int = 0,
             layer: int = 0,
             ax: plt.Axes | None = None
             ) -> None:
        '''
        Simple function to plot the environment with a single frame of odor cues.
        The starting zone is also market down with a blue contour.
        The source of the odor is marked by a red circle.

        Parameters
        ----------
        frame : int, default=0
            The frame of odor cues to print.
        layer : int, default=0
            The layer of the odor cues to print. (Ignored if the environment is not layered.)
        ax : plt.Axes, optional
            An ax on which the environment can be plot.
        '''
        # If on GPU use the CPU version to plot
        if self.on_gpu:
            self._alternate_version.plot(
                frame=frame,
                ax=ax
            )
            return # Blank return

        # TODO: Implement plotting for 3D
        assert self.dimensions == 2, "Plotting function only available for 2D environments for now..."

        if ax is None:
            _, ax = plt.subplots(1, figsize=(15,5))

        legend_elements = [[],[]]

        # Gather data frame
        data_frame: np.ndarray = self._data[layer][frame] if self.has_layers else self._data[frame]
        if not isinstance(data_frame, np.ndarray):
            data_frame = np.array(data_frame)

        if not self.data_processed:
            data_frame = _resize_array(data_frame,
                                       new_shape=self.data_shape,
                                       interpolation=self.interpolation_method.lower())

        # Odor grid
        odor = Rectangle([0,0], 1, 1, color='black', fill=True)
        frame_data = (data_frame > (self.odor_present_threshold if self.odor_present_threshold is not None else 0)).astype(float)
        environment_frame = np.zeros(self.shape, dtype=float)
        environment_frame[self.data_bounds[0,0]:self.data_bounds[0,1], self.data_bounds[1,0]:self.data_bounds[1,1]] = frame_data
        ax.imshow(environment_frame, cmap='Greys')

        legend_elements[0].append(odor)
        legend_elements[1].append(f'Frame {frame}' + ('' if not self.has_layers else f' (layer {layer})') + ' odor cues')

        # Start zone contour
        start_zone = Rectangle([0,0], 1, 1, color='blue', fill=False)
        ax.contour(self.start_probabilities, levels=[0.0], colors='blue')

        legend_elements[0].append(start_zone)
        legend_elements[1].append('Start zone')

        # Source circle
        goal_circle = Circle(self.source_position[::-1], self.source_radius, color='r', fill=False, zorder=10)
        legend_elements[0].append(goal_circle)
        legend_elements[1].append('Source')

        if self.source_radius > 0.0:
            ax.add_patch(goal_circle)
        else:
            ax.scatter(self.source_position[1], self.source_position[0], c='red')

        # Legend
        ax.legend(legend_elements[0], legend_elements[1])


    def get_observation(self,
                        pos: np.ndarray,
                        time: int | np.ndarray = 0,
                        layer: int | np.ndarray = 0
                        ) -> float | np.ndarray:
        '''
        Function to get an observation at a given position on the grid at a given time.
        A set of observations can also be requested, either at a single position for multiple timestamps or with the same amoung of positions as timestamps provided.

        Note: The position will not be checked against boundary conditions; if a position is out-of-bounds it will simply return 0.0!

        Parameters
        ----------
        pos : np.ndarray
            The position or list of positions to get observations at.
        time : int or np.ndarray, default=0
            A timestamp or list of timestamps to get the observations at.
        layer : int or np.ndarray, default=0
            A layer or list of timestamps to get the observations at.
            Note: If the environment doesnt have layers, this parameter will be ignored.

        Returns
        -------
        observation : float or np.ndarray
            A single observation or list of observations.
        '''
        xp = cp if self.on_gpu else np

        # Handling the case of a single point
        is_single_point = (len(pos.shape) == 1)
        if is_single_point:
            pos = pos[None,:]

        # Counting how many position points we are dealing with
        pos_count = len(pos)

        # Time looping
        time = time % self.timesteps

        # Determine unique layers and reindexing them if needed
        unique_layers = xp.array([layer]) if isinstance(layer, int) else xp.unique(layer)
        layer = 0 if isinstance(layer, int) else xp.where(layer == unique_layers[:,None])[0]
        layer_count = len(unique_layers)

        # Determine unique times and reindexing them if needed
        unique_times = xp.array([time]) if isinstance(time, int) else xp.unique(time)
        time = 0 if isinstance(time, int) else xp.where(time == unique_times[:,None])[0]
        time_count = len(unique_times)

        # Handling the case where the data is a sequence of slices (h5, so not numpy array)
        data = self._data

        # Selecting the required slices
        if self._data_is_numpy:
            data = data[unique_layers][:,unique_times] if self.has_layers else data[unique_times]
        else:
            # Case where we are dealing with a h5 file
            # Note: Can't use self.data_shape because we don't know whether the data is processed yet or no
            selected_slices = xp.zeros((layer_count, time_count, *self._data[0][0].shape)) if self.has_layers else xp.zeros((time_count, *self._data[0].shape))
            for i, t in enumerate(unique_times):
                if self.has_layers:
                    for j, l in enumerate(unique_layers):
                        selected_slices[j,i] = xp.array(data[int(l)][int(t)])
                else:
                    selected_slices[i] = xp.array(data[t])
            data = xp.array(selected_slices)

        # Handle the case it needs to be processed on the fly
        if not self.data_processed:
            reshaped_data = xp.zeros((layer_count, time_count, *self.data_shape)) if self.has_layers else xp.zeros((time_count, *self.data_shape))

            for i in range(time_count):
                if self.has_layers:
                    for j in range(layer_count):
                        reshaped_data[j,i] = _resize_array(data[j,i],
                                                           new_shape=self.data_shape,
                                                           interpolation=self.interpolation_method.lower())
                else:
                    reshaped_data[i] = _resize_array(data[i],
                                                     new_shape=self.data_shape,
                                                     interpolation=self.interpolation_method.lower())

            data = xp.array(reshaped_data)

        # Return 0.0 if outside of data zone
        data_pos = pos - self.margins[:,0][None,:]
        data_pos_valid = xp.all((data_pos >= 0) & (data_pos < xp.array(self.data_shape)), axis=1)
        observation = xp.zeros(pos_count, dtype=float)

        # Gathering data on layered data on not
        if self.has_layers:
            observation[data_pos_valid] = data[(layer if isinstance(layer, int) else layer[data_pos_valid]), # layer
                                               (time if isinstance(time, int) else time[data_pos_valid]), # t
                                               *data_pos[data_pos_valid,:].T] # physical position
        else:
            observation[data_pos_valid] = data[(time if isinstance(time, int) else time[data_pos_valid]), # t
                                               *data_pos[data_pos_valid,:].T] # physical position

        return float(observation[0]) if is_single_point else observation


    def source_reached(self,
                       pos: np.ndarray
                       ) -> bool | np.ndarray:
        '''
        Checks whether a given position is within the source radius.

        Parameters
        ----------
        pos : np.ndarray
            The position to check whether in the radius of the source.

        Returns
        -------
        is_at_source : bool
            Whether or not the position is within the radius of the source.
        '''
        xp = cp if self.on_gpu else np

        # Handling the case of a single point
        is_single_point = (len(pos.shape) == 1)
        if is_single_point:
            pos = pos[None,:]

        is_at_source: np.ndarray = (xp.sum((pos - self.source_position[None,:]) ** 2, axis=-1) <= (self.source_radius ** 2))

        return bool(is_at_source[0]) if is_single_point else is_at_source


    def random_start_points(self,
                            n: int = 1
                            ) -> np.ndarray:
        '''
        Function to generate n starting positions following the starting probabilities.

        Parameters
        ----------
        n : int, default=1
            How many random starting positions to generate

        Returns
        -------
        random_states_2d : np.ndarray
            The n random 2d points in a n x 2 array.
        '''
        xp = cp if self.on_gpu else np

        assert (n > 0), "n has to be a strictly positive number (>0)"

        random_states = self.rnd_state.choice(xp.arange(int(np.prod(self.shape))), size=n, replace=True, p=self.start_probabilities.ravel())
        random_states_2d = xp.array(xp.unravel_index(random_states, self.shape)).T
        return random_states_2d


    def move(self,
             pos: np.ndarray,
             movement: np.ndarray
             ) -> np.ndarray:
        '''
        Applies a movement vector to a position point and returns a new position point while respecting the boundary conditions.

        Parameters
        ----------
        pos : np.ndarray
            The start position of the movement.
        movement : np.ndarray
            A 2D movement vector.

        Returns
        -------
        new_pos : np.ndarray
            The new position after applying the movement.
        '''
        xp = cp if self.on_gpu else np

        # Applying the movement vector
        new_pos = pos + movement

        # Handling the case we are dealing with a single point.
        is_single_point = (len(pos.shape) == 1)
        if is_single_point:
            new_pos = new_pos[None,:]

        shape_array = xp.array(self.shape)[None,:]

        # Wrap boundary
        if self.boundary_condition == 'wrap':
            new_pos = xp.where(new_pos < 0, (new_pos + shape_array), new_pos)
            new_pos = xp.where(new_pos >= shape_array, (new_pos - shape_array), new_pos)

        # Stop boundary
        elif self.boundary_condition == 'stop':
            new_pos = xp.clip(new_pos, 0, (shape_array-1))

        # Special wrap - vertical only
        elif (self.dimensions == 2) and (self.boundary_condition == 'wrap_vertical'):
            height, width = self.shape

            new_pos[new_pos[:,0] < 0, 0] += height
            new_pos[new_pos[:,0] >= height, 0] -= height

            new_pos[:,1] = xp.clip(new_pos[:,1], 0, (width-1))

        # Special wrap - horizontal only
        elif (self.dimensions == 2) and (self.boundary_condition == 'wrap_horizontal'):
            height, width = self.shape

            new_pos[new_pos[:,1] < 0, 1] += width
            new_pos[new_pos[:,1] >= width, 1] -= width

            new_pos[:,0] = xp.clip(new_pos[:,0], 0, (height-1))

        return new_pos[0] if is_single_point else new_pos


    def distance_to_source(self,
                           point: np.ndarray,
                           metric: Literal['manhattan'] = 'manhattan'
                           ) -> float | np.ndarray:
        '''
        Function to compute the distance(s) between given points and the source point.

        Parameters
        ----------
        point : np.ndarray
            A single or an Nx2 array containing N points.
        metric : 'manhattan'
            The metric to use to compute the distance.

        Returns
        -------
        dist : float or np.ndarray
            A single distance or a list of distance in a 1D distance array.
        '''
        xp = cp if self.on_gpu else np

        # Handling the case we have a single point
        is_single_point = (len(point.shape) == 1)
        if is_single_point:
            point = point[None,:]

        # Computing dist
        dist = None
        if metric == 'manhattan':
            dist = xp.sum(xp.abs(self.source_position[None,:] - point), axis=-1) - self.source_radius

        if dist is None: # Meaning it was not computed
            raise NotImplementedError('This distance metric has not yet been implemented')

        return float(dist[0]) if is_single_point else dist


    def save(self,
             folder: str | None = None,
             save_arrays: bool = False,
             force: bool = False
             ) -> None:
        '''
        Function to save the environment to the memory.

        By default it saved in a new folder at the current path in a new folder with the name 'Env-<name>' where <name> is the name set when initializing an environment.
        In this folder a file "METADATA.json" is created containing all the properties of the environment.

        The numpy arrays of the environment (grid and start_probabilities) can be saved or not. If not, when the environment is loaded it needs to be reconstructed from the original data file.
        The arrays are saved to .npy files along with the METADATA file.

        If an environment of the same name is already saved, the saving will be interupted. It can however be forced with the force parameter.

        Parameters
        ----------
        folder : str, optional
            The folder to which to save the environment data. If it is not provided, it will be created in the current folder.
        save_arrays : bool, default=False
            Whether or not to save the numpy arrays to memory. (The arrays can be heavy)
        force : bool, default=False
            In case an environment of the same name is already saved, it will be overwritten.
        '''
        # If on gpu, use the cpu version to save
        if self.on_gpu:
            self._alternate_version.save(
                folder=folder,
                save_arrays=save_arrays,
                force=force
            )
            return # Blank return

        # Assert either data_file is provided or save_arrays is enabled
        assert save_arrays or ((self.data_file_path is not None) and (self.start_type is not None)), "The environment was not created from a data file so 'save_arrays' has to be set to True."

        # Adding env name to folder path
        if folder is None:
            folder = f'./Env-{self.name}'
        else:
            folder += '/Env-' + self.name

        # Checking the folder exists or creates it
        if not os.path.exists(folder):
            os.mkdir(folder)
        elif len(os.listdir(folder)) > 0:
            if force:
                shutil.rmtree(folder)
                os.mkdir(folder)
            else:
                raise Exception(f'{folder} is not empty. If you want to overwrite the saved model, enable "force".')

        # Generating the metadata arguments dictionary
        arguments = {}
        arguments['name'] = self.name

        if self.data_file_path is not None:
            arguments['data_file_path'] = self.data_file_path

        arguments['timesteps']                     = int(self.timesteps)
        arguments['data_shape']                    = self.data_shape
        arguments['dimensions']                    = self.dimensions
        arguments['margins']                       = self.margins.tolist()
        arguments['shape']                         = self.shape
        arguments['data_bounds']                   = self.data_bounds.tolist()
        arguments['original_data_source_position'] = self.original_data_source_position.tolist()
        arguments['data_source_position']          = self.data_source_position.tolist()
        arguments['layers']                        = (self.layer_labels if self.has_layers else False)
        arguments['source_position']               = self.source_position.tolist()
        arguments['source_radius']                 = self.source_radius
        arguments['interpolation_method']          = self.interpolation_method
        arguments['preprocess_data']               = self._preprocess_data
        arguments['data_processed']                = self.data_processed
        arguments['boundary_condition']            = self.boundary_condition
        arguments['start_type']                    = self.start_type
        arguments['seed']                          = self.seed

        # Check how the start probabilities were built
        if self.start_type.startswith('custom') and len(self.start_type.split('_')) == 1 and not save_arrays:
            raise Exception('Start probabilities have been set from a custom array, please enable save_arrays to be able to reconstruct the environment later.')

        if self.odor_present_threshold is not None:
            arguments['odor_present_threshold'] = self.odor_present_threshold

        # Output the arguments to a METADATA file
        with open(folder + '/METADATA.json', 'w') as json_file:
            json.dump(arguments, json_file, indent=4)

        # Output the numpy arrays
        if save_arrays:
            if isinstance(self._data, np.ndarray):
                np.save(folder + '/data.npy', self._data)
            else:
                raise NotImplementedError('The saving of data that is not a Numpy array was not implemented yet.')
            np.save(folder + '/start_probabilities.npy', self.start_probabilities)

        # Success print
        self.saved_at = os.path.abspath(folder).replace('\\', '/')
        print(f'Environment saved to: {folder}')


    @classmethod
    def load(cls,
             folder: str
             ) -> 'Environment':
        '''
        Function to load an environment from a given folder.

        Parameters
        ----------
        folder : str
            The folder of the Environment.

        Returns
        -------
        loaded_env : Environment
            The loaded environment.
        '''
        assert os.path.exists(folder), "Folder doesn't exist..."
        assert folder.split('/')[-1].startswith('Env-'), "The folder provided is not the data of en Environment object."

        # Load arguments
        arguments: dict = None
        with open(folder + '/METADATA.json', 'r') as json_file:
            arguments = json.load(json_file)

        # Check if numpy arrays are provided, if not, recreate a new environment model
        if os.path.exists(folder + '/data.npy') and os.path.exists(folder + '/start_probabilities.npy'):
            data = np.load(folder + '/data.npy')
            start_probabilities = np.load(folder + '/start_probabilities.npy')

            loaded_env = cls.__new__(cls)

            # Set the arguments
            loaded_env.name                          = arguments['name']
            loaded_env.timesteps                     = arguments['timesteps']
            loaded_env.data_shape                    = arguments['data_shape']
            loaded_env.dimensions                    = arguments['dimensions']
            loaded_env.margins                       = np.array(arguments['margins'])
            loaded_env.shape                         = arguments['shape']
            loaded_env.data_bounds                   = np.array(arguments['data_bounds'])
            loaded_env.original_data_source_position = np.array(arguments['original_data_source_position'])
            loaded_env.data_source_position          = np.array(arguments['data_source_position'])
            loaded_env.source_position               = np.array(arguments['source_position'])
            loaded_env.source_radius                 = arguments['source_radius']
            loaded_env.has_layers                    = isinstance(arguments['layers'], list)
            loaded_env.layers                        = np.arange(len(arguments['layers'])) if loaded_env.has_layers else None
            loaded_env.layer_labels                  = arguments['layers']
            loaded_env.interpolation_method          = arguments['interpolation_method']
            loaded_env._preprocess_data              = arguments['preprocess_data']
            loaded_env.data_processed                = arguments['data_processed']
            loaded_env.boundary_condition            = arguments['boundary_condition']
            loaded_env.on_gpu                        = False
            loaded_env.seed                          = arguments['seed']
            loaded_env.rnd_state                     = np.random.RandomState(arguments['seed'])

            # Optional arguments
            loaded_env.data_file_path                = arguments.get('data_file_path')
            loaded_env.odor_present_threshold        = arguments.get('odor_present_threshold')
            loaded_env.start_type                    = arguments.get('start_type')

            # Arrays
            loaded_env._data = data
            loaded_env.start_probabilities = start_probabilities

        else:
            start_zone: str = arguments['start_type']
            start_zone_boundaries = None
            if start_zone.startswith('custom'):
                start_zone_boundaries = np.array(start_zone.split('_')[1:]).reshape((arguments['dimensions'],2)).astype(int)

            loaded_env = Environment(
                data_file              = arguments['data_file_path'],
                data_source_position   = arguments['original_data_source_position'],
                source_radius          = arguments['source_radius'],
                layers                 = arguments['layers'],
                shape                  = arguments['shape'],
                margins                = arguments['margins'],
                interpolation_method   = arguments['interpolation_method'],
                preprocess_data        = arguments['preprocess_data'],
                boundary_condition     = arguments['boundary_condition'],
                start_zone             = (start_zone_boundaries if start_zone_boundaries is not None else start_zone),
                odor_present_threshold = arguments.get('odor_present_threshold'),
                name                   = arguments['name'],
                seed                   = arguments['seed']
            )

        # Folder where the environment was pulled from
        loaded_env.saved_at = os.path.abspath(folder)

        return loaded_env


    def to_gpu(self) -> 'Environment':
        '''
        Function to send the numpy arrays of the environment to the gpu memory.
        It returns a new instance of the Environment with the arrays as cupy arrays.

        Returns
        -------
        gpu_environment : Environment
            A new environment instance where the arrays are on the gpu memory.
        '''
        # Check whether the environment is already on the gpu or not
        if self.on_gpu:
            return self

        # Warn and overwrite alternate_version in case it already exists
        if self._alternate_version is not None:
            print('[warning] A GPU instance already existed and is being recreated.')
            self._alternate_version = None

        assert gpu_support, "GPU support is not enabled..."

        # Generating a new instance
        cls = self.__class__
        gpu_environment = cls.__new__(cls)

        # Copying arguments to gpu
        for arg, val in self.__dict__.items():
            if isinstance(val, np.ndarray):
                setattr(gpu_environment, arg, cp.array(val))
            elif arg == 'rnd_state':
                setattr(gpu_environment, arg, cp.random.RandomState(self.seed))
            else:
                setattr(gpu_environment, arg, val)

        # Self reference instances
        self._alternate_version = gpu_environment
        gpu_environment._alternate_version = self

        gpu_environment.on_gpu = True
        return gpu_environment


    def to_cpu(self) -> 'Environment':
        '''
        Function to send the numpy arrays of the environment to the cpu memory.
        It returns a new instance of the Environment with the arrays as numpy arrays.

        Returns
        -------
        cpu_environment : Environment
            A new environment instance where the arrays are on the cpu memory.
        '''
        # Check whether the agent is already on the cpu or not
        if not self.on_gpu:
            return self

        if self._alternate_version is not None:
            print('[warning] A CPU instance already existed and is being recreated.')
            self._alternate_version = None

        # Generating a new instance
        cls = self.__class__
        cpu_environment = cls.__new__(cls)

        # Copying arguments to gpu
        for arg, val in self.__dict__.items():
            if isinstance(val, cp.ndarray):
                setattr(cpu_environment, arg, cp.asnumpy(val))
            elif arg == 'rnd_state':
                setattr(cpu_environment, arg, np.random.RandomState(self.seed))
            else:
                setattr(cpu_environment, arg, val)

        # Self reference instances
        self._alternate_version = cpu_environment
        cpu_environment._alternate_version = self

        cpu_environment.on_gpu = True
        return cpu_environment


    @property
    def gpu_version(self) -> 'Environment':
        '''
        A version of the Environment on the GPU.
        If the environment is already on the GPU it returns itself, otherwise the to_gpu function is called to generate a new one.
        '''
        if self.on_gpu:
            return self
        else:
            if self._alternate_version is not None: # Check if an alternate version already exists
                return self._alternate_version
            else: # Generate an alternate version on the gpu
                return self.to_gpu()


    @property
    def cpu_version(self) -> 'Environment':
        '''
        A version of the Environment on the CPU.
        If the environment is already on the CPU it returns itself, otherwise the to_cpu function is called to generate a new one.
        '''
        if not self.on_gpu:
            return self
        else:
            if self._alternate_version is not None: # Check if an alternate version already exists
                return self._alternate_version
            else: # Generate an alternate version on the cpu
                return self.to_cpu()


    def modify(self,
               data_source_position: list | np.ndarray | None = None,
               source_radius: float | None = None,
               shape: list | np.ndarray | None = None,
               margins: int | list | np.ndarray | None = None,
               multiplier: list | np.ndarray | None = None,
               interpolation_method: str | None = None,
               boundary_condition: str | None = None
               ) -> 'Environment':
        '''
        Returns a copy of the environment with one or more parameters modified.

        Parameters
        ----------
        data_source_position: list or np.ndarray, optional
            A new position for the source relative to the data file.
        source_radius: float, optional
            A new source radius.
        shape: list or np.ndarray, optional
            A new shape of environment.
        margins: int or list or np.ndarray, optional
            A new set of margins.
        multiplier: list or np.ndarray, optional
            A new multiplier to be applied to the data file (this will in turn increase or reduce the margins).
        interpolation_method: str, optional
            A new interpolation method to be used.
        boundary_condition: str, optional
            New boundary conditions for how the agent should behave at the edges.

        Returns
        -------
        modified_environment
            A copy of the environment where the modified parameters have been applied.
        '''
        if self.on_gpu:
            cpu_environment = self.cpu_version
            new_cpu_environment = cpu_environment.modify(
                data_source_position = data_source_position,
                source_radius        = source_radius,
                shape                = shape,
                margins              = margins,
                multiplier           = multiplier,
                interpolation_method = interpolation_method,
                boundary_condition   = boundary_condition
            )
            return new_cpu_environment.to_gpu()

        modified_environment = Environment(
            data_file              = (self.data_file_path if (self.data_file_path is not None) else self._data),
            data_source_position   = (data_source_position if (data_source_position is not None) else self.original_data_source_position),
            source_radius          = (source_radius if (source_radius is not None) else self.source_radius),
            layers                 = (self.layer_labels if self.has_layers else False),
            shape                  = (shape if (shape is not None) else self.shape),
            margins                = (margins if (margins is not None) else self.margins),
            multiplier             = (multiplier if (multiplier is not None) else [1.0,1.0]),
            interpolation_method   = (interpolation_method if (interpolation_method is not None) else self.interpolation_method),
            preprocess_data        = self._preprocess_data,
            boundary_condition     = (boundary_condition if (boundary_condition is not None) else self.boundary_condition),
            start_zone             = self.start_type,
            odor_present_threshold = self.odor_present_threshold,
            name                   = self.name,
            seed                   = self.seed
        )
        return modified_environment


    def modify_scale(self,
                     scale_factor: float
                     ) -> 'Environment':
        '''
        Function to modify the size of the environment by a scale factor.
        Everything will be scaled this factor. This includes: shape, margins, source radius, and data shape.

        Parameters
        ----------
        scale_factor : float
            By how much to modify the size of the current environment.

        Returns
        -------
        modified_environment : Environment
            The environment with the scale factor applied.
        '''
        modified_source_radius = self.source_radius * scale_factor
        modified_shape = (np.array(self.shape) * scale_factor).astype(int)
        modified_margins = (self.margins * scale_factor).astype(int)

        modified_environment = Environment(
            data_file              = (self.data_file_path if (self.data_file_path is not None) else self._data),
            data_source_position   = self.original_data_source_position,
            source_radius          = modified_source_radius,
            layers                 = (self.layer_labels if self.has_layers else False),
            shape                  = modified_shape,
            margins                = modified_margins,
            multiplier             = [1.0,1.0],
            interpolation_method   = self.interpolation_method,
            preprocess_data        = self._preprocess_data,
            boundary_condition     = self.boundary_condition,
            start_zone             = self.start_type,
            odor_present_threshold = self.odor_present_threshold,
            name                   = self.name,
            seed                   = self.seed
        )
        return modified_environment

_data_is_numpy property

Wheter the data is a numpy array or not.

cpu_version property

A version of the Environment on the CPU. If the environment is already on the CPU it returns itself, otherwise the to_cpu function is called to generate a new one.

data property

The whole dataset with the right shape. If not preprocessed to modify its shape the data will be processed when querrying this object.

gpu_version property

A version of the Environment on the GPU. If the environment is already on the GPU it returns itself, otherwise the to_gpu function is called to generate a new one.

distance_to_source(point, metric='manhattan')

Function to compute the distance(s) between given points and the source point.

Parameters:

Name	Type	Description	Default
point	ndarray	A single or an Nx2 array containing N points.	required
metric	manhattan	The metric to use to compute the distance.	'manhattan'

Returns:

Name	Type	Description
dist	float or ndarray	A single distance or a list of distance in a 1D distance array.

Source code in olfactory_navigation/environment.py

def distance_to_source(self,
                       point: np.ndarray,
                       metric: Literal['manhattan'] = 'manhattan'
                       ) -> float | np.ndarray:
    '''
    Function to compute the distance(s) between given points and the source point.

    Parameters
    ----------
    point : np.ndarray
        A single or an Nx2 array containing N points.
    metric : 'manhattan'
        The metric to use to compute the distance.

    Returns
    -------
    dist : float or np.ndarray
        A single distance or a list of distance in a 1D distance array.
    '''
    xp = cp if self.on_gpu else np

    # Handling the case we have a single point
    is_single_point = (len(point.shape) == 1)
    if is_single_point:
        point = point[None,:]

    # Computing dist
    dist = None
    if metric == 'manhattan':
        dist = xp.sum(xp.abs(self.source_position[None,:] - point), axis=-1) - self.source_radius

    if dist is None: # Meaning it was not computed
        raise NotImplementedError('This distance metric has not yet been implemented')

    return float(dist[0]) if is_single_point else dist

get_observation(pos, time=0, layer=0)

Function to get an observation at a given position on the grid at a given time. A set of observations can also be requested, either at a single position for multiple timestamps or with the same amoung of positions as timestamps provided.

Note: The position will not be checked against boundary conditions; if a position is out-of-bounds it will simply return 0.0!

Parameters:

Name	Type	Description	Default
pos	ndarray	The position or list of positions to get observations at.	required
time	int or ndarray	A timestamp or list of timestamps to get the observations at.	0
layer	int or ndarray	A layer or list of timestamps to get the observations at. Note: If the environment doesnt have layers, this parameter will be ignored.	0

Returns:

Name	Type	Description
observation	float or ndarray	A single observation or list of observations.

Source code in olfactory_navigation/environment.py

def get_observation(self,
                    pos: np.ndarray,
                    time: int | np.ndarray = 0,
                    layer: int | np.ndarray = 0
                    ) -> float | np.ndarray:
    '''
    Function to get an observation at a given position on the grid at a given time.
    A set of observations can also be requested, either at a single position for multiple timestamps or with the same amoung of positions as timestamps provided.

    Note: The position will not be checked against boundary conditions; if a position is out-of-bounds it will simply return 0.0!

    Parameters
    ----------
    pos : np.ndarray
        The position or list of positions to get observations at.
    time : int or np.ndarray, default=0
        A timestamp or list of timestamps to get the observations at.
    layer : int or np.ndarray, default=0
        A layer or list of timestamps to get the observations at.
        Note: If the environment doesnt have layers, this parameter will be ignored.

    Returns
    -------
    observation : float or np.ndarray
        A single observation or list of observations.
    '''
    xp = cp if self.on_gpu else np

    # Handling the case of a single point
    is_single_point = (len(pos.shape) == 1)
    if is_single_point:
        pos = pos[None,:]

    # Counting how many position points we are dealing with
    pos_count = len(pos)

    # Time looping
    time = time % self.timesteps

    # Determine unique layers and reindexing them if needed
    unique_layers = xp.array([layer]) if isinstance(layer, int) else xp.unique(layer)
    layer = 0 if isinstance(layer, int) else xp.where(layer == unique_layers[:,None])[0]
    layer_count = len(unique_layers)

    # Determine unique times and reindexing them if needed
    unique_times = xp.array([time]) if isinstance(time, int) else xp.unique(time)
    time = 0 if isinstance(time, int) else xp.where(time == unique_times[:,None])[0]
    time_count = len(unique_times)

    # Handling the case where the data is a sequence of slices (h5, so not numpy array)
    data = self._data

    # Selecting the required slices
    if self._data_is_numpy:
        data = data[unique_layers][:,unique_times] if self.has_layers else data[unique_times]
    else:
        # Case where we are dealing with a h5 file
        # Note: Can't use self.data_shape because we don't know whether the data is processed yet or no
        selected_slices = xp.zeros((layer_count, time_count, *self._data[0][0].shape)) if self.has_layers else xp.zeros((time_count, *self._data[0].shape))
        for i, t in enumerate(unique_times):
            if self.has_layers:
                for j, l in enumerate(unique_layers):
                    selected_slices[j,i] = xp.array(data[int(l)][int(t)])
            else:
                selected_slices[i] = xp.array(data[t])
        data = xp.array(selected_slices)

    # Handle the case it needs to be processed on the fly
    if not self.data_processed:
        reshaped_data = xp.zeros((layer_count, time_count, *self.data_shape)) if self.has_layers else xp.zeros((time_count, *self.data_shape))

        for i in range(time_count):
            if self.has_layers:
                for j in range(layer_count):
                    reshaped_data[j,i] = _resize_array(data[j,i],
                                                       new_shape=self.data_shape,
                                                       interpolation=self.interpolation_method.lower())
            else:
                reshaped_data[i] = _resize_array(data[i],
                                                 new_shape=self.data_shape,
                                                 interpolation=self.interpolation_method.lower())

        data = xp.array(reshaped_data)

    # Return 0.0 if outside of data zone
    data_pos = pos - self.margins[:,0][None,:]
    data_pos_valid = xp.all((data_pos >= 0) & (data_pos < xp.array(self.data_shape)), axis=1)
    observation = xp.zeros(pos_count, dtype=float)

    # Gathering data on layered data on not
    if self.has_layers:
        observation[data_pos_valid] = data[(layer if isinstance(layer, int) else layer[data_pos_valid]), # layer
                                           (time if isinstance(time, int) else time[data_pos_valid]), # t
                                           *data_pos[data_pos_valid,:].T] # physical position
    else:
        observation[data_pos_valid] = data[(time if isinstance(time, int) else time[data_pos_valid]), # t
                                           *data_pos[data_pos_valid,:].T] # physical position

    return float(observation[0]) if is_single_point else observation

load(folder) classmethod

Function to load an environment from a given folder.

Parameters:

Name	Type	Description	Default
folder	str	The folder of the Environment.	required

Returns:

Name	Type	Description
loaded_env	Environment	The loaded environment.

Source code in olfactory_navigation/environment.py

@classmethod
def load(cls,
         folder: str
         ) -> 'Environment':
    '''
    Function to load an environment from a given folder.

    Parameters
    ----------
    folder : str
        The folder of the Environment.

    Returns
    -------
    loaded_env : Environment
        The loaded environment.
    '''
    assert os.path.exists(folder), "Folder doesn't exist..."
    assert folder.split('/')[-1].startswith('Env-'), "The folder provided is not the data of en Environment object."

    # Load arguments
    arguments: dict = None
    with open(folder + '/METADATA.json', 'r') as json_file:
        arguments = json.load(json_file)

    # Check if numpy arrays are provided, if not, recreate a new environment model
    if os.path.exists(folder + '/data.npy') and os.path.exists(folder + '/start_probabilities.npy'):
        data = np.load(folder + '/data.npy')
        start_probabilities = np.load(folder + '/start_probabilities.npy')

        loaded_env = cls.__new__(cls)

        # Set the arguments
        loaded_env.name                          = arguments['name']
        loaded_env.timesteps                     = arguments['timesteps']
        loaded_env.data_shape                    = arguments['data_shape']
        loaded_env.dimensions                    = arguments['dimensions']
        loaded_env.margins                       = np.array(arguments['margins'])
        loaded_env.shape                         = arguments['shape']
        loaded_env.data_bounds                   = np.array(arguments['data_bounds'])
        loaded_env.original_data_source_position = np.array(arguments['original_data_source_position'])
        loaded_env.data_source_position          = np.array(arguments['data_source_position'])
        loaded_env.source_position               = np.array(arguments['source_position'])
        loaded_env.source_radius                 = arguments['source_radius']
        loaded_env.has_layers                    = isinstance(arguments['layers'], list)
        loaded_env.layers                        = np.arange(len(arguments['layers'])) if loaded_env.has_layers else None
        loaded_env.layer_labels                  = arguments['layers']
        loaded_env.interpolation_method          = arguments['interpolation_method']
        loaded_env._preprocess_data              = arguments['preprocess_data']
        loaded_env.data_processed                = arguments['data_processed']
        loaded_env.boundary_condition            = arguments['boundary_condition']
        loaded_env.on_gpu                        = False
        loaded_env.seed                          = arguments['seed']
        loaded_env.rnd_state                     = np.random.RandomState(arguments['seed'])

        # Optional arguments
        loaded_env.data_file_path                = arguments.get('data_file_path')
        loaded_env.odor_present_threshold        = arguments.get('odor_present_threshold')
        loaded_env.start_type                    = arguments.get('start_type')

        # Arrays
        loaded_env._data = data
        loaded_env.start_probabilities = start_probabilities

    else:
        start_zone: str = arguments['start_type']
        start_zone_boundaries = None
        if start_zone.startswith('custom'):
            start_zone_boundaries = np.array(start_zone.split('_')[1:]).reshape((arguments['dimensions'],2)).astype(int)

        loaded_env = Environment(
            data_file              = arguments['data_file_path'],
            data_source_position   = arguments['original_data_source_position'],
            source_radius          = arguments['source_radius'],
            layers                 = arguments['layers'],
            shape                  = arguments['shape'],
            margins                = arguments['margins'],
            interpolation_method   = arguments['interpolation_method'],
            preprocess_data        = arguments['preprocess_data'],
            boundary_condition     = arguments['boundary_condition'],
            start_zone             = (start_zone_boundaries if start_zone_boundaries is not None else start_zone),
            odor_present_threshold = arguments.get('odor_present_threshold'),
            name                   = arguments['name'],
            seed                   = arguments['seed']
        )

    # Folder where the environment was pulled from
    loaded_env.saved_at = os.path.abspath(folder)

    return loaded_env

modify(data_source_position=None, source_radius=None, shape=None, margins=None, multiplier=None, interpolation_method=None, boundary_condition=None)

Returns a copy of the environment with one or more parameters modified.

Parameters:

Name	Type	Description	Default
data_source_position	list | ndarray | None	A new position for the source relative to the data file.	None
source_radius	float | None	A new source radius.	None
shape	list | ndarray | None	A new shape of environment.	None
margins	int | list | ndarray | None	A new set of margins.	None
multiplier	list | ndarray | None	A new multiplier to be applied to the data file (this will in turn increase or reduce the margins).	None
interpolation_method	str | None	A new interpolation method to be used.	None
boundary_condition	str | None	New boundary conditions for how the agent should behave at the edges.	None

Returns:

Type	Description
modified_environment	A copy of the environment where the modified parameters have been applied.

Source code in olfactory_navigation/environment.py

def modify(self,
           data_source_position: list | np.ndarray | None = None,
           source_radius: float | None = None,
           shape: list | np.ndarray | None = None,
           margins: int | list | np.ndarray | None = None,
           multiplier: list | np.ndarray | None = None,
           interpolation_method: str | None = None,
           boundary_condition: str | None = None
           ) -> 'Environment':
    '''
    Returns a copy of the environment with one or more parameters modified.

    Parameters
    ----------
    data_source_position: list or np.ndarray, optional
        A new position for the source relative to the data file.
    source_radius: float, optional
        A new source radius.
    shape: list or np.ndarray, optional
        A new shape of environment.
    margins: int or list or np.ndarray, optional
        A new set of margins.
    multiplier: list or np.ndarray, optional
        A new multiplier to be applied to the data file (this will in turn increase or reduce the margins).
    interpolation_method: str, optional
        A new interpolation method to be used.
    boundary_condition: str, optional
        New boundary conditions for how the agent should behave at the edges.

    Returns
    -------
    modified_environment
        A copy of the environment where the modified parameters have been applied.
    '''
    if self.on_gpu:
        cpu_environment = self.cpu_version
        new_cpu_environment = cpu_environment.modify(
            data_source_position = data_source_position,
            source_radius        = source_radius,
            shape                = shape,
            margins              = margins,
            multiplier           = multiplier,
            interpolation_method = interpolation_method,
            boundary_condition   = boundary_condition
        )
        return new_cpu_environment.to_gpu()

    modified_environment = Environment(
        data_file              = (self.data_file_path if (self.data_file_path is not None) else self._data),
        data_source_position   = (data_source_position if (data_source_position is not None) else self.original_data_source_position),
        source_radius          = (source_radius if (source_radius is not None) else self.source_radius),
        layers                 = (self.layer_labels if self.has_layers else False),
        shape                  = (shape if (shape is not None) else self.shape),
        margins                = (margins if (margins is not None) else self.margins),
        multiplier             = (multiplier if (multiplier is not None) else [1.0,1.0]),
        interpolation_method   = (interpolation_method if (interpolation_method is not None) else self.interpolation_method),
        preprocess_data        = self._preprocess_data,
        boundary_condition     = (boundary_condition if (boundary_condition is not None) else self.boundary_condition),
        start_zone             = self.start_type,
        odor_present_threshold = self.odor_present_threshold,
        name                   = self.name,
        seed                   = self.seed
    )
    return modified_environment

modify_scale(scale_factor)

Function to modify the size of the environment by a scale factor. Everything will be scaled this factor. This includes: shape, margins, source radius, and data shape.

Parameters:

Name	Type	Description	Default
scale_factor	float	By how much to modify the size of the current environment.	required

Returns:

Name	Type	Description
modified_environment	Environment	The environment with the scale factor applied.

Source code in olfactory_navigation/environment.py

def modify_scale(self,
                 scale_factor: float
                 ) -> 'Environment':
    '''
    Function to modify the size of the environment by a scale factor.
    Everything will be scaled this factor. This includes: shape, margins, source radius, and data shape.

    Parameters
    ----------
    scale_factor : float
        By how much to modify the size of the current environment.

    Returns
    -------
    modified_environment : Environment
        The environment with the scale factor applied.
    '''
    modified_source_radius = self.source_radius * scale_factor
    modified_shape = (np.array(self.shape) * scale_factor).astype(int)
    modified_margins = (self.margins * scale_factor).astype(int)

    modified_environment = Environment(
        data_file              = (self.data_file_path if (self.data_file_path is not None) else self._data),
        data_source_position   = self.original_data_source_position,
        source_radius          = modified_source_radius,
        layers                 = (self.layer_labels if self.has_layers else False),
        shape                  = modified_shape,
        margins                = modified_margins,
        multiplier             = [1.0,1.0],
        interpolation_method   = self.interpolation_method,
        preprocess_data        = self._preprocess_data,
        boundary_condition     = self.boundary_condition,
        start_zone             = self.start_type,
        odor_present_threshold = self.odor_present_threshold,
        name                   = self.name,
        seed                   = self.seed
    )
    return modified_environment

move(pos, movement)

Applies a movement vector to a position point and returns a new position point while respecting the boundary conditions.

Parameters:

Name	Type	Description	Default
pos	ndarray	The start position of the movement.	required
movement	ndarray	A 2D movement vector.	required

Returns:

Name	Type	Description
new_pos	ndarray	The new position after applying the movement.

Source code in olfactory_navigation/environment.py

def move(self,
         pos: np.ndarray,
         movement: np.ndarray
         ) -> np.ndarray:
    '''
    Applies a movement vector to a position point and returns a new position point while respecting the boundary conditions.

    Parameters
    ----------
    pos : np.ndarray
        The start position of the movement.
    movement : np.ndarray
        A 2D movement vector.

    Returns
    -------
    new_pos : np.ndarray
        The new position after applying the movement.
    '''
    xp = cp if self.on_gpu else np

    # Applying the movement vector
    new_pos = pos + movement

    # Handling the case we are dealing with a single point.
    is_single_point = (len(pos.shape) == 1)
    if is_single_point:
        new_pos = new_pos[None,:]

    shape_array = xp.array(self.shape)[None,:]

    # Wrap boundary
    if self.boundary_condition == 'wrap':
        new_pos = xp.where(new_pos < 0, (new_pos + shape_array), new_pos)
        new_pos = xp.where(new_pos >= shape_array, (new_pos - shape_array), new_pos)

    # Stop boundary
    elif self.boundary_condition == 'stop':
        new_pos = xp.clip(new_pos, 0, (shape_array-1))

    # Special wrap - vertical only
    elif (self.dimensions == 2) and (self.boundary_condition == 'wrap_vertical'):
        height, width = self.shape

        new_pos[new_pos[:,0] < 0, 0] += height
        new_pos[new_pos[:,0] >= height, 0] -= height

        new_pos[:,1] = xp.clip(new_pos[:,1], 0, (width-1))

    # Special wrap - horizontal only
    elif (self.dimensions == 2) and (self.boundary_condition == 'wrap_horizontal'):
        height, width = self.shape

        new_pos[new_pos[:,1] < 0, 1] += width
        new_pos[new_pos[:,1] >= width, 1] -= width

        new_pos[:,0] = xp.clip(new_pos[:,0], 0, (height-1))

    return new_pos[0] if is_single_point else new_pos

plot(frame=0, layer=0, ax=None)

Simple function to plot the environment with a single frame of odor cues. The starting zone is also market down with a blue contour. The source of the odor is marked by a red circle.

Parameters:

Name	Type	Description	Default
frame	int	The frame of odor cues to print.	0
layer	int	The layer of the odor cues to print. (Ignored if the environment is not layered.)	0
ax	Axes	An ax on which the environment can be plot.	None

Source code in olfactory_navigation/environment.py

def plot(self,
         frame: int = 0,
         layer: int = 0,
         ax: plt.Axes | None = None
         ) -> None:
    '''
    Simple function to plot the environment with a single frame of odor cues.
    The starting zone is also market down with a blue contour.
    The source of the odor is marked by a red circle.

    Parameters
    ----------
    frame : int, default=0
        The frame of odor cues to print.
    layer : int, default=0
        The layer of the odor cues to print. (Ignored if the environment is not layered.)
    ax : plt.Axes, optional
        An ax on which the environment can be plot.
    '''
    # If on GPU use the CPU version to plot
    if self.on_gpu:
        self._alternate_version.plot(
            frame=frame,
            ax=ax
        )
        return # Blank return

    # TODO: Implement plotting for 3D
    assert self.dimensions == 2, "Plotting function only available for 2D environments for now..."

    if ax is None:
        _, ax = plt.subplots(1, figsize=(15,5))

    legend_elements = [[],[]]

    # Gather data frame
    data_frame: np.ndarray = self._data[layer][frame] if self.has_layers else self._data[frame]
    if not isinstance(data_frame, np.ndarray):
        data_frame = np.array(data_frame)

    if not self.data_processed:
        data_frame = _resize_array(data_frame,
                                   new_shape=self.data_shape,
                                   interpolation=self.interpolation_method.lower())

    # Odor grid
    odor = Rectangle([0,0], 1, 1, color='black', fill=True)
    frame_data = (data_frame > (self.odor_present_threshold if self.odor_present_threshold is not None else 0)).astype(float)
    environment_frame = np.zeros(self.shape, dtype=float)
    environment_frame[self.data_bounds[0,0]:self.data_bounds[0,1], self.data_bounds[1,0]:self.data_bounds[1,1]] = frame_data
    ax.imshow(environment_frame, cmap='Greys')

    legend_elements[0].append(odor)
    legend_elements[1].append(f'Frame {frame}' + ('' if not self.has_layers else f' (layer {layer})') + ' odor cues')

    # Start zone contour
    start_zone = Rectangle([0,0], 1, 1, color='blue', fill=False)
    ax.contour(self.start_probabilities, levels=[0.0], colors='blue')

    legend_elements[0].append(start_zone)
    legend_elements[1].append('Start zone')

    # Source circle
    goal_circle = Circle(self.source_position[::-1], self.source_radius, color='r', fill=False, zorder=10)
    legend_elements[0].append(goal_circle)
    legend_elements[1].append('Source')

    if self.source_radius > 0.0:
        ax.add_patch(goal_circle)
    else:
        ax.scatter(self.source_position[1], self.source_position[0], c='red')

    # Legend
    ax.legend(legend_elements[0], legend_elements[1])

random_start_points(n=1)

Function to generate n starting positions following the starting probabilities.

Parameters:

Name	Type	Description	Default
n	int	How many random starting positions to generate	1

Returns:

Name	Type	Description
random_states_2d	ndarray	The n random 2d points in a n x 2 array.

Source code in olfactory_navigation/environment.py

def random_start_points(self,
                        n: int = 1
                        ) -> np.ndarray:
    '''
    Function to generate n starting positions following the starting probabilities.

    Parameters
    ----------
    n : int, default=1
        How many random starting positions to generate

    Returns
    -------
    random_states_2d : np.ndarray
        The n random 2d points in a n x 2 array.
    '''
    xp = cp if self.on_gpu else np

    assert (n > 0), "n has to be a strictly positive number (>0)"

    random_states = self.rnd_state.choice(xp.arange(int(np.prod(self.shape))), size=n, replace=True, p=self.start_probabilities.ravel())
    random_states_2d = xp.array(xp.unravel_index(random_states, self.shape)).T
    return random_states_2d

save(folder=None, save_arrays=False, force=False)

Function to save the environment to the memory.

By default it saved in a new folder at the current path in a new folder with the name 'Env-' where is the name set when initializing an environment. In this folder a file "METADATA.json" is created containing all the properties of the environment.

The numpy arrays of the environment (grid and start_probabilities) can be saved or not. If not, when the environment is loaded it needs to be reconstructed from the original data file. The arrays are saved to .npy files along with the METADATA file.

If an environment of the same name is already saved, the saving will be interupted. It can however be forced with the force parameter.

Parameters:

Name	Type	Description	Default
folder	str	The folder to which to save the environment data. If it is not provided, it will be created in the current folder.	None
save_arrays	bool	Whether or not to save the numpy arrays to memory. (The arrays can be heavy)	False
force	bool	In case an environment of the same name is already saved, it will be overwritten.	False

Source code in olfactory_navigation/environment.py

def save(self,
         folder: str | None = None,
         save_arrays: bool = False,
         force: bool = False
         ) -> None:
    '''
    Function to save the environment to the memory.

    By default it saved in a new folder at the current path in a new folder with the name 'Env-<name>' where <name> is the name set when initializing an environment.
    In this folder a file "METADATA.json" is created containing all the properties of the environment.

    The numpy arrays of the environment (grid and start_probabilities) can be saved or not. If not, when the environment is loaded it needs to be reconstructed from the original data file.
    The arrays are saved to .npy files along with the METADATA file.

    If an environment of the same name is already saved, the saving will be interupted. It can however be forced with the force parameter.

    Parameters
    ----------
    folder : str, optional
        The folder to which to save the environment data. If it is not provided, it will be created in the current folder.
    save_arrays : bool, default=False
        Whether or not to save the numpy arrays to memory. (The arrays can be heavy)
    force : bool, default=False
        In case an environment of the same name is already saved, it will be overwritten.
    '''
    # If on gpu, use the cpu version to save
    if self.on_gpu:
        self._alternate_version.save(
            folder=folder,
            save_arrays=save_arrays,
            force=force
        )
        return # Blank return

    # Assert either data_file is provided or save_arrays is enabled
    assert save_arrays or ((self.data_file_path is not None) and (self.start_type is not None)), "The environment was not created from a data file so 'save_arrays' has to be set to True."

    # Adding env name to folder path
    if folder is None:
        folder = f'./Env-{self.name}'
    else:
        folder += '/Env-' + self.name

    # Checking the folder exists or creates it
    if not os.path.exists(folder):
        os.mkdir(folder)
    elif len(os.listdir(folder)) > 0:
        if force:
            shutil.rmtree(folder)
            os.mkdir(folder)
        else:
            raise Exception(f'{folder} is not empty. If you want to overwrite the saved model, enable "force".')

    # Generating the metadata arguments dictionary
    arguments = {}
    arguments['name'] = self.name

    if self.data_file_path is not None:
        arguments['data_file_path'] = self.data_file_path

    arguments['timesteps']                     = int(self.timesteps)
    arguments['data_shape']                    = self.data_shape
    arguments['dimensions']                    = self.dimensions
    arguments['margins']                       = self.margins.tolist()
    arguments['shape']                         = self.shape
    arguments['data_bounds']                   = self.data_bounds.tolist()
    arguments['original_data_source_position'] = self.original_data_source_position.tolist()
    arguments['data_source_position']          = self.data_source_position.tolist()
    arguments['layers']                        = (self.layer_labels if self.has_layers else False)
    arguments['source_position']               = self.source_position.tolist()
    arguments['source_radius']                 = self.source_radius
    arguments['interpolation_method']          = self.interpolation_method
    arguments['preprocess_data']               = self._preprocess_data
    arguments['data_processed']                = self.data_processed
    arguments['boundary_condition']            = self.boundary_condition
    arguments['start_type']                    = self.start_type
    arguments['seed']                          = self.seed

    # Check how the start probabilities were built
    if self.start_type.startswith('custom') and len(self.start_type.split('_')) == 1 and not save_arrays:
        raise Exception('Start probabilities have been set from a custom array, please enable save_arrays to be able to reconstruct the environment later.')

    if self.odor_present_threshold is not None:
        arguments['odor_present_threshold'] = self.odor_present_threshold

    # Output the arguments to a METADATA file
    with open(folder + '/METADATA.json', 'w') as json_file:
        json.dump(arguments, json_file, indent=4)

    # Output the numpy arrays
    if save_arrays:
        if isinstance(self._data, np.ndarray):
            np.save(folder + '/data.npy', self._data)
        else:
            raise NotImplementedError('The saving of data that is not a Numpy array was not implemented yet.')
        np.save(folder + '/start_probabilities.npy', self.start_probabilities)

    # Success print
    self.saved_at = os.path.abspath(folder).replace('\\', '/')
    print(f'Environment saved to: {folder}')

source_reached(pos)

Checks whether a given position is within the source radius.

Parameters:

Name	Type	Description	Default
pos	ndarray	The position to check whether in the radius of the source.	required

Returns:

Name	Type	Description
is_at_source	bool	Whether or not the position is within the radius of the source.

Source code in olfactory_navigation/environment.py

def source_reached(self,
                   pos: np.ndarray
                   ) -> bool | np.ndarray:
    '''
    Checks whether a given position is within the source radius.

    Parameters
    ----------
    pos : np.ndarray
        The position to check whether in the radius of the source.

    Returns
    -------
    is_at_source : bool
        Whether or not the position is within the radius of the source.
    '''
    xp = cp if self.on_gpu else np

    # Handling the case of a single point
    is_single_point = (len(pos.shape) == 1)
    if is_single_point:
        pos = pos[None,:]

    is_at_source: np.ndarray = (xp.sum((pos - self.source_position[None,:]) ** 2, axis=-1) <= (self.source_radius ** 2))

    return bool(is_at_source[0]) if is_single_point else is_at_source

to_cpu()

Function to send the numpy arrays of the environment to the cpu memory. It returns a new instance of the Environment with the arrays as numpy arrays.

Returns:

Name	Type	Description
cpu_environment	Environment	A new environment instance where the arrays are on the cpu memory.

Source code in olfactory_navigation/environment.py

def to_cpu(self) -> 'Environment':
    '''
    Function to send the numpy arrays of the environment to the cpu memory.
    It returns a new instance of the Environment with the arrays as numpy arrays.

    Returns
    -------
    cpu_environment : Environment
        A new environment instance where the arrays are on the cpu memory.
    '''
    # Check whether the agent is already on the cpu or not
    if not self.on_gpu:
        return self

    if self._alternate_version is not None:
        print('[warning] A CPU instance already existed and is being recreated.')
        self._alternate_version = None

    # Generating a new instance
    cls = self.__class__
    cpu_environment = cls.__new__(cls)

    # Copying arguments to gpu
    for arg, val in self.__dict__.items():
        if isinstance(val, cp.ndarray):
            setattr(cpu_environment, arg, cp.asnumpy(val))
        elif arg == 'rnd_state':
            setattr(cpu_environment, arg, np.random.RandomState(self.seed))
        else:
            setattr(cpu_environment, arg, val)

    # Self reference instances
    self._alternate_version = cpu_environment
    cpu_environment._alternate_version = self

    cpu_environment.on_gpu = True
    return cpu_environment

to_gpu()

Function to send the numpy arrays of the environment to the gpu memory. It returns a new instance of the Environment with the arrays as cupy arrays.

Returns:

Name	Type	Description
gpu_environment	Environment	A new environment instance where the arrays are on the gpu memory.

Source code in olfactory_navigation/environment.py

def to_gpu(self) -> 'Environment':
    '''
    Function to send the numpy arrays of the environment to the gpu memory.
    It returns a new instance of the Environment with the arrays as cupy arrays.

    Returns
    -------
    gpu_environment : Environment
        A new environment instance where the arrays are on the gpu memory.
    '''
    # Check whether the environment is already on the gpu or not
    if self.on_gpu:
        return self

    # Warn and overwrite alternate_version in case it already exists
    if self._alternate_version is not None:
        print('[warning] A GPU instance already existed and is being recreated.')
        self._alternate_version = None

    assert gpu_support, "GPU support is not enabled..."

    # Generating a new instance
    cls = self.__class__
    gpu_environment = cls.__new__(cls)

    # Copying arguments to gpu
    for arg, val in self.__dict__.items():
        if isinstance(val, np.ndarray):
            setattr(gpu_environment, arg, cp.array(val))
        elif arg == 'rnd_state':
            setattr(gpu_environment, arg, cp.random.RandomState(self.seed))
        else:
            setattr(gpu_environment, arg, val)

    # Self reference instances
    self._alternate_version = gpu_environment
    gpu_environment._alternate_version = self

    gpu_environment.on_gpu = True
    return gpu_environment

SimulationHistory

Class to record the steps that happened during a simulation with the following information being saved:

• the positions the agents pass by
• the actions the agents take
• the observations the agents receive ('observations')
• the time in the simulation process

Parameters:

Name	Type	Description	Default
start_points	ndarray	The initial points of the agents in the simulation.	required
environment	Environment	The environment on which the simulation is run (can be different from the one associated with the agent).	required
agent	Agent	The agent used in the simulation.	required
time_shift	ndarray	An array of time shifts in the simulation data.	required
horizon	int	The horizon of the simulation. i.e. how many steps can be taken by the agent during the simulation before he is considered lost.	required
reward_discount	float	A discount to be applied to the rewards received by the agent. (eg: reward of 1 received at time n would be: 1 * reward_discount^n)	0.99

Attributes:

Name	Type	Description
start_points	ndarray
environment	Environment
agent	Agent
time_shift	ndarray
horizon	int
reward_discount	float
environment_dimensions	int	The amount of dimensions of the environment.
environment_shape	tuple[int]	The shape of the environment.
environment_source_position	ndarray	The position of the odor source in the environment.
environment_source_radius	float	The radius of the odor source in the environment.
environment_layer_labels	list[str] or None	A list of the layer labels if the environment has layers.
agent_thresholds	ndarray	An array of the olfaction thresholds of the agent.
n	int	The amount of simulations.
start_time	datetime	The datetime the simulations start.
actions	list[ndarray]	A list of numpy arrays. At each step of the simulation, an array of shape n by 2 is appended to this list representing the n actions as dy,dx vectors.
positions	list[ndarray]	A list of numpy arrays. At each step of the simulation, an array of shape n by 2 is appended to this list representing the n positions as y,x vectors.
observations	list[ndarray]	A list of numpy arrays. At each step of the simulation, an array of shape n is appended to this list representing the n observations received by the agents.
reached_source	ndarray	A numpy array of booleans saying whether the simulations reached the source or not.
done_at_step	ndarray	A numpy array containing n elements that records when a given simulation reaches the source (-1 is not reached).

Source code in olfactory_navigation/simulation.py

class SimulationHistory:
    '''
    Class to record the steps that happened during a simulation with the following information being saved:

    - the positions the agents pass by
    - the actions the agents take
    - the observations the agents receive ('observations')
    - the time in the simulation process


    Parameters
    ----------
    start_points : np.ndarray
        The initial points of the agents in the simulation.
    environment : Environment
        The environment on which the simulation is run (can be different from the one associated with the agent).
    agent : Agent
        The agent used in the simulation.
    time_shift : np.ndarray
        An array of time shifts in the simulation data.
    horizon : int
        The horizon of the simulation. i.e. how many steps can be taken by the agent during the simulation before he is considered lost.
    reward_discount : float, default=0.99
        A discount to be applied to the rewards received by the agent. (eg: reward of 1 received at time n would be: 1 * reward_discount^n)

    Attributes
    ----------
    start_points : np.ndarray
    environment : Environment
    agent : Agent
    time_shift : np.ndarray
    horizon : int
    reward_discount : float
    environment_dimensions : int
        The amount of dimensions of the environment.
    environment_shape : tuple[int]
        The shape of the environment.
    environment_source_position : np.ndarray
        The position of the odor source in the environment.
    environment_source_radius : float
        The radius of the odor source in the environment.
    environment_layer_labels : list[str] or None
        A list of the layer labels if the environment has layers.
    agent_thresholds : np.ndarray
        An array of the olfaction thresholds of the agent.
    n : int
        The amount of simulations.
    start_time : datetime
        The datetime the simulations start.
    actions : list[np.ndarray]
        A list of numpy arrays. At each step of the simulation, an array of shape n by 2 is appended to this list representing the n actions as dy,dx vectors.
    positions : list[np.ndarray]
        A list of numpy arrays. At each step of the simulation, an array of shape n by 2 is appended to this list representing the n positions as y,x vectors.
    observations : list[np.ndarray]
        A list of numpy arrays. At each step of the simulation, an array of shape n is appended to this list representing the n observations received by the agents.
    reached_source : np.ndarray
        A numpy array of booleans saying whether the simulations reached the source or not.
    done_at_step : np.ndarray
        A numpy array containing n elements that records when a given simulation reaches the source (-1 is not reached).
    '''
    def __init__(self,
                 start_points: np.ndarray,
                 environment: Environment,
                 agent: Agent,
                 time_shift: np.ndarray,
                 horizon: int,
                 reward_discount: float = 0.99
                 ) -> None:
        # If only on state is provided, we make it a 1x2 vector
        if len(start_points.shape) == 1:
            start_points = start_points[None,:]

        # Fixed parameters
        self.n = len(start_points)
        self.environment = environment.cpu_version
        self.agent = agent.cpu_version
        self.time_shift = time_shift if gpu_support and cp.get_array_module(time_shift) == np else cp.asnumpy(time_shift)
        self.horizon = horizon
        self.reward_discount = reward_discount
        self.start_time = datetime.now()

        # Simulation Tracking
        self.start_points = start_points if gpu_support and cp.get_array_module(start_points) == np else cp.asnumpy(start_points)
        self.actions = []
        self.positions = []
        self.observations = []
        self.timestamps: list[datetime] = []

        self._running_sims = np.arange(self.n)
        self.reached_source = np.zeros(self.n, dtype=bool)
        self.done_at_step = np.full(self.n, fill_value=-1)

        # Environment and agent attributes
        self.environment_dimensions = self.environment.dimensions
        self.environment_shape = self.environment.shape
        self.environment_source_position = self.environment.source_position
        self.environment_source_radius = self.environment.source_radius
        self.environment_layer_labels = self.environment.layer_labels
        self.agent_thresholds = self.agent.thresholds

        # Other parameters
        self._simulation_dfs = None


    def add_step(self,
                 actions: np.ndarray,
                 next_positions: np.ndarray,
                 observations: np.ndarray,
                 reached_source: np.ndarray,
                 interupt: np.ndarray
                 ) -> None:
        '''
        Function to add a step in the simulation history.

        Parameters
        ----------
        actions : np.ndarray
            The actions that were taken by the agents.
        next_positions : np.ndarray
            The positions that were reached by the agents after having taken actions.
        observations : np.ndarray
            The observations the agents receive after having taken actions.
        reached_source : np.ndarray
            A boolean array of whether each agent has reached the source or not.
        interupt : np.ndarray
            A boolean array of whether each agent has to be terminated even if it hasnt reached the source yet.
        '''
        self._simulation_dfs = None

        # Time tracking
        self.timestamps.append(datetime.now())

        # Check if environment if layered and/or 3D
        layered = 0 if self.environment_layer_labels is None else 1

        # Handle case cupy arrays are provided
        if gpu_support:
            actions = actions if cp.get_array_module(actions) == np else cp.asnumpy(actions)
            next_positions = next_positions if cp.get_array_module(next_positions) == np else cp.asnumpy(next_positions)
            observations = observations if cp.get_array_module(observations) == np else cp.asnumpy(observations)
            reached_source = reached_source if cp.get_array_module(reached_source) == np else cp.asnumpy(reached_source)
            interupt = interupt if cp.get_array_module(interupt) == np else cp.asnumpy(interupt)

        # Actions tracking
        action_all_sims = np.full((self.n, (layered + self.environment_dimensions)), fill_value=-1)
        action_all_sims[self._running_sims] = actions
        self.actions.append(action_all_sims)

        # Next states tracking
        next_position_all_sims = np.full((self.n, self.environment_dimensions), fill_value=-1)
        next_position_all_sims[self._running_sims] = next_positions
        self.positions.append(next_position_all_sims)

        # Observation tracking
        observation_all_sims = np.full((self.n,), fill_value=-1, dtype=float)
        observation_all_sims[self._running_sims] = observations
        self.observations.append(observation_all_sims)

        # Recording at which step the simulation is done if it is done and whether it reached the source
        self.done_at_step[self._running_sims[reached_source | interupt]] = len(self.positions)
        self.reached_source[self._running_sims[reached_source]] = True

        # Updating the list of running sims
        self._running_sims = self._running_sims[~reached_source & ~interupt]


    def compute_distance_to_source(self) -> np.ndarray:
        '''
        Function to compute the optimal distance to the source of each starting point according to the optimal_distance_metric attribute.

        Returns
        -------
        distance : np.ndarray
            The optimal distances to the source point.
        '''
        point = self.start_points

        # Handling the case we have a single point
        is_single_point = (len(point.shape) == 1)
        if is_single_point:
            point = point[None,:]

        # Computing dist
        dist = None
        # if self.optimal_distance_metric == 'manhattan': # TODO Allow for other metrics to be used
        dist = np.sum(np.abs(self.environment_source_position[None,:] - point), axis=-1) - self.environment_source_radius

        if dist is None: # Meaning it was not computed
            raise NotImplementedError('This distance metric has not yet been implemented')

        return float(dist[0]) if is_single_point else dist


    @property
    def runs_analysis_df(self) -> pd.DataFrame:
        '''
        A Pandas DataFrame analyzing the results of the simulations.
        It aggregates the simulations in single rows, recording:

         - <axis>:              The starting positions at the given axis
         - optimal_steps_count: The minimal amount of steps to reach the source
         - converged:           Whether or not the simulation reached the source
         - reached_horizon:     Whether the failed simulation reached to horizon
         - steps_taken:         The amount of steps the agent took to reach the source, (horizon if the simulation did not reach the source)
         - discounted_rewards:  The discounted reward received by the agent over the course of the simulation
         - extra_steps:         The amount of extra steps compared to the optimal trajectory
         - t_min_over_t:        Normalized version of the extra steps measure, where it tends to 1 the least amount of time the agent took to reach the source compared to an optimal trajectory.
        '''
        # Get axes labels
        axes_labels = None
        if self.environment_dimensions <= 3:
            axes_labels = ['z', 'y', 'x'][-self.environment_dimensions:]
        else:
            axes_labels = [f'x{i}' for i in range(self.environment_dimensions)]

        # Dataframe creation
        df = pd.DataFrame(self.start_points, columns=axes_labels)
        df['optimal_steps_count'] = self.compute_distance_to_source()
        df['converged'] = self.reached_source
        df['reached_horizon'] = np.all(self.positions[-1] != -1, axis=1) & ~self.reached_source & (len(self.positions) == self.horizon)
        df['steps_taken'] = np.where(self.done_at_step >= 0, self.done_at_step, len(self.positions))
        df['discounted_rewards'] = self.reward_discount ** df['steps_taken']
        df['extra_steps'] = df['steps_taken'] - df['optimal_steps_count']
        df['t_min_over_t'] = df['optimal_steps_count'] / df['steps_taken']

        # Reindex
        runs_list = [f'run_{i}' for i in range(self.n)]
        df.index = runs_list

        return df


    @property
    def general_analysis_df(self) -> pd.DataFrame:
        '''
        A Pandas DataFrame analyzing the results of the simulations.
        Summarizing the performance of all the simulations with the following metrics:

         - converged:           Whether or not the simulation reached the source
         - reached_horizon:     Whether the failed simulation reached to horizon
         - steps_taken:         The amount of steps the agent took to reach the source, (horizon if the simulation did not reach the source)
         - discounted_rewards:  The discounted reward received by the agent over the course of the simulation
         - extra_steps:         The amount of extra steps compared to the optimal trajectory
         - t_min_over_t:        Normalized version of the extra steps measure, where it tends to 1 the least amount of time the agent took to reach the source compared to an optimal trajectory.

        For the measures (converged, steps_taken, discounted_rewards, extra_steps, t_min_over_t), the average and standard deviations are computed in rows at the top.
        '''
        df = self.runs_analysis_df

        # Analysis aggregations
        columns_to_analyze = ['converged', 'reached_horizon', 'steps_taken', 'discounted_rewards', 'extra_steps', 't_min_over_t']
        row_names = [['mean', 'standard_deviation', 'success_mean', 'success_standard_deviation']]
        general_analysis_data = [
            df[columns_to_analyze].mean(),
            df[columns_to_analyze].std(),
            df.loc[df['converged'], columns_to_analyze].mean(),
            df.loc[df['converged'], columns_to_analyze].std()
        ]

        return pd.DataFrame(data=general_analysis_data, index=row_names, columns=columns_to_analyze)


    @property
    def done_count(self) -> int:
        '''
        Returns how many simulations are terminated (whether they reached the source or not).
        '''
        return self.n - len(self._running_sims)


    @property
    def successful_simulation(self) -> np.ndarray:
        return self.reached_source


    @property
    def success_count(self) -> int:
        '''
        Returns how many simulations reached the source.
        '''
        return int(np.sum(self.successful_simulation))


    @property
    def simulations_at_horizon(self) -> np.ndarray:
        '''
        Returns a boolean array of which simulations reached the horizon.
        '''
        last_position_exists = np.all(self.positions[-1] != -1, axis=1)
        simulation_reached_horizon = (len(self.positions) == self.horizon)
        return last_position_exists & ~self.reached_source & simulation_reached_horizon


    @property
    def summary(self) -> str:
        '''
        A string summarizing the performances of all the simulations.
        The metrics used are averages of:

         - Step count
         - Extra steps
         - Discounted rewards
         - Tmin / T

        Along with the respective the standard deviations and equally for only for the successful simulations.
        '''
        success_sim_count = self.success_count
        failed_count = self.n - success_sim_count
        reached_horizon_count = int(np.sum(self.simulations_at_horizon))
        summary_str = f'Simulations reached goal: {success_sim_count}/{self.n} ({failed_count} failures (reached horizon: {reached_horizon_count})) ({(success_sim_count*100)/self.n:.2f}% success)'

        if success_sim_count == 0:
            return summary_str

        # Metrics
        df = self.general_analysis_df

        summary_str += f"\n - {'Average step count:':<35} {df.loc['mean','steps_taken'].item():.3f} +- {df.loc['standard_deviation','steps_taken'].item():.2f} "
        summary_str += f"(Successful only: {df.loc['success_mean','steps_taken'].item():.3f} +- {df.loc['success_standard_deviation','steps_taken'].item():.2f})"

        summary_str += f"\n - {'Extra steps:':<35} {df.loc['mean','extra_steps'].item():.3f} +- {df.loc['standard_deviation','extra_steps'].item():.2f} "
        summary_str += f"(Successful only: {df.loc['success_mean','extra_steps'].item():.3f} +- {df.loc['success_standard_deviation','extra_steps'].item():.2f})"

        summary_str += f"\n - {'Average discounted rewards (ADR):':<35} {df.loc['mean','discounted_rewards'].item():.3f} +- {df.loc['standard_deviation','discounted_rewards'].item():.2f} "
        summary_str += f"(Successful only: {df.loc['success_mean','discounted_rewards'].item():.3f} +- {df.loc['success_standard_deviation','discounted_rewards'].item():.2f})"

        summary_str += f"\n - {'Tmin/T:':<35} {df.loc['mean','t_min_over_t'].item():.3f} +- {df.loc['standard_deviation','t_min_over_t'].item():.2f} "
        summary_str += f"(Successful only: {df.loc['success_mean','t_min_over_t'].item():.3f} +- {df.loc['success_standard_deviation','t_min_over_t'].item():.2f})"

        return summary_str


    @property
    def simulation_dfs(self) -> list[pd.DataFrame]:
        '''
        A list of the pandas DataFrame where each dataframe is a single simulation history.
        Each row is a different time instant of simulation process with each column being:

         - time (of the simulation data)
         - [position] (z,) y, x  OR  x0, x1, ... xn
         - (layer)
         - [movement] (dz,) dy, dx  OR  dx0, dx1, ... dxn
         - o (pure, not thresholded)
         - reached_source (boolean)
        '''
        if self._simulation_dfs is None:
            self._simulation_dfs = []

            # Converting state, actions and observation to numpy arrays
            states_array = np.array(self.positions)
            action_array = np.array(self.actions)
            observation_array = np.array(self.observations)

            # Get axes labels
            axes_labels = None
            if self.environment_dimensions <= 3:
                axes_labels = ['z', 'y', 'x'][-self.environment_dimensions:]
            else:
                axes_labels = [f'x{i}' for i in range(self.environment_dimensions)]

            # Loop through the n simulations
            for i in range(self.n):
                length = self.done_at_step[i] if self.done_at_step[i] >= 0 else len(states_array)

                # Creation of the dataframe
                df = {}
                df['time'] = np.arange(length+1) + self.time_shift[i]

                # - Position variables
                for axis_i, axis in enumerate(axes_labels):
                    df[axis] = np.hstack([self.start_points[i, axis_i], states_array[:length, i, axis_i]])

                # - Action variables
                if self.environment_layer_labels is not None:
                    df['layer'] = np.hstack([[None], action_array[:length, i, 0]])

                for axis_i, axis in enumerate(axes_labels):
                    axis_i += (0 if self.environment_layer_labels is None else 1)
                    df['d' + axis]   = np.hstack([[None], action_array[:length, i, axis_i]])

                # - Other variables
                df['o'] = np.hstack([[None], observation_array[:length, i]])
                df['reached_source'] = np.hstack([[None], np.where(np.arange(1,length+1) == self.done_at_step[i], (1 if self.reached_source[i] else 0), 0)])

                # Append
                self._simulation_dfs.append(pd.DataFrame(df))

        return self._simulation_dfs


    def __add__(self, other_hist: 'SimulationHistory'):
        # Asserting the SimulationHistory objects are compatible
        assert self.horizon == other_hist.horizon, "The 'horizon' parameters must match between the two SimulationHistory objects..."
        assert self.reward_discount == other_hist.reward_discount, "The 'reward_discount' parameters must match between the two SimulationHistory objects..."
        assert self.environment_dimensions == other_hist.environment_dimensions, "The 'environment_dimensions' parameters must match between the two SimulationHistory objects..."
        assert self.environment_shape == other_hist.environment_shape, "The 'environment_shape' parameters must match between the two SimulationHistory objects..."
        assert self.environment_layer_labels == other_hist.environment_layer_labels, "The 'environment_layer_labels' parameters must match between the two SimulationHistory objects..."
        assert all(self.environment_source_position == other_hist.environment_source_position), "The 'environment_source_position' parameters must match between the two SimulationHistory objects..."
        assert self.environment_source_radius == other_hist.environment_source_radius, "The 'environment_source_radius' parameters must match between the two SimulationHistory objects..."
        assert all(self.agent_thresholds == other_hist.agent_thresholds), "The 'agent_thresholds' parameters must match between the two SimulationHistory objects..."

        # Combining arrays
        combined_start_points = np.vstack([self.start_points,
                                           other_hist.start_points])
        combined_time_shifts = np.hstack([self.time_shift,
                                          other_hist.time_shift])
        combined_reached_source = np.hstack([self.reached_source,
                                             other_hist.reached_source])
        combined_done_at_step = np.hstack([self.done_at_step,
                                           other_hist.done_at_step])

        combined_actions = []
        combined_positions = []
        combined_observations = []
        for step_i in range(max([len(self.actions), len(other_hist.actions)])):
            self_in_range = (step_i < len(self.actions))
            other_in_range = (step_i < len(other_hist.actions))
            combined_actions.append(np.vstack([self.actions[step_i] if self_in_range else np.full_like(self.actions[0], fill_value=-1),
                                               other_hist.actions[step_i] if other_in_range else np.full_like(other_hist.actions[0], fill_value=-1)]))
            combined_positions.append(np.vstack([self.positions[step_i] if self_in_range else np.full_like(self.positions[0], fill_value=-1),
                                                 other_hist.positions[step_i] if other_in_range else np.full_like(other_hist.positions[0], fill_value=-1)]))
            combined_observations.append(np.hstack([self.observations[step_i] if self_in_range else np.full_like(self.observations[0], fill_value=-1),
                                                    other_hist.observations[step_i] if other_in_range else np.full_like(other_hist.observations[0], fill_value=-1)]))

        # Creating the combined simulation history object
        combined_hist = SimulationHistory(start_points = combined_start_points,
                                          environment = self.environment,
                                          agent = self.agent,
                                          time_shift = combined_time_shifts,
                                          horizon = self.horizon,
                                          reward_discount = self.reward_discount)

        combined_hist.start_time = self.start_time
        combined_hist.actions = combined_actions
        combined_hist.positions = combined_positions
        combined_hist.observations = combined_observations
        combined_hist.reached_source = combined_reached_source
        combined_hist.done_at_step = combined_done_at_step

        return combined_hist


    def save(self,
             file: str | None = None,
             folder: str | None = None,
             save_analysis: bool = True,
             save_components: bool = False
             ) -> None:
        '''
        Function to save the simulation history to a csv file in a given folder.
        Additionally, an analysis of the runs can be saved if the save_analysis is enabled.
        The environment and agent used can be saved in the saved folder by enabling the 'save_component' parameter.

        Parameters
        ----------
        file : str, optional
            The name of the file the simulation histories will be saved to.
            If it is not provided, it will be by default "Simulations-<env_name>-n_<sim_count>-<sim_start_timestamp>-horizon_<max_sim_length>.csv"
        folder : str, optional
            Folder to save the simulation histories to.
            If the folder name is not provided the current folder will be used.
        save_analysis : bool, default=True
            Whether to save an additional csv file with an analysis of the runs of the simulation.
            It will contain the amount of steps taken, the amount of extra steps compared to optimality, the discounted rewards and the ratio between optimal trajectory and the steps taken.
            The means and standard deviations of all the runs are also computed.
            The file will have the same name as the simulation history file with an additional '-analysis' tag at the end.
        save_components : bool, default=False
            Whether or not to save the environment and agent along with the simulation histories in the given folder.
        '''
        assert (self.environment is not None) and (self.agent is not None), "Function not available, the agent and/or the environment is not set."

        # Handle file name
        if file is None:
            env_name = f's_' + '_'.join([str(axis_shape) for axis_shape in self.environment_shape])
            file = f'Simulations-{env_name}-n_{self.n}-{self.start_time.strftime("%Y%m%d_%H%M%S")}-horizon_{len(self.positions)}.csv'

        if not file.endswith('.csv'):
            file += '.csv'

        # Handle folder
        if folder is None:
            folder = './'

        if '/' not in folder:
            folder = './' + folder

        if not os.path.exists(folder):
            os.mkdir(folder)

        if not folder.endswith('/'):
            folder += '/'

        # Save components if requested
        if save_components:
            if (self.environment.saved_at is None) or (folder not in self.environment.saved_at):
                self.environment.save(folder=folder)

            if (self.agent.saved_at is None) or (folder not in self.agent.saved_at):
                self.agent.save(folder=folder)

        # Create csv file
        combined_df = pd.concat(self.simulation_dfs)

        # Adding other useful info
        padding = [None] * len(combined_df)
        combined_df['timestamps'] = [self.start_time.strftime('%Y%m%d_%H%M%S%f')] + [ts.strftime('%H%M%S%f') for ts in self.timestamps] + padding[:-(len(self.timestamps)+1)]
        combined_df['horizon'] = [self.horizon] + padding[:-1]
        combined_df['reward_discount'] = [self.reward_discount] + padding[:-1]

        environment_info = [
            self.environment.name,
            self.environment.saved_at,
            str(self.environment_dimensions), # int
            '_'.join(str(axis_size) for axis_size in self.environment_shape),
            '_'.join(str(axis_position) for axis_position in self.environment_source_position),
            str(self.environment_source_radius), # float
            '' if (self.environment_layer_labels is None) else '&'.join(self.environment_layer_labels) # Using '&' as splitter as '_' could be used in the labels themselves
        ]
        combined_df['environment'] = (environment_info + padding[:-len(environment_info)])

        # Converting the thresholds array to a string to be saved
        thresholds_string = ''
        if len(self.agent_thresholds.shape) == 2:
            thresholds_string = '-'.join(['_'.join([str(item) for item in row]) for row_i, row in enumerate(self.agent_thresholds[:,1:-1])])
        else:
            thresholds_string = '_'.join([str(item) for item in self.agent_thresholds])

        agent_info = [
            self.agent.name,
            self.agent.class_name,
            self.agent.saved_at,
            thresholds_string
        ]
        combined_df['agent'] = (agent_info + padding[:-len(agent_info)])

        # Saving csv
        combined_df.to_csv(folder + file, index=False)

        print(f'Simulations saved to: {folder + file}')

        if save_analysis:
            runs_analysis_file_name = file.replace('.csv', '-runs_analysis.csv')
            self.runs_analysis_df.to_csv(folder + runs_analysis_file_name)
            print(f"Simulation's runs analysis saved to: {folder + runs_analysis_file_name}")

            general_analysis_file_name = file.replace('.csv', '-general_analysis.csv')
            self.general_analysis_df.to_csv(folder + general_analysis_file_name)
            print(f"Simulation's general analysis saved to: {folder + general_analysis_file_name}")


    @classmethod
    def load_from_file(cls,
                       file: str,
                       environment: bool | Environment = False,
                       agent: bool | Agent = False
                       ) -> 'SimulationHistory':
        '''
        Function to load the simulation history from a file.
        This can be useful to use the plot functions on the simulations saved in succh file.

        The environment and agent can provided as a backup in the case they cannot be loaded from the file.

        Parameters
        ----------
        file : str
            A file (with the path) of the simulation histories csv. (the analysis file cannot be used for this)
        environment : bool or Environment, default=False
            If set to True, it will try to load the environment that was used for the simulation (if the save path is available).
            Or, an environment instance to be linked with the simulation history object.
        agent : bool or Agent, default=False
            If set to True, it will try to load the agent that was used for the simulation (if the save path is available).
            An agent instance to be linked with the simulation history object.

        Returns
        -------
        hist : SimulationHistory
            The loaded instance of a simulation history object.
        '''
        # Retrieving columns
        with open(file, 'r') as f:
            header = f.readline()
        columns = header.replace('\n','').split(',')

        # Setting the datatypes of columns
        column_dtypes = {col: float for col in columns}
        column_dtypes['time'] = int
        if 'layer' in columns:
            column_dtypes['layer'] = int
        column_dtypes['timestamps'] = str
        column_dtypes['environment'] = str
        column_dtypes['agent'] = str

        # Retrieving the combined dataframe
        combined_df = pd.read_csv(file, dtype=column_dtypes)

        # Retrieving horizon and reward discount
        horizon = int(combined_df['horizon'][0])
        reward_discount = combined_df['reward_discount'][0]

        # Retrieving environment
        if (not isinstance(environment, Environment)) and (environment == True):
            environment_name = combined_df['environment'][0]
            environment_path = combined_df['environment'][1]

            environment_path_check = (environment_path is not None) and (not np.isnan(environment_path))
            assert environment_path_check, "Environment was not saved at the time of the saving of the simulation history. Input an environment to the environment parameter or toggle the parameter to False."

            try:
                environment = Environment.load(environment_path)
            except:
                print(f'Failed to retrieve "{environment_name}" environment from memory')

        # Retrieving agent
        if (not isinstance(agent, Agent)) and (agent == True):
            agent_name = combined_df['environment'][0]
            agent_class = combined_df['environment'][1]
            agent_path = combined_df['environment'][2]

            agent_path_check = (agent_path is not None) and (not np.isnan(agent_path))
            assert agent_path_check, "Agent was not saved at the time of the saving of the simulation history. Input an agent to the agent parameter or toggle the parameter to False."

            try:
                class_instance = None
                for (class_name, class_obj) in inspect.getmembers(sys.modules[__name__], inspect.isclass):
                    if class_name == agent_class:
                        class_instance = class_obj
                        break
                agent = class_instance.load(combined_df['agent'][2])
            except:
                print(f'Failed to retrieve "{agent_name}" agent from memory')

        # Other attributes
        environment_dimensions = int(combined_df['environment'][2])
        environment_shape = tuple([int(axis_shape) for axis_shape in combined_df['environment'][3].split('_')])
        environment_source_position = np.array([float(pos_axis) for pos_axis in combined_df['environment'][4].split('_')])
        environment_source_radius = float(combined_df['environment'][5])
        layer_entery = combined_df['environment'][6]
        environment_layer_labels = (None if ((not isinstance(layer_entery, str)) or (len(layer_entery) == 0)) else layer_entery.split('&'))

        # Processing the threshold string
        thresholds_string = combined_df['agent'][3]
        if '-' in thresholds_string:
            rows_thresholds_string = thresholds_string.split('-')
            layer_thresholds = []
            for row in rows_thresholds_string:
                layer_thresholds.append([float(item) for item in row.split('_')])
            agent_thresholds = np.array(layer_thresholds)

        else:
            agent_thresholds = np.array([float(item) for item in thresholds_string.split('_')])

        # Columns to retrieve
        columns = [col for col in columns if col not in ['reward_discount', 'environment', 'agent']]

        # Checking how many dimensions there are
        has_layers = (((len(columns) - 5) % 2) == 1)
        dimensions = int((len(columns) - 5) / 2)

        # Recreation of list of simulations
        sim_start_rows = np.argwhere(combined_df[['reached_source']].isnull())[1:,0].tolist()
        n = (len(sim_start_rows) + 1)

        simulation_arrays = np.split(combined_df[columns].to_numpy(), sim_start_rows)
        simulation_dfs = [pd.DataFrame(sim_array, columns=columns) for sim_array in simulation_arrays]

        # Making a combined numpy array with all the simulations
        sizes = np.array([len(sim_array) for sim_array in simulation_arrays])
        max_length = sizes.max()
        paddings = max_length - sizes

        padded_simulation_arrays = [np.pad(sim_arr, ((0,pad),(0,0)), constant_values=-1) for sim_arr, pad in zip(simulation_arrays, paddings)]
        all_simulation_arrays = np.array(padded_simulation_arrays).transpose((1,0,2))

        # Timeshift
        time_shift = all_simulation_arrays[0,:,0].astype(int)

        # Gathering start states
        start_points = all_simulation_arrays[0,:,1:(1+dimensions)].astype(int)

        # Recreating action, state and observations
        positions = all_simulation_arrays[1:, :, 1:(1+dimensions)]
        actions = all_simulation_arrays[1:, :, (1+dimensions):((1+dimensions) + (1 if has_layers else 0) + dimensions)]
        observations = all_simulation_arrays[1:, :, ((1+dimensions) + (1 if has_layers else 0) + dimensions)]
        reached_source = np.array([(df['reached_source'][len(df)-1] == 1) for df in simulation_dfs])
        done_at_step = np.where((sizes-1 < horizon), sizes-1, -1)

        # Building SimulationHistory instance
        hist = cls.__new__(cls)

        hist.n = len(start_points)
        hist.environment = environment.cpu_version if isinstance(environment, Environment) else None
        hist.agent = agent.cpu_version if isinstance(agent, Agent) else None
        hist.time_shift = time_shift
        hist.horizon = horizon
        hist.reward_discount = reward_discount
        hist.start_time = datetime.strptime(combined_df['timestamps'][0], '%Y%m%d_%H%M%S%f')

        hist.start_points = start_points
        hist._running_sims = None

        hist.positions = [*positions]
        hist.actions = [*actions]
        hist.observations = [*observations]
        hist.reached_source = reached_source
        hist.done_at_step = done_at_step
        hist.timestamps = [datetime.strptime(ts, '%H%M%S%f') for ts in combined_df['timestamps'][1:max_length]]

        # Other attributes
        hist.environment_dimensions = environment_dimensions
        hist.environment_shape = environment_shape
        hist.environment_source_position = environment_source_position
        hist.environment_source_radius = environment_source_radius
        hist.environment_layer_labels = environment_layer_labels
        hist.agent_thresholds = agent_thresholds

        # Saving simulation dfs back
        hist._simulation_dfs = simulation_dfs

        return hist


    def plot(self,
             sim_id: int = 0,
             ax: plt.Axes | None = None
             ) -> None:
        '''
        Function to plot a the trajectory of a given simulation.
        An ax can be use to plot it on.

        Parameters
        ----------
        sim_id : int, default=0
            The id of the simulation to plot.
        ax : plt.Axes, optional
            The ax on which to plot the path. (If not provided, a new axis will be created)
        '''
        # TODO: Setup 3D plotting
        assert self.environment_dimensions == 2, "Plotting function only available for 2D environments for now..."

        # Generate ax is not provided
        if ax is None:
            _, ax = plt.subplots(figsize=(18,3))

        # Retrieving sim
        sim = self.simulation_dfs[sim_id]

        # Plot setup
        env_shape = self.environment_shape
        ax.imshow(np.zeros(self.environment_shape), cmap='Greys', zorder=-100)
        ax.set_xlim(0, env_shape[1])
        ax.set_ylim(env_shape[0], 0)

        # Start
        start_coord = sim[['x', 'y']].to_numpy()[0]
        ax.scatter(start_coord[0], start_coord[1], c='green', label='Start')

        # Source circle
        goal_circle = Circle(self.environment_source_position[::-1], self.environment_source_radius, color='r', fill=False, label='Source')
        ax.add_patch(goal_circle)

        # Until step
        seq = sim[['x','y']].to_numpy()

        # Path
        ax.plot(seq[:,0], seq[:,1], zorder=-1, c='black', label='Path')

        # Layer observations
        if self.environment_layer_labels is not None:
            obs_layer = sim[['layer']][1:].to_numpy()
            layer_colors = np.array(list(colors.TABLEAU_COLORS.values()))

            for layer_i, layer_label in enumerate(self.environment_layer_labels[1:]):
                layer_i += 1
                layer_mask = (obs_layer == layer_i)[:,0] # Reshaping to a single vector and not an n by 1 array
                ax.scatter(seq[1:][layer_mask,0], seq[1:][layer_mask,1], # X, Y
                           marker='x',
                           color=layer_colors[(layer_i-1) % len(layer_colors)], # Looping over the colors in case there are more layers than colors
                           zorder=2,
                           label=layer_label)

        # Process odor cues
        odor_cues = sim['o'][1:].to_numpy()
        observation_ids = None
        if self.environment.has_layers and len(self.agent_thresholds.shape) == 2:
            layer_ids = sim[['layer']][1:].to_numpy()
            action_layer_thresholds = self.agent_thresholds[layer_ids]
            observation_ids = np.argwhere((odor_cues[:,None] >= action_layer_thresholds[:,:-1]) & (odor_cues[:,None] < action_layer_thresholds[:,1:]))[:,1]
        else:
            # Setting observation ids
            observation_ids = np.argwhere((odor_cues[:,None] >= self.agent_thresholds[:-1][None,:]) & (odor_cues[:,None] < self.agent_thresholds[1:][None,:]))[:,1]

        # Check whether the odor detection is binary or by level
        odor_bins = self.agent_thresholds.shape[-1] - 1
        if odor_bins > 2:
            odor_levels = np.arange(odor_bins - 1) + 1
            for level in odor_levels:
                cues_at_level = (observation_ids == level)
                ax.scatter(seq[1:][cues_at_level,0], seq[1:][cues_at_level,1],
                           zorder=1,
                           alpha=(level / odor_bins),
                           label=f'Sensed level {level}')
        else:
            something_sensed = (observation_ids == 1)
            ax.scatter(seq[1:][something_sensed,0], seq[1:][something_sensed,1],
                       zorder=1,
                       label='Something observed')

        # Generate legend
        ax.legend()


    def plot_runtimes(self,
                      ax: plt.Axes | None = None
                      ) -> None:
        '''
        Function to plot the runtimes over the iterations.

        Parameters
        ----------
        ax : plt.Axes, optional
            The ax on which to plot the path. (If not provided, a new axis will be created)
        '''
        # Generate ax is not provided
        if ax is None:
            _, ax = plt.subplots(figsize=(18,3))

        # Computing differences
        timestamp_differences_ms = np.diff(np.array([int(ts.strftime('%H%M%S%f')) for ts in self.timestamps])) / 1000

        # Actual plot
        ax.plot(timestamp_differences_ms)

        # Axes
        ax.set_xlabel('Iteration')
        ax.set_ylabel('Runtime (ms)')


    def plot_successes(self,
                       ax: plt.Axes | None = None
                       ) -> None:
        '''
        Function to plot a 2D map of whether a given starting point was successfull or not (and whether it died early).

        Parameters
        ----------
        ax : plt.Axes, optional
            The ax on which to plot the path. (If not provided, a new axis will be created)
        '''
        assert self.environment_dimensions == 2, "Only implemented for 2D environments..."

        # Generate ax is not provided
        if ax is None:
            _, ax = plt.subplots(figsize=(18,3))

        # Setting up an empty grid of the starting points
        start_points_grid = np.zeros(self.environment_shape)

        # Compute the successful, failed and the ones that reached the horizon
        success_points = self.start_points[self.successful_simulation]
        failed_points = self.start_points[~self.successful_simulation]
        failed_not_at_horizon_points = self.start_points[~self.successful_simulation & ~self.simulations_at_horizon]

        start_points_grid[failed_points[:,0], failed_points[:,1]] = -1
        start_points_grid[success_points[:,0], success_points[:,1]] = 1

        ax.imshow(start_points_grid, cmap='RdBu')

        # The crosses where the points did not reach the horizon
        ax.scatter(failed_not_at_horizon_points[:,1], failed_not_at_horizon_points[:,0], marker='x', color='black', s=10, label='Died early')
        ax.legend()

done_count property

Returns how many simulations are terminated (whether they reached the source or not).

general_analysis_df property

A Pandas DataFrame analyzing the results of the simulations. Summarizing the performance of all the simulations with the following metrics:

• converged: Whether or not the simulation reached the source
• reached_horizon: Whether the failed simulation reached to horizon
• steps_taken: The amount of steps the agent took to reach the source, (horizon if the simulation did not reach the source)
• discounted_rewards: The discounted reward received by the agent over the course of the simulation
• extra_steps: The amount of extra steps compared to the optimal trajectory
• t_min_over_t: Normalized version of the extra steps measure, where it tends to 1 the least amount of time the agent took to reach the source compared to an optimal trajectory.

For the measures (converged, steps_taken, discounted_rewards, extra_steps, t_min_over_t), the average and standard deviations are computed in rows at the top.

runs_analysis_df property

A Pandas DataFrame analyzing the results of the simulations. It aggregates the simulations in single rows, recording:

• : The starting positions at the given axis
• optimal_steps_count: The minimal amount of steps to reach the source
• converged: Whether or not the simulation reached the source
• reached_horizon: Whether the failed simulation reached to horizon
• steps_taken: The amount of steps the agent took to reach the source, (horizon if the simulation did not reach the source)
• discounted_rewards: The discounted reward received by the agent over the course of the simulation
• extra_steps: The amount of extra steps compared to the optimal trajectory
• t_min_over_t: Normalized version of the extra steps measure, where it tends to 1 the least amount of time the agent took to reach the source compared to an optimal trajectory.

simulation_dfs property

A list of the pandas DataFrame where each dataframe is a single simulation history. Each row is a different time instant of simulation process with each column being:

• time (of the simulation data)
• [position] (z,) y, x OR x0, x1, ... xn
• (layer)
• [movement] (dz,) dy, dx OR dx0, dx1, ... dxn
• o (pure, not thresholded)
• reached_source (boolean)

simulations_at_horizon property

Returns a boolean array of which simulations reached the horizon.

success_count property

Returns how many simulations reached the source.

summary property

A string summarizing the performances of all the simulations. The metrics used are averages of:

• Step count
• Extra steps
• Discounted rewards
• Tmin / T

Along with the respective the standard deviations and equally for only for the successful simulations.

add_step(actions, next_positions, observations, reached_source, interupt)

Function to add a step in the simulation history.

Parameters:

Name	Type	Description	Default
actions	ndarray	The actions that were taken by the agents.	required
next_positions	ndarray	The positions that were reached by the agents after having taken actions.	required
observations	ndarray	The observations the agents receive after having taken actions.	required
reached_source	ndarray	A boolean array of whether each agent has reached the source or not.	required
interupt	ndarray	A boolean array of whether each agent has to be terminated even if it hasnt reached the source yet.	required

Source code in olfactory_navigation/simulation.py

def add_step(self,
             actions: np.ndarray,
             next_positions: np.ndarray,
             observations: np.ndarray,
             reached_source: np.ndarray,
             interupt: np.ndarray
             ) -> None:
    '''
    Function to add a step in the simulation history.

    Parameters
    ----------
    actions : np.ndarray
        The actions that were taken by the agents.
    next_positions : np.ndarray
        The positions that were reached by the agents after having taken actions.
    observations : np.ndarray
        The observations the agents receive after having taken actions.
    reached_source : np.ndarray
        A boolean array of whether each agent has reached the source or not.
    interupt : np.ndarray
        A boolean array of whether each agent has to be terminated even if it hasnt reached the source yet.
    '''
    self._simulation_dfs = None

    # Time tracking
    self.timestamps.append(datetime.now())

    # Check if environment if layered and/or 3D
    layered = 0 if self.environment_layer_labels is None else 1

    # Handle case cupy arrays are provided
    if gpu_support:
        actions = actions if cp.get_array_module(actions) == np else cp.asnumpy(actions)
        next_positions = next_positions if cp.get_array_module(next_positions) == np else cp.asnumpy(next_positions)
        observations = observations if cp.get_array_module(observations) == np else cp.asnumpy(observations)
        reached_source = reached_source if cp.get_array_module(reached_source) == np else cp.asnumpy(reached_source)
        interupt = interupt if cp.get_array_module(interupt) == np else cp.asnumpy(interupt)

    # Actions tracking
    action_all_sims = np.full((self.n, (layered + self.environment_dimensions)), fill_value=-1)
    action_all_sims[self._running_sims] = actions
    self.actions.append(action_all_sims)

    # Next states tracking
    next_position_all_sims = np.full((self.n, self.environment_dimensions), fill_value=-1)
    next_position_all_sims[self._running_sims] = next_positions
    self.positions.append(next_position_all_sims)

    # Observation tracking
    observation_all_sims = np.full((self.n,), fill_value=-1, dtype=float)
    observation_all_sims[self._running_sims] = observations
    self.observations.append(observation_all_sims)

    # Recording at which step the simulation is done if it is done and whether it reached the source
    self.done_at_step[self._running_sims[reached_source | interupt]] = len(self.positions)
    self.reached_source[self._running_sims[reached_source]] = True

    # Updating the list of running sims
    self._running_sims = self._running_sims[~reached_source & ~interupt]

compute_distance_to_source()

Function to compute the optimal distance to the source of each starting point according to the optimal_distance_metric attribute.

Returns:

Name	Type	Description
distance	ndarray	The optimal distances to the source point.

Source code in olfactory_navigation/simulation.py

def compute_distance_to_source(self) -> np.ndarray:
    '''
    Function to compute the optimal distance to the source of each starting point according to the optimal_distance_metric attribute.

    Returns
    -------
    distance : np.ndarray
        The optimal distances to the source point.
    '''
    point = self.start_points

    # Handling the case we have a single point
    is_single_point = (len(point.shape) == 1)
    if is_single_point:
        point = point[None,:]

    # Computing dist
    dist = None
    # if self.optimal_distance_metric == 'manhattan': # TODO Allow for other metrics to be used
    dist = np.sum(np.abs(self.environment_source_position[None,:] - point), axis=-1) - self.environment_source_radius

    if dist is None: # Meaning it was not computed
        raise NotImplementedError('This distance metric has not yet been implemented')

    return float(dist[0]) if is_single_point else dist

load_from_file(file, environment=False, agent=False) classmethod

Function to load the simulation history from a file. This can be useful to use the plot functions on the simulations saved in succh file.

The environment and agent can provided as a backup in the case they cannot be loaded from the file.

Parameters:

Name	Type	Description	Default
file	str	A file (with the path) of the simulation histories csv. (the analysis file cannot be used for this)	required
environment	bool or Environment	If set to True, it will try to load the environment that was used for the simulation (if the save path is available). Or, an environment instance to be linked with the simulation history object.	False
agent	bool or Agent	If set to True, it will try to load the agent that was used for the simulation (if the save path is available). An agent instance to be linked with the simulation history object.	False

Returns:

Name	Type	Description
hist	SimulationHistory	The loaded instance of a simulation history object.

Source code in olfactory_navigation/simulation.py

@classmethod
def load_from_file(cls,
                   file: str,
                   environment: bool | Environment = False,
                   agent: bool | Agent = False
                   ) -> 'SimulationHistory':
    '''
    Function to load the simulation history from a file.
    This can be useful to use the plot functions on the simulations saved in succh file.

    The environment and agent can provided as a backup in the case they cannot be loaded from the file.

    Parameters
    ----------
    file : str
        A file (with the path) of the simulation histories csv. (the analysis file cannot be used for this)
    environment : bool or Environment, default=False
        If set to True, it will try to load the environment that was used for the simulation (if the save path is available).
        Or, an environment instance to be linked with the simulation history object.
    agent : bool or Agent, default=False
        If set to True, it will try to load the agent that was used for the simulation (if the save path is available).
        An agent instance to be linked with the simulation history object.

    Returns
    -------
    hist : SimulationHistory
        The loaded instance of a simulation history object.
    '''
    # Retrieving columns
    with open(file, 'r') as f:
        header = f.readline()
    columns = header.replace('\n','').split(',')

    # Setting the datatypes of columns
    column_dtypes = {col: float for col in columns}
    column_dtypes['time'] = int
    if 'layer' in columns:
        column_dtypes['layer'] = int
    column_dtypes['timestamps'] = str
    column_dtypes['environment'] = str
    column_dtypes['agent'] = str

    # Retrieving the combined dataframe
    combined_df = pd.read_csv(file, dtype=column_dtypes)

    # Retrieving horizon and reward discount
    horizon = int(combined_df['horizon'][0])
    reward_discount = combined_df['reward_discount'][0]

    # Retrieving environment
    if (not isinstance(environment, Environment)) and (environment == True):
        environment_name = combined_df['environment'][0]
        environment_path = combined_df['environment'][1]

        environment_path_check = (environment_path is not None) and (not np.isnan(environment_path))
        assert environment_path_check, "Environment was not saved at the time of the saving of the simulation history. Input an environment to the environment parameter or toggle the parameter to False."

        try:
            environment = Environment.load(environment_path)
        except:
            print(f'Failed to retrieve "{environment_name}" environment from memory')

    # Retrieving agent
    if (not isinstance(agent, Agent)) and (agent == True):
        agent_name = combined_df['environment'][0]
        agent_class = combined_df['environment'][1]
        agent_path = combined_df['environment'][2]

        agent_path_check = (agent_path is not None) and (not np.isnan(agent_path))
        assert agent_path_check, "Agent was not saved at the time of the saving of the simulation history. Input an agent to the agent parameter or toggle the parameter to False."

        try:
            class_instance = None
            for (class_name, class_obj) in inspect.getmembers(sys.modules[__name__], inspect.isclass):
                if class_name == agent_class:
                    class_instance = class_obj
                    break
            agent = class_instance.load(combined_df['agent'][2])
        except:
            print(f'Failed to retrieve "{agent_name}" agent from memory')

    # Other attributes
    environment_dimensions = int(combined_df['environment'][2])
    environment_shape = tuple([int(axis_shape) for axis_shape in combined_df['environment'][3].split('_')])
    environment_source_position = np.array([float(pos_axis) for pos_axis in combined_df['environment'][4].split('_')])
    environment_source_radius = float(combined_df['environment'][5])
    layer_entery = combined_df['environment'][6]
    environment_layer_labels = (None if ((not isinstance(layer_entery, str)) or (len(layer_entery) == 0)) else layer_entery.split('&'))

    # Processing the threshold string
    thresholds_string = combined_df['agent'][3]
    if '-' in thresholds_string:
        rows_thresholds_string = thresholds_string.split('-')
        layer_thresholds = []
        for row in rows_thresholds_string:
            layer_thresholds.append([float(item) for item in row.split('_')])
        agent_thresholds = np.array(layer_thresholds)

    else:
        agent_thresholds = np.array([float(item) for item in thresholds_string.split('_')])

    # Columns to retrieve
    columns = [col for col in columns if col not in ['reward_discount', 'environment', 'agent']]

    # Checking how many dimensions there are
    has_layers = (((len(columns) - 5) % 2) == 1)
    dimensions = int((len(columns) - 5) / 2)

    # Recreation of list of simulations
    sim_start_rows = np.argwhere(combined_df[['reached_source']].isnull())[1:,0].tolist()
    n = (len(sim_start_rows) + 1)

    simulation_arrays = np.split(combined_df[columns].to_numpy(), sim_start_rows)
    simulation_dfs = [pd.DataFrame(sim_array, columns=columns) for sim_array in simulation_arrays]

    # Making a combined numpy array with all the simulations
    sizes = np.array([len(sim_array) for sim_array in simulation_arrays])
    max_length = sizes.max()
    paddings = max_length - sizes

    padded_simulation_arrays = [np.pad(sim_arr, ((0,pad),(0,0)), constant_values=-1) for sim_arr, pad in zip(simulation_arrays, paddings)]
    all_simulation_arrays = np.array(padded_simulation_arrays).transpose((1,0,2))

    # Timeshift
    time_shift = all_simulation_arrays[0,:,0].astype(int)

    # Gathering start states
    start_points = all_simulation_arrays[0,:,1:(1+dimensions)].astype(int)

    # Recreating action, state and observations
    positions = all_simulation_arrays[1:, :, 1:(1+dimensions)]
    actions = all_simulation_arrays[1:, :, (1+dimensions):((1+dimensions) + (1 if has_layers else 0) + dimensions)]
    observations = all_simulation_arrays[1:, :, ((1+dimensions) + (1 if has_layers else 0) + dimensions)]
    reached_source = np.array([(df['reached_source'][len(df)-1] == 1) for df in simulation_dfs])
    done_at_step = np.where((sizes-1 < horizon), sizes-1, -1)

    # Building SimulationHistory instance
    hist = cls.__new__(cls)

    hist.n = len(start_points)
    hist.environment = environment.cpu_version if isinstance(environment, Environment) else None
    hist.agent = agent.cpu_version if isinstance(agent, Agent) else None
    hist.time_shift = time_shift
    hist.horizon = horizon
    hist.reward_discount = reward_discount
    hist.start_time = datetime.strptime(combined_df['timestamps'][0], '%Y%m%d_%H%M%S%f')

    hist.start_points = start_points
    hist._running_sims = None

    hist.positions = [*positions]
    hist.actions = [*actions]
    hist.observations = [*observations]
    hist.reached_source = reached_source
    hist.done_at_step = done_at_step
    hist.timestamps = [datetime.strptime(ts, '%H%M%S%f') for ts in combined_df['timestamps'][1:max_length]]

    # Other attributes
    hist.environment_dimensions = environment_dimensions
    hist.environment_shape = environment_shape
    hist.environment_source_position = environment_source_position
    hist.environment_source_radius = environment_source_radius
    hist.environment_layer_labels = environment_layer_labels
    hist.agent_thresholds = agent_thresholds

    # Saving simulation dfs back
    hist._simulation_dfs = simulation_dfs

    return hist

plot(sim_id=0, ax=None)

Function to plot a the trajectory of a given simulation. An ax can be use to plot it on.

Parameters:

Name	Type	Description	Default
sim_id	int	The id of the simulation to plot.	0
ax	Axes	The ax on which to plot the path. (If not provided, a new axis will be created)	None

Source code in olfactory_navigation/simulation.py

def plot(self,
         sim_id: int = 0,
         ax: plt.Axes | None = None
         ) -> None:
    '''
    Function to plot a the trajectory of a given simulation.
    An ax can be use to plot it on.

    Parameters
    ----------
    sim_id : int, default=0
        The id of the simulation to plot.
    ax : plt.Axes, optional
        The ax on which to plot the path. (If not provided, a new axis will be created)
    '''
    # TODO: Setup 3D plotting
    assert self.environment_dimensions == 2, "Plotting function only available for 2D environments for now..."

    # Generate ax is not provided
    if ax is None:
        _, ax = plt.subplots(figsize=(18,3))

    # Retrieving sim
    sim = self.simulation_dfs[sim_id]

    # Plot setup
    env_shape = self.environment_shape
    ax.imshow(np.zeros(self.environment_shape), cmap='Greys', zorder=-100)
    ax.set_xlim(0, env_shape[1])
    ax.set_ylim(env_shape[0], 0)

    # Start
    start_coord = sim[['x', 'y']].to_numpy()[0]
    ax.scatter(start_coord[0], start_coord[1], c='green', label='Start')

    # Source circle
    goal_circle = Circle(self.environment_source_position[::-1], self.environment_source_radius, color='r', fill=False, label='Source')
    ax.add_patch(goal_circle)

    # Until step
    seq = sim[['x','y']].to_numpy()

    # Path
    ax.plot(seq[:,0], seq[:,1], zorder=-1, c='black', label='Path')

    # Layer observations
    if self.environment_layer_labels is not None:
        obs_layer = sim[['layer']][1:].to_numpy()
        layer_colors = np.array(list(colors.TABLEAU_COLORS.values()))

        for layer_i, layer_label in enumerate(self.environment_layer_labels[1:]):
            layer_i += 1
            layer_mask = (obs_layer == layer_i)[:,0] # Reshaping to a single vector and not an n by 1 array
            ax.scatter(seq[1:][layer_mask,0], seq[1:][layer_mask,1], # X, Y
                       marker='x',
                       color=layer_colors[(layer_i-1) % len(layer_colors)], # Looping over the colors in case there are more layers than colors
                       zorder=2,
                       label=layer_label)

    # Process odor cues
    odor_cues = sim['o'][1:].to_numpy()
    observation_ids = None
    if self.environment.has_layers and len(self.agent_thresholds.shape) == 2:
        layer_ids = sim[['layer']][1:].to_numpy()
        action_layer_thresholds = self.agent_thresholds[layer_ids]
        observation_ids = np.argwhere((odor_cues[:,None] >= action_layer_thresholds[:,:-1]) & (odor_cues[:,None] < action_layer_thresholds[:,1:]))[:,1]
    else:
        # Setting observation ids
        observation_ids = np.argwhere((odor_cues[:,None] >= self.agent_thresholds[:-1][None,:]) & (odor_cues[:,None] < self.agent_thresholds[1:][None,:]))[:,1]

    # Check whether the odor detection is binary or by level
    odor_bins = self.agent_thresholds.shape[-1] - 1
    if odor_bins > 2:
        odor_levels = np.arange(odor_bins - 1) + 1
        for level in odor_levels:
            cues_at_level = (observation_ids == level)
            ax.scatter(seq[1:][cues_at_level,0], seq[1:][cues_at_level,1],
                       zorder=1,
                       alpha=(level / odor_bins),
                       label=f'Sensed level {level}')
    else:
        something_sensed = (observation_ids == 1)
        ax.scatter(seq[1:][something_sensed,0], seq[1:][something_sensed,1],
                   zorder=1,
                   label='Something observed')

    # Generate legend
    ax.legend()

plot_runtimes(ax=None)

Function to plot the runtimes over the iterations.

Parameters:

Name	Type	Description	Default
ax	Axes	The ax on which to plot the path. (If not provided, a new axis will be created)	None

Source code in olfactory_navigation/simulation.py

def plot_runtimes(self,
                  ax: plt.Axes | None = None
                  ) -> None:
    '''
    Function to plot the runtimes over the iterations.

    Parameters
    ----------
    ax : plt.Axes, optional
        The ax on which to plot the path. (If not provided, a new axis will be created)
    '''
    # Generate ax is not provided
    if ax is None:
        _, ax = plt.subplots(figsize=(18,3))

    # Computing differences
    timestamp_differences_ms = np.diff(np.array([int(ts.strftime('%H%M%S%f')) for ts in self.timestamps])) / 1000

    # Actual plot
    ax.plot(timestamp_differences_ms)

    # Axes
    ax.set_xlabel('Iteration')
    ax.set_ylabel('Runtime (ms)')

plot_successes(ax=None)

Function to plot a 2D map of whether a given starting point was successfull or not (and whether it died early).

Parameters:

Name	Type	Description	Default
ax	Axes	The ax on which to plot the path. (If not provided, a new axis will be created)	None

Source code in olfactory_navigation/simulation.py

def plot_successes(self,
                   ax: plt.Axes | None = None
                   ) -> None:
    '''
    Function to plot a 2D map of whether a given starting point was successfull or not (and whether it died early).

    Parameters
    ----------
    ax : plt.Axes, optional
        The ax on which to plot the path. (If not provided, a new axis will be created)
    '''
    assert self.environment_dimensions == 2, "Only implemented for 2D environments..."

    # Generate ax is not provided
    if ax is None:
        _, ax = plt.subplots(figsize=(18,3))

    # Setting up an empty grid of the starting points
    start_points_grid = np.zeros(self.environment_shape)

    # Compute the successful, failed and the ones that reached the horizon
    success_points = self.start_points[self.successful_simulation]
    failed_points = self.start_points[~self.successful_simulation]
    failed_not_at_horizon_points = self.start_points[~self.successful_simulation & ~self.simulations_at_horizon]

    start_points_grid[failed_points[:,0], failed_points[:,1]] = -1
    start_points_grid[success_points[:,0], success_points[:,1]] = 1

    ax.imshow(start_points_grid, cmap='RdBu')

    # The crosses where the points did not reach the horizon
    ax.scatter(failed_not_at_horizon_points[:,1], failed_not_at_horizon_points[:,0], marker='x', color='black', s=10, label='Died early')
    ax.legend()

save(file=None, folder=None, save_analysis=True, save_components=False)

Function to save the simulation history to a csv file in a given folder. Additionally, an analysis of the runs can be saved if the save_analysis is enabled. The environment and agent used can be saved in the saved folder by enabling the 'save_component' parameter.

Parameters:

Name	Type	Description	Default
file	str	The name of the file the simulation histories will be saved to. If it is not provided, it will be by default "Simulations--n_--horizon_.csv"	None
folder	str	Folder to save the simulation histories to. If the folder name is not provided the current folder will be used.	None
save_analysis	bool	Whether to save an additional csv file with an analysis of the runs of the simulation. It will contain the amount of steps taken, the amount of extra steps compared to optimality, the discounted rewards and the ratio between optimal trajectory and the steps taken. The means and standard deviations of all the runs are also computed. The file will have the same name as the simulation history file with an additional '-analysis' tag at the end.	True
save_components	bool	Whether or not to save the environment and agent along with the simulation histories in the given folder.	False

Source code in olfactory_navigation/simulation.py

def save(self,
         file: str | None = None,
         folder: str | None = None,
         save_analysis: bool = True,
         save_components: bool = False
         ) -> None:
    '''
    Function to save the simulation history to a csv file in a given folder.
    Additionally, an analysis of the runs can be saved if the save_analysis is enabled.
    The environment and agent used can be saved in the saved folder by enabling the 'save_component' parameter.

    Parameters
    ----------
    file : str, optional
        The name of the file the simulation histories will be saved to.
        If it is not provided, it will be by default "Simulations-<env_name>-n_<sim_count>-<sim_start_timestamp>-horizon_<max_sim_length>.csv"
    folder : str, optional
        Folder to save the simulation histories to.
        If the folder name is not provided the current folder will be used.
    save_analysis : bool, default=True
        Whether to save an additional csv file with an analysis of the runs of the simulation.
        It will contain the amount of steps taken, the amount of extra steps compared to optimality, the discounted rewards and the ratio between optimal trajectory and the steps taken.
        The means and standard deviations of all the runs are also computed.
        The file will have the same name as the simulation history file with an additional '-analysis' tag at the end.
    save_components : bool, default=False
        Whether or not to save the environment and agent along with the simulation histories in the given folder.
    '''
    assert (self.environment is not None) and (self.agent is not None), "Function not available, the agent and/or the environment is not set."

    # Handle file name
    if file is None:
        env_name = f's_' + '_'.join([str(axis_shape) for axis_shape in self.environment_shape])
        file = f'Simulations-{env_name}-n_{self.n}-{self.start_time.strftime("%Y%m%d_%H%M%S")}-horizon_{len(self.positions)}.csv'

    if not file.endswith('.csv'):
        file += '.csv'

    # Handle folder
    if folder is None:
        folder = './'

    if '/' not in folder:
        folder = './' + folder

    if not os.path.exists(folder):
        os.mkdir(folder)

    if not folder.endswith('/'):
        folder += '/'

    # Save components if requested
    if save_components:
        if (self.environment.saved_at is None) or (folder not in self.environment.saved_at):
            self.environment.save(folder=folder)

        if (self.agent.saved_at is None) or (folder not in self.agent.saved_at):
            self.agent.save(folder=folder)

    # Create csv file
    combined_df = pd.concat(self.simulation_dfs)

    # Adding other useful info
    padding = [None] * len(combined_df)
    combined_df['timestamps'] = [self.start_time.strftime('%Y%m%d_%H%M%S%f')] + [ts.strftime('%H%M%S%f') for ts in self.timestamps] + padding[:-(len(self.timestamps)+1)]
    combined_df['horizon'] = [self.horizon] + padding[:-1]
    combined_df['reward_discount'] = [self.reward_discount] + padding[:-1]

    environment_info = [
        self.environment.name,
        self.environment.saved_at,
        str(self.environment_dimensions), # int
        '_'.join(str(axis_size) for axis_size in self.environment_shape),
        '_'.join(str(axis_position) for axis_position in self.environment_source_position),
        str(self.environment_source_radius), # float
        '' if (self.environment_layer_labels is None) else '&'.join(self.environment_layer_labels) # Using '&' as splitter as '_' could be used in the labels themselves
    ]
    combined_df['environment'] = (environment_info + padding[:-len(environment_info)])

    # Converting the thresholds array to a string to be saved
    thresholds_string = ''
    if len(self.agent_thresholds.shape) == 2:
        thresholds_string = '-'.join(['_'.join([str(item) for item in row]) for row_i, row in enumerate(self.agent_thresholds[:,1:-1])])
    else:
        thresholds_string = '_'.join([str(item) for item in self.agent_thresholds])

    agent_info = [
        self.agent.name,
        self.agent.class_name,
        self.agent.saved_at,
        thresholds_string
    ]
    combined_df['agent'] = (agent_info + padding[:-len(agent_info)])

    # Saving csv
    combined_df.to_csv(folder + file, index=False)

    print(f'Simulations saved to: {folder + file}')

    if save_analysis:
        runs_analysis_file_name = file.replace('.csv', '-runs_analysis.csv')
        self.runs_analysis_df.to_csv(folder + runs_analysis_file_name)
        print(f"Simulation's runs analysis saved to: {folder + runs_analysis_file_name}")

        general_analysis_file_name = file.replace('.csv', '-general_analysis.csv')
        self.general_analysis_df.to_csv(folder + general_analysis_file_name)
        print(f"Simulation's general analysis saved to: {folder + general_analysis_file_name}")

run_test(agent, n=None, start_points=None, environment=None, time_shift=0, time_loop=True, horizon=1000, initialization_values={}, reward_discount=0.99, print_progress=True, print_stats=True, print_warning=True, use_gpu=False, batches=-1)

Function to run n simulations for a given agent in its environment (or a given modified environment). The simulations start either from random start points or provided trough the start_points parameter. The simulation can have shifted initial times (in the olfactory simulation).

The simulation will run for at most 'horizon' steps, after which the simulations will be considered failed.

Some statistics can be printed at end of the simulation with the 'print_stats' parameter. It will print some performance statisitcs about the simulations such as the average discounter reward. The reward discount can be set by the 'reward_discount' parameter.

To speedup the simulations, it can be run on the gpu by toggling the 'use_gpu' parameter. This will have the consequence to send the various arrays to the gpu memory. This will only work if the agent has the support for to work with cupy arrays.

This method returns a SimulationHistory object that saves all the positions the agent went through, the actions the agent took, and the observation the agent received. It also provides the possibility the save the results to a csv file and plot the various trajectories.

Parameters:

Name	Type	Description	Default
agent	Agent	The agent to be tested	required
n	int	How many simulation to run in parallel. n is optional but it needs to match with what is provided in start_points.	None
start_points	ndarray	The starting points of the simulation in 2d space. If not provided, n random points will be generated based on the start probabilities of the environment. Else, the amount of start_points need to match to n, if it is provided.	None
environment	Environment	The environment to run the simulations in. By default, the environment linked to the agent will used. This parameter is intended if the environment needs to be modified compared to environment the agent was trained on.	None
time_shift	int or ndarray	The time at which to start the olfactory simulation array. It can be either a single value, or n values.	0
time_loop	bool	Whether to loop the time if reaching the end. (starts back at 0)	True
horizon	int	The amount of steps to run the simulation for before killing the remaining simulations.	1000
initialization_values	dict	In the case the agent is to be initialized with custom values, the paramaters to be passed on the initialize_state function can be set here.	{}
reward_discount	float	How much a given reward is discounted based on how long it took to get it. It is purely used to compute the Average Discount Reward (ADR) after the simulation.	0.99
print_progress	bool	Whether to show a progress bar of what step the simulations are at.	True
print_stats	bool	Whether to print the stats at the end of the run.	True
print_warning	bool	Whether to print warnings when they occur or not.	True
use_gpu	bool	Whether to run the simulations on the GPU or not.	False
batches	int	In how many batches the simulations should be run. This is useful in the case there are too many simulations and the memory can fill up. The value of batches=-1 will make it that different batches amount are tried in increasing order if a MemoryError is encountered.	-1

Returns:

Name	Type	Description
hist	SimulationHistory	A SimulationHistory object that tracked all the positions, actions and observations.

Source code in olfactory_navigation/simulation.py

def run_test(agent: Agent,
             n: int | None = None,
             start_points: np.ndarray | None = None,
             environment: Environment | None = None,
             time_shift: int | np.ndarray = 0,
             time_loop: bool = True,
             horizon: int = 1000,
             initialization_values: dict = {},
             reward_discount: float = 0.99,
             print_progress: bool = True,
             print_stats: bool = True,
             print_warning: bool = True,
             use_gpu: bool = False,
             batches: int = -1
             ) -> SimulationHistory:
    '''
    Function to run n simulations for a given agent in its environment (or a given modified environment).
    The simulations start either from random start points or provided trough the start_points parameter.
    The simulation can have shifted initial times (in the olfactory simulation).

    The simulation will run for at most 'horizon' steps, after which the simulations will be considered failed.

    Some statistics can be printed at end of the simulation with the 'print_stats' parameter.
    It will print some performance statisitcs about the simulations such as the average discounter reward.
    The reward discount can be set by the 'reward_discount' parameter.

    To speedup the simulations, it can be run on the gpu by toggling the 'use_gpu' parameter.
    This will have the consequence to send the various arrays to the gpu memory.
    This will only work if the agent has the support for to work with cupy arrays.

    This method returns a SimulationHistory object that saves all the positions the agent went through,
    the actions the agent took, and the observation the agent received.
    It also provides the possibility the save the results to a csv file and plot the various trajectories.

    Parameters
    ----------
    agent : Agent
        The agent to be tested
    n : int, optional
        How many simulation to run in parallel.
        n is optional but it needs to match with what is provided in start_points.
    start_points : np.ndarray, optional
        The starting points of the simulation in 2d space.
        If not provided, n random points will be generated based on the start probabilities of the environment.
        Else, the amount of start_points need to match to n, if it is provided.
    environment : Environment, optional
        The environment to run the simulations in.
        By default, the environment linked to the agent will used.
        This parameter is intended if the environment needs to be modified compared to environment the agent was trained on.
    time_shift : int or np.ndarray, default=0
        The time at which to start the olfactory simulation array.
        It can be either a single value, or n values.
    time_loop : bool, default=True
        Whether to loop the time if reaching the end. (starts back at 0)
    horizon : int, default=1000
        The amount of steps to run the simulation for before killing the remaining simulations.
    initialization_values : dict, default={}
        In the case the agent is to be initialized with custom values,
        the paramaters to be passed on the initialize_state function can be set here.
    reward_discount : float, default=0.99
        How much a given reward is discounted based on how long it took to get it.
        It is purely used to compute the Average Discount Reward (ADR) after the simulation.
    print_progress : bool, default=True
        Whether to show a progress bar of what step the simulations are at.
    print_stats : bool, default=True
        Whether to print the stats at the end of the run.
    print_warning : bool, default=True
        Whether to print warnings when they occur or not.
    use_gpu : bool, default=False
        Whether to run the simulations on the GPU or not.
    batches : int, default=-1
        In how many batches the simulations should be run.
        This is useful in the case there are too many simulations and the memory can fill up.
        The value of batches=-1 will make it that different batches amount are tried in increasing order if a MemoryError is encountered.

    Returns
    -------
    hist : SimulationHistory
        A SimulationHistory object that tracked all the positions, actions and observations.
    '''
    # Gathering n
    if n is None:
        if (start_points is None) or (len(start_points.shape) == 1):
            n = 1
        else:
            n = len(start_points)

    # Handle the case an specific environment is given
    if (environment is not None) and print_warning:
        if environment.shape != agent.environment.shape:
            print("[Warning] The provided environment's shape doesn't match the environment has been trained on...")
        print('Using the provided environment, not the agent environment.')
    else:
        environment = agent.environment

    # Timeshift
    if isinstance(time_shift, int):
        time_shift = np.ones(n) * time_shift
    else:
        time_shift = np.array(time_shift)
        assert time_shift.shape == (n,), f"time_shift array has a wrong shape (Given: {time_shift.shape}, expected ({n},))"
    time_shift = time_shift.astype(int)

    # Auto batches selector where the amount of batches increases if a memory error is detected
    if batches < 0:
        all_try_batches = (2**np.arange(np.log2(11000), dtype=int))
        for try_batches in all_try_batches:
            try:
                hist = run_test(agent = agent,
                                n = n,
                                start_points = start_points,
                                environment = environment,
                                time_shift = time_shift,
                                time_loop = time_loop,
                                horizon = horizon,
                                initialization_values = initialization_values,
                                reward_discount = reward_discount,
                                print_progress = print_progress,
                                print_stats = print_stats,
                                print_warning = False, # If there was any, it would have been printed already
                                use_gpu = use_gpu,
                                batches = try_batches)
                return hist
            except MemoryError as e:
                print(f'Memory full: {e}')
                print('Increasing the amount of batches...')

    # If more than one batch is selected, split the starting point arrays by the amounts of simulations in each batch
    elif batches > 1:
        # Computing the amount of simulations to be in each batch
        n_batches = np.array([n / batches] * batches).astype(int)
        n_batches[:(n%batches)] += 1
        n_start = 0

        # Full SimulationHistory object
        combined_hist = None

        # Time tracking
        all_sim_start_ts = datetime.now()

        # Batches loop
        batch_iterator = tqdm(n_batches, desc='Batches') if print_progress else n_batches
        for b_n in batch_iterator:
            b_hist = run_test(agent = agent,
                              n = b_n,
                              start_points = start_points[n_start:n_start+b_n],
                              environment = environment,
                              time_shift = time_shift[n_start:n_start+b_n],
                              time_loop = time_loop,
                              horizon = horizon,
                              initialization_values = initialization_values,
                              reward_discount = reward_discount,
                              print_progress = print_progress,
                              print_stats = False, # Forced false to not print too many things
                              print_warning = False, # If there was any, it would have been printed already
                              use_gpu = use_gpu,
                              batches = 1)
            n_start += b_n

            # Combining SimulationHistory objects
            if combined_hist is None:
                combined_hist = b_hist
            else:
                combined_hist += b_hist

        # Print stats of the complete history is asked
        if print_stats:
            all_sim_end_ts = datetime.now()
            print(f'Simulations done in {(all_sim_end_ts - all_sim_start_ts).total_seconds():.3f}s:')
            print(combined_hist.summary)

        return combined_hist

    # Move things to GPU if needed
    xp = np
    if use_gpu:
        assert gpu_support, f"GPU support is not enabled, the use_gpu option is not available."
        xp = cp

        # Move instances to GPU
        agent = agent.gpu_version
        environment = environment.gpu_version
        time_shift = cp.array(time_shift)

        if start_points is not None:
            start_points = cp.array(start_points)

    # Set start positions
    agent_position = None
    if start_points is not None:
        assert start_points.shape == (n, environment.dimensions), f'The provided start_points are of the wrong shape (expected {environment.dimensions}; received {start_points.shape[1]})'
        agent_position = start_points
    else:
        # Generating random starts
        agent_position = environment.random_start_points(n)

    # Initialize agent's state
    agent.initialize_state(n, **initialization_values)

    # Create simulation history tracker
    hist = SimulationHistory(
        start_points=agent_position,
        environment=environment,
        agent=agent,
        time_shift=time_shift,
        horizon=horizon,
        reward_discount=reward_discount
    )

    # Track begin of simulation ts
    sim_start_ts = datetime.now()

    # Simulation loop
    iterator = trange(horizon) if print_progress else range(horizon)
    for i in iterator:
        # Letting agent choose the action to take based on it's curent state
        action = agent.choose_action()

        # Updating the agent's actual position (hidden to him)
        agent_position = environment.move(pos=agent_position,
                                          movement=(action if not environment.has_layers else action[:,1:])) # Getting only the physical component of the action vector if environment has layers.

        # Get an observation based on the new position of the agent
        observation = environment.get_observation(pos=agent_position,
                                                  time=(time_shift + i),
                                                  layer=(0 if not environment.has_layers else action[:,0])) # Getting the layer information column of the action matrix.

        # Check if the source is reached
        source_reached = environment.source_reached(agent_position)

        # Add the position to the observation if the agent is space aware
        if agent.space_aware:
            observation = xp.hstack((observation[:,None], agent_position))

        # Return the observation to the agent
        update_succeeded = agent.update_state(action=action,
                                              observation=observation,
                                              source_reached=source_reached)
        if update_succeeded is None:
            update_succeeded = xp.ones(len(source_reached) , dtype=bool)

        # Handling the case where simulations have reached the end
        sims_at_end = ((time_shift + i + 1) >= (math.inf if time_loop else environment.timesteps))

        # Agents to terminate
        to_terminate = source_reached | sims_at_end | ~update_succeeded

        # Send the values to the tracker
        hist.add_step(
            actions=action,
            next_positions=agent_position,
            observations=observation[:,0] if agent.space_aware else observation,
            reached_source=source_reached,
            interupt=to_terminate
        )

        # Interupt agents that reached the end
        agent_position = agent_position[~to_terminate]
        time_shift = time_shift[~to_terminate]
        agent.kill(simulations_to_kill=to_terminate)

        # Early stopping if all agents done
        if len(agent_position) == 0:
            break

        # Update progress bar
        if print_progress:
            done_count = hist.done_count
            success_count = hist.success_count
            success_percentage = (success_count/done_count)*100 if done_count > 0 else 100
            dead_percentage = ((done_count-success_count)/done_count)*100 if done_count > 0 else 0
            iterator.set_postfix({
                'done ': f' {done_count}/{n} ({(done_count/n)*100:.1f}%)',
                'success ': f' {success_count}/{done_count} ({success_percentage:.1f}%)',
                'dead ': f' {done_count-success_count}/{done_count} ({dead_percentage:.1f}%)'
            })

    # If requested print the simulation start
    if print_stats:
        sim_end_ts = datetime.now()
        print(f'Simulations done in {(sim_end_ts - sim_start_ts).total_seconds():.3f}s:')
        print(hist.summary)

    return hist