Fix/ase fire deform grad forces

tests/test_optimizers.py

@@ -10,7 +10,7 @@
 import torch_sim as ts
 from torch_sim.models.interface import ModelInterface
 from torch_sim.optimizers import BFGSState, FireFlavor, FireState, LBFGSState, OptimState
-from torch_sim.optimizers.cell_filters import CellLBFGSState, deform_grad
+from torch_sim.optimizers.cell_filters import CellLBFGSState, CellOptimState, deform_grad
 from torch_sim.state import SimState
 
 
@@ -1088,6 +1088,28 @@ def test_frechet_cell_fire_optimization(
     )
 
 
+@pytest.mark.parametrize("cell_filter", [ts.CellFilter.unit, ts.CellFilter.frechet])
+def test_cell_optim_state_deform_grad_forces(
+    ar_supercell_sim_state: SimState,
+    lj_model: ModelInterface,
+    cell_filter: ts.CellFilter,
+) -> None:
+    """deform_grad_forces applies forces @ deform_grad, like ASE's cell filters."""
+    state = ts.fire_init(
+        state=ar_supercell_sim_state, model=lj_model, cell_filter=cell_filter
+    )
+    assert isinstance(state, CellOptimState)
+
+    # freshly initialized, cell == reference_cell, so the transform is a no-op
+    torch.testing.assert_close(state.deform_grad_forces(), state.forces)
+
+    # against a reference cell scaled by s, deform_grad = I / s, so the
+    # transform scales the forces by 1 / s
+    scale = 0.95
+    state.reference_cell = state.cell.clone() * scale
+    torch.testing.assert_close(state.deform_grad_forces(), state.forces / scale)
+
+
 def test_frechet_lbfgs_clamps_extreme_deformation(
     ar_supercell_sim_state: SimState, lj_model: ModelInterface
 ) -> None:

torch_sim/optimizers/bfgs.py

@@ -247,25 +247,8 @@ def bfgs_step(  # noqa: C901, PLR0915
     )
 
     if isinstance(state, CellBFGSState):
-        # Get current deformation gradient
-        # reference_cell.mT: [S, 3, 3], row_vector_cell: [S, 3, 3]
-        cur_deform_grad = cell_filters.deform_grad(
-            state.reference_cell.mT, state.row_vector_cell
-        )  # [S, 3, 3]
-
-        # Transform forces to scaled coordinates
-        # forces: [N, 3], cur_deform_grad[system_idx]: [N, 3, 3]
-        forces_scaled = torch.bmm(
-            state.forces.unsqueeze(1),  # [N, 1, 3]
-            cur_deform_grad[state.system_idx],  # [N, 3, 3]
-        ).squeeze(1)  # [N, 3]
-
-        # Current fractional positions
-        # positions: [N, 3] -> frac_positions: [N, 3]
-        frac_positions = torch.linalg.solve(
-            cur_deform_grad[state.system_idx],  # [N, 3, 3]
-            state.positions.unsqueeze(-1),  # [N, 3, 1]
-        ).squeeze(-1)  # [N, 3]
+        forces_scaled = state.deform_grad_forces()  # [N, 3]
+        frac_positions = state.frac_positions()  # [N, 3]
 
         # Pack into dense tensors [N, 3] -> [S, M, 3] -> [S, D]
         # For cell state, prev_positions is already fractional (stored that way)
