Upgrade --save

CLAUDE.md

@@ -25,7 +25,7 @@ Or directly via `py src/main.py <prepath> [variable] [options]` (backward-compat
 Where `<prepath>` is a **path relative to `PSC_PLOT_DATA_DIR`, ending in the file prefix** — e.g. `pfd` (data root itself) or `run5/pfd` (subdirectory). The prefix part selects the data source: field prefixes (`pfd`, `pfd_moments`, `gauss`, `continuity`) or particle prefixes (`prt`). `Prepath` is just a `str` alias (`src/lib/file_util.py`); `split_prepath()` splits it into `(subdir, prefix)`. Examples live in `plots/check.sh` and `plots/check2.sh` and serve as the de-facto smoke tests / usage reference.
 
 Common flags:
-- `-q` quiet (don't show interactively), `-s [dir]` save output (defaults to `.` if dir omitted)
+- `-q` quiet (don't show interactively), `-s`/`--save` save output — each argument is either a path fragment `[dir/][stem][.ext]` or one of `dir=`/`name=`/`format=`; bare `-s` means cwd, a stem derived from `name_fragments`, and the default format. A directory fragment must end in `/` (or use `dir=`), otherwise it's taken as the filename stem. `--save-dpi` is deliberately separate, pending a `--dpi` that would apply to live figures too
 - `-w`/`--with [prepath::][var_key]` switch the active dataset and/or variable mid-pipeline (loads the prepath if not already in the world); `--copy [new=old|old]` duplicate a variable + its metadata; `-c`/`--compute` force full computation
 - Adaptor flags such as `--roll`, `--reduce`, `--bin`, `--scatter [name]`, `--mag`, `--nan0`, `--scale`, `--fourier`/`-f`, `--pos`, `--species`, `--transform-spherical`, `-v` (versus, sets axes), `-b` (bin)
 
@@ -58,7 +58,7 @@ pytest --mpl
 
 The figure tests live in `tests/test_plots.py` and use `pytest-mpl` for image comparison against baseline PNGs in `tests/baseline/`. Each test runs the full CLI pipeline (parsing → loading → adaptors → plotting) against small test datasets in `tests/data/` (test-2d: x=1,y=8,z=16; test-3d: x=4,y=4,z=4).
 
-Other suites (plain `pytest`, no `--mpl` needed): `test_save.py`/`test_save_filename.py` (save formats and the `name_fragments`-derived filename), `test_idx_efficient.py` (asserts `Idx` prunes dask partitions instead of filtering), `test_particle_bp_vs_h5.py` + `test_h5_species_discovery.py` (loader parity/species discovery), `test_memory.py`, `test_particle_bp_perf.py`, `test_dask_graph.py`, and `test_synthetic_particles.py` (with the `synthetic_particles.py` generator).
+Other suites (plain `pytest`, no `--mpl` needed): `test_save.py` (everything about saving, in three sections: the `--save` argument grammar via `parse_save`, the `name_fragments`-derived filename, and end-to-end pipeline saves), `test_parse_util.py`/`test_parse_errors.py` (parsing helpers and argparse error rendering), `test_idx_efficient.py` (asserts `Idx` prunes dask partitions instead of filtering), `test_particle_bp_vs_h5.py` + `test_h5_species_discovery.py` (loader parity/species discovery), `test_memory.py`, `test_particle_bp_perf.py`, `test_dask_graph.py`, and `test_synthetic_particles.py` (with the `synthetic_particles.py` generator).
 
 To regenerate baselines after intentional visual changes:
 
@@ -118,7 +118,7 @@ The code lives under `src/lib/` and is organized around **loaders** (sources), *
 2. Everything runs through a lazy, memoized **node graph** (`src/lib/data/node.py`). `compile_action_nodes` chains: `RootNode(config)` → `AdaptorNode(With(args.prepath, args.variable))` → one `AdaptorNode` per user adaptor → `PlotNode(hooks)` → action node(s): `ShowPlotNode` (`--show`, default on), `SavePlotNode` (`--save`), and/or `DaskGraphNode` (`--dask-graph`, renders the dask graph as SVG instead of plotting). Note the positional args are just an implicit **`--with`** — loading is `With`'s job, not a dedicated loader node. Each `DataProcessingNode` has a `@cache`d `pull()` and accumulates `name_fragments` (used to derive save filenames). If the adaptor list contains no `Versus`, a default `Versus(["y","z"], time_dim_rule="guess", color_dim=None)` is appended (`_with_versus`) — this is what selects axes/color/time dims and appends the plot target.
 3. Data flows as a **`DataWorld`** (`src/lib/data/data_world.py`): a frozen dataclass holding `datas: dict[Prepath, DataWithAttrs]`, an `active_prepath: Prepath | None`, `plot_targets: list[PlotTarget]`, and `config`. `RootNode.pull()` returns an empty world; each `AdaptorNode` applies a `WorldAdaptor` to it. `active_data` returns `None` when there's no active prepath, `require_active_data()` raises; mutate with `with_active(data=…, prepath=…)` (sets/replaces the active entry) or `with_data(prepath, data)` (inserts without activating — what `Loader.apply_world` does). A `__post_init__` assert enforces that `active_prepath` is a key of `datas`. Holding **multiple named datas + multiple plot targets** is the "split vars" capability this branch is named for.
 4. `PlotNode.pull()` calls `get_plot(world)` (`src/lib/plotting/get_plot.py`), which builds one `Renderer` **per `PlotTarget`**: `Field1dRenderer` when the target has no `color_dim`; `PolarFieldRenderer` / `Field2dRenderer` chosen by `SpatialDimsRTheta` / `SpatialDimsXY`; `ScatterRenderer` for `List` data with exactly 2 spatial dims. It wraps the renderer **list** in `StaticPlot` or `AnimatedPlot` based on `n_frames = max(r.get_n_frames())` — **animated iff `n_frames > 1`**, not on `time_dim` presence. Hooks are then attached via `plot.add_hook()`.
-5. `ShowPlotNode`/`SavePlotNode.pull()` call `plot.show()` / `plot.save_to_path()`, which lazily `_initialize()` the figure (see the PlotInfo/AxesManager layer below).
+5. `ShowPlotNode`/`SavePlotNode.pull()` call `plot.show()` / `plot.save_to_path()`, which lazily `_initialize()` the figure (see the PlotInfo/AxesManager layer below). `SavePlotNode` carries a `SaveSpec(dir, name, format)` (`src/lib/parsing/parse_save.py`) — the parsed `--save` arguments, each `None` when the user didn't specify it — and resolves the defaults at pull time: `dir` → cwd, `name` → `get_save_file_stem()`, `format` → `plot.default_save_format()`, warning and falling back when the requested format isn't in `plot.allowed_save_formats()`. **Resolution lives in the node, not on `SaveSpec`**, because picking a default format needs the `Plot`; that's also why an unwritable extension isn't rejected at parse time, and why `--save` accepts any extension. An explicit `name=` is used verbatim — only the derived stem goes through `sanitize_stem`.
 
 ### PlotTarget
 

src/lib/data/adaptors/bin.py

@@ -7,81 +7,82 @@
 
 from lib import var_info_registry
 from lib.data.adaptor import MetadataAdaptor
-from lib.data.data_with_attrs import Field, FieldMetadata, List
+from lib.data.data_with_attrs import Field, FieldMetadata, LazyList, List
+from lib.data.types import VarKey
 from lib.parsing import parse_util
 from lib.parsing.args_registry import arg_parser
 
 
-def _guess_bin_edgess(data: List, varname_to_nbins: dict[str, int | None]) -> list:
-    varname_to_edges: dict[str, np.array] = {}
+def _guess_bin_edgess(data: List, keys_to_nbins: dict[VarKey, int | None]) -> list:
+    keys_to_edges: dict[VarKey, np.ndarray] = {}
 
-    compute_varnames = []
+    compute_keys: list[VarKey] = []
     mins_to_compute = []
     maxs_to_compute = []
     nbins_so_far = 1
-    varnames_with_missing_nbins = []
+    keys_with_missing_nbins: list[VarKey] = []
 
     df = data.data
 
     # Calculate edges using metadata when possible
 
-    for varname, nbins in varname_to_nbins.items():
-        if varname in data.coordss():
-            coords = data.coordss()[varname]
+    for key, nbins in keys_to_nbins.items():
+        if key in data.coordss():
+            coords = data.coordss()[key]
             if nbins is None:
                 nbins = len(coords)
                 # note: use inf as right edge for convenience; it gets sliced out later
-                varname_to_edges[varname] = np.concat((coords, [np.inf]))
+                keys_to_edges[key] = np.concatenate((coords, [np.inf]))
             else:
-                varname_to_edges[varname] = np.linspace(coords[0], coords[-1] + coords[1] - coords[0], nbins + 1, endpoint=True)
-        elif varname in data.metadata.var_infos and (data.metadata.var_infos[varname].geometry == "polar:theta" or data.metadata.var_infos[varname].geometry == "spherical:phi"):
-            varname_to_edges[varname] = np.linspace(-np.pi, np.pi, nbins + 1, endpoint=True)
-        elif varname in data.metadata.var_infos and data.metadata.var_infos[varname].geometry == "spherical:theta":
-            varname_to_edges[varname] = np.linspace(0.0, np.pi, nbins + 1, endpoint=True)
+                keys_to_edges[key] = np.linspace(coords[0], coords[-1] + coords[1] - coords[0], nbins + 1, endpoint=True)
+        elif key in data.metadata.var_infos and (data.metadata.var_infos[key].geometry == "polar:theta" or data.metadata.var_infos[key].geometry == "spherical:phi"):
+            keys_to_edges[key] = np.linspace(-np.pi, np.pi, nbins + 1, endpoint=True)
+        elif key in data.metadata.var_infos and data.metadata.var_infos[key].geometry == "spherical:theta":
+            keys_to_edges[key] = np.linspace(0.0, np.pi, nbins + 1, endpoint=True)
         else:
-            compute_varnames.append(varname)
-            mins_to_compute.append(df[varname].min())
-            maxs_to_compute.append(df[varname].max())
+            compute_keys.append(key)
+            mins_to_compute.append(df[key].min())
+            maxs_to_compute.append(df[key].max())
 
         if nbins:
             nbins_so_far *= nbins
         else:
-            varnames_with_missing_nbins.append(varname)
+            keys_with_missing_nbins.append(key)
 
     # If needed, batch-compute the missing edges
 
-    if compute_varnames:
+    if compute_keys:
         if isinstance(df, dd.DataFrame):
             computed_mins, computed_maxs = dask.array.compute(mins_to_compute, maxs_to_compute)
         else:
             computed_mins, computed_maxs = mins_to_compute, maxs_to_compute
 
-        if varnames_with_missing_nbins:
+        if keys_with_missing_nbins:
             # split bins evenly across remaining dimensions
             n_data = len(df)
             mean_n_data_per_bin_so_far = n_data / nbins_so_far
             target_mean_n_data_per_bin = 10  # arbitrary number
 
-            guessed_nbins = math.ceil((mean_n_data_per_bin_so_far / target_mean_n_data_per_bin) ** (1 / len(compute_varnames)))
-            for varname_with_missing_nbins in varnames_with_missing_nbins:
-                varname_to_nbins[varname_with_missing_nbins] = guessed_nbins
+            guessed_nbins = math.ceil((mean_n_data_per_bin_so_far / target_mean_n_data_per_bin) ** (1 / len(compute_keys)))
+            for varname_with_missing_nbins in keys_with_missing_nbins:
+                keys_to_nbins[varname_with_missing_nbins] = guessed_nbins
 
-        for varname, min, max in zip(compute_varnames, computed_mins, computed_maxs):
-            nbins = varname_to_nbins[varname]
-            varname_to_edges[varname] = np.linspace(min, max, nbins + 1, endpoint=True)
+        for key, min, max in zip(compute_keys, computed_mins, computed_maxs):
+            nbins = keys_to_nbins[key]
+            keys_to_edges[key] = np.linspace(min, max, nbins + 1, endpoint=True)
 
     # ensure edges are in same order as bin values
-    edgess = [varname_to_edges[varname] for varname in varname_to_nbins]
+    edgess = [keys_to_edges[key] for key in keys_to_nbins]
     return edgess
 
 
 class Bin(MetadataAdaptor):
-    def __init__(self, varname_to_nbins: dict[str, int | None]):
-        self.varname_to_nbins = varname_to_nbins
+    def __init__(self, key_to_nbins: dict[VarKey, int | None]):
+        self.keys_to_nbins = key_to_nbins
 
     def apply_field(self, data: Field) -> Field:
         dim_names_to_bin_size = {}
-        for dim_name, nbins in self.varname_to_nbins.items():
+        for dim_name, nbins in self.keys_to_nbins.items():
             if not nbins:
                 continue
 
@@ -96,48 +97,47 @@ def apply_field(self, data: Field) -> Field:
         return data.with_active(data=data.require_active_subdata().coarsen(dim_names_to_bin_size, boundary="pad").mean())
 
     def apply_list(self, data: List) -> Field:
-        bin_edgess = _guess_bin_edgess(data, self.varname_to_nbins)
+        bin_edgess = _guess_bin_edgess(data, self.keys_to_nbins)
 
-        df = data.data
-        if isinstance(df, dd.DataFrame):
+        if isinstance(data, LazyList):
             binned_data, _ = dask.array.histogramdd(
-                [df[active_key].to_dask_array() for active_key in self.varname_to_nbins],
+                [data[key].to_dask_array() for key in self.keys_to_nbins],
                 bin_edgess,
                 density=False,
-                weights=df[data.metadata.weight_key].to_dask_array() if data.metadata.weight_key else None,
+                weights=data[data.metadata.weight_key].to_dask_array() if data.metadata.weight_key else None,
             )
         else:
             binned_data, _ = np.histogramdd(
-                [df[active_key] for active_key in self.varname_to_nbins],
+                [data[key] for key in self.keys_to_nbins],
                 bin_edgess,
                 density=False,
-                weights=df[data.metadata.weight_key] if data.metadata.weight_key else None,
+                weights=data[data.metadata.weight_key] if data.metadata.weight_key else None,
             )
 
         # note: the slice removes any infs
-        coords = dict(zip(self.varname_to_nbins.keys(), (edges[:-1] for edges in bin_edgess)))
+        coords = dict(zip(self.keys_to_nbins.keys(), (edges[:-1] for edges in bin_edgess)))
 
         da = xr.DataArray(
             binned_data,
             coords,
-            dims=self.varname_to_nbins.keys(),
+            dims=self.keys_to_nbins.keys(),
         )
 
-        f_dim = var_info_registry.lookup("prt", "f")
+        f_info = var_info_registry.lookup("prt", "f")
 
         subject = data.metadata.subject
         if subject is not None and subject.latex == r"\text{Ions}":
-            f_dim = f_dim.assign(display=f_dim.display.latex + r"_\text{i}")
+            f_info = f_info.assign(display=f_info.display.latex + r"_\text{i}")
         elif subject is not None and subject.latex == r"\text{Electrons}":
-            f_dim = f_dim.assign(display=f_dim.display.latex + r"_\text{e}")
+            f_info = f_info.assign(display=f_info.display.latex + r"_\text{e}")
         new_var_infos = {key: data.metadata.var_infos[key] for key in da.coords if key in data.metadata.var_infos}
         # want: psc-plot prt.i --derive K="ux^2+uy^2+uz^2" --bin K=128 -v K --scale log
 
-        new_var_infos["f"] = f_dim
+        new_var_infos["f"] = f_info
         return Field({"f": da}, FieldMetadata.create_from(data.metadata, active_key="f", var_infos=new_var_infos))
 
     def get_name_fragments(self) -> list[str]:
-        subfrags = "_".join(f"{varname}={nbins}" if nbins else varname for varname, nbins in self.varname_to_nbins.items())
+        subfrags = "_".join(f"{varname}={nbins}" if nbins else varname for varname, nbins in self.keys_to_nbins.items())
         return [f"bin_{subfrags}"]
 
 
@@ -152,30 +152,30 @@ def get_name_fragments(self) -> list[str]:
     nargs="+",
 )
 def parse_bin(args: list[str]) -> Bin:
-    varname_to_nbins = {}
+    keys_to_nbins = {}
     insert_bin_t = True
 
     for arg in args:
         split_arg = arg.split("=")
 
         if len(split_arg) == 2 and not split_arg[1]:
             # arg is "t=", i.e., disable implicit binning along t
-            parse_util.parse_value(split_arg[0], "active_key", ["t"])
+            parse_util.parse_value(split_arg[0], "var_key", ["t"])
             insert_bin_t = False
             continue
         elif len(split_arg) > 2:
             parse_util.fail_format(arg, _BIN_FORMAT)
 
-        [active_key, nbins_arg, *_] = split_arg + [""]
+        [key, nbins_arg, *_] = split_arg + [""]
 
-        parse_util.parse_identifier(active_key, "active_key")
+        parse_util.parse_identifier(key, "var_key")
         nbins = parse_util.parse_optional_number(nbins_arg, "nbins", int)
 
-        varname_to_nbins[active_key] = nbins
-        if active_key == "t":
+        keys_to_nbins[key] = nbins
+        if key == "t":
             insert_bin_t = False
 
     if insert_bin_t:
-        varname_to_nbins["t"] = None
+        keys_to_nbins["t"] = None
 
-    return Bin(varname_to_nbins)
+    return Bin(keys_to_nbins)

src/lib/data/compile.py

@@ -42,28 +42,17 @@ def compile_action_nodes(args: Args, config: PscPlotConfig) -> list[DataProcessi
     action_nodes = []
 
     if args.dask_graph:
-        action_nodes.append(DaskGraphNode(plot_node.input_node, save_dir=args.save, show=args.show))
+        action_nodes.append(DaskGraphNode(plot_node.input_node, save=args.save, show=args.show))
         return action_nodes
 
     if args.show:
         action_nodes.append(ShowPlotNode(plot_node))
 
-    if args.save is None and args.save_format:
-        print("error: --save-format requires --save", file=sys.stderr)
-        sys.exit(1)
-
-    if args.save_format == "mp4" and not config.ffmpeg_bin:
-        print("error: --save-format mp4 requires ffmpeg", file=sys.stderr)
-        sys.exit(1)
-
     if args.save is not None:
-        action_nodes.append(
-            SavePlotNode(
-                plot_node,
-                save_dir=args.save,
-                save_format=args.save_format,
-                save_dpi=args.save_dpi,
-            )
-        )
+        if args.save.format == "mp4" and not config.ffmpeg_bin:
+            print("error: format=mp4 requires ffmpeg", file=sys.stderr)
+            sys.exit(1)
+
+        action_nodes.append(SavePlotNode(plot_node, save=args.save, save_dpi=args.save_dpi))
 
     return action_nodes