From 05823243d1e41152d9018566a5e6a0b4d7687a20 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Mon, 23 Mar 2026 11:35:31 +0900 Subject: [PATCH 1/5] fix --- torch_sim/integrators/md.py | 2 +- torch_sim/integrators/npt.py | 15 +++++++-------- torch_sim/integrators/nvt.py | 13 +++++++------ 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/torch_sim/integrators/md.py b/torch_sim/integrators/md.py index 76ff69c1..4a593c04 100644 --- a/torch_sim/integrators/md.py +++ b/torch_sim/integrators/md.py @@ -409,7 +409,7 @@ def init_fn( Q = ( kT_batched.unsqueeze(-1) - * torch.square(tau_batched).unsqueeze(-1) ** 2 + * torch.square(tau_batched).unsqueeze(-1) * torch.ones((n_systems, chain_length), dtype=dtype, device=device) ) Q[:, 0] *= degrees_of_freedom diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py index 4a86084c..19e1fda0 100644 --- a/torch_sim/integrators/npt.py +++ b/torch_sim/integrators/npt.py @@ -1181,7 +1181,7 @@ def _npt_nose_hoover_compute_cell_force( # Compute force on cell coordinate per system # F = alpha * KE - dU/dV - P*V*d return ( - (alpha * KE_per_system) + (alpha * 2 * KE_per_system) - (internal_pressure * volume) - (external_pressure * volume * dim) ) @@ -1226,11 +1226,9 @@ def _npt_nose_hoover_inner_step( volume, volume_to_cell = _npt_nose_hoover_cell_info(state) cell = volume_to_cell(volume) - # Get model output - state.cell = cell - model_output = model(state) - # First half step: Update momenta + # Reuse forces/stress from previous step (positions and cell unchanged; + # NHC chain half-steps only scale momenta/cell_momentum, not geometry) n_atoms_per_system = torch.bincount(state.system_idx, minlength=state.n_systems) alpha = 1 + 1 / n_atoms_per_system # [n_systems] @@ -1240,7 +1238,7 @@ def _npt_nose_hoover_inner_step( positions=positions, momenta=momenta, masses=masses, - stress=model_output["stress"], + stress=state.stress, external_pressure=external_pressure, system_idx=state.system_idx, ) @@ -1406,7 +1404,7 @@ def npt_nose_hoover_init( ) # Compute total DOF for thermostat initialization and a zero KE placeholder - dof_per_system = torch.bincount(state.system_idx, minlength=n_systems) * dim + dof_per_system = torch.bincount(state.system_idx, minlength=n_systems) * dim - dim KE_thermostat = ts.calc_kinetic_energy( masses=state.masses, momenta=momenta, system_idx=state.system_idx ) @@ -1613,7 +1611,8 @@ def npt_nose_hoover_invariant( # Calculate degrees of freedom per system n_atoms_per_system = torch.bincount(state.system_idx, minlength=state.n_systems) - dof_per_system = n_atoms_per_system * state.positions.shape[-1] # n_atoms * n_dim + dim = state.positions.shape[-1] + dof_per_system = n_atoms_per_system * dim - dim # 3N - 3 (COM correction) # Initialize total energy with PE + KE e_tot = e_pot + e_kin_per_system diff --git a/torch_sim/integrators/nvt.py b/torch_sim/integrators/nvt.py index 8e74bf85..c0b41628 100644 --- a/torch_sim/integrators/nvt.py +++ b/torch_sim/integrators/nvt.py @@ -314,11 +314,11 @@ def nvt_nose_hoover_init( masses=state.masses, momenta=momenta, system_idx=state.system_idx ) - # Calculate degrees of freedom per system + # Calculate degrees of freedom per system (subtract 3 for COM motion, + # matching LAMMPS compute_temp which uses dof = 3N - 3) n_atoms_per_system = torch.bincount(state.system_idx) - dof_per_system = ( - n_atoms_per_system * state.positions.shape[-1] - ) # n_atoms * n_dimensions + dim = state.positions.shape[-1] + dof_per_system = n_atoms_per_system * dim - dim # 3N - 3 # Initialize state return NVTNoseHooverState.from_state( @@ -431,9 +431,10 @@ def nvt_nose_hoover_invariant( masses=state.masses, momenta=state.momenta, system_idx=state.system_idx ) - # Get system degrees of freedom per system + # Get system degrees of freedom per system (3N - 3 for COM correction) n_atoms_per_system = torch.bincount(state.system_idx) - dof = n_atoms_per_system * state.positions.shape[-1] # n_atoms * n_dimensions + dim = state.positions.shape[-1] + dof = n_atoms_per_system * dim - dim # Start with system energy e_tot = e_pot + e_kin From 5c04f395dbce2b3c568b6b177c2eafaab4ddb186 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Mon, 23 Mar 2026 12:18:22 +0900 Subject: [PATCH 2/5] fix --- torch_sim/integrators/npt.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py index 19e1fda0..9206767c 100644 --- a/torch_sim/integrators/npt.py +++ b/torch_sim/integrators/npt.py @@ -1179,7 +1179,7 @@ def _npt_nose_hoover_compute_cell_force( internal_pressure = torch.trace(stress).unsqueeze(0).expand(n_systems) # Compute force on cell coordinate per system - # F = alpha * KE - dU/dV - P*V*d + # F = alpha * (2 * KE) - dU/dV - P*V*d return ( (alpha * 2 * KE_per_system) - (internal_pressure * volume) From 431834dc1d345ae2fef0b41050cd5e548fe2e926 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Wed, 25 Mar 2026 17:01:45 +0900 Subject: [PATCH 3/5] use state.get_number_of_degrees_of_freedom --- torch_sim/integrators/npt.py | 16 +++++++--------- torch_sim/integrators/nvt.py | 6 ++---- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py index 9206767c..c236bd45 100644 --- a/torch_sim/integrators/npt.py +++ b/torch_sim/integrators/npt.py @@ -1227,11 +1227,10 @@ def _npt_nose_hoover_inner_step( cell = volume_to_cell(volume) # First half step: Update momenta - # Reuse forces/stress from previous step (positions and cell unchanged; - # NHC chain half-steps only scale momenta/cell_momentum, not geometry) - n_atoms_per_system = torch.bincount(state.system_idx, minlength=state.n_systems) - alpha = 1 + 1 / n_atoms_per_system # [n_systems] + # alpha = 1 + dim / degrees_of_freedom (3 * natoms - 3) + alpha = 1 + 3 / state.get_number_of_degrees_of_freedom() # [n_systems] + # Reuse stress from previous step since positions and cell unchanged cell_force_val = _npt_nose_hoover_compute_cell_force( alpha=alpha, volume=volume, @@ -1404,7 +1403,8 @@ def npt_nose_hoover_init( ) # Compute total DOF for thermostat initialization and a zero KE placeholder - dof_per_system = torch.bincount(state.system_idx, minlength=n_systems) * dim - dim + dof_per_system = state.get_number_of_degrees_of_freedom() - 3 + KE_thermostat = ts.calc_kinetic_energy( masses=state.masses, momenta=momenta, system_idx=state.system_idx ) @@ -1610,14 +1610,12 @@ def npt_nose_hoover_invariant( ) # Calculate degrees of freedom per system - n_atoms_per_system = torch.bincount(state.system_idx, minlength=state.n_systems) - dim = state.positions.shape[-1] - dof_per_system = n_atoms_per_system * dim - dim # 3N - 3 (COM correction) + dof_per_system = state.get_number_of_degrees_of_freedom() # Initialize total energy with PE + KE e_tot = e_pot + e_kin_per_system - # Add thermostat chain contributions (batched per system, DOF = n_atoms * 3) + # Add thermostat chain contributions (batched per system, DOF = 3 * n_atoms - 3) e_tot += _compute_chain_energy(state.thermostat, kT, e_tot, dof_per_system) # Add barostat chain contributions (batched per system, DOF = 1) diff --git a/torch_sim/integrators/nvt.py b/torch_sim/integrators/nvt.py index c0b41628..31dd24ad 100644 --- a/torch_sim/integrators/nvt.py +++ b/torch_sim/integrators/nvt.py @@ -318,7 +318,7 @@ def nvt_nose_hoover_init( # matching LAMMPS compute_temp which uses dof = 3N - 3) n_atoms_per_system = torch.bincount(state.system_idx) dim = state.positions.shape[-1] - dof_per_system = n_atoms_per_system * dim - dim # 3N - 3 + dof_per_system = state.get_number_of_degrees_of_freedom() - 3 # Initialize state return NVTNoseHooverState.from_state( @@ -432,9 +432,7 @@ def nvt_nose_hoover_invariant( ) # Get system degrees of freedom per system (3N - 3 for COM correction) - n_atoms_per_system = torch.bincount(state.system_idx) - dim = state.positions.shape[-1] - dof = n_atoms_per_system * dim - dim + dof = state.get_number_of_degrees_of_freedom() # Start with system energy e_tot = e_pot + e_kin From 80521f1e9d3fd708456caf42926607bede97295f Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Wed, 25 Mar 2026 19:25:44 +0900 Subject: [PATCH 4/5] fix --- torch_sim/integrators/nvt.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/torch_sim/integrators/nvt.py b/torch_sim/integrators/nvt.py index 31dd24ad..1b801727 100644 --- a/torch_sim/integrators/nvt.py +++ b/torch_sim/integrators/nvt.py @@ -316,8 +316,6 @@ def nvt_nose_hoover_init( # Calculate degrees of freedom per system (subtract 3 for COM motion, # matching LAMMPS compute_temp which uses dof = 3N - 3) - n_atoms_per_system = torch.bincount(state.system_idx) - dim = state.positions.shape[-1] dof_per_system = state.get_number_of_degrees_of_freedom() - 3 # Initialize state From c092796ef9bfa0b14084736cb49471eb0bca553e Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Mon, 30 Mar 2026 09:54:49 -0400 Subject: [PATCH 5/5] fix: updated test constants npt and nvt --- tests/test_integrators.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_integrators.py b/tests/test_integrators.py index dd9b49d0..27ea5dba 100644 --- a/tests/test_integrators.py +++ b/tests/test_integrators.py @@ -339,7 +339,7 @@ def test_nvt_nose_hoover(ar_double_sim_state: ts.SimState, lj_model: LennardJone temperatures_list = [t.tolist() for t in temperatures_tensor.T] assert torch.allclose( temperatures_tensor[-1], - torch.tensor([300.0096, 299.7024], dtype=dtype), + torch.tensor([290.3553, 289.9699], dtype=dtype), ) energies_tensor = torch.stack(energies) @@ -728,7 +728,7 @@ def test_npt_nose_hoover(ar_double_sim_state: ts.SimState, lj_model: LennardJone temperatures_list = [t.tolist() for t in temperatures_tensor.T] assert torch.allclose( temperatures_tensor[-1], - torch.tensor([298.2752, 297.9444], dtype=dtype), + torch.tensor([287.5729, 287.1330], dtype=dtype), ) energies_tensor = torch.stack(energies)