Implement the Alchemical Model from torch_alchemical repo

[abmazitov]

Current status

• Model is implemented
• Tests are implemented
• Test are working
• Fixed the Failed to downcast a Function to a GraphFunction issue
• Docs prototype is ready
• WIP example is ready
• Found out the way to pre-compute the neighborlists before the training loop
• Found a place to compute the basis_normalization_factor (ideally make it as a method of the model)
• Fixed the example
• Checked if training works on CUDA and MPS: Doesn't work on MPS because of CPU and CUDA only implementation of spheric art
• Figured out if float32 is an option for splines in torch_spex, waiting for update of splining routine
• Updated to main, including new Dataset and Dataloader

Status of tests (local run)

• Functionality
• Invariance
• Regression
• TorchScript

TODO

• Implement the domain decomposition integration - TBD in the next PR

---

📚 Documentation preview 📚: https://metatensor-models--66.org.readthedocs.build/en/66/

[frostedoyster]

Should this model (and PET) go to an experimental folder for now?

@PicoCentauri @abmazitov @serfg @Luthaf

[Luthaf]

I'd put all models including SOAP-BPNN in an experimental folder for now!

[PicoCentauri]

I agree. Let's discuss this next week on the meeting on how we implement this.

[PicoCentauri]

Please don't forget to add the model to the docs. One entry at the README plus a page at docs/src/architectures.

[PicoCentauri]

This should not be here but will be passed directly from the train function.

[abmazitov]

These lines are somehow needed for training step in test_regression.py

[PicoCentauri]

Yes these lines but the device is a global parameter and will be provided via the train function.

[frostedoyster]

Amazing work, just some small details.
Regarding the large file with structures, I think it's too large (@PicoCentauri). It doesn't train on my laptop (runs out of memory) and we already have 100 structures from the original alchemical dataset in the repo. Can we re-use them?

Bonus question: @abmazitov have you tried learning from the stresses?

[frostedoyster]

Suggested change

	with SOAP features and Alchemical Compression of the composition space. This model
	with SOAP features and alchemical compression of the composition space. This model

[frostedoyster]

We don't want this for each model in the docs, right @PicoCentauri? However, maybe we can leave the example here and run it from the test CI of the alchemical model

[Luthaf]

yeah, I would not duplicate this file. If we want this to be a test, let's make it a test and not an example.

[abmazitov]

Fine, then I'll delete this file

[frostedoyster]

We need a small description of the error

[abmazitov]

Thx, this is no longer needed

[frostedoyster]

We need a small description of the error

[abmazitov]

Same

[PicoCentauri]

Very nice work @abmazitov !!! I will take a look after @frostedoyster comments are done.

Regarding the large file with structures, I think it's too large (@PicoCentauri). It doesn't train on my laptop (runs out of memory) and we already have 100 structures from the original alchemical dataset in the repo. Can we re-use them?

Yes try to use the file we already have. If you can not test the model on a notebook it is not good xD

[abmazitov]

What do you think of adding the model_requires_neighborlists flag to ModelCapabilities or to the forward pass keywords, or to the model class attribute? Since the Alchemical Model doesn't work without NLs, there is a few places in the code (i.e. metatensor-models eval step), where the workflow breaks because of the NLs missing.

What I would like to see, is a model check performed any time the forward pass is called. In other words, if the model requires the NLs and systems don't have them, then we compute the missing NLs first, and then do the forward pass

@Luthaf @PicoCentauri @frostedoyster Let me know what do you think

[Luthaf]

What do you think of adding the model_requires_neighborlists flag to ModelCapabilities or to the forward pass keywords, or to the model class attribute? Since the Alchemical Model doesn't work without NLs, there is a few places in the code (i.e. metatensor-models eval step), where the workflow breaks because of the NLs missing.

If your model requires NL, it should define somewhere in the module tree the corresponding requested_neighbors_lists(self) -> List[NeighborsListOptions]. This is done automatically by rascaline calculators.

Then MetatensorAtomisticModels will use these functions to assemble the full list of NL, and the simulation engine (here the eval script) is responsible for computing the NL and storing them in the systems.

[Luthaf]

At the very least for tests. For train/eval, it depends how much slower check_consistency=True makes the code

[frostedoyster]

@Luthaf could you please check this?

[Luthaf]

yeah, I would not duplicate this file. If we want this to be a test, let's make it a test and not an example.

[abmazitov]

Ready for merging

[frostedoyster]

I'm having a look and I will add my comments here before we merge:

• Documentation: it's there, having a small guide on how to change the parameters wouldn't hurt, but it can be left for later. I'm going to add some basic comments next to each hyperparameter. I really dislike the citation format we're using, which doesn't even make it easy for the user to cite the relevant papers. It looks really cool, but it doesn't do what we want in our context. Bibtex needs to be visible and ready to be copied. I'll open an issue
• Model: only some small changes (4 lines) were needed for it to work on different device and dtype. However, the handling of the composition weights is wrong (I assume due to the internal rescaling of targets in the model), and the energy errors are very high as a result

EDIT: the problem was in the order that the rescaling and the composition weights were applied. Changing the order to rescaling first, then applying composition weights fixed the issue. I'm merging

[PicoCentauri]

Very nice work!!! 🥳