Repository: projectmesa/mesa
Pull request #2755: Add support for dynamic discrete spaces
Author: login: quaquel
Labels: name: enhancement
description: Release notes label
Created: 2025-04-14T17:55:26Z | Merged: 2025-04-20T08:40:49Z

Currently, all `DiscreteSpace` subclasses rely on `@cache` and `@cached_property` for neighborhood-related functionality. This is critical for performance. However, if you want to represent a dynamically changing discrete space (e.g., a dynamic network), this is currently not possible. See ##2754 for more details

# API
This PR adds explicit support for changing discrete spaces during the simulation. It adds 4 new methods to `DiscreteSpace` which means these are available to all subclasses. These methods are `add_cell`, `remove_cell`, `add_connection`, `remove_connection`. With these, you can make any modification to a discrete space that you want. For example, it allows you to dynamically change a network or impose a barrier in an orthogonal grid by removing connections between specific cells. 

# Implementation details
This adds support for adding and removing cells and connections to discrete spaces. The critical issue is that discrete spaces relies on caching neighborhoods for substantial performance gains. `Cell` already has a `connect` and `disconnect` method, so clearing the relevant cell caches is placed here. The new methods (`add_cell`, `remove_cell`, `add_connection`, `remove_connection`) are added to `DiscreteSpace`, and the relevant cache (`all_cells`) is cleared here as well. 

The clearing of caches on `connect` and `disconnect` adds some additional overhead when constructing a discrete space. It might be possible to reduce this overhead by checking if the cache is already active (long story short, this can be done by checking if `self.__dict__["name of cached property"]` exists, but this adds code complexity). 
  


Diff:
diff --git a/mesa/discrete_space/cell.py b/mesa/discrete_space/cell.py
index 2b1a42d004d..f8621ead808 100644
--- a/mesa/discrete_space/cell.py
+++ b/mesa/discrete_space/cell.py
@@ -84,6 +84,7 @@ def connect(self, other: Cell, key: Coordinate | None = None) -> None:
         """
         if key is None:
             key = other.coordinate
+        self._clear_cache()
         self.connections[key] = other
 
     def disconnect(self, other: Cell) -> None:
@@ -96,6 +97,7 @@ def disconnect(self, other: Cell) -> None:
         keys_to_remove = [k for k, v in self.connections.items() if v == other]
         for key in keys_to_remove:
             del self.connections[key]
+        self._clear_cache()
 
     def add_agent(self, agent: CellAgent) -> None:
         """Adds an agent to the cell.
@@ -210,3 +212,15 @@ def __getstate__(self):
             "connections"
         ] = {}  # replace this with empty connections to avoid infinite recursion error in pickle/deepcopy
         return state
+
+    def _clear_cache(self):
+        """Helper function to clear local cache."""
+        try:
+            self.__dict__.pop(
+                "neighborhood"
+            )  # cached properties are stored in __dict__, see functools.cached_property docs
+        except KeyError:
+            pass  # cache is not set
+        else:
+            self.get_neighborhood.cache_clear()
+            self._neighborhood.cache_clear()
diff --git a/mesa/discrete_space/discrete_space.py b/mesa/discrete_space/discrete_space.py
index 4a71673cd31..f480ff15fbe 100644
--- a/mesa/discrete_space/discrete_space.py
+++ b/mesa/discrete_space/discrete_space.py
@@ -85,6 +85,68 @@ def agents(self) -> AgentSet:
     def _connect_cells(self): ...
     def _connect_single_cell(self, cell: T): ...
 
