Fix/ase fire deform grad forces

[curtischong]

Summary

The problem lies in this code https://github.com/curtischong/torch-sim/blob/98c3ca52d5e6e45030efbbc3442ba99f91e2a335/torch_sim/optimizers/fire.py#L323-L328:

cur_deform_grad = cell_filters.deform_grad(
  state.row_vector_cell,
  getattr(state, "reference_row_vector_cell", state.row_vector_cell), # L325
)

Since the fire state (which is a child of CellOptimState) never has reference_row_vector_cell, the state.row_vector_cell is always used instead.

Here is the definition of CellOptimState, we can see it doesn't have reference_row_vector_cell.

https://github.com/curtischong/torch-sim/blob/98c3ca52d5e6e45030efbbc3442ba99f91e2a335/torch_sim/optimizers/cell_filters.py#L427

Adding the DeformGradMixin to CellOptimState gives it the reference_row_vector_cell attribute.

I found this bug bc I saw that torchsim didn't match the ASE implementation. I ran the test script below to show that our optimizer diverges from ASE's because of this bug.

This was the script used to replicate the bug on main

"""Does torch-sim's ase_fire + frechet cell filter track ASE's FIRE + FrechetCellFilter?

ASE's cell filters hand FIRE ``forces @ deform_grad`` for the atomic rows, where
``deform_grad`` maps the reference cell onto the current cell. Both runs start
from the same 5%-strained, rattled Ar supercell with matched Lennard-Jones
parameters and FIRE settings, so the trajectories should be identical.
"""

import numpy as np
import torch
from ase.build import bulk
from ase.calculators.lj import LennardJones
from ase.filters import FrechetCellFilter
from ase.optimize import FIRE

import torch_sim as ts
from torch_sim.models.lennard_jones import LennardJonesModel

lj = dict(sigma=3.405, epsilon=0.0104, cutoff=2.5 * 3.405)
dt, max_step = 1.0, 10.0  # large so the cell moves appreciably each step
pos_tol, cell_tol = 1e-8, 1e-6  # �; fp noise is ~1e-11 / ~1e-7 respectively

atoms = bulk("Ar", "fcc", a=5.26 * 1.05, cubic=True).repeat((2, 2, 2))
atoms.rattle(stdev=0.1, seed=0)
atoms.calc = LennardJones(sigma=lj["sigma"], epsilon=lj["epsilon"], rc=lj["cutoff"])
ase_fire = FIRE(FrechetCellFilter(atoms), logfile=None, dt=dt, maxstep=max_step)

model = LennardJonesModel(**lj, dtype=torch.float64, compute_stress=True)
state = ts.fire_init(
    ts.io.atoms_to_state(atoms, device="cpu", dtype=torch.float64),
    model,
    cell_filter=ts.CellFilter.frechet,
    dt_start=dt,
)

print(f"torch-sim {ts.__version__} ase_fire + CellFilter.frechet  vs  ASE FIRE + FrechetCellFilter")
print(f"{len(atoms)}-atom Ar supercell, 5% strained + rattled, LJ, {dt=}, {max_step=}\n"
print(f"{'step':>4} | {'max |pos_ts - pos_ase| (�)':>26} | {'max |cell_ts - cell_ase| (�)':>28} | status")
print("-" * 80)
all_match = True
for step in range(1, 6):
    ase_fire.step()
    state = ts.fire_step(state, model, max_step=max_step)
    pos_err = np.abs(state.positions.numpy() - atoms.positions).max()
    cell_err = np.abs(state.cell[0].numpy() - atoms.cell.array).max()
    match = pos_err < pos_tol and cell_err < cell_tol
    all_match &= match
    print(f"{step:>4} | {pos_err:>26.1e} | {cell_err:>28.1e} | {'match' if match else 'DIVERGED'}")

print(
    "\nStep 1 matches regardless: the cell still equals the reference cell, so"
    " forces @ deform_grad == forces. From step 2 on the cell has moved and the transform matters."
)
print(
    "\nRESULT: torch-sim ase_fire "
    + ("TRACKS" if all_match else "DOES NOT TRACK")
    + " ASE FIRE + FrechetCellFilter "
    + f"(tolerance: pos < {pos_tol:.0e} �, cell < {cell_tol:.0e} �)"
)

Results:

main:

step | max |pos_ts - pos_ase| (�) | max |cell_ts - cell_ase| (�) | status
--------------------------------------------------------------------------------
   1 |                    1.4e-12 |                      1.8e-12 | match
   2 |                    4.8e-05 |                      2.1e-06 | DIVERGED
   3 |                    1.6e-04 |                      8.5e-06 | DIVERGED
   4 |                    3.9e-04 |                      2.8e-05 | DIVERGED
   5 |                    8.4e-04 |                      1.1e-04 | DIVERGED

this branch:

step | max |pos_ts - pos_ase| (�) | max |cell_ts - cell_ase| (�) | status
--------------------------------------------------------------------------------
   1 |                    1.4e-12 |                      1.8e-12 | match
   2 |                    2.7e-12 |                      1.0e-08 | match
   3 |                    3.6e-12 |                      6.0e-08 | match
   4 |                    2.0e-12 |                      1.5e-07 | match
   5 |                    4.6e-11 |                      1.9e-07 | match

[curtischong]

adding this mixin is the key fix of this PR, it gives the CellOptimState the reference_row_vector_cell attribute.

[curtischong]

this is such a subtle bug that ONLY affected FIRE, and NOT the BFGS or L-BFGS optimizers since the BFGS optimizers read the reference_cell directly. whereas if you look at fire, it looked for getattr(state, "reference_row_vector_cell",

[curtischong]

I could've fixed this bug by making the fire implementation match the other 2, but it's cleaner to just add the DeformGradMixin to the CellOptimState and use shared helper functions

[curtischong]

I considered adding this function to DeformGradMixin but decided against it since it needs forces and system_idx which are missing from DeformGradMixin but CellOptimState provides

[curtischong]

ase's version of this function is here:
https://gitlab.com/ase/ase/-/blob/3.29.0/ase/_4/optimize/cellutil.py#L213

[curtischong]

fire, bfgs, and l-bfgs all individually calculate frac_positions which we will use the state.frac_positions() helper function now. I think this bug arose since we had different implementations to calculate forces_scaled, so factoring out all this logic into the same helpers is defensive programming

[curtischong]

This PR comes after curtischong#18 (it's rebased off this one for clarity), but I'd probably review that one too if you're in a reviewing spree this weekend rhys