Minor fixes

src/lib/cli.py

@@ -58,7 +58,7 @@ def _run_dask_graph(args: Args) -> None:
         sys.exit(1)
 
     save_dir = args.save or Path.cwd()
-    save_dir.mkdir(exist_ok=True)
+    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
@@ -101,7 +101,7 @@ def main():
     if args.show:
         plot.show()
     if args.save is not None:
-        args.save.mkdir(exist_ok=True)
+        args.save.mkdir(exist_ok=True, parents=True)
 
         if format not in plot.allowed_save_formats():
             if format == args.save_format:  # user actually specified this format

src/lib/data/adaptor.py

@@ -6,6 +6,7 @@
 
 from lib.data.data_with_attrs import DataWithAttrs, Field, List, Metadata
 from lib.has_name_fragments import HasNameFragments
+from lib.latex import Latex
 
 
 def _fail_apply_field(adaptor_type: type[Adaptor]):
@@ -38,11 +39,11 @@ def apply_field(self, data: Field) -> DataWithAttrs:
 class MetadataAdaptor(Adaptor):
     """Wraps `apply` to perform standard metadata mutations."""
 
-    def get_modified_display_latex(self, metadata: Metadata) -> str:
-        return metadata.active_var_info.display.latex
+    def get_modified_display_latex(self, metadata: Metadata) -> Latex:
+        return metadata.active_var_info.display
 
-    def get_modified_unit_latex(self, metadata: Metadata) -> str:
-        return metadata.active_var_info.unit.latex
+    def get_modified_unit_latex(self, metadata: Metadata) -> Latex:
+        return metadata.active_var_info.unit
 
     def apply(self, data: DataWithAttrs) -> DataWithAttrs:
         data = super().apply(data)

src/lib/data/adaptors/bin.py

@@ -10,6 +10,7 @@
 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:
