Adding support for MLIPs with charge and spin embeddings

[steffen-wedig]

Hi everyone,
thank you for developing this nice package.

Are there currently plans to support MLIPs with charge or spin embeddings? I am particularly interested in using the recently released MACE-OMOL 25. As the SimState does not contain "total_spin" and "total_charge", the forward pass of this model fails.

Would be very nice if torch-sim were to offer this capability.

atoms = molecule("H2O")
mace_raw = mace_omol(
    "MACE-omol-0-extra-large-1024.model",
    device="cuda",
    return_raw_model=True,
)
state = ts.io.atoms_to_state(atoms, device="cuda", dtype=torch.float64)
mace_model = MaceModel(model=mace_raw)
mace_model(state)

Traceback (most recent call last):
File "/home/steffen/projects/mlips/es_torch_sim.py", line 69, in
print(mace_model(state))
^^^^^^^^^^^^^^^^^
File "/home/steffen/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/steffen/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/steffen/anaconda3/lib/python3.11/site-packages/torch_sim/models/mace.py", line 333, in forward
out = self.model(
^^^^^^^^^^^
File "/home/steffen/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/steffen/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/steffen/anaconda3/lib/python3.11/site-packages/mace/modules/models.py", line 510, in forward
embedding_features[name] = data[name]
~~~~^^^^^^
KeyError: 'total_spin'

[CompRhys]

Thanks for the suggestion. Given OMOL and models trained with it this is something that makes sense to add in the near future. for total_spin and total_charge we already have _global_attributes concept to handle these but need to consider whether they should go in the base class or a special MolecularSimState

[orionarcher]

This is now supported by #373! Apologies it took so long.