From 9a651468b56d6135a5815c8061b13f3d58e12ce6 Mon Sep 17 00:00:00 2001 From: Teddy Koker Date: Fri, 5 Jun 2026 14:22:55 -0700 Subject: [PATCH 1/3] correct sign for d3 --- torch_sim/models/dispersion.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torch_sim/models/dispersion.py b/torch_sim/models/dispersion.py index a60ede7c6..2c7862ebd 100644 --- a/torch_sim/models/dispersion.py +++ b/torch_sim/models/dispersion.py @@ -178,6 +178,8 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor] # d3_out[3] is only defined if compute_virial is True # we use [-1] to index it to avoid typing errors. volumes = state.volume.unsqueeze(-1).unsqueeze(-1) - stress = (d3_out[-1] * UnitConversion.Hartree_to_eV) / volumes + # nvalchemiops returns the negative strain-gradient virial; TorchSim stress + # follows dE / dstrain / volume, matching ASE and other TorchSim models. + stress = -(d3_out[-1] * UnitConversion.Hartree_to_eV) / volumes results["stress"] = stress.to(self._dtype).detach() return results From 398c007661a03bab693069fce25ec5499426ae40 Mon Sep 17 00:00:00 2001 From: Teddy Koker Date: Fri, 5 Jun 2026 14:34:46 -0700 Subject: [PATCH 2/3] add test --- tests/models/test_dispersion.py | 105 ++++++++++++++++++++++++++++++++ 1 file changed, 105 insertions(+) diff --git a/tests/models/test_dispersion.py b/tests/models/test_dispersion.py index 5e7fe5b37..36ef18c18 100644 --- a/tests/models/test_dispersion.py +++ b/tests/models/test_dispersion.py @@ -7,6 +7,7 @@ import pytest import torch +import torch_sim as ts from tests.conftest import DEVICE, DTYPE from tests.models.conftest import make_validate_model_outputs_test @@ -69,6 +70,110 @@ def d3_model_r2scan() -> D3DispersionModel: ) +def test_d3_stress_matches_finite_strain_sign() -> None: + """Stress should match dE/dstrain/V, not the opposite virial sign.""" + row_cell = torch.tensor( + [[4.2, 0.3, 0.1], [0.2, 4.8, 0.4], [0.15, 0.35, 5.1]], + dtype=DTYPE, + device=DEVICE, + ) + positions = torch.tensor( + [[0.4, 0.5, 0.6], [1.9, 1.4, 2.3], [3.1, 2.6, 1.7]], + dtype=DTYPE, + device=DEVICE, + ) + state = ts.SimState( + positions=positions, + masses=torch.ones(3, dtype=DTYPE, device=DEVICE), + cell=row_cell.mT.unsqueeze(0), + pbc=True, + atomic_numbers=torch.tensor([6, 8, 14], dtype=torch.int64, device=DEVICE), + ) + + gen = torch.Generator(device=DEVICE) + gen.manual_seed(1234) + max_z = 14 + mesh = 5 + rcov = torch.rand(max_z + 1, generator=gen, device=DEVICE) + 0.5 + r4r2 = torch.rand(max_z + 1, generator=gen, device=DEVICE) + 0.5 + c6ab = 20.0 * ( + torch.rand( + max_z + 1, + max_z + 1, + mesh, + mesh, + generator=gen, + device=DEVICE, + ) + + 0.1 + ) + c6ab = 0.5 * (c6ab + c6ab.permute(1, 0, 3, 2)) + cn_ref = 4.0 * torch.rand( + max_z + 1, + max_z + 1, + mesh, + mesh, + generator=gen, + device=DEVICE, + ) + cn_ref = 0.5 * (cn_ref + cn_ref.permute(1, 0, 3, 2)) + + model = D3DispersionModel( + **PBE_BJ, + d3_params=D3Parameters(rcov=rcov, r4r2=r4r2, c6ab=c6ab, cn_ref=cn_ref), + cutoff=6.0, + device=DEVICE, + dtype=DTYPE, + compute_forces=True, + compute_stress=True, + ) + + stress = model(state)["stress"][0] + volume = state.volume[0] + frac_positions = torch.linalg.solve(row_cell.mT, positions.mT).mT + identity = torch.eye(3, dtype=DTYPE, device=DEVICE) + + def strained_energy(strain: torch.Tensor) -> torch.Tensor: + strained_row_cell = row_cell @ (identity + strain) + strained_state = ts.SimState( + positions=frac_positions @ strained_row_cell, + masses=state.masses, + cell=strained_row_cell.mT.unsqueeze(0), + pbc=state.pbc, + atomic_numbers=state.atomic_numbers, + system_idx=state.system_idx, + ) + return model(strained_state)["energy"][0] + + step = 1e-3 + finite_diff_stress = torch.zeros((3, 3), dtype=DTYPE, device=DEVICE) + for idx_i in range(3): + for idx_j in range(idx_i, 3): + strain = torch.zeros((3, 3), dtype=DTYPE, device=DEVICE) + if idx_i == idx_j: + strain[idx_i, idx_i] = step + energy_plus = strained_energy(strain) + strain[idx_i, idx_i] = -step + energy_minus = strained_energy(strain) + else: + strain[idx_i, idx_j] = 0.5 * step + strain[idx_j, idx_i] = 0.5 * step + energy_plus = strained_energy(strain) + strain[idx_i, idx_j] = -0.5 * step + strain[idx_j, idx_i] = -0.5 * step + energy_minus = strained_energy(strain) + stress_component = (energy_plus - energy_minus) / (2 * step * volume) + finite_diff_stress[idx_i, idx_j] = stress_component + finite_diff_stress[idx_j, idx_i] = stress_component + + torch.testing.assert_close( + stress, + finite_diff_stress, + rtol=5e-3, + atol=5e-7, + ) + + test_d3_pbe_outputs = make_validate_model_outputs_test( model_fixture_name="d3_model_pbe", device=DEVICE, dtype=DTYPE ) From 20d603268eb258a87afb343d3b390ae495789ee7 Mon Sep 17 00:00:00 2001 From: Teddy Koker Date: Sat, 6 Jun 2026 08:41:19 -0700 Subject: [PATCH 3/3] pin alchemi to 0.3.1, fix dsf, ewald, and pme stress sign, add tests --- pyproject.toml | 2 +- tests/models/test_electrostatics.py | 92 +++++++++++++++++++++++++++++ torch_sim/models/electrostatics.py | 12 +++- 3 files changed, 102 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9f473cd5d..55615ddaa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ "h5py>=3.15.1", "numpy>=1.26,<3; python_version < '3.13'", "numpy>=2.3.2,<3; python_version >= '3.13'", - "nvalchemi-toolkit-ops[torch]>=0.3.0", + "nvalchemi-toolkit-ops[torch]>=0.3.1", "tables>=3.11.1", "torch>=2", "tqdm>=4.67", diff --git a/tests/models/test_electrostatics.py b/tests/models/test_electrostatics.py index 3621bea69..8cd8bbcef 100644 --- a/tests/models/test_electrostatics.py +++ b/tests/models/test_electrostatics.py @@ -93,6 +93,98 @@ def test_dsf_nonzero_energy() -> None: assert out["energy"].abs().item() > 0 +@pytest.mark.parametrize( + ("model_cls", "kwargs"), + [ + pytest.param(DSFCoulombModel, {"cutoff": 8.0, "alpha": 0.2}, id="dsf"), + pytest.param(EwaldModel, {"cutoff": 8.0, "accuracy": 1e-6}, id="ewald"), + pytest.param( + PMEModel, + {"cutoff": 8.0, "accuracy": 1e-6, "mesh_spacing": 1.0}, + id="pme", + ), + ], +) +def test_electrostatics_stress_matches_finite_strain_sign( + model_cls: type[DSFCoulombModel | EwaldModel | PMEModel], + kwargs: dict[str, float], +) -> None: + """Electrostatic stress should match dE/dstrain/V, not the virial sign.""" + row_cell = torch.tensor( + [[5.2, 0.3, 0.1], [0.2, 5.6, 0.4], [0.15, 0.35, 6.1]], + dtype=DTYPE, + device=DEVICE, + ) + positions = torch.tensor( + [[0.4, 0.5, 0.6], [1.9, 1.4, 2.3], [3.1, 2.6, 1.7], [4.0, 3.4, 4.2]], + dtype=DTYPE, + device=DEVICE, + ) + charges = torch.tensor([0.8, -0.7, 0.4, -0.5], dtype=DTYPE, device=DEVICE) + state = ts.SimState( + positions=positions, + masses=torch.ones(4, dtype=DTYPE, device=DEVICE), + cell=row_cell.mT.unsqueeze(0), + pbc=True, + atomic_numbers=torch.tensor([11, 17, 11, 17], dtype=torch.int64, device=DEVICE), + ) + state._atom_extras["partial_charges"] = charges # noqa: SLF001 + + model = model_cls( + **kwargs, + device=DEVICE, + dtype=DTYPE, + compute_forces=True, + compute_stress=True, + ) + + stress = model(state)["stress"][0] + volume = state.volume[0] + frac_positions = torch.linalg.solve(row_cell.mT, positions.mT).mT + identity = torch.eye(3, dtype=DTYPE, device=DEVICE) + + def strained_energy(strain: torch.Tensor) -> torch.Tensor: + strained_row_cell = row_cell @ (identity + strain) + strained_state = ts.SimState( + positions=frac_positions @ strained_row_cell, + masses=state.masses, + cell=strained_row_cell.mT.unsqueeze(0), + pbc=state.pbc, + atomic_numbers=state.atomic_numbers, + system_idx=state.system_idx, + ) + strained_state._atom_extras["partial_charges"] = charges # noqa: SLF001 + return model(strained_state)["energy"][0] + + step = 1e-3 + finite_diff_stress = torch.zeros((3, 3), dtype=DTYPE, device=DEVICE) + for idx_i in range(3): + for idx_j in range(idx_i, 3): + strain = torch.zeros((3, 3), dtype=DTYPE, device=DEVICE) + if idx_i == idx_j: + strain[idx_i, idx_i] = step + energy_plus = strained_energy(strain) + strain[idx_i, idx_i] = -step + energy_minus = strained_energy(strain) + else: + strain[idx_i, idx_j] = 0.5 * step + strain[idx_j, idx_i] = 0.5 * step + energy_plus = strained_energy(strain) + strain[idx_i, idx_j] = -0.5 * step + strain[idx_j, idx_i] = -0.5 * step + energy_minus = strained_energy(strain) + stress_component = (energy_plus - energy_minus) / (2 * step * volume) + finite_diff_stress[idx_i, idx_j] = stress_component + finite_diff_stress[idx_j, idx_i] = stress_component + + torch.testing.assert_close( + stress, + finite_diff_stress, + rtol=5e-3, + atol=5e-7, + ) + + def test_ewald_pme_energy_agreement() -> None: """Ewald and PME should give the same converged Coulomb energy.""" state = _make_charged_state() diff --git a/torch_sim/models/electrostatics.py b/torch_sim/models/electrostatics.py index 092f5bfee..f9c69e620 100644 --- a/torch_sim/models/electrostatics.py +++ b/torch_sim/models/electrostatics.py @@ -146,7 +146,9 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor] results["forces"] = forces.to(self._dtype).detach() if self._compute_stress: volumes = state.volume.unsqueeze(-1).unsqueeze(-1) - stress = (out[-1] * UnitConversion.e2_per_Ang_to_eV) / volumes + # nvalchemiops returns the negative strain-gradient virial; TorchSim stress + # follows dE / dstrain / volume, matching ASE and other TorchSim models. + stress = -(out[-1] * UnitConversion.e2_per_Ang_to_eV) / volumes results["stress"] = stress.to(self._dtype).detach() return results @@ -258,7 +260,9 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor] results["forces"] = forces.to(self._dtype).detach() if self._compute_stress: volumes = state.volume.unsqueeze(-1).unsqueeze(-1) - stress = (out[-1] * UnitConversion.e2_per_Ang_to_eV) / volumes + # nvalchemiops returns the negative strain-gradient virial; TorchSim stress + # follows dE / dstrain / volume, matching ASE and other TorchSim models. + stress = -(out[-1] * UnitConversion.e2_per_Ang_to_eV) / volumes results["stress"] = stress.to(self._dtype).detach() return results @@ -387,6 +391,8 @@ def forward(self, state: SimState, **_kwargs: object) -> dict[str, torch.Tensor] results["forces"] = forces.to(self._dtype).detach() if self._compute_stress: volumes = state.volume.unsqueeze(-1).unsqueeze(-1) - stress = (out[-1] * UnitConversion.e2_per_Ang_to_eV) / volumes + # nvalchemiops returns the negative strain-gradient virial; TorchSim stress + # follows dE / dstrain / volume, matching ASE and other TorchSim models. + stress = -(out[-1] * UnitConversion.e2_per_Ang_to_eV) / volumes results["stress"] = stress.to(self._dtype).detach() return results