Add create_agent factory method to Agent (#2351)

Author: quaquel

mesa/agent.py

@@ -85,6 +85,54 @@ def step(self) -> None:
     def advance(self) -> None:  # noqa: D102
         pass
 
+    @classmethod
+    def create_agents(cls, model: Model, n: int, *args, **kwargs) -> AgentSet[Agent]:
+        """Create N agents.
+
+        Args:
+            model: the model to which the agents belong
+            args: arguments to pass onto agent instances
+                  each arg is either a single object or a sequence of length n
+            n: the number of agents to create
+            kwargs: keyword arguments to pass onto agent instances
+                   each keyword arg is either a single object or a sequence of length n
+
+        Returns:
+            AgentSet containing the agents created.
+
+        """
+
+        class ListLike:
+            """Helper class to make default arguments act as if they are in a list of length N."""
+
+            def __init__(self, value):
+                self.value = value
+
+            def __getitem__(self, i):
+                return self.value
+
+        listlike_args = []
+        for arg in args:
+            if isinstance(arg, (list | np.ndarray | tuple)) and len(arg) == n:
+                listlike_args.append(arg)
+            else:
+                listlike_args.append(ListLike(arg))
+
+        listlike_kwargs = {}
+        for k, v in kwargs.items():
+            if isinstance(v, (list | np.ndarray | tuple)) and len(v) == n:
+                listlike_kwargs[k] = v
+            else:
+                listlike_kwargs[k] = ListLike(v)
+
+        agents = []
+        for i in range(n):
+            instance_args = [arg[i] for arg in listlike_args]
+            instance_kwargs = {k: v[i] for k, v in listlike_kwargs.items()}
+            agent = cls(model, *instance_args, **instance_kwargs)
+            agents.append(agent)
+        return AgentSet(agents, random=model.random)
+
     @property
     def random(self) -> Random:
         """Return a seeded stdlib rng."""

tests/test_agent.py

@@ -170,6 +170,30 @@ def test_agent_rng():
     assert agent.rng is model.rng
 
 
+def test_agent_create():
+    """Test create agent factory method."""
+
+    class TestAgent(Agent):
+        def __init__(self, model, attr, def_attr, a=0, b=0):
+            super().__init__(model)
+            self.some_attribute = attr
+            self.some_default_value = def_attr
+            self.a = a
+            self.b = b
+
+    model = Model(seed=42)
+    n = 10
+    some_attribute = model.rng.random(n)
+    a = tuple([model.random.random() for _ in range(n)])
+    TestAgent.create_agents(model, n, some_attribute, 5, a=a, b=7)
+
+    for agent, value, a_i in zip(model.agents, some_attribute, a):
+        assert agent.some_attribute == value
+        assert agent.some_default_value == 5
+        assert agent.a == a_i
+        assert agent.b == 7
+
+
 def test_agent_add_remove_discard():
     """Test adding, removing and discarding agents from AgentSet."""
     model = Model()