Index

A

AdaptiveHelpFormatter (class in neural_lam.train_model) aggregate() (neural_lam.gnn_layers.InteractionNet method) aggregate_and_plot_metrics() (neural_lam.models.module.ForecasterModule method) all_gather_cat() (neural_lam.models.module.ForecasterModule method) apply_time_crop() (in module neural_lam.utils.time) ar_steps (neural_lam.weather_dataset.WeatherDataset attribute)

ar_steps_eval (neural_lam.config.TrainRunConfig attribute) (neural_lam.weather_dataset.WeatherDataModule attribute) ar_steps_train (neural_lam.config.TrainRunConfig attribute) (neural_lam.weather_dataset.WeatherDataModule attribute) ARCHETYPE_FUNCTIONS (in module neural_lam.create_graph_with_wmg) ARForecaster (class in neural_lam.models.forecasters.autoregressive) assert_edge_index_in_range() (in module neural_lam.create_graph)

B

base_dataset (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset attribute) BaseDatastore (class in neural_lam.datastore.base) BaseGraphEFM (class in neural_lam.models.step_predictors.graph.graph_efm) BaseGraphLatentDecoder (class in neural_lam.models.latent.base_decoder) BaseGraphModel (class in neural_lam.models.step_predictors.graph.base) BaseHiGraphModel (class in neural_lam.models.step_predictors.graph.hierarchical) BaseLatentEncoder (class in neural_lam.models.latent.base_encoder) Batch (in module neural_lam.models.module) batch_size (neural_lam.config.TrainRunConfig attribute) (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset attribute) (neural_lam.weather_dataset.WeatherDataModule attribute) boundary_mask (neural_lam.datastore.base.BaseDatastore property) (neural_lam.datastore.mdp.MDPDatastore property) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS property) (neural_lam.models.forecasters.autoregressive.ARForecaster attribute)

boundary_mean (neural_lam.models.module.ForecasterModule attribute) boundary_mean_tiled (neural_lam.models.module.ForecasterModule attribute) boundary_std (neural_lam.models.module.ForecasterModule attribute) boundary_std_tiled (neural_lam.models.module.ForecasterModule attribute) BufferList (class in neural_lam.utils.buffer_list) build_learnable_prior() (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM method) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM method) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale method) build_parser() (in module neural_lam.train_model) build_predictor() (in module neural_lam.train_model)

C

CartesianGridShape (class in neural_lam.datastore.base) check_graph_type() (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM method) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM method) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale method) checkpoint_dir (neural_lam.train_model.Run property) checkpoint_path (neural_lam.train_model.Run attribute) chunk_sizes (neural_lam.gnn_layers.SplitMLPs attribute) clamp_lower_idx (neural_lam.models.step_predictors.base.StepPredictor attribute) clamp_lower_upper_idx (neural_lam.models.step_predictors.base.StepPredictor attribute) clamp_upper_idx (neural_lam.models.step_predictors.base.StepPredictor attribute) class_name (neural_lam.datastore.npyfilesmeps.config.Projection attribute) cli() (in module neural_lam.create_graph) (in module neural_lam.create_graph_with_wmg) (in module neural_lam.datastore.npyfilesmeps.compute_standardization_stats) combine_with_latent() (neural_lam.models.latent.base_decoder.BaseGraphLatentDecoder method) (neural_lam.models.latent.graph_decoder.GraphLatentDecoder method) (neural_lam.models.latent.hi_graph_decoder.HiGraphLatentDecoder method) common_step() (neural_lam.models.module.ForecasterModule method) compute_dist_params() (neural_lam.models.latent.base_encoder.BaseLatentEncoder method) (neural_lam.models.latent.constant_encoder.ConstantLatentEncoder method) (neural_lam.models.latent.graph_encoder.GraphLatentEncoder method) (neural_lam.models.latent.hi_graph_encoder.HiGraphLatentEncoder method)

