[Feat] Add validate_graph_components() for pre-export structural consistency checks

[prajwal-tech07]

Summary

Add a validate_graph_components() function that validates the structural consistency of constructed networkx.DiGraph graph components before they are serialized to disk. This catches construction bugs at the source rather than downstream in neural-lam.

Problem

Currently, graph construction bugs are only discovered when:

• to_pyg() tries to serialize the graph and hits unexpected shapes/missing attributes
• neural-lam's load_graph() fails with cryptic FileNotFoundError or TypeError (as documented in the format mismatch audit in neural-lam#339)
• Or worse — they pass silently and produce wrong training results

Recent examples of bugs that would have been caught earlier:

• [Bug] Wrong KDTree Node Mapping in connect_nodes_across_graphs #82: Wrong KDTree node mapping in connect_nodes_across_graphs — grid nodes mapped to wrong mesh nodes
• #nx_draw_with_pos_and_attr silently overwrites custom edge_cmap due to wrong key check #93: nx_draw_with_pos_and_attr silently overwrites custom edge_cmap due to wrong key check
• [Bug] node_features_values returns a list of tensors instead of a tensor when list_from_attribute is None #95: node_features_values returns a list of tensors instead of a tensor when list_from_attribute is None
• Assert that all nodes are in g2m #42: Not all grid nodes are connected in g2m (open since Feb 2025, on v0.4.0 milestone)

