Dataloader output: dict instead of 4-tuple

README.md

@@ -60,12 +60,13 @@ loader = DataLoader(ds, batch_size=8, collate_fn=neural_collate)
 # 2) Load a pretrained model from the Hugging Face Hub.
 model = StateNet.from_pretrained("urancon/deepSTRF-statenet-gru-ns1").eval()
 
-# 3) Score it.
-stims, responses, _, _ = next(iter(loader))
-pred    = model(stims)                                      # (B, N, R=1, T)
-psth    = responses.nanmean(dim=2, keepdim=True)
-cc      = corrcoef(pred, psth, reduction='mean')
-cc_norm = normalized_corrcoef(pred, responses, method='schoppe', reduction='mean')
+# 3) Score it. Each batch is a dict: 'stims', 'responses', 'valid_mask', 'stim_meta'.
+batch     = next(iter(loader))
+responses = batch['responses']
+pred      = model(batch['stims'])                           # (B, N, R=1, T)
+psth      = responses.nanmean(dim=2, keepdim=True)
+cc        = corrcoef(pred, psth, reduction='mean')
+cc_norm   = normalized_corrcoef(pred, responses, method='schoppe', reduction='mean')
 print(f"CCraw = {cc:.3f}   CCnorm = {cc_norm:.3f}")
 ```
 

deepSTRF/datasets/neural_dataset.py

@@ -260,9 +260,14 @@ def __getitem__(self, idx):
 
         Returns
         -------
-        stim, responses, mask, stim_meta
-            Tuple for a scalar index, or 4 lists for slice / list indexing.
-            ``responses`` and ``mask`` are restricted to the selected neurons.
+        dict
+            A dict with keys ``'stims'``, ``'responses'``, ``'valid_mask'``,
+            ``'stim_meta'``. For a scalar index each value is a single item;
+            for slice / list indexing each value is a list. ``'responses'``
+            and ``'valid_mask'`` are restricted to the selected neurons. The
+            dict container (rather than a positional tuple) lets datasets add
+            extra per-trial keys later — e.g. ``'behav'`` for behavioural
+            covariates — without changing the unpacking contract.
         """
         iter_idx = self._iter_idx  # snapshot once; O(S * |I|) per call
 
@@ -292,8 +297,10 @@ def __getitem__(self, idx):
         masks = [all_masks[i][selected] for i in indices]
 
         if single:
-            return stims[0], resps[0], masks[0], metas[0]
-        return stims, resps, masks, metas
+            return {'stims': stims[0], 'responses': resps[0],
+                    'valid_mask': masks[0], 'stim_meta': metas[0]}
+        return {'stims': stims, 'responses': resps,
+                'valid_mask': masks, 'stim_meta': metas}
 
     def __repr__(self):
         return (f"{self.__class__.__name__}(N_neurons={self.get_N()}, "

deepSTRF/training/fitter.py

@@ -376,9 +376,8 @@ def _train_one_epoch(self) -> Dict[str, Any]:
         responses_list: List[torch.Tensor] = []
 
         for batch in self.train_loader:
-            stims, responses, _valid_mask, _stim_metas = batch
-            stims = stims.to(self.device)
-            responses = responses.to(self.device)
+            stims = batch['stims'].to(self.device)
+            responses = batch['responses'].to(self.device)
 
             self.optimizer.zero_grad()
             pred = self.model(stims)
@@ -417,9 +416,8 @@ def _evaluate(self, loader: DataLoader) -> Dict[str, Any]:
 
         with torch.no_grad():
             for batch in loader:
-                stims, responses, _valid_mask, _stim_metas = batch
-                stims = stims.to(self.device)
-                responses = responses.to(self.device)
+                stims = batch['stims'].to(self.device)
+                responses = batch['responses'].to(self.device)
 
                 pred = self.model(stims)
                 if hasattr(self.model, "detach"):

deepSTRF/utils/data.py

@@ -109,43 +109,45 @@ def neural_collate(batch):
 
     Parameters
     ----------
-    batch : list of 4-tuples
-        Each tuple is ``(stim, per_neuron_responses, per_neuron_mask,
-        stim_meta)`` as yielded by ``NeuralDataset.__getitem__`` for a
-        single item:
+    batch : list of dict
+        Each dict is one item as yielded by ``NeuralDataset.__getitem__``,
+        with keys:
 
-        * ``stim`` — a stim tensor of shape ``(..., T_s)`` (modality-specific
+        * ``'stims'`` — a stim tensor of shape ``(..., T_s)`` (modality-specific
           leading dims, e.g. ``(1, F, T_s)`` for audio).
-        * ``per_neuron_responses`` — list of length ``N_selected``; each
-          element is a ``(R_{s,n}, T_s)`` spike-count tensor or a
-          ``(1, 1)`` NaN sentinel.
-        * ``per_neuron_mask`` — ``(N_selected,)`` bool tensor (currently
-          ignored; the fine-grained ``valid_mask`` returned by this
-          function subsumes it).
-        * ``stim_meta`` — per-stim metadata dict.
+        * ``'responses'`` — list of length ``N_selected``; each element is a
+          ``(R_{s,n}, T_s)`` spike-count tensor or a ``(1, 1)`` NaN sentinel.
+        * ``'valid_mask'`` — ``(N_selected,)`` per-neuron bool tensor (ignored
+          here; the fine-grained batch ``'valid_mask'`` below subsumes it).
+        * ``'stim_meta'`` — per-stim metadata dict.
+
+        Extra keys (e.g. ``'behav'``) are passed through untouched: any key
+        not handled explicitly is collected into a length-``B`` list.
 
     Returns
     -------
-    stims : torch.Tensor
-        ``(B, ..., T_stim_max)`` float tensor, zero-padded along the last
-        axis. Contains no NaN.
-    responses : torch.Tensor
-        ``(B, N_selected, R_max, T_resp_max)`` float tensor. NaN-padded
-        along both the repeat (``R``) and time (``T``) axes. Fully-NaN
-        slabs mark (stim, neuron) pairs with no recorded data. The
-        response-time axis is sized to ``T_resp_max`` independently of the
-        stim-time axis: in spectrogram mode the two are equal (one bin
-        per neural sample), but in waveform mode the stim axis runs at
-        ``audio_fs`` Hz while responses stay at the dataset's neural
-        ``dt_ms`` rate.
-    valid_mask : torch.Tensor
-        ``(B, N_selected, R_max, T_resp_max)`` bool tensor.
-        ``~responses.isnan()``, cached here so downstream loss code does
-        not have to recompute.
-    stim_metas : list
-        Length-``B`` list of the per-item stim_meta dicts.
+    dict
+        A dict with keys:
+
+        * ``'stims'`` — ``(B, ..., T_stim_max)`` float tensor, zero-padded
+          along the last axis. Contains no NaN.
+        * ``'responses'`` — ``(B, N_selected, R_max, T_resp_max)`` float
+          tensor. NaN-padded along both the repeat (``R``) and time (``T``)
+          axes. Fully-NaN slabs mark (stim, neuron) pairs with no recorded
+          data. The response-time axis is sized to ``T_resp_max``
+          independently of the stim-time axis: in spectrogram mode the two
+          are equal (one bin per neural sample), but in waveform mode the
+          stim axis runs at ``audio_fs`` Hz while responses stay at the
+          dataset's neural ``dt_ms`` rate.
+        * ``'valid_mask'`` — ``(B, N_selected, R_max, T_resp_max)`` bool
+          tensor, ``~responses.isnan()``, cached here so downstream loss code
+          does not have to recompute.
+        * ``'stim_meta'`` — length-``B`` list of the per-item stim_meta dicts.
+        * any extra per-item keys — length-``B`` lists, passed through.
     """
-    stims_list, resps_list, _masks_list, metas_list = zip(*batch)
+    stims_list = [item['stims'] for item in batch]
+    resps_list = [item['responses'] for item in batch]
+    metas_list = [item['stim_meta'] for item in batch]
     B = len(stims_list)
     N = len(resps_list[0])
 
@@ -171,7 +173,21 @@ def neural_collate(batch):
     # "for free" and does not need to scan again.
     valid_mask = ~responses.isnan()
 
-    return stims, responses, valid_mask, list(metas_list)
+    out = {
+        'stims': stims,
+        'responses': responses,
+        'valid_mask': valid_mask,
+        'stim_meta': list(metas_list),
+    }
+
+    # pass through any extra per-item keys (e.g. 'behav') as length-B lists,
+    # so datasets can add covariates without touching this collate.
+    handled = {'stims', 'responses', 'valid_mask', 'stim_meta'}
+    for key in batch[0]:
+        if key not in handled:
+            out[key] = [item[key] for item in batch]
+
+    return out
 
 
 def concat_neural_datasets(datasets: Sequence[NeuralDataset],

docs/_source/md/README_datasets.md

@@ -1,37 +1,71 @@
 # Conventions
 
-One of the main contributions of this repository is the presence of several off-the-shelf electrophysiology datasets, 
-compiled from various sources, preprocessed, and represented by convenient PyTorch classes.
-
-A major limitation hindering research in the field of auditory neural response fitting, is the great variability of 
-preprocessing methods and formats, due to signals of different nature (e.g. patch-clamp, extracellular recordings),
-different groups having different habits, different experimental setups, etc. This is why we tried to adopt a single 
-"format" for all datasets present in this repository.
-
-All datasets inherit from `torch.utils.data.dataset.Dataset` and can therefore benefit from other PyTorch utilities, 
-such as random splitting, shuffling, concatenating, augmentation, and so on.
-
-Because original data differ a lot between sources, the core and attributes of each dataset class might differ a bit, 
-but they still share a lot in common and can be used in very similar ways from the exterior.
-Namely, the constructor of dataset classes must include one argument to select neurons according to a criterion (mask, 
-index, a label, brain area). By default, all neurons are selected and teh constructor will grab the data for all of them. Similarly, the ability to choose stimuli of different nature (natural, pure tone, chord, etc.) 
-should be implemented as well. If several signals are available for each neuron, then an option should be made to choose 
-between them.
-
-In this framework, each dataset contain different sets of data for each neuron, notably:
-* input stimuli (spectrograms) that were presented to this neuron
-* the corresponding response, aligned in time
-* per-stimulus metrics, like measures of variability (e.g., *TTRC*, *CC_half*) pre-computed once during dataset instanciation,
-to save time during fitting
-
-As mentioned above, the set of units to work with is defined according to the user's criterion, in the constructor. 
-This number of neurons becomes the attribute `self.N_neurons`. 
-
-An important thing to note is that depending on the dataset, neurons may have different numbers of stimulus/response pairs,
-of different durations. Because of this, it can be necessary to fit computational models *unit by unit*, and then average 
-results across units. Therefore, all datasets come with a method `select_neuron(index)` that sets the attribute `self.I` 
-to the desired `index` (0 by default at instanciation). This attribute can also be a list of multiple desired index, if
-you want to train on a *population*, rather than a single isolated unit. You can then iterate over the stimulus/response 
-pairs of the currently selected neurons.
-
-Each call of the dataset or associated dataloader will return 4 things: `spectrogram, responses, ccmax, ttrc`.
+One of the main contributions of this library is a growing zoo of off-the-shelf
+electrophysiology (and EEG) datasets, compiled from various public sources,
+preprocessed, and exposed through a single PyTorch API.
+
+A major obstacle in sensory-response fitting is the sheer variability of
+preprocessing methods and formats — different signals (extracellular spikes,
+multi-unit activity, intracellular potential, scalp EEG), different labs, and
+different experimental setups. deepSTRF hides that behind one common contract so
+the same model and training loop work across datasets.
+
+## One base class, one shape
+
+Every dataset subclasses
+{class}`~deepSTRF.datasets.neural_dataset.NeuralDataset` (itself a
+`torch.utils.data.Dataset`), so it composes with the usual PyTorch utilities
+(splitting, shuffling, concatenation, …). Internally each dataset stores, per
+stimulus:
+
+* the **stimulus** presented (a spectrogram `(1, F, T)`, or a raw waveform
+  `(1, T_audio)` when `return_waveform=True`);
+* the **trial-resolved responses**, time-aligned to the stimulus;
+* per-stimulus and per-neuron **metadata** dicts (`stim_meta`, `nrn_meta`).
+
+Datasets are sparse and ragged — neurons may not all hear every stimulus, and
+stimuli vary in duration and repeat count. deepSTRF handles this with **NaN
+sentinels** for missing (stimulus, neuron) trials and zero-/NaN-padding at
+collate time. The full contract — the canonical `(B, N, R, T)` response shape,
+the NaN-sentinel rules, and the bidirectional neuron/stimulus selection — is
+documented in {doc}`data_paradigm`. Read that before touching any
+response-path code.
+
+## What a batch looks like
+
+Pair a dataset with `neural_collate` and a `DataLoader`. Each batch is a
+**dict** (not a positional tuple), so future per-trial variables can be added as
+new keys without breaking your unpacking:
+
+```python
+from torch.utils.data import DataLoader
+from deepSTRF.utils.data import neural_collate
+
+loader = DataLoader(ds, batch_size=8, collate_fn=neural_collate)
+
+for batch in loader:
+    stims      = batch['stims']        # (B, ..., T)        zero-padded, no NaN
+    responses  = batch['responses']    # (B, N, R, T)       NaN-padded
+    valid_mask = batch['valid_mask']   # (B, N, R, T) bool  ~responses.isnan()
+    metas      = batch['stim_meta']    # length-B list of per-stim dicts
+    ...
+```
+
+Indexing the dataset directly (`ds[i]`) returns the same keys for a single item.
+`CCmax` / `TTRC`-style normalisation is **not** precomputed at the dataloader
+boundary — it is derived on demand from `responses` by the metrics in
+{doc}`metrics_paradigm` (e.g. `normalized_corrcoef`).
+
+## Selecting neurons and stimuli
+
+The set of units (and stimuli) to work with is chosen through the selection API
+rather than a fixed constructor argument — by default all neurons are selected.
+The selection drives both `len(ds)` and iteration. See {doc}`data_paradigm` for
+the exact semantics; the common entry points are:
+
+```python
+ds.select_population([0, 1, 2])                          # by index (or a single int)
+ds.select_pop_by_nrn_attr("area", "Field_L")            # by metadata label
+ds.select_pop_by_nrn_predicate(lambda n: n.get("snr", 0) > 0.5)   # by threshold
+ds.select_stims_by_attr("type", "human_speech")         # restrict the stimulus set
+```

docs/_source/md/data_paradigm.md

@@ -219,7 +219,8 @@ NaN-aware by default. The training loop therefore stays simple:
 ```python
 from deepSTRF.metrics import mse_loss, corrcoef, normalized_corrcoef
 
