"""Network-based cell space using arbitrary connection patterns.
Creates spaces where cells connect based on network relationships rather than
spatial proximity. Built on NetworkX graphs, this enables:
- Arbitrary connectivity patterns between cells
- Graph-based neighborhood definitions
- Logical rather than physical distances
- Dynamic connectivity changes
- Integration with NetworkX's graph algorithms
Useful for modeling systems like social networks, transportation systems,
or any environment where connectivity matters more than physical location.
"""
from collections.abc import Callable, Mapping
from random import Random
from typing import Any
import numpy as np
from modelly.core.exceptions import SpaceException
from modelly.space.base import DiscreteSpace
from modelly.space.cells.cell import Cell
[docs]
class Network(DiscreteSpace[Cell]):
"""A networked discrete space."""
def __init__(
self,
G: Any, # noqa: N803
capacity: int | None = None,
random: Random | None = None,
cell_klass: type[Cell] = Cell,
layout: Mapping | Callable | None = None,
) -> None:
"""A Networked grid.
Args:
G: a NetworkX Graph instance.
capacity (int) : the capacity of the cell
random (Random): a random number generator
cell_klass (type[Cell]): The base Cell class to use in the Network
layout: A dictionary mapping node IDs to physical positions (x, y),
or a callable that generates them (e.g. nx.spring_layout).
It defaults to nx.circular_layout
This ensures all nodes possess physical (x, y) positions for visualization and
spatial queries without introducing performance bottlenecks on large graphs
"""
if layout is None:
import networkx as nx # noqa: PLC0415
layout = nx.circular_layout
super().__init__(capacity=capacity, random=random, cell_klass=cell_klass)
self.G = G
node_positions = _layout_positions(self.G, layout)
self._kdtree_cells = []
self._kdtree_dirty = False # True when KDTree needs a rebuild.
positions_for_tree = self._create_cells(node_positions, capacity)
from scipy.spatial import ( # noqa: PLC0415 - scipy is heavy; import lazily
KDTree,
)
if positions_for_tree:
self._kdtree = KDTree(np.array(positions_for_tree))
else:
self._kdtree = None
self._connect_cells()
def _create_cells(self, node_positions: Mapping, capacity: int | None) -> list:
"""Create one cell per node and return the node positions, in node order."""
positions_for_tree = []
# Create cells and gather KD-Tree data simultaneously
for node_id in self.G.nodes:
try:
pos = np.array(node_positions[node_id])
except KeyError as e:
raise SpaceException(
f"Node ID '{node_id}' is missing from the provided layout dictionary."
) from e
cell = self.cell_klass(
coordinate=node_id,
capacity=capacity,
random=self.random,
position=pos,
)
self._cells[node_id] = cell
self._kdtree_cells.append(cell)
positions_for_tree.append(pos)
return positions_for_tree
def _rebuild_kdtree(self) -> None:
"""Rebuild the KD-Tree."""
if self._kdtree_cells:
from scipy.spatial import ( # noqa: PLC0415 - scipy is heavy; import lazily
KDTree,
)
positions = np.array([c._position for c in self._kdtree_cells])
self._kdtree = KDTree(positions)
else:
self._kdtree = None
self._kdtree_dirty = False # set to false after rebuilding every time
def _connect_cells(self) -> None:
for cell in self.all_cells:
self._connect_single_cell(cell)
def _connect_single_cell(self, cell: Cell):
for node_id in self.G.neighbors(cell.coordinate):
cell.connect(self._cells[node_id], node_id)
[docs]
def find_nearest_cell(self, position: np.ndarray) -> Cell:
"""Find the network node nearest to the given position.
Only works for spatial networks (networks with node positions).
Note:
The KD-Tree index is rebuilt lazily. If cells were added or removed
since the previous spatial query, this method rebuilds the KD-Tree
before performing the nearest-neighbor lookup.
Args:
position: Physical position [x, y]
Returns:
Cell: The node closest to the position
Raises:
ValueError: If network is not spatial
"""
# build the kd-tree once if the _kdtree_dirty flag is true
if self._kdtree_dirty:
self._rebuild_kdtree()
if getattr(self, "_kdtree", None) is None:
raise ValueError("No nodes with positions found in network")
_, index = self._kdtree.query(position)
return self._kdtree_cells[index]
[docs]
def add_cell(self, cell: Cell):
"""Add a cell to the space."""
super().add_cell(cell)
self.G.add_node(cell.coordinate)
if cell._position is not None:
self._kdtree_cells.append(cell)
self._kdtree_dirty = True
[docs]
def remove_cell(self, cell: Cell):
"""Remove a cell from the space."""
super().remove_cell(cell)
self.G.remove_node(cell.coordinate)
if cell._position is not None and cell in self._kdtree_cells:
self._kdtree_cells.remove(cell)
self._kdtree_dirty = True
[docs]
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)
[docs]
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)
def _layout_positions(graph: Any, layout: Mapping | Callable) -> Mapping:
"""Resolve the ``layout`` argument of ``Network`` to node positions."""
if callable(layout):
return layout(graph)
if isinstance(layout, Mapping):
return layout
raise TypeError(
"Incorrect Layout Argument.\nShould be either `Mapping` or `Callable`"
)
