From a1089ad8f825acc300a4b293aed29f0728d2a51c Mon Sep 17 00:00:00 2001 From: Sahil Chhoker Date: Tue, 4 Feb 2025 13:20:57 +0530 Subject: [PATCH 1/3] Feat: Added portrayal components for agent and property layer --- .../components/portrayal_components.py | 159 ++++++++++++++++++ mesa/visualization/mpl_space_drawing.py | 19 ++- 2 files changed, 171 insertions(+), 7 deletions(-) create mode 100644 mesa/visualization/components/portrayal_components.py diff --git a/mesa/visualization/components/portrayal_components.py b/mesa/visualization/components/portrayal_components.py new file mode 100644 index 00000000000..6ae3b29c670 --- /dev/null +++ b/mesa/visualization/components/portrayal_components.py @@ -0,0 +1,159 @@ +from dataclasses import dataclass, field, asdict +from typing import Optional, Union, ClassVar, Set, Dict +from matplotlib import pyplot as plt +import matplotlib.colors as mcolors +from numbers import Number +from matplotlib.colors import Colormap + +@dataclass +class AgentPortrayalStyle: + """ + A class to define the visual style of an agent in Mesa visualizations. + + Attributes: + color: Color specification for the agent. Can be any valid matplotlib color. + marker: Marker style for the agent (e.g., 'o', 's', '^'). + size: Size of the agent marker in points. + zorder: The drawing order of the agent relative to other elements. + + Example: + >>> def agent_portrayal(agent): + >>> portrayal = AgentPortrayalStyle( + >>> color='tab:orange' if agent.id == 1 else 'tab:blue', + >>> marker='^', + >>> size=70, + >>> zorder=1 + >>> ) + >>> return dict(portrayal) + """ + + VALID_MARKERS: ClassVar[Set[str]] = { + 'o', 's', '^', 'v', '<', '>', 'D', 'p', 'h', 'H', '8', '*', + '+', 'x', '.', ',', '1', '2', '3', '4', '|', '_', + } + + # Default values + color: Optional[Union[str, tuple]] = "tab:blue" + marker: Optional[str] = "o" + size: Optional[Number] = 50 + zorder: Optional[int] = 1 + + def __post_init__(self): + """Validate the attributes after initialization.""" + if self.color is not None: + try: + mcolors.to_rgb(self.color) + except ValueError: + raise ValueError(f"Invalid color specification: {self.color}") + + if self.marker is not None: + if not isinstance(self.marker, str): + raise ValueError(f"Marker must be a string, got {type(self.marker)}") + if self.marker not in self.VALID_MARKERS: + raise ValueError( + f"Invalid marker '{self.marker}'. Valid markers are: {', '.join(sorted(self.VALID_MARKERS))}" + ) + + if self.size is not None and not isinstance(self.size, Number): + raise ValueError(f"Size must be a number, got {type(self.size)}") + + if self.zorder is not None and not isinstance(self.zorder, int): + raise ValueError(f"Zorder must be an integer, got {type(self.zorder)}") + + def __iter__(self): + return iter(asdict(self).items()) + + +@dataclass +class PropertyLayerStyle: + """Style configuration for a single property layer. + + Args: + color: A valid matplotlib color string + colormap: A valid matplotlib colormap name or Colormap object + vmin: Minimum value for colormap scaling + vmax: Maximum value for colormap scaling + alpha: Opacity value between 0 and 1 + colorbar: Whether to show colorbar + + Note: color and colormap are mutually exclusive. + """ + vmin: Optional[float] = None + vmax: Optional[float] = None + alpha: Optional[float] = None + colorbar: Optional[bool] = None + color: Optional[str] = None + colormap: Optional[Union[str, Colormap]] = None + + def __post_init__(self): + """Validate style attributes.""" + # Validate color/colormap exclusivity + if self.color is not None and self.colormap is not None: + raise ValueError("Cannot specify both color and colormap") + if self.color is not None and self.colormap is not None: + raise ValueError("Cannot specify both color and colormap") + + # Validate color if specified + if self.color is not None: + try: + mcolors.to_rgb(self.color) + except ValueError: + raise ValueError(f"Invalid color specification: {self.color}") + + # Validate colormap if specified + if isinstance(self.colormap, str) and self.colormap is not None: + if self.colormap not in plt.colormaps(): + raise ValueError(f"Invalid colormap name: {self.colormap}") + + # Validate numeric ranges + if self.alpha is not None and (not isinstance(self.alpha, (int, float)) or not 0 <= self.alpha <= 1): + raise ValueError(f"Alpha must be between 0 and 1, got {self.alpha}") + if self.vmin is not None and self.vmax is not None and self.vmin >= self.vmax: + raise ValueError(f"vmin ({self.vmin}) must be less than vmax ({self.vmax})") + + +@dataclass +class PropertyLayerPortrayal: + """Manager for property layer styles.""" + layers: Dict[str, PropertyLayerStyle] = field(default_factory=dict) + + def add_layer( + self, + name: str, + color: Optional[str] = None, + colormap: Optional[Union[str, Colormap]] = None, + vmin: Optional[float] = None, + vmax: Optional[float] = None, + alpha: Optional[float] = None, + colorbar: Optional[bool] = None + ) -> None: + """Add a new property layer style. + + Args: + name: Name of the layer + color: A valid matplotlib color string (mutually exclusive with colormap) + colormap: A valid matplotlib colormap name or object (mutually exclusive with color) + vmin: Minimum value for colormap scaling + vmax: Maximum value for colormap scaling + alpha: Opacity value between 0 and 1 + colorbar: Whether to show colorbar + + Example: + >>> portrayal = PropertyLayerPortrayal() + >>> portrayal.add_layer('temperature', colormap='coolwarm', vmin=0, vmax=100) + >>> portrayal.add_layer('elevation', color='brown', vmin=0, vmax=100) + >>> propertylayer_portrayal = dict(portrayal) + + Note: Either color or colormap must be specified, but not both. + """ + self.layers[name] = PropertyLayerStyle( + vmin=vmin, + vmax=vmax, + alpha=alpha, + colorbar=colorbar, + color=color, + colormap=colormap + ) + + def __iter__(self): + return ((name, asdict(style)) for name, style in self.layers.items()) diff --git a/mesa/visualization/mpl_space_drawing.py b/mesa/visualization/mpl_space_drawing.py index 352f05e1c02..022abbf54d4 100644 --- a/mesa/visualization/mpl_space_drawing.py +++ b/mesa/visualization/mpl_space_drawing.py @@ -236,18 +236,22 @@ def draw_property_layers( # Get portrayal properties, or use defaults alpha = portrayal.get("alpha", 1) - vmin = portrayal.get("vmin", np.min(data)) - vmax = portrayal.get("vmax", np.max(data)) + vmin = portrayal.get("vmin") + vmax = portrayal.get("vmax") colorbar = portrayal.get("colorbar", True) + color = portrayal.get("color", None) + colormap = portrayal.get("colormap", None) + vmin = np.min(data) if vmin is None else vmin + vmax = np.max(data) if vmax is None else vmax # Prepare colormap - if "color" in portrayal: - rgba_color = to_rgba(portrayal["color"]) + if color: + rgba_color = to_rgba(color) cmap = LinearSegmentedColormap.from_list( layer_name, [(0, 0, 0, 0), (*rgba_color[:3], alpha)] ) elif "colormap" in portrayal: - cmap = portrayal.get("colormap", "viridis") + cmap = colormap if isinstance(cmap, list): cmap = LinearSegmentedColormap.from_list(layer_name, cmap) elif isinstance(cmap, str): @@ -258,7 +262,7 @@ def draw_property_layers( ) if isinstance(space, OrthogonalGrid): - if "color" in portrayal: + if color: data = data.T normalized_data = (data - vmin) / (vmax - vmin) rgba_data = np.full((*data.shape, 4), rgba_color) @@ -285,12 +289,13 @@ def draw_property_layers( norm = Normalize(vmin=vmin, vmax=vmax) colors = data.ravel() # flatten data to 1D array - if "color" in portrayal: + if color: normalized_colors = np.clip(norm(colors), 0, 1) rgba_colors = np.full((len(colors), 4), rgba_color) rgba_colors[:, 3] = normalized_colors * alpha else: rgba_colors = cmap(norm(colors)) + rgba_colors[..., 3] *= alpha # Draw hexagons collection = PolyCollection(hexagons, facecolors=rgba_colors, zorder=-1) From 17a1a9b8e998a1d2344bd6208ef220e72c3875d5 Mon Sep 17 00:00:00 2001 From: Sahil Chhoker Date: Tue, 4 Feb 2025 13:27:58 +0530 Subject: [PATCH 2/3] Doc: added necessary docstrings --- .../components/portrayal_components.py | 121 ++++++++++-------- 1 file changed, 67 insertions(+), 54 deletions(-) diff --git a/mesa/visualization/components/portrayal_components.py b/mesa/visualization/components/portrayal_components.py index 6ae3b29c670..138f55fdf6c 100644 --- a/mesa/visualization/components/portrayal_components.py +++ b/mesa/visualization/components/portrayal_components.py @@ -1,14 +1,22 @@ -from dataclasses import dataclass, field, asdict -from typing import Optional, Union, ClassVar, Set, Dict -from matplotlib import pyplot as plt -import matplotlib.colors as mcolors +"""Components for defining visual styles and portrayals in Mesa visualizations. + +This module provides classes for configuring the visual appearance of agents and +property layers in Mesa visualization components. +""" + +from collections.abc import Iterator +from dataclasses import asdict, dataclass, field from numbers import Number +from typing import ClassVar + +import matplotlib.colors as mcolors +from matplotlib import pyplot as plt from matplotlib.colors import Colormap + @dataclass class AgentPortrayalStyle: - """ - A class to define the visual style of an agent in Mesa visualizations. + """A class to define the visual style of an agent in Mesa visualizations. Attributes: color: Color specification for the agent. Can be any valid matplotlib color. @@ -24,28 +32,30 @@ class AgentPortrayalStyle: >>> size=70, >>> zorder=1 >>> ) - >>> return dict(portrayal) + >>> return dict(portrayal) """ - - VALID_MARKERS: ClassVar[Set[str]] = { - 'o', 's', '^', 'v', '<', '>', 'D', 'p', 'h', 'H', '8', '*', - '+', 'x', '.', ',', '1', '2', '3', '4', '|', '_', + + # fmt:off + VALID_MARKERS: ClassVar[set[str]] = { + "o", "s", "^", "v", "<", ">", "D", "p", "h", "H", "8", "*", + "+", "x", ".", ",", "1", "2", "3", "4", "|", "_", } + # fmt:on # Default values - color: Optional[Union[str, tuple]] = "tab:blue" - marker: Optional[str] = "o" - size: Optional[Number] = 50 - zorder: Optional[int] = 1 + color: str | tuple | None = "tab:blue" + marker: str | None = "o" + size: Number | None = 50 + zorder: int | None = 1 def __post_init__(self): """Validate the attributes after initialization.""" if self.color is not None: try: mcolors.to_rgb(self.color) - except ValueError: - raise ValueError(f"Invalid color specification: {self.color}") - + except ValueError as err: + raise ValueError(f"Invalid color specification: {self.color}") from err + if self.marker is not None: if not isinstance(self.marker, str): raise ValueError(f"Marker must be a string, got {type(self.marker)}") @@ -53,21 +63,22 @@ def __post_init__(self): raise ValueError( f"Invalid marker '{self.marker}'. Valid markers are: {', '.join(sorted(self.VALID_MARKERS))}" ) - + if self.size is not None and not isinstance(self.size, Number): raise ValueError(f"Size must be a number, got {type(self.size)}") - + if self.zorder is not None and not isinstance(self.zorder, int): raise ValueError(f"Zorder must be an integer, got {type(self.zorder)}") - - def __iter__(self): + + def __iter__(self) -> Iterator[tuple[str, any]]: + """Return an iterator of the style attributes and their values.""" return iter(asdict(self).items()) @dataclass class PropertyLayerStyle: """Style configuration for a single property layer. - + Args: color: A valid matplotlib color string colormap: A valid matplotlib colormap name or Colormap object @@ -78,35 +89,35 @@ class PropertyLayerStyle: Note: color and colormap are mutually exclusive. """ - vmin: Optional[float] = None - vmax: Optional[float] = None - alpha: Optional[float] = None - colorbar: Optional[bool] = None - color: Optional[str] = None - colormap: Optional[Union[str, Colormap]] = None - + + vmin: float | None = None + vmax: float | None = None + alpha: float | None = None + colorbar: bool | None = None + color: str | None = None + colormap: str | Colormap | None = None + def __post_init__(self): """Validate style attributes.""" # Validate color/colormap exclusivity if self.color is not None and self.colormap is not None: raise ValueError("Cannot specify both color and colormap") - if self.color is not None and self.colormap is not None: - raise ValueError("Cannot specify both color and colormap") - + # Validate color if specified if self.color is not None: try: mcolors.to_rgb(self.color) - except ValueError: - raise ValueError(f"Invalid color specification: {self.color}") - + except ValueError as err: + raise ValueError(f"Invalid color specification: {self.color}") from err + # Validate colormap if specified - if isinstance(self.colormap, str) and self.colormap is not None: - if self.colormap not in plt.colormaps(): - raise ValueError(f"Invalid colormap name: {self.colormap}") - + if isinstance(self.colormap, str) and self.colormap not in plt.colormaps(): + raise ValueError(f"Invalid colormap name: {self.colormap}") + # Validate numeric ranges - if self.alpha is not None and (not isinstance(self.alpha, (int, float)) or not 0 <= self.alpha <= 1): + if self.alpha is not None and ( + not isinstance(self.alpha, int | float) or not 0 <= self.alpha <= 1 + ): raise ValueError(f"Alpha must be between 0 and 1, got {self.alpha}") if self.vmin is not None and self.vmax is not None and self.vmin >= self.vmax: raise ValueError(f"vmin ({self.vmin}) must be less than vmax ({self.vmax})") @@ -115,27 +126,28 @@ def __post_init__(self): @dataclass class PropertyLayerPortrayal: """Manager for property layer styles.""" - layers: Dict[str, PropertyLayerStyle] = field(default_factory=dict) - + + layers: dict[str, PropertyLayerStyle] = field(default_factory=dict) + def add_layer( self, name: str, - color: Optional[str] = None, - colormap: Optional[Union[str, Colormap]] = None, - vmin: Optional[float] = None, - vmax: Optional[float] = None, - alpha: Optional[float] = None, - colorbar: Optional[bool] = None + color: str | None = None, + colormap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + alpha: float | None = None, + colorbar: bool | None = None, ) -> None: """Add a new property layer style. - + Args: name: Name of the layer color: A valid matplotlib color string (mutually exclusive with colormap) colormap: A valid matplotlib colormap name or object (mutually exclusive with color) vmin: Minimum value for colormap scaling vmax: Maximum value for colormap scaling - alpha: Opacity value between 0 and 1 + alpha: Opacity value between 0 and 1 colorbar: Whether to show colorbar Example: @@ -143,7 +155,7 @@ def add_layer( >>> portrayal.add_layer('temperature', colormap='coolwarm', vmin=0, vmax=100) >>> portrayal.add_layer('elevation', color='brown', vmin=0, vmax=100) >>> propertylayer_portrayal = dict(portrayal) - + Note: Either color or colormap must be specified, but not both. """ self.layers[name] = PropertyLayerStyle( @@ -152,8 +164,9 @@ def add_layer( alpha=alpha, colorbar=colorbar, color=color, - colormap=colormap + colormap=colormap, ) - def __iter__(self): + def __iter__(self) -> Iterator[tuple[str, dict]]: + """Return an iterator of layer names and their style configurations.""" return ((name, asdict(style)) for name, style in self.layers.items()) From 9af4323f7909c2c8d580009c5419b7cdf82c3f00 Mon Sep 17 00:00:00 2001 From: Sahil Chhoker Date: Tue, 4 Feb 2025 13:47:55 +0530 Subject: [PATCH 3/3] fix: added null checks --- mesa/visualization/components/portrayal_components.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mesa/visualization/components/portrayal_components.py b/mesa/visualization/components/portrayal_components.py index 138f55fdf6c..ede531391dd 100644 --- a/mesa/visualization/components/portrayal_components.py +++ b/mesa/visualization/components/portrayal_components.py @@ -100,6 +100,8 @@ class PropertyLayerStyle: def __post_init__(self): """Validate style attributes.""" # Validate color/colormap exclusivity + if self.color is None and self.colormap is None: + raise ValueError("Please specify either color or colormap") if self.color is not None and self.colormap is not None: raise ValueError("Cannot specify both color and colormap")