-for stims, responses, valid_mask, stim_metas in loader:
+for batch in loader:                                      # batch is a dict with keys
+    stims, responses = batch['stims'], batch['responses'] #   'stims','responses','valid_mask','stim_meta'
     pred    = model(stims)                                # (B, N, 1, T_max)
     gt_psth = responses.nanmean(dim=2, keepdim=True)      # (B, N, 1, T_max)
 

docs/_source/md/fitter.md

@@ -21,9 +21,9 @@ canonical path is the three-line loop documented in
 `metrics_paradigm.md` §7:
 
 ```python
-for stims, responses, valid_mask, stim_metas in loader:
-    pred = model(stims)                                    # (B, N, 1, T)
-    loss = mse_loss(pred, responses)                       # auto-PSTH inside
+for batch in loader:                                       # batch is a dict
+    pred = model(batch['stims'])                           # (B, N, 1, T)
+    loss = mse_loss(pred, batch['responses'])              # auto-PSTH inside
     loss.backward(); optimizer.step(); optimizer.zero_grad()
 ```
 
@@ -553,11 +553,12 @@ The legacy module exposes `optimize_multiple_seeds`,
 | `set_random_seed`            | `deepSTRF.training.set_random_seed` (kept as deprecated re-export)  |
 
 Behavioral differences worth flagging:
