Add a general torch CompositionModel

[PicoCentauri]

Closes #269

Creates an API compatible torch.nn.Module for fitting and predicting compositions weights based on the calculate_composition_weights() function.

Currently, this is the bare implementation and before we include this in the other architectures, I would like to have some feedback.

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)?
• Use in alchemical_model
• Use in GAP
• Use in SOAP-BPNN

---

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

[Luthaf]

I worked around one issue in the docs (metatensor/metatensor#728), but there seems to be another one now

      File "/home/runner/work/metatrain/metatrain/examples/programmatic/llpr/llpr.py", line 42, in <module>
        model = load_atomistic_model("model.pt", extensions_directory="extensions/")
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/home/runner/work/metatrain/metatrain/.tox/docs/lib/python3.12/site-packages/metatensor/torch/atomistic/model.py", line 46, in load_atomistic_model
        return torch.jit.load(path)
               ^^^^^^^^^^^^^^^^^^^^
      File "/home/runner/work/metatrain/metatrain/.tox/docs/lib/python3.12/site-packages/torch/jit/_serialization.py", line 170, in load
        return wrap_cpp_module(cpp_module)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/home/runner/work/metatrain/metatrain/.tox/docs/lib/python3.12/site-packages/torch/jit/_recursive.py", line 1027, in wrap_cpp_module
        return torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    AttributeError: type object 'RecursiveScriptModule' has no attribute '_construct'

I wonder if that's related to the release of torch v2.4.1.

[Luthaf]

Hmm, actually, this can be caused by PYTORCH_JIT=O as well. I'll find another wokraround

[PicoCentauri]

Very cool. Thanks for trying to finish this.

Once the CI is fixed I can give some more comments.

[PicoCentauri]

I think we can also move this inside the CompositionModel itself?

there is one solution in the LJ example of Guillaume.

I recreated the system in a code.

            # select only real atoms and discard ghosts
            if selected_atoms is not None:
                current_system_mask = selected_atoms.column("system") == system_i
                current_atoms = selected_atoms.column("atom")
                current_atoms = current_atoms[current_system_mask].to(torch.long)

                system_final = System(
                    types=system.types[current_atoms],
                    positions=system.positions[current_atoms],
                    cell=system.cell,
                )

[frostedoyster]

Thanks, I did it in a slightly different way to avoid label conflicts but it's now there

[PicoCentauri]

Yes nice implementation.

[PicoCentauri]

I think this is not needed anymore?

[frostedoyster]

Unfortunately this needs to be kept for the alchemical model (notice that it changed directories), which works in a way that doesn't allow me to change things without changing the alchemical model code itself

[PicoCentauri]

Okay!

[PicoCentauri]

NICE!