PyTorch Geometric support in MeshGraphNet

[Alexey-Kamenev]

PhysicsNeMo Pull Request

Description

The document describing the general approach is available here.

Key Changes

1. Added PyTorch Geometric (PyG) support in MeshGraphNet.
2. Added PyG-based vortex shedding example alongside the existing DGL implementation.
   • Added PyG-based VortexSheddingDataset datapipe.

Features

• The changes are non-intrusive and do not affect the existing DGL implementation which users can continue to use.
• Existing DGL vortex shedding example has been renamed to vortex_shedding_mgn_dgl.
• A new vortex shedding example is added in vortex_shedding_mgn.
  • This approach allows users to easily compare different GNN implementations.
• Any existing commands or scripts that use the DGL implementation can be used with the new PyG implementation.
• Models are interchangeable between the DGL and PyG implementations.
• The changes are tested and validated. See Known Issues section for more details.

Breaking Changes

• DGL datasets created previously as serialized DGLGraph objects will not be compatible
  with the new PyG implementation. Users will need to change the corresponding datapipe
  to load DGLGraph objects and convert them to PyG objects.

Out of scope

Advanced features are not included or not validated in this PR, such as:

• Distributed MGN.
• CuGraphOps

Known Issues

• DGL vs PyG loss curves in vortex shedding example. PyG loss curve has too much variance. 1379

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

[Alexey-Kamenev]

/blossom-ci

[Alexey-Kamenev]

/blossom-ci

[coreyjadams]

nit: The requirements.txt should include torch_scatter. I got it to work with

pip install torch-scatter -f https://data.pyg.org/whl/torch-2.7.0+cu128.html

[Alexey-Kamenev]

nit: The requirements.txt should include torch_scatter. I got it to work with

Good point - done!

[Alexey-Kamenev]

/blossom-ci

[Alexey-Kamenev]

/blossom-ci

[Alexey-Kamenev]

/blossom-ci

[Alexey-Kamenev]

/blossom-ci

[Alexey-Kamenev]

/blossom-ci

[Alexey-Kamenev]

/blossom-ci