.. _workflow: 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. .. code-block:: yaml # directory that all other paths are relative to # Ex: the repository directory input_dir: # Where to save the network save_dir: # Where your training set lives training_dir: # Location of config file with cosmology + bias parameters # This file contains all of the parameter ranges cosmo_dir : # Can be one of 'stacked_transformer' or 'combined_tracer_transformer' model_type : stacked_transformer loss_type : hyperbolic_chi2 num_cosmo_params : 5 num_nuisance_params : 3 # <- per tracer # 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 code will read in binning information from `ps_properties.npz` in your training set directroy. Alternatively, you can specify `num_kbins`, `num_zbins`, `num_tracers`, and `num_ells` in the config file. We recommend you use the first option as it's less error prone. - 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. - There are two major architecture options to choose from, those being: - `stacked_transformer`: this architecture assigns a seperate network for each tracer and redshift bin, which in our testing performs best, but also is more expensive to train. This is the default architecture in `mentat-lss`. - `combined_tracer_transformer`: this option uses two networks per redshift bin, with one handling the auto spectra and the other the cross spectra. This setup is faster to train but performs worse than stacked_transformer in our testing. If you have a large number of tracers and redshift bins, this option may be more feasible for you. .. _`Adamo et al (2026)`: https://arxiv.org/abs/2603.16003 Once you have your configs all sorted, you can train you network using the following script, .. code-block:: python 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 :doc:`tutorials/training`. 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. Optimizing the Emulator (NEW) --------------------- As of version 1.1, we have added Optuna support for optimizing your emulator! There is a new script in `scripts/optimize_with_optuna.py` that will run an Optuna project that performs hyperparameter optimization far more efficiently than a basic grid search. Testing the Emulator -------------------- We provide an example jupyter notebook for running various tests on your emulator :doc:`here `. Using the Emulator ------------------- Finally, once you are sure your emulator works, you can generate power spectrum with, .. code-block:: python 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, .. code-block:: python # 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. .. _symbolic_pofk: https://github.com/DeaglanBartlett/symbolic_pofk You can then hook up this method to your favorite MCMC sampler to run some likelihood analyses!