Add experimental.graph2mat architecture

[pfebrer]

From the creators of experimental.mace...

This one is quite experimental and I'm not sure it will ever aim at being stable, let's see.

What it is useful for

This architecture predicts sparse matrices in spherical basis from the output of any model.

Implementation

The idea of the architecture is simple enough: it takes any architecture in metatrain, asks for a spherical per-atom output, and then applies graph2mat on top of it. This is basically what I had in mind when I developed graph2mat, so I'm very happy that metatrain standardizes everything in a way that implementing this is "trivial" (you know, once you have spent a whole year in COSMO 😆).

The architecture supports multiple matrix targets, and for each target a separate graph is built, since different matrices can have different sparsity patterns. As a consequence of this, the graph used to construct each matrix is not shared by the base (featurizer) model. E.g. the base model can have a higher/lower/adaptive cutoff for neighbors, or even not have a graph at all.

The main point of friction with metatrain is that graph2mat works with a completely flattened array as the batch (due to the sparsity/raggedness of the target), while metatrain in general is more suited for uniform targets (e.g. the supported TargetInfos). For now I solved the problem by using a DiskDataset that makes metatrain happy, but then converting to graph2mat batch in a callable of the collate function.

I tested this with soap_bpnn, PET and MACE and it is working fine.

Future perspective

There is a case for thinking that the architecture is unnecessary, since one can just add graph2mat as a head in the architectures where it makes sense. However, I think it is good to first test things in this experimental architecture because things are probably going to change fast.

There are also architectures for which other approaches are probably much more efficient and using graph2mat for matrices would make things unnecessarily complex. For example, @jwa7 is working on a more native way of doing this in PET.

Still, since graph2mat is very modular and easy to modify, it is nice to have it to quickly test new approaches before moving into modifying the other architectures (non-goal of metatrain, I know haha).

Things missing.

For a proof of concept, I made up a target type (basis) that allows me to play with things. This will be changed to adapt to the target type that Joe is using in his PET implementation, since after all for graph2mat the target type is just a tool to trick metatrain into allowing its running. Therefore, this PR is likely to stay as a draft until Joe finishes his implementation.

Generating inputs

The architecture requires mainly two non-trivial inputs: the disk dataset and the basis specification. Both will be creatable using graph2mat tools, although the disk dataset will be general enough that it could be generated with whatever other tool.

To test it

[not tested on GPU, will test soon!]

The architecture can be tested with this subset of 100 QM9 structures: https://drive.google.com/file/d/1gV4QP4ZwW_BDXdSe0K-UPvu2G3NPg2Nt/view?usp=sharing, which contains the density_matrix, hamiltonian, energy_density_matrix and overlap.

Then run mtt train with the typical options yaml:

architecture:
  name: experimental.graph2mat
  model:
    # Graph2mat model options
    basis_yaml: qm9_basis.yaml
    basis_grouping: basis_shape
    # Featurizer options, this is the same as the full architecture options
    # used to train a model, i.e. you can have featurizer_architecture.model
    featurizer_architecture:
      name: soap_bpnn
  training:
    batch_size: 10
    checkpoint_interval: 20
    optimizer_kwargs:
        lr: 0.005
    loss: mae

training_set:
  systems:
    read_from: qm9_100.zip
    length_unit: angstrom
  targets:
    density_matrix:
      type: basis
    # Uncomment the following to train on more targets
    #hamiltonian:
    #  type: basis
    #overlap:
    #  type: basis
    #energy_density_matrix:
    # type: basis
  # This is needed because to rearrange things in the collate function we
  # need to know the system indices.
  extra_data:
    system_index:
      type: scalar
      per_atom: False

validation_set: 0.2

With the qm9_basis.yaml file containing the basis specification:

- type: 9
  R: [1.5940, 1.1662, 1.8989, 1.8989, 1.8989, 1.1957, 1.1957, 1.1957, 1.8989, 1.8989,
    1.8989, 1.8989, 1.8989]
  basis:
  - [2, 0, 1]
  - [2, 1, -1]
  - [1, 2, 1]
  basis_convention: siesta_spherical
