From 616583abe9ab80585a94ae16d3809c1b9973da9d Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Tue, 14 Apr 2026 00:21:33 +0530
Subject: [PATCH 01/13] feat: add create_graph_with_wmg.py CLI using
 weather-model-graphs

---
 neural_lam/create_graph.py          |   9 ++
 neural_lam/create_graph_with_wmg.py | 188 ++++++++++++++++++++++++++++
 pyproject.toml                      |   6 +-
 tests/test_graph_creation.py        | 103 +++++++++++++++
 4 files changed, 305 insertions(+), 1 deletion(-)
 create mode 100644 neural_lam/create_graph_with_wmg.py

diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py
index 19734d722..8b8eae016 100644
--- a/neural_lam/create_graph.py
+++ b/neural_lam/create_graph.py
@@ -2,6 +2,7 @@
 
 # Standard library
 import os
+import warnings
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
 from typing import Optional
 
@@ -910,6 +911,14 @@ def cli(input_args: Optional[list[str]] = None) -> None:
         Argument list forwarded to :class:`argparse.ArgumentParser`. When
         ``None``, ``sys.argv`` is used.
     """
+    warnings.warn(
+        "create_graph.py is deprecated and will be removed in a future "
+        "version. Use create_graph_with_wmg.py instead, which delegates "
+        "graph creation to weather-model-graphs (wmg). See "
+        "https://github.com/mllam/neural-lam/issues/384 for details.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     parser = ArgumentParser(
         description="Graph generation for neural-lam",
         formatter_class=ArgumentDefaultsHelpFormatter,
diff --git a/neural_lam/create_graph_with_wmg.py b/neural_lam/create_graph_with_wmg.py
new file mode 100644
index 000000000..068b8d712
--- /dev/null
+++ b/neural_lam/create_graph_with_wmg.py
@@ -0,0 +1,188 @@
+# Standard library
+import os
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+
+# Third-party
+import numpy as np
+import weather_model_graphs as wmg
+
+# Local
+from .config import load_config_and_datastore
+from .datastore.base import BaseRegularGridDatastore
+
+ARCHETYPE_FUNCTIONS = {
+    "keisler": wmg.create.archetype.create_keisler_graph,
+    "graphcast": wmg.create.archetype.create_graphcast_graph,
+    "hierarchical": wmg.create.archetype.create_oskarsson_hierarchical_graph,
+}
+
+
+def _estimate_mesh_node_distance(xy):
+    """Estimate a reasonable mesh node distance from grid coordinates.
+
+    Uses the average grid spacing to produce a mesh that is roughly 3x
+    coarser than the grid, similar to the default behaviour of the old
+    ``create_graph.py`` script.
+
+    Parameters
+    ----------
+    xy : np.ndarray
+        Grid coordinates of shape ``(N, 2)``.
+
+    Returns
+    -------
+    float
+        Estimated mesh node distance in coordinate units.
+    """
+    x_range = np.ptp(xy[:, 0])
+    y_range = np.ptp(xy[:, 1])
+    n_points = len(xy)
+    # avg grid spacing ≈ sqrt(area / n_points)
+    avg_spacing = np.sqrt(x_range * y_range / n_points)
+    # mesh is ~3x coarser than the grid
+    return float(avg_spacing * 3)
+
+
+def create_graph_from_datastore(
+    datastore,
+    output_root_path,
+    archetype="keisler",
+    mesh_node_distance=None,
+    level_refinement_factor=3,
+    max_num_levels=None,
+):
+    """Create graph using weather-model-graphs and save in neural-lam format.
+
+    Parameters
+    ----------
+    datastore : BaseRegularGridDatastore
+        Datastore providing grid coordinates.
+    output_root_path : str
+        Directory where the .pt graph files will be saved.
+    archetype : str
+        Graph archetype to create: ``"keisler"``, ``"graphcast"``, or
+        ``"hierarchical"``.
+    mesh_node_distance : float or None
+        Distance between created mesh nodes (in coordinate units). If None,
+        automatically estimated from the grid spacing.
+    level_refinement_factor : int
+        Refinement factor between mesh hierarchy levels. Only used for
+        ``"graphcast"`` and ``"hierarchical"`` archetypes.
+    max_num_levels : int or None
+        Maximum number of mesh hierarchy levels. Only used for ``"graphcast"``
+        and ``"hierarchical"`` archetypes.
+    """
+    if not isinstance(datastore, BaseRegularGridDatastore):
+        raise NotImplementedError(
+            "Only graph creation for BaseRegularGridDatastore is supported"
+        )
+
+    if archetype not in ARCHETYPE_FUNCTIONS:
+        raise ValueError(
+            f"Unknown archetype '{archetype}'. "
+            f"Must be one of: {list(ARCHETYPE_FUNCTIONS.keys())}"
+        )
+
+    xy = datastore.get_xy(category="state", stacked=False)
+
+    # wmg expects coords as 2D array of shape (num_nodes, 2), but the
+    # datastore may return a 3D array of shape (Nx, Ny, 2) when
+    # stacked=False.  Reshape to (N, 2) for wmg.
+    xy = np.array(xy)
+    if xy.ndim == 3:
+        xy = xy.reshape(-1, 2)
+
+    if mesh_node_distance is None:
+        mesh_node_distance = _estimate_mesh_node_distance(xy)
+
+    # Build keyword arguments for the archetype function
+    archetype_kwargs = dict(
+        coords=xy,
+        mesh_node_distance=mesh_node_distance,
+        return_components=True,
+    )
+
+    # Only multiscale/hierarchical archetypes accept these parameters
+    if archetype in ("graphcast", "hierarchical"):
+        archetype_kwargs["level_refinement_factor"] = level_refinement_factor
+        archetype_kwargs["max_num_levels"] = max_num_levels
+
+    archetype_fn = ARCHETYPE_FUNCTIONS[archetype]
+    graph_components = archetype_fn(**archetype_kwargs)
+
+    hierarchical = archetype == "hierarchical"
+
+    wmg.save.to_neural_lam(
+        graph_components=graph_components,
+        output_directory=output_root_path,
+        hierarchical=hierarchical,
+    )
+
+
+def cli(input_args=None):
+    """Command-line interface for graph creation using weather-model-graphs."""
+    parser = ArgumentParser(
+        description="Graph generation for neural-lam using "
+        "weather-model-graphs (wmg)",
+        formatter_class=ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--config_path",
+        type=str,
+        help="Path to neural-lam configuration file",
+    )
+    parser.add_argument(
+        "--name",
+        type=str,
+        default="multiscale",
+        help="Name to save graph as (used as subdirectory name)",
+    )
+    parser.add_argument(
+        "--archetype",
+        type=str,
+        default="keisler",
+        choices=["keisler", "graphcast", "hierarchical"],
+        help="Graph archetype to create",
+    )
+    parser.add_argument(
+        "--mesh_node_distance",
+        type=float,
+        default=None,
+        help="Distance between mesh nodes (in coordinate units). "
+        "If not set, estimated automatically from grid spacing.",
+    )
+    parser.add_argument(
+        "--level_refinement_factor",
+        type=int,
+        default=3,
+        help="Refinement factor between mesh hierarchy levels "
+        "(only used for graphcast and hierarchical)",
+    )
+    parser.add_argument(
+        "--max_num_levels",
+        type=int,
+        default=None,
+        help="Maximum number of mesh levels "
+        "(only used for graphcast and hierarchical)",
+    )
+    args = parser.parse_args(input_args)
+
+    assert (
+        args.config_path is not None
+    ), "Specify your config with --config_path"
+
+    # Load neural-lam configuration and datastore to use
+    _, datastore = load_config_and_datastore(config_path=args.config_path)
+
+    create_graph_from_datastore(
+        datastore=datastore,
+        output_root_path=os.path.join(datastore.root_path, "graph", args.name),
+        archetype=args.archetype,
+        mesh_node_distance=args.mesh_node_distance,
+        level_refinement_factor=args.level_refinement_factor,
+        max_num_levels=args.max_num_levels,
+    )
+
+
+if __name__ == "__main__":
+    cli()
diff --git a/pyproject.toml b/pyproject.toml
index 113329e13..afab8ac49 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ dependencies = [
     "matplotlib>=3.7.0",
     "plotly>=5.15.0",
     "torch>=2.3.0",
-    "torch-geometric==2.3.1",
+    "torch-geometric>=2.5.3",
     "parse>=1.20.2",
     "dataclass-wizard<0.31.0",
     "mllam-data-prep>=0.5.0",
@@ -38,6 +38,7 @@ dependencies = [
     "boto3>=1.35.32",
     "nvidia-ml-py>=13.580.82",
     "pillow>=9.0.0",
+    "weather-model-graphs[pytorch]>=0.3.0",
 ]
 requires-python = ">=3.10"
 
@@ -50,6 +51,9 @@ cpu = ["torch>=2.12,<2.13"]
 gpu = ["torch>=2.12,<2.13"]        # CUDA 13.0, default GPU build
 gpu-cu128 = ["torch>=2.11,<2.12"]  # CUDA 12.8, last torch series with cu128 wheels
 
+[project.scripts]
+create_graph_with_wmg = "neural_lam.create_graph_with_wmg:cli"
+
 [dependency-groups]
 dev = ["pre-commit>=3.8.0", "pytest>=8.3.2", "pooch>=1.8.2"]
 
diff --git a/tests/test_graph_creation.py b/tests/test_graph_creation.py
index c97bd1910..2854dbbe4 100644
--- a/tests/test_graph_creation.py
+++ b/tests/test_graph_creation.py
@@ -1,6 +1,7 @@
 # Standard library
 import importlib.util
 import tempfile
+import warnings
 from pathlib import Path
 
 # Third-party
@@ -13,6 +14,9 @@
     METAINFO_FILENAME,
     create_graph_from_datastore,
 )
+from neural_lam.create_graph_with_wmg import (
+    create_graph_from_datastore as wmg_create_graph_from_datastore,
+)
 from neural_lam.datastore import DATASTORES
 from neural_lam.datastore.base import BaseRegularGridDatastore
 from neural_lam.utils import BufferList, load_graph
@@ -341,3 +345,102 @@ def test_buffer_list_iter(buffer_list_five):
     """Iteration yields all buffers in order."""
     values = [t.item() for t in buffer_list_five]
     assert values == [0.0, 1.0, 2.0, 3.0, 4.0]
+
+
+@pytest.mark.parametrize("archetype", ["keisler", "graphcast", "hierarchical"])
+@pytest.mark.parametrize("datastore_name", DATASTORES.keys())
+def test_wmg_graph_creation(datastore_name, archetype):
+    """Check that graph creation via weather-model-graphs produces the
+    expected .pt files with correct shapes and types."""
+    datastore = init_datastore_example(datastore_name)
+
+    if not isinstance(datastore, BaseRegularGridDatastore):
+        pytest.skip(
+            f"Skipping test for {datastore_name} as it is not a regular "
+            "grid datastore."
+        )
+
+    hierarchical = archetype == "hierarchical"
+
+    required_graph_files = [
+        "m2m_edge_index.pt",
+        "g2m_edge_index.pt",
+        "m2g_edge_index.pt",
+        "m2m_features.pt",
+        "g2m_features.pt",
+        "m2g_features.pt",
+        "mesh_features.pt",
+    ]
+    if hierarchical:
+        required_graph_files.extend(
+            [
+                "mesh_up_edge_index.pt",
+                "mesh_down_edge_index.pt",
+                "mesh_up_features.pt",
+                "mesh_down_features.pt",
+            ]
+        )
+
+    d_features = 3
+    d_mesh_static = 2
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        graph_dir_path = Path(tmpdir) / "graph" / archetype
+
+        wmg_create_graph_from_datastore(
+            datastore=datastore,
+            output_root_path=str(graph_dir_path),
+            archetype=archetype,
+        )
+
+        assert graph_dir_path.exists()
+
+        # check that all the required files are present
+        for file_name in required_graph_files:
+            assert (graph_dir_path / file_name).exists()
+
+        # try to load each and ensure they have the right shape
+        for file_name in required_graph_files:
+            file_id = Path(file_name).stem
+            result = torch.load(graph_dir_path / file_name, weights_only=True)
+
+            if file_id.startswith("g2m") or file_id.startswith("m2g"):
+                assert isinstance(result, torch.Tensor)
+
+                if file_id.endswith("_index"):
+                    assert result.shape[0] == 2
+                elif file_id.endswith("_features"):
+                    assert result.shape[1] == d_features
+
+            elif file_id.startswith("m2m") or file_id.startswith("mesh"):
+                assert isinstance(result, list)
+
+                for r in result:
+                    assert isinstance(r, torch.Tensor)
+
+                    if file_id == "mesh_features":
+                        assert r.shape[1] == d_mesh_static
+                    elif file_id.endswith("_index"):
+                        assert r.shape[0] == 2
+                    elif file_id.endswith("_features"):
+                        assert r.shape[1] == d_features
+
+
+@pytest.mark.parametrize("datastore_name", DATASTORES.keys())
+def test_old_create_graph_deprecation_warning(datastore_name):
+    """Check that the old create_graph CLI emits a deprecation warning."""
+    # First-party
+    from neural_lam.create_graph import cli as old_cli
+
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        try:
+            old_cli(["--config_path", "nonexistent.yaml"])
+        except Exception:
+            pass  # We only care about the warning, not the error
+
+        deprecation_warnings = [
+            x for x in w if issubclass(x.category, DeprecationWarning)
+        ]
+        assert len(deprecation_warnings) >= 1
+        assert "create_graph_with_wmg" in str(deprecation_warnings[0].message)

From fbe388e2675e599af1cf25e567833b4c7e78f5a1 Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Tue, 14 Apr 2026 01:09:07 +0530
Subject: [PATCH 02/13] chore: point weather-model-graphs dep to PR #123 branch

Update pyproject.toml to install weather-model-graphs from the
issue-384/to-neural-lam branch of the fork, so that CI and reviewers
can test the neural-lam side against the unreleased to_neural_lam()
changes before wmg PR #123 is merged.

Will revert to a versioned PyPI dependency once PR #123 is released.
---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index afab8ac49..9b34d407e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,7 +38,7 @@ dependencies = [
     "boto3>=1.35.32",
     "nvidia-ml-py>=13.580.82",
     "pillow>=9.0.0",
-    "weather-model-graphs[pytorch]>=0.3.0",
+    "weather-model-graphs[pytorch] @ git+https://github.com/prajwal-tech07/weather-model-graphs@issue-384/to-neural-lam",
 ]
 requires-python = ">=3.10"
 

From e51483127237de80035ab637b5244a8fdbc46411 Mon Sep 17 00:00:00 2001
From: Leif Denby <leif@denby.eu>
Date: Tue, 14 Apr 2026 07:36:32 +0200
Subject: [PATCH 03/13] fix pyproject.toml ref to wmg branch

---
 pyproject.toml | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 9b34d407e..e1f17e0a6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,7 +38,7 @@ dependencies = [
     "boto3>=1.35.32",
     "nvidia-ml-py>=13.580.82",
     "pillow>=9.0.0",
-    "weather-model-graphs[pytorch] @ git+https://github.com/prajwal-tech07/weather-model-graphs@issue-384/to-neural-lam",
+    "weather-model-graphs",
 ]
 requires-python = ">=3.10"
 
@@ -189,3 +189,6 @@ exclude = [
   ".venv/",
   "venv/",
 ]
+
+[tool.uv.sources]
+weather-model-graphs = { git = "https://github.com/prajwal-tech07/weather-model-graphs", rev = "issue-384/to-neural-lam" }

From 35fee71d3c42f31a517cc143eaa2e6e285ce65bd Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Tue, 14 Apr 2026 15:20:46 +0530
Subject: [PATCH 04/13] Address review: rename to _estimate_grid_node_spacing,
 expose grid_mesh_ratio, use stacked=True, add return_components comment, add
 README entry

---
 README.md                           | 23 ++++++++++++++
 neural_lam/create_graph_with_wmg.py | 47 ++++++++++++++++-------------
 2 files changed, 49 insertions(+), 21 deletions(-)

diff --git a/README.md b/README.md
index 35d029a65..e058bd43e 100644
--- a/README.md
+++ b/README.md
@@ -399,6 +399,9 @@ python -m neural_lam.datastore.npyfilesmeps.compute_standardization_stats <path-
 
 ### Graph creation
 
+> **Note:** The `create_graph` command below is deprecated and will be removed
+> in a future release. Please use `create_graph_with_wmg` (see below) instead.
+
 Run `python -m neural_lam.create_graph` with suitable options to generate the graph you want to use (see `python -m neural_lam.create_graph --help` for a list of options).
 The graphs used for the different models in the [paper](#graph-based-neural-weather-prediction-for-limited-area-modeling) can be created as:
 
@@ -408,6 +411,26 @@ The graphs used for the different models in the [paper](#graph-based-neural-weat
 
 The graph-related files are stored in a directory called `graphs`.
 
+### Graph creation with weather-model-graphs
+
+The recommended way to create graphs is with the `create_graph_with_wmg`
+command, which delegates graph construction to
+[weather-model-graphs](https://github.com/mllam/weather-model-graphs):
+
+```bash
+python -m neural_lam.create_graph_with_wmg --config_path <neural-lam-config-path> --archetype <archetype>
+```
+
+Available archetypes:
+
+* **keisler** (default): `python -m neural_lam.create_graph_with_wmg --config_path <neural-lam-config-path> --archetype keisler`
+* **graphcast**: `python -m neural_lam.create_graph_with_wmg --config_path <neural-lam-config-path> --archetype graphcast`
+* **hierarchical**: `python -m neural_lam.create_graph_with_wmg --config_path <neural-lam-config-path> --archetype hierarchical`
+
+Run `python -m neural_lam.create_graph_with_wmg --help` for the full list of
+options (e.g. `--mesh_node_distance`, `--grid_mesh_ratio`,
+`--level_refinement_factor`, `--max_num_levels`).
+
 ## Logging your experiments
 
 ### Weights & Biases Integration
diff --git a/neural_lam/create_graph_with_wmg.py b/neural_lam/create_graph_with_wmg.py
index 068b8d712..987cf733f 100644
--- a/neural_lam/create_graph_with_wmg.py
+++ b/neural_lam/create_graph_with_wmg.py
@@ -17,12 +17,8 @@
 }
 
 
-def _estimate_mesh_node_distance(xy):
-    """Estimate a reasonable mesh node distance from grid coordinates.
-
-    Uses the average grid spacing to produce a mesh that is roughly 3x
-    coarser than the grid, similar to the default behaviour of the old
-    ``create_graph.py`` script.
+def _estimate_grid_node_spacing(xy):
+    """Estimate the average grid node spacing from grid coordinates.
 
     Parameters
     ----------
@@ -32,15 +28,13 @@ def _estimate_mesh_node_distance(xy):
     Returns
     -------
     float
-        Estimated mesh node distance in coordinate units.
+        Estimated average grid node spacing in coordinate units.
     """
     x_range = np.ptp(xy[:, 0])
     y_range = np.ptp(xy[:, 1])
     n_points = len(xy)
     # avg grid spacing ≈ sqrt(area / n_points)