+    def add_cell(self, cell: T):
+        """Add a cell to the space.
+
+        Args:
+            cell: cell to add
+
+        Note:
+            Discrete spaces rely on caching neighborhood relations for speedups. Adding or removing cells and
+            connections at runtime is possible. However, only the caches of cells directly affected will be cleared. So
+            if you rely on getting neighborhoods of cells with a radius higher than 1, these might not be cleared
+            correctly if you are adding or removing cells and connections at runtime.
+
+        """
+        self.__dict__.pop("all_cells", None)
+        self._cells[cell.coordinate] = cell
+
+    def remove_cell(self, cell: T):
+        """Remove a cell from the space.
+
+        Note:
+            Discrete spaces rely on caching neighborhood relations for speedups. Adding or removing cells and
+            connections at runtime is possible. However, only the caches of cells directly affected will be cleared. So
+            if you rely on getting neighborhoods of cells with a radius higher than 1, these might not be cleared
+            correctly if you are adding or removing cells and connections at runtime.
+
+
+        """
+        neighbors = cell.neighborhood
+        self._cells.pop(cell.coordinate)
+        self.__dict__.pop("all_cells", None)
+
+        # iterate over all neighbors
+        for neighbor in neighbors.cells:
+            neighbor.disconnect(cell)
+            cell.disconnect(neighbor)
+
+    def add_connection(self, cell1: T, cell2: T):
+        """Add a connection between the two cells.
+
+        Note:
+            Discrete spaces rely on caching neighborhood relations for speedups. Adding or removing cells and
+            connections at runtime is possible. However, only the caches of cells directly affected will be cleared. So
+            if you rely on getting neighborhoods of cells with a radius higher than 1, these might not be cleared
+            correctly if you are adding or removing cells and connections at runtime.
+
+        """
+        cell1.connect(cell2)
+        cell2.connect(cell1)
+
+    def remove_connection(self, cell1: T, cell2: T):
+        """Remove a connection between the two cells.
+
+        Note:
+            Discrete spaces rely on caching neighborhood relations for speedups. Adding or removing cells and
+            connections at runtime is possible. However, only the caches of cells directly affected will be cleared. So
+            if you rely on getting neighborhoods of cells with a radius higher than 1, these might not be cleared
+            correctly if you are adding or removing cells and connections at runtime.
+
+        """
+        cell1.disconnect(cell2)
+        cell2.disconnect(cell1)
+
     @cached_property
     def all_cells(self):
         """Return all cells in space."""
diff --git a/mesa/discrete_space/network.py b/mesa/discrete_space/network.py
index 7f692cd2eaf..3fc0bc2ff14 100644
--- a/mesa/discrete_space/network.py
+++ b/mesa/discrete_space/network.py
@@ -55,3 +55,23 @@ def _connect_cells(self) -> None:
     def _connect_single_cell(self, cell: Cell):
         for node_id in self.G.neighbors(cell.coordinate):
             cell.connect(self._cells[node_id], node_id)
+
+    def add_cell(self, cell: Cell):
+        """Add a cell to the space."""
+        super().add_cell(cell)
+        self.G.add_node(cell.coordinate)
+
+    def remove_cell(self, cell: Cell):
+        """Remove a cell from the space."""
+        super().remove_cell(cell)
+        self.G.remove_node(cell.coordinate)
+
+    def add_connection(self, cell1: Cell, cell2: Cell):
+        """Add a connection between the two cells."""
+        super().add_connection(cell1, cell2)
+        self.G.add_edge(cell1.coordinate, cell2.coordinate)
+
+    def remove_connection(self, cell1: Cell, cell2: Cell):
+        """Remove a connection between the two cells."""
+        super().remove_connection(cell1, cell2)
+        self.G.remove_edge(cell1.coordinate, cell2.coordinate)
diff --git a/tests/test_discrete_space.py b/tests/test_discrete_space.py
index 42e61d14780..28919dcc333 100644
--- a/tests/test_discrete_space.py
+++ b/tests/test_discrete_space.py
@@ -315,6 +315,42 @@ def test_orthogonal_grid_moore_1d():
         assert connection.coordinate in {(1,), (9,)}
 
 
+def test_dynamic_modifications_to_space():
+    """Test dynamic modifications to DiscreteSpace."""
+    grid = OrthogonalMooreGrid(
+        (5, 5), torus=False, capacity=1, random=random.Random(42)
+    )
+
+    cells = grid._cells
+
+    # test remove_connection
+    cell1 = cells[(2, 0)]
+    cell2 = cells[(3, 0)]
+    grid.remove_connection(cell1, cell2)
+
+    assert cell2 not in cell1.neighborhood
+    assert cell1 not in cell2.neighborhood
+
+    # test add connection
+    grid.add_connection(cell1, cell2)
+    assert cell1 in cell2.neighborhood
+    assert cell2 in cell1.neighborhood
+
+    # test remove cell
+    neighbors = cell1.neighborhood
+    grid.remove_cell(cell1)
+    for neighbor in neighbors:
+        assert cell1 not in neighbor.neighborhood
+
+    # test add_cells
+    grid.add_cell(cell1)
+    for neighbor in neighbors:
+        grid.add_connection(cell1, neighbor)
+
+    for neighbor in neighbors:
+        assert cell1 in neighbor.neighborhood
+
+
 def test_cell_neighborhood():
     """Test neighborhood method of cell in different GridSpaces."""
     # orthogonal grid
