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/tests/test_optimizers_vs_ase.py b/tests/test_optimizers_vs_ase.py index 19681d15..8e6f67b7 100644 --- a/tests/test_optimizers_vs_ase.py +++ b/tests/test_optimizers_vs_ase.py @@ -282,6 +282,16 @@ def _run_and_compare_optimizers( ), ] +STRICT_FIRE_CASES = [ + ( + "rattled_ar_supercell_stretched_sim_state", + 5, + 1e-10, + 1e-11, + "Ar stretched-rattled supercell (Frechet)", + ), +] + BFGS_CASES = [ ( "rattled_sio2_sim_state", @@ -441,6 +451,71 @@ def test_optimizer_vs_ase_parametrized( ) +@pytest.mark.parametrize( + ( + "sim_state_fixture_name", + "steps", + "position_atol", + "cell_atol", + "test_id", + ), + STRICT_FIRE_CASES, +) +def test_fire_frechet_matches_ase_strict( + sim_state_fixture_name: str, + steps: int, + position_atol: float, + cell_atol: float, + test_id: str, + ts_mace_mpa: MaceModel, + ase_mace_mpa: "MACECalculator", + request: pytest.FixtureRequest, +) -> None: + """Compare the first FIRE steps with tight position and cell tolerances.""" + initial_sim_state = request.getfixturevalue(sim_state_fixture_name) + state = ts.fire_init( + initial_sim_state.clone(), + ts_mace_mpa, + cell_filter=ts.CellFilter.frechet, + ) + + ase_atoms = ts.io.state_to_atoms( + initial_sim_state.clone().to(dtype=DTYPE, device=ts_mace_mpa.device) + )[0] + ase_atoms.calc = ase_mace_mpa + filtered_ase_atoms = FrechetCellFilter(ase_atoms) + ase_optimizer = FIRE(filtered_ase_atoms, logfile=None) + + for step in range(1, steps + 1): + state = ts.fire_step(state, ts_mace_mpa) + ase_optimizer.run(fmax=0.02, steps=1) + + ase_positions = torch.as_tensor( + filtered_ase_atoms.atoms.get_positions(), + device=state.device, + dtype=state.dtype, + ) + ase_cell = torch.as_tensor( + filtered_ase_atoms.atoms.get_cell().array, + device=state.device, + dtype=state.dtype, + ) + torch.testing.assert_close( + state.positions, + ase_positions, + rtol=0.0, + atol=position_atol, + msg=f"{test_id}: position trajectories diverged at FIRE step {step}", + ) + torch.testing.assert_close( + state.cell[0].mT, + ase_cell, + rtol=0.0, + atol=cell_atol, + msg=f"{test_id}: cell trajectories diverged at FIRE step {step}", + ) + + # TODO (AG): Can we merge these tests with the FIRE tests? diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py index 035b6911..86a3e4c5 100644 --- a/torch_sim/autobatching.py +++ b/torch_sim/autobatching.py @@ -776,9 +776,6 @@ class InFlightAutoBatcher[T: SimState]: To avoid a slow memory estimation step, set the `max_memory_scaler` to a known value. - .. image:: https://github.com/janosh/diagrams/raw/main/assets/in-flight-auto-batcher/in-flight-auto-batcher.svg - :alt: In-flight auto-batcher diagram - Attributes: model (ModelInterface): Model used for memory estimation and processing. memory_scales_with (str): Metric type used for memory estimation. 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] diff --git a/torch_sim/testing.py b/torch_sim/testing.py index f73534ad..ad312a48 100644 --- a/torch_sim/testing.py +++ b/torch_sim/testing.py @@ -253,6 +253,17 @@ def make_ar_supercell_sim_state( return ts.io.atoms_to_state(atoms, device, dtype) +def make_rattled_ar_supercell_stretched_sim_state( + device: torch.device | None = None, dtype: torch.dtype | None = None +) -> ts.SimState: + """Create a rattled FCC Argon supercell with a 5% stretched cell.""" + from ase.build import bulk + + atoms = bulk("Ar", "fcc", a=5.26 * 1.05, cubic=True).repeat([2, 2, 2]) + atoms.rattle(stdev=0.1, seed=0) + return ts.io.atoms_to_state(atoms, device, dtype) + + def make_fe_supercell_sim_state( device: torch.device | None = None, dtype: torch.dtype | None = None ) -> ts.SimState: @@ -482,6 +493,9 @@ def make_distorted_fcc_al_conventional_sim_state( "sio2_sim_state": make_sio2_sim_state, "rattled_sio2_sim_state": make_rattled_sio2_sim_state, "ar_supercell_sim_state": make_ar_supercell_sim_state, + "rattled_ar_supercell_stretched_sim_state": ( + make_rattled_ar_supercell_stretched_sim_state + ), "fe_supercell_sim_state": make_fe_supercell_sim_state, "casio3_sim_state": make_casio3_sim_state, "osn2_sim_state": make_osn2_sim_state,