compute_grid_input_dim() (in module neural_lam.utils.graph) ComputeConfig (class in neural_lam.config) config (neural_lam.datastore.base.BaseDatastore property) (neural_lam.datastore.mdp.MDPDatastore property) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS property) config_path (neural_lam.config.DatastoreSelection attribute) configure_optimizers() (neural_lam.models.module.ForecasterModule method) ConstantLatentEncoder (class in neural_lam.models.latent.constant_encoder) coords_projection (neural_lam.datastore.base.BaseDatastore property) (neural_lam.datastore.mdp.MDPDatastore property) coords_projection() (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS method) create_dataarray_from_tensor() (neural_lam.weather_dataset.WeatherDataset method) create_gif (neural_lam.config.TrainRunConfig attribute) (neural_lam.models.module.ForecasterModule attribute) create_graph() (in module neural_lam.api) (in module neural_lam.create_graph) create_graph_from_datastore() (in module neural_lam.create_graph) (in module neural_lam.create_graph_with_wmg) create_metric_log_dict() (neural_lam.models.module.ForecasterModule method) crps_gauss() (in module neural_lam.metrics) CURRENT_GRAPH_SPEC_VERSION (in module neural_lam.create_graph) CustomMLFlowLogger (class in neural_lam.custom_loggers)

D

da_boundary_forcing (neural_lam.weather_dataset.WeatherDataset attribute) da_forcing (neural_lam.weather_dataset.WeatherDataset attribute) da_state (neural_lam.weather_dataset.WeatherDataset attribute) DataConfig (class in neural_lam.config) Dataset (class in neural_lam.datastore.npyfilesmeps.config) dataset (neural_lam.datastore.npyfilesmeps.config.NpyDatastoreConfig attribute) datastore (neural_lam.models.module.ForecasterModule attribute) (neural_lam.weather_dataset.WeatherDataset attribute) datastore_boundary (neural_lam.models.module.ForecasterModule attribute) (neural_lam.weather_dataset.WeatherDataset attribute)

DATASTORE_CLASSES (in module neural_lam.datastore) DATASTORES (in module neural_lam.datastore) datastores (neural_lam.config.NeuralLAMConfig attribute) DatastoreSelection (class in neural_lam.config) decoder (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale attribute) DEFINED_METRICS (in module neural_lam.metrics) devices (neural_lam.config.ComputeConfig attribute) diff_mean (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) diff_std (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute)

E

edge_split_sections (neural_lam.models.step_predictors.graph.hi_lam_parallel.HiLAMParallel attribute) embedd_grid_and_graph() (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM method) embedd_grid_with_target() (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM method) embedd_m2m (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) embedd_mesh() (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM method) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM method) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale method) embedd_mesh_nodes() (neural_lam.models.step_predictors.graph.base.BaseGraphModel method) (neural_lam.models.step_predictors.graph.graph_lam.GraphLAM method) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel method)

encoder (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale attribute) encoding_grid_mlp (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) epochs (neural_lam.config.TrainRunConfig attribute) eval (neural_lam.config.TrainRunConfig attribute) eval_split (neural_lam.weather_dataset.WeatherDataModule attribute) evaluate() (in module neural_lam.api) example_plots() (neural_lam.train_model.Run method) expand_to_batch() (neural_lam.models.step_predictors.base.StepPredictor method) expected_dim_order() (neural_lam.datastore.base.BaseDatastore method)

F

fit() (in module neural_lam.train_model) forcing_mean (neural_lam.models.module.ForecasterModule attribute) forcing_mean_tiled (neural_lam.models.module.ForecasterModule attribute) forcing_std (neural_lam.models.module.ForecasterModule attribute) forcing_std_tiled (neural_lam.models.module.ForecasterModule attribute) Forecaster (class in neural_lam.models.forecasters.base) forecaster (neural_lam.models.module.ForecasterModule attribute) ForecasterModule (class in neural_lam.models.module) forward() (neural_lam.gnn_layers.InteractionNet method) (neural_lam.gnn_layers.SplitMLPs method) (neural_lam.models.forecasters.autoregressive.ARForecaster method) (neural_lam.models.forecasters.base.Forecaster method) (neural_lam.models.latent.base_decoder.BaseGraphLatentDecoder method) (neural_lam.models.latent.base_encoder.BaseLatentEncoder method) (neural_lam.models.step_predictors.base.StepPredictor method) (neural_lam.models.step_predictors.graph.base.BaseGraphModel method) (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM method) (neural_lam.models.step_predictors.persistence.PersistencePredictor method)

fractional_plot_bundle() (in module neural_lam.utils.plot) from_networkx_with_start_index() (in module neural_lam.create_graph)

G

