diff --git a/src/lib/data/adaptors/versus.py b/src/lib/data/adaptors/versus.py index 3c127a2..819f2f6 100644 --- a/src/lib/data/adaptors/versus.py +++ b/src/lib/data/adaptors/versus.py @@ -1,21 +1,52 @@ +from typing import Literal + from lib.data.adaptor import MetadataAdaptor from lib.data.adaptors.fourier import Fourier from lib.data.adaptors.reduce import Reduce -from lib.data.data_with_attrs import Field, List +from lib.data.data_with_attrs import DataWithAttrs, Field, List from lib.parsing import parse_util from lib.parsing.args_registry import arg_parser class Versus(MetadataAdaptor): - def __init__(self, spatial_dims: list[str], time_dim: str | None, color_dim: str | None): + def __init__( + self, + spatial_dims: list[str], + *, + time_dim_rule: str | None | Literal["guess"], + color_dim: str | None, + ): self.spatial_dims = spatial_dims - self.time_dim = time_dim + self.time_dim_rule = time_dim_rule self.color_dim = color_dim - self.all_dims = spatial_dims + ([time_dim] if time_dim else []) + ([color_dim] if color_dim else []) + + def _get_retained_dim_keys(self, data: DataWithAttrs) -> list[str]: + retained_dims = self.spatial_dims.copy() + + if time_dim := self._get_time_dim(data): + retained_dims.append(time_dim) + + if self.color_dim: + retained_dims.append(self.color_dim) + + if isinstance(data, List) and data.metadata.active_key: + retained_dims.append(data.metadata.active_key) + + return retained_dims + + def _get_time_dim(self, data: DataWithAttrs) -> str | None: + if self.time_dim_rule != "guess": + return self.time_dim_rule + + if "t" in data.dims and "t" not in self.spatial_dims and "t" != self.color_dim: + return "t" + + return None def apply_field(self, data: Field) -> Field: # 1. apply implicit coordinate transforms, as necessary - for dim_name in self.all_dims: + retained_dims = self._get_retained_dim_keys(data) + for dim_name in retained_dims: # 1a. already have the coordinate; do nothing if dim_name in data.dims: continue @@ -35,13 +66,13 @@ def apply_field(self, data: Field) -> Field: # TODO # 2. reduce remaining dimensions via arithmetic mean - reduce_dims = [dim for dim in data.dims if dim not in self.all_dims] + reduce_dims = [dim for dim in data.dims if dim not in retained_dims] reduce = Reduce(reduce_dims, "mean") data = reduce.apply(data) return data.assign_metadata( spatial_dims=self.spatial_dims.copy(), - time_dim=self.time_dim, + time_dim=self._get_time_dim(data), color_dim=self.color_dim, ) @@ -50,7 +81,7 @@ def apply_list(self, data: List) -> List: # TODO # 2. drop unused vars - keep_vars = self.all_dims + ([data.metadata.active_key] if data.metadata.active_key else []) + keep_vars = self._get_retained_dim_keys(data) drop_vars = [active_key for active_key in data.dims if active_key not in keep_vars] data = data.assign_data(data.data.drop(columns=drop_vars)) @@ -60,48 +91,44 @@ def apply_list(self, data: List) -> List: return data.assign_metadata( spatial_dims=spatial_dims, - time_dim=self.time_dim, + time_dim=self._get_time_dim(data), color_dim=self.color_dim, ) def get_name_fragments(self) -> list[str]: dims = ",".join(self.spatial_dims) - if self.time_dim: - dims += f";time={self.time_dim}" + if self.time_dim_rule != "guess": + dims += f";{_TIME_PREFIX}{self.time_dim_rule or ''}" if self.color_dim: - dims += f";color={self.color_dim}" - return [f"vs_{dims}"] + dims += f";{_COLOR_PREFIX}{self.color_dim}" + return [f"v_{dims}"] _TIME_PREFIX = "time=" _COLOR_PREFIX = "color=" -_VERSUS_FORMAT = f"dim_name | {_TIME_PREFIX}[dim_name] | {_COLOR_PREFIX}dim_name" +_VERSUS_FORMAT = f"dim_key | {_TIME_PREFIX}[dim_key] | {_COLOR_PREFIX}dim_key" @arg_parser( dest="adaptors", flags=["--versus", "-v"], metavar=_VERSUS_FORMAT, - help=f"Specifies the independent axes to plot against (automatically performs necessary Fourier and coordinate transforms, and reduces other dimensions via arithmetic mean). The optional `{_TIME_PREFIX}[dim_name]` specifies an additional dimension to use as the time axis. If unspecified, defaults to `{_TIME_PREFIX}t` unless t is used as a spatial axis. Alternatively, disable time by passing `{_TIME_PREFIX}` (with no dim_name).", + help=f"Specifies the independent axes of the plot. Remaining dimensions are reduced via arithmetic mean. Time has a special behavior: if {_TIME_PREFIX}[dim_key] is omitted, it is set to 't' if 't' is present in the data and isn't being used as a different axis. Disable this guessing by passing {_TIME_PREFIX} (with no dim_key).", nargs="+", ) def parse_versus(args: list[str]) -> Versus: spatial_dims = [] - time_dim = "t" - time_dim_is_default = True + time_dim_rule = "guess" color_dim = None for arg in args: if arg.startswith(_TIME_PREFIX): - time_dim = arg.removeprefix(_TIME_PREFIX) or None - parse_util.check_optional_identifier(time_dim, "time dim_name") - time_dim_is_default = False + time_dim_rule = arg.removeprefix(_TIME_PREFIX) or None + parse_util.check_optional_identifier(time_dim_rule, "time dim_key") elif arg.startswith(_COLOR_PREFIX): color_dim = arg.removeprefix(_COLOR_PREFIX) - parse_util.check_identifier(color_dim, "color dim_name") + parse_util.check_identifier(color_dim, "color dim_key") else: - parse_util.check_identifier(arg, "dim_name") + parse_util.check_identifier(arg, "dim_key") spatial_dims.append(arg) - if time_dim_is_default and arg == time_dim: - time_dim = None - return Versus(spatial_dims, time_dim, color_dim) + return Versus(spatial_dims, time_dim_rule=time_dim_rule, color_dim=color_dim) diff --git a/src/lib/data/compile.py b/src/lib/data/compile.py index 97a8df6..a428149 100644 --- a/src/lib/data/compile.py +++ b/src/lib/data/compile.py @@ -9,7 +9,7 @@ def compile_source(loader: DataSource, adaptors: list[Adaptor]) -> DataSource: if isinstance(adaptor, Versus): break else: - adaptors.append(Versus(["y", "z"], "t", color_dim=None)) + adaptors.append(Versus(["y", "z"], time_dim_rule="guess", color_dim=None)) pipeline = Pipeline(*adaptors) source = DataSourceWithPipeline(loader, pipeline) diff --git a/tests/test_save_filename.py b/tests/test_save_filename.py index 79b0b37..4560f79 100644 --- a/tests/test_save_filename.py +++ b/tests/test_save_filename.py @@ -13,9 +13,10 @@ def _stem(args_list: list[str]) -> str: @pytest.mark.parametrize( "args_list, expected_stem", [ - (["pfd", "hx_fc"], "pfd-hx_fc-vs_y,z;time=t"), - (["pfd", "hx_fc", "--nan0"], "pfd-hx_fc-nan0-vs_y,z;time=t"), - (["pfd", "hx_fc", "--scale", "log"], "pfd-hx_fc-vs_y,z;time=t-scale_log"), + (["pfd", "hx_fc"], "pfd-hx_fc-v_y,z"), + (["pfd", "hx_fc", "--nan0"], "pfd-hx_fc-nan0-v_y,z"), + (["pfd", "hx_fc", "--scale", "log"], "pfd-hx_fc-v_y,z-scale_log"), + (["pfd", "hx_fc", "-v", "y", "z", "time="], "pfd-hx_fc-v_y,z;time="), ], ) def test_save_file_stem(args_list, expected_stem):