Add Spin/Charge System Conditioning to PET

[JonathanSchmidt1]

Add charge/spin system conditioning module that has extra embeddings for charge and spin that are added to node embeddings at each message passing step
Add respective hyperparameters and dataset options

Contributor (creator of pull-request) checklist

• Tests updated (for new features and bugfixes)?
• Documentation updated (for new features)?
• Issue referenced (for PRs that solve an issue)?

Maintainer/Reviewer checklist

• CHANGELOG updated with public API or any other important changes?
• GPU tests passed (maintainer comment: "cscs-ci run")?

---

📚 Documentation preview 📚: https://metatrain--1040.org.readthedocs.build/en/1040/

[JonathanSchmidt1]

Ps: Sorry about the undescriptive name seems like I cant change it or I am too stupid ^^

[Luthaf]

Do we already have .bin files in the dataset?

[JonathanSchmidt1]

I think so

metatrain/examples/0-beginner/01-data_preparation.py

Line 196 in 5715627

a_path = root / "a.bin"

[Luthaf]

ok. Then this would mainly need some documentation in the user doc, saying that we'll check for data at this key and put it in the system under this name

[Luthaf]

cscs-ci run

[pfebrer]

@frostedoyster is this needed? From the last ml-devel meeting I understood that in principle changes to an architecture's checkpoint shouldn't need a version upgrade on the wrapper.

I think what needs to be done is to simply upgrade the checkpoint file with the new file, without any renaming, but let's see what Filippo says.

(same for llpr)

[pfebrer]

Also, I think passing global data in general is something very useful. It would be nice if the user could pass whatever global data and it would be automatically embedded into PET. Could we not hardcode it to spin and charge? We could still allow the user to do e.g. type: spin if spin needs special treatment.

I'm thinking that the input could look something like:

systems:
    read_from: ...
    ...
    features:
        - key: spin
          per_atom: false # Maybe not needed and can be inferred?
          type: spin # So that min-max gets automatically assigned (?)

Or maybe features would be a separate yaml section from systems, idk

[JonathanSchmidt1]

Discussed a bit with @pfebrer . We could generalize it but we have to take into account that the implementation only supports integer embeddings right now and we might also want to take the embedding specs already existing in MACE into account.

[pfebrer]

Yes, regarding this

and we might also want to take the embedding specs already existing in MACE into account.

the thing is that as they are right now, charge and spin are inputs in the base hyperparameters. So if we don't generalize, we would at least need to make sure that these inputs also make sense for other architectures. Otherwise we could make them temporarily inputs of the PET architecture until we find a good solution, but I'm not sure this is even possible (?).

I think the issue is that this situation requires something that had never happened in metatrain: we need to use additional inputs (not just coordinates and atomic numbers). Maybe extra_data can be used for this? 🤔 In any case, probably worth a discussion since this case will establish a precedent.