Add --copy

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)

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,17 +156,14 @@ 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,
             origin="lower",
             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