Util Functions

Miscelanious functions used by mentat-lss

mentat_lss.utils.calc_avg_loss(emulator, data_loader, loss_function: callable, bin_idx=None, mode='galaxy_ps')

run thru the given data set and returns the average loss value for a given sub-network, or all sub-networks in a list

Parameters:

• emulator (ps_emulator) – emulator object to calculate the average loss with

• data_loader (dataLoader) – Pytorch DataLoader object containing the data to loop over

• loss_function (callable) – loss function to use

• bin_idx (list, optional) – [ps, z] values to calculate the average loss for. If None, recuresively calls this function with all possible values of ps and z. Defaults to None

• mode (str, optional) – which type of network to calculate the loss for (CURRENTLY HAS TO BE “galaxy_ps”!). Defaults to “galaxy_ps”.

Returns:

average loss corresponding to the net with bin_idx, or list of average loss for all sub-networks.

Return type:

total_loss (float or torch.Tensor)

mentat_lss.utils.delta_chi_squared(predict: Tensor, target: Tensor, invcov: Tensor, normalized=False)

Calculates the delta chi squared of the given inputs, which is given by the equation, delta_chi2 = (predict - target)^T * invcov * (predict - target).

Depending on whether the inputs are normalized, this function expects different shapes and cauclates delta chi2 differenty. In either case however, the above equation applies.

Parameters:

• predict (torch.Tensor) – output of the emulator. Should have shape [b, 1, nl*nk] OR [nps, nz, nk, nl]

• target (torch.Tensor) – data from the training / validation / test set. Should have shape [b, 1, nl*nk] OR [nps, nz, nk, nl]

• invcov (torch.Tensor) – full inverse covariance matrix. Should have shape (z, nps*nl*nk, nps*nl*nk). Is only used if normalized == False

• normalized (bool, optional) – Whether or not predict and target are normalized. Defaults to False.

Raises:

ValueError – if predict and target have different or unexpected shapes

Returns:

delta_chi2 of the batch of inputs

Return type:

chi2 (torch.Tensor)

mentat_lss.utils.get_full_invcov(cov: Tensor, num_zbins: int)

Calculates the full (multi-tracer) inverse covariance matrix given

Parameters:

• cov (torch.Tensor) – set of covariance matrices. Should have shape (num_zbins, X, X)

• num_zbins (int) – number of redshift-bins.

Returns:

Inverse covariance matrices. Has shape (num_zbins, X, X)

Return type:

invcov (torch.Tensor)

mentat_lss.utils.get_gaussan_priors(cosmo_dict: dict)

interprets the input cosmo_dict and outputs a list of scipy.stats.norm objects for use by Nautilus

Parameters:

cosmo_dict (dict) – dictionary of cosmology + nuisance parameter values and ranges

Returns:

list of scipy.stats.norm objects corresponding to the Gaussian priors of prior_names prior_names (list): names of the corresponding priors in priors

Return type:

priors (list)

mentat_lss.utils.get_invcov_blocks(cov: Tensor, num_spectra: int, num_zbins: int, num_kbins: int, num_ells: int)

Calculates block inverse covariance matrices given a full multi-tracer covariance matrix.

Parameters:

• cov (torch.Tensor) – full block covariance matrix. Should have shape (num_zbins, num_spectra*num_kbins*num_ells, num_spectra*num_kbins*num_ells)

• num_spectra (int) – number of (auto + cross) power spectra in the corresponding covariance matrix

• num_zbins (int) – number of redshift bins in the corresponding covariance matrix

• num_kbins (int) – number of k-mode bins in the corresponding covariance matrix

• num_ells (int) – number of multipole moments in the corresponding covariance matrix

Returns:

Set of inverse block covariance matrices. Has shape (num_spectra, num_zbins, num_ells*num_kbins, num_ells*num_kbins)

Return type:

invcov_blocks (torch.Tensor)

mentat_lss.utils.get_parameter_ranges(cosmo_dict: dict)

Returns cosmology and bias parameter priors based on the input cosmo_dict