- type: 8
  R: [1.7490, 1.3119, 2.0835, 2.0835, 2.0835, 1.3451, 1.3451, 1.3451, 2.0835, 2.0835,
    2.0835, 2.0835, 2.0835]
  basis:
  - [2, 0, 1]
  - [2, 1, -1]
  - [1, 2, 1]
  basis_convention: siesta_spherical
- type: 7
  R: [1.9495, 1.5182, 2.2650, 2.2650, 2.2650, 1.5373, 1.5373, 1.5373, 2.2650, 2.2650,
    2.2650, 2.2650, 2.2650]
  basis:
  - [2, 0, 1]
  - [2, 1, -1]
  - [1, 2, 1]
  basis_convention: siesta_spherical
- type: 1
  R: [2.4919, 1.9896, 2.4919, 2.4919, 2.4919]
  basis:
  - [2, 0, 1]
  - [1, 1, -1]
  basis_convention: siesta_spherical
- type: 6
  R: [2.1635, 1.7712, 2.5773, 2.5773, 2.5773, 1.8389, 1.8389, 1.8389, 2.5773, 2.5773,
    2.5773, 2.5773, 2.5773]
  basis:
  - [2, 0, 1]
  - [2, 1, -1]
  - [1, 2, 1]
  basis_convention: siesta_spherical

Hope you think this is nice, and looking forward to having this one merged :)

Contributor (creator of pull-request) checklist

• Add your architecture to the experimental or stable folder. See the
  [docs/src/dev-docs/architecture-life-cycle.rst](Architecture life cycle)
  document for requirements. src/metatrain/experimental/<architecture_name>
• Document and provide defaults for the hyperparameters of your model.
• Added tests for your architecture. See https://docs.metatensor.org/metatrain/latest/dev-docs/new-architecture.html#testing-tests
• Added test run to the CI (file .github/workflow/architecture-tests.yml)
• Add a new dependencies entry in the optional-dependencies section in the
  pyproject.toml
• Add maintainers as codeowners in CODEOWNERS
• Trigger a GPU test by asking a maintainer to comment "cscs-ci run".

Reviewer checklist

New experimental architectures

• Capability to fit at least a single quantity and predict it, verified through CI
  tests.
• Compatibility with JIT compilation using TorchScript <https://pytorch.org/docs/stable/jit.html>_.
• Provision of reasonable default hyperparameters.
• A contact person designated as the maintainer, mentioned in __maintainers__ and the CODEOWNERS file
• All external dependencies must be pip-installable. While not required to be on
  PyPI, a public git repository or another public URL with a repository is acceptable.

New stable architectures

• Provision of regression prediction tests with a small (not exported) checkpoint
  file.
• Comprehensive architecture documentation
• If an architecture has external dependencies, all must be publicly available on
  PyPI.
• Adherence to the standard output infrastructure of metatrain, including
  logging and model save locations.

---

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

[HaoZeke]

This is two new experimental architectures plus a pile of core mutations (scaler, composition, loss, eval, SPACE). It cannot land in this shape.

The package-level breakage is in src/metatrain/utils/additive/remove.py: it imports metatrain.experimental.edge_composition, which imports elearn. elearn is not on PyPI and is not a metatrain extra. That is why tests dies at conftest collection (ModuleNotFoundError: No module named 'elearn') and why docs die the same way. A core util cannot import an experimental architecture.

src/metatrain/utils/loss.py has the same layering leak. MatrixLoss.compute imports graph2mat and experimental.graph2mat, then hardcodes matrix_name = "hamiltonian". Loss functions in utils/ have to stay architecture-agnostic.

tox.ini is also just wrong. [testenv:mace-tests] now changedirs into experimental/edge_composition/tests/, and [testenv:edge-composition-tests] points at experimental/dpa3/tests/. There is no graph2mat-tests env at all.

On the experimental lifecycle (docs/src/dev-docs/architecture-life-cycle.rst):