-- The legacy code unpacks `(spectrogram, responses, ccmax, ttrc)`. The
-  new code unpacks `(stims, responses, valid_mask, stim_metas)` from
-  `neural_collate`. `ccmax` and `ttrc` are no longer dataloader-side
-  pre-computed tensors — they are computed on demand by
-  `normalized_corrcoef` from raw `responses`.
+- The legacy code unpacks a `(spectrogram, responses, ccmax, ttrc)`
+  tuple. `neural_collate` now yields a **dict** with keys `'stims'`,
+  `'responses'`, `'valid_mask'`, `'stim_meta'` (read fields by key:
+  `batch['stims']`, `batch['responses']`, …). `ccmax` and `ttrc` are no
+  longer dataloader-side pre-computed tensors — they are computed on demand
+  by `normalized_corrcoef` from raw `responses`.
 - The legacy code calls `prediction.squeeze(-2)` to drop the R-axis
   before metrics. The new metrics expect `(B, N, 1, T)` per the model
   paradigm — no squeeze.

docs/_source/md/metrics_paradigm.md

@@ -493,7 +493,8 @@ Mirroring `data_paradigm.md` §6, the canonical loop now looks like:
 ```python
 from deepSTRF.metrics import mse_loss, corrcoef, normalized_corrcoef
 
-for stims, responses, valid_mask, stim_metas in loader:
+for batch in loader:                                             # batch is a dict
+    stims, responses = batch['stims'], batch['responses']
     pred = model(stims)                                          # (B, N, 1, T)
 
     loss = mse_loss(pred, responses)                             # auto-PSTH inside

examples/alice_eeg_tutorial.ipynb

@@ -195,17 +195,17 @@
     "    for ep in range(max_epochs):\n",
     "        model.train()\n",
     "        ep_losses = []\n",
-    "        for stims, responses, _, _ in train_loader:\n",
-    "            stims, responses = stims.to(device), responses.to(device)\n",
+    "        for batch in train_loader:\n",
+    "            stims, responses = batch['stims'].to(device), batch['responses'].to(device)\n",
     "            pred = model(stims)\n",
     "            loss = mse_loss(pred, responses)\n",
     "            opt.zero_grad(); loss.backward(); opt.step()\n",
     "            ep_losses.append(loss.item())\n",
     "        model.eval()\n",
     "        with torch.no_grad():\n",
     "            v = []\n",
-    "            for stims, responses, _, _ in val_loader:\n",
-    "                stims, responses = stims.to(device), responses.to(device)\n",
+    "            for batch in val_loader:\n",
+    "                stims, responses = batch['stims'].to(device), batch['responses'].to(device)\n",
     "                pred = model(stims); gt = responses.nanmean(dim=2, keepdim=True)\n",
     "                v.append(corrcoef(pred, gt, reduction='none').cpu())\n",
     "            val_cc = torch.stack(v).nanmean().item()\n",
@@ -220,8 +220,8 @@
     "    model.eval()\n",
     "    with torch.no_grad():\n",
     "        cc, ve = [], []\n",
-    "        for stims, responses, _, _ in test_loader:\n",
-    "            stims, responses = stims.to(device), responses.to(device)\n",
+    "        for batch in test_loader:\n",
+    "            stims, responses = batch['stims'].to(device), batch['responses'].to(device)\n",
     "            pred = model(stims); gt = responses.nanmean(dim=2, keepdim=True)\n",
     "            cc.append(corrcoef(pred, gt, reduction='none').cpu())\n",
     "            ve.append(fve(pred, gt, reduction='none').cpu())\n",