g2m_edge_index (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) g2m_embedder (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) g2m_features (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) g2m_gnn (neural_lam.models.latent.graph_decoder.GraphLatentDecoder attribute) (neural_lam.models.latent.graph_encoder.GraphLatentEncoder attribute) (neural_lam.models.latent.hi_graph_decoder.HiGraphLatentDecoder attribute) (neural_lam.models.latent.hi_graph_encoder.HiGraphLatentEncoder attribute) (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) g2m_gnn_type (neural_lam.config.ModelConfig attribute) (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) get_clamped_new_state() (neural_lam.models.step_predictors.base.StepPredictor method) get_dataarray() (neural_lam.datastore.base.BaseDatastore method) (neural_lam.datastore.mdp.MDPDatastore method) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS method) get_gnn_class() (in module neural_lam.gnn_layers) get_integer_time() (in module neural_lam.utils.time) get_lat_lon() (neural_lam.datastore.base.BaseDatastore method) (neural_lam.datastore.mdp.MDPDatastore method) get_manual_state_feature_weights() (in module neural_lam.loss_weighting) get_metric() (in module neural_lam.metrics) get_num_data_vars() (neural_lam.datastore.base.BaseDatastore method) (neural_lam.datastore.mdp.MDPDatastore method) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS method) get_num_mesh() (neural_lam.models.step_predictors.graph.base.BaseGraphModel method) (neural_lam.models.step_predictors.graph.graph_lam.GraphLAM method) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel method) get_rank() (in module neural_lam.datastore.npyfilesmeps.compute_standardization_stats) get_standardization_dataarray() (neural_lam.datastore.base.BaseDatastore method) (neural_lam.datastore.mdp.MDPDatastore method) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS method) get_state_feature_weighting() (in module neural_lam.loss_weighting) get_time_crop_slice() (in module neural_lam.utils.time)

get_time_step() (in module neural_lam.utils.time) get_uniform_state_feature_weights() (in module neural_lam.loss_weighting) get_vars_long_names() (neural_lam.datastore.base.BaseDatastore method) (neural_lam.datastore.mdp.MDPDatastore method) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS method) get_vars_names() (neural_lam.datastore.base.BaseDatastore method) (neural_lam.datastore.mdp.MDPDatastore method) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS method) get_vars_units() (neural_lam.datastore.base.BaseDatastore method) (neural_lam.datastore.mdp.MDPDatastore method) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS method) get_world_size() (in module neural_lam.datastore.npyfilesmeps.compute_standardization_stats) get_xy() (neural_lam.datastore.base.BaseDatastore method) (neural_lam.datastore.mdp.MDPDatastore method) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS method) get_xy_extent() (neural_lam.datastore.base.BaseDatastore method) GNN_TYPES (in module neural_lam.gnn_layers) graph (neural_lam.config.ModelConfig attribute) GraphEFM (class in neural_lam.models.step_predictors.graph.graph_efm) GraphEFMMultiScale (class in neural_lam.models.step_predictors.graph.graph_efm) GraphLAM (class in neural_lam.models.step_predictors.graph.graph_lam) GraphLatentDecoder (class in neural_lam.models.latent.graph_decoder) GraphLatentEncoder (class in neural_lam.models.latent.graph_encoder) grid_current_embedder (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) grid_dim (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) grid_embedder (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) GRID_HEIGHT (in module neural_lam.plot_graph) grid_input_dim (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) grid_prev_embedder (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) grid_shape_state (neural_lam.datastore.base.BaseDatastore property) (neural_lam.datastore.mdp.MDPDatastore property) (neural_lam.datastore.npyfilesmeps.config.NpyDatastoreConfig attribute) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS property) grid_static_features (neural_lam.models.step_predictors.base.StepPredictor attribute) grid_update_mlp (neural_lam.models.latent.base_decoder.BaseGraphLatentDecoder attribute)

H

has_ensemble_forcing (neural_lam.datastore.base.BaseDatastore attribute) (neural_lam.datastore.mdp.MDPDatastore attribute) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS attribute) has_working_latex() (in module neural_lam.utils.plot) hi_processor_step() (neural_lam.models.step_predictors.graph.hi_lam.HiLAM method) (neural_lam.models.step_predictors.graph.hi_lam_parallel.HiLAMParallel method) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel method) hidden_dim (neural_lam.config.ModelConfig attribute) (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute)

hidden_layers (neural_lam.config.ModelConfig attribute) (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) hierarchical (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) HiGraphLatentDecoder (class in neural_lam.models.latent.hi_graph_decoder) HiGraphLatentEncoder (class in neural_lam.models.latent.hi_graph_encoder) HiLAM (class in neural_lam.models.step_predictors.graph.hi_lam) HiLAMParallel (class in neural_lam.models.step_predictors.graph.hi_lam_parallel)

I