Parameters:

cosmo_dict (dict) – dictionary of cosmology + nuisance parameter values and ranges

Returns:

name of parameters that are varied in the input cosmo_dict priors (np.array): min and max bounds of parameters varied in cosmo_dict. Has shape [len(param_names), 2]

Return type:

params (list)

mentat_lss.utils.hyperbolic_chi2_loss(predict: Tensor, target: Tensor, invcov: Tensor, normalized=False)

Calculates the hyperbolic delta chi2 of the given inputs, which is given by the equation L = <sqrt(1 + 2(delta_chi_squared)> - 1

Parameters:

• predict (torch.Tensor) – output of the network

• target (torch.Tensor) – (batch of) elements in the training set. Should have the same shape as predict

• invcov (torch.Tensor) – either a full or block inverse covariance matrix, depending on whether the input is normalizeed or not

• normalized (bool, optional) – whether or not the inputs are normalized. Defaults to False.

Returns:

mean hyperbolic chi2 of the given batch of inputs

Return type:

hyperbolic_chi2 (torch.Tensor)

mentat_lss.utils.hyperbolic_loss(predict, target, **args)

Calculates the hyperbolic loss of the inputs given by <sqrt(1 + 2(predict - target)**2)> - 1

Parameters:

• predict (torch.Tensor) – output of the network

• target (torch.Tensor) – (batch of) elements in the training set. Should have the same shape as predict

• **args – extra arguments (needed by interface of ps_emulator)

Returns:

hyperbolic loss of the given inputs

Return type:

hyperbolic_loss

mentat_lss.utils.is_in_hypersphere(priors, params)

Returns whether or not the given params are within a hypersphere with edges defined by bounds

mentat_lss.utils.load_config_file(config_file: str)

loads in the emulator config file as a dictionary object

Parameters:

config_file – Config file path and name to laod

Raises:

IOError – If config_file could not be read in

mentat_lss.utils.make_hypersphere(priors: array, dim: int, N: int)

Generates a hypersphere of N samples using the method from https://arxiv.org/abs/2405.01396v1

Parameters:

• priors – (np.array): array of parameter minima and maxima. Should have shape (dim, 2)

• dim (int) – number of parameters to generate a hypersphere with

• N (int) – number of samples to uniformly generate within the hypersphere

Returns:

list of parameters uniformly sampled within a hypersphere. Has shape (N, dim)

Return type:

sphere_points (np.array)

mentat_lss.utils.make_latin_hypercube(priors, N)

Generates a latin hypercube of N samples with lower and upper bounds given by priors

mentat_lss.utils.mse_loss(predict: Tensor, target: Tensor, **args)

Calculates the mean-squared-error loss of the inputs.

Parameters:

• predict (torch.Tensor) – output of the network

• target (torch.Tensor) – (batch of) elements in the training set. Should have the same shape as predict

• **args – extra arguments (needed by interface of ps_emulator)

Returns:

mean-squared-error loss of the given inputs

Return type:

mse_loss

mentat_lss.utils.normalize_cosmo_params(params: Tensor, normalizations: Tensor)

Linearly normalizes input cosmology + bias parameters to lie within the range [0,1]

Parameters:

• params (torch.Tensor) – batch of input parameters to normalize. Should have shape [batch, num_spectra*num_zbins, num_cosmo_params + (num_nuisance_params)].

• normalizations (torch.Tensor) – Tensor of parameter minima and maxima. Should have shape [2, num_spectra*num_zbins, num_cosmo_params + (num_nuisance_params)

Returns:

batch of normalized input parameters. has shape [batch, num_spectra*num_zbins, num_cosmo_params + (num_nuisance_params)

Return type:

norm_params (torch.Tensor)

mentat_lss.utils.normalize_power_spectrum(ps_raw: Tensor, ps_fid: Tensor, sqrt_eigvals: Tensor, Q: Tensor)

Normalizes the given galaxy power spectrum multipoles using the method described in http://arxiv.org/abs/2403.12337

Parameters:

• ps_raw (torch.Tensor) – batch of power spectra in units of (Mpc/h)^3 to normalize. Should have shape [b, nps, z, nk*nl]

• ps_fid (torch.Tensor) – fiducial power spectrum multipoles in units of (Mpc/h)^3 used for normalization. Should have shape [nps, z, nk*nl]

• sqrt_eigvals (torch.Tensor) – set of sqrt eigenvalues used for normalization. Should have shape [ps, z, nk*nl]

• Q (torch.Tensor) – set of eigenvectors used for normalization. Should have shape [ps, z, nk*nl, nk*nl]

Returns:

normalized power spectrum multipoles. Has shape [b, nps, z, nk*nl]

Return type:

ps_norm

mentat_lss.utils.organize_training_set(training_dir: str, train_frac: float, valid_frac: float, test_frac: float, param_dim, num_zbins, num_spectra, num_ells, k_dim, remove_old_files=True)

Takes a set of matrices and reorganizes them into training, validation, and tests sets

Parameters:

• training_dir – Directory contaitning matrices to organize

• train_frac – Fraction of dataset to partition as the training set

• valid_frac – Fraction of dataset to partition as the validation set

• test_frac – Fraction of dataset to partition as the test set

• param_dim – Dimension of input parameter arrays

• mat_dim – Dimention of power spectra

• remove_old_files – If True, deletes old data files after loading data into memory and before re-organizing. Default True.

mentat_lss.utils.prepare_emu_inputs(sample: dict, cosmo_dict: dict, num_tracers: int, num_zbins: int, required_emu_params: dict)

takes a set of parameters and oragnizes them to the format expected by mentat-lss

Parameters:

• sample (dict) – dictionary of (param_name, param_value)

• cosmo_dict (dict)

• num_tracers (int) – number of correlated tracers to calculate

• num_zbins (int) – number of independent redshift bins to calculate

• required_emu_params (dict) – dictionary of parameter names required by the emulator

Returns:

1D list of parameters that can be directly passed to mentat-lss

Return type:

param_vector (np.array)

mentat_lss.utils.prepare_ps_inputs(sample: dict, cosmo_dict: dict, num_tracers: int, num_zbins: int)

takes a set of parameters and oragnizes them to the format expected by ps_theory_calculator

Parameters:

• sample (dict) – dictionary of (param_name, param_value)

• cosmo_dict (dict)

• num_tracers (int) – number of correlated tracers to calculate

• num_zbins (int) – number of independent redshift bins to calculate

Returns:

1D list of parameters that can be directly passed to ps_theory_calculator

Return type:

param_vector (np.array)

mentat_lss.utils.un_normalize_power_spectrum(ps_raw: Tensor, ps_fid: Tensor, sqrt_eigvals: Tensor, Q: Tensor, Q_inv: Tensor)

Reverses normalization of a batch of output power spectru based on the method developed by http://arxiv.org/abs/2403.12337

Parameters:

• ps (torch.Tensor) – power spectrum to reverse normalization. Expected shape is either [nb, nps, nz, nk*nl] or [nb, 1, nk*nl]

• ps_fid (torch.Tensor) – fiducial power spectrum used to reverse normalization. Expected shape is [nps*nz, nk*nl]

• sqrt_eigvals (torch.Tensor) – square root eigenvalues of the inverse covariance matrix. Expected shape is [nps*nz, nk*nl]

• Q (torch.Tensor) – eigenvectors of the inverse covariance matrix. Expected shape is [nps*nz, nk*nl, nk*nl]

• Q_inv (torch.Tensor) – inverse eigenvectors of the inverse covariance matrix. Expected shape is [nps*nz, nk*nl, nk*nl]

• net_idx (torch.Tensor) – (optional) index specifying the specific sub-network output to reverse normalization. Default None. If not specified, will reverse normalization for the entire emulator output

Returns:

galaxy power spectrum multipoles in units of (Mpc/h)^3 in the same shape as ps

Return type:

ps_new (torch.Tensor)

Raises:

IndexError – If the given shape of ps_raw is invalid.