Does TorchSim support different spins and multiplicities?

[niklashoelter]

It appears that io.atoms_to_state does not pass any of the atoms object properties to the SimState object. In FairChem, these properties are provided via atoms.info, which is a dictionary containing values such as "charge" and "spin".

Could you clarify if the batch mode also supports different spins and multiplicities? If not, would it be possible to add support for passing these properties (e.g., via atoms.info or a similar mechanism) to the SimState object?

This would be helpful for workflows that require specifying different quantum numbers for each atom or molecule in the batch.

Thank you!

[CompRhys]

Related issue: #253

Short answer is at this point no but it should be possible to add via _global_attributes in SimState and correct plumbing on the ModelInterface.

[niklashoelter]

Sorry for the duplicate issue - is there any implementation of this planned in the near future? Or are you looking for contributions?

[CompRhys]

Contributions always super welcome! Maybe @orionarcher could prioritise this given a few requests now if you don't feel like trying to tackle it sooner?

[orionarcher]

Yeah I can bump this up in my priorities, clearly an in-demand feature! You are also welcome to take a crack at it @niklashoelter! Best place to start would be a proposal for how to tackle it.

[Nekkrad]

Just an heads up, charge and spin are also required by force fields trained on OMOL dataset such as OrbMol and MACE-OMOL

[thomasloux]

Should be added in _system_attributes and I would suggest to keep SimState and not create a new dataclass Molecular. Then either the value used is None, but at the moment this choice as been excluded. A default value could be used, it's the choice in fairchem model:
https://github.com/facebookresearch/fairchem/blob/main/src/fairchem/core/calculate/ase_calculator.py#L270

but does not seem to be a default value for OrbMol:
https://github.com/orbital-materials/orb-models/blob/main/orb_models/forcefield/nn_util.py#L137
https://github.com/orbital-materials/orb-models/blob/main/orb_models/forcefield/calculator.py#L100

[thomasloux]

Main reason against creating a new dataclass is simply that it would be a pain for subclass, at least given how it's implemented at the moment. But if there is a nice way to implement that, it would probably be better than None/default value

[curtischong]

I think this was why I initially wanted to make system_attributes a dictionary of tensors in #228. (rather than just a set of attribute names) Because I saw orb do this and it's a general interface for people to add new attributes. I think we need to add 3 dictionaries to SimState, for atom_attributes, system_attributes, and global_attributes to support models that output additional properties (mag moments, even hessians)

I've been thinking about this for a few weeks now and I needed to say this and get it off my chest.

[niklashoelter]

Are there any news on this issue/enhancement (and #253) already or any agreement on which proposed solution might be best?

[orionarcher]

The answer is now yes! I added support for charge and spin in #373!

If more model properties continue being added we may need a general solution, but for now I just added them directly to SimState.

[niklashoelter]

Thanks @orionarcher! While taking a closer a look into it, I noticed that passing of the spin and charge parameters seems to work for a single model iteration, but the informatin gets lost in optimization runs - most likely in the init_fn (maybe also the step_fn?) of the optimizer (e.g. I was using fire, and fire_init does not copy the new SimState attributes).

The second thing I was wondering - it seems like Fairchem is using "multiplicitiy" and "spin" synonymously (ref https://fair-chem.github.io/uma_tutorials/uma_tutorial.html#illustrative-examples) - whereas you used the correct conversion in the Fairchem tests. Should not matter for functionality, but I just wanted to raise this issue for awareness.

[orionarcher]

but the informatin gets lost in optimization runs - most likely in the init_fn (maybe also the step_fn?) of the optimizer (e.g. I was using fire, and fire_init does not copy the new SimState attributes).

Ahh hmm, that sounds like a bug to me. Would you mind raising an issue?

The second thing I was wondering - it seems like Fairchem is using "multiplicitiy" and "spin" synonymously

That's a good callout, are there any changes you'd suggest to the language on the TorchSim side? We should clarify as appropriate.

[niklashoelter]

Issue raised.

For consistency, I'd propose to keep the "spin" label also in TorchSim, but definitely highlight this in the documentation and docstrings and adapt the tests accordingly.