New adaptors: --with, --recenter, --diff

src/lib/data/adaptors/bin.py

@@ -141,14 +141,14 @@ def get_name_fragments(self) -> list[str]:
         return [f"bin_{subfrags}"]
 
 
-_BIN_FORMAT = "active_key[=nbins]"
+_BIN_FORMAT = "var_key[=nbins]"
 
 
 @arg_parser(
     dest="adaptors",
     flags=["--bin", "-b"],
     metavar=_BIN_FORMAT,
-    help="Bin the data along these variables, which serve as axes. If nbins is unspecified, it is guessed. Note that t is implicitly binned; disable by passing t= (with no nbins).",
+    help="Bin the data along each of the given variables, which become coordinates. If nbins is unspecified, it is guessed. Note that t is implicitly binned; disable by passing t= (with no nbins).",
     nargs="+",
 )
 def parse_bin(args: list[str]) -> Bin:

src/lib/data/adaptors/derive.py

@@ -183,15 +183,15 @@ def _resolve_from_registry(self, name: str):
 
 _DERIVE_PARSER = Lark(_DERIVE_GRAMMAR)
 
-_DERIVE_FORMAT = "active_key[=expression]"
+_DERIVE_FORMAT = "new_var_key[=expression]"
 _EXPRESSION_DESCRIPTION = "The expression can be any mathematical expression using the standard operators (+, -, *, /, ^), parentheses, signed floating point numbers, and existing variable names."
 
 
 @arg_parser(
     dest="adaptors",
     flags="--derive",
     metavar=_DERIVE_FORMAT,
-    help=f"Create a new variable with the given name. {_EXPRESSION_DESCRIPTION} If the expression is omitted, the variable name is derived via the registry of derivable variables.",
+    help=f"Create a new variable with the given name. {_EXPRESSION_DESCRIPTION} If the expression is omitted, the variable is derived via the registry of derivable variables.",
 )
 def parse_derive(arg: str) -> Derive:
     return Derive(arg)

src/lib/data/adaptors/diff.py

@@ -0,0 +1,78 @@
+from typing import Literal
+
+import xarray as xr
+
+from lib.data.adaptor import BareAdaptor
+from lib.parsing import parse_util
+from lib.parsing.args_registry import arg_parser
+
+type Boundary = Literal["periodic", "pad"]
+
+
+def _diff_one(da: xr.DataArray, dim: str, interp_dir: int, boundary: Boundary) -> xr.DataArray:
+    shifted = da.roll({dim: -interp_dir}, roll_coords=False)
+
+    if boundary == "pad":
+        boundary_idx = 0 if interp_dir == -1 else -1
+        shifted = shifted.copy()
+        shifted[{dim: boundary_idx}] = da[{dim: boundary_idx}]
+
+    return interp_dir * (shifted - da)
+
+
+class Diff(BareAdaptor):
+    def __init__(self, specs: list[tuple[str, int, Boundary]]):
+        self.specs = specs
+
+    def apply_field_bare(self, da: xr.DataArray) -> xr.DataArray:
+        for dim, interp_dir, boundary in self.specs:
+            da = _diff_one(da, dim, interp_dir, boundary)
+        return da
+
+    def get_name_fragments(self) -> list[str]:
+        parts = []
+        prev_boundary: Boundary | None = None
+        for dim, interp_dir, boundary in self.specs:
+            if boundary != prev_boundary:
+                parts.append(boundary)
+                prev_boundary = boundary
+            sign = "+" if interp_dir > 0 else "-"
+            parts.append(f"{dim}={sign}{abs(interp_dir)}")
+        return [f"diff_{'_'.join(parts)}"]
+
+
+PERIODIC_MARKER = "periodic"
+PAD_MARKER = "pad"
+DIR_TO_SHIFT = {"+": 1, "-": -1}
+DIFF_FORMAT = f"[{PERIODIC_MARKER} | {PAD_MARKER}] dim_name[,dim_name...]={set(DIR_TO_SHIFT)} [...]"
+
+
+@arg_parser(
+    dest="adaptors",
+    flags="--diff",
+    metavar=DIFF_FORMAT,
+    help=f"Take the forward ('+') or backward ('-') difference along the given dimension(s). '{PAD_MARKER}'/'{PERIODIC_MARKER}' markers determine how to handle boundaries for subsequent specs (default: {PERIODIC_MARKER}).",
+    nargs="+",
+)
+def parse(args: list[str]) -> Diff:
+    specs: list[tuple[str, int, Boundary]] = []
+    boundary: Boundary = "periodic"
+
+    for arg in args:
+        if arg == PERIODIC_MARKER:
+            boundary = "periodic"
+            continue
+        if arg == PAD_MARKER:
+            boundary = "pad"
+            continue
+
+        dims_arg, interp_dir_arg = parse_util.parse_assignment(arg, DIFF_FORMAT)
+
+        parse_util.check_value(interp_dir_arg, "dir", DIR_TO_SHIFT.keys())
+        interp_dir = DIR_TO_SHIFT[interp_dir_arg]
+
+        for dim in parse_util.parse_comma_separated_list(dims_arg):
+            parse_util.check_identifier(dim, "dim_name")
+            specs.append((dim, interp_dir, boundary))
+
+    return Diff(specs)

