Add CHGNet model support #288
Conversation
- Integrates CHGNet (Crystal Hamiltonian Graph Neural Network), a neural network potential developed by Ceder's Group. Similar to MACE, CHGNet is a message-passing network trained on the Materials Project database, however, it uses a different architecture. Some of these differences include: (1) CHGNet includes magnetic moment calculations, (2) uses charge-informed graph features, (3) trained on a different set of Materials Project data. - CHGNetModel wrapper inherits from ModelInterface - Add comprehensive test suite - Create example script demonstrating CHGNet usage with MACE comparison - Add chgnet>=0.4.2 as optional dependency in pyproject.toml - Support batched calculations, PBC, and TorchSim workflows - Include proper error handling for missing dependencies
|
Thanks for the contribution! I left some comments and happy to merge once those are added. Also, it would be good to the model in CI in |
There was a problem hiding this comment.
Thanks for the contribution! CHGNet will be great to have.
Would you mind working through the steps described in Adding New Models? You've already nailed 1 and 2 but are missing a few steps to make sure the interface is fully tested.
3. Add torch_sim.models.tests.make_validate_model_outputs_test and torch_sim.models.tests.make_model_calculator_consistency_test as models tests. See any of the other model tests for examples.
4. Update test.yml to include proper installation and testing of the relevant model.
5. Pull the model import up to torch_sim.models by adding import to torch_sim.models.__init__.py in try except clause.
6. Update docs/conf.py to include model in autodoc_mock_imports = [...]
orionarcher
added
ecosystem
|
Hi, Thank you for all the comments and suggestions, happy to contribute. I've been working on those and will update the PR as soon as possible. Best. |
- Expanded CHGNet model wrapper to include stress and magnetic moments. - Add custom ASE calculator for consistency testing. - Integrate CHGNet into CI pipeline and documentation. - Enhanced example script with stress/magnetic moment output and reduced printing with a single clean table. - Rewrote tests to follow MatterSim, removing redundant tests.
|
Looks like the test for chgnet are failing. Good to merge once the comment is resolved and the failing tests are fixed. |
There was a problem hiding this comment.
Just a couple comments. As @abhijeetgangan pointed out, the existing ASE CHGNet calculator should be used and the tests should be fixed. Otherwise, nice work!
- Replace custom CHGNetCalculator with official one from chgnet.model.dynamics - Add stress unit conversion (GPa to eV/A^3) to match ASE calculator - Remove redundant atomic numbers validation - Remove redundant test test_chgnet_missing_atomic_numbers - Fixed failing tests
|
Looks like the CHGnet is moved to matgl as per the latest changes on that repo. Maybe we should adapt to that repo? Also, the tests are failing on mac os. |
3 participants
This PR adds a wrapper for the CHGNet model (Deng et al., Nat. Mach. Intell 5, 1031–1041, 2023) to TorchSim. CHGNet is a pretrained universal neural network potential for charge-informed atomistic modeling.
Changes Made
torch_sim/models/chgnet.py- CHGNetModel wrappertests/models/test_chgnet.py- Comprehensive test suiteexamples/scripts/1_Introduction/1.4_CHGNet.py- Demonstrationchgnet>=0.4.2topyproject.tomlFeatures
Testing
python -m pytest tests/models/test_chgnet.py -vNotes
torch_sim/models/chgnet.py)How to Test
pip install chgnet>=0.4.2 python -m pytest tests/models/test_chgnet.py -v python examples/scripts/1_Introduction/1.4_CHGNet.py