Cross-directory loading in --with

src/lib/config.py

@@ -20,7 +20,7 @@ def parse_optional[T](s: str | None, parser: Callable[[str], T]) -> T | None:
 @dataclass
 class PscPlotConfig:
     _: KW_ONLY
-    data_dir: Path = field(default_factory=Path.cwd)
+    data_root: Path = field(default_factory=Path.cwd)
     ffmpeg_bin: Path | None = None
     dask_num_workers: int = 1
     dask_chunk_size: int = 1_000_000
@@ -30,7 +30,7 @@ class PscPlotConfig:
     def from_env(cls) -> Self:
         config = cls()
 
-        config.data_dir = parse_optional(os.environ.get(_DATA_DIR_KEY), Path) or config.data_dir
+        config.data_root = parse_optional(os.environ.get(_DATA_DIR_KEY), Path) or config.data_root
         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

src/lib/data/adaptor.py

@@ -29,7 +29,7 @@ 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))
+        return world.with_active(data=self.apply(world.active_data))
 
     def apply(self, data: DataWithAttrs) -> DataWithAttrs:
         if isinstance(data, List):
@@ -59,25 +59,22 @@ def get_modified_unit_latex(self, metadata: Metadata) -> Latex:
     def apply(self, data: DataWithAttrs) -> DataWithAttrs:
         data = super().apply(data)
 
-        var_infos = data.metadata.var_infos
-        if data.metadata.active_key is not None and data.metadata.active_key in var_infos:
+        if info := data.active_info:
             display_latex = self.get_modified_display_latex(data.metadata)
             unit_latex = self.get_modified_unit_latex(data.metadata)
-            old_dim = var_infos[data.metadata.active_key]
-            new_dim = old_dim.assign(display=display_latex, unit=unit_latex)
-            var_infos = {**var_infos, data.metadata.active_key: new_dim}
+            data = data.with_active(info=info.assign(display=display_latex, unit=unit_latex))
 
-        return data.assign_metadata(var_infos=var_infos)
+        return data
 
 
 class BareAdaptor(MetadataAdaptor):
     """An adaptor that works with the raw data, no metadata required."""
 
     def apply_field(self, data: Field) -> DataWithAttrs:
-        return data.with_active_data(self.apply_field_bare(data.active_data))
+        return data.with_active(data=self.apply_field_bare(data.require_active_subdata()))
 
     def apply_list(self, data: List) -> DataWithAttrs:
-        return data.with_active_data(self.apply_list_bare(data.active_data))
+        return data.with_active(data=self.apply_list_bare(data.require_active_subdata()))
 
     def apply_field_bare(self, da: xr.DataArray) -> xr.DataArray:
         _fail_apply_field(self.__class__)

src/lib/data/adaptors/bin.py

@@ -93,7 +93,7 @@ def apply_field(self, data: Field) -> Field:
 
             dim_names_to_bin_size[dim_name] = bin_size
 
-        return data.with_active_data(data.active_data.coarsen(dim_names_to_bin_size, boundary="pad").mean())
+        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)

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.with_active_data(data.active_data.compute())
+        return data.with_active(data=data.require_active_subdata().compute())

src/lib/data/adaptors/derive.py

@@ -3,15 +3,9 @@
 
 from lib import var_info_registry
 from lib.data.adaptor import WorldAdaptor
-from lib.data.data_with_attrs import Field, List
-from lib.data.loader import get_loader
-from lib.derived_field_variables.derived_field_variable import (
-    DERIVED_FIELD_VARIABLES,
-    derive_field_variable,
-)
-from lib.derived_particle_variables.derived_particle_variable import (
-    derive_particle_variable,
-)
+from lib.data.data_world import DataWorld
+from lib.data.ensure_derived import ensure_derived
+from lib.data.loader import load
 from lib.parsing.args_registry import arg_parser
 
 
@@ -21,66 +15,15 @@ def __init__(self, expression: str):
         self.ast = _DERIVE_PARSER.parse(expression)
 
     def apply_world(self, world):
-        active = world.active_data
-        scoped_prefixes = _collect_scoped_prefixes(self.ast)
-
-        if isinstance(active, List):
-            if scoped_prefixes:
-                raise ValueError("--derive: cross-prefix references (prefix::key) are not supported for particle data.")
-            return world.with_active_data(AssignNewVariable(active).transform(self.ast))
-
-        if isinstance(active, Field):
-            siblings = _load_siblings(world, scoped_prefixes)
-            return world.with_active_data(AssignNewFieldVariable(active, siblings).transform(self.ast))
-
-        raise ValueError("--derive requires an active variable to derive into; specify one as a positional argument.")
+        return AssignNewVariable(world).transform(self.ast)
 
     def get_name_fragments(self):
-        if self.ast.data == "assign_default":
-            return []
         return [f'derive_"{self.expression}"']
 
 
