lab-cosmo / librascal

A scalable and versatile library to generate representations for atomic-scale learning
https://lab-cosmo.github.io/librascal/
GNU Lesser General Public License v2.1
80 stars 17 forks source link

Gradients of (λ-)SOAP features #402

Closed kzinovjev closed 2 years ago

kzinovjev commented 2 years ago

Hi, I'm using compute_lambda_soap from rascal.utils to compute both SOAP and λ-SOAP features - it's a single line of code and does the job well. I was wondering, is there a similarly concise way to also get the gradients of the features w.r.t. the atomic positions? Or at least w.r.t. the spherical expansion coefs? Thanks!

ceriottm commented 2 years ago

Dear Kirill, to compute SOAP gradients, you simply need to set "compute_gradients": True in the hypers dictionary. Extracting the gradients is not entirely trivial because (for obvious reasons) they are stored in a sparse format, i.e. only for the neighbors of each atom. I can't find an example right now, hopefully someone else will chime in. As far as I remember though, gradients won't work for lambda-SOAP, which we are currently moving to compute with a different infrastructure.

kzinovjev commented 2 years ago

Ah, that's a pity - would need the gradients for λ as well. Not in a rush though, so good to know that you are going to work on that at some point. An example of getting the explicit gradients from the sparse ones would indeed be handy! Having it in the docs or an example notebook would make sense as well. Thanks!

Luthaf commented 2 years ago

For spherical expansion, you can use something like this for gradients:

from rascal.representations  import SphericalExpansion

calculator = SphericalExpansion(**hypers, compute_gradients=True)
managers = calculator.transform(frames)

features = managers.get_features(calculator)
# features is a n_atoms x n_features array

gradients = managers.get_features_gradient(calculator)
# gradients is a 3 x n_neighbors x n_features array

metadata = managers.get_gradients_info()

metadata is a n_neighbors x 5 array, where the columns correspond to structure index, central atom index, neighbor atom index, central atom species, and neighbor atom species. The ith row in metadata contains the metadata from 3 * i to 3 * i + 2 rows in gradients (3 rows for gradients in x/y/z directions), and indicate which of the row in feature we are taking the gradient of (structure index + central atom index) and with respect to which neighbor atom the gradient is taken.