Slicer: select channels by regex against CoordinateAxis labels

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])

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])