init_datastore() (in module neural_lam.datastore) INIT_STEPS (neural_lam.weather_dataset.WeatherDataset attribute) init_training_logger_metrics() (in module neural_lam.utils.logging) InteractionNet (class in neural_lam.gnn_layers) interior_mask (neural_lam.models.forecasters.autoregressive.ARForecaster attribute) interior_mask_bool (neural_lam.models.module.ForecasterModule attribute) intra_down_gnns (neural_lam.models.latent.hi_graph_decoder.HiGraphLatentDecoder attribute) intra_level_gnns (neural_lam.models.latent.hi_graph_encoder.HiGraphLatentEncoder attribute) intra_up_gnns (neural_lam.models.latent.hi_graph_decoder.HiGraphLatentDecoder attribute) InvalidConfigError

inverse_sigmoid() (in module neural_lam.utils.tensor) inverse_softplus() (in module neural_lam.utils.tensor) is_ensemble (neural_lam.datastore.base.BaseDatastore attribute) (neural_lam.datastore.mdp.MDPDatastore attribute) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS attribute) is_forecast (neural_lam.datastore.base.BaseDatastore attribute) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS attribute) is_on_rectilinear_grid (neural_lam.datastore.base.BaseDatastore property) (neural_lam.datastore.mdp.MDPDatastore property) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS attribute)

K

kind (neural_lam.config.DatastoreSelection attribute)

kwargs (neural_lam.datastore.npyfilesmeps.config.Projection attribute)

L

latent_dim (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) latent_embedder (neural_lam.models.latent.base_decoder.BaseGraphLatentDecoder attribute) latent_param_map (neural_lam.models.latent.graph_encoder.GraphLatentEncoder attribute) (neural_lam.models.latent.hi_graph_encoder.HiGraphLatentEncoder attribute) latent_spatial_dim (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM property) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM property) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale property) LEGACY_GRAPH_SPEC_VERSION (in module neural_lam.utils.graph) level_mesh_sizes (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) load (neural_lam.config.TrainRunConfig attribute) load_and_register_graph() (in module neural_lam.utils.graph) load_config_and_datastore() (in module neural_lam.config) load_forecaster_module_from_checkpoint() (in module neural_lam.train_model)

load_graph() (in module neural_lam.utils.graph) load_single_member (neural_lam.config.DataConfig attribute) (neural_lam.weather_dataset.WeatherDataModule attribute) (neural_lam.weather_dataset.WeatherDataset attribute) log_image() (neural_lam.custom_loggers.CustomMLFlowLogger method) log_on_rank_zero() (in module neural_lam.utils.logging) logger (neural_lam.config.LoggingConfig attribute) logger_project (neural_lam.config.LoggingConfig attribute) logger_run_name (neural_lam.config.LoggingConfig attribute) LoggingConfig (class in neural_lam.config) loss (neural_lam.config.TrainRunConfig attribute) (neural_lam.models.module.ForecasterModule attribute) lower (neural_lam.config.OutputClamping attribute) lr (neural_lam.config.TrainRunConfig attribute)

M

