Multiplotting overhaul

src/lib/cli.py

@@ -1,117 +1,25 @@
-import sys
-import warnings
-import webbrowser
-from pathlib import Path
-
 import dask
-import matplotlib.pyplot as plt
-
-from lib import parsing
-from lib.config import CONFIG
-from lib.parsing.args import Args
-from lib.plotting.plot import SaveFormat
-
-
-def _resolve_save_format(args: Args) -> SaveFormat | None:
-    if args.save is None:
-        if args.save_format is not None:
-            print("error: --save-format requires --save", file=sys.stderr)
-            sys.exit(1)
-        return None
-
-    if args.save_format == "mp4":
-        if not CONFIG.ffmpeg_bin:
-            print("error: --save-format mp4 requires ffmpeg", file=sys.stderr)
-            sys.exit(1)
-        return "mp4"
-
-    if args.save_format == "gif":
-        return "gif"
-
-    # save_format is None: try mp4, fall back to gif
-    if CONFIG.ffmpeg_bin:
-        return "mp4"
-
-    message = "ffmpeg not found; will save animations as gif instead of mp4"
-    warnings.warn(message)
-    return "gif"
-
-
-def _run_dask_graph(args: Args) -> None:
-    data = args.get_data()
 
-    collections = data.dask_collections()
-    if not collections:
-        print(
-            f"error: --dask-graph requires dask-backed data; pipeline produced eager {type(data).__name__}",
-            file=sys.stderr,
-        )
-        sys.exit(1)
-
-    try:
-        import graphviz  # noqa: F401
-    except ImportError:
-        print(
-            "error: --dask-graph requires the 'graphviz' package; install with `pip install -e \".[dask-graph]\"`",
-            file=sys.stderr,
-        )
-        sys.exit(1)
-
-    save_dir = args.save or Path.cwd()
-    save_dir.mkdir(exist_ok=True, parents=True)
-    path = save_dir / f"{args.get_save_file_stem()}.daskgraph.svg"
-    # dask.visualize's optimize_graph flag only lowers legacy HLG collections
-    # (e.g. dask Arrays), not new-style Expr ones (dask DataFrames) — without
-    # pre-optimizing the latter, un-lowered nodes (e.g. Concat from dd.concat)
-    # fail with NotImplementedError in _layer.
-    collections = [c.optimize() if hasattr(c, "optimize") else c for c in collections]
-    dask.visualize(*collections, filename=str(path), optimize_graph=True)
-    print(f"wrote to {path}")
-
-    if args.show:
-        webbrowser.open(path.absolute().as_uri())
+from lib.config import PscPlotConfig
+from lib.data.compile import compile_action_nodes
+from lib.parsing.parse import parse_args
 
 
 def main():
-    dask.config.set(num_workers=CONFIG.dask_num_workers)
-    if CONFIG.dask_scheduler == "distributed":
+    config = PscPlotConfig.from_env()
+
+    dask.config.set(num_workers=config.dask_num_workers)
+    if config.dask_scheduler == "distributed":
         from dask.distributed import Client, LocalCluster
 
-        cluster = LocalCluster(n_workers=CONFIG.dask_num_workers, threads_per_worker=1, processes=True)
+        cluster = LocalCluster(n_workers=config.dask_num_workers, threads_per_worker=1, processes=True)
         Client(cluster)
-    elif CONFIG.dask_scheduler:
-        dask.config.set(scheduler=CONFIG.dask_scheduler)
-
-    args = parsing.get_parsed_args()
-
-    if args.dask_graph:
-        if args.save_format is not None:
-            warnings.warn("--save-format is ignored with --dask-graph")
-        _run_dask_graph(args)
-        return
-
-    # resolve format BEFORE applying pipeline in order to fail early
-    format = _resolve_save_format(args)
-
-    if format == "mp4":
-        plt.rcParams["animation.ffmpeg_path"] = str(CONFIG.ffmpeg_bin)
-
-    plot = args.get_animation()
-
-    if args.show:
-        plot.show()
-    if args.save is not None:
-        args.save.mkdir(exist_ok=True, parents=True)
+    elif config.dask_scheduler:
+        dask.config.set(scheduler=config.dask_scheduler)
 