• graph2mat has no ArchitectureTests / TorchscriptTests. The only test (tests/test_transforms.py) reads /home/febrer/COSMO_disk/.../scfbench_main_100.zip.
• MetaGraph2Mat.export is NotImplementedError (model.py:614). JIT is a listed experimental requirement.
• pyproject.toml has no graph2mat extra (graph2mat is on PyPI 0.0.13; e3nn and sphericart-torch are also imported at module level).
• edge-composition extras list only sphericart-torch. The model imports elearn and featomic.torch at the top of model.py.
• CODEOWNERS is untouched. Changelog Unreleased is empty. Contributor checklist is all [ ].

The core side-effects look like debug leftovers, not a reviewable contract:

• composition/model.py comments out densify_atomic_basis_dataset_info and the eval-mode sparsify, and leaves if not self.training: ....
• composition/trainer.py comments out atomic_basis_transform.
• experimental/space/model.py stops skipping mtt::aux:: outputs.
• additive/remove.py has if True or isinstance(...) / if False and not isinstance(...).

experimental.graph2mat as a wrapper around a featurizer is a reasonable experiment. The PR description already says the target type is made up, export is unfinished, GPU is untested, and this might stay a draft until Joe's PET matrix path exists. That is fine as a branch. It is not fine as a merge into main.

What I would do:

1. Revert every change under src/metatrain/{cli,composition,scaler,utils,experimental/space} and tox.ini except things that have their own PR and tests (the atom_pair helpers maybe).
2. Split edge_composition out. It is a second architecture and it currently poisons the default import graph.
3. Give graph2mat the usual experimental skeleton: extra, CODEOWNERS, tox + architecture-tests job, ArchitectureTests + TorchscriptTests on an in-repo fixture, working export(), no star import, no MACE copy-paste leftovers (model.py:46, restart error string, cutoff = 80.999 + i * 0.03).
4. Keep MatrixLoss / Blocks2Matrix inside experimental.graph2mat until they do not name hamiltonian.

I did not run the suite. The GitHub checks are already red on lint, docs, tests, and every architecture job, and the collection error above is enough.

Not approving this.

[HaoZeke]

This is the collection killer. utils.additive is imported from the default test path, so this pulls experimental.edge_composition (and elearn, which is not on PyPI) into every mtt install. Core utils cannot depend on an experimental architecture.

[HaoZeke]

if True or isinstance(...) and the matching if False and not isinstance(...) below look like debug leftovers. They also flip the train/eval contract that the comment two lines down is still describing.

[HaoZeke]

utils.loss.MatrixLoss importing experimental.graph2mat and hardcoding matrix_name = "hamiltonian" is not a public loss API. Keep this inside the architecture until it can name an arbitrary matrix target.

[HaoZeke]

mace-tests now runs edge_composition tests (and the new edge-composition-tests env changedirs into dpa3/tests/). Restore the MACE env and add a real graph2mat-tests one.

[HaoZeke]

This comments out the densify path that #1182 / #1115 put in so composition checkpoints stay portable. The eval-mode sparsify a bit lower is also commented out, and if not self.training: ... is a no-op. This is a stable architecture; do not gut it from a graph2mat PR.

[HaoZeke]

Experimental gate in architecture-life-cycle.rst is JIT / export(). A NotImplementedError here means this is not an experimental architecture yet, it is a training-only prototype. Same file still says "Interface of MACE for metatrain" and uses leftover i in cutoff=80.999 + i * 0.03.

[HaoZeke]

This is not a CI test. Point it at something in tests/resources/ (there is already a scfbench_2_bidirectional_edges.zip in this PR) or drop it until there is an in-repo fixture. Also missing ArchitectureTests / TorchscriptTests.

[HaoZeke]

elearn is not on PyPI. Experimental criterion 6 is "pip-installable, public URL". Even then this architecture does not belong in the same PR as graph2mat, and it must not be imported from utils/.

[HaoZeke]

Unrelated SPACE change: mtt::aux:: outputs will now leak into return_dict. Leave SPACE alone in this PR.

[pfebrer]

I will track progress in pfebrer#15. Anyway this will probably end up being a hook #1209 instead of a full architecture