Mohit-Lakra · sumit08kumar · Mar 29, 2026 · Mar 29, 2026
diff --git a/.gitignore b/.gitignore
@@ -89,6 +89,7 @@ activate
 # Build artifacts
 dist/
 build/
+docs/_build/
 *.egg-info/
 
 tests/test_outputs/
diff --git a/docs/Makefile b/docs/Makefile
@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?= -W
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
diff --git a/docs/_static/.gitkeep b/docs/_static/.gitkeep
diff --git a/docs/api/generated/neural_lam.config.rst b/docs/api/generated/neural_lam.config.rst
@@ -0,0 +1,29 @@
+neural\_lam.config
+==================
+
+.. automodule:: neural_lam.config
+
+
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      load_config_and_datastore
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      DatastoreSelection
+      ManualStateFeatureWeighting
+      NeuralLAMConfig
+      OutputClamping
+      TrainingConfig
+      UniformFeatureWeighting
+
+   .. rubric:: Exceptions
+
+   .. autosummary::
+
+      InvalidConfigError
+
diff --git a/docs/api/generated/neural_lam.create_graph.rst b/docs/api/generated/neural_lam.create_graph.rst
@@ -0,0 +1,21 @@
+neural\_lam.create\_graph
+=========================
+
+.. automodule:: neural_lam.create_graph
+
+
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      cli
+      create_graph
+      create_graph_from_datastore
+      from_networkx_with_start_index
+      mk_2d_graph
+      plot_graph
+      prepend_node_index
+      save_edges
+      save_edges_list
+      sort_nodes_internally
+
diff --git a/docs/api/generated/neural_lam.datastore.base.rst b/docs/api/generated/neural_lam.datastore.base.rst
@@ -0,0 +1,14 @@
+neural\_lam.datastore.base
+==========================
+
+.. automodule:: neural_lam.datastore.base
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      BaseDatastore
+      BaseRegularGridDatastore
+      CartesianGridShape
+
diff --git a/docs/api/generated/neural_lam.datastore.mdp.rst b/docs/api/generated/neural_lam.datastore.mdp.rst
@@ -0,0 +1,12 @@
+neural\_lam.datastore.mdp
+=========================
+
+.. automodule:: neural_lam.datastore.mdp
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      MDPDatastore
+
diff --git a/...i/generated/neural_lam.datastore.npyfilesmeps.compute_standardization_stats.rst b/...i/generated/neural_lam.datastore.npyfilesmeps.compute_standardization_stats.rst
@@ -0,0 +1,23 @@
+neural\_lam.datastore.npyfilesmeps.compute\_standardization\_stats
+==================================================================
+
+.. automodule:: neural_lam.datastore.npyfilesmeps.compute_standardization_stats
+
+
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      cli
+      get_rank
+      get_world_size
+      main
+      save_stats
+      setup
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      PaddedWeatherDataset
+
diff --git a/docs/api/generated/neural_lam.datastore.npyfilesmeps.config.rst b/docs/api/generated/neural_lam.datastore.npyfilesmeps.config.rst
@@ -0,0 +1,14 @@
+neural\_lam.datastore.npyfilesmeps.config
+=========================================
+
+.. automodule:: neural_lam.datastore.npyfilesmeps.config
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      Dataset
+      NpyDatastoreConfig
+      Projection
+
diff --git a/docs/api/generated/neural_lam.datastore.npyfilesmeps.rst b/docs/api/generated/neural_lam.datastore.npyfilesmeps.rst
@@ -0,0 +1,15 @@
+neural\_lam.datastore.npyfilesmeps
+==================================
+
+.. automodule:: neural_lam.datastore.npyfilesmeps
+
+
+.. rubric:: Modules
+
+.. autosummary::
+   :toctree:
+   :recursive:
+
+   compute_standardization_stats
+   config
+   store
diff --git a/docs/api/generated/neural_lam.datastore.npyfilesmeps.store.rst b/docs/api/generated/neural_lam.datastore.npyfilesmeps.store.rst
@@ -0,0 +1,12 @@
+neural\_lam.datastore.npyfilesmeps.store
+========================================
+
+.. automodule:: neural_lam.datastore.npyfilesmeps.store
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      NpyFilesDatastoreMEPS
+
diff --git a/docs/api/generated/neural_lam.datastore.plot_example.rst b/docs/api/generated/neural_lam.datastore.plot_example.rst
@@ -0,0 +1,12 @@
+neural\_lam.datastore.plot\_example
+===================================
+
+.. automodule:: neural_lam.datastore.plot_example
+
+
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      plot_example_from_datastore
+
diff --git a/docs/api/generated/neural_lam.datastore.rst b/docs/api/generated/neural_lam.datastore.rst
@@ -0,0 +1,22 @@
+neural\_lam.datastore
+=====================
+
+.. automodule:: neural_lam.datastore
+
+
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      init_datastore
+
+.. rubric:: Modules
+
+.. autosummary::
+   :toctree:
+   :recursive:
+
+   base
+   mdp
+   npyfilesmeps
+   plot_example
diff --git a/docs/api/generated/neural_lam.interaction_net.rst b/docs/api/generated/neural_lam.interaction_net.rst
@@ -0,0 +1,13 @@
+neural\_lam.interaction\_net
+============================
+
+.. automodule:: neural_lam.interaction_net
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      InteractionNet
+      SplitMLPs
+
diff --git a/docs/api/generated/neural_lam.loss_weighting.rst b/docs/api/generated/neural_lam.loss_weighting.rst
@@ -0,0 +1,14 @@
+neural\_lam.loss\_weighting
+===========================
+
+.. automodule:: neural_lam.loss_weighting
+
+
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      get_manual_state_feature_weights
+      get_state_feature_weighting
+      get_uniform_state_feature_weights
+
diff --git a/docs/api/generated/neural_lam.metrics.rst b/docs/api/generated/neural_lam.metrics.rst
@@ -0,0 +1,19 @@
+neural\_lam.metrics
+===================
+
+.. automodule:: neural_lam.metrics
+
+
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      crps_gauss
+      get_metric
+      mae
+      mask_and_reduce_metric
+      mse
+      nll
+      wmae
+      wmse
+
diff --git a/docs/api/generated/neural_lam.models.ar_model.rst b/docs/api/generated/neural_lam.models.ar_model.rst
@@ -0,0 +1,12 @@
+neural\_lam.models.ar\_model
+============================
+
+.. automodule:: neural_lam.models.ar_model
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      ARModel
+
diff --git a/docs/api/generated/neural_lam.models.base_graph_model.rst b/docs/api/generated/neural_lam.models.base_graph_model.rst
@@ -0,0 +1,12 @@
+neural\_lam.models.base\_graph\_model
+=====================================
+
+.. automodule:: neural_lam.models.base_graph_model
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      BaseGraphModel
+
diff --git a/docs/api/generated/neural_lam.models.base_hi_graph_model.rst b/docs/api/generated/neural_lam.models.base_hi_graph_model.rst
@@ -0,0 +1,12 @@
+neural\_lam.models.base\_hi\_graph\_model
+=========================================
+
+.. automodule:: neural_lam.models.base_hi_graph_model
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      BaseHiGraphModel
+
diff --git a/docs/api/generated/neural_lam.models.graph_lam.rst b/docs/api/generated/neural_lam.models.graph_lam.rst
@@ -0,0 +1,12 @@
+neural\_lam.models.graph\_lam
+=============================
+
+.. automodule:: neural_lam.models.graph_lam
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      GraphLAM
+
diff --git a/docs/api/generated/neural_lam.models.hi_lam.rst b/docs/api/generated/neural_lam.models.hi_lam.rst
@@ -0,0 +1,12 @@
+neural\_lam.models.hi\_lam
+==========================
+
+.. automodule:: neural_lam.models.hi_lam
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      HiLAM
+
diff --git a/docs/api/generated/neural_lam.models.hi_lam_parallel.rst b/docs/api/generated/neural_lam.models.hi_lam_parallel.rst
@@ -0,0 +1,12 @@
+neural\_lam.models.hi\_lam\_parallel
+====================================
+
+.. automodule:: neural_lam.models.hi_lam_parallel
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      HiLAMParallel
+
diff --git a/docs/api/generated/neural_lam.models.rst b/docs/api/generated/neural_lam.models.rst
@@ -0,0 +1,18 @@
+neural\_lam.models
+==================
+
+.. automodule:: neural_lam.models
+
+
+.. rubric:: Modules
+
+.. autosummary::
+   :toctree:
+   :recursive:
+
+   ar_model
+   base_graph_model
+   base_hi_graph_model
+   graph_lam
+   hi_lam
+   hi_lam_parallel
diff --git a/docs/api/generated/neural_lam.train_model.rst b/docs/api/generated/neural_lam.train_model.rst
@@ -0,0 +1,12 @@
+neural\_lam.train\_model
+========================
+
+.. automodule:: neural_lam.train_model
+
+
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      main
+