Emulator API

This page contains the main functions you need to interact with mentat-lss. Specifically, these functions create and use the galaxy power spectrum emulator.

mentat_lss.emulator.compile_multiple_device_training_results(save_dir: str, config_dir: str, num_gpus: int)

takes networks saved on seperate ranks and combines them to the same format as when training on one device

Parameters:

• save_dir (string) – base save directory, where each rank was saved in its own sub-directory

• config_dir (string) – path+name of the original network config file

• num_gpus (int) – number of gpus to compile results of

Returns:

emulator object with all training data combined together.

Return type:

full_emulator (ps_emulator)

class mentat_lss.emulator.ps_emulator(net_dir: str, mode: str = 'train', device: device = None)

Class defining the neural network emulator.

check_kbins_are_compatible(test_kbins: array)

Tests whether the passed test_kbins is the same as the emulator k-bins

Parameters:

test_kbins (np.array) – k-array to check

Returns:

Whether or not the given k-bins are compatible

Return type:

is_compatible (bool)

get_emulated_power_spectrum(params, extrapolate: bool = False, raw_output: bool = False)

Gets the power spectra corresponding to the given input params by passing them though the emulator

Parameters:

• params – 1D or 2D numpy array, torch Tensor, or dictionary containing a list of cosmology + galaxy bias parameters.

• array (if params is a 2D)

• simultaniously (this function generates a batch of power spectra)

• extrapolate (bool) – Whether or not to pass through the emulator if the given input parameters are outside the range it was trained on.

• False (Default)

• raw_output – bool specifying whether or not to return the raw network output without undoing normalization. Default False

Returns:

emulated galaxy power spectrum multipoles (P_tree + P_1loop). If given a batch of parameters, has shape [nb, nps, nz, nk, nl]. Otherwise, has shape [nps, nz, nk, nl]. If extrapolate is false and the given input parameters are out of bounds, then this function returns an array of all zeros.

Return type:

galaxy_ps (np.array)

get_power_spectra(params, extrapolate: bool = False, raw_output: bool = False)

Gets the full galaxy power spectrum multipoles (emulated and analytically calculated)

Parameters:

• params – 1D or 2D numpy array, torch Tensor, or dictionary containing a list of cosmology + galaxy bias parameters. if params is a 2D array, this function generates a batch of power spectra simultaniously

• extrapolate (bool) – Whether or not to pass through the emulator if the given input parameters are outside the range it was trained on. Default False

• raw_output (bool) – Whether or not to return the raw network output without undoing normalization. Default False

Returns:

Emulated galaxy power spectrum multipoles. If raw_output = False, has shape [nps, nz, nk, nl] or [nb, nps, nz, nk, nl]. Else has shape [nb, nps, nz, nk*nl]

Return type:

galaxy_ps (np.array)

get_required_analytic_parameters()

Returns a list of input parameters used by our analytic eft model, not directly emulated.

NOTE: These parameters are currently hard-coded.

Returns:

list of input (counterterm + stoch) parameters.

Return type:

required_analytic_params (list)

get_required_emu_parameters()

Returns a list of input parameters needed by the emulator.

Currently, mentat-lss requires input parameters to be in the same order as given by the return value of this function. For example. If the return list is [‘h’, ‘omch2’], you should pass in [h, omch2] to get_power_spectra in that order.

Returns:

list of input cosmology + bias parameters required by the emulator.

Return type:

required_emu_params (list)

load_data(key: str, data_frac=1.0, return_dataloader=True, data_dir='')

loads and returns the training / validation / test dataset into memory

Parameters:

• key – one of [“training”, “validation”, “testing”] that specifies what type of data-set to laod

• data_frac – fraction of the total data-set to load in. Default 1

• return_dataloader – Determines what object type to return the data as. Default True If true: returns data as a pytorch.utils.data.DataLoader object. If false: returns data as a pk_galaxy_dataset object.

• data_dir – location of the data-set on disk. Default “”

Returns:

The desired data-set in either a pk_galaxy_dataset or DataLoader object.

Return type:

data

Raises:

• KeyError – If key is an incorrect value.

• ValueError – If some property of the loaded dataset does not match with the emulator.

load_trained_model(path)

loads the pre-trained network from file into the current model, as well as all relavent information needed for normalization. This function is called by the constructor, but can also be called directly by the user if desired.

Parameters:

path – The directory+filename of the trained network to load.