-        if format not in plot.allowed_save_formats():
-            if format == args.save_format:  # user actually specified this format
-                message = f"{format} is incompatible with the data; reverting to default ({plot.default_save_format()})"
-                warnings.warn(message)
-            else:
-                assert args.save_format is None
+    args = parse_args()
 
-            format = plot.default_save_format()
+    actions = compile_action_nodes(args, config)
 
-        path = args.save / f"{args.get_save_file_stem()}.{format}"
-        plot.save_to_path(path, dpi=args.save_dpi)
-        print(f"wrote to {path}")
+    for action in actions:
+        action.pull()

src/lib/config.py

@@ -1,6 +1,6 @@
 import os
 import shutil
-from dataclasses import dataclass
+from dataclasses import KW_ONLY, dataclass, field
 from pathlib import Path
 from typing import Callable, Self
 
@@ -19,32 +19,21 @@ def parse_optional[T](s: str | None, parser: Callable[[str], T]) -> T | None:
 
 @dataclass
 class PscPlotConfig:
-    data_dir: Path
-    ffmpeg_bin: Path | None
-    dask_num_workers: int
-    dask_chunk_size: int
-    dask_scheduler: str | None
+    _: KW_ONLY
+    data_dir: Path = field(default_factory=Path.cwd)
+    ffmpeg_bin: Path | None = None
+    dask_num_workers: int = 1
+    dask_chunk_size: int = 1_000_000
+    dask_scheduler: str | None = None
 
     @classmethod
-    def _load(cls) -> Self:
-        data_dir = parse_optional(os.environ.get(_DATA_DIR_KEY), Path)
-        if not data_dir:
-            message = f"Path to data not specified. Set the {_DATA_DIR_KEY} environment variable to specify."
-            raise RuntimeError(message)
+    def from_env(cls) -> Self:
+        config = cls()
 
-        ffmpeg_bin = parse_optional(os.environ.get(_FFMPEG_BIN_KEY, shutil.which("ffmpeg")), Path)
+        config.data_dir = parse_optional(os.environ.get(_DATA_DIR_KEY), Path) or config.data_dir
+        config.ffmpeg_bin = parse_optional(os.environ.get(_FFMPEG_BIN_KEY, shutil.which("ffmpeg")), Path) or config.ffmpeg_bin
+        config.dask_num_workers = parse_optional(os.environ.get(_DASK_NUM_WORKERS_KEY), int) or os.cpu_count() or config.dask_num_workers
+        config.dask_chunk_size = parse_optional(os.environ.get(_DASK_CHUNK_SIZE_KEY), int) or config.dask_chunk_size
+        config.dask_scheduler = os.environ.get(_DASK_SCHEDULER_KEY) or config.dask_scheduler
 
-        dask_num_workers = parse_optional(os.environ.get(_DASK_NUM_WORKERS_KEY), int)
-        if not dask_num_workers:
-            dask_num_workers = os.cpu_count() or 1
-
-        dask_chunk_size = parse_optional(os.environ.get(_DASK_CHUNK_SIZE_KEY), int)
-        if not dask_chunk_size:
-            dask_chunk_size = 1_000_000
-
-        dask_scheduler = os.environ.get(_DASK_SCHEDULER_KEY) or None
-
-        return cls(data_dir, ffmpeg_bin, dask_num_workers, dask_chunk_size, dask_scheduler)
-
-
-CONFIG = PscPlotConfig._load()
+        return config

src/lib/data/adaptor.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
+
 import dask.dataframe as dd
 import pandas as pd
 import xarray as xr
 
 from lib.data.data_with_attrs import DataWithAttrs, Field, List, Metadata
+from lib.data.data_world import DataWorld
 from lib.has_name_fragments import HasNameFragments
 from lib.latex import Latex
 
@@ -19,14 +22,22 @@ def _fail_apply_list(adaptor_type: type[Adaptor]):
     raise RuntimeError(message)
 
 