m2g_edge_index (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) m2g_embedder (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) m2g_features (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) m2g_gnn (neural_lam.models.latent.graph_decoder.GraphLatentDecoder attribute) (neural_lam.models.latent.hi_graph_decoder.HiGraphLatentDecoder attribute) (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) m2g_gnn_type (neural_lam.config.ModelConfig attribute) (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) m2m_edge_index (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale attribute) (neural_lam.models.step_predictors.graph.graph_lam.GraphLAM attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) m2m_embedder (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale attribute) (neural_lam.models.step_predictors.graph.graph_lam.GraphLAM attribute) m2m_features (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale attribute) (neural_lam.models.step_predictors.graph.graph_lam.GraphLAM attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) m2m_gnns (neural_lam.models.latent.graph_decoder.GraphLatentDecoder attribute) (neural_lam.models.latent.graph_encoder.GraphLatentEncoder attribute) mae() (in module neural_lam.metrics) main() (in module neural_lam.datastore.npyfilesmeps.compute_standardization_stats) (in module neural_lam.plot_graph) (in module neural_lam.train_model) make_down_gnns() (neural_lam.models.step_predictors.graph.hi_lam.HiLAM method) make_gnn_seq() (in module neural_lam.utils.networks) make_mlp() (in module neural_lam.utils.networks) make_same_gnns() (neural_lam.models.step_predictors.graph.hi_lam.HiLAM method) make_up_gnns() (neural_lam.models.step_predictors.graph.hi_lam.HiLAM method) ManualStateFeatureWeighting (class in neural_lam.config) mask_and_reduce_metric() (in module neural_lam.metrics) matched_metrics (neural_lam.models.module.ForecasterModule attribute) MDPDatastore (class in neural_lam.datastore.mdp) mesh_aggr (neural_lam.config.ModelConfig attribute) (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) mesh_down_edge_index (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) mesh_down_embedders (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) mesh_down_features (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) mesh_down_gnn_type (neural_lam.config.ModelConfig attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) mesh_down_gnns (neural_lam.models.latent.hi_graph_decoder.HiGraphLatentDecoder attribute) (neural_lam.models.step_predictors.graph.hi_lam.HiLAM attribute) mesh_down_same_gnns (neural_lam.models.step_predictors.graph.hi_lam.HiLAM attribute) mesh_down_step() (neural_lam.models.step_predictors.graph.hi_lam.HiLAM method) mesh_embedder (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale attribute) (neural_lam.models.step_predictors.graph.graph_lam.GraphLAM attribute) mesh_embedders (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) MESH_HEIGHT (in module neural_lam.plot_graph) mesh_init_gnns (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) MESH_LEVEL_DIST (in module neural_lam.plot_graph) mesh_read_gnns (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) mesh_same_embedders (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) mesh_static_features (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale attribute) (neural_lam.models.step_predictors.graph.graph_lam.GraphLAM attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) mesh_up_edge_index (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) mesh_up_embedders (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) mesh_up_features (neural_lam.models.step_predictors.graph.graph_efm.GraphEFM attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) mesh_up_gnn_type (neural_lam.config.ModelConfig attribute) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute)

mesh_up_gnns (neural_lam.models.latent.hi_graph_decoder.HiGraphLatentDecoder attribute) (neural_lam.models.latent.hi_graph_encoder.HiGraphLatentEncoder attribute) (neural_lam.models.step_predictors.graph.hi_lam.HiLAM attribute) mesh_up_same_gnns (neural_lam.models.step_predictors.graph.hi_lam.HiLAM attribute) mesh_up_step() (neural_lam.models.step_predictors.graph.hi_lam.HiLAM method) message() (neural_lam.gnn_layers.InteractionNet method) (neural_lam.gnn_layers.PropagationNet method) METAINFO_FILENAME (in module neural_lam.create_graph) metrics_watch (neural_lam.config.TrainRunConfig attribute) mk_2d_graph() (in module neural_lam.create_graph) mlp_blueprint_end (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) mlps (neural_lam.gnn_layers.SplitMLPs attribute) model (neural_lam.config.ModelConfig attribute) ModelConfig (class in neural_lam.config) MODELS (in module neural_lam.models) module neural_lam neural_lam.api neural_lam.config neural_lam.create_graph neural_lam.create_graph_with_wmg neural_lam.custom_loggers neural_lam.datastore neural_lam.datastore.base neural_lam.datastore.mdp neural_lam.datastore.npyfilesmeps neural_lam.datastore.npyfilesmeps.compute_standardization_stats neural_lam.datastore.npyfilesmeps.config neural_lam.datastore.npyfilesmeps.store neural_lam.datastore.plot_example neural_lam.gnn_layers neural_lam.loss_weighting neural_lam.metrics neural_lam.models neural_lam.models.forecasters neural_lam.models.forecasters.autoregressive neural_lam.models.forecasters.base neural_lam.models.latent neural_lam.models.latent.base_decoder neural_lam.models.latent.base_encoder neural_lam.models.latent.constant_encoder neural_lam.models.latent.graph_decoder neural_lam.models.latent.graph_encoder neural_lam.models.latent.hi_graph_decoder neural_lam.models.latent.hi_graph_encoder neural_lam.models.module neural_lam.models.step_predictors neural_lam.models.step_predictors.base neural_lam.models.step_predictors.graph neural_lam.models.step_predictors.graph.base neural_lam.models.step_predictors.graph.graph_efm neural_lam.models.step_predictors.graph.graph_lam neural_lam.models.step_predictors.graph.hi_lam neural_lam.models.step_predictors.graph.hi_lam_parallel neural_lam.models.step_predictors.graph.hierarchical neural_lam.models.step_predictors.persistence neural_lam.plot_graph neural_lam.train_model neural_lam.utils neural_lam.utils.buffer_list neural_lam.utils.graph neural_lam.utils.logging neural_lam.utils.networks neural_lam.utils.plot neural_lam.utils.tensor neural_lam.utils.time neural_lam.vis neural_lam.weather_dataset mse() (in module neural_lam.metrics) multiprocessing_context (neural_lam.weather_dataset.WeatherDataModule attribute)

