fix(api): 🐛 fix draw actors methods

SongshGeoLab · Feb 16, 2025 · fff471b · fff471b
1 parent a87bf6a
commit fff471b
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 73 deletions.
diff --git a/abses/tools/viz.py b/abses/tools/viz.py
diff --git a/abses/viz/viz_actors.py b/abses/viz/viz_actors.py
@@ -5,8 +5,8 @@
 # GitHub   : https://github.com/SongshGeo
 # Website: https://cv.songshgeo.com/
 
-"""Visualize ActorsList
-"""
+"""Visualize ActorsList"""
+
 from __future__ import annotations
 
 import contextlib
@@ -28,7 +28,6 @@
     from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    from abses.actor import Actor, GeoType
     from abses.main import MainModel
     from abses.sequences import ActorsList
 
@@ -51,7 +50,7 @@ def hist(
         """Plot hist."""
         df = self.actors.summary(attrs=attr)
         if palette is None:
-            palette = self._style_dict("color", "blue")
+            palette = self._style_dict(request="color", default="blue")
         sns.histplot(df, x=attr, ax=ax, hue="breed", palette=palette)
         if savefig:
             plt.savefig(savefig)
@@ -60,11 +59,9 @@ def hist(
 
     def _style_dict(self, request: str, default: Any, **kwargs) -> StyleDict:
         styles = {}
-        for breed in self.actors.to_dict().keys():
-            # TODO: 这里需要改成 agent_types
-            breed_cls = self.model.breeds[breed]
-            style_dict = breed_cls.viz_attrs(render_marker=True, **kwargs)
-            styles[breed] = style_dict.get(request, default)
+        for breed_cls in self.model.agents_by_type.keys():
+            style_dict = breed_cls.viz_attrs(**kwargs)
+            styles[breed_cls.__name__] = style_dict.get(request, default)
         return styles
 
     @with_axes
@@ -83,9 +80,7 @@ def display(
         subset = self.actors.select(geo_type="Shape")
         if not subset:
             return ax
-        data = gpd.GeoSeries(
-            subset.array("geometry"), crs=self.model.nature.crs
-        )
+        data = gpd.GeoSeries(subset.array("geometry"), crs=self.model.nature.crs)
         obj = data.boundary if boundary else data
         obj.plot(ax=ax, alpha=alpha, **kwargs)
         return ax
@@ -108,9 +103,7 @@ def positions(
         y_axis = "y" if coords else "row"
         x_axis = "x" if coords else "col"
         count_data = (
-            data.groupby([y_axis, x_axis, "breed"])
-            .size()
-            .reset_index(name="count")
+            data.groupby([y_axis, x_axis, "breed"]).size().reset_index(name="count")
         )
         if hue == "breed":
             palette = self._style_dict("color", "blue")

diff --git a/abses/viz/viz_nature.py b/abses/viz/viz_nature.py
@@ -5,19 +5,19 @@
 # GitHub   : https://github.com/SongshGeo
 # Website: https://cv.songshgeo.com/
 
-"""Viz natural module.
-"""
+"""Viz natural module."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Dict, Optional, cast
 
 from matplotlib.axes import Axes
 
 from abses.tools.func import with_axes
-from abses.tools.viz import COLOR_BAR
+
+COLOR_BAR = {"fraction": 0.03, "pad": 0.04}
 
 if TYPE_CHECKING:
-    from abses.actor import Actor
     from abses.nature import PatchModule
 
 
@@ -50,9 +50,8 @@ def show(
         else:
             xda = self.module.get_xarray(attr)
             xda.plot(ax=ax, cbar_kwargs=COLOR_BAR, alpha=0.8)
-        if self.model.breeds and with_actors:
-            self.model.actors.plot.show(ax=ax, **scatter_kwargs)
         ax.axes.set_aspect("equal")
-        if self.model.breeds:
+        if with_actors:
+            self.model.actors.plot.show(ax=ax, **scatter_kwargs)
             ax.legend(**legend_kwargs)
         return ax