-class Adaptor(HasNameFragments):
+class WorldAdaptor(ABC, HasNameFragments):
+    @abstractmethod
+    def apply_world(self, world: DataWorld) -> DataWorld: ...
+
+
+class Adaptor(WorldAdaptor):
+    def apply_world(self, world: DataWorld) -> DataWorld:
+        return world.with_active_data(self.apply(world.active_data))
+
     def apply(self, data: DataWithAttrs) -> DataWithAttrs:
         if isinstance(data, List):
             return self.apply_list(data)
         elif isinstance(data, Field):
             return self.apply_field(data)
         else:
-            message = f"unrecognized data type: {data.__class__:r}"
+            message = f"unrecognized data type: {data.__class__!r}"
             raise Exception(message)
 
     def apply_list(self, data: List) -> DataWithAttrs:

src/lib/data/adaptors/bin.py

@@ -10,7 +10,6 @@
 from lib.data.data_with_attrs import Field, FieldMetadata, List
 from lib.parsing import parse_util
 from lib.parsing.args_registry import arg_parser
-from lib.var_info import VarInfo
 
 
 def _guess_bin_edgess(data: List, varname_to_nbins: dict[str, int | None]) -> list:
@@ -81,7 +80,6 @@ def __init__(self, varname_to_nbins: dict[str, int | None]):
         self.varname_to_nbins = varname_to_nbins
 
     def apply_field(self, data: Field) -> Field:
-        da = data.data
         dim_names_to_bin_size = {}
         for dim_name, nbins in self.varname_to_nbins.items():
             if not nbins:
@@ -95,8 +93,7 @@ def apply_field(self, data: Field) -> Field:
 
             dim_names_to_bin_size[dim_name] = bin_size
 
-        da = da.coarsen(dim_names_to_bin_size, boundary="pad").mean()
-        return data.assign_data(da)
+        return data.with_active_data(data.active_data.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)
@@ -137,7 +134,7 @@ def apply_list(self, data: List) -> Field:
         # 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
-        return Field(da.to_dataset(name="f"), FieldMetadata.create_from(data.metadata, active_key="f", var_infos=new_var_infos))
+        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())
@@ -163,15 +160,15 @@ def parse_bin(args: list[str]) -> Bin:
 
         if len(split_arg) == 2 and not split_arg[1]:
             # arg is "t=", i.e., disable implicit binning along t
