From db467513d6b58b153838fa4035667fc38ae3704d Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 18 Jun 2026 09:13:14 -0400 Subject: [PATCH 1/3] field_2d: provide axes --- src/lib/plotting/renderers/field_2d.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/lib/plotting/renderers/field_2d.py b/src/lib/plotting/renderers/field_2d.py index f27c1de..b9600f7 100644 --- a/src/lib/plotting/renderers/field_2d.py +++ b/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) From 3b48dcf5896706e627f5386a6e243942d44114a7 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 18 Jun 2026 09:14:32 -0400 Subject: [PATCH 2/3] polar_field: provide axes --- src/lib/plotting/renderers/polar_field.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/lib/plotting/renderers/polar_field.py b/src/lib/plotting/renderers/polar_field.py index 4c04bc4..ba31317 100644 --- a/src/lib/plotting/renderers/polar_field.py +++ b/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) From e67a974abaa0901ac8c8d86ed1805e92c189dfbc Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 18 Jun 2026 08:57:51 -0400 Subject: [PATCH 3/3] +hooks/show_com --- src/lib/plotting/hooks/show_com.py | 46 ++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 src/lib/plotting/hooks/show_com.py diff --git a/src/lib/plotting/hooks/show_com.py b/src/lib/plotting/hooks/show_com.py new file mode 100644 index 0000000..606be16 --- /dev/null +++ b/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()