N

n_boundary_points (neural_lam.config.DatastoreSelection attribute) n_buffers (neural_lam.utils.buffer_list.BufferList attribute) n_example_pred (neural_lam.config.TrainRunConfig attribute) (neural_lam.models.module.ForecasterModule attribute) name (neural_lam.datastore.npyfilesmeps.config.Dataset attribute) neural_lam module neural_lam.api module neural_lam.config module neural_lam.create_graph module neural_lam.create_graph_with_wmg module neural_lam.custom_loggers module neural_lam.datastore module neural_lam.datastore.base module neural_lam.datastore.mdp module neural_lam.datastore.npyfilesmeps module neural_lam.datastore.npyfilesmeps.compute_standardization_stats module neural_lam.datastore.npyfilesmeps.config module neural_lam.datastore.npyfilesmeps.store module neural_lam.datastore.plot_example module neural_lam.gnn_layers module neural_lam.loss_weighting module neural_lam.metrics module neural_lam.models module neural_lam.models.forecasters module neural_lam.models.forecasters.autoregressive module neural_lam.models.forecasters.base module neural_lam.models.latent module neural_lam.models.latent.base_decoder module neural_lam.models.latent.base_encoder module neural_lam.models.latent.constant_encoder module neural_lam.models.latent.graph_decoder module neural_lam.models.latent.graph_encoder module neural_lam.models.latent.hi_graph_decoder module neural_lam.models.latent.hi_graph_encoder module neural_lam.models.module module neural_lam.models.step_predictors module neural_lam.models.step_predictors.base module neural_lam.models.step_predictors.graph module

neural_lam.models.step_predictors.graph.base module neural_lam.models.step_predictors.graph.graph_efm module neural_lam.models.step_predictors.graph.graph_lam module neural_lam.models.step_predictors.graph.hi_lam module neural_lam.models.step_predictors.graph.hi_lam_parallel module neural_lam.models.step_predictors.graph.hierarchical module neural_lam.models.step_predictors.persistence module neural_lam.plot_graph module neural_lam.train_model module neural_lam.utils module neural_lam.utils.buffer_list module neural_lam.utils.graph module neural_lam.utils.logging module neural_lam.utils.networks module neural_lam.utils.plot module neural_lam.utils.tensor module neural_lam.utils.time module neural_lam.vis module neural_lam.weather_dataset module NeuralLAMConfig (class in neural_lam.config) nll() (in module neural_lam.metrics) node_residual_target() (neural_lam.gnn_layers.InteractionNet method) (neural_lam.gnn_layers.PropagationNet method) NpyDatastoreConfig (class in neural_lam.datastore.npyfilesmeps.config) NpyFilesDatastoreMEPS (class in neural_lam.datastore.npyfilesmeps.store) num_ensemble_members (neural_lam.datastore.npyfilesmeps.config.Dataset attribute) num_forcing_features (neural_lam.datastore.npyfilesmeps.config.Dataset attribute) num_future_boundary_steps (neural_lam.config.DataConfig attribute) (neural_lam.weather_dataset.WeatherDataModule attribute) (neural_lam.weather_dataset.WeatherDataset attribute) num_future_forcing_steps (neural_lam.config.DataConfig attribute) (neural_lam.weather_dataset.WeatherDataModule attribute) (neural_lam.weather_dataset.WeatherDataset attribute) num_grid_points (neural_lam.datastore.base.BaseDatastore property) (neural_lam.datastore.mdp.MDPDatastore property) num_levels (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel attribute) num_mesh_nodes (neural_lam.models.latent.constant_encoder.ConstantLatentEncoder attribute) num_nodes (neural_lam.config.ComputeConfig attribute) num_past_boundary_steps (neural_lam.config.DataConfig attribute) (neural_lam.weather_dataset.WeatherDataModule attribute) (neural_lam.weather_dataset.WeatherDataset attribute) num_past_forcing_steps (neural_lam.config.DataConfig attribute) (neural_lam.weather_dataset.WeatherDataModule attribute) (neural_lam.weather_dataset.WeatherDataset attribute) num_rec (neural_lam.gnn_layers.InteractionNet attribute) num_sanity_val_steps (neural_lam.config.TrainRunConfig attribute) num_state_vars (neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM attribute) num_timesteps (neural_lam.datastore.npyfilesmeps.config.Dataset attribute) num_workers (neural_lam.config.DataConfig attribute) (neural_lam.weather_dataset.WeatherDataModule attribute)

