Source code for modelly.space.continuous.space

"""A Continuous Space class."""

from __future__ import annotations

import warnings
from collections.abc import Iterable
from random import Random
from typing import TYPE_CHECKING

import numpy as np
from numpy.typing import ArrayLike

from modelly.core.agent import Agent, AgentSet
from modelly.space.continuous.queries import NeighborQueriesMixin

if TYPE_CHECKING:
    from modelly.space.continuous.agents import ContinuousSpaceAgent


[docs] class ContinuousSpace[A: ContinuousSpaceAgent](NeighborQueriesMixin): """Continuous space where each agent can have an arbitrary position.""" @property def x_min(self): # noqa: D102 # compatibility with solara_viz return self.dimensions[0, 0] @property def x_max(self): # noqa: D102 # compatibility with solara_viz return self.dimensions[0, 1] @property def y_min(self): # noqa: D102 # compatibility with solara_viz return self.dimensions[1, 0] @property def y_max(self): # noqa: D102 # compatibility with solara_viz return self.dimensions[1, 1] @property def width(self): # noqa: D102 # compatibility with solara_viz return self.size[0] @property def height(self): # noqa: D102 # compatibility with solara_viz return self.size[1] def __init__( self, dimensions: ArrayLike, torus: bool = False, random: Random | None = None, n_agents: int = 100, ) -> None: """Create a new continuous space. Args: dimensions: a numpy array like object where each row specifies the minimum and maximum value of that dimension. torus: boolean for whether the space wraps around or not random: a seeded stdlib random.Random instance n_agents: the expected number of agents in the space Internally, a numpy array is used to store the positions of all agents. This is resized if needed, but you can control the initial size explicitly by passing n_agents. """ if random is None: warnings.warn( "Random number generator not specified, this can make models non-reproducible. Please pass a random number generator explicitly", UserWarning, stacklevel=2, ) random = Random() # noqa: MDL001 (a space takes no model: only the rng passed in seeds it) self.random = random self.dimensions: np.array = np.asanyarray(dimensions) self.ndims: int = self.dimensions.shape[0] self.size: np.array = self.dimensions[:, 1] - self.dimensions[:, 0] self.center: np.array = np.sum(self.dimensions, axis=1) / 2 self.torus: bool = torus self._init_storage(n_agents) def _init_storage(self, n_agents: int) -> None: """Allocate the position array and the agent bookkeeping.""" # self._agent_positions is the array containing all agent positions # plus potential extra empty rows # agent_positions is a view into _agent_positions containing only the filled rows self._agent_positions: np.array = np.empty( (n_agents, self.dimensions.shape[0]), dtype=float ) # Initialize agent_positions as an empty view (zero rows) into _agent_positions. # This avoids copying data and ensures agent_positions always reflects the current # active slice of _agent_positions as agents are added or removed. self.agent_positions = self._agent_positions[0:0] # the list of agents in the space self.active_agents = [] self._n_agents = 0 # the number of active agents in the space # a mapping from agents to index and vice versa self._index_to_agent: dict[int, Agent] = {} # (radius, n, hash of positions) -> neighbour lists, see neighbors_within self._neighbors_cache: tuple[tuple, dict] | None = None @property def agents(self) -> AgentSet[A]: """Return an AgentSet with the agents in the space.""" return AgentSet(self.active_agents, random=self.random) def _add_agent(self, agent: Agent) -> int: """Helper method for adding an agent to the space. This method manages the numpy array with the agent positions and ensuring it is enlarged if and when needed. It is called automatically by ContinousSpaceAgent when created. """ index = self._n_agents self._n_agents += 1 current_capacity = self._agent_positions.shape[0] if current_capacity <= index: # expand using pre-allocation with a growth factor growth_factor = 0.2 n = max(int(current_capacity * growth_factor), 20) new_capacity = current_capacity + n new_positions = np.empty( (new_capacity, self.dimensions.shape[0]), dtype=float ) new_positions[:current_capacity] = self._agent_positions self._agent_positions = new_positions agent._modelly_index = index self._index_to_agent[index] = agent # we want to maintain a view rather than a copy on the active agents and positions # this is essential for the performance of the rest of this code self.active_agents.append(agent) self.agent_positions = self._agent_positions[0 : self._n_agents] return index def _remove_agent(self, agent: Agent) -> None: """Remove an agent from the space. This method is automatically called by ContinuousSpaceAgent.remove. """ index = agent._modelly_index last_index = self._n_agents - 1 # If the removed agent isn't already the last one, swap the last one into its place if index != last_index: last_agent = self.active_agents[last_index] # Swap in active_agents list self.active_agents[index] = last_agent # Swap in numpy array self._agent_positions[index] = self._agent_positions[last_index] last_agent._modelly_index = index self._index_to_agent[index] = last_agent # Pop the last elements self.active_agents.pop() self._index_to_agent.pop(last_index, None) self._n_agents -= 1 self.agent_positions = self._agent_positions[0 : self._n_agents] def __getstate__(self) -> dict: """Drop the view of the filled rows and the neighbour cache. Pickled on its own the view would become a copy, and agent positions written through it would no longer reach ``_agent_positions``. """ state = self.__dict__.copy() del state["agent_positions"] state["_neighbors_cache"] = None return state def __setstate__(self, state: dict) -> None: """Restore state and make ``agent_positions`` a view again.""" self.__dict__.update(state) self.agent_positions = self._agent_positions[0 : self._n_agents]
[docs] def calculate_difference_vector(self, point: np.ndarray, agents=None) -> np.ndarray: """Calculate the difference vector between the point and all agenents. Args: point: the point to calculate the difference vector for agents: the agents to calculate the difference vector of point with. By default, all agents are considered. """ point = np.asanyarray(point) positions = ( self.agent_positions if agents is None else self._agent_positions[[a._modelly_index for a in agents]] ) offset = positions - point if not self.torus: return offset # Across the wrap the same pair is ``size`` closer or farther; keep, per # component, whichever of the two has the smaller magnitude. wrapped = offset - np.sign(offset) * self.size return np.where(np.abs(offset) < np.abs(wrapped), offset, wrapped)
[docs] def calculate_distances( self, point: ArrayLike, agents: Iterable[Agent] | None = None, **kwargs ) -> tuple[np.ndarray, list]: """Calculate the distance between the point and all agents. Args: point: the point to calculate the difference vector for agents: the agents to calculate the difference vector of point with. By default, all agents are considered. kwargs: any additional keyword arguments are passed to scipy's cdist, which is used only if torus is False. This allows for non-Euclidian distance measures. """ point = np.asanyarray(point) if agents is None: positions = self.agent_positions agents = self.active_agents else: positions = self._agent_positions[[a._modelly_index for a in agents]] agents = np.asarray(agents) if self.torus: delta = np.abs(point - positions) delta = np.minimum(delta, self.size - delta, out=delta) # + is much faster than np.sum or array.sum dists = delta[:, 0] ** 2 for i in range(1, self.ndims): dists += delta[:, i] ** 2 dists = np.sqrt(dists) else: from scipy.spatial.distance import cdist # noqa: PLC0415 - scipy is heavy dists = cdist(point[np.newaxis, :], positions, **kwargs)[0, :] return dists, agents
[docs] def in_bounds(self, point: ArrayLike) -> bool: """Check if point is inside the bounds of the space.""" return bool( ( (np.asanyarray(point) >= self.dimensions[:, 0]) & (point <= self.dimensions[:, 1]) ).all() )
[docs] def torus_correct(self, point: ArrayLike) -> np.ndarray: """Apply a torus correction to the point.""" return self.dimensions[:, 0] + np.mod( np.asanyarray(point) - self.dimensions[:, 0], self.size )