Fix ase_fire parity with ASE and temperature/pressure shape handling

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(

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 = [

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

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
         )

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)

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."