Agent Dataclass Reference

Agent module for the AMR Hub ABM simulation.

agent

Module to represent an agent in the AMR Hub ABM simulation.

This module defines the Agent class, which represents an individual agent in the simulation. The Agent class includes attributes for the agent's location, heading, tasks, and infection status, as well as methods for moving the agent, performing tasks, and recording the agent's state over time. The module also defines related classes and enumerations, such as AgentType, InfectionStatus, and Record, to support the functionality of the Agent class.

Agent dataclass

Representation of an agent in the AMR Hub ABM simulation.

Source code in src/amr_hub_abm/agent/agent.py

@dataclass
class Agent:
    """Representation of an agent in the AMR Hub ABM simulation."""

    # ------------------------------------------------------------------------------
    idx: int
    location: Location
    heading_rad: float
    rng_generator: Generator

    interaction_radius: float = field(default=0.01)
    tasks: list[Task] = field(default_factory=list)
    agent_type: AgentType = field(default=AgentType.GENERIC)
    infection_status: InfectionStatus = field(default=InfectionStatus.SUSCEPTIBLE)
    infection_details: dict = field(default_factory=dict)

    movement_speed: float = field(default=0.001)  # units per time step
    stochasticity: float = field(default=5.0)  # degrees of randomness in movement

    trajectory_length: int = field(default=0)
    trajectory: Record = field(init=False)

    stationary: bool = field(default=False, init=False)

    # NG Added for GPU compatibility placeholder
    use_gpu: bool = field(default=False)
    target_x: float = field(default=0.0)
    target_y: float = field(default=0.0)

    # --8<--- [end:Agent]
    # ------------------------------------------------------------------------------

    @property
    def heading_degrees(self) -> float:
        """
        Get the heading of the agent in degrees.

        Returns
        -------
        float
            Heading of the agent in degrees [0, 360).

        """
        return math.degrees(self.heading_rad)

    @heading_degrees.setter
    def heading_degrees(self, value: float) -> None:
        """
        Set the heading of the agent in degrees.

        Converts to radians and ensures it is between 0 and 360 degrees.

        Parameters
        ----------
        value : float
            The heading of the agent in degrees.

        """
        self.heading_rad = math.radians(value) % (2 * math.pi)

    def __post_init__(self) -> None:
        """
        Post-initialization to setup the trajectory record and validate parameters.

        Raises
        ------
        ValueError
            If the trajectory length is negative.

        """
        # Ensure heading is between 0 and 360 degrees
        self.heading_rad = self.heading_rad % (2 * math.pi)

        logger.debug(
            "Created Agent id %s of type %s at location %s with heading %s",
            self.idx,
            self.agent_type,
            self.location,
            self.heading_rad,
        )

        if self.trajectory_length < 0:
            msg = "trajectory_length must be non-negative."
            raise NonNegativeValueError(msg)

        if self.trajectory_length > 0:
            self.trajectory = Record(total_time=self.trajectory_length)

    def __repr__(self) -> str:
        """
        Return a string representation of the agent.

        Returns
        -------
        str
            A string representation of the agent, including its index, location,
            heading, interaction radius, type, and infection status.

        """
        return (
            f"Agent(idx={self.idx}, {self.location}, "
            f"{math.degrees(self.heading_rad):.2f}°, "
            f"{self.interaction_radius}, {self.agent_type.value}, "
            f"{self.infection_status.value})"
        )

    def add_task(
        self,
        time: int,
        time_needed: int,
        location: Location,
        event_type: str,
        additional_info: dict | None = None,
    ) -> None:
        """
        Add a task to the agent's task list and log the addition.

        Parameters
        ----------
        time : int
            The time at which the task is added.
        location : Location
            The location associated with the task.
        event_type : str
            The type of task to add.
        additional_info : dict | None, optional
            Additional information required for certain task types.

        Raises
        ------
        SimulationModeError
            If the event_type is invalid.

        """
        if event_type not in TASK_TYPES:
            msg = f"Invalid task type: {event_type}. Must be one of {TASK_TYPES}."
            raise SimulationModeError(msg)
        task_type = TaskType[event_type.upper()]
        if task_type == TaskType.GENERIC:
            msg = "Task type GENERIC not implemented yet."
            raise NotImplementedError(msg)

        context = build_task_context(
            time=time,
            time_needed=time_needed,
            location=location,
            additional_info=additional_info,
        )

        task = TASK_BUILDERS[task_type](context)
        self.tasks.append(task)

    def perform_task(
        self, current_time: int, engine: SpatialQuery, *, record: bool = False
    ) -> None:
        """
        Perform the agent's current task if it's due.

        Parameters
        ----------
        current_time : int
            The current time step in the simulation.
        engine : SpatialQuery
            The engine instance used to resolve geometry queries and bounds.
        record : bool, optional
            Whether to record the agent's state at the current time step.

        """
        if record:
            logger.info(
                "Recording state for Agent id %s at time %s: location=%s",
                self.idx,
                current_time,
                self.location,
            )
            record_state(agent=self, current_time=current_time)

        if not self.tasks:
            return

        task_handlers = (
            perform_in_progress_task,
            perform_moving_to_task_location,
            perform_suspended_task,
            perform_to_be_started_task,
        )

        for handler in task_handlers:
            # We now pass the engine dynamically down into the task handlers!
            if handler(self, current_time=current_time, engine=engine):
                return

    def attempt_task_insertion(
        self,
        next_task: Task,
        next_task_move_time: float,
        current_time: int,
        engine: SpatialQuery,
        task_durations: TaskDurationConfig,
    ) -> None:
        """
        Attempt to insert a task to occupy an empty chair.

        Parameters
        ----------
        next_task : Task
            The next task for which to attempt insertion of an occupy_content task.
        next_task_move_time : float
            The time at which the next task is scheduled to move to the next stage.
        current_time : int
            The current time step in the simulation.
        engine : SpatialQuery
            The SpatialQuery instance for resolving spatial queries.

        """
        if isinstance(next_task, TaskOccupyContent):
            return
        if self.stationary:
            return

        room = engine.get_room(self)
        if room is None:
            logger.info(
                "Agent id %s is not located in any room. Cannot check for "
                "empty chairs for task %s.",
                self.idx,
                next_task.task_type.name,
            )
            return

        empty_chairs = [
            content
            for content in room.contents
            if content.content_type == ContentType.CHAIR and content.available
        ]

        logger.info(
            "Agent id %s found %s empty chairs in room %s for task %s.",
            self.idx,
            len(empty_chairs),
            room.name,
            next_task.task_type.name,
        )

        if empty_chairs:
            chair = empty_chairs[0]
            estimated_time_to_chair = engine.estimate_time_to_reach_location(
                self, chair.location
            )
            if current_time + estimated_time_to_chair < next_task_move_time:
                self.add_task(
                    time=current_time,
                    time_needed=task_durations.time_needed_occupy_content,
                    location=chair.location,
                    event_type="occupy_content",
                    additional_info={
                        "content_type": ContentType.CHAIR,
                        "room": room,
                        "content": chair,
                    },
                )
                logger.info(
                    """
                    Current time: %s
                    Estimated time to chair: %s
                    Next task move time: %s.
                    Agent id %s inserted occupy_content task for chair at location %s
                    to be performed before next task move time %s.
                    """,
                    current_time,
                    estimated_time_to_chair,
                    next_task_move_time,
                    self.idx,
                    chair.location,
                    next_task_move_time,
                )
                logger.warning(
                    """
                    Length of task list %s at time %s.
                    """,
                    len(self.tasks),
                    current_time,
                )

