diff --git a/src/lib/data/adaptors/diff.py b/src/lib/data/adaptors/diff.py index 41e3bb3..940549e 100644 --- a/src/lib/data/adaptors/diff.py +++ b/src/lib/data/adaptors/diff.py @@ -1,78 +1,91 @@ +from dataclasses import dataclass from typing import Literal 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 -type Boundary = Literal["periodic", "pad"] +type Boundary = Literal["truncate", "periodic", "pad"] +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) +@dataclass +class _Diff1d: + dim_key: str + dir: Literal[-1, 1] + boundary: Boundary - if boundary == "pad": - boundary_idx = 0 if interp_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) + 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 + diff[{self.dim_key: boundary_idx}] = 0.0 - return interp_dir * (shifted - da) + return diff class Diff(BareAdaptor): - def __init__(self, specs: list[tuple[str, int, Boundary]]): - 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(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, interp_dir, boundary in self.specs: - da = _diff_one(da, dim, interp_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, 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)}") + 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)}"] -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_key[,dim_key...]={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" + diffs_1d: list[_Diff1d] = [] + boundary = BOUNDARY_KEYS[0] for arg in args: - if arg == PERIODIC_MARKER: - boundary = "periodic" - continue - if arg == PAD_MARKER: - boundary = "pad" + if arg in BOUNDARY_KEYS: + 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)) + parse_util.check_identifier(dim, "dim_key") + diffs_1d.append(_Diff1d(dim, dir, boundary)) - return Diff(specs) + return Diff(diffs_1d)