src/lib/data/adaptors/nan0.py

@@ -15,4 +15,4 @@ def apply_field_bare(self, da: xr.DataArray) -> xr.DataArray:
         return da.where(da != 0, np.nan)
 
     def get_name_fragments(self) -> list[str]:
-        return [f"nan0"]
+        return ["nan0"]

src/lib/data/adaptors/recenter.py

@@ -0,0 +1,78 @@
+from typing import Literal
+
+import xarray as xr
+
+from lib.data.adaptor import BareAdaptor
+from lib.parsing import parse_util
+from lib.parsing.args_registry import arg_parser
+
+type Boundary = Literal["periodic", "pad"]
+
+
+def _recenter_one(da: xr.DataArray, dim: str, interp_dir: int, boundary: Boundary) -> xr.DataArray:
+    shifted = da.roll({dim: -interp_dir}, roll_coords=False)
+
+    if boundary == "pad":
+        boundary_idx = 0 if interp_dir == -1 else -1
+        shifted = shifted.copy()
+        shifted[{dim: boundary_idx}] = da[{dim: boundary_idx}]
+
+    return 0.5 * (da + shifted)
+
+
+class Recenter(BareAdaptor):
+    def __init__(self, specs: list[tuple[str, int, Boundary]]):
+        self.specs = specs
+
+    def apply_field_bare(self, da: xr.DataArray) -> xr.DataArray:
+        for dim, interp_dir, boundary in self.specs:
+            da = _recenter_one(da, dim, interp_dir, boundary)
+        return da
+
+    def get_name_fragments(self) -> list[str]:
+        parts = []
+        prev_boundary: Boundary | None = None
+        for dim, interp_dir, boundary in self.specs:
+            if boundary != prev_boundary:
+                parts.append(boundary)
+                prev_boundary = boundary
+            sign = "+" if interp_dir > 0 else "-"
+            parts.append(f"{dim}={sign}{abs(interp_dir)}")
+        return [f"recenter_{'_'.join(parts)}"]
+
+
+PERIODIC_MARKER = "periodic"
+PAD_MARKER = "pad"
+DIR_TO_SHIFT = {"+": 1, "-": -1}
+RECENTER_FORMAT = f"[{PERIODIC_MARKER} | {PAD_MARKER}] dim_name[,dim_name...]={set(DIR_TO_SHIFT)} [...]"
+
+
+@arg_parser(
+    dest="adaptors",
+    flags="--recenter",
+    metavar=RECENTER_FORMAT,
+    help=f"Average each value with its neighbor in the given direction per dimension(s). '{PAD_MARKER}'/'{PERIODIC_MARKER}' markers determine how to handle boundaries for subsequent specs (default: {PERIODIC_MARKER}).",
+    nargs="+",
+)
+def parse(args: list[str]) -> Recenter:
+    specs: list[tuple[str, int, Boundary]] = []
+    boundary: Boundary = "periodic"
+
+    for arg in args:
+        if arg == PERIODIC_MARKER:
+            boundary = "periodic"
+            continue
+        if arg == PAD_MARKER:
+            boundary = "pad"
+            continue
+
+        dims_arg, interp_dir_arg = parse_util.parse_assignment(arg, RECENTER_FORMAT)
+
+        parse_util.check_value(interp_dir_arg, "dir", DIR_TO_SHIFT.keys())
+        interp_dir = DIR_TO_SHIFT[interp_dir_arg]
+
+        for dim in parse_util.parse_comma_separated_list(dims_arg):
+            parse_util.check_identifier(dim, "dim_name")
+            specs.append((dim, interp_dir, boundary))
+
+    return Recenter(specs)

