diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py index 3f419f68..7e77146f 100644 --- a/tests/test_optimizers.py +++ b/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: diff --git a/torch_sim/optimizers/bfgs.py b/torch_sim/optimizers/bfgs.py index 855e05db..4a4c5595 100644 --- a/torch_sim/optimizers/bfgs.py +++ b/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] diff --git a/torch_sim/optimizers/cell_filters.py b/torch_sim/optimizers/cell_filters.py index 8a58f8d4..25b1b60f 100644 --- a/torch_sim/optimizers/cell_filters.py +++ b/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): diff --git a/torch_sim/optimizers/fire.py b/torch_sim/optimizers/fire.py index d92da8d4..a71f3cd1 100644 --- a/torch_sim/optimizers/fire.py +++ b/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) diff --git a/torch_sim/optimizers/lbfgs.py b/torch_sim/optimizers/lbfgs.py index 56651598..a36f2a04 100644 --- a/torch_sim/optimizers/lbfgs.py +++ b/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]