Qiu et al. 2026 Dataset

Mouse retinal ganglion cell axonal bouton responses to natural movies, recorded with chronic two-photon calcium imaging in the superficial layers of the superior colliculus of awake mice, originally published in Qiu et al. (2026): The functional organization of retinal input to the mouse superior colliculus, bioRxiv.

Dataset: huggingface.co/datasets/open-retina/open-retina — ten sessions from three mice (~17k boutons after quality masking), one zip per session under franke_lab/qiu_2026/.

Train on it with uv run openretina train --config-name qiu_2026_core_readout.

Dataset layout

Unlike the other datasets shipped with openretina, qiu_2026 arrives in the Sensorium FileTree format: one .npy per trial at data/<stream>/{i}.npy for the four streams videos, responses, behavior and pupil_center, positionally aligned to the rows of meta/trials/*.npy and NaN-padded at the end to a length that is uniform within a session but differs between them (300, 450 or 960 frames); a clip trial's valid portion is always its first 300 frames, i.e. 10 s at 30 Hz. The trial helpers hold the shared trimming, splitting and session-discovery logic that keeps all four streams frame-aligned.

Four things set this dataset apart from the others:

Behavior is folded into the input channels. The model input is (3, T, 36, 64): one grayscale video channel plus two behavior channels (pupil size and locomotion) broadcast across every pixel. The video channel is z-scored with the session's shipped scalar statistics, each behavior channel with its own, so the norm_mean/norm_std on the returned MoviesTrainTestSplit describe the video channel only.
Targets are raw calcium fluorescence, masked per session with the shipped neurons_fluor_good array (an index array, not a boolean mask). spike_inference selects the transform: "raw" leaves the trace untouched, "subtract_min" shifts each neuron by its train-derived minimum to give the Poisson loss a non-negative target, and "cascade" is reserved for inferred spike rates and not yet implemented.
Pupil center is loaded separately by pupil. It is not a model target: it feeds the shifter, which corrects each neuron's readout center for eye movements. Because openretina's shared DataPoint carries only (inputs, targets), this subpackage defines its own QiuDataPoint and QiuMovieDataSet rather than changing the shared one.
Test clips are repeated within a session, but the design is not the same in every session: six sessions share one set of six clips at 14–20 repeats each, while the four with the least training data were tested on 94–96 distinct clips at 1–10 repeats (supersets of those same six). Repeats are keyed by condition_hash. The video is identical across a condition's repeats but the behavior and eye position are not, and the model is nonlinear in both, so with per_repeat_test_behavior=True the test dataloader puts the repeats in the batch dimension and averages the predictions rather than the inputs. Non-clip trials in the same test tier — chirp/moving bar movies in five sessions, monet2/trippy noise in the other four — are excluded by the stimulus_type="clip" filter. The dataset page has the per-session breakdown.

Dataloaders

dataloaders

qiu_2026 dataloaders that carry pupil_center alongside inputs/targets.

open-retina's shared DataPoint is a 2-field (inputs, targets) namedtuple; extending it would break every existing dataset (see the integration plan §0.4). Instead this module defines a qiu-local QiuDataPoint and a QiuMovieDataSet that slices a per-frame pupil trace with the same chunking as the movie/responses, injected into :func:get_movie_dataloader via its dataset_cls/extra_arrays hook.

qiu_2026_dataloaders mirrors :func:multiple_movies_dataloaders but (a) splits the pupil trace identically to the movie via :func:generate_movie_splits, (b) honors each session's curated validation split (the validation_clip_indices stashed in session_kwargs) so validation trials are carved out of training rather than chosen at random, (c) evaluates each test condition as one whole clip, and (d) by default puts that condition's repeats in the batch dimension so the model runs once per repeat with that repeat's own behavior, and the predictions -- not the inputs -- are averaged.

QiuMovieDataSet

QiuMovieDataSet(
    movies,
    responses,
    roi_ids,
    roi_coords,
    group_assignment,
    split,
    chunk_size,
    pupil_center,
)

Bases: MovieDataSet

MovieDataSet that additionally carries a per-frame pupil_center (2, T) trace.

pupil_center is sliced on its time axis (axis 1) with the exact chunk the movie (axis 1) and responses (axis 0) are sliced with, and returned as the third field of a :class:QiuDataPoint.

Source code in openretina/data_io/qiu_2026/dataloaders.py

def __init__(
    self,
    movies,
    responses,
    roi_ids,
    roi_coords,
    group_assignment,
    split,
    chunk_size,
    pupil_center,
):
    super().__init__(movies, responses, roi_ids, roi_coords, group_assignment, split, chunk_size)
    self.pupil_center = torch.as_tensor(pupil_center, dtype=torch.float)  # (2, T)

QiuRepeatTestDataSet

QiuRepeatTestDataSet(
    movies,
    responses,
    roi_ids,
    roi_coords,
    group_assignment,
    split,
    chunk_size,
    behavior_by_repeat,
    pupil_by_repeat,
)

Bases: MovieDataSet

Test dataset whose index dimension is the repeat, not a time chunk.

A test condition is one clip presented 15-20 times. The movie is identical on every repeat but the animal's behavior is not, so each item re-assembles the (C, T, H, W) input from one shared video channel plus that repeat's own behavior channels, and pairs it with that repeat's own pupil trace. The target is the trial-averaged response, the same for every item: the model sees R inputs and :meth:BaseCoreReadout.test_step averages the R predictions before scoring them, which is what the trial-averaged response actually is -- the average over repeats of what the retina did while the animal behaved as it did on that repeat.

Storing R assembled input tensors instead would cost ~R times the test movie memory for a video that never changes (several GB across the full dataset), so the video is kept once and broadcast here.

Source code in openretina/data_io/qiu_2026/dataloaders.py

def __init__(
    self,
    movies,
    responses,
    roi_ids,
    roi_coords,
    group_assignment,
    split,
    chunk_size,
    behavior_by_repeat,
    pupil_by_repeat,
):
    super().__init__(movies, responses, roi_ids, roi_coords, group_assignment, split, chunk_size)
    self.behavior_by_repeat = torch.as_tensor(behavior_by_repeat, dtype=torch.float)  # (R, n_beh, T)
    self.pupil_by_repeat = torch.as_tensor(pupil_by_repeat, dtype=torch.float)  # (R, 2, T)
    # Channel 0 of the repeat-averaged test movie is the video, identical on every repeat. A view.
    self.video = self.samples[0][:1]  # (1, T, H, W)

    n_repeats, _, n_time = self.behavior_by_repeat.shape
    if self.pupil_by_repeat.shape[0] != n_repeats:
        raise ValueError(
            f"Behavior has {n_repeats} repeats but pupil has {self.pupil_by_repeat.shape[0]}; "
            "both must come from the same test condition."
        )
    if n_time != self.video.shape[1] or self.pupil_by_repeat.shape[2] != self.video.shape[1]:
        raise ValueError(
            f"Per-repeat traces are not frame-aligned with the movie: behavior {n_time} frames, "
            f"pupil {self.pupil_by_repeat.shape[2]}, movie {self.video.shape[1]}."
        )

qiu_2026_dataloaders

qiu_2026_dataloaders(
    neuron_data_dictionary: dict[
        str, ResponsesTrainTestSplit
    ],
    movies_dictionary: dict[str, Any],
    pupil_dictionary: dict[str, dict[str, Any]],
    train_chunk_size: int = 50,
    batch_size: int = 32,
    seed: int = 42,
    clip_length: int = 300,
    num_val_clips: int = 10,
    val_clip_indices: list[int] | None = None,
    allow_over_boundaries: bool = False,
    release_movies: bool = False,
    per_repeat_test_behavior: bool = True,
) -> dict[str, dict[str, DataLoader]]

Build dict[split][session] -> DataLoader yielding :class:QiuDataPoint batches.

Validation-clip selection per session, in priority order: the explicit val_clip_indices argument, then the session's curated session_kwargs["validation_clip_indices"], then a random draw of num_val_clips. The curated path is the intended one — it removes the dataset's designated validation trials from the training subset. Each test condition is evaluated as a single whole clip.

PARAMETER DESCRIPTION

per_repeat_test_behavior

Evaluate each test condition by feeding every repeat's behavior and pupil trace through the model separately and averaging the R predictions, instead of averaging the R inputs and running the model once. The video is the same clip on every repeat, but the behavior channels and the pupil centre are not, and the model is nonlinear in both, so f(video, mean_r b_r) != mean_r f(video, b_r). Test dataloaders then yield :class:QiuRepeatDataPoint batches of R repeats; train and validation are unaffected either way. Set to False to reproduce the repeat-averaged-input numbers.

TYPE: bool DEFAULT: True

release_movies

Destructive. Drop each session's entry from movies_dictionary as soon as its splits have been taken, so already-consumed source movies are freed while the loop is still running. The full dataset is ~28 GB of movies and the splits duplicate all of it, so without this the caller holds both copies of every session at once; releasing as we go is what keeps peak RSS down (freeing afterwards would only help the steady state). The caller's dictionary is left EMPTY, so anything needing the raw movies — notably :func:~openretina.data_io.base.compute_data_info, which reads .train.shape — must run before this function, as openretina/cli/train.py and the qiu notebooks both do.

Only frees anything if the caller passes its own dictionary directly. Going through hydra.utils.instantiate(cfg.dataloader, movies_dictionary=...) hands this function an OmegaConf-rebuilt copy of the container and of each wrapper, so dropping entries from it leaves the caller's references — and therefore the arrays — alive. Instantiate the builder with _partial_=True and call it normally instead; both CLI entry points do.

TYPE: bool DEFAULT: False

Source code in openretina/data_io/qiu_2026/dataloaders.py

def qiu_2026_dataloaders(
    neuron_data_dictionary: dict[str, ResponsesTrainTestSplit],
    movies_dictionary: dict[str, Any],
    pupil_dictionary: dict[str, dict[str, Any]],
    train_chunk_size: int = 50,
    batch_size: int = 32,
    seed: int = 42,
    clip_length: int = 300,
    num_val_clips: int = 10,
    val_clip_indices: list[int] | None = None,
    allow_over_boundaries: bool = False,
    release_movies: bool = False,
    per_repeat_test_behavior: bool = True,
) -> dict[str, dict[str, DataLoader]]:
    """Build ``dict[split][session] -> DataLoader`` yielding :class:`QiuDataPoint` batches.

    Validation-clip selection per session, in priority order: the explicit ``val_clip_indices`` argument, then
    the session's curated ``session_kwargs["validation_clip_indices"]``, then a random draw of
    ``num_val_clips``. The curated path is the intended one — it removes the dataset's designated validation
    trials from the training subset. Each test condition is evaluated as a single whole clip.

    Args:
        per_repeat_test_behavior: Evaluate each test condition by feeding every repeat's behavior and pupil
            trace through the model separately and averaging the R *predictions*, instead of averaging the
            R inputs and running the model once. The video is the same clip on every repeat, but the
            behavior channels and the pupil centre are not, and the model is nonlinear in both, so
            ``f(video, mean_r b_r) != mean_r f(video, b_r)``. Test dataloaders then yield
            :class:`QiuRepeatDataPoint` batches of R repeats; train and validation are unaffected either
            way. Set to ``False`` to reproduce the repeat-averaged-input numbers.

        release_movies: **Destructive.** Drop each session's entry from ``movies_dictionary`` as soon as its
            splits have been taken, so already-consumed source movies are freed while the loop is still
            running. The full dataset is ~28 GB of movies and the splits duplicate all of it, so without
            this the caller holds both copies of every session at once; releasing as we go is what keeps
            peak RSS down (freeing afterwards would only help the steady state). The caller's dictionary is
            left EMPTY, so anything needing the raw movies — notably
            :func:`~openretina.data_io.base.compute_data_info`, which reads ``.train.shape`` — must run
            before this function, as ``openretina/cli/train.py`` and the qiu notebooks both do.

            Only frees anything if the caller passes its own dictionary directly. Going through
            ``hydra.utils.instantiate(cfg.dataloader, movies_dictionary=...)`` hands this function an
            OmegaConf-rebuilt copy of the container and of each wrapper, so dropping entries from it
            leaves the caller's references — and therefore the arrays — alive. Instantiate the builder
            with ``_partial_=True`` and call it normally instead; both CLI entry points do.
    """
    assert set(neuron_data_dictionary) == set(movies_dictionary) == set(pupil_dictionary), (
        "neuron_data_dictionary, movies_dictionary and pupil_dictionary must share the same session keys."
    )

    dataloaders: dict[str, Any] = collections.defaultdict(dict)
    for session_key, session_data in tqdm(neuron_data_dictionary.items(), desc="Creating qiu movie dataloaders"):
        num_clips = movies_dictionary[session_key].train.shape[1] // clip_length

        if val_clip_indices is not None:
            val_clip_idx = list(val_clip_indices)
        elif "validation_clip_indices" in session_data.session_kwargs:
            val_clip_idx = list(session_data.session_kwargs["validation_clip_indices"])
        else:
            rnd = np.random.RandomState(seed)
            val_clip_idx = list(rnd.choice(num_clips, num_val_clips, replace=False))

        movie_train_subset, movie_val, movie_test_dict = generate_movie_splits(
            movies_dictionary[session_key].train,
            movies_dictionary[session_key].test_dict,
            val_clip_idc=val_clip_idx,
            num_clips=num_clips,
            clip_length=clip_length,
        )
        # Read before the release_movies pop below, which empties the caller's dictionary.
        behavior_by_repeat_dict = getattr(movies_dictionary[session_key], "test_behavior_by_repeat_dict", {}) or {}
        # Last use of the raw movie for this session: everything below works off the splits above.
        # Safe to mutate mid-loop because we iterate `neuron_data_dictionary`, not this dictionary.
        if release_movies:
            movies_dictionary.pop(session_key, None)

        pupil_train_subset, pupil_val, pupil_test_dict = _split_pupil(
            pupil_dictionary[session_key]["train"],
            pupil_dictionary[session_key]["test_dict"],
            val_clip_idx,
            num_clips,
            clip_length,
        )
        neuron_data = NeuronDataSplit(
            responses=session_data, val_clip_idx=val_clip_idx, num_clips=num_clips, clip_length=clip_length
        )

        for fold, movie, pupil, chunk_size in [
            ("train", movie_train_subset, pupil_train_subset, train_chunk_size),
            ("validation", movie_val, pupil_val, clip_length),
        ]:
            dataloaders[fold][session_key] = get_movie_dataloader(
                movie=movie,
                responses=neuron_data.response_dict[fold],
                split=fold,
                chunk_size=chunk_size,
                batch_size=batch_size,
                scene_length=clip_length,
                allow_over_boundaries=allow_over_boundaries,
                dataset_cls=QiuMovieDataSet,
                extra_arrays={"pupil_center": pupil},
            )

        pupil_by_repeat_dict = pupil_dictionary[session_key].get("test_by_repeat_dict", {})
        if per_repeat_test_behavior:
            missing = [
                name
                for name in movie_test_dict
                if name not in behavior_by_repeat_dict or name not in pupil_by_repeat_dict
            ]
            if missing:
                raise ValueError(
                    f"Session {session_key}: per_repeat_test_behavior=True, but no per-repeat behavior/pupil "
                    f"was loaded for test conditions {missing}. Load stimuli and pupil with the qiu_2026 "
                    f"loaders (both populate them), or set dataloader.per_repeat_test_behavior=false to "
                    f"evaluate with repeat-averaged inputs."
                )

        # Each test condition is a single whole clip (300 or 450 frames): evaluate it in one chunk.
        for name, movie in movie_test_dict.items():
            test_chunk_size = movie.shape[1]
            if per_repeat_test_behavior:
                n_repeats = int(behavior_by_repeat_dict[name].shape[0])
                dataloaders[name][session_key] = get_movie_dataloader(
                    movie=movie,
                    responses=neuron_data.response_dict_test[name],
                    split="test",
                    chunk_size=test_chunk_size,
                    # One batch holding every repeat, so test_step can average the predictions over it.
                    batch_size=n_repeats,
                    # Explicit: the default start indices walk the time axis, but here the dataset is
                    # indexed by repeat, so leaving it to the default would yield only the first repeat.
                    start_indices=list(range(n_repeats)),
                    scene_length=clip_length,
                    allow_over_boundaries=True,
                    dataset_cls=QiuRepeatTestDataSet,
                    extra_arrays={
                        "behavior_by_repeat": behavior_by_repeat_dict[name],
                        "pupil_by_repeat": pupil_by_repeat_dict[name],
                    },
                )
            else:
                test_batch_size = _compute_test_batch_size(batch_size, train_chunk_size, test_chunk_size)
                dataloaders[name][session_key] = get_movie_dataloader(
                    movie=movie,
                    responses=neuron_data.response_dict_test[name],
                    split="test",
                    chunk_size=test_chunk_size,
                    batch_size=test_batch_size,
                    scene_length=clip_length,
                    allow_over_boundaries=True,
                    dataset_cls=QiuMovieDataSet,
                    extra_arrays={"pupil_center": pupil_test_dict[name]},
                )

    return dataloaders

Stimuli

stimuli

qiu_2026 stimulus loader.

Builds a (C, T, H, W) input where C = 1 video + len(behavior_channels) behavior channels: the grayscale movie plus the selected behavior traces (pupil size + locomotion) broadcast spatially across every pixel. The video channel is z-scored with the session's shipped video statistics (a shared scalar); each behavior channel is z-scored with its own shipped statistics inside the tensor, so the norm_mean/norm_std stored on the container describe only the video channel.

Repeated test stimuli additionally keep their behavior per repeat (test_behavior_by_repeat_dict), since the movie is identical across a condition's repeats but the animal's behavior is not.

load_stimuli_for_session

load_stimuli_for_session(
    session_path: str | PathLike,
    *,
    behavior_channels: tuple[int, ...] = BEHAVIOR_CHANNELS,
    stimulus_type: str = "clip",
    video_deviation_tolerance: float = 0.1,
) -> MoviesTrainTestSplit

Build a :class:MoviesTrainTestSplit for one qiu_2026 session.

The continuous train movie concatenates all train+validation trials (each trimmed to its valid length) in file-index order. test_dict holds one clip per condition_hash, averaged over the condition's repeats in both the video and the behavior channels.

Averaging the video is near enough to a no-op: the repeats are the same clip, differing at most in which of the pipeline's two temporal-resampling renderings of it they got (see :func:_warn_on_video_deviation). Averaging the behavior is not — the animal behaves differently on every repeat, and the model is nonlinear in those channels. So test_behavior_by_repeat_dict additionally keeps each repeat's behavior as (repeats, n_behavior, T), letting the dataloader feed them through the model one at a time and average the predictions instead (see qiu_2026_dataloaders).

Only trials matching stimulus_type (see :func:~openretina.data_io.qiu_2026.trials.read_stimulus_type) are included; some sessions' test tier additionally contains non-clip functional-characterization trials with per-repeat frame-count jitter that would otherwise break the repeat-stacking below.

PARAMETER	DESCRIPTION
`video_deviation_tolerance`	Warn when a condition's repeats differ by more than this mean absolute deviation, in units of the video std. The default is calibrated for `stimulus_type="clip"`. See :func:`_warn_on_video_deviation`. TYPE: `float` DEFAULT: `0.1`

Source code in openretina/data_io/qiu_2026/stimuli.py

def load_stimuli_for_session(
    session_path: str | os.PathLike,
    *,
    behavior_channels: tuple[int, ...] = C.BEHAVIOR_CHANNELS,
    stimulus_type: str = "clip",
    video_deviation_tolerance: float = 0.1,
) -> MoviesTrainTestSplit:
    """Build a :class:`MoviesTrainTestSplit` for one qiu_2026 session.

    The continuous train movie concatenates all train+validation trials (each trimmed to its valid
    length) in file-index order. ``test_dict`` holds one clip per ``condition_hash``, averaged over the
    condition's repeats in both the video and the behavior channels.

    Averaging the video is near enough to a no-op: the repeats are the same clip, differing at most in
    which of the pipeline's two temporal-resampling renderings of it they got (see
    :func:`_warn_on_video_deviation`). Averaging the *behavior* is not — the animal behaves differently on
    every repeat, and the model is nonlinear in those channels. So ``test_behavior_by_repeat_dict``
    additionally keeps each repeat's behavior as ``(repeats, n_behavior, T)``, letting the dataloader feed
    them through the model one at a time and average the *predictions* instead (see
    ``qiu_2026_dataloaders``).

    Only trials matching ``stimulus_type`` (see
    :func:`~openretina.data_io.qiu_2026.trials.read_stimulus_type`) are included; some sessions' ``test``
    tier additionally contains non-clip functional-characterization trials with per-repeat frame-count
    jitter that would otherwise break the repeat-stacking below.

    Args:
        video_deviation_tolerance: Warn when a condition's repeats differ by more than this mean
            absolute deviation, in units of the video std. The default is calibrated for
            ``stimulus_type="clip"``. See :func:`_warn_on_video_deviation`.
    """
    session_path = resolve_session_path(session_path)
    tiers = read_tiers(session_path)
    condition_hash = read_condition_hash(session_path)
    stimulus_types = read_stimulus_type(session_path)
    video_mean, video_std, behavior_mean, behavior_std = _read_norm_stats(session_path)

    def load_video(i: int) -> tuple[np.ndarray, int]:
        """Return the trial's normalized ``(1, T, H, W)`` video and its valid frame count."""
        video = load_trial_array(session_path, "videos", i)
        n_valid = valid_length(video, STREAM_TIME_AXIS["videos"])
        video = trim_time(video, STREAM_TIME_AXIS["videos"], n_valid)
        return _normalize_video(video, video_mean, video_std), n_valid

    def load_behavior(i: int, n_valid: int) -> np.ndarray:
        """Return the trial's normalized ``(n_behavior, T)`` behavior, trimmed to ``n_valid`` frames."""
        behavior = load_trial_array(session_path, "behavior", i)
        n_behavior_valid = valid_length(behavior, STREAM_TIME_AXIS["behavior"])
        if n_behavior_valid != n_valid:
            raise ValueError(
                f"Trial {i} of session {session_path.name}: behavior has {n_behavior_valid} valid frames "
                f"but the video has {n_valid}. The streams must be frame-aligned to stack repeats."
            )
        behavior = trim_time(behavior, STREAM_TIME_AXIS["behavior"], n_valid)
        return _normalize_behavior(behavior, behavior_mean, behavior_std, behavior_channels)

    def build_trial(i: int) -> np.ndarray:
        video_norm, n_valid = load_video(i)
        return _assemble_input_tensor(video_norm, load_behavior(i, n_valid))

    train = np.concatenate([build_trial(i) for i in train_val_indices(tiers, stimulus_types, stimulus_type)], axis=1)

    test_dict: dict[str, np.ndarray] = {}
    test_behavior_by_repeat_dict: dict[str, np.ndarray] = {}
    for condition, indices in test_conditions(tiers, condition_hash, stimulus_types, stimulus_type).items():
        videos: list[np.ndarray] = []
        behaviors: list[np.ndarray] = []
        for i in indices:
            video_norm, n_valid = load_video(i)
            if videos and video_norm.shape != videos[0].shape:
                raise ValueError(
                    f"Session {session_path.name}, condition {condition}: trial {i} has video shape "
                    f"{video_norm.shape} but trial {indices[0]} has {videos[0].shape}; repeats of one "
                    f"condition must share a frame count to be stacked."
                )
            videos.append(video_norm)
            behaviors.append(load_behavior(i, n_valid))

        behavior_by_repeat = np.stack(behaviors, axis=0)  # (R, n_behavior, T)
        _warn_on_video_deviation(session_path, condition, indices, videos, video_deviation_tolerance)
        test_dict[condition] = _assemble_input_tensor(np.mean(videos, axis=0), behavior_by_repeat.mean(axis=0))
        test_behavior_by_repeat_dict[condition] = behavior_by_repeat.astype(np.float32)

    return MoviesTrainTestSplit(
        train=train,
        test_dict=test_dict,
        stim_id="qiu_2026",
        norm_mean=video_mean,
        norm_std=video_std,
        test_behavior_by_repeat_dict=test_behavior_by_repeat_dict,
    )

load_all_stimuli

load_all_stimuli(
    base_data_path: str | PathLike,
    *,
    behavior_channels: tuple[int, ...] = BEHAVIOR_CHANNELS,
    stimulus_type: str = "clip",
    sessions: list[str] | None = None,
    video_deviation_tolerance: float = 0.1,
) -> dict[str, MoviesTrainTestSplit]

Load every discovered qiu_2026 session into a {session_key: MoviesTrainTestSplit} dict.

Source code in openretina/data_io/qiu_2026/stimuli.py

def load_all_stimuli(
    base_data_path: str | os.PathLike,
    *,
    behavior_channels: tuple[int, ...] = C.BEHAVIOR_CHANNELS,
    stimulus_type: str = "clip",
    sessions: list[str] | None = None,
    video_deviation_tolerance: float = 0.1,
) -> dict[str, MoviesTrainTestSplit]:
    """Load every discovered qiu_2026 session into a ``{session_key: MoviesTrainTestSplit}`` dict."""
    base_data_path = Path(get_local_file_path(str(base_data_path)))
    stimuli_all_sessions: dict[str, MoviesTrainTestSplit] = {}
    for path in tqdm(discover_sessions(base_data_path, sessions), desc="qiu_2026 stimuli"):
        stimuli_all_sessions[session_key(path)] = load_stimuli_for_session(
            path,
            behavior_channels=behavior_channels,
            stimulus_type=stimulus_type,
            video_deviation_tolerance=video_deviation_tolerance,
        )
    return stimuli_all_sessions

Responses

responses

qiu_2026 response loader.

Loads raw calcium fluorescence, applies the per-session neurons_fluor_good index mask, trims NaN padding, and assembles a continuous train trace plus per-condition_hash test dictionaries (both the trial-averaged test_dict and the per-repeat test_by_trial_dict).

spike_inference controls the response transform: - "raw": fluorescence untouched (default; used by fast smoke tests). - "subtract_min": subtract each neuron's train-derived minimum → non-negative, trainable interim target for the Poisson loss until CASCADE is wired in. - "cascade": CASCADE inferred spike rates (Rupprecht 2021) — not yet implemented.

load_responses_for_session

load_responses_for_session(
    session_path: str | PathLike,
    good_idx: ndarray | None = None,
    *,
    spike_inference: SpikeInference = "raw",
    stimulus_type: str = "clip",
) -> ResponsesTrainTestSplit

Build a :class:ResponsesTrainTestSplit for one qiu_2026 session.

good_idx is the neurons_fluor_good index array; None keeps all neurons. Only trials matching stimulus_type (see :func:~openretina.data_io.qiu_2026.trials.read_stimulus_type) are included; some sessions' test tier additionally contains non-clip functional-characterization trials with per-repeat frame-count jitter that would otherwise break the repeat-stacking below.

Source code in openretina/data_io/qiu_2026/responses.py

def load_responses_for_session(
    session_path: str | os.PathLike,
    good_idx: np.ndarray | None = None,
    *,
    spike_inference: SpikeInference = "raw",
    stimulus_type: str = "clip",
) -> ResponsesTrainTestSplit:
    """Build a :class:`ResponsesTrainTestSplit` for one qiu_2026 session.

    ``good_idx`` is the ``neurons_fluor_good`` index array; ``None`` keeps all neurons. Only trials
    matching ``stimulus_type`` (see
    :func:`~openretina.data_io.qiu_2026.trials.read_stimulus_type`) are included; some sessions' ``test``
    tier additionally contains non-clip functional-characterization trials with per-repeat frame-count
    jitter that would otherwise break the repeat-stacking below.
    """
    session_path = resolve_session_path(session_path)
    tiers = read_tiers(session_path)
    condition_hash = read_condition_hash(session_path)
    stimulus_types = read_stimulus_type(session_path)
    time_axis = STREAM_TIME_AXIS["responses"]

    def load_trimmed(i: int) -> np.ndarray:
        responses = load_trial_array(session_path, "responses", i)  # (N, T)
        if good_idx is not None:
            responses = responses[good_idx]
        n_valid = valid_length(responses, time_axis)
        return trim_time(responses, time_axis, n_valid).astype(np.float32)

    train_indices = train_val_indices(tiers, stimulus_types, stimulus_type)
    train = np.concatenate([load_trimmed(i) for i in train_indices], axis=1)  # (N, T_total)
    offset = _spike_offset(train, spike_inference)
    train = (train - offset).astype(np.float32)

    test_dict: dict[str, np.ndarray] = {}
    test_by_trial_dict: dict[str, np.ndarray] = {}
    for condition, indices in test_conditions(tiers, condition_hash, stimulus_types, stimulus_type).items():
        by_trial = np.stack([load_trimmed(i) for i in indices], axis=0)  # (repeats, N, T_cond)
        by_trial = (by_trial - offset[None]).astype(np.float32)
        test_by_trial_dict[condition] = by_trial
        test_dict[condition] = by_trial.mean(axis=0).astype(np.float32)  # (N, T_cond)

    session_kwargs = {
        "validation_clip_indices": validation_clip_indices(tiers, stimulus_types, stimulus_type),
    }
    return ResponsesTrainTestSplit(
        train=train,
        test_dict=test_dict,
        test_by_trial_dict=test_by_trial_dict,
        stim_id="qiu_2026",
        session_kwargs=session_kwargs,
    )

load_all_responses

load_all_responses(
    base_data_path: str | PathLike,
    *,
    apply_quality_mask: bool = True,
    spike_inference: SpikeInference = "raw",
    stimulus_type: str = "clip",
    sessions: list[str] | None = None,
) -> dict[str, ResponsesTrainTestSplit]

Load every discovered qiu_2026 session into a {session_key: ResponsesTrainTestSplit} dict.

Source code in openretina/data_io/qiu_2026/responses.py

def load_all_responses(
    base_data_path: str | os.PathLike,
    *,
    apply_quality_mask: bool = True,
    spike_inference: SpikeInference = "raw",
    stimulus_type: str = "clip",
    sessions: list[str] | None = None,
) -> dict[str, ResponsesTrainTestSplit]:
    """Load every discovered qiu_2026 session into a ``{session_key: ResponsesTrainTestSplit}`` dict."""
    base_data_path = Path(get_local_file_path(str(base_data_path)))
    session_paths = discover_sessions(base_data_path, sessions)
    masks = discover_quality_masks(base_data_path) if apply_quality_mask else {}

    responses_all_sessions: dict[str, ResponsesTrainTestSplit] = {}
    for path in tqdm(session_paths, desc="qiu_2026 responses"):
        key = session_key(path)
        good_idx = match_quality_mask(key, masks) if apply_quality_mask else None
        responses_all_sessions[key] = load_responses_for_session(
            path,
            good_idx=good_idx,
            spike_inference=spike_inference,
            stimulus_type=stimulus_type,
        )
    return responses_all_sessions

Pupil

pupil

qiu_2026 pupil-center loader.

Loads the (2, T) eye-position trace per trial, z-scores it with the shipped pupil statistics, trims NaN padding, and splits it with the same train/test logic used by the stimulus and response loaders so the trace stays frame-aligned. There is no container for pupil data (it is not a model target); the per-session dict returned here is consumed by the qiu dataloader, which slices it into the shifter input.

Layout mirrors the response split: train is the continuous train+validation trace; test_dict holds one trace per condition_hash, averaged over the condition's repeats to pair with the trial-averaged test responses, and test_by_repeat_dict keeps the individual repeats, which differ because the eye moves differently on every presentation of the same clip.

load_pupil_for_session

load_pupil_for_session(
    session_path: str | PathLike,
    *,
    stimulus_type: str = "clip",
) -> dict[str, ndarray | dict[str, ndarray]]

Return the train trace plus both test layouts for one session.

{"train": (2, T_total), "test_dict": {condition: (2, T_cond)}, "test_by_repeat_dict": {condition: (repeats, 2, T_cond)}}.

Only trials matching stimulus_type (see :func:~openretina.data_io.qiu_2026.trials.read_stimulus_type) are included; some sessions' test tier additionally contains non-clip functional-characterization trials with per-repeat frame-count jitter that would otherwise break the repeat-stacking below.

Source code in openretina/data_io/qiu_2026/pupil.py

def load_pupil_for_session(
    session_path: str | os.PathLike, *, stimulus_type: str = "clip"
) -> dict[str, np.ndarray | dict[str, np.ndarray]]:
    """Return the train trace plus both test layouts for one session.

    ``{"train": (2, T_total), "test_dict": {condition: (2, T_cond)},
    "test_by_repeat_dict": {condition: (repeats, 2, T_cond)}}``.

    Only trials matching ``stimulus_type`` (see
    :func:`~openretina.data_io.qiu_2026.trials.read_stimulus_type`) are included; some sessions' ``test``
    tier additionally contains non-clip functional-characterization trials with per-repeat frame-count
    jitter that would otherwise break the repeat-stacking below.
    """
    session_path = resolve_session_path(session_path)
    tiers = read_tiers(session_path)
    condition_hash = read_condition_hash(session_path)
    stimulus_types = read_stimulus_type(session_path)
    time_axis = STREAM_TIME_AXIS["pupil_center"]

    stat = session_path / "meta" / "statistics" / "pupil_center" / "all"
    pupil_mean = np.asarray(np.load(stat / "mean.npy")[..., 0], dtype=np.float32)[:, None]  # (2, 1)
    pupil_std = np.asarray(np.load(stat / "std.npy")[..., 0], dtype=np.float32)[:, None]

    def load_trimmed(i: int) -> np.ndarray:
        pupil = load_trial_array(session_path, "pupil_center", i).astype(np.float32)  # (2, T)
        n_valid = valid_length(pupil, time_axis)
        pupil = trim_time(pupil, time_axis, n_valid)
        return (pupil - pupil_mean) / pupil_std

    train = np.concatenate(
        [load_trimmed(i) for i in train_val_indices(tiers, stimulus_types, stimulus_type)], axis=1
    )  # (2, T_total)

    test_dict: dict[str, np.ndarray] = {}
    test_by_repeat_dict: dict[str, np.ndarray] = {}
    for condition, indices in test_conditions(tiers, condition_hash, stimulus_types, stimulus_type).items():
        repeats = np.stack([load_trimmed(i) for i in indices], axis=0)  # (repeats, 2, T_cond)
        # Eye position genuinely differs across repeats, so the average is not the trace the shifter
        # saw on any one of them. Both are kept: `test_dict` is the average (used when the dataloader
        # evaluates a condition as a single clip), `test_by_repeat_dict` keeps every repeat so the
        # per-repeat path can shift each prediction with its own trace before averaging predictions.
        test_by_repeat_dict[condition] = repeats.astype(np.float32)
        test_dict[condition] = repeats.mean(axis=0).astype(np.float32)

    return {"train": train, "test_dict": test_dict, "test_by_repeat_dict": test_by_repeat_dict}

load_all_pupil

load_all_pupil(
    base_data_path: str | PathLike,
    *,
    stimulus_type: str = "clip",
    sessions: list[str] | None = None,
) -> dict[str, dict[str, ndarray | dict[str, ndarray]]]

Load every discovered qiu_2026 session's pupil trace into a {session_key: {...}} dict.

Source code in openretina/data_io/qiu_2026/pupil.py

def load_all_pupil(
    base_data_path: str | os.PathLike,
    *,
    stimulus_type: str = "clip",
    sessions: list[str] | None = None,
) -> dict[str, dict[str, np.ndarray | dict[str, np.ndarray]]]:
    """Load every discovered qiu_2026 session's pupil trace into a ``{session_key: {...}}`` dict."""
    base_data_path = Path(get_local_file_path(str(base_data_path)))
    pupil_all_sessions: dict[str, dict[str, np.ndarray | dict[str, np.ndarray]]] = {}
    for path in tqdm(discover_sessions(base_data_path, sessions), desc="qiu_2026 pupil"):
        pupil_all_sessions[session_key(path)] = load_pupil_for_session(path, stimulus_type=stimulus_type)
    return pupil_all_sessions

Trial helpers

trials

Shared trial-splitting, NaN-trimming and session-discovery helpers for the qiu_2026 loaders.

Every per-trial array lives at data/<stream>/{i}.npy for i in 0..n_trials-1, aligned positionally to row i of every meta/trials/*.npy array (NOT to the value stored in trial_idx.npy, which carries the original DataJoint ids with gaps). NaN padding is strictly trailing (whole-frame NaN) and identical across the four streams for a given trial, so the same trim/split logic must be applied to all of them — hence these shared helpers, imported by stimuli.py, responses.py and pupil.py so their outputs stay frame-aligned.