snnlab.sim.datasets

def load_dataset(
    name,
    max_samples=None,
    split=False,
    dt_ms=None,
    t_ms=None,
    evaluation_split="validation",
    evaluation_only=False,
):
    """Load full dataset as (X, y) numpy arrays in [0, 1] / int64.

    Args:
        name: "mnist" (dense pixel rows) or "shd" (object array of events)
        max_samples: optional cap on the official training partition; the
               official MNIST test partition is never subsampled through this
               argument
        split: if True, return (X_train, X_eval, y_train, y_eval)
        evaluation_split: for MNIST, select either the deterministic validation
               split used during training or the untouched official test split.
               SHD retains its official train/test contract.
        dt_ms, t_ms: currently unused; kept for call-site compatibility
        evaluation_only: skip the unused MNIST training partition; requires
               split=True and evaluation_split='test'. Training outputs are None.

    Single canonical loader used by train, infer, and image paths so
    "first digit-0 sample" means the same physical sample everywhere.
    """
    if evaluation_only and (not split or evaluation_split != "test"):
        raise ValueError("evaluation_only requires split=True and evaluation_split='test'")
    if name == "mnist":
        from torchvision import datasets, transforms

        mnist_test = datasets.MNIST(
            root="/tmp/mnist",
            train=False,
            download=True,
            transform=transforms.ToTensor(),
        )
        X_test = mnist_test.data.numpy().reshape(-1, 784).astype(np.float32) / 255.0
        y_test = mnist_test.targets.numpy().astype(np.int64)
        if evaluation_only:
            return None, X_test, None, y_test
        mnist_train = datasets.MNIST(
            root="/tmp/mnist",
            train=True,
            download=True,
            transform=transforms.ToTensor(),
        )
        X = mnist_train.data.numpy().reshape(-1, 784).astype(np.float32) / 255.0
        y = mnist_train.targets.numpy().astype(np.int64)
    elif name == "shd":
        # Event data. Keep SHD's OFFICIAL train/test split — it holds out two
        # speakers unseen in training, and every published SHD number is reported
        # on that test set, so resplitting would break
        # comparability. Returns object arrays of per-sample (units, times); the
        # caller bins them to spikes lazily at the run's dt.
        X_tr, y_tr = _read_shd_split("train")
        X_te, y_te = _read_shd_split("test")
        X_tr, y_tr = _subsample(X_tr, y_tr, max_samples)
        X_te, y_te = _subsample(X_te, y_te, max_samples)
        if split:
            return X_tr, X_te, y_tr, y_te
        return np.concatenate([X_tr, X_te]), np.concatenate([y_tr, y_te])
    else:
        raise ValueError(f"Unknown dataset: {name}")

    if max_samples is not None and max_samples < len(X):
        idx = np.random.RandomState(MNIST_SPLIT_SEED).choice(
            len(X), max_samples, replace=False
        )
        X, y = X[idx], y[idx]

    if split:
        if evaluation_split == "test":
            return X, X_test, y, y_test
        from sklearn.model_selection import train_test_split

        validation_size = max(
            int(np.ceil(len(y) * MNIST_VALIDATION_FRACTION)),
            len(np.unique(y)),
        )
        X_train, X_validation, y_train, y_validation = train_test_split(
            X,
            y,
            test_size=validation_size,
            random_state=MNIST_SPLIT_SEED,
            stratify=y,
        )
        if evaluation_split == "validation":
            return X_train, X_validation, y_train, y_validation
        raise ValueError(
            "MNIST evaluation_split must be 'validation' or 'test', "
            f"got {evaluation_split!r}"
        )
    return X, y