From ceec1151eab2f91bfcd422f3702ffcf88980aa49 Mon Sep 17 00:00:00 2001 From: James McClung Date: Wed, 17 Jun 2026 08:27:16 -0400 Subject: [PATCH 1/6] diff: mod display --- src/lib/data/adaptors/diff.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/lib/data/adaptors/diff.py b/src/lib/data/adaptors/diff.py index 41e3bb3..2e14510 100644 --- a/src/lib/data/adaptors/diff.py +++ b/src/lib/data/adaptors/diff.py @@ -3,6 +3,8 @@ import xarray as xr from lib.data.adaptor import BareAdaptor +from lib.data.data_with_attrs import Metadata +from lib.latex import Latex from lib.parsing import parse_util from lib.parsing.args_registry import arg_parser @@ -24,6 +26,10 @@ class Diff(BareAdaptor): def __init__(self, specs: list[tuple[str, int, Boundary]]): self.specs = specs + def get_modified_display_latex(self, metadata: Metadata) -> Latex: + dims = ",".join(spec[0] for spec in self.specs) + return Latex(f"\\Delta_{{{dims}}}{metadata.active_var_info.display}") + 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) From 9e89c45ed214da7903a60e854c4b32d8e4a133e7 Mon Sep 17 00:00:00 2001 From: James McClung Date: Wed, 17 Jun 2026 08:51:28 -0400 Subject: [PATCH 2/6] diff: +BOUNDARY_KEYS --- src/lib/data/adaptors/diff.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/src/lib/data/adaptors/diff.py b/src/lib/data/adaptors/diff.py index 2e14510..80c9c7e 100644 --- a/src/lib/data/adaptors/diff.py +++ b/src/lib/data/adaptors/diff.py @@ -9,6 +9,7 @@ from lib.parsing.args_registry import arg_parser type Boundary = Literal["periodic", "pad"] +BOUNDARY_KEYS: tuple[Boundary, ...] = Boundary.__value__.__args__ def _diff_one(da: xr.DataArray, dim: str, interp_dir: int, boundary: Boundary) -> xr.DataArray: @@ -47,29 +48,24 @@ def get_name_fragments(self) -> list[str]: 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)} [...]" +DIFF_FORMAT = f"[{' | '.join(BOUNDARY_KEYS)}] 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}).", + help=f"Take the forward ('+') or backward ('-') difference along the given dimension(s). {'/'.join(BOUNDARY_KEYS)} markers determine how to handle boundaries for subsequent specs (default: {BOUNDARY_KEYS[0]}).", nargs="+", ) def parse(args: list[str]) -> Diff: specs: list[tuple[str, int, Boundary]] = [] - boundary: Boundary = "periodic" + boundary = BOUNDARY_KEYS[0] for arg in args: - if arg == PERIODIC_MARKER: - boundary = "periodic" - continue - if arg == PAD_MARKER: - boundary = "pad" + if args in BOUNDARY_KEYS: + boundary = arg continue dims_arg, interp_dir_arg = parse_util.parse_assignment(arg, DIFF_FORMAT) From 014c9ac02b9927cec1f57249b2a155158f1e9395 Mon Sep 17 00:00:00 2001 From: James McClung Date: Wed, 17 Jun 2026 08:54:40 -0400 Subject: [PATCH 3/6] diff: ->dir --- src/lib/data/adaptors/diff.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/src/lib/data/adaptors/diff.py b/src/lib/data/adaptors/diff.py index 80c9c7e..640e4be 100644 --- a/src/lib/data/adaptors/diff.py +++ b/src/lib/data/adaptors/diff.py @@ -12,15 +12,15 @@ BOUNDARY_KEYS: tuple[Boundary, ...] = Boundary.__value__.__args__ -def _diff_one(da: xr.DataArray, dim: str, interp_dir: int, boundary: Boundary) -> xr.DataArray: - shifted = da.roll({dim: -interp_dir}, roll_coords=False) +def _diff_one(da: xr.DataArray, dim: str, dir: int, boundary: Boundary) -> xr.DataArray: + shifted = da.roll({dim: -dir}, roll_coords=False) if boundary == "pad": - boundary_idx = 0 if interp_dir == -1 else -1 + boundary_idx = 0 if dir == -1 else -1 shifted = shifted.copy() shifted[{dim: boundary_idx}] = da[{dim: boundary_idx}] - return interp_dir * (shifted - da) + return dir * (shifted - da) class Diff(BareAdaptor): @@ -32,19 +32,19 @@ def get_modified_display_latex(self, metadata: Metadata) -> Latex: return Latex(f"\\Delta_{{{dims}}}{metadata.active_var_info.display}") 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) + for dim, dir, boundary in self.specs: + da = _diff_one(da, dim, 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: + for dim, 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)}") + sign = "+" if dir > 0 else "-" + parts.append(f"{dim}={sign}{abs(dir)}") return [f"diff_{'_'.join(parts)}"] @@ -68,13 +68,13 @@ def parse(args: list[str]) -> Diff: boundary = arg continue - dims_arg, interp_dir_arg = parse_util.parse_assignment(arg, DIFF_FORMAT) + dims_arg, 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] + parse_util.check_value(dir_arg, "dir", DIR_TO_SHIFT.keys()) + dir = DIR_TO_SHIFT[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)) + specs.append((dim, dir, boundary)) return Diff(specs) From 8ebf867c71415bcd86cecbb8cde09256078b73e7 Mon Sep 17 00:00:00 2001 From: James McClung Date: Wed, 17 Jun 2026 08:56:32 -0400 Subject: [PATCH 4/6] diff: +DiffSpec --- src/lib/data/adaptors/diff.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/lib/data/adaptors/diff.py b/src/lib/data/adaptors/diff.py index 640e4be..ed543ca 100644 --- a/src/lib/data/adaptors/diff.py +++ b/src/lib/data/adaptors/diff.py @@ -10,9 +10,12 @@ type Boundary = Literal["periodic", "pad"] BOUNDARY_KEYS: tuple[Boundary, ...] = Boundary.__value__.__args__ +type DimKey = str +type Dir = Literal[-1, 1] +type DiffSpec = tuple[DimKey, Dir, Boundary] -def _diff_one(da: xr.DataArray, dim: str, dir: int, boundary: Boundary) -> xr.DataArray: +def _diff_one(da: xr.DataArray, dim: DimKey, dir: Dir, boundary: Boundary) -> xr.DataArray: shifted = da.roll({dim: -dir}, roll_coords=False) if boundary == "pad": @@ -24,7 +27,7 @@ def _diff_one(da: xr.DataArray, dim: str, dir: int, boundary: Boundary) -> xr.Da class Diff(BareAdaptor): - def __init__(self, specs: list[tuple[str, int, Boundary]]): + def __init__(self, specs: list[DiffSpec]): self.specs = specs def get_modified_display_latex(self, metadata: Metadata) -> Latex: @@ -60,7 +63,7 @@ def get_name_fragments(self) -> list[str]: nargs="+", ) def parse(args: list[str]) -> Diff: - specs: list[tuple[str, int, Boundary]] = [] + specs: list[DiffSpec] = [] boundary = BOUNDARY_KEYS[0] for arg in args: From 92ef56ce19804d8c582bf9b7c3f3666235b35152 Mon Sep 17 00:00:00 2001 From: James McClung Date: Wed, 17 Jun 2026 09:03:21 -0400 Subject: [PATCH 5/6] diff: +_Diff1d --- src/lib/data/adaptors/diff.py | 58 +++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 27 deletions(-) diff --git a/src/lib/data/adaptors/diff.py b/src/lib/data/adaptors/diff.py index ed543ca..49391e6 100644 --- a/src/lib/data/adaptors/diff.py +++ b/src/lib/data/adaptors/diff.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from typing import Literal import xarray as xr @@ -10,49 +11,52 @@ type Boundary = Literal["periodic", "pad"] BOUNDARY_KEYS: tuple[Boundary, ...] = Boundary.__value__.__args__ -type DimKey = str -type Dir = Literal[-1, 1] -type DiffSpec = tuple[DimKey, Dir, Boundary] -def _diff_one(da: xr.DataArray, dim: DimKey, dir: Dir, boundary: Boundary) -> xr.DataArray: - shifted = da.roll({dim: -dir}, roll_coords=False) +@dataclass +class _Diff1d: + dim_key: str + dir: Literal[-1, 1] + boundary: Boundary - if boundary == "pad": - boundary_idx = 0 if dir == -1 else -1 - shifted = shifted.copy() - shifted[{dim: boundary_idx}] = da[{dim: boundary_idx}] + def apply_field_bare(self, da: xr.DataArray) -> xr.DataArray: + shifted = da.roll({self.dim_key: -self.dir}, roll_coords=False) + + if self.boundary == "pad": + boundary_idx = 0 if self.dir == -1 else -1 + shifted = shifted.copy() + shifted[{self.dim_key: boundary_idx}] = da[{self.dim_key: boundary_idx}] - return dir * (shifted - da) + return self.dir * (shifted - da) class Diff(BareAdaptor): - def __init__(self, specs: list[DiffSpec]): - self.specs = specs + def __init__(self, diffs_1d: list[_Diff1d]): + self.diffs_1d = diffs_1d def get_modified_display_latex(self, metadata: Metadata) -> Latex: - dims = ",".join(spec[0] for spec in self.specs) + dims = ",".join(diff_1d.dim_key for diff_1d in self.diffs_1d) return Latex(f"\\Delta_{{{dims}}}{metadata.active_var_info.display}") def apply_field_bare(self, da: xr.DataArray) -> xr.DataArray: - for dim, dir, boundary in self.specs: - da = _diff_one(da, dim, dir, boundary) + for diff_1d in self.diffs_1d: + da = diff_1d.apply_field_bare(da) return da def get_name_fragments(self) -> list[str]: parts = [] prev_boundary: Boundary | None = None - for dim, dir, boundary in self.specs: - if boundary != prev_boundary: - parts.append(boundary) - prev_boundary = boundary - sign = "+" if dir > 0 else "-" - parts.append(f"{dim}={sign}{abs(dir)}") + for diff_1d in self.diffs_1d: + if diff_1d.boundary != prev_boundary: + parts.append(diff_1d.boundary) + prev_boundary = diff_1d.boundary + sign = "+" if diff_1d.dir > 0 else "-" + parts.append(f"{diff_1d.dim_key}={sign}{abs(diff_1d.dir)}") return [f"diff_{'_'.join(parts)}"] DIR_TO_SHIFT = {"+": 1, "-": -1} -DIFF_FORMAT = f"[{' | '.join(BOUNDARY_KEYS)}] dim_name[,dim_name...]={set(DIR_TO_SHIFT)} [...]" +DIFF_FORMAT = f"[{' | '.join(BOUNDARY_KEYS)}] dim_key[,dim_key...]={set(DIR_TO_SHIFT)} [...]" @arg_parser( @@ -63,11 +67,11 @@ def get_name_fragments(self) -> list[str]: nargs="+", ) def parse(args: list[str]) -> Diff: - specs: list[DiffSpec] = [] + diffs_1d: list[_Diff1d] = [] boundary = BOUNDARY_KEYS[0] for arg in args: - if args in BOUNDARY_KEYS: + if arg in BOUNDARY_KEYS: boundary = arg continue @@ -77,7 +81,7 @@ def parse(args: list[str]) -> Diff: dir = DIR_TO_SHIFT[dir_arg] for dim in parse_util.parse_comma_separated_list(dims_arg): - parse_util.check_identifier(dim, "dim_name") - specs.append((dim, dir, boundary)) + parse_util.check_identifier(dim, "dim_key") + diffs_1d.append(_Diff1d(dim, dir, boundary)) - return Diff(specs) + return Diff(diffs_1d) From 414b57daaa9f6fba17999dfa1cc27e86651fe640 Mon Sep 17 00:00:00 2001 From: James McClung Date: Wed, 17 Jun 2026 09:24:09 -0400 Subject: [PATCH 6/6] diff: +truncate (new default) --- src/lib/data/adaptors/diff.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/lib/data/adaptors/diff.py b/src/lib/data/adaptors/diff.py index 49391e6..940549e 100644 --- a/src/lib/data/adaptors/diff.py +++ b/src/lib/data/adaptors/diff.py @@ -9,7 +9,7 @@ from lib.parsing import parse_util from lib.parsing.args_registry import arg_parser -type Boundary = Literal["periodic", "pad"] +type Boundary = Literal["truncate", "periodic", "pad"] BOUNDARY_KEYS: tuple[Boundary, ...] = Boundary.__value__.__args__ @@ -21,13 +21,17 @@ class _Diff1d: def apply_field_bare(self, da: xr.DataArray) -> xr.DataArray: shifted = da.roll({self.dim_key: -self.dir}, roll_coords=False) + diff = self.dir * (shifted - da) + + if self.boundary == "truncate": + isel = slice(1, None) if self.dir == -1 else slice(0, -1) + return diff.isel({self.dim_key: isel}) if self.boundary == "pad": boundary_idx = 0 if self.dir == -1 else -1 - shifted = shifted.copy() - shifted[{self.dim_key: boundary_idx}] = da[{self.dim_key: boundary_idx}] + diff[{self.dim_key: boundary_idx}] = 0.0 - return self.dir * (shifted - da) + return diff class Diff(BareAdaptor):