.gitignore

@@ -2,3 +2,4 @@
 __pycache__/
 .DS_Store
 outputs/
+ai/

config.toml

@@ -21,3 +21,13 @@ ensemble_size = 30
 droprate = 0.05
 learning_rate = 0.0001
 num_epochs = 25
+deterministic = false # force deterministic cuDNN kernels (slower, exact reproduction)
+
+
+[evaluation]
+# Gaussian noise standard deviations applied to images during noisy evaluation
+# (step 6). 0.0 is the clean baseline and should be kept first.
+noise_levels = [0.0, 0.02, 0.05, 0.1, 0.2]
+# Monte-Carlo forward passes used to estimate Bayesian predictive/model
+# uncertainty (step 5/6). Ignored for the deterministic ensemble.
+mc_passes = 30

evaluation/__init__.py

@@ -0,0 +1,31 @@
+"""Evaluation output layer.
+
+This package owns the canonical on-disk format for model evaluations. Steps 4-6
+of the pipeline (evaluate normal / Bayesian / noisy) all produce data in the
+single schema defined in ``evaluation.schema`` so that step 7 (analysis) can be
+written as pure functions over the resulting netCDF files.
+"""
+
+from evaluation.schema import (
+    CLASS_NAMES,
+    MODEL_KIND_BAYESIAN,
+    MODEL_KIND_NORMAL,
+    SCHEMA_VERSION,
+    EvaluationAccumulator,
+    build_evaluation_dataset,
+    compute_uncertainty,
+    load_evaluation,
+    save_evaluation,
+)
+
+__all__ = [
+    "CLASS_NAMES",
+    "MODEL_KIND_BAYESIAN",
+    "MODEL_KIND_NORMAL",
+    "SCHEMA_VERSION",
+    "EvaluationAccumulator",
+    "build_evaluation_dataset",
+    "compute_uncertainty",
+    "load_evaluation",
+    "save_evaluation",
+]

evaluation/schema.py

@@ -0,0 +1,429 @@
+"""Canonical netCDF / xarray schema for pipeline evaluations (steps 4-6).
+
+This module is the single source of truth for how model outputs are stored on
+disk. Locking it down means every evaluator (normal, Bayesian, noisy) writes the
+exact same dimensions, variables, coordinates, and dtypes, so the analysis stage
+(step 7) only ever has to understand one format.
+
+--------------------------------------------------------------------------------
+Schema (xarray.Dataset)
+--------------------------------------------------------------------------------
+Dimensions
+    model         ensemble member index (0..M-1)
+    noise_level   Gaussian noise sigma applied to the image; ALWAYS present.
+                  The clean baseline is stored as noise_level == 0.0, so noised
+                  and un-noised evaluations share one file per model family.
+    sample        fixed test-set sample index (from a NON-shuffled loader)
+    class_name    output classes, ordered ["AD", "NL"]
+    mc_pass       (optional) raw Monte-Carlo draw index, Bayesian only
+
+Data variables
+    prob                (model, noise_level, sample, class_name) float32
+                        Mean predictive probability. For Bayesian models this is
+                        the mean over MC passes; for normal models it is the
+                        single softmax output.
+    pred                (model, noise_level, sample) int8
+                        argmax over class_name of ``prob``.
+    predictive_entropy  (model, noise_level, sample) float32
+                        Total/predictive uncertainty: entropy of the mean
+                        predictive distribution (bayesian_torch.predictive_entropy).
+    mutual_information  (model, noise_level, sample) float32
+                        Epistemic/model uncertainty (bayesian_torch.mutual_information).
+                        Identically ~0 for deterministic (K=1) models.
+    correct             (model, noise_level, sample) int8 (0/1)
+                        Whether ``pred`` matches the true class.
+    mc_prob             (model, noise_level, sample, class_name, mc_pass) float32
+                        OPTIONAL raw per-draw probabilities, only if retained.
+
+Coordinates
+    model_kind   (model)   "normal" | "bayesian"
+    image_id     (sample)  ADNI Image Data ID (join key to clinical data)
+    ptid         (sample)  patient id (enables patient-level analysis)
+    true_class   (sample)  ground-truth class label, one of class_name
+    noise_level  (noise_level) float sigma values
+    class_name   (class_name)  ["AD", "NL"]
+
+Attributes
+    schema_version, created (ISO-8601 UTC), seed, git_commit, n_mc,
+    config (TOML string snapshot)
+"""
+
+from __future__ import annotations
+
+import datetime as _dt
+import subprocess
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
+
+import numpy as np
+import xarray as xr
+
+# bayesian-torch's own uncertainty utilities, so definitions match the library
+# the Bayesian models were trained with (operate on numpy mc_preds).
+from bayesian_torch.utils.util import mutual_information as _bt_mutual_information
+from bayesian_torch.utils.util import predictive_entropy as _bt_predictive_entropy
+
+SCHEMA_VERSION = "1.0"
+
+CLASS_NAMES: Tuple[str, ...] = ("AD", "NL")
+
+MODEL_KIND_NORMAL = "normal"
+MODEL_KIND_BAYESIAN = "bayesian"
+
+# netCDF engine that supports the full type set (float32, int8, VLEN strings).
+_NETCDF_ENGINE = "netcdf4"
+
+# Encodings applied on write. netCDF has no native bool, so ``correct`` is int8.
+_ENCODINGS: Dict[str, Dict[str, Any]] = {
+    "prob": {"dtype": "float32", "zlib": True, "complevel": 4},
+    "pred": {"dtype": "int8", "zlib": True, "complevel": 4},
+    "predictive_entropy": {"dtype": "float32", "zlib": True, "complevel": 4},
+    "mutual_information": {"dtype": "float32", "zlib": True, "complevel": 4},
+    "correct": {"dtype": "int8", "zlib": True, "complevel": 4},
+    "mc_prob": {"dtype": "float32", "zlib": True, "complevel": 4},
+}
+
+
+# ---------------------------------------------------------------------------
+# Uncertainty computation
+# ---------------------------------------------------------------------------
+def compute_uncertainty(
+    mc_probs: np.ndarray,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Reduce Monte-Carlo predictive probabilities to summary statistics.
+
+    Args:
+        mc_probs: Array of shape ``(K, S, C)`` -- K Monte-Carlo passes, S
+            samples, C classes. For a deterministic model use ``K == 1``.
+
+    Returns:
+        ``(mean_prob, predictive_entropy, mutual_information)`` with shapes
+        ``(S, C)``, ``(S,)`` and ``(S,)`` respectively. ``mutual_information``
+        is ~0 when ``K == 1``.
+    """
+    mc = np.asarray(mc_probs, dtype=np.float64)
+    if mc.ndim != 3:
+        raise ValueError(
+            f"mc_probs must be (K, S, C); got shape {mc.shape} with ndim {mc.ndim}."
+        )
+
+    mean_prob = mc.mean(axis=0)
+    pred_entropy = _bt_predictive_entropy(mc)
+    # mutual_information == predictive_entropy - mean(per-pass entropy). With a
+    # single pass the two entropy terms are identical, so this is exactly 0.
+    mut_info = _bt_mutual_information(mc)
+
+    return (
+        mean_prob.astype(np.float32),
+        np.asarray(pred_entropy, dtype=np.float32),
+        np.asarray(mut_info, dtype=np.float32),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Dataset construction
+# ---------------------------------------------------------------------------
+def build_evaluation_dataset(
+    *,
+    prob: np.ndarray,
+    predictive_entropy: np.ndarray,
+    mutual_information: np.ndarray,
+    model_kinds: Sequence[str],
+    noise_levels: Sequence[float],
+    image_ids: Sequence[int],
+    ptids: Sequence[str],
+    true_classes: Sequence[str],
+    mc_prob: Optional[np.ndarray] = None,
+    attrs: Optional[Mapping[str, Any]] = None,
+) -> xr.Dataset:
+    """Assemble a schema-compliant ``xarray.Dataset``.
+
+    Args:
+        prob: ``(M, L, S, C)`` mean predictive probabilities.
+        predictive_entropy: ``(M, L, S)`` total uncertainty.
+        mutual_information: ``(M, L, S)`` model uncertainty.
+        model_kinds: length-M sequence of "normal"/"bayesian".
+        noise_levels: length-L sequence of sigma values (include 0.0 baseline).
+        image_ids: length-S sequence of ADNI Image Data IDs.
+        ptids: length-S sequence of patient ids.
+        true_classes: length-S ground-truth labels drawn from ``CLASS_NAMES``.
+        mc_prob: optional ``(M, L, S, C, K)`` raw MC draws.
+        attrs: optional extra dataset attributes (merged over the defaults).
+
+    Returns:
+        A validated ``xarray.Dataset`` following the module schema.
+    """
+    prob = np.asarray(prob, dtype=np.float32)
+    predictive_entropy = np.asarray(predictive_entropy, dtype=np.float32)
+    mutual_information = np.asarray(mutual_information, dtype=np.float32)
+
+    m, ell, s, c = _validate_shapes(
+        prob,
+        predictive_entropy,
+        mutual_information,
+        model_kinds,
+        noise_levels,
+        image_ids,
+        ptids,
+        true_classes,
+    )
+
+    # Derived variables: predicted class index and correctness.
+    pred = prob.argmax(axis=-1).astype(np.int8)  # (M, L, S)
+    class_index = {name: i for i, name in enumerate(CLASS_NAMES)}
+    try:
+        true_idx = np.array([class_index[t] for t in true_classes], dtype=np.int8)
+    except KeyError as exc:
+        raise ValueError(
+            f"true_classes contains a label not in CLASS_NAMES {CLASS_NAMES}: {exc}"
+        ) from None
+    # Broadcast true index (S,) against pred (M, L, S).
+    correct = (pred == true_idx[None, None, :]).astype(np.int8)
+
+    data_vars: Dict[str, Any] = {
+        "prob": (("model", "noise_level", "sample", "class_name"), prob),
+        "pred": (("model", "noise_level", "sample"), pred),
+        "predictive_entropy": (
+            ("model", "noise_level", "sample"),
+            predictive_entropy,
+        ),
+        "mutual_information": (
+            ("model", "noise_level", "sample"),
+            mutual_information,
+        ),
+        "correct": (("model", "noise_level", "sample"), correct),
+    }
+
+    if mc_prob is not None:
+        mc_prob = np.asarray(mc_prob, dtype=np.float32)
+        if mc_prob.ndim != 5 or mc_prob.shape[:4] != (m, ell, s, c):
+            raise ValueError(
+                f"mc_prob must be (M, L, S, C, K)=({m}, {ell}, {s}, {c}, K); "
+                f"got {mc_prob.shape}."
+            )
+        data_vars["mc_prob"] = (
+            ("model", "noise_level", "sample", "class_name", "mc_pass"),
+            mc_prob,
+        )
+
+    coords: Dict[str, Any] = {
+        "model": np.arange(m, dtype=np.int32),
+        "model_kind": ("model", np.asarray(list(model_kinds), dtype=object)),
+        "noise_level": np.asarray(list(noise_levels), dtype=np.float32),
+        "sample": np.arange(s, dtype=np.int32),
+        "image_id": ("sample", np.asarray(list(image_ids), dtype=np.int64)),
+        "ptid": ("sample", np.asarray(list(ptids), dtype=object)),
+        "true_class": ("sample", np.asarray(list(true_classes), dtype=object)),
+        "class_name": np.asarray(CLASS_NAMES, dtype=object),
+    }
+
+    ds = xr.Dataset(data_vars=data_vars, coords=coords)
+    ds.attrs.update(_default_attrs())
+    if attrs:
+        ds.attrs.update(dict(attrs))
+    return ds
+
+
+def _validate_shapes(
+    prob: np.ndarray,
+    predictive_entropy: np.ndarray,
+    mutual_information: np.ndarray,
+    model_kinds: Sequence[str],
+    noise_levels: Sequence[float],
+    image_ids: Sequence[int],
+    ptids: Sequence[str],
+    true_classes: Sequence[str],
+) -> Tuple[int, int, int, int]:
+    if prob.ndim != 4:
+        raise ValueError(f"prob must be (M, L, S, C); got shape {prob.shape}.")
+    m, ell, s, c = prob.shape
+
+    if c != len(CLASS_NAMES):
+        raise ValueError(
+            f"prob's class axis has length {c}, expected {len(CLASS_NAMES)} "
+            f"for CLASS_NAMES {CLASS_NAMES}."
+        )
+    for name, arr in (
+        ("predictive_entropy", predictive_entropy),
+        ("mutual_information", mutual_information),
+    ):
+        if arr.shape != (m, ell, s):
+            raise ValueError(
+                f"{name} must be (M, L, S)=({m}, {ell}, {s}); got {arr.shape}."
+            )
+    for name, seq, expected in (
+        ("model_kinds", model_kinds, m),
+        ("noise_levels", noise_levels, ell),
+        ("image_ids", image_ids, s),
+        ("ptids", ptids, s),
+        ("true_classes", true_classes, s),
+    ):
+        if len(seq) != expected:
+            raise ValueError(
+                f"{name} must have length {expected}; got {len(seq)}."
+            )
+    unknown_kinds = set(model_kinds) - {MODEL_KIND_NORMAL, MODEL_KIND_BAYESIAN}
+    if unknown_kinds:
+        raise ValueError(
+            f"model_kinds contains unknown values {unknown_kinds}; "
+            f"expected only {MODEL_KIND_NORMAL!r}/{MODEL_KIND_BAYESIAN!r}."
+        )
+    return m, ell, s, c
+
+
+def _default_attrs() -> Dict[str, Any]:
+    return {
+        "schema_version": SCHEMA_VERSION,
+        "created": _dt.datetime.now(_dt.timezone.utc).isoformat(),
+        "git_commit": _git_commit(),
+    }
+
+
+def _git_commit() -> str:
+    try:
+        out = subprocess.run(
+            ["git", "rev-parse", "HEAD"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        return out.stdout.strip() if out.returncode == 0 else "unknown"
+    except Exception:
+        return "unknown"
+
+
+# ---------------------------------------------------------------------------
+# I/O
+# ---------------------------------------------------------------------------
+def save_evaluation(ds: xr.Dataset, path: str) -> None:
+    """Write an evaluation dataset to netCDF using the locked-down encodings."""
+    encoding = {name: _ENCODINGS[name] for name in ds.data_vars if name in _ENCODINGS} # pyright: ignore
+    ds.to_netcdf(path, engine=_NETCDF_ENGINE, encoding=encoding)
+
+
+def load_evaluation(path: str) -> xr.Dataset:
+    """Load an evaluation dataset previously written by :func:`save_evaluation`."""
+    return xr.open_dataset(path, engine=_NETCDF_ENGINE)
+
+
+# ---------------------------------------------------------------------------
+# Incremental builder
+# ---------------------------------------------------------------------------
+class EvaluationAccumulator:
+    """Collect per-(model, noise_level) results, then finalize to a Dataset.
+
+    This is the intended entry point for the (not-yet-implemented) evaluation
+    tasks. Sample-axis metadata (image ids, ptids, true classes) is fixed by the
+    non-shuffled test loader and supplied once; each ``add`` call contributes one
+    model's Monte-Carlo probabilities at one noise level.
+
+    Example (pseudocode for a future evaluate task)::
+
+        acc = EvaluationAccumulator(image_ids, ptids, true_classes,
+                                    noise_levels=[0.0, 0.05])
+        for m, model in enumerate(models):
+            for sigma in [0.0, 0.05]:
+                mc = run_mc_passes(model, loader, sigma, k=K)  # (K, S, C)
+                acc.add(model_index=m, model_kind="bayesian",
+                        noise_level=sigma, mc_probs=mc)
+        ds = acc.to_dataset(attrs={"seed": seed, "n_mc": K})
+        save_evaluation(ds, out_path)
+    """
+
+    def __init__(
+        self,
+        image_ids: Sequence[int],
+        ptids: Sequence[str],
+        true_classes: Sequence[str],
+        noise_levels: Sequence[float],
+    ) -> None:
+        self._image_ids = list(image_ids)
+        self._ptids = list(ptids)
+        self._true_classes = list(true_classes)
+        self._noise_levels = list(noise_levels)
+        self._s = len(self._image_ids)
+
+        if not (len(self._ptids) == len(self._true_classes) == self._s):
+            raise ValueError(
+                "image_ids, ptids, and true_classes must all have the same length."
+            )
+        self._noise_index = {round(float(n), 12): i for i, n in enumerate(self._noise_levels)}
+
+        # keyed by model index -> dict with kind and per-noise arrays
+        self._models: Dict[int, Dict[str, Any]] = {}
+
+    def add(
+        self,
+        *,
+        model_index: int,
+        model_kind: str,
+        noise_level: float,
+        mc_probs: np.ndarray,
+    ) -> None:
+        """Add one model's MC probabilities ``(K, S, C)`` at one noise level."""
+        key = round(float(noise_level), 12)
+        if key not in self._noise_index:
+            raise ValueError(
+                f"noise_level {noise_level} not in declared levels {self._noise_levels}."
+            )
+        mean_prob, pred_ent, mut_info = compute_uncertainty(mc_probs)
+        if mean_prob.shape[0] != self._s:
+            raise ValueError(
+                f"mc_probs sample axis {mean_prob.shape[0]} != declared S {self._s}."
+            )
+
+        entry = self._models.setdefault(
+            model_index,
+            {
+                "kind": model_kind,
+                "prob": {},
+                "pe": {},
+                "mi": {},
+            },
+        )
+        if entry["kind"] != model_kind:
+            raise ValueError(
+                f"model_index {model_index} already registered as {entry['kind']!r}, "
+                f"got {model_kind!r}."
+            )
+        entry["prob"][key] = mean_prob
+        entry["pe"][key] = pred_ent
+        entry["mi"][key] = mut_info
+
+    def to_dataset(self, attrs: Optional[Mapping[str, Any]] = None) -> xr.Dataset:
+        """Stack all collected results into a schema-compliant Dataset."""
+        if not self._models:
+            raise ValueError("No results were added to the accumulator.")
+
+        model_indices = sorted(self._models)
+        m = len(model_indices)
+        ell = len(self._noise_levels)
+        c = len(CLASS_NAMES)
+
+        prob = np.zeros((m, ell, self._s, c), dtype=np.float32)
+        pe = np.zeros((m, ell, self._s), dtype=np.float32)
+        mi = np.zeros((m, ell, self._s), dtype=np.float32)
+        model_kinds: List[str] = []
+
+        for mi_idx, model_index in enumerate(model_indices):
+            entry = self._models[model_index]
+            model_kinds.append(entry["kind"])
+            for key, li in self._noise_index.items():
+                if key not in entry["prob"]:
+                    raise ValueError(
+                        f"model_index {model_index} is missing noise level {key}."
+                    )
+                prob[mi_idx, li] = entry["prob"][key]
+                pe[mi_idx, li] = entry["pe"][key]
+                mi[mi_idx, li] = entry["mi"][key]
+
+        return build_evaluation_dataset(
+            prob=prob,
+            predictive_entropy=pe,
+            mutual_information=mi,
+            model_kinds=model_kinds,
+            noise_levels=self._noise_levels,
+            image_ids=self._image_ids,
+            ptids=self._ptids,
+            true_classes=self._true_classes,
+            attrs=attrs,
+        )

model/dataset.py

@@ -1,7 +1,8 @@
+import math
 import pathlib as pl
 import random
 import re
-from typing import Callable, Dict, Iterator, List, Tuple
+from typing import Callable, Dict, Iterator, List, Tuple, cast
 
 import nibabel as nib
 import pandas as pd
@@ -10,6 +11,15 @@ import torch.utils.data as data
 from jaxtyping import Float
 from torch.utils.data import DataLoader, Subset
 
+from util.seeding import torch_generator
+
+# Tokens that appear in MRI filenames to indicate the diagnostic class.
+# On disk the cognitively-normal class is spelled "NL" (not "CN").
+CLASS_TOKENS: Dict[str, torch.Tensor] = {
+    "AD": torch.tensor([1.0, 0.0]),
+    "NL": torch.tensor([0.0, 1.0]),
+}
+
 
 def _row_to_float_tensor(row: pd.DataFrame, *, image_id: int) -> torch.Tensor:
     values = row.drop(columns=["Image Data ID"]).iloc[0]
@@ -28,12 +38,30 @@ def xls_pre(df: pd.DataFrame) -> pd.DataFrame:
     """
     Preprocess the Excel DataFrame.
     This function can be customized to filter or modify the DataFrame as needed.
+
+    Returns a DataFrame with columns [Image Data ID, Sex, Age (current)] where
+    Sex is encoded 0/1. "Image Data ID" is deliberately kept as a column (not an
+    index) because downstream lookups filter on it directly
+    (``xls_values["Image Data ID"] == img_id``).
     """
 
-    data: pd.DataFrame = df[["Image Data ID", "Sex", "Age (current)"]]  # pyright: ignore
-    data["Sex"] = data["Sex"].str.strip()  # type: ignore
-    data = data.replace({"M": 0, "F": 1})  # type: ignore
-    data.set_index("Image Data ID")  # type: ignore
+    # Work on an explicit copy so column assignment writes back reliably instead
+    # of raising SettingWithCopyWarning against a view of ``df``.
+    #
+    # The ``cast``s are needed because pandas ships no type stubs: basedpyright
+    # widens ``.copy()``/``.replace()``/``__getitem__`` to
+    # ``DataFrame | Series | Unknown``, which then has no ``.str`` accessor and
+    # isn't assignable to the declared return type. ``cast`` asserts the concrete
+    # pandas type at each step without any runtime effect.
+    data = cast(pd.DataFrame, df[["Image Data ID", "Sex", "Age (current)"]].copy())
+
+    # ``.str`` is pandas' vectorized string accessor: it applies a str method
+    # element-wise over every value in the Series. Here it strips surrounding
+    # whitespace from each Sex entry (e.g. " M " -> "M") before encoding.
+    sex_col = cast(pd.Series, data["Sex"])
+    data["Sex"] = sex_col.str.strip()
+
+    data = cast(pd.DataFrame, data.replace({"M": 0, "F": 1}))
 
     return data
 
@@ -136,16 +164,18 @@ def load_adni_data_from_file(
 
         file_mri_data = torch.from_numpy(nib.load(file).get_fdata())  # type: ignore # type checking does not work well with nibabel
 
-        # Read the filename to determine the expected class
-        file_expected_class = torch.tensor([0.0, 0.0])  # Default to a tensor of zeros
+        # Read the filename to determine the expected class. On disk the classes
+        # are spelled "AD" and "NL"; see CLASS_TOKENS.
+        file_expected_class: torch.Tensor | None = None
+        for token, one_hot in CLASS_TOKENS.items():
+            if token in filename:
+                file_expected_class = one_hot.clone()
+                break
 
-        if "AD" in filename:
-            file_expected_class = torch.tensor([1.0, 0.0])
-        elif "NL" in filename:
-            file_expected_class = torch.tensor([0.0, 1.0])
-        else:
+        if file_expected_class is None:
             raise ValueError(
-                f"Filename {filename} does not contain a valid class identifier (AD or CN)."
+                f"Filename {filename} does not contain a valid class identifier "
+                f"({' or '.join(CLASS_TOKENS)})."
             )
 
         mri_data_unstacked.append(file_mri_data)
@@ -196,8 +226,8 @@ def divide_dataset(
     Returns:
         Result[List[data.Subset[ADNIDataset]], str]: A Result object containing the subsets or an error message.
     """
-    if sum(ratios) != 1.0:
-        raise ValueError(f"Ratios must sum to 1.0, got {ratios}.")
+    if not math.isclose(sum(ratios), 1.0, rel_tol=1e-9, abs_tol=1e-9):
+        raise ValueError(f"Ratios must sum to 1.0, got {ratios} (sum={sum(ratios)}).")
 
     # Set the random seed for reproducibility
     gen = torch.Generator().manual_seed(seed)
@@ -207,27 +237,52 @@ def divide_dataset(
 def initalize_dataloaders(
     datasets: List[Subset[ADNIDataset]],
     batch_size: int = 64,
+    seed: int | None = None,
 ) -> List[DataLoader[ADNIDataset]]:
     """
-    Initializes the DataLoader for the given datasets.
+    Initializes the DataLoaders for the [train, val, test] datasets.
+
+    The three splits are configured differently:
+      - train: shuffled (for SGD) with a seeded generator for reproducibility,
+        and ``drop_last=True`` so a size-1 final batch cannot crash BatchNorm1d.
+      - val / test: NOT shuffled, so sample order is stable. Stable order is
+        required for evaluation, where saved model outputs must align with a
+        fixed sample axis (see the netCDF evaluation schema).
 
     Args:
-        datasets (List[Subset[ADNIDataset]]): List of datasets to create DataLoaders for.
-        batch_size (int): The batch size for the DataLoader.
+        datasets: Subsets in ``[train, val, test]`` order.
+        batch_size: The batch size for the DataLoaders.
+        seed: Seed for the training shuffle generator. ``None`` leaves shuffling
+            to global RNG state.
 
     Returns:
-        List[DataLoader[ADNIDataset]]: A list of DataLoaders for the datasets.
+        List[DataLoader[ADNIDataset]]: DataLoaders in the same order as ``datasets``.
     """
+    loader_configs = [
+        {"shuffle": True, "drop_last": True},  # train
+        {"shuffle": False, "drop_last": False},  # val
+        {"shuffle": False, "drop_last": False},  # test
+    ]
+    if len(datasets) != len(loader_configs):
+        raise ValueError(
+            f"Expected 3 datasets ([train, val, test]), got {len(datasets)}."
+        )
+
     pin_memory = torch.cuda.is_available()
-    return [
-        DataLoader(
-            dataset,
-            batch_size=batch_size,
-            shuffle=True,
-            pin_memory=pin_memory,
+    loaders: List[DataLoader[ADNIDataset]] = []
+    for dataset, cfg in zip(datasets, loader_configs):
+        generator = torch_generator(seed) if cfg["shuffle"] else None
+        loaders.append(
+            DataLoader(
+                dataset,
+                batch_size=batch_size,
+                shuffle=cfg["shuffle"],
+                drop_last=cfg["drop_last"],
+                pin_memory=pin_memory,
+                generator=generator,
+            )
         )
-        for dataset in datasets
-    ]
+    return loaders
 
 
 def divide_dataset_by_patient_id(
@@ -253,8 +308,8 @@ def divide_dataset_by_patient_id(
         This split is grouped by PTID, so all images from the same patient are assigned
         to exactly one partition to avoid patient-level leakage across train/val/test.
     """
-    if sum(ratios) != 1.0:
-        raise ValueError(f"Ratios must sum to 1.0, got {ratios}.")
+    if not math.isclose(sum(ratios), 1.0, rel_tol=1e-9, abs_tol=1e-9):
+        raise ValueError(f"Ratios must sum to 1.0, got {ratios} (sum={sum(ratios)}).")
 
     if not ptids:
         raise ValueError("ptids list cannot be empty.")

model/dnn_mod.py

@@ -36,6 +36,19 @@ from __future__ import absolute_import, division, print_function
 import bayesian_torch.layers as bayesian_layers
 from bayesian_torch.utils.util import get_rho
 
+# Canonical prior/posterior parameters for the DNN->BNN conversion. Training and
+# model loading MUST use the same values so a saved Bayesian state_dict maps onto
+# an identically-structured converted model.
+DEFAULT_BNN_PRIOR_PARAMETERS = {
+    "prior_mu": 0.0,
+    "prior_sigma": 1.0,
+    "posterior_mu_init": 0.0,
+    "posterior_rho_init": -3.0,
+    "type": "Reparameterization",
+    "moped_enable": False,
+    "moped_delta": 0.5,
+}
+
 # --------------------------------------------------------------------------------
 # Parameters used to define BNN layyers.
 #    bnn_prior_parameters = {

model/layers.py

@@ -44,30 +44,34 @@ class SplitCNVBlock(nn.Module):
 
         self.split_dim = split_dim
 
-        self.leftcnv_1 = SepCNVBlock(
-            in_channels // 2, mid_channels // 2, (3, 4, 3), droprate=drop_rate
+        # Build both branches once here. Previously these Sequentials were
+        # (re)assigned to ``self`` inside forward(), which registered
+        # ``leftblock``/``rightblock`` as submodules at runtime and duplicated
+        # their parameters into the state_dict. A freshly constructed model (no
+        # forward yet) lacked those keys, so a strict load_state_dict of a trained
+        # checkpoint failed with "Unexpected key(s)" at evaluation time.
+        self.leftblock = nn.Sequential(
+            SepCNVBlock(
+                in_channels // 2, mid_channels // 2, (3, 4, 3), droprate=drop_rate
+            ),
+            SepCNVBlock(
+                mid_channels // 2, out_channels // 2, (3, 4, 3), droprate=drop_rate
+            ),
         )
-        self.rightcnv_1 = SepCNVBlock(
-            in_channels // 2, mid_channels // 2, (3, 4, 3), droprate=drop_rate
-        )
-
-        self.leftcnv_2 = SepCNVBlock(
-            mid_channels // 2, out_channels // 2, (3, 4, 3), droprate=drop_rate
-        )
-        self.rightcnv_2 = SepCNVBlock(
-            mid_channels // 2, out_channels // 2, (3, 4, 3), droprate=drop_rate
+        self.rightblock = nn.Sequential(
+            SepCNVBlock(
+                in_channels // 2, mid_channels // 2, (3, 4, 3), droprate=drop_rate
+            ),
+            SepCNVBlock(
+                mid_channels // 2, out_channels // 2, (3, 4, 3), droprate=drop_rate
+            ),
         )
 
     def forward(self, x: Float[torch.Tensor, "N C D H W"]):
-        (left, right) = torch.tensor_split(x, 2, dim=self.split_dim)
-
-        self.leftblock = nn.Sequential(self.leftcnv_1, self.leftcnv_2)
-        self.rightblock = nn.Sequential(self.rightcnv_1, self.rightcnv_2)
-
+        left, right = torch.tensor_split(x, 2, dim=self.split_dim)
         left = self.leftblock(left)
         right = self.rightblock(right)
-        a = torch.cat((left, right), dim=self.split_dim)
-        return a
+        return torch.cat((left, right), dim=self.split_dim)
 
 
 class MidFlowBlock(nn.Module):
@@ -84,8 +88,8 @@ class MidFlowBlock(nn.Module):
             channels, channels, (3, 3, 3), droprate=drop_rate, padding="same"
         )
 
-        # self.block = nn.Sequential(self.cnv1, self.cnv2, self.cnv3)
-        self.block = self.cnv1
+        self.block = nn.Sequential(self.cnv1, self.cnv2, self.cnv3)
+        #self.block = self.cnv1
 
     def forward(self, x: Float[torch.Tensor, "N C D H W"]):
         a = nn.ELU()(self.block(x) + x)

model/training.py

@@ -1,5 +1,4 @@
 import pathlib as pl
-import time
 from threading import Event
 from typing import Callable, Tuple, cast
 
@@ -10,6 +9,7 @@ import torch.nn as nn
 from torch.utils.data import DataLoader
 
 from model.dataset import ADNIDataset
+from util.control import check_control_events as _check_control_events
 from util.progress import ProgressTracker
 from util.ui_logger import PipelineLogger
 
@@ -62,18 +62,6 @@ def _batch_correct_and_total(
     return correct, total
 
 
-def _check_control_events(
-    stop_event: Event | None,
-    pause_event: Event | None,
-) -> None:
-    if stop_event is not None and stop_event.is_set():
-        raise InterruptedError("Pipeline execution stopped by user.")
-    while pause_event is not None and pause_event.is_set():
-        time.sleep(0.5)
-        if stop_event is not None and stop_event.is_set():
-            raise InterruptedError("Pipeline execution stopped by user while paused.")
-
-
 def test_model(
     model: nn.Module,
     test_loader: DataLoader[ADNIDataset],
@@ -102,20 +90,21 @@ def test_model(
 
     test_progress = progress.get_sub_tracker("Testing Batches")
     test_progress.update(total=len(test_loader), advance=0)
-    for _, (mri, xls, targets, _) in enumerate(test_loader):
-        _check_control_events(stop_event=stop_event, pause_event=pause_event)
-        mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
-        outputs = model((mri, xls))
-        loss = criterion(outputs, targets)
-        batch_size = mri.size(0)
-        test_loss += loss.item() * batch_size
-        total_samples += batch_size
+    with torch.no_grad():
+        for _, (mri, xls, targets, _) in enumerate(test_loader):
+            _check_control_events(stop_event=stop_event, pause_event=pause_event)
+            mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
+            outputs = model((mri, xls))
+            loss = criterion(outputs, targets)
+            batch_size = mri.size(0)
+            test_loss += loss.item() * batch_size
+            total_samples += batch_size
 
-        # Calculate accuracy
-        batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
-        correct += batch_correct
-        total += batch_total
-        test_progress.update(total=len(test_loader), advance=1)
+            # Calculate accuracy
+            batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
+            correct += batch_correct
+            total += batch_total
+            test_progress.update(total=len(test_loader), advance=1)
 
     test_loss = test_loss / total_samples if total_samples > 0 else 0.0
     test_acc = correct / total if total > 0 else 0.0

pyproject.toml

@@ -21,8 +21,15 @@ dependencies = [
     "jaxtyping",
     "textual",
     "toml",
+    "xarray",
+    "netcdf4>=1.7.4",
 ]
 
 
 [tool.mypy]
 exclude = [".venv/**"]
+
+[dependency-groups]
+dev = [
+    "pandas-stubs>=3.0.3.260530",
+]

tasks/__init__.py

@@ -1,58 +1,28 @@
-import time
-from typing import Any, Dict
+"""Pipeline task registry and scenario definitions.
 
-from util.progress import ProgressTracker
-from util.ui_logger import PipelineLogger
+A *task* is a callable ``task(tracker, logger, config, state) -> None``. A
+*scenario* is an ordered list of task ids the pipeline runs. Tasks communicate
+through the mutable ``state`` dict (dataloaders, model paths, control events).
 
+Pipeline stages (see ai/ARCHITECTURE.md):
+    1. load_data        -> implemented
+    2. train_regular    -> implemented
+    3. train_bayesian   -> implemented
+    4. evaluate_regular -> implemented
+    5. evaluate_bayesian-> implemented
+    6. evaluate_noisy   -> implemented
+    load_models         -> implemented; discovers saved .pt ensembles on disk and
+                           records their paths in ``state`` so evaluation is
+                           decoupled from training (training frees models from VRAM
+                           after saving). Evaluation loads one model at a time.
+"""
+
+from . import evaluate
 from . import load_data
+from . import load_models
 from . import train_bayesian
 from . import train_normal
 
-
-def dummy_task(
-    tracker: ProgressTracker,
-    logger: PipelineLogger,
-    config: Dict[str, Any],
-    state: Dict[str, Any],
-):
-    stop_event = state["events"]["stop"]
-    pause_event = state["events"]["pause"]
-
-    steps = 5
-    # The title is now set gracefully by the parent injecting the sub_tracker,
-    # so we just initialize the total.
-    tracker.update(total=steps, advance=0)
-
-    logger.info("Initializing process...")
-
-    for i in range(steps):
-        if stop_event.is_set():
-            raise InterruptedError("Pipeline execution stopped by user.")
-
-        while pause_event.is_set():
-            time.sleep(0.5)
-            if stop_event.is_set():
-                raise InterruptedError(
-                    "Pipeline execution stopped by user while paused."
-                )
-
-        # Child Process (Sub Progress Tracker)
-        sub_steps = 10
-        sub_tracker = tracker.get_sub_tracker(f"Batch {i + 1}")
-        sub_tracker.update(total=sub_steps, advance=0)
-
-        for j in range(sub_steps):
-            if stop_event.is_set():
-                raise InterruptedError("Pipeline execution stopped by user.")
-            time.sleep(0.05)
-            sub_tracker.update(advance=1)  # Advance sub task
-
-        tracker.update(advance=1)  # Advance main task (clears sub task)
-
-        if i == 2:
-            logger.info("Halfway through current task execution...")
-
-
 PIPELINE_TASKS = {
     "load_data": {
         "task_name": "Load Image and ADNIMERGE",
@@ -66,13 +36,21 @@ PIPELINE_TASKS = {
         "task_name": "Train Bayesian Models",
         "task_func": train_bayesian.train_bayesian_task,
     },
+    "load_models": {
+        "task_name": "Load Saved Models",
+        "task_func": load_models.load_models_task,
+    },
     "evaluate_regular": {
         "task_name": "Evaluate Regular Models",
-        "task_func": dummy_task,
+        "task_func": evaluate.evaluate_normal_task,
     },
     "evaluate_bayesian": {
         "task_name": "Evaluate Bayesian Models",
-        "task_func": dummy_task,
+        "task_func": evaluate.evaluate_bayesian_task,
+    },
+    "evaluate_noisy": {
+        "task_name": "Evaluate Models on Noised Data",
+        "task_func": evaluate.evaluate_noisy_task,
     },
 }
 
@@ -83,16 +61,29 @@ SCENARIOS = {
             "load_data",
             "train_regular",
             "train_bayesian",
+            "load_models",
             "evaluate_regular",
             "evaluate_bayesian",
+            "evaluate_noisy",
         ],
     },
     "scen_load_all": {
         "label": "2. Load, Evaluate, & Noise Analysis",
-        "tasks": ["load_data", "evaluate_regular", "evaluate_bayesian"],
+        "tasks": [
+            "load_data",
+            "load_models",
+            "evaluate_regular",
+            "evaluate_bayesian",
+            "evaluate_noisy",
+        ],
     },
     "scen_load_eval": {
         "label": "3. Load & Evaluate (Skip Noise)",
-        "tasks": ["load_data", "evaluate_regular", "evaluate_bayesian"],
+        "tasks": [
+            "load_data",
+            "load_models",
+            "evaluate_regular",
+            "evaluate_bayesian",
+        ],
     },
 }

tasks/load_data.py

@@ -5,6 +5,7 @@ import pandas as pd
 
 import model.dataset as ds
 from util.progress import ProgressTracker
+from util.seeding import seed_everything
 from util.ui_logger import PipelineLogger
 
 
@@ -15,6 +16,21 @@ def load_data_task(
     state: Dict[str, Any],
 ):
 
+    if config["data"]["seed"] is None:
+        log.info("Seed is not defined - using default seed of 0")
+        config["data"]["seed"] = 0
+
+    # Seed everything once at the start of the pipeline so data loading, weight
+    # init, and shuffling are reproducible. Per-member seeds are derived from
+    # this base inside the training tasks.
+    base_seed = int(config["data"]["seed"])
+    deterministic = bool(config.get("training", {}).get("deterministic", False))
+    applied = seed_everything(base_seed, deterministic=deterministic)
+    log.info(
+        f"Seeded RNGs with base seed {applied}"
+        + (" (deterministic cuDNN)" if deterministic else "")
+    )
+
     log.info("Loading Files")
     mri_files = pl.Path(config["data"]["mri_files_path"]).glob("*.nii")
     xls_file = pl.Path(config["data"]["xls_file_path"])
@@ -26,10 +42,6 @@ def load_data_task(
         xls_preprocessor=ds.xls_pre,
     )
 
-    if config["data"]["seed"] is None:
-        log.info("Seed is not defined - using default seed of 0")
-        config["data"]["seed"] = 0
-
     ptid_df = pd.read_csv(xls_file)
     ptid_df.columns = ptid_df.columns.str.strip()
 
@@ -42,6 +54,9 @@ def load_data_task(
 
     ptids = list(zip(ptid_df["Image Data ID"].tolist(), ptid_df["PTID"].tolist()))
 
+    # Mapping used later by evaluation to attach patient ids to the sample axis.
+    state["image_to_ptid"] = {int(iid): str(pid) for iid, pid in ptids}
+
     # Split is grouped by PTID to prevent patient-level leakage across partitions.
     datasets = ds.divide_dataset_by_patient_id(
         dataset,
@@ -50,9 +65,12 @@ def load_data_task(
         seed=config["data"]["seed"],
     )
 
-    # Initialize the dataloaders
+    # Initialize the dataloaders. The train shuffle generator is seeded from the
+    # base seed so shuffling is reproducible; val/test are left unshuffled.
     train_loader, val_loader, test_loader = ds.initalize_dataloaders(
-        datasets, batch_size=config["training"]["batch_size"]
+        datasets,
+        batch_size=config["training"]["batch_size"],
+        seed=base_seed,
     )
 
     log.info("Dataloaders initalized")

tasks/train_bayesian.py

@@ -1,7 +1,4 @@
-import gc
 import pathlib as pl
-import time
-from threading import Event
 from typing import Any, Dict
 
 import torch
@@ -11,29 +8,16 @@ from torch.utils.data import DataLoader
 import model.dataset as ds
 import model.training as tn
 from model.cnn import CNN3D
-from model.dnn_mod import dnn_to_bnn_mod, get_kl_loss
+from model.dnn_mod import DEFAULT_BNN_PRIOR_PARAMETERS, dnn_to_bnn_mod, get_kl_loss
+from util.control import check_control_events as _check_control_events
+from util.control import release_torch_memory as _release_torch_memory
 from util.progress import ProgressTracker
+from util.seeding import derive_seed, seed_everything
 from util.ui_logger import PipelineLogger
 
-
-def _release_torch_memory(device: str) -> None:
-    gc.collect()
-    if device.startswith("cuda"):
-        torch.cuda.empty_cache()
-    elif device.startswith("mps") and hasattr(torch, "mps"):
-        torch.mps.empty_cache()
-
-
-def _check_control_events(
-    stop_event: Event | None,
-    pause_event: Event | None,
-) -> None:
-    if stop_event is not None and stop_event.is_set():
-        raise InterruptedError("Pipeline execution stopped by user.")
-    while pause_event is not None and pause_event.is_set():
-        time.sleep(0.5)
-        if stop_event is not None and stop_event.is_set():
-            raise InterruptedError("Pipeline execution stopped by user while paused.")
+# RNG stream id for the Bayesian ensemble; distinct from the normal ensemble so
+# bayesian-member-N and normal-member-N do not share an RNG stream.
+_BAYESIAN_SEED_STREAM = 1
 
 
 def train_bayesian_task(
@@ -63,20 +47,20 @@ def train_bayesian_task(
     train_progress = track.get_sub_tracker("Training Progress")
     track.update(total=config["training"]["ensemble_size"], advance=0)
 
-    bnn_prior_parameters = {
-        "prior_mu": 0.0,
-        "prior_sigma": 1.0,
-        "posterior_mu_init": 0.0,
-        "posterior_rho_init": -3.0,
-        "type": "Reparameterization",
-        "moped_enable": False,
-        "moped_delta": 0.5,
-    }
+    # Shared with the model loader (evaluation) so saved state_dicts map onto an
+    # identically-converted model.
+    bnn_prior_parameters = DEFAULT_BNN_PRIOR_PARAMETERS
+
+    base_seed = int(config["data"]["seed"])
 
     for model_num in range(config["training"]["ensemble_size"]):
         _check_control_events(stop_event=stop_event, pause_event=pause_event)
+        # Seed per member (distinct stream from the normal ensemble).
+        member_seed = derive_seed(base_seed, model_num, stream=_BAYESIAN_SEED_STREAM)
+        seed_everything(member_seed)
         log.info(
-            f"Training model {model_num + 1}/{config['training']['ensemble_size']}..."
+            f"Training model {model_num + 1}/{config['training']['ensemble_size']} "
+            f"(seed {member_seed})..."
         )
         model = CNN3D(
             image_channels=config["data"]["image_channels"],

tasks/train_normal.py

@@ -1,7 +1,4 @@
 import pathlib as pl
-import time
-import gc
-from threading import Event
 from typing import Any, Dict
 
 import torch
@@ -11,28 +8,15 @@ from torch.utils.data import DataLoader
 import model.dataset as ds
 import model.training as tn
 from model.cnn import CNN3D
+from util.control import check_control_events as _check_control_events
+from util.control import release_torch_memory as _release_torch_memory
 from util.progress import ProgressTracker
+from util.seeding import derive_seed, seed_everything
 from util.ui_logger import PipelineLogger
 
-
-def _release_torch_memory(device: str) -> None:
-    gc.collect()
-    if device.startswith("cuda"):
-        torch.cuda.empty_cache()
-    elif device.startswith("mps") and hasattr(torch, "mps"):
-        torch.mps.empty_cache()
-
-
-def _check_control_events(
-    stop_event: Event | None,
-    pause_event: Event | None,
-) -> None:
-    if stop_event is not None and stop_event.is_set():
-        raise InterruptedError("Pipeline execution stopped by user.")
-    while pause_event is not None and pause_event.is_set():
-        time.sleep(0.5)
-        if stop_event is not None and stop_event.is_set():
-            raise InterruptedError("Pipeline execution stopped by user while paused.")
+# RNG stream id for the normal ensemble, keeps its seeds distinct from the
+# Bayesian ensemble's (see util.seeding.derive_seed).
+_NORMAL_SEED_STREAM = 0
 
 
 def train_normal_task(
@@ -59,13 +43,19 @@ def train_normal_task(
         intermediate_model_dir.mkdir(parents=True, exist_ok=True)
     log.info(f"Intermediate models will be saved to {intermediate_model_dir}")
 
-    models = []
+    base_seed = int(config["data"]["seed"])
+
     train_progress = track.get_sub_tracker("Training Progress")
     track.update(total=config["training"]["ensemble_size"], advance=0)
     for model_num in range(config["training"]["ensemble_size"]):
         _check_control_events(stop_event=stop_event, pause_event=pause_event)
+        # Seed per member so weight init / dropout are reproducible AND distinct
+        # across ensemble members.
+        member_seed = derive_seed(base_seed, model_num, stream=_NORMAL_SEED_STREAM)
+        seed_everything(member_seed)
         log.info(
-            f"Training model {model_num + 1}/{config['training']['ensemble_size']}..."
+            f"Training model {model_num + 1}/{config['training']['ensemble_size']} "
+            f"(seed {member_seed})..."
         )
         # Train the model
         model = (
@@ -79,8 +69,6 @@ def train_normal_task(
             .to(config["training"]["device"])
         )
 
-        models.append(model)
-
         optimizer = optim.Adam(
             model.parameters(), lr=config["training"]["learning_rate"]
         )

util/seeding.py

@@ -0,0 +1,78 @@
+"""Centralized seeding helpers for reproducible training and evaluation.
+
+Reproducibility matters here because the whole point of the harness is to
+compare ensemble members and quantify uncertainty. If weight initialization,
+dropout, dataloader shuffling, and BNN sampling are not seeded, "ensemble
+diversity" is accidental and results are not reproducible across runs.
+"""
+
+import os
+import random
+from typing import Optional
+
+import numpy as np
+import torch
+
+# Numpy requires seeds in the uint32 range; torch/python are more permissive,
+# so we clamp everything to this range for consistency.
+_UINT32 = 2**32
+
+
+def seed_everything(seed: int, deterministic: bool = False) -> int:
+    """Seed all relevant RNGs (python, numpy, torch, CUDA).
+
+    Args:
+        seed: The master seed. Clamped into the uint32 range.
+        deterministic: If True, force cuDNN into deterministic mode. This makes
+            convolutions reproducible at some throughput cost. Leave False for
+            training speed; set True when exact reproduction is required.
+
+    Returns:
+        The (clamped) seed actually applied, for logging.
+    """
+    seed = int(seed) % _UINT32
+
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+    if deterministic:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        # Opt into deterministic algorithms where available; warn_only so a
+        # missing deterministic kernel degrades gracefully instead of crashing.
+        torch.use_deterministic_algorithms(True, warn_only=True)
+
+    return seed
+
+
+def derive_seed(base_seed: int, index: int, stream: int = 0) -> int:
+    """Derive a stable, distinct seed for ensemble member ``index``.
+
+    ``stream`` separates otherwise-colliding families (e.g. normal vs. Bayesian
+    ensembles) so that normal-member-0 and bayesian-member-0 do not share an
+    identical RNG stream.
+
+    Args:
+        base_seed: The experiment master seed.
+        index: The ensemble member index (0-based).
+        stream: An offset identifying the RNG family.
+
+    Returns:
+        A seed in the uint32 range, deterministic in all three arguments.
+    """
+    return (int(base_seed) + stream * 100_000 + int(index)) % _UINT32
+
+
+def torch_generator(seed: Optional[int]) -> Optional[torch.Generator]:
+    """Return a seeded ``torch.Generator`` for DataLoader shuffling, or None.
+
+    Giving the training DataLoader its own generator makes the shuffle order a
+    pure function of ``seed`` rather than of global RNG state mutated elsewhere.
+    """
+    if seed is None:
+        return None
+    return torch.Generator().manual_seed(int(seed) % _UINT32)

uv.lock

@@ -2,7 +2,8 @@ version = 1
 revision = 3
 requires-python = ">=3.14"
 resolution-markers = [
-    "sys_platform == 'win32'",
+    "platform_machine == 'ARM64' and sys_platform == 'win32'",
+    "platform_machine != 'ARM64' and sys_platform == 'win32'",
     "sys_platform == 'emscripten'",
     "sys_platform != 'emscripten' and sys_platform != 'win32'",
 ]
@@ -25,6 +26,7 @@ dependencies = [
     { name = "jaxtyping" },
     { name = "jupyterlab" },
     { name = "matplotlib" },
+    { name = "netcdf4" },
     { name = "nibabel" },
     { name = "numpy" },
     { name = "pandas" },
@@ -39,12 +41,18 @@ dependencies = [
     { name = "xarray" },
 ]
 
+[package.dev-dependencies]
+dev = [
+    { name = "pandas-stubs" },
+]
+
 [package.metadata]
 requires-dist = [
     { name = "bayesian-torch" },
     { name = "jaxtyping" },
     { name = "jupyterlab" },
     { name = "matplotlib" },
+    { name = "netcdf4", specifier = ">=1.7.4" },
     { name = "nibabel" },
     { name = "numpy" },
     { name = "pandas" },
@@ -59,6 +67,9 @@ requires-dist = [
     { name = "xarray" },
 ]
 
+[package.metadata.requires-dev]
+dev = [{ name = "pandas-stubs", specifier = ">=3.0.3.260530" }]
+
 [[package]]
 name = "anyio"
 version = "4.14.1"
@@ -260,6 +271,30 @@ wheels = [
 ]
 
 [[package]]
+name = "cftime"
+version = "1.6.5"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/65/dc/470ffebac2eb8c54151eb893055024fe81b1606e7c6ff8449a588e9cd17f/cftime-1.6.5.tar.gz", hash = "sha256:8225fed6b9b43fb87683ebab52130450fc1730011150d3092096a90e54d1e81e", size = 326605, upload-time = "2025-10-13T18:56:26.352Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ea/6c/a9618f589688358e279720f5c0fe67ef0077fba07334ce26895403ebc260/cftime-1.6.5-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:c69ce3bdae6a322cbb44e9ebc20770d47748002fb9d68846a1e934f1bd5daf0b", size = 502725, upload-time = "2025-10-13T18:56:19.424Z" },
+    { url = "https://files.pythonhosted.org/packages/d8/e3/da3c36398bfb730b96248d006cabaceed87e401ff56edafb2a978293e228/cftime-1.6.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:e62e9f2943e014c5ef583245bf2e878398af131c97e64f8cd47c1d7baef5c4e2", size = 485445, upload-time = "2025-10-13T18:56:20.853Z" },
+    { url = "https://files.pythonhosted.org/packages/32/93/b05939e5abd14bd1ab69538bbe374b4ee2a15467b189ff895e9a8cdaddf6/cftime-1.6.5-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7da5fdaa4360d8cb89b71b8ded9314f2246aa34581e8105c94ad58d6102d9e4f", size = 1584434, upload-time = "2025-10-13T19:39:17.084Z" },
+    { url = "https://files.pythonhosted.org/packages/7f/89/648397f9936e0b330999c4e776ebf296ec3c6a65f9901687dbca4ab820da/cftime-1.6.5-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bff865b4ea4304f2744a1ad2b8149b8328b321dd7a2b9746ef926d229bd7cd49", size = 1609812, upload-time = "2025-10-13T18:56:21.971Z" },
+    { url = "https://files.pythonhosted.org/packages/e7/0f/901b4835aa67ad3e915605d4e01d0af80a44b114eefab74ae33de6d36933/cftime-1.6.5-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:e552c5d1c8a58f25af7521e49237db7ca52ed2953e974fe9f7c4491e95fdd36c", size = 1669768, upload-time = "2025-10-13T18:56:24.027Z" },
+    { url = "https://files.pythonhosted.org/packages/22/d5/e605e4b28363e7a9ae98ed12cabbda5b155b6009270e6a231d8f10182a17/cftime-1.6.5-cp314-cp314-win_amd64.whl", hash = "sha256:e645b095dc50a38ac454b7e7f0742f639e7d7f6b108ad329358544a6ff8c9ba2", size = 463818, upload-time = "2025-10-13T18:56:25.376Z" },
+    { url = "https://files.pythonhosted.org/packages/3d/89/a8f85ae697ff10206ec401c2621f5ca9f327554f586d62f244739ceeb347/cftime-1.6.5-cp314-cp314-win_arm64.whl", hash = "sha256:b9044d7ac82d3d8af189df1032fdc871bbd3f3dd41a6ec79edceb5029b71e6e0", size = 459862, upload-time = "2026-01-02T20:45:02.625Z" },
+    { url = "https://files.pythonhosted.org/packages/ab/05/7410e12fd03a0c52717e74e6a1b49958810807dda212e23b65d43ea99676/cftime-1.6.5-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:9ef56460cb0576e1a9161e1428c9e1a633f809a23fa9d598f313748c1ae5064e", size = 533781, upload-time = "2026-01-02T20:45:04.818Z" },
+    { url = "https://files.pythonhosted.org/packages/44/ba/10e3546426d3ed9f9cc82e4a99836bb6fac1642c7830f7bdd0ac1c3f0805/cftime-1.6.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:4f4873d38b10032f9f3111c547a1d485519ae64eee6a7a2d091f1f8b08e1ba50", size = 515218, upload-time = "2026-01-02T20:45:06.788Z" },
+    { url = "https://files.pythonhosted.org/packages/bd/68/efa11eae867749e921bfec6a865afdba8166e96188112dde70bb8bb49254/cftime-1.6.5-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ccce0f4c9d3f38dd948a117e578b50d0e0db11e2ca9435fb358fd524813e4b61", size = 1579932, upload-time = "2026-01-02T20:45:11.194Z" },
+    { url = "https://files.pythonhosted.org/packages/9d/6c/0971e602c1390a423e6621dfbad9f1d375186bdaf9c9c7f75e06f1fbf355/cftime-1.6.5-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:19cbfc5152fb0b34ce03acf9668229af388d7baa63a78f936239cb011ccbe6b1", size = 1555894, upload-time = "2026-01-02T20:45:16.351Z" },
+    { url = "https://files.pythonhosted.org/packages/ad/fc/8475a15b7c3209a4a68b563dfc5e01ce74f2d8b9822372c3d30c68ab7f39/cftime-1.6.5-cp314-cp314t-win_amd64.whl", hash = "sha256:4470cd5ef3c2514566f53efbcbb64dd924fa0584637d90285b2f983bd4ee7d97", size = 513027, upload-time = "2026-01-02T20:45:20.023Z" },
+    { url = "https://files.pythonhosted.org/packages/f7/80/4ecbda8318fbf40ad4e005a4a93aebba69e81382e5b4c6086251cd5d0ee8/cftime-1.6.5-cp314-cp314t-win_arm64.whl", hash = "sha256:034c15a67144a0a5590ef150c99f844897618b148b87131ed34fda7072614662", size = 469065, upload-time = "2026-01-02T20:45:23.398Z" },
+]
+
+[[package]]
 name = "charset-normalizer"
 version = "3.4.7"
 source = { registry = "https://pypi.org/simple" }
@@ -1203,6 +1238,33 @@ wheels = [
 ]
 
 [[package]]
+name = "netcdf4"
+version = "1.7.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "certifi" },
+    { name = "cftime" },
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/34/b6/0370bb3af66a12098da06dc5843f3b349b7c83ccbdf7306e7afa6248b533/netcdf4-1.7.4.tar.gz", hash = "sha256:cdbfdc92d6f4d7192ca8506c9b3d4c1d9892969ff28d8e8e1fc97ca08bf12164", size = 838352, upload-time = "2026-01-05T02:27:38.593Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/38/de/38ed7e1956943d28e8ea74161e97c3a00fb98d6d08943b4fd21bae32c240/netcdf4-1.7.4-cp311-abi3-macosx_13_0_x86_64.whl", hash = "sha256:dec70e809cc65b04ebe95113ee9c85ba46a51c3a37c058d2b2b0cadc4d3052d8", size = 23427499, upload-time = "2026-01-05T02:27:06.568Z" },
+    { url = "https://files.pythonhosted.org/packages/e5/70/2f73c133b71709c412bc81d8b721e28dc6237ba9d7dad861b7bfbb70408a/netcdf4-1.7.4-cp311-abi3-macosx_14_0_arm64.whl", hash = "sha256:75cf59100f0775bc4d6b9d4aca7cbabd12e2b8cf3b9a4fb16d810b92743a315a", size = 22847667, upload-time = "2026-01-05T02:27:09.421Z" },
+    { url = "https://files.pythonhosted.org/packages/77/ce/43a3c0c41a6e2e940d87feea79d29aa88302211ac122604838f8a5a48de6/netcdf4-1.7.4-cp311-abi3-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ddfc7e9d261125c74708119440c85ea288b5fee41db676d2ba1ce9be11f96932", size = 10274769, upload-time = "2026-01-05T21:31:19.243Z" },
+    { url = "https://files.pythonhosted.org/packages/7b/7a/a8d32501bb95ecff342004a674720164f95ad616f269450b3bc13dc88ae3/netcdf4-1.7.4-cp311-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a72c9f58767779ec14cb7451c3b56bdd8fdc027a792fac2062b14e090c5617f3", size = 10123122, upload-time = "2026-01-05T21:31:22.773Z" },
+    { url = "https://files.pythonhosted.org/packages/18/68/e89b4fa9242e59326c849c39ce0f49eb68499603c639405a8449900a4f15/netcdf4-1.7.4-cp311-abi3-win_amd64.whl", hash = "sha256:9476e1f23161ae5159cd1548c50c8a37922e77d76583e247133f256ef7b825fc", size = 21299637, upload-time = "2026-01-05T02:27:11.856Z" },
+    { url = "https://files.pythonhosted.org/packages/6c/fc/edd41a3607241027aa4533e7f18e0cd647e74dde10a63274c65350f59967/netcdf4-1.7.4-cp311-abi3-win_arm64.whl", hash = "sha256:876ad9d58f09c98741c066c726164c45a098a58fb90e5fac9e74de4bb8a793fd", size = 2386377, upload-time = "2026-01-05T02:27:13.808Z" },
+    { url = "https://files.pythonhosted.org/packages/d8/2b/684b15dd4791f8be295b2f6fa97377bbc07a768478a63b7d3c4951712e36/netcdf4-1.7.4-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a5841de0735e8e4875b367c668e81d334287858d64dd9f3e3e2261e808c84922", size = 10395635, upload-time = "2026-01-05T02:27:19.655Z" },
+    { url = "https://files.pythonhosted.org/packages/37/dc/44d21524cf1b1c64254f92e22395a7a10f70c18f3a13a18ac9db258760f7/netcdf4-1.7.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:86fac03a8c5b250d57866e7d98918a64742e4b0de1681c5c86bac5726bab8aee", size = 10237725, upload-time = "2026-01-05T02:27:22.298Z" },
+    { url = "https://files.pythonhosted.org/packages/d4/9d/c3ddf54296ad8f18f02f77f23452bdb0971aece1b87e84bab9d734bf72cc/netcdf4-1.7.4-cp314-cp314t-macosx_13_0_x86_64.whl", hash = "sha256:ad083d260301b5add74b1669c75ab0df03bdf986decfcc092cb45eec2615b5f1", size = 23515258, upload-time = "2026-01-05T02:27:24.837Z" },
+    { url = "https://files.pythonhosted.org/packages/dd/44/bc0346e995d436d03fab682b7fbd2a9adcf0db6a05790b8f24853bf08170/netcdf4-1.7.4-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:7f22014092cc9da3f056b0368e2e38c42afd5725c87ad4843eb2f467e16dd4f6", size = 22910171, upload-time = "2026-01-05T02:27:27.166Z" },
+    { url = "https://files.pythonhosted.org/packages/30/6b/f9bc3f43c55e2dac72ee9f98d77860789bdd5d50c29adf164a6bdb303078/netcdf4-1.7.4-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:224a15434c165a5e0225e5831f591edf62533044b1ce62fdfee815195bbd077d", size = 10567579, upload-time = "2026-01-05T02:27:29.382Z" },
+    { url = "https://files.pythonhosted.org/packages/6d/d5/e7685c66b7f011c73cd746127f986358a26c642a4e4a1aa5ab51481b6586/netcdf4-1.7.4-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:31a2318305de6831a18df25ad0df9f03b6d68666af0356d4f6057d66c02ffeb6", size = 10255032, upload-time = "2026-01-05T02:27:31.744Z" },
+    { url = "https://files.pythonhosted.org/packages/a6/14/7506738bb6c8bc373b01e5af8f3b727f83f4f496c6b108490ea2609dc2cf/netcdf4-1.7.4-cp314-cp314t-win_amd64.whl", hash = "sha256:6c4a0aa9446c3a616ef3be015b629dc6173643f8b09546de26a4e40e272cd1ed", size = 22289653, upload-time = "2026-01-05T02:27:34.294Z" },
+    { url = "https://files.pythonhosted.org/packages/af/2e/39d5e9179c543f2e6e149a65908f83afd9b6d64379a90789b323111761db/netcdf4-1.7.4-cp314-cp314t-win_arm64.whl", hash = "sha256:034220887d48da032cb2db5958f69759dbb04eb33e279ec6390571d4aea734fe", size = 2531682, upload-time = "2026-01-05T02:27:37.062Z" },
+]
+
+[[package]]
 name = "networkx"
 version = "3.6.1"
 source = { registry = "https://pypi.org/simple" }
@@ -1456,6 +1518,18 @@ wheels = [
 ]
 
 [[package]]
+name = "pandas-stubs"
+version = "3.0.3.260530"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/3d/aa/c41a8a0ff86fd85dbb3ec0c1f3fa488ca64a8b5f82654ae1b07d84acefe5/pandas_stubs-3.0.3.260530.tar.gz", hash = "sha256:d1efe47b2e5a312c047d7feabec5cb7a55365747983420077e9fcbe9ab74f714", size = 113183, upload-time = "2026-05-30T17:47:40.34Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/0b/e0/99ec5b02203c4e9ce878bc63d8caa06ac1f891e4d63bded9a5ced70fcb4f/pandas_stubs-3.0.3.260530-py3-none-any.whl", hash = "sha256:a6277eb1c8cebf48d9b2413fcd2e9a6b4ff479c934a223c29eacbc3058c4cb55", size = 173780, upload-time = "2026-05-30T17:47:39.13Z" },
+]
+
+[[package]]
 name = "pandocfilters"
 version = "1.5.1"
 source = { registry = "https://pypi.org/simple" }