O

on_after_batch_transfer() (neural_lam.models.module.ForecasterModule method) on_load_checkpoint() (neural_lam.models.module.ForecasterModule method) on_test_epoch_end() (neural_lam.models.module.ForecasterModule method) on_validation_epoch_end() (neural_lam.models.module.ForecasterModule method) OPEN_WATER_FILENAME_FORMAT (in module neural_lam.datastore.npyfilesmeps.store) output_clamping (neural_lam.config.TrainingConfig attribute)

output_dist (neural_lam.models.latent.base_encoder.BaseLatentEncoder attribute) output_map (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) output_std (neural_lam.config.ModelConfig attribute) (neural_lam.models.latent.base_decoder.BaseGraphLatentDecoder attribute) (neural_lam.models.step_predictors.base.StepPredictor attribute) OutputClamping (class in neural_lam.config)

P

padded_samples (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset attribute) PaddedWeatherDataset (class in neural_lam.datastore.npyfilesmeps.compute_standardization_stats) param_map (neural_lam.models.latent.base_decoder.BaseGraphLatentDecoder attribute) per_var_std (neural_lam.models.module.ForecasterModule attribute) PersistencePredictor (class in neural_lam.models.step_predictors.persistence) plot_dir (neural_lam.train_model.Run property) plot_error_heatmap() (in module neural_lam.vis) plot_error_map() (in module neural_lam.vis) plot_example_from_datastore() (in module neural_lam.datastore.plot_example) plot_examples() (neural_lam.models.module.ForecasterModule method) plot_graph() (in module neural_lam.create_graph) (in module neural_lam.plot_graph) plot_on_axis() (in module neural_lam.vis) plot_prediction() (in module neural_lam.vis) plot_spatial_error() (in module neural_lam.vis) plotted_examples (neural_lam.models.module.ForecasterModule attribute)

precision (neural_lam.config.ComputeConfig attribute) predictor (neural_lam.models.forecasters.autoregressive.ARForecaster attribute) predicts_std (neural_lam.models.forecasters.autoregressive.ARForecaster property) (neural_lam.models.forecasters.base.Forecaster property) (neural_lam.models.step_predictors.base.StepPredictor property) prepare_clamping_params() (neural_lam.models.step_predictors.base.StepPredictor method) prepend_node_index() (in module neural_lam.create_graph) process_step() (neural_lam.models.step_predictors.graph.base.BaseGraphModel method) (neural_lam.models.step_predictors.graph.graph_lam.GraphLAM method) (neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel method) processor (neural_lam.models.step_predictors.graph.graph_lam.GraphLAM attribute) processor_layers (neural_lam.config.ModelConfig attribute) (neural_lam.models.step_predictors.graph.base.BaseGraphModel attribute) Projection (class in neural_lam.datastore.npyfilesmeps.config) projection (neural_lam.datastore.npyfilesmeps.config.NpyDatastoreConfig attribute) PropagationNet (class in neural_lam.gnn_layers)

R

real_sample_masks() (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset method) remove_state_features_with_index (neural_lam.datastore.npyfilesmeps.config.Dataset attribute) restore_opt (neural_lam.config.TrainRunConfig attribute) (neural_lam.models.module.ForecasterModule attribute) rmse_plot (neural_lam.train_model.Run property) root_path (neural_lam.datastore.base.BaseDatastore property) (neural_lam.datastore.mdp.MDPDatastore property) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS property)

Run (class in neural_lam.train_model) run() (in module neural_lam.train_model) run_dir (neural_lam.train_model.Run attribute) runs_root (neural_lam.config.LoggingConfig attribute)

S

save_dir (neural_lam.custom_loggers.CustomMLFlowLogger property) save_edges() (in module neural_lam.create_graph) save_edges_list() (in module neural_lam.create_graph) save_stats() (in module neural_lam.datastore.npyfilesmeps.compute_standardization_stats) seed (neural_lam.config.ComputeConfig attribute) setup() (in module neural_lam.datastore.npyfilesmeps.compute_standardization_stats) (neural_lam.weather_dataset.WeatherDataModule method) setup_training_logger() (in module neural_lam.utils.logging) SHORT_NAME (neural_lam.datastore.mdp.MDPDatastore attribute) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS attribute) sigmoid_lower_lims (neural_lam.models.step_predictors.base.StepPredictor attribute) sigmoid_upper_lims (neural_lam.models.step_predictors.base.StepPredictor attribute) softplus_lower_lims (neural_lam.models.step_predictors.base.StepPredictor attribute) softplus_upper_lims (neural_lam.models.step_predictors.base.StepPredictor attribute) sort_nodes_internally() (in module neural_lam.create_graph) spatial_coordinates (neural_lam.datastore.base.BaseDatastore attribute) (neural_lam.datastore.mdp.MDPDatastore attribute)