heading_degrees property writable

Get the heading of the agent in degrees.

Returns:

Type	Description
float	Heading of the agent in degrees [0, 360).

__post_init__()

Post-initialization to setup the trajectory record and validate parameters.

Raises:

Type	Description
ValueError	If the trajectory length is negative.

Source code in src/amr_hub_abm/agent/agent.py

def __post_init__(self) -> None:
    """
    Post-initialization to setup the trajectory record and validate parameters.

    Raises
    ------
    ValueError
        If the trajectory length is negative.

    """
    # Ensure heading is between 0 and 360 degrees
    self.heading_rad = self.heading_rad % (2 * math.pi)

    logger.debug(
        "Created Agent id %s of type %s at location %s with heading %s",
        self.idx,
        self.agent_type,
        self.location,
        self.heading_rad,
    )

    if self.trajectory_length < 0:
        msg = "trajectory_length must be non-negative."
        raise NonNegativeValueError(msg)

    if self.trajectory_length > 0:
        self.trajectory = Record(total_time=self.trajectory_length)

__repr__()

Return a string representation of the agent.

Returns:

Type	Description
str	A string representation of the agent, including its index, location, heading, interaction radius, type, and infection status.

Source code in src/amr_hub_abm/agent/agent.py

def __repr__(self) -> str:
    """
    Return a string representation of the agent.

    Returns
    -------
    str
        A string representation of the agent, including its index, location,
        heading, interaction radius, type, and infection status.

    """
    return (
        f"Agent(idx={self.idx}, {self.location}, "
        f"{math.degrees(self.heading_rad):.2f}°, "
        f"{self.interaction_radius}, {self.agent_type.value}, "
        f"{self.infection_status.value})"
    )

