KeyError 'stress' when using GemNet or EquiformerV2 in ts.optimize

[R1j1t]

Hi Team, I am trying to use fairchem models for optimization, however, I get the following error: KeyError: 'stress'

Below is my code:

from torch_sim.models.fairchem import FairChemModel
from fairchem.core.models.model_registry import model_name_to_local_file

# checkpoint_path = model_name_to_local_file('GemNet-OC-S2EFS-OC20+OC22', local_cache='mlip_checkpoints/')
checkpoint_path = model_name_to_local_file('EquiformerV2-lE4-lF100-S2EFS-OC22', local_cache='mlip_checkpoints/')

# run natively on gpus
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

foundational_potential = FairChemModel(model=checkpoint_path, cpu=False, seed=42, pbc=True, compute_stress=False)

trajectory_files = [f"bulk_traj_{i}.h5md" for i in range(len(bulk_systems_map))]

# relax all of the high temperature states
relaxed_state = ts.optimize(
    system=[i for i in bulk_systems_map.values()],
    model=foundational_potential,
    optimizer=ts.frechet_cell_fire,
    autobatcher=False,
)

print(relaxed_state.energy)

The entire stack trace sits here:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[21], line 4
      1 trajectory_files = [f"bulk_traj_{i}.h5md" for i in range(len(bulk_systems_map))]
      3 # relax all of the high temperature states
----> 4 relaxed_state = ts.optimize(
      5     system=[i for i in bulk_systems_map.values()],
      6     model=foundational_potential,
      7     optimizer=ts.frechet_cell_fire,
      8     autobatcher=False,
      9 )
     11 print(relaxed_state.energy)

File [~/miniconda3/envs/esurf/lib/python3.11/site-packages/torch_sim/runners.py:414](http://localhost:8888/home/ec2-user/miniconda3/envs/esurf/lib/python3.11/site-packages/torch_sim/runners.py#line=413), in optimize(system, model, optimizer, convergence_fn, trajectory_reporter, autobatcher, max_steps, steps_between_swaps, pbar, **optimizer_kwargs)
    407 autobatcher = _configure_in_flight_autobatcher(
    408     model, state, autobatcher, max_attempts
    409 )
    411 if not isinstance(
    412     state, (FireState, UnitCellFireState, UnitCellGDState, FrechetCellFIREState)
    413 ):
--> 414     state = _chunked_apply(
    415         init_fn,
    416         state,
    417         model,
    418         max_memory_scaler=autobatcher.max_memory_scaler,
    419         memory_scales_with=autobatcher.memory_scales_with,
    420     )
    421 autobatcher.load_states(state)
    422 trajectory_reporter = _configure_reporter(
    423     trajectory_reporter,
    424     properties=["potential_energy"],
    425 )

File [~/miniconda3/envs/esurf/lib/python3.11/site-packages/torch_sim/runners.py:281](http://localhost:8888/home/ec2-user/miniconda3/envs/esurf/lib/python3.11/site-packages/torch_sim/runners.py#line=280), in _chunked_apply(fn, states, model, **batcher_kwargs)
    278 autobatcher.load_states(states)
    279 initialized_states = []
--> 281 initialized_states = [fn(batch) for batch in autobatcher]
    283 ordered_states = autobatcher.restore_original_order(initialized_states)
    284 return concatenate_states(ordered_states)

File [~/miniconda3/envs/esurf/lib/python3.11/site-packages/torch_sim/runners.py:281](http://localhost:8888/home/ec2-user/miniconda3/envs/esurf/lib/python3.11/site-packages/torch_sim/runners.py#line=280), in <listcomp>(.0)
    278 autobatcher.load_states(states)
    279 initialized_states = []
--> 281 initialized_states = [fn(batch) for batch in autobatcher]
    283 ordered_states = autobatcher.restore_original_order(initialized_states)
    284 return concatenate_states(ordered_states)

File [~/miniconda3/envs/esurf/lib/python3.11/site-packages/torch_sim/optimizers.py:1109](http://localhost:8888/home/ec2-user/miniconda3/envs/esurf/lib/python3.11/site-packages/torch_sim/optimizers.py#line=1108), in frechet_cell_fire.<locals>.fire_init(state, cell_factor, scalar_pressure, dt_start, alpha_start)
   1107 energy = model_output["energy"]  # [n_batches]
   1108 forces = model_output["forces"]  # [n_total_atoms, 3]
-> 1109 stress = model_output["stress"]  # [n_batches, 3, 3]
   1111 # Calculate initial cell positions using matrix logarithm
   1112 # Calculate current deformation gradient (identity matrix at start)
   1113 cur_deform_grad = DeformGradMixin._deform_grad(  # noqa: SLF001
   1114     state.row_vector_cell, state.row_vector_cell
   1115 )  # shape: (n_batches, 3, 3)

KeyError: 'stress'

Based on the above error I tried using

foundational_potential = FairChemModel(model=checkpoint_path, cpu=False, seed=42, pbc=True, compute_stress=False)

NotImplementedError                       Traceback (most recent call last)
Cell In[22], line 12
      7 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
      9 # mace = mace_mp(model="small", return_raw_model=True, default_dtype=torch.float64)
     10 # mace_model = MaceModel(model=mace, device=device, dtype = torch.float64)
---> 12 foundational_potential = FairChemModel(model=checkpoint_path, cpu=False, seed=42, pbc=True, compute_stress=True)

File [~/miniconda3/envs/esurf/lib/python3.11/site-packages/torch_sim/models/fairchem.py:301](http://localhost:8888/home/ec2-user/miniconda3/envs/esurf/lib/python3.11/site-packages/torch_sim/models/fairchem.py#line=300), in FairChemModel.__init__(self, model, neighbor_list_fn, config_yml, model_name, local_cache, trainer, cpu, seed, dtype, compute_stress, pbc, disable_amp)
    299 stress_output = "stress" in self.implemented_properties
    300 if not stress_output and compute_stress:
--> 301     raise NotImplementedError("Stress output not implemented for this model")

NotImplementedError: Stress output not implemented for this model

Above error aligns with gemnet models where model outputs are:

 {'energy': {'level': 'system'},
 'forces': {'level': 'atom',
  'train_on_free_atoms': True,
  'eval_on_free_atoms': True}}

Does this mean, they are not compatible with torch-sim fire opt? I was able to run it with gradient_descent though?

Package Details:

torch                     2.4.1
torch-geometric           2.6.1
torch_scatter             2.1.2+pt24cu124
torch-sim-atomistic       0.2.2
fairchem-core             1.10.0

[curtischong]

Thank you for pointing this out! I've forwarded this to someone on our team who have had a similar issue. Hopefully when they're less busy we'll get back to you on this!