With multiple mesh_layout variants being added (rectilinear in #81, triangular in #92, icosahedral in #76, prebuilt in #91), the surface area for construction bugs is growing. We need a systematic way to validate graph components.

Relationship to existing work

• neural-lam#323 validates the on-disk .pt tensor format after export. This proposal validates the NetworkX DiGraph objects before export — they are complementary.
• Assert that all nodes are in g2m #42 ("Assert that all nodes are in g2m") is a specific instance of one check. This proposal generalizes it into a reusable validation framework alongside other checks.
• Strengthen test_save_to_pyg with output assertions and pytest.skip #88 ("Strengthen test_save_to_pyg with output assertions") focuses on the serialization test. This proposal tests the graph objects themselves.
• neural-lam#339 (the wmg↔neural-lam bridge RFC) identifies that format mismatches between the two repos are a key pain point. Catching structural issues before serialization would prevent many of these mismatches from occurring.

Proposed API

from weather_model_graphs import validate_graph_components

components = wmg.create.archetype.create_keisler_graph(
    coords=xy, return_components=True
)

report = validate_graph_components(
    components=components,      # {"g2m": DiGraph, "m2m": DiGraph, "m2g": DiGraph}
    hierarchical=False,
    graph_crs=None,             # optional: enables CRS-specific checks
)

report.passed      # True/False
report.summary()   # prints human-readable summary
report.checks      # list of CheckResult(name, passed, details)

Checks to implement

1. Bidirectional grid coverage (generalizes Assert that all nodes are in g2m #42)
   Every grid node that appears in g2m should also appear in m2g (or be explicitly marked as boundary-only). Flags grid nodes that are "encoded to mesh" but never "decoded from mesh", which would cause silent data loss.

2. Edge feature consistency
   All edges within a component have the same set of attributes (len, vdiff, component, etc.) and consistent dimensions. For example, vdiff should be 2D for projected CRS and 3D for geographic CRS — not mixed within a single component.

3. Mesh hierarchy completeness (hierarchical only)
   For hierarchical graphs: every level-L mesh node has at least one up edge to level L+1 and one down edge from L+1. No orphan mesh nodes at any level.

4. No degenerate structures

   • No self-loops (edge from a node to itself)
   • No empty components (a component with zero edges)
   • No disconnected subgraphs within a single component (e.g., mesh nodes that form isolated clusters)

5. Coordinate sanity
   All node pos attributes are finite (no NaN, no inf), and within reasonable bounds for the declared CRS (e.g., latitude in [-90, 90] for geographic).

6. Component labeling consistency
   Edge component attributes match the expected values (g2m, m2m, m2g) and are uniform within each subgraph.

Usage in CI / tests

def test_keisler_graph_is_valid():
    components = create_keisler_graph(coords=xy, return_components=True)
    report = validate_graph_components(components, hierarchical=False)
    assert report.passed, report.summary()

def test_hierarchical_graph_is_valid():
    components = create_oskarsson_hierarchical_graph(coords=xy, return_components=True)
    report = validate_graph_components(components, hierarchical=True)
    assert report.passed, report.summary()

This also becomes valuable for validating mesh_layout="prebuilt" (#91) graphs where users provide their own mesh — we can verify it meets structural requirements before attempting serialization.

Implementation plan

1. Add weather_model_graphs/validation.py with validate_graph_components() and individual check functions
2. Add corresponding tests in tests/test_validation.py
3. Optionally integrate as an automatic step in to_pyg() / to_neural_lam() (with a validate=True flag)
4. Add a brief section to documentation

No new dependencies required — purely operates on networkx.DiGraph objects using existing networkx APIs.

[Ashu-Rajput08]

Hi! This looks like a very valuable addition to the repository, especially with the increasing number of graph construction pathways (mesh_layout variants, different connectivity strategies, hierarchical meshes, etc.). As the surface area of graph construction grows, having a systematic validation step before export makes a lot of sense.

While reviewing the codebase and investigating issues like #82 (KDTree node mapping), I noticed that many of the current bugs originate during graph construction, but are only discovered much later during serialization (to_pyg()) or downstream in neural-lam. A validation layer that operates directly on the NetworkX graph objects before export would significantly improve robustness and debugging.

My understanding of the intended role of validate_graph_components() is to perform structural consistency checks on the graph components returned by create_all_graph_components(..., return_components=True), i.e.:

{
  "g2m": networkx.DiGraph,
  "m2m": networkx.DiGraph,
  "m2g": networkx.DiGraph
}

before they are serialized or used downstream.

Proposed Structure

I think this could be implemented as a dedicated module such as:

weather_model_graphs/
    validation.py

with a main entry function:

validate_graph_components(
    components: dict[str, nx.DiGraph],
    hierarchical: bool = False,
    graph_crs=None
)

which returns a structured validation report.

For example:

report = validate_graph_components(components)

report.passed        # overall True/False
report.summary()     # human-readable description
report.checks        # list of individual CheckResult objects

Validation Report Design

Each validation could return an object such as:

CheckResult(
    name="grid_bidirectional_coverage",
    passed=True,
    details=None
)

A simple structure could look like:

class CheckResult:
    def __init__(self, name, passed, details=None):
        self.name = name
        self.passed = passed
        self.details = details


class ValidationReport:
    def __init__(self, checks):
        self.checks = checks

    @property
    def passed(self):
        return all(c.passed for c in self.checks)

    def summary(self):
        lines = []
        for c in self.checks:
            status = "PASS" if c.passed else "FAIL"
            lines.append(f"{c.name}: {status}")
            if c.details:
                lines.append(f"  {c.details}")
        return "\n".join(lines)

Proposed Validation Checks

Based on the description in this issue and the current repository architecture, the following checks seem particularly important:

---

1. Grid Coverage Consistency (g2m ↔ m2g)

Every grid node encoded into the mesh (g2m) should also appear in the decode graph (m2g), unless intentionally masked.

Example logic:

g2m_targets = {v for _, v in g2m.edges()}
m2g_targets = {v for _, v in m2g.edges()}

missing = g2m_targets - m2g_targets

If missing is non-empty, this indicates grid nodes that are encoded but never decoded, which would lead to silent data loss in the model pipeline.

---

2. Edge Attribute Consistency

Edges within each component should share the same attribute structure.

For example:

• len
• vdiff
• component

Checks would verify:

• every edge has the same attribute keys
• attribute dimensionality is consistent (e.g. vdiff shape)

This helps prevent serialization mismatches.

---

3. Degenerate Structures

Basic structural integrity checks:

• no self-loops (u == v)
• component contains edges
• no disconnected clusters within a component

These checks help catch accidental graph fragmentation.

---

4. Coordinate Sanity

Validate node pos attributes:

• no NaN
• no infinite values
• values within plausible CRS bounds (if graph_crs provided)

Example:

np.isfinite(pos).all()

---

5. Hierarchical Mesh Consistency

For hierarchical graphs:

• each level-L node should connect to level-L+1 via up edges
• and receive connections from down edges

This prevents orphan nodes in the hierarchy.

---

Tests

Tests could be placed in:

tests/test_validation.py

Example:

def test_keisler_graph_is_valid():
    components = create_keisler_graph(coords=xy, return_components=True)
    report = validate_graph_components(components)
    assert report.passed

and for hierarchical graphs:

def test_hierarchical_graph_is_valid():
    components = create_oskarsson_hierarchical_graph(coords=xy, return_components=True)
    report = validate_graph_components(components, hierarchical=True)
    assert report.passed

---

Integration

This validator could optionally be integrated into export functions such as:

to_pyg(validate=True)

so that validation runs automatically before serialization.

---

Implementation Plan

If this direction makes sense, my approach would be:

1. Implement validation.py with the base validation framework (CheckResult, ValidationReport, validate_graph_components).
2. Add the first set of structural checks described above.
3. Add unit tests covering both valid and intentionally malformed graphs.
4. Optionally explore integrating validation as a pre-export step.

Happy to help implement this or contribute individual checks once the overall API direction is finalized.

[prajwal-tech07]

Thanks @Ashu-Rajput08 — glad the design makes sense. The structure you outlined aligns with what I had in mind. I'd like to implement this myself since it ties into the mesh_layout work I'm already doing (#81, #92, #91), where validating different layout variants before export is particularly important. Will wait for maintainer feedback on the API direction before starting.

[Ashu-Rajput08]

@prajwal-tech07 Thanks for the clarification! That makes sense, especially since it connects closely with the mesh layout work you mentioned.

While investigating issue #82, I noticed that problems like incorrect node mappings during graph construction can be quite difficult to detect until much later in the pipeline. A validation layer like validate_graph_components() would be really useful for catching these structural issues earlier.

Looking forward to seeing how the API direction evolves. Happy to help with testing or validation checks if needed.