to_neural_lam() to wmg.save for neural-lam tensor-on-disk format

[prajwal-tech07]

Describe your changes

Add a to_neural_lam() function to wmg.save that serializes graph components (from wmg.create.archetype.*(..., return_components=True)) to the tensor-on-disk .pt format expected by neural_lam.utils.load_graph().

This is the first step in bridging weather-model-graphs and neural-lam — putting the serialization logic in wmg so that neural-lam's CLI can remain a thin wrapper.

What it does:

• Accepts the {"g2m", "m2m", "m2g"} dict returned by any archetype with return_components=True
• Writes 7 core .pt files (g2m/m2g as single tensors, m2m/mesh_features as lists)
• Writes 4 additional files for hierarchical graphs (mesh_up/mesh_down edge_index and features)
• Edge features are written raw (unnormalized) — neural-lam normalizes at load time
• Mesh node positions are pre-normalized by max(abs(pos)), matching the existing neural-lam convention

Supports all three archetypes:

• Keisler (flat single-scale)
• GraphCast (flat multiscale)
• Oskarsson (hierarchical)

Also adds two internal helpers (_graph_to_edge_tensors, _graph_to_node_features) that cleanly strip heterogeneous node/edge attributes before PyG conversion.

Tests: 22 new tests across 4 test classes covering all archetypes, shape/dtype validation, normalization checks, and edge cases.

No new dependencies — uses the existing [pytorch] optional extras (torch, torch-geometric).

Issue Link

Refs mllam/neural-lam#384

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the documentation to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging)

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section
  reflecting type of change (add section where missing):
  • added: when you have added new functionality
  • changed: when default behaviour of the code has been changed
  • fixes: when your contribution fixes a bug

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• author has added an entry to the changelog (and designated the change as added, changed or fixed)
• Once the PR is ready to be merged, squash commits and merge the PR.

[prajwal-tech07]

Fixed the pre-commit linting failure in a1ec5b5:

• Applied black formatting to save.py and test_save.py
• Removed unused import numpy as np from test_save.py (flake8 F401)

All pre-commit hooks and tests (24/24) pass locally now. The CI should be green on the next run.

[prajwal-tech07]

Reverted the strict torch-geometric==2.3.1) pin back to the flexible >=2.5.3 requirement (matching upstream main) in ef36af0. This way we don't restrict users to a specific version, consistent with how neural-lam handles version requirements.

All pre-commit hooks and tests (24/24) pass.

[prajwal-tech07]

Renamed to_neural_lam → to_torch_tensors_on_disk as suggested — this better describes what the function actually does (save torch tensors to disk) without namespacing it to neural-lam specifically. All test references and class names updated accordingly. All 24 tests pass.

The corresponding neural-lam PR (mllam/neural-lam#596) has been updated to use the new function name as well.