gpmp.dataloader module

The gpmp.dataloader module provides optional containers and finite mini-batch iterators for observations. Most GPmp routines can be called directly with arrays xi, zi. Dataloaders are useful when observations are already split into shards or when a covariance-parameter selection criterion should be evaluated from batches.

Data model

Dataset stores observation points and scalar or vector observations:

from gpmp.dataloader import Dataset, DataLoader

dataset = Dataset(xi, zi)
loader = DataLoader(dataset, batch_size=100, shuffle=False)

xi has shape (n, d). zi has shape (n,), (n, 1), or another first dimension equal to n. Dataset also accepts lists of arrays. In that case, each list element is treated as a shard and samples are fetched without concatenating all shards first.

DataLoader iterates over a Dataset and yields pairs (x_batch, z_batch). The arrays use the active gpmp.num backend. If batch_size=None, one iteration contains the full dataset.

Use in parameter selection

Selection methods accept either explicit arrays xi, zi or a dataloader. Do not pass both in the same call.

model, info = gp.kernel.select_parameters_sigma2_rho_with_reml(
    model, dataloader=loader, info=True
)

With a dataloader, a criterion evaluation is computed from batches. For a loader with batches \(b_\ell\) and batch sizes \(n_\ell\), the batch-averaged criterion has the form

\[\overline J(\theta) = \frac{\sum_\ell n_\ell J_\ell(\theta)} {\sum_\ell n_\ell},\]

where \(J_\ell(\theta)\) is the criterion evaluated on batch \(b_\ell\).

Some selection functions expose batches_per_eval. With batches_per_eval=0, one criterion call consumes the whole loader. With a positive value, one criterion call uses only that many successive batches and cycles through the loader across calls. This gives a stochastic or semi-stochastic criterion and should be used with care when interpreting optimizer convergence.

Iterator options

DataLoader supports:

  • batch_size: number of samples per batch.

  • shuffle: whether to shuffle indices at each epoch.

  • drop_last: whether to drop an incomplete last batch.

  • seed and set_epoch: deterministic shuffling across epochs.

  • infinite: cycle indefinitely.

When reproducibility matters, set the GPmp seed through gpmp.num.set_seed and pass a seed to DataLoader if shuffling is enabled.

Dataset

class gpmp.dataloader.Dataset(x: ndarray[tuple[Any, ...], dtype[floating]] | List[ndarray[tuple[Any, ...], dtype[floating]]], z: ndarray[tuple[Any, ...], dtype[floating]] | List[ndarray[tuple[Any, ...], dtype[floating]]])[source]

Dataset storing covariates x and observations z.

x and z may each be a single array or a list of arrays (shards) that share the same first-dimension length.

Shards are retained at construction and indexing is performed lazily, without concatenation, with O(log(#shards)) index lookup.

static k_fold_indices(n_samples: int, n_splits: int, seed: int | None = None) List[Tuple[ndarray[tuple[Any, ...], dtype[floating]], ndarray[tuple[Any, ...], dtype[floating]]]][source]

Return exactly k (train, val) index tuples for k-fold CV.

Each split has approximately the same size.

static repeated_k_fold_indices(n_samples: int, n_splits: int, n_repeats: int, seed: int | None = None) List[Tuple[ndarray[tuple[Any, ...], dtype[floating]], ndarray[tuple[Any, ...], dtype[floating]]]][source]

Return n_repeats × k shuffled k-fold splits.

Each repetition is independently shuffled.

static split(dataset: Dataset, ratios: Tuple[float, float, float] = (0.8, 0.1, 0.1), seed: int | None = None) Tuple[Dataset, Dataset, Dataset][source]

Return (train, val, test) datasets according to ratios.

Samples are randomly shuffled before splitting.

subset(indices: ndarray[tuple[Any, ...], dtype[floating]]) Dataset[source]

Return a dataset restricted to indices.

Shard structure is preserved if possible.

DataLoader

class gpmp.dataloader.DataLoader(dataset: Dataset, batch_size: int | None = None, shuffle: bool = True, drop_last: bool = False, seed: int | None = None, infinite: bool = False)[source]

Mini-batch generator over a Dataset.

Supports optional shuffling, infinite iteration, and deterministic seeding across epochs for reproducibility (especially useful with distributed training).

Fetching is shard-aware and avoids full dataset concatenation.

Behaviour when the last batch is incomplete:

  • drop_last=True – discard it.

  • drop_last=False – yield a smaller batch.

Notes

  • Call set_epoch() before each epoch to ensure deterministic shuffling across distributed workers.

  • If infinite=True the iterator cycles forever (useful for GANs).

reduce_mean(func) ndarray[tuple[Any, ...], dtype[floating]][source]

Compute the weighted mean of func(x_batch, z_batch) over batches.

Each batch output is weighted by its batch size.

Parameters:

func (callable) – A function (x_batch, z_batch) -> scalar or array (batch result).

Returns:

mean – Weighted mean over all samples.

Return type:

Array

set_epoch(epoch: int) None[source]

Manually set current epoch (affects shuffling).

Scaling objects

class gpmp.dataloader.Normalizer(mean: ndarray[tuple[Any, ...], dtype[floating]], std: ndarray[tuple[Any, ...], dtype[floating]])[source]

Standardize covariates to zero mean and unit variance.

class gpmp.dataloader.RobustScaler(median: ndarray[tuple[Any, ...], dtype[floating]], iqr: ndarray[tuple[Any, ...], dtype[floating]])[source]

Scale covariates by median and interquartile range.

class gpmp.dataloader.MinMaxScaler(x_min: ndarray[tuple[Any, ...], dtype[floating]], x_max: ndarray[tuple[Any, ...], dtype[floating]])[source]

Rescale covariates to lie between 0 and 1.

class gpmp.dataloader.ObservationScaler(mean: float, std: float)[source]

Standardise scalar observations to zero mean and unit variance.