-def _collect_scoped_prefixes(ast) -> set[str]:
-    """Distinct prefixes referenced via `prefix::key` anywhere in the expression."""
-    prefixes = set()
-    for tree in ast.find_data("variable"):
-        if len(tree.children) == 2:
-            prefixes.add(str(tree.children[0]))
-    return prefixes
-
-
-def _load_siblings(world, prefixes: set[str]) -> dict[str, Field]:
-    """Resolve each scoped prefix to a loaded Field, reusing any already in the world
-    and auto-loading the rest the same way `--with` does."""
-    siblings: dict[str, Field] = {}
-    for prefix in prefixes:
-        if prefix in world.datas:
-            siblings[prefix] = world.datas[prefix]
-        else:
-            loaded = get_loader(world.config.data_dir, prefix, None).apply_world(world)
-            siblings[prefix] = loaded.datas[prefix]
-    return siblings
-
-
-def _resolve_field_variable(field: Field, key: str) -> Field:
-    """Return a Field guaranteed to contain `key`, deriving it via the prefix's
-    registry when it isn't already in the dataset."""
-    if key in field.data:
-        return field
-    prefix = field.metadata.prefix
-    if prefix is None:
-        raise ValueError(f"--derive cannot resolve '{key}': field metadata has no prefix.")
-    if key not in DERIVED_FIELD_VARIABLES.get(prefix, {}):
-        raise ValueError(
-            f"--derive: '{key}' is not in the '{prefix}' dataset and not in its derived-variable registry {list(DERIVED_FIELD_VARIABLES.get(prefix, {}))}. Note that earlier adaptors (e.g. --downsample) may have dropped variables that became incompatible with the active grid; consider moving --derive earlier in the pipeline."
-        )
-    return derive_field_variable(field, key, prefix)
-
-
 class AssignNewVariable(Transformer_InPlace):
-    def __init__(self, data: List):
-        self._data = data
+    def __init__(self, world: DataWorld):
+        self.world = world
         super().__init__(visit_tokens=True)
 
     def number(self, toks: list):
@@ -89,66 +32,28 @@ def number(self, toks: list):
 
     def new_variable(self, toks: list):
         [tok] = toks
-        return tok
+        return str(tok)
 
     def variable(self, toks: list):
-        [tok] = toks
-        return self._data.data[tok]
-
-    def addition(self, toks: list):
-        [lhs, rhs] = toks
-        return lhs + rhs
-
-    def subtraction(self, toks: list):
-        [lhs, rhs] = toks
-        return lhs - rhs
-
-    def multiplication(self, toks: list):
-        [lhs, rhs] = toks
-        return lhs * rhs
-
-    def division(self, toks: list):
-        [lhs, rhs] = toks
-        return lhs / rhs
-
-    def exponentiation(self, toks: list):
-        [lhs, rhs] = toks
-        return lhs**rhs
-
-    def assign_default(self, toks: list):
-        [new_variable] = toks
-        return derive_particle_variable(self._data, new_variable, "prt")
-
-    def assignment(self, toks: list):
-        [new_variable, val] = toks
-        df = self._data.data
-        df = df.assign(**{new_variable: val})
-        return self._data.assign_data(df)
-
+        key = str(toks[0])
+        data = self.world.require_active_data()
+        data = ensure_derived(data, key)
+        return data[key]
 
-class AssignNewFieldVariable(Transformer_InPlace):
-    def __init__(self, data: Field, siblings: dict[str, Field]):
-        self._data = data
-        self._siblings = siblings
-        super().__init__(visit_tokens=True)
+    def prepath(self, toks: list):
+        return "/".join(str(tok) for tok in toks)
 
-    def number(self, toks: list):
-        [tok] = toks
-        return float(tok)
+    def scoped_variable(self, toks: list):
+        [prepath, key] = [str(tok) for tok in toks]
+        if prepath not in self.world.datas:
+            data = load(self.world.config, prepath)
+        else:
+            data = self.world.datas[prepath]
 
-    def new_variable(self, toks: list):
-        [tok] = toks
-        return str(tok)
+        data = ensure_derived(data, key)
+        self.world = self.world.with_data(prepath, data)
 
-    def variable(self, toks: list):
-        if len(toks) == 2:
-            prefix, key = str(toks[0]), str(toks[1])
-            sibling = _resolve_field_variable(self._siblings[prefix], key)
-            self._siblings[prefix] = sibling
-            return sibling.data[key]
-        key = str(toks[0])
-        self._data = _resolve_field_variable(self._data, key)
-        return self._data.data[key]
+        return data[key]
 
     def addition(self, toks: list):
         [lhs, rhs] = toks
@@ -170,32 +75,25 @@ def exponentiation(self, toks: list):
         [lhs, rhs] = toks
         return lhs**rhs
 
-    def assign_default(self, toks: list):
-        [new_variable] = toks
-        self._data = _resolve_field_variable(self._data, new_variable)
-        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_metadata(active_key=new_variable, var_infos=new_var_infos)
-
     def assignment(self, toks: list):