spatial_loss_maps (neural_lam.models.module.ForecasterModule attribute) split (neural_lam.weather_dataset.WeatherDataset attribute) SplitMLPs (class in neural_lam.gnn_layers) stack_grid_coords() (neural_lam.datastore.base.BaseDatastore method) state_feature_weighting (neural_lam.config.TrainingConfig attribute) state_feature_weights_values (neural_lam.datastore.base.BaseDatastore property) STATE_FILENAME_FORMAT (in module neural_lam.datastore.npyfilesmeps.store) state_mean (neural_lam.models.module.ForecasterModule attribute) (neural_lam.models.step_predictors.base.StepPredictor attribute) state_std (neural_lam.models.module.ForecasterModule attribute) (neural_lam.models.step_predictors.base.StepPredictor attribute) step_length (neural_lam.datastore.base.BaseDatastore property) (neural_lam.datastore.mdp.MDPDatastore property) (neural_lam.datastore.npyfilesmeps.config.Dataset attribute) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS property) StepPredictor (class in neural_lam.models.step_predictors.base)

T

test_dataloader() (neural_lam.weather_dataset.WeatherDataModule method) test_dataset (neural_lam.weather_dataset.WeatherDataModule attribute) test_metrics (neural_lam.models.module.ForecasterModule attribute) test_step() (neural_lam.models.module.ForecasterModule method) TOA_SW_DOWN_FLUX_FILENAME_FORMAT (in module neural_lam.datastore.npyfilesmeps.store) total_samples (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset attribute) train() (in module neural_lam.api)

train_dataloader() (neural_lam.weather_dataset.WeatherDataModule method) train_dataset (neural_lam.weather_dataset.WeatherDataModule attribute) train_steps_to_log (neural_lam.config.TrainRunConfig attribute) trainable (neural_lam.models.step_predictors.base.StepPredictor property) training (neural_lam.config.NeuralLAMConfig attribute) training_step() (neural_lam.models.module.ForecasterModule method) TrainingConfig (class in neural_lam.config) TrainRunConfig (class in neural_lam.config)

U

UniformFeatureWeighting (class in neural_lam.config) unstack_grid_coords() (neural_lam.datastore.base.BaseDatastore method)

update_edges (neural_lam.gnn_layers.InteractionNet attribute) upper (neural_lam.config.OutputClamping attribute)

V

val_dataloader() (neural_lam.weather_dataset.WeatherDataModule method) val_dataset (neural_lam.weather_dataset.WeatherDataModule attribute) val_interval (neural_lam.config.TrainRunConfig attribute) val_metrics (neural_lam.models.module.ForecasterModule attribute) val_steps_to_log (neural_lam.config.TrainRunConfig attribute)

validation_step() (neural_lam.models.module.ForecasterModule method) var_leads_metrics_watch (neural_lam.config.TrainRunConfig attribute) var_longnames (neural_lam.datastore.npyfilesmeps.config.Dataset attribute) var_names (neural_lam.datastore.npyfilesmeps.config.Dataset attribute) var_units (neural_lam.datastore.npyfilesmeps.config.Dataset attribute)

W

wandb_id (neural_lam.config.LoggingConfig attribute) WeatherDataModule (class in neural_lam.weather_dataset) WeatherDataset (class in neural_lam.weather_dataset)

weights (neural_lam.config.ManualStateFeatureWeighting attribute) wmae() (in module neural_lam.metrics) wmse() (in module neural_lam.metrics) world_size (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset attribute)

X

x (neural_lam.datastore.base.CartesianGridShape attribute)

Y

y (neural_lam.datastore.base.CartesianGridShape attribute)

Z

zero_index_edge_index() (in module neural_lam.utils.graph) zero_index_g2m() (in module neural_lam.utils.graph)

zero_index_m2g() (in module neural_lam.utils.graph) zero_offset_edge_index() (in module neural_lam.create_graph)