src/lib/data/adaptors/with_.py

@@ -0,0 +1,28 @@
+from lib.data.adaptor import Adaptor
+from lib.data.data_with_attrs import DataWithAttrs
+from lib.parsing.args_registry import arg_parser
+
+
+class With(Adaptor):
+    def __init__(self, key: str):
+        self.key = key
+
+    def apply(self, data: DataWithAttrs) -> DataWithAttrs:
+        return data.assign_metadata(
+            active_key=self.key,
+            name_fragments=data.metadata.name_fragments + self.get_name_fragments(),
+        )
+
+    def get_name_fragments(self) -> list[str]:
+        return [f"with_{self.key}"]
+
+
+@arg_parser(
+    dest="adaptors",
+    flags=["--with", "-w"],
+    metavar="key",
+    help="set the active variable to the given key",
+    nargs="just one",
+)
+def parse_with(arg: str) -> With:
+    return With(arg)

src/lib/plotting/hooks/grid.py

@@ -4,7 +4,7 @@
 import numpy.typing as npt
 from matplotlib.axes import Axes
 
-from lib.data.data_with_attrs import DataWithAttrs, List
+from lib.data.data_with_attrs import DataWithAttrs
 from lib.parsing import parse_util
 from lib.parsing.args_registry import arg_parser
 from lib.plotting.frame_data_traits import HasAxes, HasData, assert_impl
@@ -91,11 +91,11 @@ def get_axis_id(data: DataWithAttrs, dim_name: str) -> Literal["x", "y"]:
 
 
 @arg_parser(
+    dest="hooks",
     flags="--grid",
-    metavar=(GRID_FORMAT),
+    metavar=GRID_FORMAT,
     nargs="+",
     help="Mark the given axes with major and/or minor gridlines. The spacing between major grid lines can be specified, as can the number of minor lines between major lines. If unspecified, major/minor lines are placed at preexisting major/minor tick locations. Note that linear axes have no minor ticks by default.",
-    dest="hooks",
 )
 def parse_grid(args: list[str]) -> Grid:
     major: MajorGridParams = {}

src/lib/plotting/hooks/scale.py

@@ -3,7 +3,7 @@
 from matplotlib.colors import SymLogNorm
 from matplotlib.scale import SymmetricalLogScale
 
-from lib.data.data_with_attrs import DataWithAttrs, List
+from lib.data.data_with_attrs import DataWithAttrs
 from lib.parsing import parse_util
 from lib.parsing.args_registry import arg_parser
 from lib.plotting import plt_util

src/lib/plotting/static_plot.py

@@ -34,4 +34,4 @@ def allowed_save_formats(self) -> list[SaveFormat]:
 
     def _save_to_path(self, path: Path):
         self._initialize()
-        self.fig.savefig(path)
+        self.fig.savefig(path, dpi=300.0)