@@ -451,6 +487,23 @@ def test_networkgrid():
 
     pickle.loads(pickle.dumps(grid))  # noqa: S301
 
+    cell = Cell(10)  # n = 10, so 10 + 1
+    grid.add_cell(cell)
+    grid.add_connection(cell, grid._cells[0])
+    assert cell in grid._cells[0].neighborhood
+    assert grid._cells[0] in cell.neighborhood
+
+    grid.remove_connection(cell, grid._cells[0])
+    assert cell not in grid._cells[0].neighborhood
+    assert grid._cells[0] not in cell.neighborhood
+
+    cell = Cell(10)  # n = 10, so 10 + 1
+    grid.add_cell(cell)
+    grid.add_connection(cell, grid._cells[0])
+    grid.remove_cell(cell)  # this also removes all connections
+    assert cell not in grid._cells[0].neighborhood
+    assert grid._cells[0] not in cell.neighborhood
+
 
 def test_voronoigrid():
     """Test VoronoiGrid."""


Reviews:
user: login: Sahil-Chhoker
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T05:38:23Z
user: login: quaquel
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-15T06:01:48Z
user: login: EwoutH
body: This is really cool, and a remarkably clean PR!
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-16T20:01:12Z
user: login: tpike3
body: This is great, I think we can worry about decreasing the `init` later. I know we have at least one user who needs this functionality. 
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-17T10:01:44Z

Review comments:
diff_hunk: @@ -85,6 +85,36 @@ def agents(self) -> AgentSet:
     def _connect_cells(self): ...
     def _connect_single_cell(self, cell: T): ...
 
+    def add_cell(self, cell: T):
+        """Add a cell to the space.
+
+        Args:
+            cell: cell to add
+
+        """
+        self.all_cells.clear_cache()
+        self._cells[cell.coordinate] = cell
+
+    def remove_cell(self, cell: T):
+        """Remove a cell from the space."""
+        neighbors = cell.neighborhood
+        self._cells.pop(cell.coordinate)
+        self.all_cells.clear_cache()
+
+        # iterate over all neighbors
+        for neighbor in neighbors.cells:
+            neighbor.disconnect(cell)
+
+    def add_connection(self, cell1: T, cell2: T):
+        """Add a connection between the two cells."""
+        cell1.connect(cell2)
+        cell2.connect(cell1)
+
+    def remove_connection(self, cell1: T, cell2: T):
+        """Remove a connection between the two cells."""
+        cell1.connect(cell2)
+        cell2.connect(cell1)
path: mesa/discrete_space/discrete_space.py
user: login: Sahil-Chhoker
body: Why does `remove_connection` contain `cell.connect`, am I missing something?
created_at: 2025-04-15T05:38:23Z
author_association: COLLABORATOR
diff_hunk: @@ -85,6 +85,36 @@ def agents(self) -> AgentSet:
     def _connect_cells(self): ...
     def _connect_single_cell(self, cell: T): ...
 
+    def add_cell(self, cell: T):
+        """Add a cell to the space.
+
+        Args:
+            cell: cell to add
+
+        """
+        self.all_cells.clear_cache()
+        self._cells[cell.coordinate] = cell
+
+    def remove_cell(self, cell: T):
+        """Remove a cell from the space."""
+        neighbors = cell.neighborhood
+        self._cells.pop(cell.coordinate)
+        self.all_cells.clear_cache()
+
+        # iterate over all neighbors
+        for neighbor in neighbors.cells:
+            neighbor.disconnect(cell)
+
+    def add_connection(self, cell1: T, cell2: T):
+        """Add a connection between the two cells."""
+        cell1.connect(cell2)
+        cell2.connect(cell1)
+
+    def remove_connection(self, cell1: T, cell2: T):
+        """Remove a connection between the two cells."""
+        cell1.connect(cell2)
+        cell2.connect(cell1)
path: mesa/discrete_space/discrete_space.py
user: login: quaquel
body: good point, that should be disconnect. Fixing it.
created_at: 2025-04-15T06:01:48Z
author_association: MEMBER

PR comments:
user: login: coderabbitai[bot]
created_at: 2025-04-14T17:55:34Z
body: <!-- This is an auto-generated comment: summarize by coderabbit.ai -->
<!-- This is an auto-generated comment: skip review by coderabbit.ai -->