@@ -495,15 +478,7 @@ def bfgs_step(  # noqa: C901, PLR0915
         # Apply position step in fractional space, then convert to Cartesian
         new_frac = frac_positions + flat_step  # [N, 3]
 
-        new_deform_grad = cell_filters.deform_grad(
-            state.reference_cell.mT, state.row_vector_cell
-        )  # [S, 3, 3]
-        # new_positions = new_frac @ deform_grad^T
-        new_positions = torch.bmm(
-            new_frac.unsqueeze(1),  # [N, 1, 3]
-            new_deform_grad[state.system_idx].transpose(-2, -1),  # [N, 3, 3]
-        ).squeeze(1)  # [N, 3]
-        state.set_constrained_positions(new_positions)  # [N, 3]
+        state.set_constrained_positions(state.positions_from_frac(new_frac))  # [N, 3]
     else:
         state.prev_positions = state.positions.clone()  # [N, 3]
         state.prev_forces = state.forces.clone()  # [N, 3]

torch_sim/optimizers/cell_filters.py

@@ -17,7 +17,7 @@
 import torch_sim.math as tsm
 from torch_sim.models.interface import ModelInterface
 from torch_sim.optimizers.state import BFGSState, FireState, LBFGSState, OptimState
-from torch_sim.state import SimState
+from torch_sim.state import DeformGradMixin, SimState
 
 
 MAX_LOG_DEFORM = 2.0
@@ -299,7 +299,7 @@ def unit_cell_step[T: AnyCellState](state: T, cell_lr: float | torch.Tensor) ->
         cell_lr = cell_lr.expand(state.n_systems)
 
     # Get current deformation gradient
-    cur_deform_grad = deform_grad(state.reference_cell.mT, state.row_vector_cell)
+    cur_deform_grad = state.deform_grad()
 
     # Calculate cell positions from current deformation gradient
     cell_factor_expanded = state.cell_factor.expand(state.n_systems, 3, 1)
@@ -371,7 +371,7 @@ def compute_cell_forces[T: AnyCellState](
 
     if is_frechet:
         # Frechet cell force computation
-        cur_deform_grad = deform_grad(state.reference_cell.mT, state.row_vector_cell)
+        cur_deform_grad = state.deform_grad()
         ucf_cell_grad = torch.bmm(
             virial, torch.linalg.inv(torch.transpose(cur_deform_grad, 1, 2))
         )
@@ -392,7 +392,7 @@ def compute_cell_forces[T: AnyCellState](
     else:  # Unit cell force computation
         # Note (AG): ASE transforms virial as:
         # virial = np.linalg.solve(cur_deform_grad, virial.T).T
-        cur_deform_grad = deform_grad(state.reference_cell.mT, state.row_vector_cell)
+        cur_deform_grad = state.deform_grad()
         virial_transformed = torch.linalg.solve(
             cur_deform_grad, virial.transpose(-2, -1)
         ).transpose(-2, -1)
@@ -424,10 +424,9 @@ def get_cell_filter(cell_filter: "CellFilter | tuple") -> CellFilterFuncs:
 
 
 @dataclass(kw_only=True)
-class CellOptimState(OptimState):
+class CellOptimState(OptimState, DeformGradMixin):
     """State class for cell optimization."""
 
-    reference_cell: torch.Tensor
     cell_filter: CellFilterFuncs
     cell_factor: torch.Tensor = field(default_factory=lambda: None)
     pressure: torch.Tensor = field(default_factory=lambda: None)
@@ -453,6 +452,45 @@ class CellOptimState(OptimState):
         "frechet_method",
     }
 
+    def deform_grad_forces(self) -> torch.Tensor:
+        """Atomic forces in deformation gradient space, ``forces @ deform_grad``.
+
+        Mirrors the transform ASE's ``get_forces_unitcellfilter`` and
+        ``get_forces_frechet`` apply to the atomic forces. Equals ``forces`` when
+        the cell is undeformed relative to the reference cell.
+
+        Returns:
+            The transformed atomic forces, shape (n_atoms, 3)
+        """
+        # per-atom row vector @ its system's deform_grad:
+        # (n_atoms, 1, 3) @ (n_atoms, 3, 3) -> (n_atoms, 1, 3) -> (n_atoms, 3)
+        return torch.bmm(
+            self.forces.unsqueeze(1), self.deform_grad()[self.system_idx]
+        ).squeeze(1)
+
+    def frac_positions(self) -> torch.Tensor:
+        """Atomic positions in the reference cell frame, ``solve(deform_grad, r)``.
+
+        Returns:
+            The reference-frame positions, shape (n_atoms, 3)
+        """
+        return torch.linalg.solve(
+            self.deform_grad()[self.system_idx], self.positions.unsqueeze(-1)
+        ).squeeze(-1)
+
+    def positions_from_frac(self, frac_positions: torch.Tensor) -> torch.Tensor:
+        """Cartesian positions from reference-frame positions, ``frac @ deform_grad.mT``.
+
+        Args:
+            frac_positions: Reference-frame positions, shape (n_atoms, 3)
+
+        Returns:
+            The Cartesian positions, shape (n_atoms, 3)
+        """
+        return torch.bmm(
+            frac_positions.unsqueeze(1), self.deform_grad()[self.system_idx].mT
+        ).squeeze(1)
+
 
 @dataclass(kw_only=True)
 class CellFireState(CellOptimState, FireState):

torch_sim/optimizers/fire.py

@@ -318,17 +318,13 @@ def _ase_fire_step[T: "FireState | CellFireState"](  # noqa: C901, PLR0915
             (n_systems,), alpha_start.item(), device=device, dtype=dtype
         )
 
-        # Transform forces for cell optimization
-        if isinstance(state, CellFireState):
-            cur_deform_grad = cell_filters.deform_grad(
-                state.row_vector_cell,
-                getattr(state, "reference_row_vector_cell", state.row_vector_cell),
-            )
-            forces = torch.bmm(
-                state.forces.unsqueeze(1), cur_deform_grad[state.system_idx]
-            ).squeeze(1)
-        else:
-            forces = state.forces
+        # 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.deform_grad_forces()
+            if isinstance(state, CellFireState)
+            else state.forces
+        )
 
         # Calculate power (newly zeroed systems will have power=0 → neg_mask)
         system_power = tsm.batched_vdot(forces, state.velocities, state.system_idx)
@@ -409,15 +405,8 @@ def _ase_fire_step[T: "FireState | CellFireState"](  # noqa: C901, PLR0915
     if isinstance(state, CellFireState):
         # For cell optimization, handle both atomic and cell position updates
         # This follows the ASE FIRE implementation pattern
-        # Transform atomic positions to fractional coordinates
-        cur_deform_grad = cell_filters.deform_grad(
-            state.reference_cell.mT, state.row_vector_cell
-        )
-        frac_positions = torch.linalg.solve(
-            cur_deform_grad[state.system_idx], state.positions.unsqueeze(-1)
-        ).squeeze(-1)
         # Store fractional positions (will transform to Cartesian after cell update)
-        new_frac_positions = frac_positions + dr_atom
+        new_frac_positions = state.frac_positions() + dr_atom
 
         # Update cell positions directly based on stored cell filter type
         if hasattr(state, "cell_filter") and state.cell_filter is not None:
@@ -454,9 +443,7 @@ def _ase_fire_step[T: "FireState | CellFireState"](  # noqa: C901, PLR0915
             # pre-adjustment value.  Without this, any constraint that
             # modifies the cell (e.g. FixSymmetry) causes a zigzag where the
             # optimizer repeatedly proposes from a stale cell_positions.
-            adjusted_deform_grad = cell_filters.deform_grad(
-                state.reference_cell.mT, state.row_vector_cell
-            )
+            adjusted_deform_grad = state.deform_grad()
             if is_frechet:
                 cell_factor_reshaped = state.cell_factor.view(state.n_systems, 1, 1)
                 state.cell_positions = (
@@ -471,16 +458,7 @@ def _ase_fire_step[T: "FireState | CellFireState"](  # noqa: C901, PLR0915
                 )
 
         # Transform fractional positions to Cartesian using NEW deformation gradient
-        new_deform_grad = cell_filters.deform_grad(
-            state.reference_cell.mT, state.row_vector_cell
-        )
-
-        state.set_constrained_positions(
-            torch.bmm(
-                new_frac_positions.unsqueeze(1),
-                new_deform_grad[state.system_idx].transpose(-2, -1),
-            ).squeeze(1)
-        )
+        state.set_constrained_positions(state.positions_from_frac(new_frac_positions))
     else:
         state.set_constrained_positions(state.positions + dr_atom)
 

torch_sim/optimizers/lbfgs.py

@@ -334,25 +334,8 @@ def lbfgs_step(  # noqa: PLR0915, C901
         ext_mask = atom_mask  # [S, M]
 
     if isinstance(state, CellLBFGSState):
-        # Get current deformation gradient
-        # reference_cell.mT: [S, 3, 3], row_vector_cell: [S, 3, 3]
-        cur_deform_grad = deform_grad(
-            state.reference_cell.mT, state.row_vector_cell
-        )  # [S, 3, 3]
-
-        # Transform forces to scaled coordinates
-        # forces: [N, 3], cur_deform_grad[system_idx]: [N, 3, 3] -> [N, 3]
-        forces_scaled = torch.bmm(
-            state.forces.unsqueeze(1),  # [N, 1, 3]
-            cur_deform_grad[state.system_idx],  # [N, 3, 3]
-        ).squeeze(1)  # [N, 3]
-
-        # Current fractional positions
-        # positions: [N, 3] -> frac_positions: [N, 3]
-        frac_positions = torch.linalg.solve(
-            cur_deform_grad[state.system_idx],  # [N, 3, 3]
-            state.positions.unsqueeze(-1),  # [N, 3, 1]
-        ).squeeze(-1)  # [N, 3]
+        forces_scaled = state.deform_grad_forces()  # [N, 3]
+        frac_positions = state.frac_positions()  # [N, 3]
 
         # Convert to padded per-system format: [S, M, 3]
         g_atoms = _atoms_to_padded(-forces_scaled, state.system_idx, n_systems, max_atoms)
@@ -505,9 +488,7 @@ def lbfgs_step(  # noqa: PLR0915, C901
         # pre-adjustment value.  Without this, any constraint that
         # modifies the cell (e.g. FixSymmetry) causes a zigzag where the
         # optimizer repeatedly proposes from a stale cell_positions.
-        adjusted_deform_grad = deform_grad(
-            state.reference_cell.mT, state.row_vector_cell
-        )  # [S, 3, 3]
+        adjusted_deform_grad = state.deform_grad()  # [S, 3, 3]
         if is_frechet:
             cell_factor_reshaped = state.cell_factor.view(n_systems, 1, 1)
             state.cell_positions = (
@@ -524,25 +505,12 @@ def lbfgs_step(  # noqa: PLR0915, C901
         # they are consistent with the next step's start-of-step deformation
         # gradient.  Without this, the LBFGS history vectors (s, y) mix two
         # different coordinate frames, corrupting the Hessian estimate.
-        state.prev_positions = torch.linalg.solve(
-            adjusted_deform_grad[state.system_idx],
-            state.positions.unsqueeze(-1),
-        ).squeeze(-1)  # [N, 3] (fractional in adjusted frame)
-        state.prev_forces = torch.bmm(
-            state.forces.unsqueeze(1),
-            adjusted_deform_grad[state.system_idx],
-        ).squeeze(1)  # [N, 3] (scaled in adjusted frame)
+        state.prev_positions = state.frac_positions()  # [N, 3]
+        state.prev_forces = state.deform_grad_forces()  # [N, 3]
 
         # Apply position step in fractional space, then convert to Cartesian
         new_frac = frac_positions + step_positions  # [N, 3]
-
-        new_deform_grad = adjusted_deform_grad  # already computed above
-        # new_positions = new_frac @ deform_grad^T
-        new_positions = torch.bmm(
-            new_frac.unsqueeze(1),  # [N, 1, 3]
-            new_deform_grad[state.system_idx].transpose(-2, -1),  # [N, 3, 3]
-        ).squeeze(1)  # [N, 3]
-        state.set_constrained_positions(new_positions)  # [N, 3]
+        state.set_constrained_positions(state.positions_from_frac(new_frac))  # [N, 3]
     else:
         state.prev_positions = state.positions.clone()  # [N, 3]
         state.prev_forces = state.forces.clone()  # [N, 3]
@@ -568,19 +536,8 @@ def lbfgs_step(  # noqa: PLR0915, C901
     # s = position difference, y = gradient difference
     if isinstance(state, CellLBFGSState):
         # Get new scaled forces and fractional positions for history
-        new_deform_grad = deform_grad(
-            state.reference_cell.mT, state.row_vector_cell
-        )  # [S, 3, 3]
-        # new_forces: [N, 3] -> new_forces_scaled: [N, 3]
-        new_forces_scaled = torch.bmm(
-            new_forces.unsqueeze(1),  # [N, 1, 3]
-            new_deform_grad[state.system_idx],  # [N, 3, 3]
-        ).squeeze(1)  # [N, 3]
-        # positions: [N, 3] -> new_frac_positions: [N, 3]
-        new_frac_positions = torch.linalg.solve(
-            new_deform_grad[state.system_idx],  # [N, 3, 3]
-            state.positions.unsqueeze(-1),  # [N, 3, 1]
-        ).squeeze(-1)  # [N, 3]
+        new_forces_scaled = state.deform_grad_forces()  # [N, 3]
+        new_frac_positions = state.frac_positions()  # [N, 3]
 
         # s_new_pos = frac_pos_new - frac_pos_old: [N, 3] -> [S, M, 3]
         s_new_pos_atoms = new_frac_positions - state.prev_positions  # [N, 3]