add_task(time, time_needed, location, event_type, additional_info=None)

Add a task to the agent's task list and log the addition.

Parameters:

Name	Type	Description	Default
time	int	The time at which the task is added.	required
location	Location	The location associated with the task.	required
event_type	str	The type of task to add.	required
additional_info	dict | None	Additional information required for certain task types.	None

Raises:

Type	Description
SimulationModeError	If the event_type is invalid.

Source code in src/amr_hub_abm/agent/agent.py

def add_task(
    self,
    time: int,
    time_needed: int,
    location: Location,
    event_type: str,
    additional_info: dict | None = None,
) -> None:
    """
    Add a task to the agent's task list and log the addition.

    Parameters
    ----------
    time : int
        The time at which the task is added.
    location : Location
        The location associated with the task.
    event_type : str
        The type of task to add.
    additional_info : dict | None, optional
        Additional information required for certain task types.

    Raises
    ------
    SimulationModeError
        If the event_type is invalid.

    """
    if event_type not in TASK_TYPES:
        msg = f"Invalid task type: {event_type}. Must be one of {TASK_TYPES}."
        raise SimulationModeError(msg)
    task_type = TaskType[event_type.upper()]
    if task_type == TaskType.GENERIC:
        msg = "Task type GENERIC not implemented yet."
        raise NotImplementedError(msg)

    context = build_task_context(
        time=time,
        time_needed=time_needed,
        location=location,
        additional_info=additional_info,
    )

    task = TASK_BUILDERS[task_type](context)
    self.tasks.append(task)

attempt_task_insertion(next_task, next_task_move_time, current_time, engine, task_durations)

Attempt to insert a task to occupy an empty chair.

Parameters:

Name	Type	Description	Default
next_task	Task	The next task for which to attempt insertion of an occupy_content task.	required
next_task_move_time	float	The time at which the next task is scheduled to move to the next stage.	required
current_time	int	The current time step in the simulation.	required
engine	SpatialQuery	The SpatialQuery instance for resolving spatial queries.	required

Source code in src/amr_hub_abm/agent/agent.py

def attempt_task_insertion(
    self,
    next_task: Task,
    next_task_move_time: float,
    current_time: int,
    engine: SpatialQuery,
    task_durations: TaskDurationConfig,
) -> None:
    """
    Attempt to insert a task to occupy an empty chair.

    Parameters
    ----------
    next_task : Task
        The next task for which to attempt insertion of an occupy_content task.
    next_task_move_time : float
        The time at which the next task is scheduled to move to the next stage.
    current_time : int
        The current time step in the simulation.
    engine : SpatialQuery
        The SpatialQuery instance for resolving spatial queries.

    """
    if isinstance(next_task, TaskOccupyContent):
        return
    if self.stationary:
        return

    room = engine.get_room(self)
    if room is None:
        logger.info(
            "Agent id %s is not located in any room. Cannot check for "
            "empty chairs for task %s.",
            self.idx,
            next_task.task_type.name,
        )
        return

    empty_chairs = [
        content
        for content in room.contents
        if content.content_type == ContentType.CHAIR and content.available
    ]

    logger.info(
        "Agent id %s found %s empty chairs in room %s for task %s.",
        self.idx,
        len(empty_chairs),
        room.name,
        next_task.task_type.name,
    )

    if empty_chairs:
        chair = empty_chairs[0]
        estimated_time_to_chair = engine.estimate_time_to_reach_location(
            self, chair.location
        )
        if current_time + estimated_time_to_chair < next_task_move_time:
            self.add_task(
                time=current_time,
                time_needed=task_durations.time_needed_occupy_content,
                location=chair.location,
                event_type="occupy_content",
                additional_info={
                    "content_type": ContentType.CHAIR,
                    "room": room,
                    "content": chair,
                },
            )
            logger.info(
                """
                Current time: %s
                Estimated time to chair: %s
                Next task move time: %s.
                Agent id %s inserted occupy_content task for chair at location %s
                to be performed before next task move time %s.
                """,
                current_time,
                estimated_time_to_chair,
                next_task_move_time,
                self.idx,
                chair.location,
                next_task_move_time,
            )
            logger.warning(
                """
                Length of task list %s at time %s.
                """,
                len(self.tasks),
                current_time,
            )

