Index

A

activation_function (class in mentat_lss.models.blocks)

analytic_eft_model (class in mentat_lss.models.analytic_terms)

B

block_addnorm (class in mentat_lss.models.blocks)

block_resnet (class in mentat_lss.models.blocks) block_transformer_encoder (class in mentat_lss.models.blocks)

C

calc_avg_loss() (in module mentat_lss.utils) calculate_pk_lin() (mentat_lss.models.analytic_terms.analytic_eft_model method)

check_kbins_are_compatible() (mentat_lss.emulator.ps_emulator method) compile_multiple_device_training_results() (in module mentat_lss.emulator)

D

delta_chi_squared() (in module mentat_lss.utils)

F

forward() (mentat_lss.models.blocks.activation_function method) (mentat_lss.models.blocks.block_addnorm method) (mentat_lss.models.blocks.block_resnet method) (mentat_lss.models.blocks.block_transformer_encoder method) (mentat_lss.models.blocks.linear_with_channels method) (mentat_lss.models.blocks.multi_headed_attention method) (mentat_lss.models.stacked_mlp.single_mlp method) (mentat_lss.models.stacked_mlp.stacked_mlp method) (mentat_lss.models.stacked_transformer.single_transformer method) (mentat_lss.models.stacked_transformer.stacked_transformer method)

G

get_analytic_terms() (mentat_lss.models.analytic_terms.analytic_eft_model method) get_ctr_terms() (mentat_lss.models.analytic_terms.analytic_eft_model method) get_damping_factor() (mentat_lss.models.analytic_terms.analytic_eft_model method) get_emulated_power_spectrum() (mentat_lss.emulator.ps_emulator method) get_full_invcov() (in module mentat_lss.utils) get_gaussan_priors() (in module mentat_lss.utils) get_invcov_blocks() (in module mentat_lss.utils) get_k_nl() (mentat_lss.models.analytic_terms.analytic_eft_model method) get_normalized_galaxy_power_spectra() (mentat_lss.dataset.pk_galaxy_dataset method)

get_parameter_ranges() (in module mentat_lss.utils) get_pk_lin() (mentat_lss.models.analytic_terms.analytic_eft_model method) get_pk_lin_irres_rsd() (mentat_lss.models.analytic_terms.analytic_eft_model method) get_power_spectra() (mentat_lss.emulator.ps_emulator method) get_required_analytic_parameters() (mentat_lss.emulator.ps_emulator method) get_required_emu_parameters() (mentat_lss.emulator.ps_emulator method) get_stochastic_terms() (mentat_lss.models.analytic_terms.analytic_eft_model method) get_tree_term() (mentat_lss.models.analytic_terms.analytic_eft_model method) get_true_galaxy_power_spectra() (mentat_lss.dataset.pk_galaxy_dataset method)

H

hyperbolic_chi2_loss() (in module mentat_lss.utils)

hyperbolic_loss() (in module mentat_lss.utils)

I

initialize_params() (mentat_lss.models.blocks.linear_with_channels method)

is_in_hypersphere() (in module mentat_lss.utils)

L

linear_with_channels (class in mentat_lss.models.blocks) load_config_file() (in module mentat_lss.utils)

load_data() (mentat_lss.emulator.ps_emulator method) load_trained_model() (mentat_lss.emulator.ps_emulator method)

M

make_hypersphere() (in module mentat_lss.utils) make_latin_hypercube() (in module mentat_lss.utils) mentat_lss.cosmo_utils module mentat_lss.dataset module mentat_lss.emulator module mentat_lss.models.analytic_terms module mentat_lss.models.blocks module mentat_lss.models.stacked_mlp module mentat_lss.models.stacked_transformer module

mentat_lss.training_loops module mentat_lss.utils module module mentat_lss.cosmo_utils mentat_lss.dataset mentat_lss.emulator mentat_lss.models.analytic_terms mentat_lss.models.blocks mentat_lss.models.stacked_mlp mentat_lss.models.stacked_transformer mentat_lss.training_loops mentat_lss.utils mse_loss() (in module mentat_lss.utils) multi_headed_attention (class in mentat_lss.models.blocks)

N

normalize_cosmo_params() (in module mentat_lss.utils)

normalize_data() (mentat_lss.dataset.pk_galaxy_dataset method) normalize_power_spectrum() (in module mentat_lss.utils)

O

organize_parameters() (mentat_lss.models.stacked_mlp.stacked_mlp method) (mentat_lss.models.stacked_transformer.stacked_transformer method)

organize_training_set() (in module mentat_lss.utils)

P

pk_galaxy_dataset (class in mentat_lss.dataset) prepare_emu_inputs() (in module mentat_lss.utils)

prepare_ps_inputs() (in module mentat_lss.utils) ps_emulator (class in mentat_lss.emulator)

S

set_ir_resum_params() (mentat_lss.models.analytic_terms.analytic_eft_model method) set_params() (mentat_lss.models.analytic_terms.analytic_eft_model method) single_mlp (class in mentat_lss.models.stacked_mlp)

single_transformer (class in mentat_lss.models.stacked_transformer) stacked_mlp (class in mentat_lss.models.stacked_mlp) stacked_transformer (class in mentat_lss.models.stacked_transformer)

T

to() (mentat_lss.dataset.pk_galaxy_dataset method) train_galaxy_ps_one_epoch() (in module mentat_lss.training_loops) train_on_multiple_devices() (in module mentat_lss.training_loops)

train_on_single_device() (in module mentat_lss.training_loops) transpose_output() (mentat_lss.models.blocks.multi_headed_attention method) transpose_qkv() (mentat_lss.models.blocks.multi_headed_attention method)

U

un_normalize_power_spectrum() (in module mentat_lss.utils)