Add --show-com (center of mass) hook

src/lib/plotting/hooks/show_com.py

@@ -0,0 +1,46 @@
+from lib.data.data_with_attrs import Field
+from lib.parsing.args_registry import const_arg
+from lib.plotting.frame_data_traits import HasAxes, HasFieldData, assert_impl
+from lib.plotting.hook import Hook
+
+
+def _get_center(field: Field, dim: str) -> float:
+    other_dims = set(field.dims) - {dim}
+    summed = field.active_data.sum(other_dims)
+    return (summed * field.coordss[dim]).sum(dim) / summed.sum(dim)
+
+
+def _get_centers(field: Field) -> list[float]:
+    return [_get_center(field, dim) for dim in field.dims]
+
+
+class ShowCom(Hook):
+    class PostInitData(HasFieldData, HasAxes): ...
+
+    def post_init_fig(self, init_data):
+        init_data = assert_impl(init_data, ShowCom.PostInitData)
+
+        data = init_data.data
+        assert len(data.dims) == 2
+        self.scatter = init_data.axes.scatter(*_get_centers(data), marker="x", label="center of mass")
+
+        init_data.axes.legend()
+
+    class PostUpdateData(HasFieldData): ...
+
+    def post_update_fig(self, update_data):
+        update_data = assert_impl(update_data, ShowCom.PostUpdateData)
+
+        self.scatter.set_offsets([_get_centers(update_data.data)])
+
+    def get_name_fragments(self) -> list[str]:
+        return ["show_com"]
+
+
+@const_arg(
+    flags="--show-com",
+    help="show the center of mass",
+    dest="hooks",
+)
+def get() -> ShowCom:
+    return ShowCom()

src/lib/plotting/renderers/field_2d.py

@@ -6,7 +6,7 @@
 
 from lib.data.data_with_attrs import Field
 from lib.plotting import plt_util
-from lib.plotting.frame_data_traits import HasColorNorm, HasFieldData, HasSpatialScales
+from lib.plotting.frame_data_traits import HasAxes, HasColorNorm, HasFieldData, HasSpatialScales
 from lib.plotting.renderer import Renderer
 
 
@@ -18,10 +18,10 @@ def get_extent(da: xr.DataArray, dim: str) -> tuple[float, float]:
 
 class Field2dRenderer(Renderer[Field]):
     @dataclass(kw_only=True)
-    class InitData(HasFieldData, HasSpatialScales, HasColorNorm): ...
+    class InitData(HasFieldData, HasSpatialScales, HasColorNorm, HasAxes): ...
 
     @dataclass(kw_only=True)
-    class UpdateData(HasFieldData): ...
+    class UpdateData(HasFieldData, HasAxes): ...
 
     def _transpose(self, data: Field) -> Field:
         spatial_dims = data.metadata.spatial_dims
@@ -33,6 +33,7 @@ def make_init_data(self, fig: Figure, ax: Axes, frame_data: Field) -> InitData:
             spatial_scales=["linear", "linear"],
             color_norm="linear",
             color_is_dependent=True,
+            axes=ax,
         )
 
     def init(self, fig: Figure, ax: Axes, full_data: Field, frame_data: Field, init_data: InitData) -> None:
@@ -63,7 +64,7 @@ def init(self, fig: Figure, ax: Axes, full_data: Field, frame_data: Field, init_
         ax.set_ylabel(frame_data.metadata.var_infos[spatial_dims[1]].to_axis_label())
 
     def make_update_data(self, ax: Axes, frame_data: Field) -> UpdateData:
-        return self.UpdateData(data=frame_data)
+        return self.UpdateData(data=frame_data, axes=ax)
 
     def draw(self, ax: Axes, frame_data: Field, update_data: UpdateData) -> None:
         frame_data = self._transpose(frame_data)

src/lib/plotting/renderers/polar_field.py

@@ -7,16 +7,16 @@
 
 from lib.data.data_with_attrs import Field
 from lib.plotting import plt_util
-from lib.plotting.frame_data_traits import HasColorNorm, HasFieldData, HasSpatialScales
+from lib.plotting.frame_data_traits import HasAxes, HasColorNorm, HasFieldData, HasSpatialScales
 from lib.plotting.renderer import Renderer
 
 
 class PolarFieldRenderer(Renderer[Field]):
     @dataclass(kw_only=True)
-    class InitData(HasFieldData, HasSpatialScales, HasColorNorm): ...
+    class InitData(HasFieldData, HasSpatialScales, HasColorNorm, HasAxes): ...
 
     @dataclass(kw_only=True)
-    class UpdateData(HasFieldData): ...
+    class UpdateData(HasFieldData, HasAxes): ...
 
     def subplot_kw(self) -> dict[str, Any]:
         return {"projection": "polar"}
@@ -27,6 +27,7 @@ def make_init_data(self, fig: Figure, ax: Axes, frame_data: Field) -> InitData:
             spatial_scales=["linear", "linear"],
             color_norm="linear",
             color_is_dependent=True,
+            axes=ax,
         )
 
     def init(self, fig: Figure, ax: Axes, full_data: Field, frame_data: Field, init_data: InitData) -> None:
@@ -66,7 +67,7 @@ def init(self, fig: Figure, ax: Axes, full_data: Field, frame_data: Field, init_
         # ax.set_ylabel(frame_data.metadata.var_info[spatial_dims[0]].to_axis_label())
 
     def make_update_data(self, ax: Axes, frame_data: Field) -> UpdateData:
-        return self.UpdateData(data=frame_data)
+        return self.UpdateData(data=frame_data, axes=ax)
 
     def draw(self, ax: Axes, frame_data: Field, update_data: UpdateData) -> None:
         self.im.set_array(frame_data.active_data)