diff --git a/src/lib/cli.py b/src/lib/cli.py index 98f49e2..a7d4f47 100644 --- a/src/lib/cli.py +++ b/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 diff --git a/src/lib/data/adaptor.py b/src/lib/data/adaptor.py index 486dcaf..07ddcbe 100644 --- a/src/lib/data/adaptor.py +++ b/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) diff --git a/src/lib/data/adaptors/bin.py b/src/lib/data/adaptors/bin.py index 517fb22..30885b5 100644 --- a/src/lib/data/adaptors/bin.py +++ b/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)) diff --git a/src/lib/data/adaptors/fourier.py b/src/lib/data/adaptors/fourier.py index 27525e9..128d675 100644 --- a/src/lib/data/adaptors/fourier.py +++ b/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)}"] diff --git a/src/lib/data/adaptors/idx.py b/src/lib/data/adaptors/idx.py index b410f4d..ce70778 100644 --- a/src/lib/data/adaptors/idx.py +++ b/src/lib/data/adaptors/idx.py @@ -16,19 +16,17 @@ 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 @@ -36,6 +34,11 @@ def apply_list(self, data: List) -> List: 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]: diff --git a/src/lib/data/adaptors/imag.py b/src/lib/data/adaptors/imag.py index 8d7c4f8..9123a61 100644 --- a/src/lib/data/adaptors/imag.py +++ b/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}]") diff --git a/src/lib/data/adaptors/mag.py b/src/lib/data/adaptors/mag.py index e0c5aa7..f31d5ed 100644 --- a/src/lib/data/adaptors/mag.py +++ b/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}|") diff --git a/src/lib/data/adaptors/math_op.py b/src/lib/data/adaptors/math_op.py index 029aafd..8b5f2f8 100644 --- a/src/lib/data/adaptors/math_op.py +++ b/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 = [ diff --git a/src/lib/data/adaptors/real.py b/src/lib/data/adaptors/real.py index 263c86c..2b23eeb 100644 --- a/src/lib/data/adaptors/real.py +++ b/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}]") diff --git a/src/lib/data/adaptors/reduce.py b/src/lib/data/adaptors/reduce.py index 87ad9bb..3d40243 100644 --- a/src/lib/data/adaptors/reduce.py +++ b/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: