Data Loading

Utilities

from_feather(file_name)[source]

Load a (theta, x) pair from a feather file.

Parameters:

  • file_name (Path) – Path to the .feather file.

  • nuisance_filter – fnmatch patterns for parameters to exclude from theta. None returns all parameters.

Return type:

tuple[ndarray[tuple[Any, ...], dtype[float32]], ndarray[tuple[Any, ...], dtype[float32]]]

Returns:

Tuple of (theta, x) as float32 numpy arrays.

Raises:

FileNotFoundError – If file_name does not exist.

to_feather(file_name, theta_values, x_values)[source]

Write a (theta, x) pair to a feather file.

Parameters:

  • file_name (Path) – Destination path. Must end in .feather.

  • theta_values (ndarray[tuple[Any, ...], dtype[float32]]) – Parameter array of shape (n_samples, n_params).

  • x_values (ndarray[tuple[Any, ...], dtype[float32]]) – Observable array of shape (n_samples, n_bins).

Raises:

ValueError – If file_name does not have a .feather suffix.

Return type:

None

Lightning Module

class SBIDataModule(dataset, config)[source]

Lightning data module over a (theta, x) map-style dataset.

Accepts any Dataset that returns (theta, x) tensor pairs by index — for example a lazily-loading, memory-mapped LazyFeatherDataset, or a pre-loaded TensorDataset for small datasets.

Train/validation is a deterministic contiguous split – the first 1 - validation_fraction of rows train, the tail validates – so every rank derives identical index sets without needing to agree on a seed. Decorrelation is the sampler’s job, not the split’s.

Parameters:
setup(stage=None)[source]

Called at the beginning of fit (train + validate), validate, test, or predict. This is a good hook when you need to build models dynamically or adjust something about them. This hook is called on every process when using DDP.

Return type:

None

Parameters:

stage (str | None)

Args:

stage: either 'fit', 'validate', 'test', or 'predict'

Example:

class LitModel(...):
    def __init__(self):
        self.l1 = None

    def prepare_data(self):
        download_data()
        tokenize()

        # don't do this
        self.something = else

    def setup(self, stage):
        data = load_data(...)
        self.l1 = nn.Linear(28, data.num_classes)
train_dataloader()[source]

Training data loader

Return type:

DataLoader

val_dataloader()[source]

Validation data loader

Return type:

DataLoader