Inference

class InferenceHandler(prior_path)

Parameters:

prior_path (Path)

set_dataset(data_folder)

Point the handler at a folder of .feather simulation files.

Parameters:

data_folder (Path) – Directory containing .feather files.

Return type:

None

load_training_data(verbose=True)

…

Return type:

None

Parameters:

verbose (bool)

fit_x_compressor(compressor, **kwargs)

Compress X dim

Parameters:

compressor (str)

fit_theta_compressor(compressor, **kwargs)

Compress theta dim

Parameters:

compressor (str)

create_posterior(config)

Build the NPE inference object and density estimator network.

Only the kwargs that the chosen model family actually accepts are forwarded to posterior_nn; unsupported kwargs (e.g. num_blocks for zuko-backed flows) are dropped with a DEBUG log line rather than raising a TypeError at runtime.

Parameters:

config (PosteriorConfig) – Architecture and hyperparameter settings.

Return type:

None

train_posterior(config, model_config=None)

Train the density estimator from scratch using PyTorch Lightning.

Requires load_training_data() and create_posterior() to have been called first.

Parameters:

• config (TrainingConfig) – Training loop settings.

• model_config (PosteriorConfig | None) – Architecture config embedded in every checkpoint.

Raises:

ValueError – If training data or the NPE object are missing.

Return type:

None

get_log_likelihood(theta, x, **kwargs)

Evaluate the log-likelihood of theta given observed data x.

Parameters:

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

• x (list[float] | ndarray) – Observed data vector x_o.

Return type:

Tensor

Returns:

Log-probability tensor of shape (n_samples,).

load_posterior(checkpoint_path)

Load a trained density estimator from a checkpoint for inference only.

The PosteriorConfig is read from the checkpoint’s "model_config" key. _build_posterior_nn_kwargs filtering applies, so loading a zuko checkpoint works even if num_blocks is present in the stored config (it will simply be dropped).

Parameters:

checkpoint_path (Path) – Path to a .pt / .ckpt checkpoint.

Raises:

• FileNotFoundError – If checkpoint_path does not exist.

• ValueError – If no model config can be determined.

Trainer

lightning_module

alias of <module ‘mach3sbitools.inference.lightning_module’ from ‘/home/runner/work/MaCh3SbiTools/MaCh3SbiTools/src/mach3sbitools/inference/lightning_module.py’>