MACE for NMR shielding prediction - help needed

[jan-vavrin]

Hi all,

I would like to use MACE to predict NMR parameters - specifically, NMR shielding tensors. I have already tried to implement it at https://github.com/jan-vavrin/NMR_mace, but I got to a point where I'm not sure how to proceed further. I generally followed the changes in the mu_alpha branch (#1079) and tried to match them to NMR shieldings. Would anyone be willing to help me with this?

The NMR shielding tensor is a general 3x3 tensor (i.e. "1x0e+1x1o+1x2e" in spherical irreps) for every atom. An input file for the model would therefore look something like test.extxyz.txt - we have to flatten the 3x3 tensor in the xyz file. [Tiny note: the tensor is sometimes assumed to be symmetric, but it need not be in general - however, the antisymmetric "1o" part is usually quite small.]

It is basically impossible to calculate analytical derivatives of NMR shielding tensors and numerical gradients are very expensive - the model would therefore have to be trained only on the actual values.

The things I'm struggling with currently are:

• calculating the loss correctly
• calculating the NMR shielding derivatives
• how to set it up so as not to require energies in the input file

An additional point is that when comparing NMR shielding calculations, it is usually useful to distinguish the RMSEs per element - i.e. to have a separate loss for hydrogen, carbon, etc. This is because their shieldings generally cover different ranges, calculations (and experiments) for heavier elements are less accurate, etc. Would something like this be possible in MACE as well? (I'm not really sure how one would access element information in the table outputs.)

There are possibilities for extending this in the future:

If this proved to be successful, I would like to generalise the model to predict electric field gradient (EFG) tensors as well - these are traceless symmetric tensors ("2e" only), so the input would have to only use the "2e" part of a 3x3 tensor. I suppose there is a nice way to do this, but I'm not that familiar with e3nn to know how.

One more NMR parameter of interest are J coupling constants - however, these are a scalar value for every pair of atoms. I'm not sure whether it would even be possible to predict such a property with MACE in a properly invariant way.

[ilyes319]

Hey,

All these things are in principle possible, but need a decent amount of coding that is quite similar to what have been done to make the Dielectric MACE model:

mace/mace/modules/models.py

Line 827 in 42cd495

class AtomicDielectricMACE(torch.nn.Module):

I know several people already have working implementations of these things and are working on a PR for it @gelzinyte @Nilsgoe, and I would recommend you contact them.

[Nilsgoe]

@jan-vavrin, I would highly recommend that you take a look at https://github.com/gelzinyte/mace/tree/efg-mace from @gelzinyte. She solved some of your problems in this branch. You should take a look at some of the other branches as well, they might be helpful.