-    avg_spacing = np.sqrt(x_range * y_range / n_points)
-    # mesh is ~3x coarser than the grid
-    return float(avg_spacing * 3)
+    return float(np.sqrt(x_range * y_range / n_points))
 
 
 def create_graph_from_datastore(
@@ -48,6 +42,7 @@ def create_graph_from_datastore(
     output_root_path,
     archetype="keisler",
     mesh_node_distance=None,
+    grid_mesh_ratio=3.0,
     level_refinement_factor=3,
     max_num_levels=None,
 ):
@@ -64,7 +59,10 @@ def create_graph_from_datastore(
         ``"hierarchical"``.
     mesh_node_distance : float or None
         Distance between created mesh nodes (in coordinate units). If None,
-        automatically estimated from the grid spacing.
+        automatically estimated as ``grid_mesh_ratio * grid_spacing``.
+    grid_mesh_ratio : float
+        Ratio of mesh node distance to grid node spacing. Only used when
+        ``mesh_node_distance`` is None. Default is 3.0.
     level_refinement_factor : int
         Refinement factor between mesh hierarchy levels. Only used for
         ``"graphcast"`` and ``"hierarchical"`` archetypes.
@@ -83,19 +81,17 @@ def create_graph_from_datastore(
             f"Must be one of: {list(ARCHETYPE_FUNCTIONS.keys())}"
         )
 
-    xy = datastore.get_xy(category="state", stacked=False)
-
-    # wmg expects coords as 2D array of shape (num_nodes, 2), but the
-    # datastore may return a 3D array of shape (Nx, Ny, 2) when
-    # stacked=False.  Reshape to (N, 2) for wmg.
+    xy = datastore.get_xy(category="state", stacked=True)
     xy = np.array(xy)
-    if xy.ndim == 3:
-        xy = xy.reshape(-1, 2)
 
     if mesh_node_distance is None:
-        mesh_node_distance = _estimate_mesh_node_distance(xy)
+        grid_spacing = _estimate_grid_node_spacing(xy)
+        mesh_node_distance = grid_spacing * grid_mesh_ratio
 
-    # Build keyword arguments for the archetype function
+    # Build keyword arguments for the archetype function.
+    # return_components=True is required because wmg.save.to_neural_lam()
+    # expects the graph as separate g2m, m2g and m2m sub-graph components
+    # rather than a single merged graph.
     archetype_kwargs = dict(
         coords=xy,
         mesh_node_distance=mesh_node_distance,
@@ -149,7 +145,15 @@ def cli(input_args=None):
         type=float,
         default=None,
         help="Distance between mesh nodes (in coordinate units). "
-        "If not set, estimated automatically from grid spacing.",
+        "If not set, estimated automatically from grid spacing "
+        "and --grid_mesh_ratio.",
+    )
+    parser.add_argument(
+        "--grid_mesh_ratio",
+        type=float,
+        default=3.0,
+        help="Ratio of mesh node distance to grid node spacing. "
+        "Only used when --mesh_node_distance is not set.",
     )
     parser.add_argument(
         "--level_refinement_factor",
@@ -179,6 +183,7 @@ def cli(input_args=None):
         output_root_path=os.path.join(datastore.root_path, "graph", args.name),
         archetype=args.archetype,
         mesh_node_distance=args.mesh_node_distance,
+        grid_mesh_ratio=args.grid_mesh_ratio,
         level_refinement_factor=args.level_refinement_factor,
         max_num_levels=args.max_num_levels,
     )

From 4f273424a538b0fdac57173e3bacfdc5f0fda67b Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Wed, 15 Apr 2026 17:08:52 +0530
Subject: [PATCH 05/13] Rename grid_mesh_ratio to grid_mesh_spacing_ratio,
 revert torch-geometric to ==2.3.1

---
 README.md                           |  2 +-
 neural_lam/create_graph_with_wmg.py | 14 +++++++-------
 pyproject.toml                      |  2 +-
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index e058bd43e..94eb74298 100644
--- a/README.md
+++ b/README.md
@@ -428,7 +428,7 @@ Available archetypes:
 * **hierarchical**: `python -m neural_lam.create_graph_with_wmg --config_path <neural-lam-config-path> --archetype hierarchical`
 
 Run `python -m neural_lam.create_graph_with_wmg --help` for the full list of
-options (e.g. `--mesh_node_distance`, `--grid_mesh_ratio`,
+options (e.g. `--mesh_node_distance`, `--grid_mesh_spacing_ratio`,
 `--level_refinement_factor`, `--max_num_levels`).
 
 ## Logging your experiments
diff --git a/neural_lam/create_graph_with_wmg.py b/neural_lam/create_graph_with_wmg.py
index 987cf733f..a5a2b12ca 100644
--- a/neural_lam/create_graph_with_wmg.py
+++ b/neural_lam/create_graph_with_wmg.py
@@ -42,7 +42,7 @@ def create_graph_from_datastore(
     output_root_path,
     archetype="keisler",
     mesh_node_distance=None,
-    grid_mesh_ratio=3.0,
+    grid_mesh_spacing_ratio=3.0,
     level_refinement_factor=3,
     max_num_levels=None,
 ):
@@ -59,8 +59,8 @@ def create_graph_from_datastore(
         ``"hierarchical"``.
     mesh_node_distance : float or None
         Distance between created mesh nodes (in coordinate units). If None,
-        automatically estimated as ``grid_mesh_ratio * grid_spacing``.
-    grid_mesh_ratio : float
+        automatically estimated as ``grid_mesh_spacing_ratio * grid_spacing``.
+    grid_mesh_spacing_ratio : float
         Ratio of mesh node distance to grid node spacing. Only used when
         ``mesh_node_distance`` is None. Default is 3.0.
     level_refinement_factor : int
@@ -86,7 +86,7 @@ def create_graph_from_datastore(
 
     if mesh_node_distance is None:
         grid_spacing = _estimate_grid_node_spacing(xy)
-        mesh_node_distance = grid_spacing * grid_mesh_ratio
+        mesh_node_distance = grid_spacing * grid_mesh_spacing_ratio
 
     # Build keyword arguments for the archetype function.
     # return_components=True is required because wmg.save.to_neural_lam()
@@ -146,10 +146,10 @@ def cli(input_args=None):
         default=None,
         help="Distance between mesh nodes (in coordinate units). "
         "If not set, estimated automatically from grid spacing "
-        "and --grid_mesh_ratio.",
+        "and --grid_mesh_spacing_ratio.",
     )
     parser.add_argument(
-        "--grid_mesh_ratio",
+        "--grid_mesh_spacing_ratio",
         type=float,
         default=3.0,
         help="Ratio of mesh node distance to grid node spacing. "
@@ -183,7 +183,7 @@ def cli(input_args=None):
         output_root_path=os.path.join(datastore.root_path, "graph", args.name),
         archetype=args.archetype,
         mesh_node_distance=args.mesh_node_distance,
-        grid_mesh_ratio=args.grid_mesh_ratio,
+        grid_mesh_spacing_ratio=args.grid_mesh_spacing_ratio,
         level_refinement_factor=args.level_refinement_factor,
         max_num_levels=args.max_num_levels,
     )
diff --git a/pyproject.toml b/pyproject.toml
index e1f17e0a6..fedac13f1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ dependencies = [
     "matplotlib>=3.7.0",
     "plotly>=5.15.0",
     "torch>=2.3.0",
-    "torch-geometric>=2.5.3",
+    "torch-geometric==2.3.1",
     "parse>=1.20.2",
     "dataclass-wizard<0.31.0",
     "mllam-data-prep>=0.5.0",

From deff22be01f1a4a8b2aec432425d4e1e209c50da Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Wed, 15 Apr 2026 18:40:37 +0530
Subject: [PATCH 06/13] Switch tests to use wmg-based
 create_graph_from_datastore

Migrated all test files that used the old create_graph_from_datastore()
from neural_lam.create_graph to use the new wmg-based version from
neural_lam.create_graph_with_wmg instead.

Changes:
- test_datasets.py: Use wmg create_graph_from_datastore with archetype='keisler'
- test_clamping.py: Use wmg create_graph_from_datastore with archetype='keisler'
- test_plotting.py: Use wmg create_graph_from_datastore with archetype='keisler'
- test_training.py: Use wmg create_graph_from_datastore with archetype='keisler'
- test_plot_graph.py: Use wmg create_graph_from_datastore with keisler and
  hierarchical archetypes. Removed multiscale (graphcast) parametrization
  since the graphcast archetype produces multi-level m2m edges without
  up/down edges, which is not yet compatible with utils.load_graph().
  Graphcast graph creation is separately tested in test_graph_creation.py.
---
 tests/test_clamping.py   |  4 ++--
 tests/test_datasets.py   |  4 ++--
 tests/test_plot_graph.py | 28 +++++++++++++++-------------
 tests/test_training.py   |  4 ++--
 4 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/tests/test_clamping.py b/tests/test_clamping.py
index 8c44b5688..bdd148f44 100644
--- a/tests/test_clamping.py
+++ b/tests/test_clamping.py
@@ -6,7 +6,7 @@
 
 # First-party
 from neural_lam import config as nlconfig
-from neural_lam.create_graph import create_graph_from_datastore
+from neural_lam.create_graph_with_wmg import create_graph_from_datastore
 from neural_lam.datastore.mdp import MDPDatastore
 from neural_lam.models import GraphLAM
 from tests.conftest import init_datastore_example
@@ -23,7 +23,7 @@ def test_clamping():
         create_graph_from_datastore(
             datastore=datastore,
             output_root_path=str(graph_dir_path),
-            n_max_levels=1,
+            archetype="keisler",
         )
 
     class ModelArgs:
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 4b35840ec..01148ac48 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -9,7 +9,7 @@
 
 # First-party
 from neural_lam import config as nlconfig
-from neural_lam.create_graph import create_graph_from_datastore
+from neural_lam.create_graph_with_wmg import create_graph_from_datastore
 from neural_lam.datastore import DATASTORES
 from neural_lam.datastore.base import BaseRegularGridDatastore
 from neural_lam.models import ForecasterModule
@@ -200,7 +200,7 @@ def _create_graph():
             create_graph_from_datastore(
                 datastore=datastore,
                 output_root_path=str(graph_dir_path),
-                n_max_levels=1,
+                archetype="keisler",
             )
 
     if not isinstance(datastore, BaseRegularGridDatastore):
diff --git a/tests/test_plot_graph.py b/tests/test_plot_graph.py
index 87bd3a3fa..55bd5f441 100644
--- a/tests/test_plot_graph.py
+++ b/tests/test_plot_graph.py
@@ -7,19 +7,24 @@
 
 # First-party
 from neural_lam import utils
-from neural_lam.create_graph import create_graph_from_datastore
+from neural_lam.create_graph_with_wmg import create_graph_from_datastore
 from neural_lam.plot_graph import (
     plot_graph,
 )
 from tests.dummy_datastore import DummyDatastore
 
 
-@pytest.fixture(scope="module", params=["1level", "multiscale", "hierarchical"])
+@pytest.fixture(scope="module", params=["1level", "hierarchical"])
 def graph_fixture(request, tmp_path_factory):
     """Create a graph from a DummyDatastore and load it back.
 
-    Parametrized over graph types: 1level (flat), multiscale (flat multi-level),
-    and hierarchical.
+    Parametrized over graph types: 1level (flat, keisler archetype)
+    and hierarchical (multi-level with up/down edges).
+
+    Note: The graphcast archetype is not included here because it produces
+    multi-level m2m edges without up/down edges, which is not yet
+    compatible with ``utils.load_graph``. Graphcast graph creation is
+    tested separately in ``test_graph_creation.py``.
 
     Returns
     -------
@@ -30,14 +35,11 @@ def graph_fixture(request, tmp_path_factory):
     datastore = DummyDatastore()
 
     if graph_name == "hierarchical":
-        hierarchical = True
-        n_max_levels = 3
-    elif graph_name == "multiscale":
-        hierarchical = False
-        n_max_levels = 3
+        archetype = "hierarchical"
+        max_num_levels = 3
     elif graph_name == "1level":
-        hierarchical = False
-        n_max_levels = 1
+        archetype = "keisler"
+        max_num_levels = None
     else:
         raise ValueError(f"Unknown graph_name: {graph_name}")
 
@@ -45,8 +47,8 @@ def graph_fixture(request, tmp_path_factory):
     create_graph_from_datastore(
         datastore=datastore,
         output_root_path=str(graph_dir_path),
-        hierarchical=hierarchical,
-        n_max_levels=n_max_levels,
+        archetype=archetype,
+        max_num_levels=max_num_levels,
     )
 
     grid_xy_extent = datastore.get_xy_extent(category="state")
diff --git a/tests/test_training.py b/tests/test_training.py
index bf1a5884a..5ef77522e 100644
--- a/tests/test_training.py
+++ b/tests/test_training.py
@@ -10,7 +10,7 @@
 
 # First-party
 from neural_lam import config as nlconfig
-from neural_lam.create_graph import create_graph_from_datastore
+from neural_lam.create_graph_with_wmg import create_graph_from_datastore
 from neural_lam.datastore import DATASTORES
 from neural_lam.datastore.base import BaseRegularGridDatastore
 from neural_lam.models import ForecasterModule
@@ -85,7 +85,7 @@ def run_simple_training(
         create_graph_from_datastore(
             datastore=datastore,
             output_root_path=str(graph_dir_path),
-            n_max_levels=1,
+            archetype="keisler",
         )
 
     data_module = WeatherDataModule(

From 3ee144ee122aa7f2221d1404611a7d9ae7581178 Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Sat, 18 Apr 2026 11:35:01 +0530
Subject: [PATCH 07/13] Update wmg call: to_neural_lam ->
 to_torch_tensors_on_disk

Follow rename in weather-model-graphs (mllam/weather-model-graphs#123).
---
 neural_lam/create_graph_with_wmg.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/neural_lam/create_graph_with_wmg.py b/neural_lam/create_graph_with_wmg.py
index a5a2b12ca..632f3398a 100644
--- a/neural_lam/create_graph_with_wmg.py
+++ b/neural_lam/create_graph_with_wmg.py
@@ -89,8 +89,9 @@ def create_graph_from_datastore(
         mesh_node_distance = grid_spacing * grid_mesh_spacing_ratio
 
     # Build keyword arguments for the archetype function.
-    # return_components=True is required because wmg.save.to_neural_lam()
-    # expects the graph as separate g2m, m2g and m2m sub-graph components
+    # return_components=True is required because
+    # wmg.save.to_torch_tensors_on_disk() expects the graph as
+    # separate g2m, m2g and m2m sub-graph components
     # rather than a single merged graph.
     archetype_kwargs = dict(
         coords=xy,
@@ -108,7 +109,7 @@ def create_graph_from_datastore(
 
     hierarchical = archetype == "hierarchical"
 
-    wmg.save.to_neural_lam(
+    wmg.save.to_torch_tensors_on_disk(
         graph_components=graph_components,
         output_directory=output_root_path,
         hierarchical=hierarchical,

From 1d90155be50a570aa44d70dfdf417a00e98817f3 Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Fri, 3 Jul 2026 09:14:33 +0530
Subject: [PATCH 08/13] Add on-disk validation test for wmg graphs; finalize
 rebase onto main

- test_graph_creation: validate create_graph_with_wmg output on disk with the
  graph-storage validator (docs/validate_graph.py from #323) and assert the
  metainfo.yaml spec_version -- the end-to-end format contract for the bridge
- create_graph_with_wmg: add module docstring (interrogate 100% coverage)
- pyproject: merge the weather-model-graphs pin into a single [tool.uv.sources]
  table (rebase left two declarations, which is invalid TOML)
- uv.lock: regenerate so 'uv sync --locked' matches the wmg PR #123 branch pin
---
 neural_lam/create_graph_with_wmg.py |  7 +++++++
 pyproject.toml                      |  4 +---
 tests/test_graph_creation.py        | 21 +++++++++++++++++++++
 uv.lock                             | 18 ++++++++++++++++++
 4 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/neural_lam/create_graph_with_wmg.py b/neural_lam/create_graph_with_wmg.py
index 632f3398a..b4418631b 100644
--- a/neural_lam/create_graph_with_wmg.py
+++ b/neural_lam/create_graph_with_wmg.py
@@ -1,3 +1,10 @@
+"""Create neural-lam graphs by delegating construction to weather-model-graphs.
+
+Builds the g2m/m2m/m2g graph components with weather-model-graphs (wmg) and
+saves them to disk in neural-lam's tensor-on-disk format, replacing the
+duplicated logic in ``create_graph.py``.
+"""
+
 # Standard library
 import os
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
diff --git a/pyproject.toml b/pyproject.toml
index fedac13f1..8eb9154ab 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -85,6 +85,7 @@ torch = [
     { index = "pytorch-cu130", extra = "gpu" },
     { index = "pytorch-cu128", extra = "gpu-cu128" },
 ]
+weather-model-graphs = { git = "https://github.com/prajwal-tech07/weather-model-graphs", rev = "issue-384/to-neural-lam" }
 
 [tool.black]
 line-length = 80
@@ -189,6 +190,3 @@ exclude = [
   ".venv/",
   "venv/",
 ]
-
-[tool.uv.sources]
-weather-model-graphs = { git = "https://github.com/prajwal-tech07/weather-model-graphs", rev = "issue-384/to-neural-lam" }
diff --git a/tests/test_graph_creation.py b/tests/test_graph_creation.py
index 2854dbbe4..ef6f27f3a 100644
--- a/tests/test_graph_creation.py
+++ b/tests/test_graph_creation.py
@@ -370,6 +370,7 @@ def test_wmg_graph_creation(datastore_name, archetype):
         "g2m_features.pt",
         "m2g_features.pt",
         "mesh_features.pt",
+        METAINFO_FILENAME,
     ]
     if hierarchical:
         required_graph_files.extend(
@@ -399,8 +400,28 @@ def test_wmg_graph_creation(datastore_name, archetype):
         for file_name in required_graph_files:
             assert (graph_dir_path / file_name).exists()
 
+        # Third-party
+        import yaml
+
+        meta = yaml.safe_load(
+            (graph_dir_path / METAINFO_FILENAME).read_text(encoding="utf-8")
+        )
+        assert meta is not None
+        assert meta["spec_version"] == CURRENT_GRAPH_SPEC_VERSION
+
+        # Validate the wmg-created graph on disk against the graph-storage
+        # spec and validator introduced in #323. This is the end-to-end
+        # contract check: a graph built through create_graph_with_wmg (using
+        # weather-model-graphs' to_torch_tensors_on_disk) must pass the same
+        # validator neural-lam ships for the on-disk graph format.
+        validator = _load_validator_module()
+        report, _, _ = validator.validate_graph_directory(graph_dir_path)
+        assert not report.has_fails(), report.summarize()
+
         # try to load each and ensure they have the right shape
         for file_name in required_graph_files:
+            if file_name == METAINFO_FILENAME:
+                continue
             file_id = Path(file_name).stem
             result = torch.load(graph_dir_path / file_name, weights_only=True)
 
diff --git a/uv.lock b/uv.lock
index b9d5a7ed5..778fdfd35 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2612,6 +2612,7 @@ dependencies = [
     { name = "torch-geometric" },
     { name = "tueplots" },
     { name = "wandb" },
+    { name = "weather-model-graphs" },
 ]
 
 [package.optional-dependencies]
@@ -2658,6 +2659,7 @@ requires-dist = [
     { name = "torch-geometric", specifier = "==2.3.1" },
     { name = "tueplots", specifier = ">=0.0.8" },
     { name = "wandb", specifier = ">=0.13.10" },
+    { name = "weather-model-graphs", git = "https://github.com/prajwal-tech07/weather-model-graphs?rev=issue-384%2Fto-neural-lam" },
 ]
 provides-extras = ["cpu", "gpu", "gpu-cu128"]
 
@@ -5565,6 +5567,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/68/5a/199c59e0a824a3db2b89c5d2dade7ab5f9624dbf6448dc291b46d5ec94d3/wcwidth-0.6.0-py3-none-any.whl", hash = "sha256:1a3a1e510b553315f8e146c54764f4fb6264ffad731b3d78088cdb1478ffbdad", size = 94189, upload-time = "2026-02-06T19:19:39.646Z" },
 ]
 
+[[package]]
+name = "weather-model-graphs"
+version = "0.3.0"
+source = { git = "https://github.com/prajwal-tech07/weather-model-graphs?rev=issue-384%2Fto-neural-lam#23fb3058b639935d3e59469b923138798dd4f841" }
+dependencies = [
+    { name = "loguru" },
+    { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
+    { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
+    { name = "pyproj", version = "3.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
+    { name = "pyproj", version = "3.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
+    { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
+    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
+]
+
 [[package]]
 name = "werkzeug"
 version = "3.1.7"

From 7cf64bcb85e43595811c3504bae928748e3a549a Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Tue, 28 Jul 2026 21:02:44 +0530
Subject: [PATCH 09/13] Depend on released weather-model-graphs>=0.4.0

weather-model-graphs 0.4.0 is released on PyPI and includes
save.to_torch_tensors_on_disk, so the temporary git dependency on the
PR #123 branch is no longer needed.

- pyproject: bump to weather-model-graphs>=0.4.0 and drop the
  [tool.uv.sources] git pin
- uv.lock: regenerate so it resolves 0.4.0 from PyPI

The [pytorch] extra is intentionally not used: it requires
torch-geometric>=2.5.3, which conflicts with neural-lam's
torch-geometric==2.3.1 pin. It is not needed here, since neural-lam
already provides torch/torch-geometric and the tensor-on-disk save path
does not import torch-geometric.
---
 pyproject.toml | 3 +--
 uv.lock        | 9 ++++++---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8eb9154ab..7485bbeb5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,7 +38,7 @@ dependencies = [
     "boto3>=1.35.32",
     "nvidia-ml-py>=13.580.82",
     "pillow>=9.0.0",
-    "weather-model-graphs",
+    "weather-model-graphs>=0.4.0",
 ]
 requires-python = ">=3.10"
 
@@ -85,7 +85,6 @@ torch = [
     { index = "pytorch-cu130", extra = "gpu" },
     { index = "pytorch-cu128", extra = "gpu-cu128" },
 ]
-weather-model-graphs = { git = "https://github.com/prajwal-tech07/weather-model-graphs", rev = "issue-384/to-neural-lam" }
 
 [tool.black]
 line-length = 80
diff --git a/uv.lock b/uv.lock
index 778fdfd35..4e18287f4 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2659,7 +2659,7 @@ requires-dist = [
     { name = "torch-geometric", specifier = "==2.3.1" },
     { name = "tueplots", specifier = ">=0.0.8" },
     { name = "wandb", specifier = ">=0.13.10" },
-    { name = "weather-model-graphs", git = "https://github.com/prajwal-tech07/weather-model-graphs?rev=issue-384%2Fto-neural-lam" },
+    { name = "weather-model-graphs", specifier = ">=0.4.0" },
 ]
 provides-extras = ["cpu", "gpu", "gpu-cu128"]
 
@@ -5569,8 +5569,8 @@ wheels = [
 
 [[package]]
 name = "weather-model-graphs"
-version = "0.3.0"
-source = { git = "https://github.com/prajwal-tech07/weather-model-graphs?rev=issue-384%2Fto-neural-lam#23fb3058b639935d3e59469b923138798dd4f841" }
+version = "0.4.0"
+source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "loguru" },
     { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
@@ -5582,6 +5582,9 @@ dependencies = [
     { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
     { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu') or (extra == 'extra-10-neural-lam-cpu' and extra == 'extra-10-neural-lam-gpu-cu128') or (extra == 'extra-10-neural-lam-gpu' and extra == 'extra-10-neural-lam-gpu-cu128')" },
 ]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/10/fc/de383b841fb9fe75960d55102fd16108ea5d3b245471956eac2ec1899eee/weather_model_graphs-0.4.0-py3-none-any.whl", hash = "sha256:261cee5d72bac06b6f36f8e3d2f5c7850ba56976d5759923710353d9459639dc", size = 44677, upload-time = "2026-07-28T09:23:59.658Z" },
+]
 
 [[package]]
 name = "werkzeug"

From bba00b3d18a7068aade7dee5cca905a8826b45aa Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Tue, 28 Jul 2026 21:12:27 +0530
Subject: [PATCH 10/13] Add changelog entry for create_graph_with_wmg CLI

---
 CHANGELOG.md | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4ff85253b..751ab7d85 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Add latent encoder/decoder modules and the `GraphEFM` (hierarchical) / `GraphEFMMultiScale` (flat) step predictors for the Graph-EFM ensemble forecasting model. [\#648](https://github.com/mllam/neural-lam/pull/648) @Sir-Sloth-The-Lazy
 
+- Add `neural_lam.create_graph_with_wmg` CLI which builds graphs with
+  [weather-model-graphs](https://github.com/mllam/weather-model-graphs) instead
+  of duplicating the graph-construction logic in neural-lam, supporting the
+  `keisler`, `graphcast` and `hierarchical` archetypes. The graphs it writes
+  follow the graph storage specification and are checked against
+  `docs/validate_graph.py` in the tests. `neural_lam.create_graph` still works
+  but is deprecated and will be removed in a future release.
+  [\#596](https://github.com/mllam/neural-lam/pull/596) @prajwal-tech07
+
 - Add `--num_sanity_val_steps` CLI argument to control sanity validation steps before training (#694)
 
 - Add `--train_steps_to_log` CLI option to log training loss for individual unroll steps, and deduplicate common prediction and loss computation steps across loops [\#674](https://github.com/mllam/neural-lam/issues/674) @GiGiKoneti

From 2e863da61e6219db871152d00af94e6bc41ca9a7 Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Wed, 12 Aug 2026 23:36:39 +0530
Subject: [PATCH 11/13] Address review: restore multiscale graph test, shorten
 changelog entry

Restore the multiscale (GC-LAM) case in the plot_graph fixture, using the
graphcast archetype. It was dropped when non-hierarchical m2m edges were
still written split by level, which load_graph misread as hierarchical.
Non-hierarchical graphs are written as a single merged m2m level since
to_torch_tensors_on_disk was conformed to the graph storage spec, so all
three paper graph types load again.

Shorten the create_graph_with_wmg changelog entry to a single sentence.
---
 CHANGELOG.md             |  9 +--------
 tests/test_plot_graph.py | 15 +++++++--------
 2 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 751ab7d85..a7fdb382a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,14 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Add latent encoder/decoder modules and the `GraphEFM` (hierarchical) / `GraphEFMMultiScale` (flat) step predictors for the Graph-EFM ensemble forecasting model. [\#648](https://github.com/mllam/neural-lam/pull/648) @Sir-Sloth-The-Lazy
 
-- Add `neural_lam.create_graph_with_wmg` CLI which builds graphs with
-  [weather-model-graphs](https://github.com/mllam/weather-model-graphs) instead
-  of duplicating the graph-construction logic in neural-lam, supporting the
-  `keisler`, `graphcast` and `hierarchical` archetypes. The graphs it writes
-  follow the graph storage specification and are checked against
-  `docs/validate_graph.py` in the tests. `neural_lam.create_graph` still works
-  but is deprecated and will be removed in a future release.
-  [\#596](https://github.com/mllam/neural-lam/pull/596) @prajwal-tech07
+- Add `neural_lam.create_graph_with_wmg` CLI which builds `keisler`, `graphcast` and `hierarchical` graphs with [weather-model-graphs](https://github.com/mllam/weather-model-graphs), deprecating `neural_lam.create_graph`. [\#596](https://github.com/mllam/neural-lam/pull/596) @prajwal-tech07
 
 - Add `--num_sanity_val_steps` CLI argument to control sanity validation steps before training (#694)
 
diff --git a/tests/test_plot_graph.py b/tests/test_plot_graph.py
index 55bd5f441..7a66f5a5e 100644
--- a/tests/test_plot_graph.py
+++ b/tests/test_plot_graph.py
@@ -14,17 +14,13 @@
 from tests.dummy_datastore import DummyDatastore
 
 
-@pytest.fixture(scope="module", params=["1level", "hierarchical"])
+@pytest.fixture(scope="module", params=["1level", "multiscale", "hierarchical"])
 def graph_fixture(request, tmp_path_factory):
     """Create a graph from a DummyDatastore and load it back.
 
-    Parametrized over graph types: 1level (flat, keisler archetype)
-    and hierarchical (multi-level with up/down edges).
-
-    Note: The graphcast archetype is not included here because it produces
-    multi-level m2m edges without up/down edges, which is not yet
-    compatible with ``utils.load_graph``. Graphcast graph creation is
-    tested separately in ``test_graph_creation.py``.
+    Parametrized over graph types: 1level (flat, keisler archetype),
+    multiscale (flat multi-level, graphcast archetype) and hierarchical
+    (multi-level with up/down edges).
 
     Returns
     -------
@@ -37,6 +33,9 @@ def graph_fixture(request, tmp_path_factory):
     if graph_name == "hierarchical":
         archetype = "hierarchical"
         max_num_levels = 3
+    elif graph_name == "multiscale":
+        archetype = "graphcast"
+        max_num_levels = 3
     elif graph_name == "1level":
         archetype = "keisler"
         max_num_levels = None

From eaa9865dbbcd82a522085edc97caa374816291c9 Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Thu, 13 Aug 2026 00:12:20 +0530
Subject: [PATCH 12/13] Address review: rely on the graph validator in the wmg
 creation test

Drop the hand-rolled file presence, spec version, container type and
shape assertions from test_wmg_graph_creation; the validator from #323
already covers all of them, and pinning edge features to exactly 3 was
stricter than the spec, which allows 3 or 4.

Keep an explicit check that mesh_up/mesh_down files are present exactly
when the archetype is hierarchical. The validator infers hierarchy from
the graph contents, so it cannot tell whether the requested archetype
was honoured.
---
 tests/test_graph_creation.py | 83 +++++++-----------------------------
 1 file changed, 15 insertions(+), 68 deletions(-)

diff --git a/tests/test_graph_creation.py b/tests/test_graph_creation.py
index ef6f27f3a..0b72b66b8 100644
--- a/tests/test_graph_creation.py
+++ b/tests/test_graph_creation.py
@@ -350,8 +350,8 @@ def test_buffer_list_iter(buffer_list_five):
 @pytest.mark.parametrize("archetype", ["keisler", "graphcast", "hierarchical"])
 @pytest.mark.parametrize("datastore_name", DATASTORES.keys())
 def test_wmg_graph_creation(datastore_name, archetype):
-    """Check that graph creation via weather-model-graphs produces the
-    expected .pt files with correct shapes and types."""
+    """Check that graph creation via weather-model-graphs produces a graph
+    that conforms to the graph storage specification."""
     datastore = init_datastore_example(datastore_name)
 
     if not isinstance(datastore, BaseRegularGridDatastore):
@@ -362,29 +362,6 @@ def test_wmg_graph_creation(datastore_name, archetype):
 
     hierarchical = archetype == "hierarchical"
 
-    required_graph_files = [
-        "m2m_edge_index.pt",
-        "g2m_edge_index.pt",
-        "m2g_edge_index.pt",
-        "m2m_features.pt",
-        "g2m_features.pt",
-        "m2g_features.pt",
-        "mesh_features.pt",
-        METAINFO_FILENAME,
-    ]
-    if hierarchical:
-        required_graph_files.extend(
-            [
-                "mesh_up_edge_index.pt",
-                "mesh_down_edge_index.pt",
-                "mesh_up_features.pt",
-                "mesh_down_features.pt",
-            ]
-        )
-
-    d_features = 3
-    d_mesh_static = 2
-
     with tempfile.TemporaryDirectory() as tmpdir:
         graph_dir_path = Path(tmpdir) / "graph" / archetype
 
@@ -394,57 +371,27 @@ def test_wmg_graph_creation(datastore_name, archetype):
             archetype=archetype,
         )
 
-        assert graph_dir_path.exists()
-
-        # check that all the required files are present
-        for file_name in required_graph_files:
-            assert (graph_dir_path / file_name).exists()
-
-        # Third-party
-        import yaml
-
-        meta = yaml.safe_load(
-            (graph_dir_path / METAINFO_FILENAME).read_text(encoding="utf-8")
-        )
-        assert meta is not None
-        assert meta["spec_version"] == CURRENT_GRAPH_SPEC_VERSION
-
         # Validate the wmg-created graph on disk against the graph-storage
         # spec and validator introduced in #323. This is the end-to-end
         # contract check: a graph built through create_graph_with_wmg (using
         # weather-model-graphs' to_torch_tensors_on_disk) must pass the same
-        # validator neural-lam ships for the on-disk graph format.
+        # validator neural-lam ships for the on-disk graph format. It covers
+        # file presence, spec version, container types, edge-index shapes and
+        # feature dimensions, so those are not re-checked here.
         validator = _load_validator_module()
         report, _, _ = validator.validate_graph_directory(graph_dir_path)
         assert not report.has_fails(), report.summarize()
 
-        # try to load each and ensure they have the right shape
-        for file_name in required_graph_files:
-            if file_name == METAINFO_FILENAME:
-                continue
-            file_id = Path(file_name).stem
-            result = torch.load(graph_dir_path / file_name, weights_only=True)
-
-            if file_id.startswith("g2m") or file_id.startswith("m2g"):
-                assert isinstance(result, torch.Tensor)
-
-                if file_id.endswith("_index"):
-                    assert result.shape[0] == 2
-                elif file_id.endswith("_features"):
-                    assert result.shape[1] == d_features
-
-            elif file_id.startswith("m2m") or file_id.startswith("mesh"):
-                assert isinstance(result, list)
-
-                for r in result:
-                    assert isinstance(r, torch.Tensor)
-
-                    if file_id == "mesh_features":
-                        assert r.shape[1] == d_mesh_static
-                    elif file_id.endswith("_index"):
-                        assert r.shape[0] == 2
-                    elif file_id.endswith("_features"):
-                        assert r.shape[1] == d_features
+        # The validator infers whether a graph is hierarchical from its
+        # contents, so it cannot tell whether the requested archetype was
+        # honoured. Check that separately.
+        for file_name in (
+            "mesh_up_edge_index.pt",
+            "mesh_down_edge_index.pt",
+            "mesh_up_features.pt",
+            "mesh_down_features.pt",
+        ):
+            assert (graph_dir_path / file_name).exists() == hierarchical
 
 
 @pytest.mark.parametrize("datastore_name", DATASTORES.keys())

From 8665d28fa0004881575535247b369f30449dfeee Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Thu, 13 Aug 2026 22:11:47 +0530
Subject: [PATCH 13/13] Address review: use "distance" terminology and log the
 estimated value

Settle on "distance" over "spacing" throughout, matching the
mesh_node_distance parameter that weather-model-graphs already uses.
Rename _estimate_grid_node_spacing to _estimate_grid_node_distance and
the local grid_spacing to grid_node_distance.

Rename the --grid_mesh_spacing_ratio argument to
--mesh_grid_distance_ratio. The name read as grid:mesh while the code
computes mesh:grid, so at the default of 3.0 the mesh nodes are 3x
further apart rather than 3x closer.

Log the estimated grid node distance, the ratio applied and the
resulting mesh node distance when mesh_node_distance is not given, and
say in the docstring that it is estimated automatically.
---
 README.md                           |  2 +-
 neural_lam/create_graph_with_wmg.py | 38 +++++++++++++++++------------
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/README.md b/README.md
index 94eb74298..996d51e7e 100644
--- a/README.md
+++ b/README.md
@@ -428,7 +428,7 @@ Available archetypes:
 * **hierarchical**: `python -m neural_lam.create_graph_with_wmg --config_path <neural-lam-config-path> --archetype hierarchical`
 
 Run `python -m neural_lam.create_graph_with_wmg --help` for the full list of
-options (e.g. `--mesh_node_distance`, `--grid_mesh_spacing_ratio`,
+options (e.g. `--mesh_node_distance`, `--mesh_grid_distance_ratio`,
 `--level_refinement_factor`, `--max_num_levels`).
 
 ## Logging your experiments
diff --git a/neural_lam/create_graph_with_wmg.py b/neural_lam/create_graph_with_wmg.py
index b4418631b..02512e030 100644
--- a/neural_lam/create_graph_with_wmg.py
+++ b/neural_lam/create_graph_with_wmg.py
@@ -12,6 +12,7 @@
 # Third-party
 import numpy as np
 import weather_model_graphs as wmg
+from loguru import logger
 
 # Local
 from .config import load_config_and_datastore
@@ -24,8 +25,8 @@
 }
 
 
-def _estimate_grid_node_spacing(xy):
-    """Estimate the average grid node spacing from grid coordinates.
+def _estimate_grid_node_distance(xy):
+    """Estimate the average grid node distance from grid coordinates.
 
     Parameters
     ----------
@@ -35,12 +36,12 @@ def _estimate_grid_node_spacing(xy):
     Returns
     -------
     float
-        Estimated average grid node spacing in coordinate units.
+        Estimated average grid node distance in coordinate units.
     """
     x_range = np.ptp(xy[:, 0])
     y_range = np.ptp(xy[:, 1])
     n_points = len(xy)
-    # avg grid spacing ≈ sqrt(area / n_points)
+    # avg grid node distance ≈ sqrt(area / n_points)
     return float(np.sqrt(x_range * y_range / n_points))
 
 
@@ -49,7 +50,7 @@ def create_graph_from_datastore(
     output_root_path,
     archetype="keisler",
     mesh_node_distance=None,
-    grid_mesh_spacing_ratio=3.0,
+    mesh_grid_distance_ratio=3.0,
     level_refinement_factor=3,
     max_num_levels=None,
 ):
@@ -66,9 +67,10 @@ def create_graph_from_datastore(
         ``"hierarchical"``.
     mesh_node_distance : float or None
         Distance between created mesh nodes (in coordinate units). If None,
-        automatically estimated as ``grid_mesh_spacing_ratio * grid_spacing``.
-    grid_mesh_spacing_ratio : float
-        Ratio of mesh node distance to grid node spacing. Only used when
+        the grid node distance is estimated automatically from the grid
+        coordinates and multiplied by ``mesh_grid_distance_ratio``.
+    mesh_grid_distance_ratio : float
+        Ratio of mesh node distance to grid node distance. Only used when
         ``mesh_node_distance`` is None. Default is 3.0.
     level_refinement_factor : int
         Refinement factor between mesh hierarchy levels. Only used for
@@ -92,8 +94,14 @@ def create_graph_from_datastore(
     xy = np.array(xy)
 
     if mesh_node_distance is None:
-        grid_spacing = _estimate_grid_node_spacing(xy)
-        mesh_node_distance = grid_spacing * grid_mesh_spacing_ratio
+        grid_node_distance = _estimate_grid_node_distance(xy)
+        mesh_node_distance = grid_node_distance * mesh_grid_distance_ratio
+        logger.info(
+            f"mesh_node_distance not given; estimated grid node distance "
+            f"{grid_node_distance:.2f} x mesh_grid_distance_ratio "
+            f"{mesh_grid_distance_ratio} -> mesh_node_distance "
+            f"{mesh_node_distance:.2f}"
+        )
 
     # Build keyword arguments for the archetype function.
     # return_components=True is required because
@@ -153,14 +161,14 @@ def cli(input_args=None):
         type=float,
         default=None,
         help="Distance between mesh nodes (in coordinate units). "
-        "If not set, estimated automatically from grid spacing "
-        "and --grid_mesh_spacing_ratio.",
+        "If not set, estimated automatically from the grid node distance "
+        "and --mesh_grid_distance_ratio.",
     )
     parser.add_argument(
-        "--grid_mesh_spacing_ratio",
+        "--mesh_grid_distance_ratio",
         type=float,
         default=3.0,
-        help="Ratio of mesh node distance to grid node spacing. "
+        help="Ratio of mesh node distance to grid node distance. "
         "Only used when --mesh_node_distance is not set.",
     )
     parser.add_argument(
@@ -191,7 +199,7 @@ def cli(input_args=None):
         output_root_path=os.path.join(datastore.root_path, "graph", args.name),
         archetype=args.archetype,
         mesh_node_distance=args.mesh_node_distance,
-        grid_mesh_spacing_ratio=args.grid_mesh_spacing_ratio,
+        mesh_grid_distance_ratio=args.mesh_grid_distance_ratio,
         level_refinement_factor=args.level_refinement_factor,
         max_num_levels=args.max_num_levels,
     )