diff --git a/src/lib/data/adaptors/copy.py b/src/lib/data/adaptors/copy.py new file mode 100644 index 0000000..5310b08 --- /dev/null +++ b/src/lib/data/adaptors/copy.py @@ -0,0 +1,49 @@ +from lib.data.adaptor import WorldAdaptor +from lib.data.data_with_attrs import DataWithAttrs +from lib.parsing import parse_util +from lib.parsing.args_registry import arg_parser + + +def _get_default_new_key(data: DataWithAttrs, old_key: str) -> str: + new_key_base = old_key + "_copy" + n_copies = 1 + while (new_key := new_key_base + str(n_copies)) in data.metadata.var_infos: + n_copies += 1 + return new_key + + +class Copy(WorldAdaptor): + def __init__(self, new_key: str | None, old_key: str): + self.new_key = new_key + self.old_key = old_key + + def apply_world(self, world): + active = world.require_active_data() + old = active[self.old_key] + old_info = active.metadata.var_infos[self.old_key] + new_key = self.new_key or _get_default_new_key(active, self.old_key) + return world.with_active(data=active.with_active(key=new_key, data=old, info=old_info)) + + def get_name_fragments(self): + if self.new_key: + return [f"copy_{self.new_key}={self.old_key}"] + return [f"copy_{self.old_key}"] + + +_COPY_FORMAT = "new_key=old_key | old_key" + + +@arg_parser( + dest="adaptors", + flags="--copy", + metavar=_COPY_FORMAT, + help=f"Copy a variable and all its metadata.", +) +def parse_copy(arg: str) -> Copy: + if "=" in arg: + new_key, old_key = parse_util.parse_assignment(arg, _COPY_FORMAT) + else: + new_key, old_key = None, arg + new_key = parse_util.parse_optional_identifier(new_key, "new_key") + old_key = parse_util.parse_identifier(old_key, "old_key") + return Copy(new_key, old_key) diff --git a/src/lib/plotting/setup_fig.py b/src/lib/plotting/setup_fig.py index 857775b..73c7f6e 100644 --- a/src/lib/plotting/setup_fig.py +++ b/src/lib/plotting/setup_fig.py @@ -1,5 +1,6 @@ +import math from abc import ABC, abstractmethod -from typing import Iterable +from typing import Iterable, Literal import numpy as np from matplotlib import pyplot as plt @@ -39,6 +40,21 @@ def _setup_axes(figure: Figure, plot_infos: list[PlotInfo]) -> dict[AxesIdx, tup return ret +def _get_aspect(info: PlotInfo2D) -> Literal["auto", "equal"]: + if info.dim_units[info.x_dim] != info.dim_units[info.y_dim]: + return "auto" + + x_lo, x_hi = info.dim_bounds[info.x_dim] + y_lo, y_hi = info.dim_bounds[info.y_dim] + if None in [x_lo, x_hi, y_lo, y_hi]: + return "auto" + + if math.isclose(x_hi - x_lo, y_hi - y_lo): + return "equal" + + return "auto" + + def _one_or_none[T](objs: Iterable[T]) -> T | None: one = None for obj in objs: @@ -140,10 +156,6 @@ def setup_data(self): class AxesManagerSingleImage(AxesManagerSingle2D[ImageInfo]): - def setup(self): - super().setup() - self.ax.set_aspect(1 / self.ax.get_data_ratio()) - def setup_data(self): image = self.ax.imshow( self.info.data, @@ -151,6 +163,7 @@ def setup_data(self): extent=(*self.info.dim_bounds[self.info.x_dim], *self.info.dim_bounds[self.info.y_dim]), norm=self.info.dim_scales[self.info.color_dim].to_color_norm(), interpolation="nearest", + aspect=_get_aspect(self.info), ) self.info._setter_callbacks["data"] = image.set_data @@ -160,10 +173,6 @@ def setup_data(self): class AxesManagerSingleScatter(AxesManagerSingle2D[ScatterInfo]): - def setup(self): - super().setup() - self.ax.set_aspect(1 / self.ax.get_data_ratio()) - def setup_data(self): if self.info.color_dim: scatter = self.ax.scatter( @@ -185,6 +194,7 @@ def setup_data(self): color=self.ax._get_lines.get_next_color(), s=0.5, ) + self.ax.set_aspect(_get_aspect(self.info)) update_data = lambda _=None: scatter.set_offsets(self.info.xy_data) self.info._setter_callbacks["xy_data"] = update_data @@ -451,10 +461,10 @@ def setup_data(self): image = self.image_ax.imshow( self.image_info.data, origin="lower", - aspect="auto", extent=(*self.image_info.dim_bounds[self.image_info.x_dim], *self.image_info.dim_bounds[self.image_info.y_dim]), norm=self.image_info.dim_scales[self.image_info.color_dim].to_color_norm(), interpolation="nearest", + aspect=_get_aspect(self.image_info), ) self.image_info._setter_callbacks["data"] = image.set_data diff --git a/tests/baseline/test_animated_2d.png b/tests/baseline/test_animated_2d.png index 8982834..c9d64d9 100644 Binary files a/tests/baseline/test_animated_2d.png and b/tests/baseline/test_animated_2d.png differ diff --git a/tests/baseline/test_animated_2d_binned_phase.png b/tests/baseline/test_animated_2d_binned_phase.png index 88d1f32..cca5e37 100644 Binary files a/tests/baseline/test_animated_2d_binned_phase.png and b/tests/baseline/test_animated_2d_binned_phase.png differ diff --git a/tests/baseline/test_animated_2d_derived.png b/tests/baseline/test_animated_2d_derived.png index d249c3f..6c8d058 100644 Binary files a/tests/baseline/test_animated_2d_derived.png and b/tests/baseline/test_animated_2d_derived.png differ diff --git a/tests/baseline/test_animated_scatter_ion_phase.png b/tests/baseline/test_animated_scatter_ion_phase.png index 49e8659..4e155a9 100644 Binary files a/tests/baseline/test_animated_scatter_ion_phase.png and b/tests/baseline/test_animated_scatter_ion_phase.png differ diff --git a/tests/baseline/test_animated_scatter_mul.png b/tests/baseline/test_animated_scatter_mul.png index 8b3e217..695d97a 100644 Binary files a/tests/baseline/test_animated_scatter_mul.png and b/tests/baseline/test_animated_scatter_mul.png differ diff --git a/tests/baseline/test_animated_scatter_with_variable.png b/tests/baseline/test_animated_scatter_with_variable.png index c752eca..e760eee 100644 Binary files a/tests/baseline/test_animated_scatter_with_variable.png and b/tests/baseline/test_animated_scatter_with_variable.png differ diff --git a/tests/baseline/test_crossdata_derive.png b/tests/baseline/test_crossdata_derive.png index a05c0e1..9d30c68 100644 Binary files a/tests/baseline/test_crossdata_derive.png and b/tests/baseline/test_crossdata_derive.png differ diff --git a/tests/baseline/test_hamscan.png b/tests/baseline/test_hamscan.png index 6e77f64..f07af39 100644 Binary files a/tests/baseline/test_hamscan.png and b/tests/baseline/test_hamscan.png differ diff --git a/tests/baseline/test_scale_symlog.png b/tests/baseline/test_scale_symlog.png index a0ed288..37c3c85 100644 Binary files a/tests/baseline/test_scale_symlog.png and b/tests/baseline/test_scale_symlog.png differ diff --git a/tests/baseline/test_spectrum_3d.png b/tests/baseline/test_spectrum_3d.png index a5ca5ba..bb35f6c 100644 Binary files a/tests/baseline/test_spectrum_3d.png and b/tests/baseline/test_spectrum_3d.png differ diff --git a/tests/baseline/test_static_2d.png b/tests/baseline/test_static_2d.png index e2a0add..f727ea2 100644 Binary files a/tests/baseline/test_static_2d.png and b/tests/baseline/test_static_2d.png differ diff --git a/tests/baseline/test_static_2d_spectogram.png b/tests/baseline/test_static_2d_spectogram.png index 2519158..784e9b6 100644 Binary files a/tests/baseline/test_static_2d_spectogram.png and b/tests/baseline/test_static_2d_spectogram.png differ diff --git a/tests/baseline/test_vline.png b/tests/baseline/test_vline.png index 75902af..aeafc59 100644 Binary files a/tests/baseline/test_vline.png and b/tests/baseline/test_vline.png differ