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?

[Fraser-Birks]

@wcwitt PR is ready - the reason the tests are failing is because extract_mace_data.py now also needs to extract linear_up_l0_inv, which is required later for the symmetrix descriptors to match the mace-torch versions. If you regenerate the .json used in the tests:

MODEL_FILE = "MACE-OFF23_small-1-8.json"
if not os.path.exists(MODEL_FILE):
    urlretrieve(
        "https://www.dropbox.com/scl/fi/zbg122s1zeeb1j6ogheok/MACE-OFF23_small-1-8.json?rlkey=mqb7cje9y3l0smwf75cfoahr7&st=iabk9093&dl=1",
        MODEL_FILE,
    )

With the new version of extract_mace_data.py, all tests should pass smoothly (I hope!)

[Fraser-Birks]

Just to note what this PR does - compute_symmetrix_mace_atom.cpp (and _kokkos version) are pretty straightforward, they just generate the per-atom mace descriptor built from concatenated invariants (i.e, matching what one gets from mace-torch with extract_mace_descriptor()). To set up this compute (it works totally separately from the symmetrix/mace pair_style), one needs to pass the following to LAMMPS:

compute         macedesc all symmetrix/mace/atom {MODEL_FILE} H O

To extract it into python (in the correct order):

desc = lmp.numpy.extract_compute(
            "macedesc", lmpmod.LMP_STYLE_ATOM, lmpmod.LMP_TYPE_ARRAY)
ids = lmp.numpy.extract_atom("id")
order = np.argsort(ids)
desc = np.array(desc, copy=True)[order]

compute_symmetrix_maced_atom.cpp (and _kokkos version) are a bit more complex. By default, it extracts the full global descriptor Jacobian,

$$ J_{k,i\alpha} \equiv \frac{\partial q_k}{\partial r_{i\alpha}} = \sum_{j=1}^{N} \frac{\partial d_{jk}}{\partial r_{i\alpha}} $$

where the global descriptor $\mathbf{q} \in \mathbb{R}^{K}$ is the sum over local descriptors $\mathbf{d}_{i}$. This is fairly expensive, as it requires $K$ backward passes. There is, however, a cheaper alternative built in. If one has some collective variable which is a function of the global descriptor, $\eta = f(\mathbf{q})$, then all one needs to obtain the gradient of that with respect to the Cartesian atomic positions $\partial \eta / \partial r_{i\alpha}$ is the vector-Jacobian product (VJP):

$$ \frac{\partial \eta}{\partial r_{i\alpha}} = \sum_{k=1}^{K} \frac{\partial f}{\partial q_k} \frac{\partial q_k}{\partial r_{i\alpha}} = \sum_{k=1}^{K} v_k J_{k,i\alpha} \qquad v_k \equiv \frac{\partial f}{\partial q_k} $$

Provided the user can supply a LAMMPS compute id for $v_k$ (the derivative of $\eta$ wrt the global descriptor $q$), computing the vjp only requires 1 backwards pass.

Syntax is:

compute         macedescgrad all symmetrix/maced/atom {MODEL_FILE} H O # Full Jacobian mode
compute         macedescvjp all symmetrix/maced/atom {MODEL_FILE} H O vjp <compute_v_id> # VJP mode

Extract in Python with:

jac = lmp.numpy.extract_compute(
            "macedescgrad", lmpmod.LMP_STYLE_ATOM, lmpmod.LMP_TYPE_ARRAY)
ids = lmp.numpy.extract_atom("id")
order = np.argsort(ids)
jac = np.array(jac, copy=True)[order]
# layout: jac[atom][{0,1,2}*(2*num_channels) + col],
# col in [0, num_channels) -> d(h1_restored[col])/d{x,y,z}[atom]
# col in [num_channels, 2*num_channels) -> d(H2[col-C])/d{x,y,z}[atom]
two_c = 2 * NUM_CHANNELS
assert jac.shape == (len(ATOMS), 3 * two_c)

vjp = lmp.numpy.extract_compute(
    "macedescvjp", lmpmod.LMP_STYLE_ATOM, lmpmod.LMP_TYPE_ARRAY)
ids = lmp.numpy.extract_atom("id")
order = np.argsort(ids)
vjp = np.array(vjp, copy=True)[order]
assert vjp.shape == (len(ATOMS), 3)

Both jac and vjp are covered with finite difference tests against the separate python version - CPU and GPU.