Computing \(\lambda\)-SOAP features

In a previous example, we computed the SOAP scalar descriptor. This tutorial focuses on the equivariant generalization of SOAP to tensorial objects, commonly known as \(\lambda\)-SOAP. The key difference is that the SOAP scalar is the inner product of two spherical expansions, whereas \(\lambda\)-SOAP contracts two spherical expansions using a Clebsch-Gordan coefficient:

\[\rho^{\otimes 2}_{z_1 z_2, n_1 n_2, l_1, l_2, \lambda \mu} (\{\mathbf{r}_i\}) = \sum_{m_1, m_2} C_{l_1 l_2, m_1 m_2}^{\lambda\mu} \rho_{z_1,n_1 l_1 m_1}(\{\mathbf{r}_i\}) \rho_{z_2,n_2 l_2 m_2}(\{\mathbf{r}_i\})\]

import chemfiles
import numpy as np
from metatensor import Labels

from featomic import LodeSphericalExpansion, SphericalExpansion
from featomic.clebsch_gordan import EquivariantPowerSpectrum

Let’s see how to compute the \(\lambda\)-SOAP descriptor using featomic.

First we can read the input systems using chemfiles. You can download the dataset for this example from our website.

with chemfiles.Trajectory("dataset.xyz") as trajectory:
    systems = [s for s in trajectory]

Featomic also handles systems read by ASE:

systems = ase.io.read("dataset.xyz", ":").

Next, define the hyperparameters for the spherical expansion:

HYPERPARAMETERS = {
    "cutoff": {
        "radius": 5.0,
        "smoothing": {"type": "ShiftedCosine", "width": 0.5},
    },
    "density": {
        "type": "Gaussian",
        "width": 0.3,
    },
    "basis": {
        "type": "TensorProduct",
        "max_angular": 2,
        "radial": {"type": "Gto", "max_radial": 2},
    },
}

Create the spherical expansion calculator. The SphericalExpansion class uses the hyperparameters above. Then, wrap it with EquivariantPowerSpectrum to compute the Clebsch-Gordan contraction for \(\lambda\)-SOAP.

spex_calculator = SphericalExpansion(**HYPERPARAMETERS)
calculator = EquivariantPowerSpectrum(spex_calculator)

Run the actual calculation

power_spectrum = calculator.compute(systems, neighbors_to_properties=True)

The result is a TensorMap whose keys encode symmetry:

power_spectrum.keys

Labels(
    o3_lambda  o3_sigma  center_type
        0         1           1
        1         1           1
        2         1           1
        1         -1          1
        2         -1          1
        3         1           1
        3         -1          1
        4         1           1
        0         1           6
        1         1           6
        2         1           6
        1         -1          6
        2         -1          6
        3         1           6
        3         -1          6
        4         1           6
        0         1           7
        1         1           7
        2         1           7
        1         -1          7
        2         -1          7
        3         1           7
        3         -1          7
        4         1           7
        0         1           8
        1         1           8
        2         1           8
        1         -1          8
        2         -1          8
        3         1           8
        3         -1          8
        4         1           8
)

Often, you only need specific \(\lambda\) values. For example, if the target property is the polarizability tensor, (a rank-2 symmetric Cartesian tensor), you can restrict the output to \(\lambda=0\) and \(\lambda=2\) (with \(\sigma=1\) to discard pseudotensors) using the selected_keys parameter:

power_spectrum_0_2 = calculator.compute(
    systems,
    neighbors_to_properties=True,
    selected_keys=Labels(["o3_lambda", "o3_sigma"], np.array([[0, 1], [2, 1]])),
)
power_spectrum_0_2.keys

Labels(
    o3_lambda  o3_sigma  center_type
        0         1           1
        2         1           1
        0         1           6
        2         1           6
        0         1           7
        2         1           7
        0         1           8
        2         1           8
)

You can also compute a \(\lambda\)-SOAP-like descriptor using two different expansions. For instance, combine a standarad spherical expansion with a long-range LodeSphericalExpansion:

LODE_HYPERPARAMETERS = {
    "density": {
        "type": "SmearedPowerLaw",
        "smearing": 0.3,
        "exponent": 1,
    },
    "basis": {
        "type": "TensorProduct",
        "max_angular": 2,
        "radial": {"type": "Gto", "max_radial": 3, "radius": 1.0},
    },
}
lode_calculator = LodeSphericalExpansion(**LODE_HYPERPARAMETERS)
calculator = EquivariantPowerSpectrum(spex_calculator, lode_calculator)
power_spectrum = calculator.compute(systems, neighbors_to_properties=True)
power_spectrum.keys

Labels(
    o3_lambda  o3_sigma  center_type
        0         1           1
        1         1           1
        2         1           1
        1         -1          1
        2         -1          1
        3         1           1
        3         -1          1
        4         1           1
        0         1           6
        1         1           6
        2         1           6
        1         -1          6
        2         -1          6
        3         1           6
        3         -1          6
        4         1           6
        0         1           7
        1         1           7
        2         1           7
        1         -1          7
        2         -1          7
        3         1           7
        3         -1          7
        4         1           7
        0         1           8
        1         1           8
        2         1           8
        1         -1          8
        2         -1          8
        3         1           8
        3         -1          8
        4         1           8
)