Repository: projectmesa/mesa
Pull request #1943 (2024-01-07 06:44:16)

refactor: Move Matplotlib-specific Solara components to separate file
Contains #1941, to avoid merge conflict. Depends on that PR to be merged first.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/projectmesa/mesa/pull/1943?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) Report
Attention: `83 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`0bd4429`)](https://app.codecov.io/gh/projectmesa/mesa/commit/0bd44293a76b103be4c3321cfaa9673ee655d825?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) 79.87% compared to head [(`91b85a4`)](https://app.codecov.io/gh/projectmesa/mesa/pull/1943?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) 79.34%.

| [Files](https://app.codecov.io/gh/projectmesa/mesa/pull/1943?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) | Patch % | Lines |
|---|---|---|
| [mesa/experimental/components/matplotlib.py](https://app.codecov.io/gh/projectmesa/mesa/pull/1943?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa#diff-bWVzYS9leHBlcmltZW50YWwvY29tcG9uZW50cy9tYXRwbG90bGliLnB5) | 13.18% | [79 Missing :warning: ](https://app.codecov.io/gh/projectmesa/mesa/pull/1943?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) |
| [mesa/experimental/jupyter\_viz.py](https://app.codecov.io/gh/projectmesa/mesa/pull/1943?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa#diff-bWVzYS9leHBlcmltZW50YWwvanVweXRlcl92aXoucHk=) | 33.33% | [4 Missing :warning: ](https://app.codecov.io/gh/projectmesa/mesa/pull/1943?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1943      +/-   ##
==========================================
- Coverage   79.87%   79.34%   -0.54%     
==========================================
  Files          15       16       +1     
  Lines        1267     1278      +11     
  Branches      277      255      -22     
==========================================
+ Hits         1012     1014       +2     
- Misses        216      225       +9     
  Partials       39       39              
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/projectmesa/mesa/pull/1943?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).

Corvince: Conflicts need to be resolved then its ready to merge

Patch:
diff --git a/mesa/agent.py b/mesa/agent.py
index 109a9dcd..8688dc70 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -8,6 +8,7 @@ Core Objects: Agent
 from __future__ import annotations
 
 import contextlib
+import copy
 import operator
 import warnings
 import weakref
@@ -124,9 +125,7 @@ class AgentSet(MutableSet, Sequence):
                 stacklevel=2,
             )
 
-        self._agents = weakref.WeakKeyDictionary()
-        for agent in agents:
-            self._agents[agent] = None
+        self._agents = weakref.WeakKeyDictionary({agent: None for agent in agents})
 
     def __len__(self) -> int:
         """Return the number of agents in the AgentSet."""
@@ -161,20 +160,21 @@ class AgentSet(MutableSet, Sequence):
             AgentSet: A new AgentSet containing the selected agents, unless inplace is True, in which case the current AgentSet is updated.
         """
 
-        def agent_generator():
+        if filter_func is None and agent_type is None and n == 0:
+            return self if inplace else copy.copy(self)
+
+        def agent_generator(filter_func=None, agent_type=None, n=0):
             count = 0
             for agent in self:
-                if filter_func and not filter_func(agent):
-                    continue
-                if agent_type and not isinstance(agent, agent_type):
-                    continue
-                yield agent
-                count += 1
-                # default of n is zero, zo evaluates to False
-                if n and count >= n:
-                    break
-
-        agents = agent_generator()
+                if (not filter_func or filter_func(agent)) and (
+                    not agent_type or isinstance(agent, agent_type)
+                ):
+                    yield agent
+                    count += 1
+                    if 0 < n <= count:
+                        break
+
+        agents = agent_generator(filter_func, agent_type, n)
 
         return AgentSet(agents, self.model) if not inplace else self._update(agents)
 
@@ -229,11 +229,8 @@ class AgentSet(MutableSet, Sequence):
         """Update the AgentSet with a new set of agents.
         This is a private method primarily used internally by other methods like select, shuffle, and sort.
         """
-        _agents = weakref.WeakKeyDictionary()
-        for agent in agents:
-            _agents[agent] = None
 
-        self._agents = _agents
+        self._agents = weakref.WeakKeyDictionary({agent: None for agent in agents})
         return self
 
     def do(
diff --git a/mesa/experimental/components/matplotlib.py b/mesa/experimental/components/matplotlib.py
new file mode 100644
index 00000000..e2389b5c
--- /dev/null
+++ b/mesa/experimental/components/matplotlib.py
@@ -0,0 +1,114 @@
+from typing import Optional
+
+import networkx as nx
+import solara
+from matplotlib.figure import Figure
+from matplotlib.ticker import MaxNLocator
+
+import mesa
+
+
+@solara.component
+def SpaceMatplotlib(model, agent_portrayal, dependencies: Optional[list[any]] = None):
+    space_fig = Figure()
+    space_ax = space_fig.subplots()
+    space = getattr(model, "grid", None)
+    if space is None:
+        # Sometimes the space is defined as model.space instead of model.grid
+        space = model.space
+    if isinstance(space, mesa.space.NetworkGrid):
+        _draw_network_grid(space, space_ax, agent_portrayal)
+    elif isinstance(space, mesa.space.ContinuousSpace):
+        _draw_continuous_space(space, space_ax, agent_portrayal)
+    else:
+        _draw_grid(space, space_ax, agent_portrayal)
+    space_ax.set_axis_off()
+    solara.FigureMatplotlib(space_fig, format="png", dependencies=dependencies)
+
+
+def _draw_grid(space, space_ax, agent_portrayal):
+    def portray(g):
+        x = []
+        y = []
+        s = []  # size
+        c = []  # color
+        for i in range(g.width):
+            for j in range(g.height):
+                content = g._grid[i][j]
+                if not content:
+                    continue
+                if not hasattr(content, "__iter__"):
+                    # Is a single grid
+                    content = [content]
+                for agent in content:
+                    data = agent_portrayal(agent)
+                    x.append(i)
+                    y.append(j)
+                    if "size" in data:
+                        s.append(data["size"])
+                    if "color" in data:
+                        c.append(data["color"])
+        out = {"x": x, "y": y}
+        if len(s) > 0:
+            out["s"] = s
+        if len(c) > 0:
+            out["c"] = c
+        return out
+
+    space_ax.scatter(**portray(space))
+
+
+def _draw_network_grid(space, space_ax, agent_portrayal):
+    graph = space.G
+    pos = nx.spring_layout(graph, seed=0)
+    nx.draw(
+        graph,
+        ax=space_ax,
+        pos=pos,
+        **agent_portrayal(graph),
+    )
+
+
+def _draw_continuous_space(space, space_ax, agent_portrayal):
+    def portray(space):
+        x = []
+        y = []
+        s = []  # size
+        c = []  # color
+        for agent in space._agent_to_index:
+            data = agent_portrayal(agent)
+            _x, _y = agent.pos
+            x.append(_x)
+            y.append(_y)
+            if "size" in data:
+                s.append(data["size"])
+            if "color" in data:
+                c.append(data["color"])
+        out = {"x": x, "y": y}
+        if len(s) > 0:
+            out["s"] = s
+        if len(c) > 0:
+            out["c"] = c
+        return out
+
+    space_ax.scatter(**portray(space))
+
+
+def make_plot(model, measure):
+    fig = Figure()
+    ax = fig.subplots()
+    df = model.datacollector.get_model_vars_dataframe()
+    if isinstance(measure, str):
+        ax.plot(df.loc[:, measure])
+        ax.set_ylabel(measure)
+    elif isinstance(measure, dict):
+        for m, color in measure.items():
+            ax.plot(df.loc[:, m], label=m, color=color)
+        fig.legend()
+    elif isinstance(measure, (list, tuple)):
+        for m in measure:
+            ax.plot(df.loc[:, m], label=m)
+        fig.legend()
+    # Set integer x axis
+    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
+    solara.FigureMatplotlib(fig)
diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py
index f4ac1271..579391c5 100644
--- a/mesa/experimental/jupyter_viz.py
+++ b/mesa/experimental/jupyter_viz.py
@@ -3,14 +3,11 @@ import threading
 from typing import Optional
 
 import matplotlib.pyplot as plt
-import networkx as nx
 import reacton.ipywidgets as widgets
 import solara
-from matplotlib.figure import Figure
-from matplotlib.ticker import MaxNLocator
 from solara.alias import rv
 
-import mesa
+import mesa.experimental.components.matplotlib as components_matplotlib
 
 # Avoid interactive backend
 plt.switch_backend("agg")
@@ -72,7 +69,7 @@ def JupyterViz(
                 rv.CardTitle(children=["Space"])
                 if space_drawer == "default":
                     # draw with the default implementation
-                    SpaceMatplotlib(
+                    components_matplotlib.SpaceMatplotlib(
                         model, agent_portrayal, dependencies=[current_step.value]
                     )
                 elif space_drawer:
@@ -85,7 +82,7 @@ def JupyterViz(
                     # Is a custom object
                     measure(model)
                 else:
-                    make_plot(model, measure)
+                    components_matplotlib.make_plot(model, measure)
         return main
 
     # 3. Set up UI
@@ -106,7 +103,7 @@ def JupyterViz(
             # 4. Space
             if space_drawer == "default":
                 # draw with the default implementation
-                SpaceMatplotlib(
+                components_matplotlib.SpaceMatplotlib(
                     model, agent_portrayal, dependencies=[current_step.value]
                 )
             elif space_drawer:
@@ -121,7 +118,7 @@ def JupyterViz(
                     # Is a custom object
                     measure(model)
                 else:
-                    make_plot(model, measure)
+                    components_matplotlib.make_plot(model, measure)
 
     def render_in_browser():
         # if space drawer is disabled, do not include it
@@ -182,7 +179,7 @@ def ModelController(model, play_interval, current_step, reset_counter):
     def do_step():
         model.step()
         previous_step.value = current_step.value
-        current_step.value = model.schedule.steps
+        current_step.value += 1
 
     def do_play():
         model.running = True
@@ -316,112 +313,6 @@ def UserInputs(user_params, on_change=None):
             raise ValueError(f"{input_type} is not a supported input type")
 
 
-@solara.component
-def SpaceMatplotlib(model, agent_portrayal, dependencies: Optional[list[any]] = None):
-    space_fig = Figure()
-    space_ax = space_fig.subplots()
-    space = getattr(model, "grid", None)
-    if space is None:
-        # Sometimes the space is defined as model.space instead of model.grid
-        space = model.space
-    if isinstance(space, mesa.space.NetworkGrid):
-        _draw_network_grid(space, space_ax, agent_portrayal)
-    elif isinstance(space, mesa.space.ContinuousSpace):
-        _draw_continuous_space(space, space_ax, agent_portrayal)
-    else:
-        _draw_grid(space, space_ax, agent_portrayal)
-    space_ax.set_axis_off()
-    solara.FigureMatplotlib(space_fig, format="png", dependencies=dependencies)
-
-
-def _draw_grid(space, space_ax, agent_portrayal):
-    def portray(g):
-        x = []
-        y = []
-        s = []  # size
-        c = []  # color
-        for i in range(g.width):
-            for j in range(g.height):
-                content = g._grid[i][j]
-                if not content:
-                    continue
-                if not hasattr(content, "__iter__"):
-                    # Is a single grid
-                    content = [content]
-                for agent in content:
-                    data = agent_portrayal(agent)
-                    x.append(i)
-                    y.append(j)
-                    if "size" in data:
-                        s.append(data["size"])
-                    if "color" in data:
-                        c.append(data["color"])
-        out = {"x": x, "y": y}
-        if len(s) > 0:
-            out["s"] = s
-        if len(c) > 0:
-            out["c"] = c
-        return out
-
-    space_ax.scatter(**portray(space))
-
-
-def _draw_network_grid(space, space_ax, agent_portrayal):
-    graph = space.G
-    pos = nx.spring_layout(graph, seed=0)
-    nx.draw(
-        graph,
-        ax=space_ax,
-        pos=pos,
-        **agent_portrayal(graph),
-    )
-
-
-def _draw_continuous_space(space, space_ax, agent_portrayal):
-    def portray(space):
-        x = []
-        y = []
-        s = []  # size
-        c = []  # color
-        for agent in space._agent_to_index:
-            data = agent_portrayal(agent)
-            _x, _y = agent.pos
-            x.append(_x)
-            y.append(_y)
-            if "size" in data:
-                s.append(data["size"])
-            if "color" in data:
-                c.append(data["color"])
-        out = {"x": x, "y": y}
-        if len(s) > 0:
-            out["s"] = s
-        if len(c) > 0:
-            out["c"] = c
-        return out
-
-    space_ax.scatter(**portray(space))
-
-
-def make_plot(model, measure):
-    fig = Figure()
-    ax = fig.subplots()
-    df = model.datacollector.get_model_vars_dataframe()
-    if isinstance(measure, str):
-        ax.plot(df.loc[:, measure])
-        ax.set_ylabel(measure)
-    elif isinstance(measure, dict):
-        for m, color in measure.items():
-            ax.plot(df.loc[:, m], label=m, color=color)
-        fig.legend()
-    elif isinstance(measure, (list, tuple)):
-        for m in measure:
-            ax.plot(df.loc[:, m], label=m)
-        fig.legend()
-    # Set integer x axis
-    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
-    solara.FigureMatplotlib(fig)
-
-
 def make_text(renderer):
     def function(model):
         solara.Markdown(renderer(model))
diff --git a/mesa/space.py b/mesa/space.py
index 969f1838..b730c2da 100644
--- a/mesa/space.py
+++ b/mesa/space.py
@@ -326,8 +326,10 @@ class _Grid:
             at most 9 if Moore, 5 if Von-Neumann
             (8 and 4 if not including the center).
         """
-        neighborhood = self.get_neighborhood(pos, moore, include_center, radius)
-        return self.iter_cell_list_contents(neighborhood)
+        default_val = self.default_val()
+        for x, y in self.get_neighborhood(pos, moore, include_center, radius):
+            if (cell := self._grid[x][y]) != default_val:
+                yield cell
 
     def get_neighbors(
         self,
@@ -385,11 +387,10 @@ class _Grid:
             An iterator of the agents contained in the cells identified in `cell_list`.
         """
         # iter_cell_list_contents returns only non-empty contents.
-        return (
-            cell
-            for x, y in cell_list
-            if (cell := self._grid[x][y]) != self.default_val()
-        )
+        default_val = self.default_val()
+        for x, y in cell_list:
+            if (cell := self._grid[x][y]) != default_val:
+                yield cell
 
     @accept_tuple_argument
     def get_cell_list_contents(self, cell_list: Iterable[Coordinate]) -> list[Agent]:
@@ -1045,6 +1046,17 @@ class MultiGrid(_PropertyGrid):
             self._empty_mask[agent.pos] = False
         agent.pos = None
 
+    def iter_neighbors(
+        self,
+        pos: Coordinate,
+        moore: bool,
+        include_center: bool = False,
+        radius: int = 1,
+    ) -> Iterator[Agent]:
+        return itertools.chain.from_iterable(
+            super().iter_neighbors(pos, moore, include_center, radius)
+        )
+
     @accept_tuple_argument
     def iter_cell_list_contents(
         self, cell_list: Iterable[Coordinate]
@@ -1058,10 +1070,9 @@ class MultiGrid(_PropertyGrid):
         Returns:
             An iterator of the agents contained in the cells identified in `cell_list`.
         """
+        default_val = self.default_val()
         return itertools.chain.from_iterable(
-            cell
-            for x, y in cell_list
-            if (cell := self._grid[x][y]) != self.default_val()
+            cell for x, y in cell_list if (cell := self._grid[x][y]) != default_val
         )
 
 


Test patch:
diff --git a/tests/test_jupyter_viz.py b/tests/test_jupyter_viz.py
index dd125ffc..6f88c7d1 100644
--- a/tests/test_jupyter_viz.py
+++ b/tests/test_jupyter_viz.py
@@ -82,7 +82,7 @@ class TestMakeUserInput(unittest.TestCase):
 
 
 class TestJupyterViz(unittest.TestCase):
-    @patch("mesa.experimental.jupyter_viz.SpaceMatplotlib")
+    @patch("mesa.experimental.components.matplotlib.SpaceMatplotlib")
     def test_call_space_drawer(self, mock_space_matplotlib):
         mock_model_class = Mock()
         agent_portrayal = {