-        [new_variable, val] = toks
-        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)
+        [key, subdata] = toks
+        data = self.world.require_active_data()
+        info = var_info_registry.lookup(data.metadata.prepath, key)
+        data = data.with_active(data=subdata, key=key, info=info)
+        return self.world.with_active(data=data)
 
 
 _DERIVE_GRAMMAR = r"""
-?start : assign_default | assignment
+?start : assignment
 
-assign_default : new_variable
-assignment     : new_variable "=" expression
-new_variable   : CNAME
+assignment   : new_variable "=" expression
+new_variable : CNAME
 
 ?expression : _expression_3
 
 _expression_0 : "(" expression ")"
               | variable
+              | scoped_variable
               | number
 _expression_1 : _expression_0
               | exponentiation
@@ -212,8 +110,13 @@ def assignment(self, toks: list):
 addition       : _expression_3 "+" _expression_2
 subtraction    : _expression_3 "-" _expression_2
 
-variable : (CNAME "::")? CNAME
-number   : SIGNED_NUMBER
+variable        : CNAME
+scoped_variable : prepath "::" CNAME
+number          : SIGNED_NUMBER
+
+prepath : PREFIX
+# DIR    : /[^\/]+/ TODO figure out a way to make dirs work; as is, the arbitrary chars are incompatible with math symbols, especially /
+PREFIX : /[.\w\d]+/
 
 %import common.SIGNED_NUMBER
 %import common.CNAME
@@ -223,16 +126,14 @@ def assignment(self, toks: list):
 
 
 _DERIVE_PARSER = Lark(_DERIVE_GRAMMAR)
-
-_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. A name may be scoped to another prefix as prefix::key (that prefix is auto-loaded)."
+_DERIVE_FORMAT = "new_var_key=expression"
 
 
 @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 is derived via the registry of derivable variables.",
+    help=f"Create a new variable with the given name. The expression can be any mathematical expression using the standard operators (+, -, *, /, ^), parentheses, signed floating point numbers, and other variable names. A name may be scoped to another prefix as `prepath::key` (similar to `--with`).",
 )
 def parse_derive(arg: str) -> Derive:
     return Derive(arg)

src/lib/data/adaptors/display.py

@@ -6,27 +6,25 @@
 class Display(Adaptor):
     """Override the display-LaTeX of the active variable or of a dimension."""
 
-    def __init__(self, target: str | None, value: str):
-        self.target = target
-        self.value = value
+    def __init__(self, key: str | None, display: str):
+        self.key = key
+        self.display = display
 
     def apply(self, data: DataWithAttrs) -> DataWithAttrs:
         metadata = data.metadata
 
-        target = self.target or metadata.active_key
-        if target is None:
+        key = self.key or metadata.active_key
+        if key is None:
             raise ValueError("--display requires a target; specify a variable as a positional argument or use --display TARGET=VALUE")
 
-        if target not in metadata.var_infos:
-            raise ValueError(f"--display target {target!r} is not a known key ({sorted(metadata.var_infos)})")
+        if key not in metadata.var_infos:
+            raise ValueError(f"--display target {key!r} is not a known key ({sorted(metadata.var_infos)})")
 
-        old_dim = metadata.var_infos[target]
-        new_dim = old_dim.assign(display=self.value)
-        new_var_infos = {**metadata.var_infos, target: new_dim}
-        return data.assign_metadata(var_infos=new_var_infos)
+        info = metadata.var_infos[key].assign(display=self.display)
+        return data.with_info(key, info)
 
     def get_name_fragments(self) -> list[str]:
-        return [f"display_{self.target or 'active'}={self.value}"]
+        return [f"display_{self.key or 'active'}={self.display}"]
 
 
 _DISPLAY_FORMAT = "[name=]display_latex"
@@ -40,6 +38,6 @@ def get_name_fragments(self) -> list[str]:
 )
 def parse_display(arg: str) -> Display:
     if "=" in arg:
-        name, value = arg.split("=", 1)
-        return Display(target=name, value=value)
-    return Display(target=None, value=arg)
+        key, display = arg.split("=", 1)
+        return Display(key=key, display=display)
+    return Display(key=None, display=arg)

src/lib/data/adaptors/fourier.py

@@ -36,7 +36,7 @@ def __init__(self, dim_keys: str | list[str]):
     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
+        da = data.require_active_subdata()
         new_var_infos = data.metadata.var_infos.copy()
 
         for key in self.dim_keys:
@@ -45,11 +45,11 @@ def apply_field(self, data: Field) -> Field:
             new_var_infos[f_info.key] = f_info
             da = toggle_fourier(da, info)
 
-        old_active_info = data.metadata.active_var_info
+        old_active_info = data.active_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)
+        return data.with_active(data=da).assign(var_infos=new_var_infos)
 
     def get_name_fragments(self) -> list[str]:
         return [f"fourier_{','.join(self.dim_keys)}"]