-            parse_util.check_value(split_arg[0], "active_key", ["t"])
+            parse_util.parse_value(split_arg[0], "active_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 + [""]
 
-        parse_util.check_identifier(active_key, "active_key")
+        parse_util.parse_identifier(active_key, "active_key")
         nbins = parse_util.parse_optional_number(nbins_arg, "nbins", int)
 
         varname_to_nbins[active_key] = nbins

src/lib/data/adaptors/compute.py

@@ -13,4 +13,4 @@ def apply_list(self, data: List) -> FullList:
         return data.compute()
 
     def apply_field(self, data: Field) -> Field:
-        return data.map_data(lambda da: da.compute())
+        return data.with_active_data(data.active_data.compute())

src/lib/data/adaptors/derive.py

@@ -96,7 +96,7 @@ def variable(self, toks: list):
         [tok] = toks
         name = str(tok)
         ds = self._data.data
-        if name not in ds.variables:
+        if name not in ds:
             self._resolve_from_registry(name)
         return self._data.data[name]
 
@@ -129,7 +129,7 @@ def assign_default(self, toks: list):
 
     def assignment(self, toks: list):
         [new_variable, val] = toks
-        new_ds = self._data.data.assign({new_variable: val})
+        new_ds = self._data.data | {new_variable: val}
         dim = var_info_registry.lookup(self._data.metadata.prefix, new_variable)
         new_var_infos = {**self._data.metadata.var_infos, new_variable: dim}
         return self._data.assign(new_ds, active_key=new_variable, var_infos=new_var_infos)

src/lib/data/adaptors/diff.py

@@ -81,11 +81,11 @@ def parse(args: list[str]) -> Diff:
 
         dims_arg, dir_arg = parse_util.parse_assignment(arg, DIFF_FORMAT)
 
-        parse_util.check_value(dir_arg, "dir", DIR_TO_SHIFT.keys())
+        parse_util.parse_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_key")
+            parse_util.parse_identifier(dim, "dim_key")
             diffs_1d.append(_Diff1d(dim, dir, boundary))
 
     return Diff(diffs_1d)

src/lib/data/adaptors/downsample.py

@@ -33,7 +33,7 @@ def parse(args: list[str]) -> Downsample:
     for arg in args:
         [dim_name, bin_size_arg] = parse_util.parse_assignment(arg, DOWNSAMPLE_FORMAT)
 
-        parse_util.check_identifier(dim_name, "dim_name")
+        parse_util.parse_identifier(dim_name, "dim_name")
         bin_size = parse_util.parse_number(bin_size_arg, BIN_SIZE, int)
 
         dim_names_to_bin_size[dim_name] = bin_size

src/lib/data/adaptors/fourier.py

@@ -9,20 +9,20 @@
 from lib.var_info import VarInfo
 
 
-def toggle_fourier(da: xr.DataArray, dim: VarInfo) -> xr.DataArray:
+def toggle_fourier(da: xr.DataArray, info: VarInfo) -> xr.DataArray:
     temp_prefix = "temp_"
-    f_dim = dim.toggle_fourier()
+    f_info = info.toggle_fourier()
 
     # multiply/or divide coords by 2pi to go from frequency <-> angular frequency
 
-    if dim.is_fourier():
-        da = da.assign_coords({dim.key: da.coords[dim.key] / (2 * np.pi)})
-        da = xrft.ifft(da, dim=dim.key, prefix=temp_prefix, lag=0.0)
-        da = da.rename({temp_prefix + dim.key: f_dim.key})
+    if info.is_fourier():
+        da = da.assign_coords({info.key: da.coords[info.key] / (2 * np.pi)})
+        da = xrft.ifft(da, dim=info.key, prefix=temp_prefix, lag=0.0)
+        da = da.rename({temp_prefix + info.key: f_info.key})
     else:
-        da = xrft.fft(da, dim=dim.key, prefix=temp_prefix)
-        da = da.rename({temp_prefix + dim.key: f_dim.key})
-        da = da.assign_coords({f_dim.key: da.coords[f_dim.key] * (2 * np.pi)})
+        da = xrft.fft(da, dim=info.key, prefix=temp_prefix)
+        da = da.rename({temp_prefix + info.key: f_info.key})
+        da = da.assign_coords({f_info.key: da.coords[f_info.key] * (2 * np.pi)})
 
     return da
 
@@ -37,18 +37,17 @@ def apply_field(self, data: Field) -> Field:
         pre_dim_latexs = [data.metadata.var_infos[key].display.latex for key in self.dim_keys]
 
         da = data.active_data
-        new_var_infos = dict(data.metadata.var_infos)
+        new_var_infos = data.metadata.var_infos.copy()
+
         for key in self.dim_keys:
-            dim = new_var_infos[key]
-            f_dim = dim.toggle_fourier()
-            da = toggle_fourier(da, dim)
-            del new_var_infos[key]
-            new_var_infos[f_dim.key] = f_dim
+            info = new_var_infos[key]
+            f_info = info.toggle_fourier()
+            new_var_infos[f_info.key] = f_info
+            da = toggle_fourier(da, info)
 
-        if data.metadata.active_key is not None and data.metadata.active_key in new_var_infos:
-            old_active = new_var_infos[data.metadata.active_key]
-            new_display = f"\\mathcal{{F}}_{{{','.join(pre_dim_latexs)}}}[{old_active.display}]"
-            new_var_infos[data.metadata.active_key] = old_active.assign(display=new_display)
+        old_active_info = data.metadata.active_var_info
+        new_display = f"\\mathcal{{F}}_{{{','.join(pre_dim_latexs)}}}[{old_active_info.display}]"
+        new_var_infos[data.metadata.active_key] = old_active_info.assign(display=new_display)
 
         return data.with_active_data(da).assign_metadata(var_infos=new_var_infos)
 
@@ -68,6 +67,6 @@ def get_name_fragments(self) -> list[str]:
 )
 def parse_fourier(args: list[str]) -> Fourier:
     for dim_name in args:
-        parse_util.check_identifier(dim_name, "dim_name")
+        parse_util.parse_identifier(dim_name, "dim_name")
 
     return Fourier(args)