Source code for ts3.models.transductive.ndt_stitch.ndt_stitch_extractor

"""NDT-stitch's unit-embedding extractor.

The embedding is the unit's pair of stitcher columns, ``[in.T, out]`` concatenated to 2D.
Those are per-session parameters, so they exist only for the sessions a
checkpoint was trained on, and the units come from the recording rather than from a lookup:
the stitcher has no vocabulary, only a column order matching ``units.id``.
"""

import numpy as np
import torch

from core.dataset import WholeSessionSpikeDataset
from core.utils.logger import get_cli_logger
from ts3.models.transductive.base import TransductiveExtractor

logger = get_cli_logger()




class NDTStitchExtractor(TransductiveExtractor):


    def read(
        self, state_dict: dict, dataset: WholeSessionSpikeDataset, uids: set[str]
    ) -> tuple[torch.Tensor, np.ndarray]:
        in_keys = [k for k in state_dict if k.startswith("in_stitcher") and k.endswith("weight")]
        out_keys = [k for k in state_dict if k.startswith("out_stitcher") and k.endswith("weight")]
        assert {k.split(".")[1] for k in in_keys} == {k.split(".")[1] for k in out_keys}, (
            "In and out stitcher keys must have the same session ids"
        )

        embs, out_uids = [], []
        for rec_id in sorted({k.split(".")[1] for k in in_keys}):
            if rec_id not in dataset.recording_ids:
                logger.warning(f"Recording {rec_id} not found in dataset")
                continue
            in_weight = state_dict[f"in_stitcher.{rec_id}.weight"].T  # (v, D)
            out_weight = state_dict[f"out_stitcher.{rec_id}.weight"]  # (v, D)
            assert in_weight.shape == out_weight.shape, (
                f"In and out stitcher weights must have the same shape, got "
                f"{in_weight.shape} and {out_weight.shape}"
            )
            embs.append(torch.cat([in_weight, out_weight], dim=1))  # (v, 2D)

            rec_uids = dataset.get_recording(rec_id).units.id.astype(str)  # (v,)
            assert len(rec_uids) == in_weight.shape[0], (
                f"Number of units in recording {rec_id} ({len(rec_uids)}) does not match "
                f"the number of units in the stitcher weights ({in_weight.shape[0]})"
            )
            out_uids.extend(rec_uids)

        return torch.cat(embs, dim=0), np.array(out_uids)