Adding LAMMPS computes for MACE descriptors and descriptor gradients.

[Fraser-Birks]

A useful feature for symmetrix would be the ability to use a LAMMPS compute to directly extract the per-atom descriptors and descriptor gradients, following the precedent set by compute sna/atom .

To be clear, by 'descriptor' what I mean is the invariant part of the final layer, the part which is typically used as the input to the MLP readout.

This PR will aim to

• Implement a CPU compute for descriptor extraction (similar to compute sna/atom).
• Implement a CPU compute for descriptor gradient extraction (similar to compute snad/atom).
• Write good tests for both the above.

Before merging, would be good to:

• Sort out the file location (computes are currently placed in pair_symmetrix/ for ease of use with install.sh).
• Decide on a good name (right now commands will be compute symmetrix/mace/atom and compute symmetrix/maced/atom)
• Discuss if a Kokkos implementation is necessary.
• Potentially discuss moving duplicated code to different files.

[wcwitt]

Definitely open to this, thanks.

invariant part of the final layer

Are you sure this is what you'd want? I would have guessed concatenated invariants from all layers.

[Fraser-Birks]

Are you sure this is what you'd want? I would have guessed concatenated invariants from all layers.

I'm not sure! Is there a precedent for what makes a 'MACE descriptor'?

[bernstei]

Maybe take a look at what the torch ASE calculator returns? https://github.com/ACEsuit/mace/blob/main/mace/calculators/mace.py#L580

[Fraser-Birks]

Maybe take a look at what the torch ASE calculator returns? https://github.com/ACEsuit/mace/blob/main/mace/calculators/mace.py#L580

Yep - it's the concatenated invariants from all layers. I'll update the computes!

[chiang-yuan]

Just dropped by this issue. I think this is exciting because one potentially could use these descriptor as the CVs for enhanced MD (metadynamics, umbrella sampling, biased MD, etc.) Ideas similar to what has been done somewhere lately https://pubs.acs.org/doi/full/10.1021/acs.jctc.5c01767

[Fraser-Birks]

Descriptors now include the first layer invariants, exactly matching what's returned by get_descriptors() in python. I've also updated the gradient compute to include the first layer part and checked it matches finite differences.

However, to get descriptors that matched the python I had to undo some of @wcwitt 's clever fusion with the first layer invariants. That means I'm now exporting and reading an extra property from the model json file. As that property doesn't exist in the test json files, I've broken the C.I.

@wcwitt - any chance I could request that you re-generate the .json files used in the tests with the new version of extract_mace_data.py?