> [!IMPORTANT]
> ## Review skipped
> 
> Auto reviews are disabled on this repository.
> 
> Please check the settings in the CodeRabbit UI or the `.coderabbit.yaml` file in this repository. To trigger a single review, invoke the `@coderabbitai review` command.
> 
> You can disable this status message by setting the `reviews.review_status` to `false` in the CodeRabbit configuration file.

<!-- end of auto-generated comment: skip review by coderabbit.ai -->
<!-- tips_start -->

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

<details>
<summary>❤️ Share</summary>

- [X](https://twitter.com/intent/tweet?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A&url=https%3A//coderabbit.ai)
- [Mastodon](https://mastodon.social/share?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A%20https%3A%2F%2Fcoderabbit.ai)
- [Reddit](https://www.reddit.com/submit?title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&text=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code.%20Check%20it%20out%3A%20https%3A//coderabbit.ai)
- [LinkedIn](https://www.linkedin.com/sharing/share-offsite/?url=https%3A%2F%2Fcoderabbit.ai&mini=true&title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&summary=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code)

</details>

<details>
<summary>🪧 Tips</summary>

### Chat

There are 3 ways to chat with [CodeRabbit](https://coderabbit.ai?utm_source=oss&utm_medium=github&utm_campaign=projectmesa/mesa&utm_content=2755):

- Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  - `I pushed a fix in commit <commit_id>, please review it.`
  - `Generate unit testing code for this file.`
  - `Open a follow-up GitHub issue for this discussion.`
- Files and specific lines of code (under the "Files changed" tab): Tag `@coderabbitai` in a new review comment at the desired location with your query. Examples:
  - `@coderabbitai generate unit testing code for this file.`
  -	`@coderabbitai modularize this function.`
- PR comments: Tag `@coderabbitai` in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  - `@coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.`
  - `@coderabbitai read src/utils.ts and generate unit testing code.`
  - `@coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.`
  - `@coderabbitai help me debug CodeRabbit configuration file.`

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

### CodeRabbit Commands (Invoked using PR comments)

- `@coderabbitai pause` to pause the reviews on a PR.
- `@coderabbitai resume` to resume the paused reviews.
- `@coderabbitai review` to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
- `@coderabbitai full review` to do a full review from scratch and review all the files again.
- `@coderabbitai summary` to regenerate the summary of the PR.
- `@coderabbitai generate docstrings` to [generate docstrings](https://docs.coderabbit.ai/finishing-touches/docstrings) for this PR.
- `@coderabbitai resolve` resolve all the CodeRabbit review comments.
- `@coderabbitai configuration` to show the current CodeRabbit configuration for the repository.
- `@coderabbitai help` to get help.

### Other keywords and placeholders

- Add `@coderabbitai ignore` anywhere in the PR description to prevent this PR from being reviewed.
- Add `@coderabbitai summary` to generate the high-level summary at a specific location in the PR description.
- Add `@coderabbitai` anywhere in the PR title to generate the title automatically.

### CodeRabbit Configuration File (`.coderabbit.yaml`)

- You can programmatically configure CodeRabbit by adding a `.coderabbit.yaml` file to the root of your repository.
- Please see the [configuration documentation](https://docs.coderabbit.ai/guides/configure-coderabbit) for more information.
- If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: `# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json`

### Documentation and Community

- Visit our [Documentation](https://docs.coderabbit.ai) for detailed information on how to use CodeRabbit.
- Join our [Discord Community](http://discord.gg/coderabbit) to get help, request features, and share feedback.
- Follow us on [X/Twitter](https://twitter.com/coderabbitai) for updates and announcements.

</details>

<!-- tips_end -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-14T18:03:38Z
body: Performance benchmarks:

| Model           | Size   | Init time [95% CI]         | Run time [95% CI]       |
|-----------------|--------|----------------------------|-------------------------|
| BoltzmannWealth | small  | 🔴 +29.6% [+28.5%, +30.7%] | 🔵 +0.5% [+0.3%, +0.6%] |
| BoltzmannWealth | large  | 🔴 +36.1% [+17.7%, +45.5%] | 🔵 +0.8% [+0.0%, +1.4%] |
| Schelling       | small  | 🔴 +55.1% [+54.8%, +55.4%] | 🔵 +0.3% [+0.1%, +0.5%] |
| Schelling       | large  | 🔴 +58.0% [+55.5%, +61.2%] | 🔵 -1.5% [-2.0%, -1.0%] |
| WolfSheep       | small  | 🔴 +24.7% [+24.2%, +25.3%] | 🔵 -0.3% [-0.5%, -0.1%] |
| WolfSheep       | large  | 🔴 +27.8% [+27.3%, +28.2%] | 🔴 +4.5% [+3.6%, +5.4%] |
| BoidFlockers    | small  | 🔵 +1.6% [+1.1%, +2.1%]    | 🔵 -0.8% [-1.0%, -0.6%] |
| BoidFlockers    | large  | 🔵 +1.3% [+0.9%, +1.7%]    | 🔵 -0.5% [-0.7%, -0.2%] |

author_association: NONE
user: login: quaquel
created_at: 2025-04-14T18:05:52Z
body: The increase in init_time is due to the clearing of non-existing caches on all `cell.connect` calls when instantiating the spaces. It's possible to resolve this but the numbers or tiny and it adds code complexity if we want to do that. 
author_association: MEMBER
user: login: Sahil-Chhoker
created_at: 2025-04-16T18:09:57Z
body: > The increase in init_time is due to the clearing of non-existing caches on all `cell.connect` calls when instantiating the spaces. It's possible to resolve this but the numbers or tiny and it adds code complexity if we want to do that.

I am not asking for the complete implementation, but can you give me just a initial idea of how can the `init` time be reduced?
author_association: COLLABORATOR
user: login: quaquel
created_at: 2025-04-16T18:13:34Z
body: > I am not asking for the complete implementation, but can you give me just a initial idea of how can the init time be reduced?

In `Cell`, there are 3 places were caching happens: `neighborhood`, `_neighborhood`, and `get_neighborhood`. These are tied together with `neighborhood` as the basis. `neighborhood ` is a cached property. The cached value is stored in `self.__dict__["neighborhood"]`. So, if this field thas not exist, the other caches also won't exist and thus won't have to be cleared. No idea how much runtime it would save. 
author_association: MEMBER
user: login: Sahil-Chhoker
created_at: 2025-04-16T18:16:26Z
body: But won't that eliminate the caching entirely from cell?
author_association: COLLABORATOR
user: login: quaquel
created_at: 2025-04-16T18:32:00Z
body: No, but why do you think that it might?
author_association: MEMBER
user: login: Sahil-Chhoker
created_at: 2025-04-16T20:27:32Z
body: > No, but why do you think that it might?

I am not sure about the current structure, I will do some individual research later. Thanks again for these discussions, learning a lot from these.
author_association: COLLABORATOR
user: login: quaquel
created_at: 2025-04-17T14:08:35Z
body: I'll expand the docstrings a bit more probably tomorrow and merge this afterwards. 
author_association: MEMBER
user: login: github-actions[bot]
created_at: 2025-04-20T07:44:50Z
body: Performance benchmarks:

| Model           | Size   | Init time [95% CI]         | Run time [95% CI]       |
|-----------------|--------|----------------------------|-------------------------|
| BoltzmannWealth | small  | 🔴 +30.6% [+28.8%, +32.6%] | 🔵 +0.0% [-0.2%, +0.2%] |
| BoltzmannWealth | large  | 🔴 +57.3% [+38.1%, +93.7%] | 🔵 -4.1% [-6.1%, -2.0%] |
| Schelling       | small  | 🔴 +48.4% [+48.0%, +48.8%] | 🔵 +0.3% [-0.0%, +0.6%] |
| Schelling       | large  | 🔴 +42.0% [+34.4%, +46.6%] | 🔵 -2.6% [-5.7%, +0.1%] |
| WolfSheep       | small  | 🔴 +18.7% [+18.3%, +19.1%] | 🔵 -0.5% [-0.8%, -0.3%] |
| WolfSheep       | large  | 🔴 +19.8% [+19.1%, +20.5%] | 🟢 -5.3% [-7.1%, -3.6%] |
| BoidFlockers    | small  | 🔵 -1.9% [-2.6%, -1.2%]    | 🔵 +0.2% [+0.1%, +0.3%] |
| BoidFlockers    | large  | 🔵 -0.7% [-1.2%, -0.3%]    | 🔵 +0.3% [+0.0%, +0.5%] |

author_association: NONE