Index A | B | C | D | E | F | G | H | I | K | L | M | N | O | P | R | S | T | U | V | W | X | Y A aggregate() (neural_lam.interaction_net.InteractionNet method) aggregate_and_plot_metrics() (neural_lam.models.ar_model.ARModel method) all_gather_cat() (neural_lam.models.ar_model.ARModel method) ar_steps (neural_lam.weather_dataset.WeatherDataset attribute) ar_steps_eval (neural_lam.weather_dataset.WeatherDataModule attribute) ar_steps_train (neural_lam.weather_dataset.WeatherDataModule attribute) args (neural_lam.models.ar_model.ARModel attribute) ARModel (class in neural_lam.models.ar_model) B base_dataset (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset attribute) BaseDatastore (class in neural_lam.datastore.base) BaseGraphModel (class in neural_lam.models.base_graph_model) BaseHiGraphModel (class in neural_lam.models.base_hi_graph_model) BaseRegularGridDatastore (class in neural_lam.datastore.base) batch_size (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) BufferList (class in neural_lam.utils) C CARTESIAN_COORDS (neural_lam.datastore.base.BaseRegularGridDatastore attribute) (neural_lam.datastore.mdp.MDPDatastore attribute) CartesianGridShape (class in neural_lam.datastore.base) chunk_sizes (neural_lam.interaction_net.SplitMLPs attribute) class_name (neural_lam.datastore.npyfilesmeps.config.Projection attribute), [1] cli() (in module neural_lam.create_graph) (in module neural_lam.datastore.npyfilesmeps.compute_standardization_stats) common_step() (neural_lam.models.ar_model.ARModel method) 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), [1] configure_optimizers() (neural_lam.models.ar_model.ARModel method) 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_graph() (in module neural_lam.create_graph) create_graph_from_datastore() (in module neural_lam.create_graph) create_metric_log_dict() (neural_lam.models.ar_model.ARModel method) crps_gauss() (in module neural_lam.metrics) CustomMLFlowLogger (class in neural_lam.custom_loggers) D da_forcing (neural_lam.weather_dataset.WeatherDataset attribute) da_state (neural_lam.weather_dataset.WeatherDataset attribute) Dataset (class in neural_lam.datastore.npyfilesmeps.config) dataset (neural_lam.datastore.npyfilesmeps.config.NpyDatastoreConfig attribute), [1] datastore (neural_lam.config.NeuralLAMConfig attribute), [1] (neural_lam.weather_dataset.WeatherDataset attribute) DATASTORE_CLASSES (in module neural_lam.datastore) DATASTORES (in module neural_lam.datastore) DatastoreSelection (class in neural_lam.config) DEFINED_METRICS (in module neural_lam.metrics) E edge_split_sections (neural_lam.models.hi_lam_parallel.HiLAMParallel attribute) embedd_mesh_nodes() (neural_lam.models.base_graph_model.BaseGraphModel method) (neural_lam.models.base_hi_graph_model.BaseHiGraphModel method) (neural_lam.models.graph_lam.GraphLAM method) encoding_grid_mlp (neural_lam.models.base_graph_model.BaseGraphModel attribute) eval_split (neural_lam.weather_dataset.WeatherDataModule attribute) expand_to_batch() (neural_lam.models.ar_model.ARModel static method) expected_dim_order() (neural_lam.datastore.base.BaseDatastore method) F feature_weights (neural_lam.models.ar_model.ARModel attribute) forward() (neural_lam.interaction_net.InteractionNet method) (neural_lam.interaction_net.SplitMLPs method) fractional_plot_bundle() (in module neural_lam.utils) from_networkx_with_start_index() (in module neural_lam.create_graph) G g2m_embedder (neural_lam.models.base_graph_model.BaseGraphModel attribute) g2m_gnn (neural_lam.models.base_graph_model.BaseGraphModel attribute) get_clamped_new_state() (neural_lam.models.base_graph_model.BaseGraphModel method) get_dataarray() (neural_lam.datastore.base.BaseDatastore method) (neural_lam.datastore.mdp.MDPDatastore method) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS method) get_integer_time() (in module neural_lam.utils) 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.base_graph_model.BaseGraphModel method) (neural_lam.models.base_hi_graph_model.BaseHiGraphModel method) (neural_lam.models.graph_lam.GraphLAM method) get_original_indices() (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset method) get_original_window_indices() (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset 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_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.base.BaseRegularGridDatastore method) (neural_lam.datastore.mdp.MDPDatastore method) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS method) get_xy_extent() (neural_lam.datastore.base.BaseDatastore method) GraphLAM (class in neural_lam.models.graph_lam) grid_dim (neural_lam.models.ar_model.ARModel attribute) grid_embedder (neural_lam.models.base_graph_model.BaseGraphModel attribute) GRID_HEIGHT (in module neural_lam.plot_graph) grid_shape_state (neural_lam.datastore.base.BaseRegularGridDatastore property) (neural_lam.datastore.mdp.MDPDatastore property) (neural_lam.datastore.npyfilesmeps.config.NpyDatastoreConfig attribute), [1] (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS property) H has_working_latex() (in module neural_lam.utils) hi_processor_step() (neural_lam.models.base_hi_graph_model.BaseHiGraphModel method) (neural_lam.models.hi_lam.HiLAM method) (neural_lam.models.hi_lam_parallel.HiLAMParallel method) HiLAM (class in neural_lam.models.hi_lam) HiLAMParallel (class in neural_lam.models.hi_lam_parallel) I init_datastore() (in module neural_lam.datastore) init_training_logger_metrics() (in module neural_lam.utils) InteractionNet (class in neural_lam.interaction_net) interior_mask_bool (neural_lam.models.ar_model.ARModel property) InvalidConfigError inverse_sigmoid() (in module neural_lam.utils) inverse_softplus() (in module neural_lam.utils) is_ensemble (neural_lam.datastore.base.BaseDatastore attribute) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS attribute) is_forecast (neural_lam.datastore.base.BaseDatastore attribute) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS attribute) K kind (neural_lam.config.DatastoreSelection attribute), [1] kwargs (neural_lam.datastore.npyfilesmeps.config.Projection attribute), [1] L level_mesh_sizes (neural_lam.models.base_hi_graph_model.BaseHiGraphModel attribute) load_config_and_datastore() (in module neural_lam.config) load_graph() (in module neural_lam.utils) log_image() (neural_lam.custom_loggers.CustomMLFlowLogger method) loss (neural_lam.models.ar_model.ARModel attribute) lower (neural_lam.config.OutputClamping attribute), [1] M m2g_embedder (neural_lam.models.base_graph_model.BaseGraphModel attribute) m2g_gnn (neural_lam.models.base_graph_model.BaseGraphModel attribute) m2m_embedder (neural_lam.models.graph_lam.GraphLAM 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.hi_lam.HiLAM method) make_mlp() (in module neural_lam.utils) make_same_gnns() (neural_lam.models.hi_lam.HiLAM method) make_up_gnns() (neural_lam.models.hi_lam.HiLAM method) ManualStateFeatureWeighting (class in neural_lam.config) mask_and_reduce_metric() (in module neural_lam.metrics) MDPDatastore (class in neural_lam.datastore.mdp) mesh_down_embedders (neural_lam.models.base_hi_graph_model.BaseHiGraphModel attribute) mesh_down_gnns (neural_lam.models.hi_lam.HiLAM attribute) mesh_down_same_gnns (neural_lam.models.hi_lam.HiLAM attribute) mesh_down_step() (neural_lam.models.hi_lam.HiLAM method) mesh_embedder (neural_lam.models.graph_lam.GraphLAM attribute) mesh_embedders (neural_lam.models.base_hi_graph_model.BaseHiGraphModel attribute) MESH_HEIGHT (in module neural_lam.plot_graph) mesh_init_gnns (neural_lam.models.base_hi_graph_model.BaseHiGraphModel attribute) MESH_LEVEL_DIST (in module neural_lam.plot_graph) mesh_read_gnns (neural_lam.models.base_hi_graph_model.BaseHiGraphModel attribute) mesh_same_embedders (neural_lam.models.base_hi_graph_model.BaseHiGraphModel attribute) mesh_up_embedders (neural_lam.models.base_hi_graph_model.BaseHiGraphModel attribute) mesh_up_gnns (neural_lam.models.hi_lam.HiLAM attribute) mesh_up_same_gnns (neural_lam.models.hi_lam.HiLAM attribute) mesh_up_step() (neural_lam.models.hi_lam.HiLAM method) message() (neural_lam.interaction_net.InteractionNet method) mk_2d_graph() (in module neural_lam.create_graph) mlp_blueprint_end (neural_lam.models.base_graph_model.BaseGraphModel attribute) mlps (neural_lam.interaction_net.SplitMLPs attribute) MODELS (in module neural_lam.train_model) module neural_lam neural_lam.config neural_lam.create_graph 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.interaction_net neural_lam.loss_weighting neural_lam.metrics neural_lam.models neural_lam.models.ar_model neural_lam.models.base_graph_model neural_lam.models.base_hi_graph_model neural_lam.models.graph_lam neural_lam.models.hi_lam neural_lam.models.hi_lam_parallel neural_lam.plot_graph neural_lam.train_model neural_lam.utils neural_lam.vis neural_lam.weather_dataset mse() (in module neural_lam.metrics) multiprocessing_context (neural_lam.weather_dataset.WeatherDataModule attribute) N n_buffers (neural_lam.utils.BufferList attribute) n_example_pred (neural_lam.models.ar_model.ARModel attribute) name (neural_lam.datastore.npyfilesmeps.config.Dataset attribute), [1] neural_lam module neural_lam.config module neural_lam.create_graph 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.interaction_net module neural_lam.loss_weighting module neural_lam.metrics module neural_lam.models module neural_lam.models.ar_model module neural_lam.models.base_graph_model module neural_lam.models.base_hi_graph_model module neural_lam.models.graph_lam module neural_lam.models.hi_lam module neural_lam.models.hi_lam_parallel module neural_lam.plot_graph module neural_lam.train_model module neural_lam.utils module neural_lam.vis module neural_lam.weather_dataset module NeuralLAMConfig (class in neural_lam.config) nll() (in module neural_lam.metrics) 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), [1] num_future_forcing_steps (neural_lam.weather_dataset.WeatherDataModule attribute) (neural_lam.weather_dataset.WeatherDataset attribute) num_grid_points (neural_lam.datastore.base.BaseDatastore property) (neural_lam.datastore.base.BaseRegularGridDatastore property) num_levels (neural_lam.models.base_hi_graph_model.BaseHiGraphModel attribute) num_past_forcing_steps (neural_lam.weather_dataset.WeatherDataModule attribute) (neural_lam.weather_dataset.WeatherDataset attribute) num_rec (neural_lam.interaction_net.InteractionNet attribute) num_timesteps (neural_lam.datastore.npyfilesmeps.config.Dataset attribute) num_workers (neural_lam.weather_dataset.WeatherDataModule attribute) O on_load_checkpoint() (neural_lam.models.ar_model.ARModel method) on_test_epoch_end() (neural_lam.models.ar_model.ARModel method) on_validation_epoch_end() (neural_lam.models.ar_model.ARModel method) OPEN_WATER_FILENAME_FORMAT (in module neural_lam.datastore.npyfilesmeps.store) original_indices (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset attribute) output_clamping (neural_lam.config.TrainingConfig attribute) output_map (neural_lam.models.base_graph_model.BaseGraphModel attribute) output_std (neural_lam.models.ar_model.ARModel attribute) OutputClamping (class in neural_lam.config) P padded_indices (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset attribute) padded_samples (neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset attribute) PaddedWeatherDataset (class in neural_lam.datastore.npyfilesmeps.compute_standardization_stats) plot_error_map() (in module neural_lam.vis) plot_example_from_datastore() (in module neural_lam.datastore.plot_example) plot_examples() (neural_lam.models.ar_model.ARModel method) plot_graph() (in module neural_lam.create_graph) plot_prediction() (in module neural_lam.vis) plot_spatial_error() (in module neural_lam.vis) plotted_examples (neural_lam.models.ar_model.ARModel attribute) predict_step() (neural_lam.models.ar_model.ARModel method) (neural_lam.models.base_graph_model.BaseGraphModel method) prepare_clamping_params() (neural_lam.models.base_graph_model.BaseGraphModel method) prepend_node_index() (in module neural_lam.create_graph) process_step() (neural_lam.models.base_graph_model.BaseGraphModel method) (neural_lam.models.base_hi_graph_model.BaseHiGraphModel method) (neural_lam.models.graph_lam.GraphLAM method) processor (neural_lam.models.graph_lam.GraphLAM attribute) Projection (class in neural_lam.datastore.npyfilesmeps.config) projection (neural_lam.datastore.npyfilesmeps.config.NpyDatastoreConfig attribute), [1] R rank_zero_print() (in module neural_lam.utils) remove_state_features_with_index (neural_lam.datastore.npyfilesmeps.config.Dataset attribute) restore_opt (neural_lam.models.ar_model.ARModel attribute) root_path (neural_lam.datastore.base.BaseDatastore property) (neural_lam.datastore.mdp.MDPDatastore property) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS property) 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) setup() (in module neural_lam.datastore.npyfilesmeps.compute_standardization_stats) (neural_lam.weather_dataset.WeatherDataModule method) setup_training_logger() (in module neural_lam.utils) SHORT_NAME (neural_lam.datastore.mdp.MDPDatastore attribute) (neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS attribute) sort_nodes_internally() (in module neural_lam.create_graph) spatial_loss_maps (neural_lam.models.ar_model.ARModel attribute) split (neural_lam.weather_dataset.WeatherDataset attribute) SplitMLPs (class in neural_lam.interaction_net) stack_grid_coords() (neural_lam.datastore.base.BaseRegularGridDatastore method) standardize (neural_lam.weather_dataset.WeatherDataModule attribute) (neural_lam.weather_dataset.WeatherDataset attribute) state_feature_weighting (neural_lam.config.TrainingConfig attribute), [1] state_feature_weights_values (neural_lam.datastore.base.BaseDatastore property) STATE_FILENAME_FORMAT (in module neural_lam.datastore.npyfilesmeps.store) 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) T test_dataloader() (neural_lam.weather_dataset.WeatherDataModule method) test_dataset (neural_lam.weather_dataset.WeatherDataModule attribute) test_metrics (neural_lam.models.ar_model.ARModel attribute) test_step() (neural_lam.models.ar_model.ARModel 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_dataloader() (neural_lam.weather_dataset.WeatherDataModule method) train_dataset (neural_lam.weather_dataset.WeatherDataModule attribute) training (neural_lam.config.NeuralLAMConfig attribute), [1] training_step() (neural_lam.models.ar_model.ARModel method) TrainingConfig (class in neural_lam.config) U UniformFeatureWeighting (class in neural_lam.config) unroll_prediction() (neural_lam.models.ar_model.ARModel method) unstack_grid_coords() (neural_lam.datastore.base.BaseRegularGridDatastore method) update_edges (neural_lam.interaction_net.InteractionNet attribute) upper (neural_lam.config.OutputClamping attribute), [1] V val_dataloader() (neural_lam.weather_dataset.WeatherDataModule method) val_dataset (neural_lam.weather_dataset.WeatherDataModule attribute) val_metrics (neural_lam.models.ar_model.ARModel attribute) validation_step() (neural_lam.models.ar_model.ARModel method) var_longnames (neural_lam.datastore.npyfilesmeps.config.Dataset attribute), [1] var_names (neural_lam.datastore.npyfilesmeps.config.Dataset attribute), [1] var_units (neural_lam.datastore.npyfilesmeps.config.Dataset attribute), [1] W WeatherDataModule (class in neural_lam.weather_dataset) WeatherDataset (class in neural_lam.weather_dataset) weights (neural_lam.config.ManualStateFeatureWeighting attribute), [1] 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)