Made NWBSlicer manufacturer-aware.

src/ezmsg/nwb/slicer.py

@@ -67,6 +67,72 @@ def _extract_timeseries_from_container(container, address: str | None = None) ->
     return timeseries
 
 
+def _electrode_group_manufacturer(child: pynwb.TimeSeries) -> str:
+    """Return the manufacturer of the Device backing this TimeSeries.
+
+    Only meaningful for ElectricalSeries / FeatureExtraction-style containers
+    whose ``electrodes`` DynamicTableRegion links back to an ElectrodeGroup →
+    Device. Returns ``""`` when the link is absent, unreadable, or the
+    Device's ``manufacturer`` attribute is unset/"unknown".
+    """
+    try:
+        electrodes = getattr(child, "electrodes", None)
+        if electrodes is None:
+            return ""
+        # ``electrodes`` is a DynamicTableRegion; .table → ElectrodesTable.
+        table = getattr(electrodes, "table", None)
+        if table is None:
+            return ""
+        # Grab any electrode row's group (they all share the device for an
+        # ElectricalSeries with a single ElectrodeGroup).
+        group_col = table["group"]
+        if len(group_col) == 0:
+            return ""
+        group = group_col[0]
+        device = getattr(group, "device", None)
+        if device is None:
+            return ""
+        manufacturer = getattr(device, "manufacturer", "") or ""
+        if manufacturer.lower() == "unknown":
+            return ""
+        return str(manufacturer)
+    except (AttributeError, KeyError, IndexError, TypeError):
+        return ""
+
+
+def _match_stream_key(child: pynwb.TimeSeries, stream_keys: list[str] | None) -> str | None:
+    """Resolve the user-facing key for a TimeSeries given the *stream_keys* filter.
+
+    Returns:
+    * ``child.name`` if *stream_keys* is None (no filter — keep the literal
+      NWB container name).
+    * ``child.name`` if it appears in *stream_keys* (exact match).
+    * The bare ``<key>`` portion if ``child.name == f"{manufacturer}_{key}"``
+      for some ``<key>`` in *stream_keys* and the linked Device's
+      ``manufacturer`` attribute equals ``<manufacturer>``. This keeps
+      configs that name the bare device (e.g. ``["NPLAY"]``) working when
+      the NWB writer prefixes containers with the manufacturer (e.g.
+      ``"CereLink_NPLAY"``). The returned ``<key>`` is then used as the
+      stream's storage key and template ``.key``, so downstream consumers
+      see the name they asked for.
+    * ``None`` when nothing matches and the stream should be skipped.
+    """
+    if stream_keys is None:
+        return child.name
+    if child.name in stream_keys:
+        return child.name
+    manufacturer = _electrode_group_manufacturer(child)
+    if not manufacturer:
+        return None
+    prefix = f"{manufacturer}_"
+    if not child.name.startswith(prefix):
+        return None
+    stripped = child.name[len(prefix) :]
+    if stripped in stream_keys:
+        return stripped
+    return None
+
+
 class NWBSlicer:
     """Shared NWB file handling: open, discover streams, and slice data.
 
@@ -191,91 +257,100 @@ def _load(self) -> None:
         for address, child in all_timeseries:
             if type(child) is pynwb.misc.Units:
                 ez.logger.warning("Units found in NWB file. Not yet supported.")
-            elif isinstance(child, pynwb.TimeSeries) and (self._stream_keys is None or child.name in self._stream_keys):
-                if child.data.size == 0:
-                    ez.logger.warning(f"Skipping empty TimeSeries: {child.name} {type(child)}")
-                    continue
-
-                has_timestamps = hasattr(child, "timestamps") and child.timestamps is not None
-
-                if has_timestamps:
-                    # Determine nominal rate
-                    if hasattr(child, "rate") and child.rate is not None:
-                        rate = child.rate
-                    elif "rate" in child.timestamps.attrs:
-                        rate = child.timestamps.attrs["rate"]
-                    else:
-                        dts = np.diff(child.timestamps[:])
-                        if np.var(dts) < 1e-3 or np.var(dts) < 0.05 * np.median(dts):
-                            rate = 1 / np.median(dts)
-                        else:
-                            rate = 0.0
-
-                    t0_val = child.timestamps[0]
-                    start_time = min(start_time, self._ts_off + t0_val)
-                    gain = 1 / rate if rate != 0 else 1.0
-                    stop_time = max(stop_time, self._ts_off + child.timestamps[-1] + gain)
-                    stop_time = max(
-                        stop_time,
-                        self._ts_off + t0_val + (child.data.shape[0] + 1) * gain,
-                    )
-                    tvec = child.timestamps
-                else:
-                    rate = child.rate
-                    t0_val = child.starting_time
-                    gain = 1 / rate if rate != 0 else 1.0
-                    start_time = min(start_time, self._ts_off + t0_val)
-                    stop_time = max(
-                        stop_time,
-                        self._ts_off + t0_val + (child.data.shape[0] + 1) * gain,
-                    )
-                    tvec = child.starting_time + np.arange(child.data.shape[0]) / rate
+                continue
+            if not isinstance(child, pynwb.TimeSeries):
+                continue
+            matched_key = _match_stream_key(child, self._stream_keys)
+            if matched_key is None:
+                continue
+            if child.data.size == 0:
+                ez.logger.warning(f"Skipping empty TimeSeries: {child.name} {type(child)}")
+                continue
+
+            has_timestamps = hasattr(child, "timestamps") and child.timestamps is not None
 
-                # Build axes metadata
-                axes: dict[str, typing.Any] = {}
-                if math.isclose(rate, 0.0):
-                    axes["time"] = AxisArray.CoordinateAxis(data=np.array([]), dims=["time"], unit="s")
+            if has_timestamps:
+                # Determine nominal rate
+                if hasattr(child, "rate") and child.rate is not None:
+                    rate = child.rate
+                elif "rate" in child.timestamps.attrs:
+                    rate = child.timestamps.attrs["rate"]
                 else:
-                    axes["time"] = AxisArray.LinearAxis.create_time_axis(fs=rate, offset=self._ts_off)
-                if hasattr(child, "electrodes") and child.electrodes is not None:
-                    # ``child.electrodes`` is a DynamicTableRegion whose
-                    # ``.data`` holds the positional indices into the full
-                    # electrodes table. Subset with iloc so the returned
-                    # channel labels line up 1:1 with the data columns —
-                    # otherwise an ElectricalSeries that references a
-                    # strict subset of the electrodes table produces a
-                    # ch-axis whose length does not match data.shape[1].
-                    region_idx = np.asarray(child.electrodes.data)
-                    full_df = child.electrodes.table.to_dataframe()
-                    el_df = full_df.iloc[region_idx]
-                    if "label" in el_df.columns:
-                        ch_labels = el_df["label"].values.tolist()
+                    dts = np.diff(child.timestamps[:])
+                    if np.var(dts) < 1e-3 or np.var(dts) < 0.05 * np.median(dts):
+                        rate = 1 / np.median(dts)
                     else:
-                        ch_labels = [f"ch_{idx}" for idx in el_df.index.tolist()]
-                    axes["ch"] = AxisArray.CoordinateAxis(data=np.array(ch_labels), dims=["ch"])
-
-                self._streams[child.name] = StreamInfo(
-                    dset=child.data,
-                    template=AxisArray(
-                        data=np.zeros((0,) + child.data.shape[1:], dtype=child.data.dtype),
-                        dims=(["time", "ch"] + [f"dim_{_}" for _ in range(2, child.data.ndim)])
-                        if child.data.ndim > 1
-                        else ["time"],
-                        axes=axes,
-                        key=child.name,
-                    ),
-                    fs=rate,
-                    t0=(
-                        child.starting_time
-                        if (hasattr(child, "starting_time") and child.starting_time is not None)
-                        else child.timestamps[0]
-                    ),
-                    n_samples=child.data.shape[0],
-                    timestamps=tvec if has_timestamps else None,
-                    has_timestamps=has_timestamps,
-                    is_event=False,
-                    table_ref=None,
+                        rate = 0.0
+
+                t0_val = child.timestamps[0]
+                start_time = min(start_time, self._ts_off + t0_val)
+                gain = 1 / rate if rate != 0 else 1.0
+                stop_time = max(stop_time, self._ts_off + child.timestamps[-1] + gain)
+                stop_time = max(
+                    stop_time,
+                    self._ts_off + t0_val + (child.data.shape[0] + 1) * gain,
+                )
+                tvec = child.timestamps
+            else:
+                rate = child.rate
+                t0_val = child.starting_time
+                gain = 1 / rate if rate != 0 else 1.0
+                start_time = min(start_time, self._ts_off + t0_val)
+                stop_time = max(
+                    stop_time,
+                    self._ts_off + t0_val + (child.data.shape[0] + 1) * gain,
                 )
+                tvec = child.starting_time + np.arange(child.data.shape[0]) / rate
+
+            # Build axes metadata
+            axes: dict[str, typing.Any] = {}
+            if math.isclose(rate, 0.0):
+                axes["time"] = AxisArray.CoordinateAxis(data=np.array([]), dims=["time"], unit="s")
+            else:
+                axes["time"] = AxisArray.LinearAxis.create_time_axis(fs=rate, offset=self._ts_off)
+            if hasattr(child, "electrodes") and child.electrodes is not None:
+                # ``child.electrodes`` is a DynamicTableRegion whose
+                # ``.data`` holds the positional indices into the full
+                # electrodes table. Subset with iloc so the returned
+                # channel labels line up 1:1 with the data columns —
+                # otherwise an ElectricalSeries that references a
+                # strict subset of the electrodes table produces a
+                # ch-axis whose length does not match data.shape[1].
+                region_idx = np.asarray(child.electrodes.data)
+                full_df = child.electrodes.table.to_dataframe()
+                el_df = full_df.iloc[region_idx]
+                if "label" in el_df.columns:
+                    ch_labels = el_df["label"].values.tolist()
+                else:
+                    ch_labels = [f"ch_{idx}" for idx in el_df.index.tolist()]
+                axes["ch"] = AxisArray.CoordinateAxis(data=np.array(ch_labels), dims=["ch"])
+
+            # ``matched_key`` is the user-facing key — equal to ``child.name``
+            # when there's no stream_keys filter or an exact match, or the
+            # bare ``<key>`` portion when the user requested an unprefixed
+            # device name and the container is ``"<manufacturer>_<key>"``.
+            self._streams[matched_key] = StreamInfo(
+                dset=child.data,
+                template=AxisArray(
+                    data=np.zeros((0,) + child.data.shape[1:], dtype=child.data.dtype),
+                    dims=(["time", "ch"] + [f"dim_{_}" for _ in range(2, child.data.ndim)])
+                    if child.data.ndim > 1
+                    else ["time"],
+                    axes=axes,
+                    key=matched_key,
+                ),
+                fs=rate,
+                t0=(
+                    child.starting_time
+                    if (hasattr(child, "starting_time") and child.starting_time is not None)
+                    else child.timestamps[0]
+                ),
+                n_samples=child.data.shape[0],
+                timestamps=tvec if has_timestamps else None,
+                has_timestamps=has_timestamps,
+                is_event=False,
+                table_ref=None,
+            )
 
         self._start_time = start_time
         self._stop_time = stop_time

tests/test_slicer.py

@@ -42,6 +42,142 @@ def test_stream_discovery_filter(test_nwb_path):
     s.close()
 
 
+def _build_manufacturer_prefixed_nwb(path):
+    """Build a tiny NWB with an ElectricalSeries named ``CereLink_NPLAY``
+    backed by a Device whose ``manufacturer`` attribute is ``"CereLink"``.
+
+    Mirrors the layout Orion writes: storage container path prefixed by the
+    manufacturer, with the manufacturer also stamped on the Device so
+    downstream readers can reconstruct the bare device name.
+    """
+    import datetime
+
+    from pynwb import NWBHDF5IO, NWBFile
+    from pynwb.ecephys import ElectricalSeries
+    from pynwb.file import Subject
+
+    rng = np.random.default_rng(0)
+    nwbfile = NWBFile(
+        session_description="prefix-test",
+        identifier="manufacturer-prefix-test",
+        session_start_time=datetime.datetime.now(datetime.timezone.utc),
+        subject=Subject(subject_id="sub1", age="P30Y", sex="U"),
+    )
+    device = nwbfile.create_device(name="CereLink_NPLAY", manufacturer="CereLink")
+    group = nwbfile.create_electrode_group(
+        name="CereLink_NPLAY", description="prefix-test group", location="cortex", device=device
+    )
+    nwbfile.add_electrode_column(name="label", description="Electrode label")
+    for i in range(4):
+        nwbfile.add_electrode(x=float(i), y=0.0, z=0.0, location="cortex", group=group, label=f"e{i}")
+    region = nwbfile.create_electrode_table_region(region=list(range(4)), description="all")
+
+    bb_n = 100
+    es = ElectricalSeries(
+        name="CereLink_NPLAY",
+        data=rng.standard_normal((bb_n, 4)).astype(np.float32),
+        starting_time=0.0,
+        rate=1000.0,
+        electrodes=region,
+    )
+    nwbfile.add_acquisition(es)
+
+    with NWBHDF5IO(str(path), "w") as io:
+        io.write(nwbfile)
+
+
+def test_stream_keys_match_via_manufacturer_prefix(tmp_path):
+    """stream_keys=['NPLAY'] matches a container 'CereLink_NPLAY' whose
+    Device.manufacturer == 'CereLink'. The stream is exposed under the
+    user-requested bare key, and messages carry key='NPLAY' so downstream
+    fitters keyed by the request find their data."""
+    nwb_path = tmp_path / "prefix.nwb"
+    _build_manufacturer_prefixed_nwb(nwb_path)
+
+    s = NWBSlicer(
+        filepath=str(nwb_path),
+        reference_clock=ReferenceClockType.UNKNOWN,
+        stream_keys=["NPLAY"],
+    )
+    try:
+        assert s.stream_names == ["NPLAY"], f"expected the matched key to be the bare request, got {s.stream_names}"
+        info = s.get_stream_info("NPLAY")
+        assert info.template.key == "NPLAY"
+        msg = s.read_by_index("NPLAY", 0, 10)
+        assert msg.key == "NPLAY"
+        assert msg.data.shape == (10, 4)
+    finally:
+        s.close()
+
+
+def test_stream_keys_exact_match_wins_over_manufacturer(tmp_path):
+    """Exact-match takes precedence: stream_keys=['CereLink_NPLAY'] yields
+    the literal container name."""
+    nwb_path = tmp_path / "prefix.nwb"
+    _build_manufacturer_prefixed_nwb(nwb_path)
+
+    s = NWBSlicer(
+        filepath=str(nwb_path),
+        reference_clock=ReferenceClockType.UNKNOWN,
+        stream_keys=["CereLink_NPLAY"],
+    )
+    try:
+        assert s.stream_names == ["CereLink_NPLAY"]
+        msg = s.read_by_index("CereLink_NPLAY", 0, 5)
+        assert msg.key == "CereLink_NPLAY"
+    finally:
+        s.close()
+
+
+def test_stream_keys_manufacturer_unknown_does_not_match(tmp_path):
+    """A bare request like 'NPLAY' must NOT match 'CereLink_NPLAY' when the
+    Device has no real manufacturer (unset or 'unknown'). Otherwise legacy
+    files without manufacturer metadata could accidentally match unrelated
+    streams that happen to share a suffix."""
+    import datetime
+
+    from pynwb import NWBHDF5IO, NWBFile
+    from pynwb.ecephys import ElectricalSeries
+    from pynwb.file import Subject
+
+    rng = np.random.default_rng(0)
+    path = tmp_path / "unknown_mfg.nwb"
+    nwbfile = NWBFile(
+        session_description="x",
+        identifier="x",
+        session_start_time=datetime.datetime.now(datetime.timezone.utc),
+        subject=Subject(subject_id="s", age="P30Y", sex="U"),
+    )
+    # No manufacturer (or "unknown") on the device → suffix match must be
+    # rejected.
+    device = nwbfile.create_device(name="CereLink_NPLAY", manufacturer="unknown")
+    group = nwbfile.create_electrode_group(name="CereLink_NPLAY", description="x", location="cortex", device=device)
+    nwbfile.add_electrode_column(name="label", description="Electrode label")
+    for i in range(2):
+        nwbfile.add_electrode(x=0.0, y=0.0, z=0.0, location="cortex", group=group, label=f"e{i}")
+    region = nwbfile.create_electrode_table_region(region=list(range(2)), description="all")
+    es = ElectricalSeries(
+        name="CereLink_NPLAY",
+        data=rng.standard_normal((50, 2)).astype(np.float32),
+        starting_time=0.0,
+        rate=500.0,
+        electrodes=region,
+    )
+    nwbfile.add_acquisition(es)
+    with NWBHDF5IO(str(path), "w") as io:
+        io.write(nwbfile)
+
+    s = NWBSlicer(
+        filepath=str(path),
+        reference_clock=ReferenceClockType.UNKNOWN,
+        stream_keys=["NPLAY"],
+    )
+    try:
+        assert s.stream_names == [], f"unknown manufacturer should not enable suffix match, got {s.stream_names}"
+    finally:
+        s.close()
+
+
 def test_stream_info_continuous(slicer):
     """Continuous timestamped stream metadata is correct."""
     info = slicer.get_stream_info("Broadband")