diff --git a/pyproject.toml b/pyproject.toml index 90430e1..4272a3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ dependencies = [ "ezmsg-baseproc>=1.9.0", "ezmsg-event", "ezmsg>=3.6.1", - "pycbsdk>=9.9.1", + "pycbsdk>=9.10.0", "scipy", ] diff --git a/src/ezmsg/blackrock/__init__.py b/src/ezmsg/blackrock/__init__.py index c0855fa..23e7ff1 100644 --- a/src/ezmsg/blackrock/__init__.py +++ b/src/ezmsg/blackrock/__init__.py @@ -9,6 +9,7 @@ CereLinkSpikeProducer, CereLinkSpikeSettings, CereLinkSpikeSource, + ChannelSelection, DeviceConfig, DeviceStatus, SliceConfig, @@ -43,6 +44,7 @@ "CerePlexImpedance", "CerePlexImpedanceProcessor", "CerePlexImpedanceSettings", + "ChannelSelection", "CHANNEL_DTYPE", "ChannelMapProcessor", "ChannelMapSettings", diff --git a/src/ezmsg/blackrock/cerelink.py b/src/ezmsg/blackrock/cerelink.py index 8ee592a..ce64454 100644 --- a/src/ezmsg/blackrock/cerelink.py +++ b/src/ezmsg/blackrock/cerelink.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import enum import logging import threading import time @@ -15,7 +16,7 @@ from ezmsg.baseproc.stateful import BaseStatefulProducer from ezmsg.baseproc.units import BaseProducerUnit from ezmsg.util.messages.axisarray import AxisArray, replace -from pycbsdk import ChannelType, DeviceType, SampleRate, Session +from pycbsdk import ChanInfoField, ChannelType, DeviceType, SampleRate, Session from .channel_map import CHANNEL_DTYPE, ChannelMapSettings @@ -55,6 +56,22 @@ class CcfConfig: path: str +class ChannelSelection(enum.Enum): + """Sentinel selections for :attr:`SliceConfig.channels` (the cases an + explicit channel-ID list can't express).""" + + ALL = "all" + """Every channel matching ``channel_type``; the others are disabled.""" + + ENABLED = "enabled" + """Only the channels the device **already** has enabled for this stream; the + enabled set is left unchanged. What counts as "enabled" is stream-specific: + continuous sample-group membership for the signal source (channels are + retuned to its rate/coupling but never enabled/disabled), and spike-extraction + state for the spike source (extraction is left exactly as-is, and the source + subscribes to whatever already has it on).""" + + @dataclass class SliceConfig: """Programmatic per-slice device configuration owned by this source. @@ -68,8 +85,21 @@ class SliceConfig: responsibility. """ - channels: list[int] | None = None - """1-based channel IDs to configure. ``None`` = all matching ``channel_type``.""" + channels: list[int] | ChannelSelection = ChannelSelection.ALL + """Which channels this slice targets — one field, three intents: + + - ``list[int]`` — enable exactly these 1-based channel IDs; the other + channels of ``channel_type`` are disabled. A provided list is always + respected. + - :attr:`ChannelSelection.ALL` (default) — enable every channel matching + ``channel_type``; others disabled. + - :attr:`ChannelSelection.ENABLED` — leave the device's enabled set as-is and + only consume it. Signal source: retune the already-streaming channels + (``disable_others=False``, so an unused front-end bank stays off). Spike + source: leave spike extraction untouched and subscribe to whatever already + has it on (``enable_spiking`` is ignored in this mode). An empty enabled + set yields nothing and warns. + """ channel_type: ChannelType = ChannelType.FRONTEND @@ -79,8 +109,18 @@ class SliceConfig: """ enable_spiking: bool = False - """Enable spike extraction on this slice (FRONTEND only). Honored by - :class:`CereLinkSpikeSource`; ignored by signal sources.""" + """Enable spike extraction on the selected channels (FRONTEND only). Honored + by :class:`CereLinkSpikeSource`; ignored by signal sources, and ignored when + ``channels`` is :attr:`ChannelSelection.ENABLED` (which leaves extraction + exactly as the device has it).""" + + def __post_init__(self): + if not isinstance(self.channels, (list, ChannelSelection)): + raise TypeError( + "SliceConfig.channels must be a list of 1-based channel IDs or a " + f"ChannelSelection member, not {type(self.channels).__name__}. " + "Did you mean ChannelSelection.ALL or ChannelSelection.ENABLED?" + ) DeviceConfig = CcfConfig | SliceConfig | None @@ -170,7 +210,7 @@ class _CereLinkSharedState: """State fields common to signal and spike producers.""" session: Session | None = None - ch_positions: dict | None = None # ch_id -> (col, row, bank, elec) + ch_positions: dict | None = None # ch_id -> (x, y, size, headstage, bank_num, term) @processor_state @@ -305,15 +345,59 @@ def _apply_slice_configure(self, cfg: SliceConfig) -> None: """Subclass hook: stream-specific slice config (sample-group OR spike-extract).""" raise NotImplementedError + def _enabled_channels(self, channel_type: ChannelType) -> list[int]: + """1-based IDs of *channel_type* channels the device already has enabled + for this stream — the set :attr:`ChannelSelection.ENABLED` resolves to. + + Base implementation (the signal-stream meaning): channels currently in a + continuous sample group, i.e. already streaming. The per-channel + ``SMPGROUP`` field reads 0 for raw-group channels, so we union the + membership of every continuous group (``SampleRate`` 1..6) rather than + trust that field, then intersect with the type filter. The spike producer + overrides this with its own meaning (spike-extraction state). Call after a + :meth:`sync` so the device state is fresh. + """ + sess = self.state.session + enabled: set[int] = set() + for sr in SampleRate: + if sr == SampleRate.NONE: + continue + enabled.update(sess.get_group_channels(int(sr))) + matching = set(sess.get_matching_channel_ids(channel_type)) + return sorted(enabled & matching) + + def _resolve_channels(self, cfg: SliceConfig) -> list[int]: + """Concrete 1-based channel IDs this slice targets, resolving the + :attr:`SliceConfig.channels` selection against the device's current + state (see :class:`ChannelSelection`). ``ENABLED`` :meth:`sync`\\ s first + so the group membership is fresh.""" + sess = self.state.session + if cfg.channels is ChannelSelection.ALL: + return list(sess.get_matching_channel_ids(cfg.channel_type)) + if cfg.channels is ChannelSelection.ENABLED: + sess.sync() + return self._enabled_channels(cfg.channel_type) + return list(cfg.channels) + def _apply_channel_maps(self) -> None: for cmp_cfg in self.settings.cmp_configs: if cmp_cfg.filepath: self.state.session.load_channel_map(cmp_cfg.filepath, cmp_cfg.start_chan, cmp_cfg.hs_id) def _cache_channel_metadata(self) -> None: - all_ids = self.state.session.get_matching_channel_ids(ChannelType.FRONTEND) - all_pos = self.state.session.get_channels_positions(ChannelType.FRONTEND) - self.state.ch_positions = dict(zip(all_ids, all_pos)) + # ``position[]`` carries ``(x, y, size, headstage_id)`` since CereLink + # #184 (CerebusOSS/CereLink#184); ``bank``/``term`` come from chaninfo + # (``ChanInfoField``) rather than position, where they were previously + # encoded. ``x``/``y``/``size`` are in micrometers. + sess = self.state.session + all_ids = sess.get_matching_channel_ids(ChannelType.FRONTEND) + all_pos = sess.get_channels_positions(ChannelType.FRONTEND) + banks = sess.get_channels_field(ChannelType.FRONTEND, ChanInfoField.BANK) + terms = sess.get_channels_field(ChannelType.FRONTEND, ChanInfoField.TERM) + self.state.ch_positions = { + cid: (pos[0], pos[1], pos[2], pos[3], bank, term) + for cid, pos, bank, term in zip(all_ids, all_pos, banks, terms) + } def _build_ch_info(self, channels: list[int]) -> np.ndarray: n_ch = len(channels) @@ -321,11 +405,13 @@ def _build_ch_info(self, channels: list[int]) -> np.ndarray: for i, ch_id in enumerate(channels): label = self.state.session.get_channel_label(ch_id) ch_info[i]["label"] = label or f"ch{ch_id}" - pos = self.state.ch_positions.get(ch_id, (0, 0, 0, 0)) - ch_info[i]["x"] = pos[0] - ch_info[i]["y"] = pos[1] - ch_info[i]["bank"] = chr(ord("A") + pos[2] - 1) if pos[2] > 0 else "" - ch_info[i]["elec"] = pos[3] + x, y, size, headstage, bank_num, term = self.state.ch_positions.get(ch_id, (0, 0, 0, 0, 0, 0)) + ch_info[i]["x"] = x + ch_info[i]["y"] = y + ch_info[i]["size"] = size + ch_info[i]["bank"] = chr(ord("A") + bank_num - 1) if bank_num > 0 else "" + ch_info[i]["elec"] = term + ch_info[i]["headstage"] = headstage return ch_info def _device_name(self) -> str: @@ -378,10 +464,28 @@ class CereLinkSignalProducer(_CereLinkBaseProducer[CereLinkSignalSettings, CereL """Streams one continuous sample-group as :class:`AxisArray`.""" def _apply_slice_configure(self, cfg: SliceConfig) -> None: + sess = self.state.session + if sess is None: + return rate = self.settings.subscribe_rate - if self.state.session is not None: - self.state.session.set_sample_group(cfg.channels, cfg.channel_type, rate, disable_others=True) - self.state.session.set_ac_input_coupling(cfg.channels, cfg.channel_type, cfg.ac_input_coupling) + channels = self._resolve_channels(cfg) + if cfg.channels is ChannelSelection.ENABLED: + # Retune only the channels the device already streams; never enable + # or disable one (disable_others=False), so an unused bank stays off. + if not channels: + logger.warning( + "CereLink: channels=ChannelSelection.ENABLED but no %s channels are " + "currently enabled on the device; nothing will stream. Enable channels " + "upstream, or use channels=ChannelSelection.ALL to configure all of them.", + cfg.channel_type.name, + ) + disable_others = False + else: + # Explicit list or ALL: the resolved set is authoritative; everything + # else of this channel_type is disabled. + disable_others = True + sess.set_sample_group(channels, cfg.channel_type, rate, disable_others=disable_others) + sess.set_ac_input_coupling(channels, cfg.channel_type, cfg.ac_input_coupling) def _setup_subscription(self, loop: asyncio.AbstractEventLoop) -> None: rate = self.settings.subscribe_rate @@ -558,6 +662,7 @@ async def device_status(self) -> typing.AsyncGenerator: _SPIKE_FS = 30000 # device spike clock — fixed by the protocol _NS_PER_SECOND = 1_000_000_000 _UNIT_LABELS = np.array(["unsorted", "1", "2", "3", "4", "5", "noise"], dtype="U8") +_SPKOPTS_EXTRACT = 1 # cbAINPSPK_EXTRACT bit in SPKOPTS — spike extraction enabled @processor_state @@ -602,17 +707,28 @@ def __init__(self, *args, **kwargs) -> None: self._buffer_lock = threading.Lock() def _apply_slice_configure(self, cfg: SliceConfig) -> None: + if cfg.channels is ChannelSelection.ENABLED: + # Leave spike extraction exactly as the device has it; _setup_subscription + # subscribes to whatever already has it on. (enable_spiking is ignored.) + return if cfg.enable_spiking: - self.state.session.set_spike_extraction(cfg.channels, cfg.channel_type, True) + channels = self._resolve_channels(cfg) + self.state.session.set_spike_extraction(channels, cfg.channel_type, True) + + def _enabled_channels(self, channel_type: ChannelType) -> list[int]: + """Spike-stream meaning of "enabled" (see base): channels whose spike + extraction is currently on (the ``cbAINPSPK_EXTRACT`` bit in SPKOPTS), + sorted ascending. Call after a :meth:`sync` so SPKOPTS is fresh.""" + sess = self.state.session + ids = sess.get_matching_channel_ids(channel_type) + opts = sess.get_channels_field(channel_type, ChanInfoField.SPKOPTS) + return sorted(cid for cid, o in zip(ids, opts) if o & _SPKOPTS_EXTRACT) def _setup_subscription(self, loop: asyncio.AbstractEventLoop) -> None: cfg = self.settings.configure if isinstance(cfg, SliceConfig): channel_type = cfg.channel_type - if cfg.channels is not None: - channels = list(cfg.channels) - else: - channels = self.state.session.get_matching_channel_ids(channel_type) + channels = self._resolve_channels(cfg) else: # CcfConfig or None: subscribe to all FRONTEND. The device's CCF # (or whatever's already configured) decides which of these diff --git a/src/ezmsg/blackrock/channel_map.py b/src/ezmsg/blackrock/channel_map.py index f2ca026..6a12f54 100644 --- a/src/ezmsg/blackrock/channel_map.py +++ b/src/ezmsg/blackrock/channel_map.py @@ -1,7 +1,9 @@ """Attach Blackrock ``.cmp`` channel-map metadata to an ``AxisArray``'s ``ch`` axis. The output ``ch`` axis is a structured ``CoordinateAxis`` with fields -``x``, ``y``, ``label``, ``bank``, ``elec`` for every input channel. +``x``, ``y``, ``size``, ``label``, ``bank``, ``elec``, ``headstage`` for every +input channel. ``x``/``y``/``size`` are in micrometers; ``headstage`` is the +1-based headstage id (``0`` = none/auto). :class:`ChannelMapUnit` takes the *complete* set of per-headstage overlays in one settings object (:class:`ChannelMapUnitSettings`, a tuple of @@ -12,15 +14,30 @@ Each reset proceeds in three phases: -1. **Base layer** — labels are pulled from the incoming ``ch`` axis. +1. **Base layer** — labels are pulled from the incoming ``ch`` axis. When the + incoming axis already carries structured geometry (e.g. a CereLink source + that read it from device chaninfo), its ``x``/``y``/``size``/``bank``/ + ``elec``/``headstage`` are copied through verbatim, so a map already present + upstream needs no ``.cmp`` file at all. A channel counts as positioned (and + so is skipped by the auto-grid) when it has a non-origin coordinate, or it + is the *first* channel sitting at the ``(0, 0)`` origin — a lone origin + electrode is a legitimate corner, but the device parks every *unmapped* + channel at the origin, so origin pile-ups beyond the first fall through to + the auto-grid. A companion ``src_mask`` records the positioned indices. 2. **CMP overlays** — for each :class:`ChannelMapSettings` in ``cmp_configs``, - entries from :func:`pycbsdk.cmp.parse_cmp` are written at indices - ``chan_id - 1`` (with ``chan_id`` offset by ``start_chan``). A companion - ``cmp_mask`` records which indices were set so the auto-grid pass can avoid - them. -3. **Auto-grid fill** — positions/bank/elec for indices NOT covered by any - CMP, laid out below and to the right of the CMP geometry so they don't - collide with CMP positions. + entries from :func:`pycbsdk.cmp.parse_cmp` are written at their channel + index, overriding any source geometry there. ``parse_cmp`` + (CerebusOSS/CereLink#184) returns entries keyed by device ``(bank, term)`` + with flat ``x``/``y``/``size``/``headstage`` fields (``x``/``y`` in + micrometers) and verbatim labels; the channel index is + ``(bank - 1) * 32 + (term - 1)`` — ``start_chan`` is already folded into + ``bank`` via its ``// 32`` offset. A companion ``cmp_mask`` records which + indices were set so the auto-grid pass can avoid them. +3. **Auto-grid fill** — positions/bank/elec for indices covered by neither a + CMP overlay nor a source position, laid out below and to the right of the + placed geometry so they don't collide with it. The grid step matches the + placed electrode pitch (inferred from its coordinates), so auto-laid + channels share the same micrometer scale. The same :class:`ChannelMapSettings` record is also used as a per-headstage entry in :attr:`CereLinkSignalSettings.cmp_configs`. @@ -40,12 +57,13 @@ CHANNEL_DTYPE = np.dtype( [ - ("x", "f4"), - ("y", "f4"), + ("x", "i4"), # electrode x, µm (int32, matching cbPKT_CHANINFO.position) + ("y", "i4"), # electrode y, µm + ("size", "i4"), # electrode size, µm (0 = unspecified) ("label", "U16"), ("bank", "U1"), ("elec", "i4"), - ("device", "U16"), + ("headstage", "i4"), # 1-based headstage id (0 = none/auto) ] ) @@ -60,8 +78,10 @@ class ChannelMapSettings(ez.Settings): Mirrors :meth:`pycbsdk.Session.load_channel_map`.""" hs_id: int = 0 - """Headstage identifier; labels are prefixed ``"hs{hs_id}-"`` when nonzero. - Pass ``0`` (the default) to leave labels un-prefixed.""" + """Headstage identifier, passed through to :func:`pycbsdk.cmp.parse_cmp`, + where it sets each entry's ``headstage`` field. Labels are taken verbatim + (no ``"hs{hs_id}-"`` prefix); ``bank``/``elec`` disambiguate channels that + reuse a label across headstages. Pass ``0`` for single-headstage rigs.""" class ChannelMapUnitSettings(ez.Settings): @@ -73,7 +93,8 @@ class ChannelMapUnitSettings(ez.Settings): @processor_state class ChannelMapState: channel_axis: CoordinateAxis | None = None - cmp_mask: np.ndarray | None = None # bool, length == channel_axis.data length + cmp_mask: np.ndarray | None = None # bool, indices set by a CMP overlay + src_mask: np.ndarray | None = None # bool, indices positioned by the incoming axis class ChannelMapProcessor(BaseStatefulTransformer[ChannelMapUnitSettings, AxisArray, AxisArray, ChannelMapState]): @@ -92,17 +113,17 @@ def _reset_state(self, message: AxisArray) -> None: ch_dim_idx = message.dims.index("ch") n_total = message.data.shape[ch_dim_idx] - # Base layer: labels from incoming; positions filled by the overlays - # and auto-grid below. + # Base layer: labels from incoming; positions seeded from the incoming + # structured axis when the source already carries them (else filled by + # the overlays and auto-grid below). ch_data = np.zeros(n_total, dtype=CHANNEL_DTYPE) for i, label in enumerate(self._incoming_labels(message, n_total)): ch_data[i]["label"] = label - # Carry through CHANNEL_DTYPE fields other than label (above) and those - # the CMP overlays set (below) — e.g. ``device``. - incoming = getattr(message.axes.get("ch"), "data", None) - if incoming is not None and incoming.dtype.names is not None and "device" in incoming.dtype.names: - for i in range(min(n_total, len(incoming))): - ch_data[i]["device"] = incoming[i]["device"] + + # Source geometry: copy x/y/size/bank/elec/headstage straight from the + # incoming axis and record which channels were positioned. The auto-grid + # skips these; a CMP overlay (below) still overrides them. + src_mask = self._apply_incoming_positions(message, ch_data, n_total) # CMP overlays: write each headstage's entries at chan_id-1 and mark # them in cmp_mask so the auto-grid skips them. @@ -123,53 +144,119 @@ def _reset_state(self, message: AxisArray) -> None: exc, ) continue - for chan_id, entry in parsed.items(): - idx = chan_id - 1 + for (bank, term), entry in parsed.items(): + # parse_cmp keys by device (bank, term); start_chan is already + # folded into bank via its // 32 offset, so the channel index is + # a direct (bank, term) → row mapping (32 terminals per bank). + idx = (bank - 1) * 32 + (term - 1) if not (0 <= idx < n_total): continue - col, row, bank_idx, elec = entry.position - ch_data[idx]["x"] = float(col) - ch_data[idx]["y"] = float(row) - ch_data[idx]["label"] = entry.label - ch_data[idx]["bank"] = chr(ord("A") + bank_idx - 1) - ch_data[idx]["elec"] = elec + ch_data[idx]["x"] = int(entry.x) + ch_data[idx]["y"] = int(entry.y) + ch_data[idx]["size"] = int(entry.size) + ch_data[idx]["label"] = entry.label # verbatim (no hs{N}- prefix) + ch_data[idx]["bank"] = chr(ord("A") + bank - 1) + ch_data[idx]["elec"] = term + ch_data[idx]["headstage"] = entry.headstage cmp_mask[idx] = True self.state.channel_axis = CoordinateAxis(data=ch_data, dims=["ch"], unit="struct") self.state.cmp_mask = cmp_mask + # CMP wins over source geometry: a CMP-claimed index is "placed" by the + # overlay, not the source. + self.state.src_mask = src_mask & ~cmp_mask - # Auto-grid: position/bank/elec for indices no CMP claimed, offset - # below the CMP geometry so they don't overlap. + # Auto-grid: position/bank/elec for indices neither a CMP nor the source + # claimed, offset below the placed geometry so they don't overlap. self._fill_auto_grid() + @staticmethod + def _apply_incoming_positions(message: AxisArray, ch_data: np.ndarray, n_total: int) -> np.ndarray: + """Copy structured geometry from the incoming ``ch`` axis into ``ch_data``. + + Returns a bool ``src_mask`` of channels that carry a usable source + position — any non-origin coordinate, plus the *first* channel at the + ``(0, 0)`` origin. Origin pile-ups beyond the first are the device's + "unmapped" sentinel; they are left ``False`` so the auto-grid claims + them (their copied zero coordinates get overwritten there). When the + incoming axis is unstructured or lacks ``x``/``y`` (e.g. a label-only + or plain TimeSeries source), nothing is copied and the mask is all + ``False`` — the original pure auto-grid behavior. + """ + src_mask = np.zeros(n_total, dtype=bool) + ch_axis = message.axes.get("ch") + incoming = getattr(ch_axis, "data", None) + names = getattr(getattr(incoming, "dtype", None), "names", None) + if incoming is None or not names or not ({"x", "y"} <= set(names)): + return src_mask + + copy_fields = [f for f in ("x", "y", "size", "bank", "elec", "headstage") if f in names] + seen_origin = False + for i in range(min(n_total, incoming.shape[0])): + for f in copy_fields: + ch_data[i][f] = incoming[f][i] + at_origin = int(incoming["x"][i]) == 0 and int(incoming["y"][i]) == 0 + if at_origin: + if seen_origin: + continue # duplicate origin → leave to the auto-grid + seen_origin = True + src_mask[i] = True + return src_mask + def _fill_auto_grid(self) -> None: ch_data = self.state.channel_axis.data - cmp_mask = self.state.cmp_mask - auto_idx = np.flatnonzero(~cmp_mask) + # "Placed" = positioned by a CMP overlay or the incoming source axis. + placed_mask = self.state.cmp_mask | self.state.src_mask + auto_idx = np.flatnonzero(~placed_mask) if auto_idx.size == 0: return - if cmp_mask.any(): - max_row = float(ch_data["y"][cmp_mask].max()) + # Step matches the placed electrode pitch (≈400 µm) so the auto-grid + # sits on the same scale as the real geometry. Without this, the µm + # coordinates would dwarf a unit-spaced auto-grid. Falls back to 1 when + # nothing is placed (pure auto-grid from the origin). + step = self._placed_pitch(ch_data, placed_mask) + + if placed_mask.any(): + max_row = int(ch_data["y"][placed_mask].max()) max_bank_ord = max( - (ord(str(b)) for b in ch_data["bank"][cmp_mask] if str(b)), + (ord(str(b)) for b in ch_data["bank"][placed_mask] if str(b)), default=ord("A") - 1, ) else: - # No CMP yet — start auto-grid at the origin with bank A. - # max_row = -2 makes start_row = 0 below. - max_row = -2.0 + # Nothing placed yet — start auto-grid at the origin with bank A. + # max_row = -2*step makes start_row = 0 below. + max_row = -2 * step max_bank_ord = ord("A") - 1 - start_row = max_row + 2 + start_row = max_row + 2 * step next_bank_ord = max_bank_ord + 1 grid_size = max(1, math.ceil(math.sqrt(auto_idx.size))) for i, idx in enumerate(auto_idx): - ch_data[idx]["x"] = float(i % grid_size) - ch_data[idx]["y"] = float(start_row + i // grid_size) + ch_data[idx]["x"] = (i % grid_size) * step + ch_data[idx]["y"] = start_row + (i // grid_size) * step + ch_data[idx]["size"] = step # synthetic electrodes sized to the grid pitch ch_data[idx]["bank"] = chr(next_bank_ord + i // 32) ch_data[idx]["elec"] = (i % 32) + 1 + ch_data[idx]["headstage"] = 0 # auto-grid channels have no headstage + + @staticmethod + def _placed_pitch(ch_data: np.ndarray, placed_mask: np.ndarray) -> int: + """Smallest positive spacing among the placed channels' distinct x/y. + + This is the electrode pitch in micrometers (≈400 for a default Utah + array). Defaults to ``1`` when nothing is placed or the geometry is + degenerate, giving the pure auto-grid unit spacing from the origin.""" + if not placed_mask.any(): + return 1 + deltas: list[int] = [] + for field in ("x", "y"): + vals = np.unique(ch_data[field][placed_mask]) + if vals.size > 1: + deltas.append(int(np.diff(vals).min())) + positive = [d for d in deltas if d > 0] + return min(positive) if positive else 1 @staticmethod def _incoming_labels(message: AxisArray, n_total: int) -> list[str]: diff --git a/tests/test_channel_map.py b/tests/test_channel_map.py index 28da045..227ac6f 100644 --- a/tests/test_channel_map.py +++ b/tests/test_channel_map.py @@ -53,7 +53,11 @@ def test_single_file(self): assert len(ch_ax.data) == 128 def test_structured_fields(self): - """First and last channels have correct label, bank, elec, x, y.""" + """First and last channels have correct label, bank, elec, x, y. + + Coordinates are in micrometers: the default index grid (col 0-15, + row 0-7) is scaled by the 400 µm Utah-array pitch by parse_cmp. + """ proc = _make_processor(CMP_FILE) out = proc(_make_message(128)) data = out.axes["ch"].data @@ -61,14 +65,18 @@ def test_structured_fields(self): assert first["label"] == "chan1" assert first["bank"] == "A" assert first["elec"] == 1 - assert first["x"] == pytest.approx(0.0) - assert first["y"] == pytest.approx(7.0) + assert first["x"] == 0 # col 0 * 400 µm + assert first["y"] == 2800 # row 7 * 400 µm + assert first["size"] == 400 # Utah-array pitch + assert first["headstage"] == 0 # hs_id default last = data[-1] assert last["label"] == "chan128" assert last["bank"] == "D" assert last["elec"] == 32 - assert last["x"] == pytest.approx(15.0) - assert last["y"] == pytest.approx(0.0) + assert last["x"] == 6000 # col 15 * 400 µm + assert last["y"] == 0 # row 0 * 400 µm + assert last["size"] == 400 + assert last["headstage"] == 0 def test_original_data_preserved(self): proc = _make_processor(CMP_FILE) @@ -100,6 +108,8 @@ def test_field_access_patterns(self): assert data["elec"].shape == (128,) assert data["x"].shape == (128,) assert data["y"].shape == (128,) + assert data["size"].shape == (128,) + assert data["headstage"].shape == (128,) def test_electrode_values(self): """Electrode IDs cycle 1-32 within each bank.""" @@ -193,23 +203,29 @@ def test_overlay_preserves_n_total(self): out = _make_processor(CMP_FILE)(_make_message(256)) assert len(out.axes["ch"].data) == 256 - def test_overlay_out_of_range_chan_id_skipped(self): - """CMP entries whose chan_id-1 ≥ n_total are ignored, not appended.""" - proc = _make_processor(CMP_FILE, start_chan=200) # would write 200..327 + def test_overlay_out_of_range_index_skipped(self): + """CMP entries whose (bank, term) index ≥ n_total are skipped, not appended.""" + # start_chan=200 → bank offset 200//32 = 6, so CMP banks A-D land on + # device banks G-J → indices (7-1)*32 .. (10-1)*32+31 = 192..319. + proc = _make_processor(CMP_FILE, start_chan=200) out = proc(_make_message(256)) data = out.axes["ch"].data assert len(data) == 256 - # Indices 199..255 carry the CMP's first 57 entries; rest skipped. - assert data[199]["label"] == "chan1" - assert data[255]["label"] == "chan57" + # chan1 → bank G term 1 → idx 192; banks G-H (chan1..64) fit in 0..255, + # banks I-J (idx 256..319) are out of range and dropped. + assert data[192]["label"] == "chan1" + assert data[255]["label"] == "chan64" - def test_hs_id_prefixes_labels(self): - """hs_id != 0 prefixes labels with 'hs{hs_id}-'.""" + def test_hs_id_does_not_prefix_labels(self): + """Labels are verbatim; hs_id sets the entry's headstage, not the label.""" proc = _make_processor(CMP_FILE, hs_id=2) out = proc(_make_message(128)) - labels = out.axes["ch"].data["label"] - assert labels[0] == "hs2-chan1" - assert labels[127] == "hs2-chan128" + data = out.axes["ch"].data + assert data["label"][0] == "chan1" + assert data["label"][127] == "chan128" + # hs_id lands in the headstage field instead. + assert data["headstage"][0] == 2 + assert data["headstage"][127] == 2 def test_missing_cmp_keeps_base_layer(self): """A bad CMP path warns and leaves the existing axis intact.""" @@ -291,13 +307,13 @@ def test_empty_cmp_configs_rebuilds_base(self): class TestAutoGridPlacement: def test_auto_grid_y_offset_below_cmp(self): - """Auto-grid rows start below the CMP's max y (gap of +2).""" + """Auto-grid rows start below the CMP's max y (gap of 2 pitches = 800 µm).""" proc = _make_processor(CMP_FILE) out = proc(_make_message(256)) data = out.axes["ch"].data cmp_y_max = float(data["y"][:128].max()) auto_y_min = float(data["y"][128:].min()) - assert auto_y_min >= cmp_y_max + 2 + assert auto_y_min >= cmp_y_max + 800 def test_auto_grid_bank_starts_past_cmp(self): """Auto-grid banks start one letter past the CMP's highest bank.""" @@ -311,12 +327,114 @@ def test_auto_grid_bank_starts_past_cmp(self): def test_auto_grid_skips_cmp_indices_with_offset_start_chan(self): """With start_chan=129 the CMP fills 128..255; auto-grid lays out 0..127 below.""" + # start_chan=129 → bank offset 4 → CMP banks A-D land on device banks + # E-H, i.e. indices 128..255. proc = _make_processor(CMP_FILE, start_chan=129) out = proc(_make_message(256)) data = out.axes["ch"].data - # CMP positions on the upper half. - assert data[128]["bank"] == "A" # CMP's first sorted bank - # Auto-grid (lower half) is offset below CMP's max y (=7) → y >= 9. - assert data[0]["y"] >= 9.0 - # And uses banks past CMP's max (D) → starts at E. - assert data[0]["bank"] == "E" + # CMP positions on the upper half, starting at device bank E. + assert data[128]["bank"] == "E" + # Auto-grid (lower half) sits below CMP's max y (2800 µm) by 2 pitches. + assert data[0]["y"] >= 3600.0 + # And uses banks past CMP's max (H) → starts at I. + assert data[0]["bank"] == "I" + + +# --------------------------------------------------------------------------- +# Source-provided geometry (no CMP file needed) +# --------------------------------------------------------------------------- + + +def _make_message_passthrough(n_channels: int, positions: np.ndarray, n_time: int = 5) -> AxisArray: + """Build a message whose ``ch`` axis already carries CHANNEL_DTYPE geometry, + mimicking a CereLink source that read positions from device chaninfo. + + ``positions`` is an ``(n_channels, 2)`` array of (x, y) in micrometers. + """ + ch_data = np.zeros(n_channels, dtype=CHANNEL_DTYPE) + for i in range(n_channels): + ch_data[i]["x"] = positions[i, 0] + ch_data[i]["y"] = positions[i, 1] + ch_data[i]["size"] = 400 + ch_data[i]["label"] = f"src{i}" + ch_data[i]["bank"] = chr(ord("A") + i // 32) + ch_data[i]["elec"] = (i % 32) + 1 + ch_data[i]["headstage"] = 1 + return AxisArray( + data=np.zeros((n_time, n_channels)), + dims=["time", "ch"], + axes={ + "time": LinearAxis(offset=0.0, gain=0.001), + "ch": CoordinateAxis(data=ch_data, dims=["ch"]), + }, + ) + + +class TestSourceGeometry: + def test_source_positions_preserved_without_cmp(self): + """Structured x/y/size/bank/elec/headstage from the source flow through + verbatim when no CMP file is configured.""" + pos = np.column_stack([(np.arange(64) % 8) * 400, (np.arange(64) // 8) * 400]) + proc = _make_processor(None) # empty cmp_configs + out = proc(_make_message_passthrough(64, pos)) + data = out.axes["ch"].data + np.testing.assert_array_equal(data["x"], pos[:, 0]) + np.testing.assert_array_equal(data["y"], pos[:, 1]) + assert data[0]["label"] == "src0" + assert data[0]["size"] == 400 + assert data[0]["headstage"] == 1 + # Every channel is source-placed → nothing auto-gridded. + assert proc.state.src_mask.all() + assert not proc.state.cmp_mask.any() + + def test_cmp_overrides_source_positions(self): + """A CMP overlay wins over source geometry at the indices it claims.""" + pos = np.column_stack([(np.arange(128) % 16) * 999, (np.arange(128) // 16) * 999]) + proc = _make_processor(CMP_FILE) # start_chan=1 → indices 0..127 + out = proc(_make_message_passthrough(128, pos)) + data = out.axes["ch"].data + # CMP labels/coords replaced the source ones. + assert data[0]["label"] == "chan1" + assert data[0]["x"] == 0 + assert data[127]["label"] == "chan128" + assert proc.state.cmp_mask.all() + # src_mask cleared where the CMP took over. + assert not (proc.state.src_mask & proc.state.cmp_mask).any() + + def test_unmapped_origin_pileup_falls_through_to_auto_grid(self): + """The first (0, 0) channel is kept; later origin channels (the device's + 'unmapped' sentinel) are auto-gridded below the real geometry.""" + pos = np.array([[0, 0], [400, 0], [0, 0], [0, 0]]) + proc = _make_processor(None) + out = proc(_make_message_passthrough(4, pos)) + data = out.axes["ch"].data + # ch0: lone-kept origin electrode; ch1: real position. + assert (data[0]["x"], data[0]["y"]) == (0, 0) + assert (data[1]["x"], data[1]["y"]) == (400, 0) + assert proc.state.src_mask[0] and proc.state.src_mask[1] + # ch2/ch3: duplicate origins → auto-gridded, not left stacked at (0,0). + assert not proc.state.src_mask[2] + assert not proc.state.src_mask[3] + auto_y = min(int(data[2]["y"]), int(data[3]["y"])) + assert auto_y >= int(data["y"][proc.state.src_mask].max()) + 2 * 400 + + def test_auto_grid_sits_below_source_geometry(self): + """With a partial source map, the auto-grid offsets below the source's + max y, same as it does below CMP geometry.""" + # 4 source-placed channels on a 400 µm grid; 4 unmapped origin dupes. + pos = np.array([[0, 0], [400, 0], [0, 400], [400, 400], [0, 0], [0, 0], [0, 0], [0, 0]]) + proc = _make_processor(None) + out = proc(_make_message_passthrough(8, pos)) + data = out.axes["ch"].data + src_y_max = int(data["y"][proc.state.src_mask].max()) + auto_y_min = int(data["y"][~proc.state.src_mask].min()) + assert auto_y_min >= src_y_max + 800 + + def test_unstructured_incoming_still_auto_grids_from_origin(self): + """Label-only / plain incoming axes keep the original pure auto-grid.""" + proc = _make_processor(None) + out = proc(_make_message(64)) # ch_data = np.arange(64), unstructured + data = out.axes["ch"].data + assert data[0]["x"] == 0 + assert data[0]["y"] == 0 + assert not proc.state.src_mask.any() diff --git a/tests/test_integration.py b/tests/test_integration.py index c76bec3..86cb36c 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -115,16 +115,16 @@ def test_monotonic_timestamps(self, nplayserver, tmp_path): ), f"non-monotonic at (idx, prev, next): {violations}; offsets[0]={offsets[0]}, offsets[-1]={offsets[-1]}" def test_all_channels(self, nplayserver, tmp_path): - """`SliceConfig(channels=None)` configures all matching FRONTEND channels. - NPLAY emulates the full NSP channel layout (256 FRONTEND), so even - with a 4-channel recording the AxisArray has 256 columns — the - unrecorded channels carry zeros but are still configured.""" + """Default `SliceConfig` (channels=ChannelSelection.ALL) configures all + matching FRONTEND channels. NPLAY emulates the full NSP channel layout + (256 FRONTEND), so even with a 4-channel recording the AxisArray has 256 + columns — the unrecorded channels carry zeros but are still configured.""" expected_n_ch = 256 # NPLAY's nominal FRONTEND count messages = _run_signal_source( CereLinkSignalSettings( device_type=DeviceType.NPLAY, subscribe_rate=SampleRate.SR_30kHz, - configure=SliceConfig(channel_type=ChannelType.FRONTEND), # channels=None => all + configure=SliceConfig(channel_type=ChannelType.FRONTEND), # default ALL microvolts=False, cbtime=True, ), diff --git a/tests/test_settings.py b/tests/test_settings.py index 72425d2..05b165d 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -1,14 +1,18 @@ """Unit tests for CereLink settings validation (no hardware needed).""" import pickle +from unittest.mock import MagicMock import pytest from pycbsdk import ChannelType, SampleRate from ezmsg.blackrock.cerelink import ( CcfConfig, + CereLinkSignalProducer, CereLinkSignalSettings, + CereLinkSpikeProducer, CereLinkSpikeSettings, + ChannelSelection, SliceConfig, ) @@ -76,11 +80,18 @@ def test_slice_defaults(self): ) sc = s.configure assert isinstance(sc, SliceConfig) - assert sc.channels is None # all matching + assert sc.channels is ChannelSelection.ALL # default: all matching assert sc.channel_type == ChannelType.FRONTEND assert sc.ac_input_coupling is False assert sc.enable_spiking is False + def test_invalid_channels_type_rejected(self): + """A bare string/int for channels is rejected (must be list or enum).""" + with pytest.raises(TypeError, match="must be a list"): + SliceConfig(channels="all") + with pytest.raises(TypeError, match="must be a list"): + SliceConfig(channels=3) + def test_slice_explicit_channels(self): s = CereLinkSignalSettings( subscribe_rate=SampleRate.SR_2kHz, @@ -90,6 +101,135 @@ def test_slice_explicit_channels(self): assert s.configure.ac_input_coupling is True +class TestSliceConfigureChannels: + """``CereLinkSignalProducer._apply_slice_configure`` channel selection, + driven by a mock pycbsdk Session (no hardware).""" + + @staticmethod + def _producer_with_session(session) -> CereLinkSignalProducer: + # Build without ezmsg machinery; we only exercise _apply_slice_configure. + prod = CereLinkSignalProducer.__new__(CereLinkSignalProducer) + prod.settings = CereLinkSignalSettings(subscribe_rate=SampleRate.SR_RAW) + state = MagicMock() + state.session = session + prod.state = state + return prod + + @staticmethod + def _session_with_enabled(enabled: list[int], total: int = 256): + """Session where ``enabled`` channels sit in the raw group and the + device exposes ``total`` FRONTEND channels.""" + sess = MagicMock() + sess.get_matching_channel_ids.return_value = list(range(1, total + 1)) + sess.get_group_channels.side_effect = lambda g: (list(enabled) if g == int(SampleRate.SR_RAW) else []) + return sess + + def test_enabled_configures_only_enabled(self): + sess = self._session_with_enabled(list(range(1, 129)), total=256) + prod = self._producer_with_session(sess) + prod._apply_slice_configure(SliceConfig(channel_type=ChannelType.FRONTEND, channels=ChannelSelection.ENABLED)) + chans, ctype, rate = sess.set_sample_group.call_args.args + assert chans == list(range(1, 129)) # only the enabled bank + assert rate == SampleRate.SR_RAW + # disable_others=False → the unused bank is left untouched. + assert sess.set_sample_group.call_args.kwargs["disable_others"] is False + # Reads the device state before configuring. + sess.sync.assert_called_once() + + def test_enabled_unions_continuous_groups(self): + # Enabled split across two groups → both should be configured. + sess = MagicMock() + sess.get_matching_channel_ids.return_value = list(range(1, 257)) + sess.get_group_channels.side_effect = lambda g: { + int(SampleRate.SR_30kHz): [1, 2], + int(SampleRate.SR_RAW): [5, 6], + }.get(g, []) + prod = self._producer_with_session(sess) + prod._apply_slice_configure(SliceConfig(channel_type=ChannelType.FRONTEND, channels=ChannelSelection.ENABLED)) + assert sess.set_sample_group.call_args.args[0] == [1, 2, 5, 6] + + def test_enabled_excludes_non_matching_type(self): + # A channel enabled but not FRONTEND is dropped by the type intersection. + sess = MagicMock() + sess.get_matching_channel_ids.return_value = [1, 2, 3] # FRONTEND only + sess.get_group_channels.side_effect = lambda g: ([1, 2, 3, 99] if g == int(SampleRate.SR_RAW) else []) + prod = self._producer_with_session(sess) + prod._apply_slice_configure(SliceConfig(channel_type=ChannelType.FRONTEND, channels=ChannelSelection.ENABLED)) + assert sess.set_sample_group.call_args.args[0] == [1, 2, 3] + + def test_enabled_empty_warns_but_does_not_raise(self, caplog): + sess = self._session_with_enabled([], total=256) + prod = self._producer_with_session(sess) + prod._apply_slice_configure(SliceConfig(channel_type=ChannelType.FRONTEND, channels=ChannelSelection.ENABLED)) + assert sess.set_sample_group.call_args.args[0] == [] + assert any("no FRONTEND channels are currently enabled" in r.message for r in caplog.records) + + def test_all_configures_all_and_disables_others(self): + sess = self._session_with_enabled(list(range(1, 129)), total=256) + prod = self._producer_with_session(sess) + # Default channels=ChannelSelection.ALL. + prod._apply_slice_configure(SliceConfig(channel_type=ChannelType.FRONTEND)) + # ALL resolves to every matching channel; others get disabled. + assert sess.set_sample_group.call_args.args[0] == list(range(1, 257)) + assert sess.set_sample_group.call_args.kwargs["disable_others"] is True + # ALL does not inspect the enabled set. + sess.get_group_channels.assert_not_called() + + def test_explicit_list_respected_and_disables_others(self): + sess = self._session_with_enabled(list(range(1, 129)), total=256) + prod = self._producer_with_session(sess) + prod._apply_slice_configure(SliceConfig(channels=[1, 2, 3], channel_type=ChannelType.FRONTEND)) + # Explicit list passed through verbatim; others disabled; enabled set not consulted. + assert sess.set_sample_group.call_args.args[0] == [1, 2, 3] + assert sess.set_sample_group.call_args.kwargs["disable_others"] is True + sess.get_group_channels.assert_not_called() + + +class TestSpikeSliceConfigure: + """``CereLinkSpikeProducer`` channel selection / extraction handling, + driven by a mock pycbsdk Session (no hardware).""" + + @staticmethod + def _producer_with_session(session) -> CereLinkSpikeProducer: + prod = CereLinkSpikeProducer.__new__(CereLinkSpikeProducer) + prod.settings = CereLinkSpikeSettings() + state = MagicMock() + state.session = session + prod.state = state + return prod + + @staticmethod + def _session_with_extraction(extracting: list[int], total: int = 8): + """Session exposing ``total`` FRONTEND channels, of which ``extracting`` + have the SPKOPTS extract bit set.""" + sess = MagicMock() + ids = list(range(1, total + 1)) + sess.get_matching_channel_ids.return_value = ids + sess.get_channels_field.return_value = [1 if cid in extracting else 0 for cid in ids] + return sess + + def test_enabled_leaves_extraction_untouched(self): + """channels=ENABLED never calls set_spike_extraction, even with enable_spiking.""" + sess = self._session_with_extraction([2, 5]) + prod = self._producer_with_session(sess) + prod._apply_slice_configure(SliceConfig(channels=ChannelSelection.ENABLED, enable_spiking=True)) + sess.set_spike_extraction.assert_not_called() + + def test_enabled_resolves_to_extraction_enabled_channels(self): + """The spike-stream meaning of 'enabled' is the SPKOPTS extract bit.""" + sess = self._session_with_extraction([2, 5]) + prod = self._producer_with_session(sess) + assert prod._enabled_channels(ChannelType.FRONTEND) == [2, 5] + + def test_enable_spiking_with_list_sets_extraction(self): + sess = self._session_with_extraction([]) + prod = self._producer_with_session(sess) + prod._apply_slice_configure(SliceConfig(channels=[1, 3], enable_spiking=True)) + chans, ctype, enabled = sess.set_spike_extraction.call_args.args + assert chans == [1, 3] + assert enabled is True + + class TestPickleRoundtrip: """Settings + their nested configure types must survive pickling so they can flow through ezmsg's INPUT_SETTINGS message stream."""