Workflow Example

This page walks through an expected workflow using mentat-lss, split up into generating a training set, training the network, and testing / using the network

Generating a Training Set

This stage is where you decide what power spectrum model you want to emulate, and the corresponding parameter ranges it will be valid in. Assuming you have access to ps_1loop, we have provided a script in scripts/make_training_set_ps_1loop.py that will attempt to create a training set in the right format for you.

If you are creating your own training set with a different model, the following are the required ingredients and format you should use:

  • Cosmology and survey parameters config files. We provide some example files in the configs directory.

  • Your data should be stored in pk-training.npz, pk-validation.npz, and pk-testing.npz files, each of which have:

    • a params array with shape [N, num_params].

    • a galaxy power spectrum array with shape [N, num_auto_plus_cross_spectra, num_redshift_bins, num_kbins, num_ells].

  • You should have a ‘ps_properties.npz’ with the following properties:

    • ells: the multipoles emulated (ex: [0, 2])

    • ndens: the number densities of each tracer in each redshift bin (shape [num_tracers, num_redshift_bins])

    • z_eff: the effective redshift of each redshift bin (shape [num_redshift_bins])

    • k_emu: the k-centers used to generate the training set power spectra (shape [num_kbins])

  • You should have a cov.npy file with a valid covariance matrix (ex: from CovaPT or TheCov).

Training the Emulator

One you have a training set, you’ll need to provide a configuration file specifying the specific network architecture you want to use, as well as all relavent hyperparameters. The following is an example of such a file.

# directory that all other paths are relative to
# Ex: the repository directory
input_dir: <path_to_your_directory>
# Where to save the network
save_dir: <path_to_save_directory>
# Where your training set lives
training_dir: <path_to_training_set_directory>
# Location of config file with cosmology + bias parameters
# This file contains all of the parameter ranges
cosmo_dir : <path/to/cosmo_config.yaml>

model_type : stacked_transformer
loss_type : hyperbolic_chi2

num_cosmo_params    : 5
num_nuisance_params : 3 # <- per tracer
num_tracers : 2
num_zbins   : 2
num_ells    : 2
num_kbins   : 25

# specifications are for each network - will be repeated for each sample / redshift bin
galaxy_ps_emulator:
    # mlp parameters
    num_mlp_blocks      : 4
    num_block_layers    : 2
    use_skip_connection : True

    # transformer parameters
    num_transformer_blocks : 1
    split_dim              : 5
    split_size             : 20

# Training parameters
num_epochs: 500
galaxy_ps_learning_rate: 0.005
batch_size: 1000
training_set_fraction : 1.0
early_stopping_epochs: 25
weight_initialization: He
optimizer_type : Adam

# whether to re-calculate the training-set loss at the end of each epoch
# Setting to true gives more accurate loss stats, but is slower
# Validation set loss is not changed by this option
recalculate_train_loss : False
# whether to attempt training on GPU
# NOTE: Small networks will possibly train faster on CPU!
use_gpu : True

Here are some important considerations to make before training:

  • You’ll need to decide whether to try training your network on a CPU or GPU (assuming one is available to you). In general, networks with transformers train significantly faster on GPUs, so we recommend you try training on GPUs whenever possible. By defualt, mentat-lss will attempt to use a GPU if it is available.

  • The specific binning information (num_zbins, num_ells, etc) must match those found in the training set. The code will throw an error on startup if they don’t.

  • The above specifications are the optimized values found in Adamo et al (2026). The optimal setup will potentially be different for your case, but these values should provide a good starting point.

Once you have your configs all sorted, you can train you network using the following script,

from mentat_lss.emulator import pk_emulator
import mentat_lss.training_loops as training_loops
import logging

config_file = "/path/to/config_file.yaml"

# Used for printing output during training. If you don't want any
# output. set to logging.WARNING
logging.basicConfig(level = logging.INFO)

t1 = time.time()
emulator = pk_emulator(config_file, "train")

# train on a single cpu / gpu
training_loops.train_on_single_device(emulator)

We have also provided a more robust script in scripts/train_emulator.py that also handles training on multiple GPUs. For more details, see Training on GPU(s).

During the actual training process, mentat-lss will loop through each subnet, each of which correspond to a single tracer / redshift bin. It will then print out the average training set and validation set loss values, as well as the number of epochs elapses since the validation loss improved.:

`Net idx : [ps, z], epoch: N, avg train loss: l1, avg validation loss: l2 (epochs_since_improved)`

This will repeat until either the validation loss for all sub-nets hasn’t improved for 25 epochs, or if max_epochs is reached.

Testing the Emulator

We provide an example jupyter notebook for running various tests on your emulator here.

Using the Emulator

Finally, once you are sure your emulator works, you can generate power spectrum with,

emulator = ps_emulator(emu_dir, "eval")
pk_predict = ps_emulator.get_power_spectra(input_params)

which will output power spectrum multipoles as a numpy array with shape [nps, nz, nk, nl]. You can access the required input parameters (and their order) with,

# REQUIRED: input_params should be in the same order as emu_params
emu_params = ps_emulator.get_required_emu_params()
# OPTIONAL: should be concatenated at the end of emu_params
analytic_params = ps_emulator.get_required_analytic_parameters()

Note that if you include analytic_params, the code will compute the EFT counterterms and stochastic contributions. Under the hood, the code calls symbolic_pofk to quickly calculate the linear matter power spectrum.

You can then hook up this method to your favorite MCMC sampler to run some likelihood analyses!