From 26e57b11de5662f5f6f9332ee555a2e36edd4b46 Mon Sep 17 00:00:00 2001 From: abhijeetgangan Date: Sun, 25 Oct 2026 07:16:36 -0700 Subject: [PATCH] Fix ase_fire parity with ASE and temperature/pressure shape handling --- tests/test_optimizers.py | 5 +++ tests/test_optimizers_vs_ase.py | 7 +++++ tests/test_runners.py | 23 ++++++++++++++ torch_sim/integrators/npt.py | 7 +++++ torch_sim/optimizers/fire.py | 55 +++++++++++++++++---------------- torch_sim/runners.py | 6 ++-- 6 files changed, 73 insertions(+), 30 deletions(-) diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py index 390e44df..3021549f 100644 --- a/tests/test_optimizers.py +++ b/tests/test_optimizers.py @@ -785,6 +785,11 @@ def test_fire_nan_velocities_dont_affect_other_systems( "System 0 cell differs when system 1 has NaN velocities" ) + # System 1 now takes a first step, which leaves dt, alpha and n_pos unchanged + if fire_flavor == "ase_fire": + for attr in ("dt", "alpha", "n_pos"): + assert getattr(state_mixed, attr)[1] == getattr(state, attr)[1], attr + @pytest.mark.parametrize("cell_filter", [ts.CellFilter.unit, ts.CellFilter.frechet]) def test_vv_fire_nan_cell_velocities_with_finite_atomic_velocities( diff --git a/tests/test_optimizers_vs_ase.py b/tests/test_optimizers_vs_ase.py index 8e6f67b7..7867e037 100644 --- a/tests/test_optimizers_vs_ase.py +++ b/tests/test_optimizers_vs_ase.py @@ -290,6 +290,13 @@ def _run_and_compare_optimizers( 1e-11, "Ar stretched-rattled supercell (Frechet)", ), + ( + "rattled_ar_supercell_stretched_sim_state", + 10, + 1e-10, + 5e-11, + "Ar stretched-rattled supercell past first alpha decay (Frechet)", + ), ] BFGS_CASES = [ diff --git a/tests/test_runners.py b/tests/test_runners.py index b7044328..80446301 100644 --- a/tests/test_runners.py +++ b/tests/test_runners.py @@ -239,6 +239,29 @@ def test_integrate_double_nvt_multiple_temperatures( init_kwargs={}, ) + # n_steps == n_systems: 1D temperatures are per system, 2D are (n_steps, n_systems) + init, real_step = ts.integrators.INTEGRATOR_REGISTRY[ts.Integrator.nvt_langevin] + step_kTs: list[torch.Tensor] = [] + + def spy_step(**kwargs: object) -> MDState: + step_kTs.append(kwargs["kT"]) + return real_step(**kwargs) + + temps = torch.tensor([100.0, 200.0], dtype=DTYPE, device=DEVICE) + for temperature in (temps, temps.expand(2, 2)): + step_kTs.clear() + ts.integrate( + system=ar_double_sim_state, + model=lj_model, + integrator=(init, spy_step), + n_steps=2, + temperature=temperature, + timestep=0.001, + ) + assert len(step_kTs) == 2 + for kT in step_kTs: + torch.testing.assert_close(kT, temps * ts.units.MetalUnits.temperature) + def test_integrate_double_nvt_with_reporter( ar_double_sim_state: SimState, lj_model: LennardJonesModel, tmp_path: Path diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py index b028f74f..4a4bca3d 100644 --- a/torch_sim/integrators/npt.py +++ b/torch_sim/integrators/npt.py @@ -576,6 +576,13 @@ def npt_langevin_anisotropic_step( if external_pressure_tensor.ndim == 0: external_pressure_tensor = external_pressure_tensor.expand(state.n_systems, 3) elif external_pressure_tensor.ndim == 1 and external_pressure_tensor.shape[0] == 3: + if state.n_systems == 3: + warnings.warn( + "external_pressure of shape (3,) is ambiguous for 3 systems and is read " + "as the (x, y, z) pressure of every system. Pass shape (1, 3) for that, " + "or shape (3, 3) for per-system pressures.", + stacklevel=2, + ) external_pressure_tensor = external_pressure_tensor.unsqueeze(0).expand( state.n_systems, 3 ) diff --git a/torch_sim/optimizers/fire.py b/torch_sim/optimizers/fire.py index c037d1d1..88d565eb 100644 --- a/torch_sim/optimizers/fire.py +++ b/torch_sim/optimizers/fire.py @@ -301,22 +301,27 @@ def _ase_fire_step[T: "FireState | CellFireState"]( # noqa: C901, PLR0915 n_systems = state.dt.numel() - # NaN velocities mark newly initialized atoms. Skip adaptation for a wholly - # new batch; replacing only part of an autobatch must not skip the others. - nan_velocities = state.velocities.isnan().any(dim=1) + # NaN velocities mark systems that haven't stepped yet, including ones swapped in + # by an autobatcher; reduce per system without a GPU-syncing Python branch + nan_velocities = state.velocities.isnan().any(dim=1).to(torch.int32) + is_first_step = ( + torch.zeros(n_systems, device=state.device, dtype=torch.int32) + .scatter_add_(0, state.system_idx, nan_velocities) + .bool() + ) state.velocities.nan_to_num_(nan=0.0) if isinstance(state, CellFireState): state.cell_velocities.nan_to_num_(nan=0.0) - # All atoms have NaN velocities when every system in the batch is taking its - # first step. Keep this as a tensor: a Python branch would synchronize the GPU. - is_first_step = nan_velocities.all() - # Only cell states have a reference cell to define deform_grad; ASE's cell # filters hand FIRE `forces @ deform_grad`, a plain FireState uses raw forces. forces = state.forces if isinstance(state, CellFireState): - forces = torch.where(is_first_step, forces, state.deform_grad_forces()) + forces = torch.where( + is_first_step[state.system_idx].unsqueeze(-1), + forces, + state.deform_grad_forces(), + ) # Calculate power. Newly zeroed systems have zero power. system_power = tsm.batched_vdot( @@ -325,25 +330,15 @@ def _ase_fire_step[T: "FireState | CellFireState"]( # noqa: C901, PLR0915 if isinstance(state, CellFireState): system_power += (state.cell_forces * state.cell_velocities).sum(dim=(1, 2)) - # Update dt, alpha, n_pos pos_mask_system = system_power > 0.0 neg_mask_system = ~pos_mask_system - # Only update dt, alpha, and n_pos after the first step. + # As in ASE, each system's first step leaves dt, alpha and n_pos unchanged pos_mask_system &= ~is_first_step neg_mask_system &= ~is_first_step - inc_mask = (state.n_pos > n_min) & pos_mask_system - state.dt = torch.where(inc_mask, torch.minimum(state.dt * f_inc, dt_max), state.dt) - state.alpha = torch.where(inc_mask, state.alpha * f_alpha, state.alpha) - state.n_pos = state.n_pos + pos_mask_system.to(state.n_pos.dtype) - - state.dt = torch.where(neg_mask_system, state.dt * f_dec, state.dt) - state.alpha = torch.where(neg_mask_system, alpha_start, state.alpha) - state.n_pos = torch.where(neg_mask_system, 0, state.n_pos) - - # Mix velocities before acceleration (ASE ordering). The first step keeps - # the initialized velocities, so the selections below also skip mixing. + # As in ASE, mix with the current alpha before adapting it and before acceleration; + # uphill and first-step systems restart from rest v_scaling_system = tsm.batched_vdot( state.velocities, state.velocities, state.system_idx, n_systems=n_systems ) @@ -360,26 +355,32 @@ def _ase_fire_step[T: "FireState | CellFireState"]( # noqa: C901, PLR0915 v_mixing_cell = state.cell_forces / (f_scaling_cell + eps) * v_scaling_cell alpha_cell_bc = state.alpha.view(n_systems, 1, 1) - mixed_cell_velocities = torch.where( + state.cell_velocities = torch.where( pos_mask_system.view(n_systems, 1, 1), (1.0 - alpha_cell_bc) * state.cell_velocities + alpha_cell_bc * v_mixing_cell, torch.zeros_like(state.cell_velocities), ) - state.cell_velocities = torch.where( - ~is_first_step, mixed_cell_velocities, state.cell_velocities - ) v_scaling_atom = torch.sqrt(v_scaling_system[state.system_idx].unsqueeze(-1)) f_scaling_atom = torch.sqrt(f_scaling_system[state.system_idx].unsqueeze(-1)) v_mixing_atom = forces * (v_scaling_atom / (f_scaling_atom + eps)) alpha_atom = state.alpha[state.system_idx].unsqueeze(-1) - mixed_velocities = torch.where( + state.velocities = torch.where( pos_mask_system[state.system_idx].unsqueeze(-1), (1.0 - alpha_atom) * state.velocities + alpha_atom * v_mixing_atom, torch.zeros_like(state.velocities), ) - state.velocities = torch.where(~is_first_step, mixed_velocities, state.velocities) + + # Update dt, alpha, n_pos + inc_mask = (state.n_pos > n_min) & pos_mask_system + state.dt = torch.where(inc_mask, torch.minimum(state.dt * f_inc, dt_max), state.dt) + state.alpha = torch.where(inc_mask, state.alpha * f_alpha, state.alpha) + state.n_pos = state.n_pos + pos_mask_system.to(state.n_pos.dtype) + + state.dt = torch.where(neg_mask_system, state.dt * f_dec, state.dt) + state.alpha = torch.where(neg_mask_system, alpha_start, state.alpha) + state.n_pos = torch.where(neg_mask_system, 0, state.n_pos) # Acceleration (single forward-Euler, no mass for ASE FIRE) state.velocities += forces * state.dt[state.system_idx].unsqueeze(-1) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index 21eac389..4278dce1 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -198,16 +198,16 @@ def _normalize_temperature_tensor( # A single value return temps.expand(n_steps) - if initial_state.n_systems == n_steps: + if temps.ndim == 1 and temps.shape[0] == initial_state.n_systems == n_steps: msg = ( "n_systems is equal to n_steps. Interpreting temperature array of length " "n_systems as temperatures for each system, broadcasted over steps." ) warnings.warn(msg, stacklevel=2) logger.warning(msg) - return temps.expand(n_steps) + return temps.unsqueeze(0).expand(n_steps, -1) # (n_steps, n_systems) - if temps.shape[0] == initial_state.n_systems: + if temps.shape[0] == initial_state.n_systems and initial_state.n_systems != n_steps: if temps.ndim == 2: raise ValueError( "If temperature tensor is 2D, first dimension must be n_steps."