perform_task(current_time, engine, *, record=False)

Perform the agent's current task if it's due.

Parameters:

Name	Type	Description	Default
current_time	int	The current time step in the simulation.	required
engine	SpatialQuery	The engine instance used to resolve geometry queries and bounds.	required
record	bool	Whether to record the agent's state at the current time step.	False

Source code in src/amr_hub_abm/agent/agent.py

def perform_task(
    self, current_time: int, engine: SpatialQuery, *, record: bool = False
) -> None:
    """
    Perform the agent's current task if it's due.

    Parameters
    ----------
    current_time : int
        The current time step in the simulation.
    engine : SpatialQuery
        The engine instance used to resolve geometry queries and bounds.
    record : bool, optional
        Whether to record the agent's state at the current time step.

    """
    if record:
        logger.info(
            "Recording state for Agent id %s at time %s: location=%s",
            self.idx,
            current_time,
            self.location,
        )
        record_state(agent=self, current_time=current_time)

    if not self.tasks:
        return

    task_handlers = (
        perform_in_progress_task,
        perform_moving_to_task_location,
        perform_suspended_task,
        perform_to_be_started_task,
    )

    for handler in task_handlers:
        # We now pass the engine dynamically down into the task handlers!
        if handler(self, current_time=current_time, engine=engine):
            return

enums

Module defining enumerations for agent types and infection statuses.

AgentType

Bases: IntEnum

Enumeration of possible agent types.

Source code in src/amr_hub_abm/agent/enums.py

class AgentType(IntEnum):
    """Enumeration of possible agent types."""

    GENERIC = 0
    PATIENT = 1
    HEALTHCARE_WORKER = 2

InfectionStatus

Bases: IntEnum

Enumeration of possible infection statuses.

Source code in src/amr_hub_abm/agent/enums.py

class InfectionStatus(IntEnum):
    """Enumeration of possible infection statuses."""

    SUSCEPTIBLE = 0
    EXPOSED = 1
    INFECTED = 2
    RECOVERED = 3

kinematics

Module defining agent kinematics configuration for the AMR Hub ABM simulation.

AgentKinematicsConfig dataclass

Agent kinematics parameters read from the simulation configuration file.

Parameters:

Name	Type	Description	Default
movement_speed	float	Agent movement speed, in units per time step.	required
stochasticity	float	Degrees of randomness applied to an agent's heading at each step.	required
interaction_radius	float	Distance within which an agent is considered to have reached a target.	required
max_movement_attempts	int	Maximum number of attempts to find a movement step that avoids wall intersections.	required

Source code in src/amr_hub_abm/agent/kinematics.py

@dataclass(frozen=True)
class AgentKinematicsConfig:
    """
    Agent kinematics parameters read from the simulation configuration file.

    Parameters
    ----------
    movement_speed : float
        Agent movement speed, in units per time step.
    stochasticity : float
        Degrees of randomness applied to an agent's heading at each step.
    interaction_radius : float
        Distance within which an agent is considered to have reached a target.
    max_movement_attempts : int
        Maximum number of attempts to find a movement step that avoids wall
        intersections.

    """

    movement_speed: float
    stochasticity: float
    interaction_radius: float
    max_movement_attempts: int

    @classmethod
    def from_config(cls, config_data: dict[str, Any]) -> AgentKinematicsConfig:
        """
        Build an ``AgentKinematicsConfig`` from parsed YAML configuration data.

        There are deliberately no fallback defaults here: agent kinematics must
        be defined explicitly in the simulation config file.

        Parameters
        ----------
        config_data : dict[str, Any]
            The parsed simulation configuration.

        Returns
        -------
        AgentKinematicsConfig
            The agent kinematics parameters read from the config.

        Raises
        ------
        InvalidDefinitionError
            If one or more required agent kinematics keys are missing from
            the config.

        """
        missing_keys = [
            key for key in REQUIRED_KINEMATICS_KEYS if key not in config_data
        ]
        if missing_keys:
            msg = (
                "Missing required agent kinematics config key(s): "
                f"{', '.join(missing_keys)}."
            )
            raise InvalidDefinitionError(msg)

        return cls(
            movement_speed=config_data["agent_movement_speed"],
            stochasticity=config_data["agent_stochasticity"],
            interaction_radius=config_data["agent_interaction_radius"],
            max_movement_attempts=config_data["agent_max_movement_attempts"],
        )

from_config(config_data) classmethod

Build an AgentKinematicsConfig from parsed YAML configuration data.

There are deliberately no fallback defaults here: agent kinematics must be defined explicitly in the simulation config file.

Parameters:

Name	Type	Description	Default
config_data	dict[str, Any]	The parsed simulation configuration.	required

Returns:

Type	Description
AgentKinematicsConfig	The agent kinematics parameters read from the config.

Raises:

Type	Description
InvalidDefinitionError	If one or more required agent kinematics keys are missing from the config.

Source code in src/amr_hub_abm/agent/kinematics.py

@classmethod
def from_config(cls, config_data: dict[str, Any]) -> AgentKinematicsConfig:
    """
    Build an ``AgentKinematicsConfig`` from parsed YAML configuration data.

    There are deliberately no fallback defaults here: agent kinematics must
    be defined explicitly in the simulation config file.

    Parameters
    ----------
    config_data : dict[str, Any]
        The parsed simulation configuration.

    Returns
    -------
    AgentKinematicsConfig
        The agent kinematics parameters read from the config.

    Raises
    ------
    InvalidDefinitionError
        If one or more required agent kinematics keys are missing from
        the config.

    """
    missing_keys = [
        key for key in REQUIRED_KINEMATICS_KEYS if key not in config_data
    ]
    if missing_keys:
        msg = (
            "Missing required agent kinematics config key(s): "
            f"{', '.join(missing_keys)}."
        )
        raise InvalidDefinitionError(msg)

    return cls(
        movement_speed=config_data["agent_movement_speed"],
        stochasticity=config_data["agent_stochasticity"],
        interaction_radius=config_data["agent_interaction_radius"],
        max_movement_attempts=config_data["agent_max_movement_attempts"],
    )

output

Module containing functions related to agent output.

Record dataclass

Representation of a record of an agent's state at a given time step.

Parameters:

Name	Type	Description	Default
total_time	int	The total number of time steps for which to record the agent's state.	required

Source code in src/amr_hub_abm/agent/output.py

@dataclass(slots=True)
class Record:
    """
    Representation of a record of an agent's state at a given time step.

    Parameters
    ----------
    total_time : int
        The total number of time steps for which to record the agent's state.

    """

    total_time: int

    building: npt.NDArray[np.int8] = field(init=False)
    floor: npt.NDArray[np.int8] = field(init=False)
    position: npt.NDArray[np.float64] = field(init=False)
    heading: npt.NDArray[np.float64] = field(init=False)
    infection_status: npt.NDArray[np.int8] = field(init=False)

    def __post_init__(self) -> None:
        """
        Simply initialises the numpy arrays of the class to be empty.

        This includes the building, floor, position, heading, and infection status.
        All of these are stored as integers and floats (and not strings) for memory
        efficiency
        """
        self.building = np.empty(self.total_time, dtype=np.int8)
        self.floor = np.empty(self.total_time, dtype=np.int8)
        self.position = np.full((self.total_time, 2), np.nan, dtype=np.float64)
        self.heading = np.empty((self.total_time, 1), dtype=np.float64)
        self.infection_status = np.empty(self.total_time, dtype=np.int8)

    def push(  # noqa: PLR0913
        self,
        time: int,
        building_idx: int,
        floor: int,
        pos_x: float,
        pos_y: float,
        heading: float,
        infection_status: InfectionStatus,
    ) -> None:
        """
        Push a new record of the agent's state at a given time step.

        Parameters
        ----------
        time : int
            The time step for which to record the agent's state.
        building_idx : int
            The index of the building in which the agent is located.
        floor : int
            The floor number on which the agent is located.
        pos_x : float
            The x-coordinate of the agent's position.
        pos_y : float
            The y-coordinate of the agent's position.
        heading : float
            The heading of the agent in radians.
        infection_status : InfectionStatus
            The infection status of the agent.

        Raises
        ------
        ValueError
            If the time step exceeds the total_time for the record.

        """
        if time >= self.total_time:
            msg = f"Time {time} exceeds total_time {self.total_time} for record."
            raise ValueError(msg)

        self.building[time] = building_idx
        self.floor[time] = floor
        self.heading[time] = heading
        self.position[time] = [pos_x, pos_y]
        self.infection_status[time] = infection_status.value

