Fix L-BFGS history consistency for state updates

[falletta]

set_constrained_forces and the energy/stress state updates were being applied after computing the L-BFGS history vectors. As a result, the (s, y) pairs were built using stale state. To fix this, move all state updates (forces, energy, stress) before the history computation, so y_new is calculated from the finalized forces and remains consistent with prev_forces from the previous step.

[abhijeetgangan]

This looks correct. Could you try with history of 1 and see if that causes any issue?

[falletta]

This example with max_history = 1 runs fine:

"""Optimize a rattled Si structure with LBFGS and history 1."""

import numpy as np
import torch
from ase.build import bulk
from mace.calculators.foundations_models import mace_mp

import torch_sim as ts
from torch_sim.models.mace import MaceModel, MaceUrls

# Create rattled diamond-cubic Si structure
rng = np.random.default_rng(seed=42)
si_atoms = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2))
si_atoms.positions += 0.1 * rng.standard_normal(si_atoms.positions.shape)

# Load MACE model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float32
loaded_model = mace_mp(
    model=MaceUrls.mace_mpa_medium,
    return_raw_model=True,
    default_dtype=str(dtype).removeprefix("torch."),
    device=str(device),
)
model = MaceModel(
    model=loaded_model,
    device=device,
    compute_forces=True,
    compute_stress=True,
    dtype=dtype,
)

# Optimize with LBFGS
final_state = ts.optimize(
    system=si_atoms,
    model=model,
    optimizer=ts.Optimizer.lbfgs,
    max_steps=200,
    max_history=1,
)

print(f"Final energy: {final_state.energy.item():.4f} eV")