mesa_frames.agent

 1from typing import TYPE_CHECKING, Optional
 2
 3import numpy as np
 4import pandas as pd
 5from numpy.random import randint
 6
 7if TYPE_CHECKING:
 8    from mesa_frames.model import ModelDF
 9
10
11class AgentDF:
12    """The AgentDF class is the base class for other agents.
13    It should be used as inherited class for new agents classes.
14
15    Attributes
16    ----------
17    dtypes : dict[str, str]
18        The attributes of the Agent as a dictionary of columns and data types. It contains:
19        - id : int64
20            -- The unique id of the Agent.
21        - type : str
22            -- The type of the Agent.
23    model : Optional['ModelDF']
24        The model of the simulation where the Agent is used. See src/model.py. Default: None
25    mask : pd.Series | None
26        The mask of the agents dataframe in the model which corresponds to the Agent class.
27        Initialized when model is created. Default: None
28    """
29
30    dtypes: dict[str, str] = {
31        "id": "int64",
32        "type": "str",
33    }
34    model: Optional["ModelDF"] = None
35    mask: pd.Series | None = None
36
37    @classmethod
38    def __init__(cls):
39        """Initializes the Agent class.
40        Assigns a 64-bit random id to each agent in the model.
41        """
42        if cls.mask is None or cls.model is None or cls.model.agents is None:
43            raise ValueError(
44                "The Agent classes have not been initialized. Please use the model.create_agents() method to initialize the mask."
45            )
46        cls.model.agents.loc[cls.mask, "id"] = np.random.randint(
47            low=-9223372036854775808,
48            high=9223372036854775807,
49            size=cls.mask.sum(),
50            dtype="int64",
51        )
52
53    @classmethod
54    def step(cls):
55        """The step method of the Agent class.
56        It should be decorated as @classmethod and should act on the mask portion of the agents dataframe of the model.
57        """
58        pass
59
60
61class GeoAgentDF(AgentDF):
62    """The GeoAgentDF extends the AgentDF class to include a geometry attribute.
63    The agents will be stored in a GeoDataFrame.
64
65    Attributes
66    ----------
67    dtypes : dict[str, str]
68        The attributes of the Agent as a dictionary of columns and data types. It contains:
69        - geometry : geometry
70            -- The geometry of the Agent.
71    """
72
73    dtypes: dict[str, str] = {"geometry": "geometry"}
74
75    @classmethod
76    def step(cls):
77        """The step method of the GeoAgentDF class.
78        It should act on the mask portion of the agents dataframe of the model.
79        """
80        pass