__post_init__()

Simply initialises the numpy arrays of the class to be empty.

This includes the building, floor, position, heading, and infection status. All of these are stored as integers and floats (and not strings) for memory efficiency

Source code in src/amr_hub_abm/agent/output.py

def __post_init__(self) -> None:
    """
    Simply initialises the numpy arrays of the class to be empty.

    This includes the building, floor, position, heading, and infection status.
    All of these are stored as integers and floats (and not strings) for memory
    efficiency
    """
    self.building = np.empty(self.total_time, dtype=np.int8)
    self.floor = np.empty(self.total_time, dtype=np.int8)
    self.position = np.full((self.total_time, 2), np.nan, dtype=np.float64)
    self.heading = np.empty((self.total_time, 1), dtype=np.float64)
    self.infection_status = np.empty(self.total_time, dtype=np.int8)

push(time, building_idx, floor, pos_x, pos_y, heading, infection_status)

Push a new record of the agent's state at a given time step.

Parameters:

Name	Type	Description	Default
time	int	The time step for which to record the agent's state.	required
building_idx	int	The index of the building in which the agent is located.	required
floor	int	The floor number on which the agent is located.	required
pos_x	float	The x-coordinate of the agent's position.	required
pos_y	float	The y-coordinate of the agent's position.	required
heading	float	The heading of the agent in radians.	required
infection_status	InfectionStatus	The infection status of the agent.	required

Raises:

Type	Description
ValueError	If the time step exceeds the total_time for the record.

Source code in src/amr_hub_abm/agent/output.py

def push(  # noqa: PLR0913
    self,
    time: int,
    building_idx: int,
    floor: int,
    pos_x: float,
    pos_y: float,
    heading: float,
    infection_status: InfectionStatus,
) -> None:
    """
    Push a new record of the agent's state at a given time step.

    Parameters
    ----------
    time : int
        The time step for which to record the agent's state.
    building_idx : int
        The index of the building in which the agent is located.
    floor : int
        The floor number on which the agent is located.
    pos_x : float
        The x-coordinate of the agent's position.
    pos_y : float
        The y-coordinate of the agent's position.
    heading : float
        The heading of the agent in radians.
    infection_status : InfectionStatus
        The infection status of the agent.

    Raises
    ------
    ValueError
        If the time step exceeds the total_time for the record.

    """
    if time >= self.total_time:
        msg = f"Time {time} exceeds total_time {self.total_time} for record."
        raise ValueError(msg)

    self.building[time] = building_idx
    self.floor[time] = floor
    self.heading[time] = heading
    self.position[time] = [pos_x, pos_y]
    self.infection_status[time] = infection_status.value

record_state(agent, current_time)

Push a record of the agent's current state to the trajectory.

Parameters:

Name	Type	Description	Default
current_time	int	The current time step in the simulation for which to record the agent's state.	required

Raises:

Type	Description
ValueError	If the current_time exceeds the trajectory_length of the agent.

Source code in src/amr_hub_abm/agent/output.py

def record_state(agent: Agent, current_time: int) -> None:
    """
    Push a record of the agent's current state to the trajectory.

    Parameters
    ----------
    current_time : int
        The current time step in the simulation for which to record the agent's
        state.

    Raises
    ------
    ValueError
        If the current_time exceeds the trajectory_length of the agent.

    """
    if current_time >= agent.trajectory_length:
        msg = f"Current time {current_time} "
        msg += f"exceeds trajectory length {agent.trajectory_length}."
        raise ValueError(msg)

    if agent.location.building is None:
        msg = "Agent is not located in any building. Cannot record state."
        raise ValueError(msg)

    building_idx = hash(agent.location.building) % 128  # Convert to int8 range

    agent.trajectory.push(
        time=current_time,
        building_idx=building_idx,
        floor=agent.location.floor,
        pos_x=agent.location.x,
        pos_y=agent.location.y,
        heading=agent.heading_rad,
        infection_status=agent.infection_status,
    )

plotter

Module containing plotting functions for agents.

plot_agent(agent, ax, *, show_tags=True)

Plot the agent on the given axes.

Parameters:

Name	Type	Description	Default
ax	Axes	The axes on which to plot the agent.	required
show_tags	bool	Whether to show tags with the agent's type and index.	True

Source code in src/amr_hub_abm/agent/plotter.py

def plot_agent(agent: Agent, ax: Axes, *, show_tags: bool = True) -> None:
    """
    Plot the agent on the given axes.

    Parameters
    ----------
    ax : Axes
        The axes on which to plot the agent.
    show_tags : bool, optional
        Whether to show tags with the agent's type and index.

    """
    ring_colour = INFECTION_RING_COLOUR[agent.infection_status]
    if ring_colour is not None:
        ax.plot(
            agent.location.x,
            agent.location.y,
            marker="o",
            markersize=12,  # bigger than the inner dot
            markerfacecolor="none",  # hollow ring
            markeredgecolor=ring_colour,
            markeredgewidth=2,
            zorder=2,
        )

    ax.plot(
        agent.location.x,
        agent.location.y,
        marker="o",
        markersize=5,
        color=ROLE_COLOUR_MAP[agent.agent_type],
    )

    if not show_tags:
        return

    # Build the multi-line label
    role = agent.agent_type.name.replace("_", " ").title()
    lines = [f"{role} {agent.idx}"]

    if agent.infection_status != InfectionStatus.SUSCEPTIBLE:
        lines.append(f"({agent.infection_status.name.lower()})")

    # Find what to display: in-progress task takes priority, else next NOT_STARTED
    in_progress = next(
        (t for t in agent.tasks if t.progress == TaskProgress.IN_PROGRESS),
        None,
    )
    moving = next(
        (t for t in agent.tasks if t.progress == TaskProgress.MOVING_TO_LOCATION),
        None,
    )
    upcoming = [t for t in agent.tasks if t.progress == TaskProgress.NOT_STARTED]
    next_upcoming = (
        min(upcoming, key=lambda t: (t.time_due, t.priority.value))
        if upcoming
        else None
    )

    display_task = in_progress or moving or next_upcoming
    if display_task is not None:
        task_name = display_task.task_type.name.lower()
        if isinstance(display_task, TaskAttendPatient):
            assert display_task.patient is not None  # noqa: S101
            task_name += f" → patient {display_task.patient.idx}"

        if display_task.progress == TaskProgress.IN_PROGRESS:
            prefix = "doing"
        elif display_task.progress == TaskProgress.MOVING_TO_LOCATION:
            prefix = "moving to"
        else:
            prefix = "next"
        lines.append(f"[{prefix}: {task_name}]")

    ax.text(
        agent.location.x + 0.1,
        agent.location.y + 0.05,
        "\n".join(lines),
        fontsize=7,
        ha="left",
        va="bottom",
    )

plot_agents_in_room(room, ax, agents, *, trajectory=False)

Plot agents located inside a room.

Source code in src/amr_hub_abm/agent/plotter.py

def plot_agents_in_room(
    room: Room,
    ax: Axes,
    agents: list[Agent],
    *,
    trajectory: bool = False,
) -> None:
    """Plot agents located inside a room."""
    if not room.walls:
        msg = "Cannot plot room without walls. Please provide wall definitions."
        raise InvalidRoomError(msg)

    for agent in agents:
        if (
            agent.location.building == room.building
            and agent.location.floor == room.floor
            and room.contains_point((agent.location.x, agent.location.y))
        ):
            plot_agent(agent, ax)

            if trajectory and agent.agent_type == AgentType.HEALTHCARE_WORKER:
                plot_trajectory(agent, ax)

plot_trajectory(agent, ax, current_time=None)

Plot the agent's trajectory on the given axes.

Parameters:

Name	Type	Description	Default
ax	Axes	The axes on which to plot the agent's trajectory.	required

Source code in src/amr_hub_abm/agent/plotter.py

def plot_trajectory(agent: Agent, ax: Axes, current_time: int | None = None) -> None:
    """
    Plot the agent's trajectory on the given axes.

    Parameters
    ----------
    ax : Axes
        The axes on which to plot the agent's trajectory.

    """
    if agent.trajectory_length == 0:
        msg = "Cannot plot trajectory for agent with trajectory_length of 0."
        raise SimulationModeError(msg)

    end = current_time if current_time is not None else agent.trajectory_length
    if end <= 0:
        return

    ax.plot(
        agent.trajectory.position[:, 0],
        agent.trajectory.position[:, 1],
        linestyle="-",
        linewidth=1.5,
        color=ROLE_COLOUR_MAP[agent.agent_type],
        alpha=0.7,
    )

utils

Collection of utility functions for the agent module.

add_agent_occupancy(agent, content, current_time, engine)

Add the agent's occupancy to the specified content.

Parameters:

Name	Type	Description	Default
agent	Agent	The agent whose occupancy is to be added.	required
content	Content	The content to which the agent will occupy.	required
current_time	int	The current time in the simulation, used for logging purposes.	required
engine	SpatialQuery	The engine instance used to resolve geometry queries.	required

Returns:

Type	Description
bool	Whether occupancy was acquired.

Source code in src/amr_hub_abm/agent/utils.py

def add_agent_occupancy(
    agent: Agent, content: Content, current_time: int, engine: SpatialQuery
) -> bool:
    """
    Add the agent's occupancy to the specified content.

    Parameters
    ----------
    agent : Agent
        The agent whose occupancy is to be added.
    content : Content
        The content to which the agent will occupy.
    current_time : int
        The current time in the simulation, used for logging purposes.
    engine : SpatialQuery
        The engine instance used to resolve geometry queries.

    Returns
    -------
    bool
        Whether occupancy was acquired.

    """
    agent_id = (agent.idx, agent.agent_type)
    if not content.try_occupy(agent_id):
        logger.warning(
            "Agent id %s could not occupy content id %s because it is claimed "
            "by another agent.",
            agent.idx,
            content.content_id,
        )
        return False

    agent.stationary = True

    room = engine.get_room(agent)
    room_name = "unknown" if room is None else room.name

    logger.info(
        """
        Agent id %s added occupancy to content id %s of type %s
        in room %s at time %d.
        """,
        agent.idx,
        content.content_id,
        content.content_type,
        room_name,
        current_time,
    )
    return True

remove_agent_occupancy(agent, current_time, engine)

Remove the agent's occupancy from any content they are currently occupying.

Parameters:

Name	Type	Description	Default
agent	Agent	The agent whose occupancy is to be removed.	required
current_time	int	The current time in the simulation, used for logging purposes.	required
engine	SpatialQuery	The engine instance used to resolve geometry queries.	required

Source code in src/amr_hub_abm/agent/utils.py

def remove_agent_occupancy(
    agent: Agent, current_time: int, engine: SpatialQuery
) -> None:
    """
    Remove the agent's occupancy from any content they are currently occupying.

    Parameters
    ----------
    agent : Agent
        The agent whose occupancy is to be removed.
    current_time : int
        The current time in the simulation, used for logging purposes.
    engine : SpatialQuery
        The engine instance used to resolve geometry queries.

    """
    room = engine.get_room(agent)
    if room is None:
        return
    for content in room.contents:
        if content.release_occupancy((agent.idx, agent.agent_type)):
            agent.stationary = False
            logger.info(
                """
                Agent id %s removed occupancy from content id %s of type %s
                in room %s at time %d.
                """,
                agent.idx,
                content.content_id,
                content.content_type,
                room.name,
                current_time,
            )
            return