@@ -133,6 +134,8 @@ def apply_list(self, data: List) -> Field:
         elif subject is not None and subject.latex == r"\text{Electrons}":
             f_dim = f_dim.assign(display=f_dim.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
         return Field(da.to_dataset(name="f"), FieldMetadata.create_from(data.metadata, active_key="f", var_infos=new_var_infos))
 

src/lib/data/adaptors/fourier.py

@@ -52,12 +52,6 @@ def apply_field(self, data: Field) -> Field:
 
         return data.with_active_data(da).assign_metadata(var_infos=new_var_infos)
 
-    def get_modified_display_latex(self, metadata) -> str:
-        return metadata.active_var_info.display.latex
-
-    def get_modified_unit_latex(self, metadata) -> str:
-        return metadata.active_var_info.unit.latex
-
     def get_name_fragments(self) -> list[str]:
         return [f"fourier_{','.join(self.dim_keys)}"]
 

src/lib/data/adaptors/idx.py

@@ -16,26 +16,29 @@ def apply_list(self, data: List) -> List:
         coordss = data.coordss.copy()
         df = data.data
 
-        pdim = data.metadata.partition_dim
-        pranges = data.metadata.partition_ranges
         for dim, isel in self.dim_names_to_isel.items():
             if dim not in coordss:
                 raise ValueError(f"Data has no coordinate information for dimension {dim}")
 
-            if dim == pdim and pranges is not None:
+            if dim == data.metadata.partition_dim and data.metadata.partition_ranges is not None:
                 # Dask-native partition pruning along the partition dim.
-                all_steps = list(range(len(pranges)))
+                all_steps = list(range(len(data.metadata.partition_ranges)))
                 selected_steps = all_steps[isel]
                 if isinstance(selected_steps, int):
                     selected_steps = [selected_steps]
-                partition_indices = [p for step in selected_steps for p in range(*pranges[step])]
+                partition_indices = [p for step in selected_steps for p in range(*data.metadata.partition_ranges[step])]
                 df = df.partitions[partition_indices]
                 coordss[dim] = coordss[dim][isel] if isinstance(isel, slice) else float(coordss[dim][isel])
                 continue
 
             if isinstance(isel, int):
                 pos = float(coordss[dim][isel])
                 df = df[df[dim] == pos]
+                if len(df) == 0:
+                    import warnings
+
+                    message = f"--idx {dim}={isel} on list data requires exact coordinate match, and returned an empty list. Try --idx {dim}={isel}:{isel + 1} instead."
+                    warnings.warn(message)
                 coordss[dim] = pos
             else:
                 if isel.start not in [None, 0]:

src/lib/data/adaptors/imag.py

@@ -3,12 +3,14 @@
 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.args_registry import const_arg
 
 
 @const_arg(
     dest="adaptors",
-    flags="--image",
+    flags="--imag",
     help="take the imaginary part",
 )
 class Imaginary(BareAdaptor):
@@ -21,5 +23,5 @@ def apply_list_bare(self, series: pd.Series | dd.Series) -> pd.Series | dd.Serie
     def get_name_fragments(self) -> list[str]:
         return ["imag"]
 
-    def get_modified_display_latex(self, metadata) -> str:
-        return f"\\text{{Im}}[{metadata.active_var_info.display}]"
+    def get_modified_display_latex(self, metadata: Metadata) -> Latex:
+        return Latex(f"\\text{{Im}}[{metadata.active_var_info.display}]")

src/lib/data/adaptors/mag.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.args_registry import const_arg
 
 
@@ -21,5 +23,5 @@ def apply_list_bare(self, series: pd.Series | dd.Series) -> pd.Series | dd.Serie
     def get_name_fragments(self) -> list[str]:
         return ["mag"]
 
-    def get_modified_display_latex(self, metadata) -> str:
-        return f"|{metadata.active_var_info.display}|"
+    def get_modified_display_latex(self, metadata: Metadata) -> Latex:
+        return Latex(f"|{metadata.active_var_info.display}|")

src/lib/data/adaptors/math_op.py

@@ -5,6 +5,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
 
@@ -43,8 +45,8 @@ def apply_list_bare(self, series: pd.Series | dd.Series) -> pd.Series | dd.Serie
     def get_name_fragments(self) -> list[str]:
         return [f"{self.name_abbrev}_{self.rhs}"]
 
-    def get_modified_display_latex(self, metadata) -> str:
-        return f"({metadata.active_var_info.display}){self.symbol}{{{float_to_latex_str(self.rhs)}}}"
+    def get_modified_display_latex(self, metadata: Metadata) -> Latex:
+        return Latex(f"({metadata.active_var_info.display}){self.symbol}{{{float_to_latex_str(self.rhs)}}}")
 
 
 op_params = [

src/lib/data/adaptors/real.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.args_registry import const_arg
 
 
@@ -21,5 +23,5 @@ def apply_list_bare(self, series: pd.Series | dd.Series) -> pd.Series | dd.Serie
     def get_name_fragments(self) -> list[str]:
         return ["real"]
 
-    def get_modified_display_latex(self, metadata) -> str:
-        return f"\\text{{Re}}[{metadata.active_var_info.display}]"
+    def get_modified_display_latex(self, metadata: Metadata) -> Latex:
+        return Latex(f"\\text{{Re}}[{metadata.active_var_info.display}]")

src/lib/data/adaptors/reduce.py

@@ -4,6 +4,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
 
@@ -46,10 +48,10 @@ def apply_field_bare(self, da: xr.DataArray) -> xr.DataArray:
             return da
         return REDUCE_FUNCS[self.func_name][0](da, self.dim_names)
 
-    def get_modified_display_latex(self, metadata) -> str:
+    def get_modified_display_latex(self, metadata: Metadata) -> Latex:
         if not self.dim_names:
-            return metadata.active_var_info.display.latex
-        return REDUCE_FUNCS[self.func_name][1](metadata.active_var_info.display.latex, self.dim_names)
+            return metadata.active_var_info.display
+        return Latex(REDUCE_FUNCS[self.func_name][1](metadata.active_var_info.display.latex, self.dim_names))
 
     def get_name_fragments(self) -> list[str]:
         if not self.dim_names: