From 64d744be6e5cb9427750855ecb9a50abedceaeb4 Mon Sep 17 00:00:00 2001 From: abhijeetgangan Date: Mon, 14 Apr 2025 08:01:20 -0700 Subject: [PATCH 1/4] Add per atom energies and stresses for batched LJ --- .../1_Introduction/1.1_Lennard_Jones.py | 37 +++++++++++++++++++ torch_sim/models/lennard_jones.py | 17 ++++++++- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/examples/scripts/1_Introduction/1.1_Lennard_Jones.py b/examples/scripts/1_Introduction/1.1_Lennard_Jones.py index 91d8d26f..860834c6 100644 --- a/examples/scripts/1_Introduction/1.1_Lennard_Jones.py +++ b/examples/scripts/1_Introduction/1.1_Lennard_Jones.py @@ -10,6 +10,7 @@ import torch +from torch_sim.models.lennard_jones import LennardJonesModel from torch_sim.unbatched.models.lennard_jones import UnbatchedLennardJonesModel @@ -69,6 +70,8 @@ dtype=dtype, compute_forces=True, compute_stress=True, + per_atom_energies=True, + per_atom_stresses=True, ) # Print system information @@ -88,3 +91,37 @@ print(f"Energy: {results['energy']}") print(f"Forces: {results['forces']}") print(f"Stress: {results['stress']}") +print(f"Energies: {results['energies']}") +print(f"Stresses: {results['stresses']}") + +# Batched model +batched_model = LennardJonesModel( + use_neighbor_list=True, + cutoff=2.5 * 3.405, + sigma=3.405, + epsilon=0.0104, + device=device, + dtype=dtype, + compute_forces=True, + compute_stress=True, + per_atom_energies=True, + per_atom_stresses=True, +) + +# Batched state +state = dict( + positions=positions, + cell=cell.unsqueeze(0), + atomic_numbers=atomic_numbers, + pbc=True, +) + +# Run the simulation and get results +results = batched_model(state) + +# Print the results +print(f"Energy: {results['energy']}") +print(f"Forces: {results['forces']}") +print(f"Stress: {results['stress']}") +print(f"Energies: {results['energies']}") +print(f"Stresses: {results['stresses']}") diff --git a/torch_sim/models/lennard_jones.py b/torch_sim/models/lennard_jones.py index 85f691c1..bd47038f 100644 --- a/torch_sim/models/lennard_jones.py +++ b/torch_sim/models/lennard_jones.py @@ -275,7 +275,7 @@ def unbatched_forward( return results - def forward(self, state: SimState | StateDict) -> dict[str, torch.Tensor]: + def forward(self, state: SimState | StateDict) -> dict[str, torch.Tensor]: # noqa: C901 """Compute Lennard-Jones energies, forces, and stresses for a system. Main entry point for Lennard-Jones calculations that handles batched states by @@ -293,7 +293,10 @@ def forward(self, state: SimState | StateDict) -> dict[str, torch.Tensor]: compute_forces=True) - "stress": Stress tensor with shape [n_batches, 3, 3] (if compute_stress=True) - - May include additional outputs based on configuration + - "energies": Per-atom energies with shape [n_atoms] (if + per_atom_energies=True) + - "stresses": Per-atom stresses with shape [n_atoms, 3, 3] (if + per_atom_stresses=True) Raises: ValueError: If batch cannot be inferred for multi-cell systems. @@ -307,6 +310,8 @@ def forward(self, state: SimState | StateDict) -> dict[str, torch.Tensor]: energy = results["energy"] # Shape: [n_batches] forces = results["forces"] # Shape: [n_atoms, 3] stress = results["stress"] # Shape: [n_batches, 3, 3] + energies = results["energies"] # Shape: [n_atoms] + stresses = results["stresses"] # Shape: [n_atoms, 3, 3] """ if isinstance(state, dict): state = SimState(**state, masses=torch.ones_like(state["positions"])) @@ -328,4 +333,12 @@ def forward(self, state: SimState | StateDict) -> dict[str, torch.Tensor]: if key in properties: results[key] = torch.cat([out[key] for out in outputs], dim=0) + for key in ("energies",): + if key in properties: + results[key] = torch.cat([out[key] for out in outputs], dim=0) + + for key in ("stresses",): + if key in properties: + results[key] = torch.cat([out[key] for out in outputs], dim=0) + return results From f957668de416f4150a0dc1a1a2001bea26b07181 Mon Sep 17 00:00:00 2001 From: Orion Cohen <27712051+orionarcher@users.noreply.github.com> Date: Mon, 14 Apr 2025 15:03:30 -0400 Subject: [PATCH 2/4] update changelog for v0.2.0 (#147) * update changelog for v0.2.0 * minor modification for PR template * formatting fixes * formatting and typos * remove contributors bc they aren't linked --- .github/PULL_REQUEST_TEMPLATE.md | 2 +- CHANGELOG.md | 37 +++++++++++++++++++++++++------- 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 6dd0e293..d450f442 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -7,7 +7,7 @@ Include a summary of major changes in bullet points: ## Checklist -Work-in-progress pull requests are encouraged, but please enable the draft status on your PR. + Before a pull request can be merged, the following items must be checked: diff --git a/CHANGELOG.md b/CHANGELOG.md index 5142f977..9366a4a9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,10 +1,31 @@ -## v0.0.1 +## v0.2.0 -Initial release. +### Bug Fixes ๐Ÿ› +* Fix integrate reporting kwarg to arg error, https://github.com/Radical-AI/torch-sim/issues/113 (raised by @hn-yu) +* Allow runners to take large initial batches, https://github.com/Radical-AI/torch-sim/issues/128 (raised by @YutackPark) +* Add Fairchem model support for PBC, https://github.com/Radical-AI/torch-sim/issues/111 (raised by @ryanliu30) + +### Enhancements ๐Ÿ›  +* **[breaking]** Rename `HotSwappingAutobatcher` to `InFlightAutobatcher` and `ChunkingAutoBatcher` to `BinningAutoBatcher`, https://github.com/Radical-AI/torch-sim/pull/143 @orionarcher +* Support for Orbv3, https://github.com/Radical-AI/torch-sim/pull/140, @AdeeshKolluru +* Support metatensor models, https://github.com/Radical-AI/torch-sim/pull/141, @frostedoyter @Luthaf +* Support for graph-pes models, https://github.com/Radical-AI/torch-sim/pull/118 @jla-gardner +* Support MatterSim and fix ASE cell convention issues, https://github.com/Radical-AI/torch-sim/pull/112 @CompRhys +* Implement positions only FIRE optimization, https://github.com/Radical-AI/torch-sim/pull/139 @abhijeetgangan +* Allow different temperatures in batches, https://github.com/Radical-AI/torch-sim/pull/123 @orionarcher +* FairChem model updates: PBC handling, test on OMat24 e-trained model, https://github.com/Radical-AI/torch-sim/pull/126 @AdeeshKolluru +* FairChem model from_data_list support, https://github.com/Radical-AI/torch-sim/pull/138 @ryanliu30 +* New correlation function module, https://github.com/Radical-AI/torch-sim/pull/115 @stefanbringuier + +### Documentation ๐Ÿ“– +* Imoved model documentation, https://github.com/Radical-AI/torch-sim/pull/121 @orionarcher +* Plot of TorchSim module graph in docs, https://github.com/Radical-AI/torch-sim/pull/132 @janosh -[contributors]: <> (CONTRIBUTOR SECTION) -[abhijeetgangan]: https://github.com/abhijeetgangan -[orionarcher]: https://github.com/orionarcher -[janosh]: https://github.com/janosh -[AdeeshKolluru]: https://github.com/AdeeshKolluru -[CompRhys]: https://github.com/CompRhys +### House-Keeping ๐Ÿงน +* Only install HF for fairchem tests, https://github.com/Radical-AI/torch-sim/pull/134 @CompRhys +* Don't download MBD in CI, https://github.com/Radical-AI/torch-sim/pull/135 @orionarcher +* Tigthen graph-pes test bounds, https://github.com/Radical-AI/torch-sim/pull/143 @orionarcher + +## v0.1.0 + +Initial release. From f1fb38f750c8eb1cef12519e71541031c4a24a0a Mon Sep 17 00:00:00 2001 From: abhijeetgangan Date: Mon, 14 Apr 2025 19:11:17 -0700 Subject: [PATCH 3/4] Add tests --- CHANGELOG.md | 2 +- tests/models/test_lennard_jones.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9366a4a9..aea2ef38 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,7 +24,7 @@ ### House-Keeping ๐Ÿงน * Only install HF for fairchem tests, https://github.com/Radical-AI/torch-sim/pull/134 @CompRhys * Don't download MBD in CI, https://github.com/Radical-AI/torch-sim/pull/135 @orionarcher -* Tigthen graph-pes test bounds, https://github.com/Radical-AI/torch-sim/pull/143 @orionarcher +* Tighten graph-pes test bounds, https://github.com/Radical-AI/torch-sim/pull/143 @orionarcher ## v0.1.0 diff --git a/tests/models/test_lennard_jones.py b/tests/models/test_lennard_jones.py index 2d6ec15c..3a7fa1f5 100644 --- a/tests/models/test_lennard_jones.py +++ b/tests/models/test_lennard_jones.py @@ -155,6 +155,8 @@ def models( "dtype": torch.float64, "compute_forces": True, "compute_stress": True, + "per_atom_energies": True, + "per_atom_stresses": True, } cutoff = 2.5 * 3.405 # Standard LJ cutoff * sigma @@ -178,6 +180,14 @@ def test_energy_match( assert torch.allclose(results_nl["energy"], results_direct["energy"], rtol=1e-10) +def test_per_atom_energy_match( + models: tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]], +) -> None: + """Test that per-atom energy matches between neighbor list and direct calculations.""" + results_nl, results_direct = models + assert torch.allclose(results_nl["energies"], results_direct["energies"], rtol=1e-10) + + def test_forces_match( models: tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]], ) -> None: @@ -194,6 +204,15 @@ def test_stress_match( assert torch.allclose(results_nl["stress"], results_direct["stress"], rtol=1e-10) +def test_per_atom_stress_match( + models: tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]], +) -> None: + """Test that per-atom stress tensors match between neighbor list + and direct calculations.""" + results_nl, results_direct = models + assert torch.allclose(results_nl["stresses"], results_direct["stresses"], rtol=1e-10) + + def test_force_conservation( models: tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]], ) -> None: From dc05593add64f25d35d3f8a1f60223b07dcada97 Mon Sep 17 00:00:00 2001 From: abhijeetgangan Date: Mon, 14 Apr 2025 19:16:14 -0700 Subject: [PATCH 4/4] simplify results --- torch_sim/models/lennard_jones.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/torch_sim/models/lennard_jones.py b/torch_sim/models/lennard_jones.py index bd47038f..4910bea4 100644 --- a/torch_sim/models/lennard_jones.py +++ b/torch_sim/models/lennard_jones.py @@ -275,7 +275,7 @@ def unbatched_forward( return results - def forward(self, state: SimState | StateDict) -> dict[str, torch.Tensor]: # noqa: C901 + def forward(self, state: SimState | StateDict) -> dict[str, torch.Tensor]: """Compute Lennard-Jones energies, forces, and stresses for a system. Main entry point for Lennard-Jones calculations that handles batched states by @@ -329,15 +329,7 @@ def forward(self, state: SimState | StateDict) -> dict[str, torch.Tensor]: # no for key in ("stress", "energy"): if key in properties: results[key] = torch.stack([out[key] for out in outputs]) - for key in ("forces",): - if key in properties: - results[key] = torch.cat([out[key] for out in outputs], dim=0) - - for key in ("energies",): - if key in properties: - results[key] = torch.cat([out[key] for out in outputs], dim=0) - - for key in ("stresses",): + for key in ("forces", "energies", "stresses"): if key in properties: results[key] = torch.cat([out[key] for out in outputs], dim=0)