diff --git a/src/ezmsg/sigproc/slicer.py b/src/ezmsg/sigproc/slicer.py index cdd4ca8b..9b4beb44 100644 --- a/src/ezmsg/sigproc/slicer.py +++ b/src/ezmsg/sigproc/slicer.py @@ -1,5 +1,7 @@ """Select a subset of data along a named axis using slice notation.""" +import re + import ezmsg.core as ez import numpy as np import numpy.typing as npt @@ -34,13 +36,17 @@ def parse_slice( - "5" (or any integer) -> (5,). Take only that item. applying this to a ndarray or AxisArray will drop the dimension. - A comma-separated list of the above -> a tuple of slices | ints - - A comma-separated list of values and axinfo is provided and is a CoordinateAxis -> a tuple of ints + - A comma-separated list of values and axinfo is provided and is a CoordinateAxis -> a tuple of ints. + Each value is first compared against the axis labels for an exact match; failing + that, it is treated as a regular expression and full-matched against the labels + (e.g. "C[34]" or "Ch.*"). Note: tokens containing ":" are parsed as slices, so + regexes may not contain ":". Args: s: The string representation of the slice. axinfo: (Optional) If provided, and of type CoordinateAxis, and `s` is a comma-separated list of values, then the values - in s will be checked against the values in axinfo.data. + in s will be matched (exactly, then as regex) against the values in axinfo.data. Returns: A tuple of slice objects and/or ints. @@ -52,7 +58,19 @@ def parse_slice( if len(parts) == 1: if axinfo is not None and hasattr(axinfo, "data") and parts[0] in axinfo.data: return tuple(np.where(axinfo.data == parts[0])[0]) - return (int(parts[0]),) + try: + return (int(parts[0]),) + except ValueError: + if axinfo is not None and hasattr(axinfo, "data"): + pattern = re.compile(parts[0]) + hits = tuple(ix for ix, label in enumerate(axinfo.data) if pattern.fullmatch(str(label))) + if hits: + return hits + raise ValueError( + f"Selection {parts[0]!r} matched no labels on the target axis " + f"(neither exactly nor as a regex)." + ) from None + raise return (slice(*(int(part.strip()) if part else None for part in parts)),) suplist = [parse_slice(_, axinfo=axinfo) for _ in s.split(",")] return tuple([item for sublist in suplist for item in sublist]) diff --git a/tests/unit/test_slicer.py b/tests/unit/test_slicer.py index bc546fa1..56d3b53b 100644 --- a/tests/unit/test_slicer.py +++ b/tests/unit/test_slicer.py @@ -149,3 +149,41 @@ def test_slicer_empty_first(): result = proc(empty) check_empty_result(result) check_state_not_corrupted(proc, normal) + + +def test_parse_slice_regex(): + ax = AxisArray.CoordinateAxis(data=np.array(["Fp1", "Fp2", "C3", "C4", "Cz", "O1", "O2"]), dims=["ch"]) + # Exact match still wins and returns that single index. + assert parse_slice("C3", axinfo=ax) == (2,) + # Regex full-match: all central channels. + assert parse_slice("C[34z]", axinfo=ax) == (2, 3, 4) + # Prefix pattern. + assert parse_slice("Fp.*", axinfo=ax) == (0, 1) + # Comma-separated mix of exact labels and patterns concatenates in order. + assert parse_slice("O.*, C3", axinfo=ax) == (5, 6, 2) + # fullmatch semantics: a bare prefix is not a match. + with pytest.raises(ValueError, match="matched no labels"): + parse_slice("Fp", axinfo=ax) + # Without axinfo, non-numeric tokens are still an error (int parse). + with pytest.raises(ValueError): + parse_slice("C[34z]") + + +def test_slicer_regex_selection(): + n_times = 20 + labels = np.array(["Fp1", "Fp2", "C3", "C4", "Cz", "O1", "O2"]) + n_chans = len(labels) + in_dat = np.arange(n_times * n_chans, dtype=float).reshape(n_times, n_chans) + msg_in = AxisArray( + in_dat, + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0, offset=0.0), + "ch": AxisArray.CoordinateAxis(data=labels, dims=["ch"]), + }, + key="test_slicer_regex_selection", + ) + xformer = SlicerTransformer(SlicerSettings(selection="C.*", axis="ch")) + msg_out = xformer(msg_in) + assert np.array_equal(msg_out.data, in_dat[:, 2:5]) + assert np.array_equal(msg_out.axes["ch"].data, labels[2:5])