From afd04a404ecafb20ff89a7d0c455d8ef18eeebed Mon Sep 17 00:00:00 2001 From: Joseph Abbott Date: Mon, 28 Apr 2025 15:48:06 +0200 Subject: [PATCH 1/2] Add tutorials for high-level CG calculators Co-Authored-By: ppegolo --- docs/src/how-to/computing-lambda-soap.rst | 22 ++ docs/src/how-to/density-correlations.rst | 16 + docs/src/how-to/index.rst | 3 + .../how-to/per-pair-equivariant-features.rst | 15 + .../featomic/examples/compute-lambda-soap.py | 109 +++++++ .../featomic/examples/density-correlations.py | 300 ++++++++++++++++++ .../examples/per-pair-equivariant-features.py | 85 +++++ tox.ini | 2 + 8 files changed, 552 insertions(+) create mode 100644 docs/src/how-to/computing-lambda-soap.rst create mode 100644 docs/src/how-to/density-correlations.rst create mode 100644 docs/src/how-to/per-pair-equivariant-features.rst create mode 100644 python/featomic/examples/compute-lambda-soap.py create mode 100644 python/featomic/examples/density-correlations.py create mode 100644 python/featomic/examples/per-pair-equivariant-features.py diff --git a/docs/src/how-to/computing-lambda-soap.rst b/docs/src/how-to/computing-lambda-soap.rst new file mode 100644 index 000000000..1a013a548 --- /dev/null +++ b/docs/src/how-to/computing-lambda-soap.rst @@ -0,0 +1,22 @@ +.. _userdoc-how-to-computing-lambda-soap: + +Computing :math:`\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 +:math:`\lambda`-SOAP. The key difference is that the SOAP scalar is the inner product +of two spherical expansions, whereas :math:`\lambda`-SOAP contracts two spherical +expansions using a Clebsch-Gordan coefficient: + +.. math:: + + \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\}) + +.. include:: ../examples/compute-lambda-soap.rst + :start-after: start-body + :end-before: end-body \ No newline at end of file diff --git a/docs/src/how-to/density-correlations.rst b/docs/src/how-to/density-correlations.rst new file mode 100644 index 000000000..74f9f038f --- /dev/null +++ b/docs/src/how-to/density-correlations.rst @@ -0,0 +1,16 @@ +.. _userdoc-how-to-density-correlations: + +Computing density correlations +============================== + +In a `previous example +`_, we +computed the lambda-SOAP **equivariant** descriptor. This tutorial focuses on the +computation of arbitrary body-order descriptors by means of an auto-correlations of +spherical expansion (or 'density'). Lambda-SOAP is the body-order 3 version of this, but +we will also explore the bispectrum, and how to combine iterative Clebsch-Gordan tensor +products with feature space contractions. + +.. include:: ../examples/density-correlations.rst + :start-after: start-body + :end-before: end-body \ No newline at end of file diff --git a/docs/src/how-to/index.rst b/docs/src/how-to/index.rst index e349cc897..1fd4a122a 100644 --- a/docs/src/how-to/index.rst +++ b/docs/src/how-to/index.rst @@ -11,8 +11,11 @@ are a total beginner, you can go to :ref:`userdoc-get-started` section. :maxdepth: 1 computing-soap + computing-lambda-soap + density-correlations sample-selection property-selection keys-selection splined-radial-integral long-range + per-pair-equivariant-features diff --git a/docs/src/how-to/per-pair-equivariant-features.rst b/docs/src/how-to/per-pair-equivariant-features.rst new file mode 100644 index 000000000..ac0811c5a --- /dev/null +++ b/docs/src/how-to/per-pair-equivariant-features.rst @@ -0,0 +1,15 @@ +.. _userdoc-how-to-per-pair-equivariant-features: + +Computing per-pair equivariant features +======================================= + +In a `previous example +`_, we +computed the :math:`\lambda`-SOAP equivariant descriptor. This tutorial focuses on the +generalization of :math:`\lambda`-SOAP to pairs of atoms rather than single centers. +Per-pair equivariant features are useful as descriptors for `two-center quantities such +as the Hamiltonian matrix `_. + +.. include:: ../examples/per-pair-equivariant-features.rst + :start-after: start-body + :end-before: end-body \ No newline at end of file diff --git a/python/featomic/examples/compute-lambda-soap.py b/python/featomic/examples/compute-lambda-soap.py new file mode 100644 index 000000000..f11681894 --- /dev/null +++ b/python/featomic/examples/compute-lambda-soap.py @@ -0,0 +1,109 @@ +""" +Computing λ-SOAP features +========================== + +.. start-body +""" + +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 :math:`\lambda`-SOAP descriptor using Featomic. +# +# Read systems using Chemfiles. You can download the dataset for this example from our +# :download:`website <../../static/dataset.xyz>`. + +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 :class:`~featomic.SphericalExpansion` +# class uses the hyperparameters above. Then, wrap it with +# :class:`~featomic.clebsch_gordan.EquivariantPowerSpectrum` to compute the +# Clebsch-Gordan contraction for :math:`\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 :class:`~metatensor.TensorMap` whose keys encode symmetry: + +power_spectrum.keys + +# %% +# Often, you only need specific :math:`\lambda` values. For example, if the +# `target property is the polarizability tensor +# `_, +# (a rank-2 symmetric Cartesian tensor), you can restrict the output to +# :math:`\lambda=0` and :math:`\lambda=2` (with :math:`\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 + +# %% +# You can also compute a :math:`\lambda`-SOAP-like descriptor using two different +# expansions. For instance, combine a standarad spherical expansion with a long-range +# :class:`~featomic.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 + +# %% +# +# .. end-body diff --git a/python/featomic/examples/density-correlations.py b/python/featomic/examples/density-correlations.py new file mode 100644 index 000000000..6ac206a98 --- /dev/null +++ b/python/featomic/examples/density-correlations.py @@ -0,0 +1,300 @@ +""" +Computing density correlations +============================== + +.. start-body +""" + +import chemfiles +import metatensor as mts +import numpy as np + +from featomic import SphericalExpansion +from featomic.clebsch_gordan import DensityCorrelations, EquivariantPowerSpectrum + + +# %% +# +# Compute the spherical expansion +# ------------------------------- +# +# Read systems using chemfiles. You can obtain the dataset used in this +# example from our :download:`website <../../static/dataset.xyz>`. + +with chemfiles.Trajectory("dataset.xyz") as trajectory: + systems = [s for s in trajectory] + +# %% +# +# Featomic can also handles systems read by `ASE +# `_ using +# +# ``systems = ase.io.read("dataset.xyz", ":")``. +# +# We can now define the spherical expansion hyper parameters and compute the density. +# This will be used throughout the remainder of the tutorial. + +HYPER_PARAMETERS = { + "cutoff": { + "radius": 5.0, + "smoothing": {"type": "ShiftedCosine", "width": 0.5}, + }, + "density": { + "type": "Gaussian", + "width": 0.3, + }, + "basis": { + "type": "TensorProduct", + "max_angular": 3, + "radial": {"type": "Gto", "max_radial": 4}, + }, +} + +# Initialize the SphericalExpansion calculator and compute the density +spex_calc = SphericalExpansion(**HYPER_PARAMETERS) +density = spex_calc.compute(systems) + +# %% +# +# Move to "neighbor_type" to properties, i.e. remove sparsity in this dimension. +# +# Note: this method should be called with care when computing on a subset of systems +# from a larger dataset. If the ``systems`` being computed contain a subset of the +# global atom types, an inconsistent feature dimension will be created. In this case, +# the argument "keys_to_move" should be specified as a ``Labels`` object with all global +# atom types. + +density = density.keys_to_properties("neighbor_type") + +# average number of features per block +print("total number of features:", np.sum([len(block.properties) for block in density])) + +# %% +# +# Density correlations to build a lambda-SOAP +# ------------------------------------------- +# +# Now we can use the DensityCorrelations calculator and specify that we want to take a +# single CG tensor product, i.e. ``n_correlations=1``. +# +# Initializing the calculator computes and stores the CG coefficients. As the density +# expansion is up to ``o3_lambda=3`` and we are doing a single contraction, we need CG +# coefficients computed up to ``o3_lambda=6`` in order to do a full contraction. Hence, +# we set ``max_angular=6``. + +dens_corr_calc = DensityCorrelations( + n_correlations=1, + max_angular=6, +) + +# %% +# +# This outputs an equivariant power spectrum descriptor of body-order 3, i.e. +# "lambda-SOAP" features. + +lambda_soap = dens_corr_calc.compute(density) + +# %% +# +# This is not quite equivalent to the result seen in the previous tutorial on +# :ref:`computing lambda-SOAP `. The keys contain dimensions +# ``"l_1"`` and ``"l_2"`` which for a given block track the angular order of the blocks +# combined to create it. +# +# Keeping these dimensions in the keys is useful to allow for further CG products to be +# taken, building more complex descriptors. This is covered in a later tutorial on the +# :ref:`use of ClebschGordanProduct `. +# +# For now, we can move these key dimensions to properties. Inspect the metadata before +# and after moving these dimensions. + +print("lambda-SOAP before keys_to_properties:", lambda_soap) +print("lambda-SOAP before keys_to_properties, first block:", lambda_soap[0]) + +lambda_soap = lambda_soap.keys_to_properties(["l_1", "l_2"]) + +print("lambda-SOAP after keys_to_properties:", lambda_soap) +print("lambda-SOAP after keys_to_properties, first block:", lambda_soap[0]) + +# %% +# +# This obtains a result that is equivalent to the lambda-SOAP seen in the previous +# tutorial. We can confirm this by computing an ``EquivariantPowerSpectrum`` and +# checking for consistency. + +pow_spec_calc = EquivariantPowerSpectrum(spex_calc) +equi_pow_spec = pow_spec_calc.compute(systems, neighbors_to_properties=True) + +assert mts.equal(lambda_soap, equi_pow_spec) + +# %% +# +# Computing the bispectrum +# ======================== +# +# Higher body order descriptors can be computed by increasing the ``n_correlations`` +# parameter. The ``max_angular`` should also be increased to account for the increased +# combinations in angular momenta. +# +# With more iterations, the cost of the computation scales unfavourably. Let's use a +# density with small hyper parameters to demonstrate calculation of the bispectrum, a +# body-order 4 equivariant descriptor. + +HYPER_PARAMETERS_SMALL = { + "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}, + }, +} + +# %% +# +# Taking two CG combinations of a density expanded to ``o3_lambda=2`` requires CG +# coefficients computed up to ``max_angular=6``. This is given by ``(n_iterations + 1) * +# max_angular_density``. + +# Initialize the SphericalExpansion calculator and compute the density +spex_calc = SphericalExpansion(**HYPER_PARAMETERS_SMALL) +density = spex_calc.compute(systems) +density = density.keys_to_properties("neighbor_type") + +# Initialize DensityCorrelations calculator +dens_corr_calc = DensityCorrelations( + n_correlations=2, + max_angular=6, +) + +# Compute the bispectrum +bispectrum = dens_corr_calc.compute(density) + +# %% +# +# Inspect the features of the bispectrum. There are now ``"neighbor_x_type"`` and +# ``"n_x"`` (which track the radial channel indices) dimensions created by the product +# of feature spaces of the 3 density blocks combined to make each bispectrum block. +# +# For each block, its key contains dimensions tracking the angular order of blocks +# combined to create it, namely ``["l_1", "l_2", "k_2", "l_3"]``. The ``"l_"`` +# dimensions track the angular order of the blocks from the original density, while +# ``"k_"`` dimensions track the angular order of intermediate blocks. +# +# Let's look at an example. Take the block indexed by key: +# +# ``LabelsEntry(o3_lambda=3, o3_sigma=1, center_type=7, l_1=1, l_2=2, k_2=1, l_3=2)`` +# +# This was created in the following way. +# +# First, a block from ``density`` of angular order ``l_1=1`` was combined with a block +# of order ``l_2=2``. Angular momenta coupling rules state that non-zero combinations +# can only be created for output blocks with order ``| l_1 - l_2 | <= k_2 <= | l_1 + l_2 +# |``, corresponding to ``[1, 2, 3]``. In this case, a block of order ``k_2=1`` was +# created. +# +# Next, this intermediate block of order ``k_2=1`` was then combined with a block from +# the original density of order ``l_3=2``. This can again create combinations ``[1, 2, +# 3]``, and in this case has been combined to create the output angular order +# ``o3_lambda=3``. + +print("bispectrum:", bispectrum) +print("bispectrum first block:", bispectrum[0]) + +# %% +# +# As before, move these symmetry keys to properties and inspect the metadata and the +# mean size of the features. + +bispectrum = bispectrum.keys_to_properties(["l_1", "l_2", "k_2", "l_3"]) + +print("bispectrum:", bispectrum) +print("bispectrum first block:", bispectrum[0]) +print( + "total number of features:", np.sum([len(block.properties) for block in bispectrum]) +) + +# %% +# +# Instead of computing the full CG product, a threshold can be defined to limit the +# maximum angular order of blocks computed at each step in the iterative CG coupling +# steps. +# +# This is controlled by the ``angular_cutoff`` parameter, and allows us to initialize +# the calculator with a lower ``max_angular``. +# +# Note that any truncation of the angular channels away from the maximal allowed by +# angular momenta coupling rules results in some loss of information. +# +# Let's truncate to an angular cutoff of 3: + +angular_cutoff = 3 +dens_corr_calc = DensityCorrelations( + n_correlations=2, + max_angular=angular_cutoff, +) +bispectrum_truncated = dens_corr_calc.compute(density, angular_cutoff=angular_cutoff) + +# Move the "l_" and "k_" keys to properties +bispectrum_truncated = bispectrum_truncated.keys_to_properties( + ["l_1", "l_2", "k_2", "l_3"] +) + +print("truncated bispectrum:", bispectrum_truncated) +print("truncated bispectrum first block:", bispectrum_truncated[0]) +print( + "total number of features:", + np.sum([len(block.properties) for block in bispectrum_truncated]), +) + +# %% +# +# To compute a descriptor that matches the symmetry of a given target property, the +# ``selected_keys`` argument can be passed to the ``compute`` method. This was also seen +# in the previous tutorial on :ref:`computing lambda-SOAP `. +# +# Following this example, to compute the truncated bispectrum for a polarizability +# tensor: +bispectrum_truncated_key_select = dens_corr_calc.compute( + density, + angular_cutoff=angular_cutoff, + selected_keys=mts.Labels( + ["o3_lambda", "o3_sigma"], + np.array([[0, 1], [2, 1]]), + ), +) + +# Move the "l_" and "k_" keys to properties +bispectrum_truncated_key_select = bispectrum_truncated_key_select.keys_to_properties( + ["l_1", "l_2", "k_2", "l_3"] +) + +print("truncated bispectrum with selected keys:", bispectrum_truncated_key_select) +print( + "total number of features:", + np.sum([len(block.properties) for block in bispectrum_truncated_key_select]), +) + + +# %% +# +# Conclusion +# ========== +# +# ``DensityCorrelations`` can be used to build equivariants of arbitrary body order from +# a spherical expansion of decorated atomic point cloud data. +# +# A key limitation of this approach is an exploding feature size. To reduce the number +# of output blocks, the ``angular_cutoff`` parameter can be used. This can be controlled +# by introducing contractions after each CG combination in the iteration. This will be +# covered in the next tutorial, :ref:`Clebsch-Gordan products +# `. + +# %% +# .. end-body diff --git a/python/featomic/examples/per-pair-equivariant-features.py b/python/featomic/examples/per-pair-equivariant-features.py new file mode 100644 index 000000000..a82af8dd4 --- /dev/null +++ b/python/featomic/examples/per-pair-equivariant-features.py @@ -0,0 +1,85 @@ +""" +Computing per-pair equivariant features +======================================= + +.. start-body +""" + +import chemfiles +from metatensor import Labels + +from featomic import SphericalExpansion, SphericalExpansionByPair +from featomic.clebsch_gordan import EquivariantPowerSpectrumByPair + + +# %% +# +# Let's see how to compute the per-pair :math:`\lambda`-SOAP descriptor using Featomic. +# +# Read systems using Chemfiles. You can download the dataset for this example from our +# :download:`website <../../static/dataset.xyz>`. + +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 a spherical expansion and a spherical expansion by pair calculator. +# The :class:`~featomic.SphericalExpansion` and +# :class:`~featomic.SphericalExpansionByPair` classes use the hyperparameters above. +# Then, wrap them with :class:`~featomic.clebsch_gordan.EquivariantPowerSpectrumByPair` +# to compute the Clebsch-Gordan contraction for the per-pair :math:`\lambda`-SOAP. + +spex_calculator = SphericalExpansion(**HYPERPARAMETERS) +per_pair_spex_calculator = SphericalExpansionByPair(**HYPERPARAMETERS) +calculator = EquivariantPowerSpectrumByPair(spex_calculator, per_pair_spex_calculator) + +# %% +# Run the actual calculation + +per_pair_power_spectrum = calculator.compute(systems, neighbors_to_properties=True) + +# %% +# The result is a :class:`~metatensor.TensorMap` whose keys encode symmetry and the +# species of the atoms involved: + +per_pair_power_spectrum.keys + +# %% +# Often, you only need specific :math:`\lambda` values. For example, if the +# `target property is the Hamiltonian matrix on a minimal basis +# `_, you can restrict the output to :math:`\lambda` values +# up to :math:`\lambda=2` using the ``selected_keys`` parameter: + +per_pair_power_spectrum_minimal_basis = calculator.compute( + systems, + neighbors_to_properties=True, + selected_keys=Labels.range("o3_lambda", 3), +) +per_pair_power_spectrum_minimal_basis.keys + +# %% +# .. end-body diff --git a/tox.ini b/tox.ini index db7d5c7f1..db5c4ae76 100644 --- a/tox.ini +++ b/tox.ini @@ -122,6 +122,8 @@ description = Build the package documentation. deps = -r docs/requirements.txt {[testenv]packaging_deps} + metatensor-operations + metatensor-learn allowlist_externals = bash From 8e6b6cd1d3ba5a77456bfab44df34d643c95f8b4 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Mon, 14 Jul 2025 17:33:14 +0200 Subject: [PATCH 2/2] Multiple small updates to the tutorials --- docs/src/how-to/density-correlations.rst | 16 +- .../featomic/examples/compute-lambda-soap.py | 8 +- .../featomic/examples/density-correlations.py | 161 +++++++++--------- 3 files changed, 97 insertions(+), 88 deletions(-) diff --git a/docs/src/how-to/density-correlations.rst b/docs/src/how-to/density-correlations.rst index 74f9f038f..7e8360da1 100644 --- a/docs/src/how-to/density-correlations.rst +++ b/docs/src/how-to/density-correlations.rst @@ -4,13 +4,15 @@ Computing density correlations ============================== In a `previous example -`_, we -computed the lambda-SOAP **equivariant** descriptor. This tutorial focuses on the -computation of arbitrary body-order descriptors by means of an auto-correlations of -spherical expansion (or 'density'). Lambda-SOAP is the body-order 3 version of this, but -we will also explore the bispectrum, and how to combine iterative Clebsch-Gordan tensor -products with feature space contractions. +`_, +we computed the :math:`\lambda`-SOAP **equivariant** descriptor. This tutorial +focuses on the computation of arbitrary body-order descriptors by means of an +auto-correlations of spherical expansion (or 'density'). + +We will explore both :math:`\lambda`-SOAP --- a version of auto-correlations +with body-order 3 (taking into account one central atom and two neighbors) --- +an the bispectrum which is a representation with body-order 4. .. include:: ../examples/density-correlations.rst :start-after: start-body - :end-before: end-body \ No newline at end of file + :end-before: end-body diff --git a/python/featomic/examples/compute-lambda-soap.py b/python/featomic/examples/compute-lambda-soap.py index f11681894..1165431ef 100644 --- a/python/featomic/examples/compute-lambda-soap.py +++ b/python/featomic/examples/compute-lambda-soap.py @@ -1,4 +1,6 @@ """ +.. _compute-lambda-soap: + Computing λ-SOAP features ========================== @@ -15,10 +17,10 @@ # %% # -# Let's see how to compute the :math:`\lambda`-SOAP descriptor using Featomic. +# Let's see how to compute the :math:`\lambda`-SOAP descriptor using featomic. # -# Read systems using Chemfiles. You can download the dataset for this example from our -# :download:`website <../../static/dataset.xyz>`. +# First we can read the input systems using chemfiles. You can download the dataset for +# this example from our :download:`website <../../static/dataset.xyz>`. with chemfiles.Trajectory("dataset.xyz") as trajectory: systems = [s for s in trajectory] diff --git a/python/featomic/examples/density-correlations.py b/python/featomic/examples/density-correlations.py index 6ac206a98..f55a7c586 100644 --- a/python/featomic/examples/density-correlations.py +++ b/python/featomic/examples/density-correlations.py @@ -5,7 +5,7 @@ .. start-body """ -import chemfiles +import ase.io import metatensor as mts import numpy as np @@ -15,23 +15,10 @@ # %% # -# Compute the spherical expansion -# ------------------------------- +# Computing the spherical expansion +# --------------------------------- # -# Read systems using chemfiles. You can obtain the dataset used in this -# example from our :download:`website <../../static/dataset.xyz>`. - -with chemfiles.Trajectory("dataset.xyz") as trajectory: - systems = [s for s in trajectory] - -# %% -# -# Featomic can also handles systems read by `ASE -# `_ using -# -# ``systems = ase.io.read("dataset.xyz", ":")``. -# -# We can now define the spherical expansion hyper parameters and compute the density. +# We can define the spherical expansion hyper parameters and compute the density. # This will be used throughout the remainder of the tutorial. HYPER_PARAMETERS = { @@ -50,9 +37,11 @@ }, } +systems = ase.io.read("dataset.xyz", ":") + # Initialize the SphericalExpansion calculator and compute the density -spex_calc = SphericalExpansion(**HYPER_PARAMETERS) -density = spex_calc.compute(systems) +spherical_expansion = SphericalExpansion(**HYPER_PARAMETERS) +density = spherical_expansion.compute(systems) # %% # @@ -61,8 +50,8 @@ # Note: this method should be called with care when computing on a subset of systems # from a larger dataset. If the ``systems`` being computed contain a subset of the # global atom types, an inconsistent feature dimension will be created. In this case, -# the argument "keys_to_move" should be specified as a ``Labels`` object with all global -# atom types. +# the argument to ``keys_to_properties`` should be specified as a ``Labels`` object with +# all global atom types. density = density.keys_to_properties("neighbor_type") @@ -71,18 +60,18 @@ # %% # -# Density correlations to build a lambda-SOAP -# ------------------------------------------- +# Density correlations to build a :math:`\lambda`-SOAP +# ---------------------------------------------------- # -# Now we can use the DensityCorrelations calculator and specify that we want to take a -# single CG tensor product, i.e. ``n_correlations=1``. +# We can now use the ``DensityCorrelations`` calculator and specify that we want to take +# a single Clebsch-Gordan (CG) tensor product, i.e. ``n_correlations=1``. # -# Initializing the calculator computes and stores the CG coefficients. As the density -# expansion is up to ``o3_lambda=3`` and we are doing a single contraction, we need CG -# coefficients computed up to ``o3_lambda=6`` in order to do a full contraction. Hence, -# we set ``max_angular=6``. +# During initialization, the calculator computes and stores the CG coefficients. As the +# density expansion is up to ``o3_lambda=3`` and we are doing a single contraction, we +# need CG coefficients computed up to ``o3_lambda=6`` in order to do a full contraction. +# Hence, we set ``max_angular=6``. -dens_corr_calc = DensityCorrelations( +density_correlations = DensityCorrelations( n_correlations=1, max_angular=6, ) @@ -90,42 +79,45 @@ # %% # # This outputs an equivariant power spectrum descriptor of body-order 3, i.e. -# "lambda-SOAP" features. +# :math:`\lambda`-SOAP features. -lambda_soap = dens_corr_calc.compute(density) +lambda_soap = density_correlations.compute(density) # %% # # This is not quite equivalent to the result seen in the previous tutorial on -# :ref:`computing lambda-SOAP `. The keys contain dimensions -# ``"l_1"`` and ``"l_2"`` which for a given block track the angular order of the blocks -# combined to create it. +# :ref:`computing λ-SOAP `. The keys contain +# dimensions ``"l_1"`` and ``"l_2"`` which for a given block track the angular order of +# the blocks from the input combined to create the block in the output +# :class:`~metatensor.TensorMap`. # # Keeping these dimensions in the keys is useful to allow for further CG products to be -# taken, building more complex descriptors. This is covered in a later tutorial on the -# :ref:`use of ClebschGordanProduct `. -# -# For now, we can move these key dimensions to properties. Inspect the metadata before -# and after moving these dimensions. +# taken, building more complex descriptors. For now, we can move these key dimensions to +# properties. Inspect the metadata before and after moving these dimensions. -print("lambda-SOAP before keys_to_properties:", lambda_soap) -print("lambda-SOAP before keys_to_properties, first block:", lambda_soap[0]) +print("λ-SOAP before keys_to_properties:", lambda_soap.keys) +print("first block:", lambda_soap[0]) + +# %% +# lambda_soap = lambda_soap.keys_to_properties(["l_1", "l_2"]) -print("lambda-SOAP after keys_to_properties:", lambda_soap) -print("lambda-SOAP after keys_to_properties, first block:", lambda_soap[0]) +print("λ-SOAP after keys_to_properties:", lambda_soap.keys) +print("first block:", lambda_soap[0]) # %% # -# This obtains a result that is equivalent to the lambda-SOAP seen in the previous -# tutorial. We can confirm this by computing an ``EquivariantPowerSpectrum`` and -# checking for consistency. +# This obtains a result that is equivalent to the :math:`\lambda`-SOAP seen in the +# previous tutorial. We can confirm this by computing an ``EquivariantPowerSpectrum`` +# and checking for consistency. -pow_spec_calc = EquivariantPowerSpectrum(spex_calc) -equi_pow_spec = pow_spec_calc.compute(systems, neighbors_to_properties=True) +equivariant_ps_calculator = EquivariantPowerSpectrum(spherical_expansion) +equivariant_ps = equivariant_ps_calculator.compute( + systems, neighbors_to_properties=True +) -assert mts.equal(lambda_soap, equi_pow_spec) +assert mts.equal(lambda_soap, equivariant_ps) # %% # @@ -163,33 +155,39 @@ # max_angular_density``. # Initialize the SphericalExpansion calculator and compute the density -spex_calc = SphericalExpansion(**HYPER_PARAMETERS_SMALL) -density = spex_calc.compute(systems) +spherical_expansion = SphericalExpansion(**HYPER_PARAMETERS_SMALL) +density = spherical_expansion.compute(systems) density = density.keys_to_properties("neighbor_type") # Initialize DensityCorrelations calculator -dens_corr_calc = DensityCorrelations( +density_correlations = DensityCorrelations( n_correlations=2, max_angular=6, ) # Compute the bispectrum -bispectrum = dens_corr_calc.compute(density) +bispectrum = density_correlations.compute(density) # %% # -# Inspect the features of the bispectrum. There are now ``"neighbor_x_type"`` and -# ``"n_x"`` (which track the radial channel indices) dimensions created by the product -# of feature spaces of the 3 density blocks combined to make each bispectrum block. +# There are now ``"neighbor_x_type"`` and ``"n_x"`` (which track the radial channel +# indices) dimensions created by the product of feature spaces of the 3 density blocks +# combined to make each bispectrum block. # # For each block, its key contains dimensions tracking the angular order of blocks # combined to create it, namely ``["l_1", "l_2", "k_2", "l_3"]``. The ``"l_"`` # dimensions track the angular order of the blocks from the original density, while # ``"k_"`` dimensions track the angular order of intermediate blocks. + +print("bispectrum first block:", bispectrum[0]) + +# %% # -# Let's look at an example. Take the block indexed by key: -# -# ``LabelsEntry(o3_lambda=3, o3_sigma=1, center_type=7, l_1=1, l_2=2, k_2=1, l_3=2)`` +# Let's look at an example. Take the block at index 156: + +print(bispectrum.keys[156]) + +# %% # # This was created in the following way. # @@ -204,20 +202,21 @@ # 3]``, and in this case has been combined to create the output angular order # ``o3_lambda=3``. -print("bispectrum:", bispectrum) -print("bispectrum first block:", bispectrum[0]) # %% # -# As before, move these symmetry keys to properties and inspect the metadata and the -# mean size of the features. +# As before, we can move these symmetry keys to properties and inspect the metadata and +# the total size of the features. bispectrum = bispectrum.keys_to_properties(["l_1", "l_2", "k_2", "l_3"]) -print("bispectrum:", bispectrum) -print("bispectrum first block:", bispectrum[0]) +print("first block:", bispectrum[0]) + +# %% +# print( - "total number of features:", np.sum([len(block.properties) for block in bispectrum]) + "total number of features:", + np.sum([len(block.properties) for block in bispectrum]), ) # %% @@ -235,19 +234,26 @@ # Let's truncate to an angular cutoff of 3: angular_cutoff = 3 -dens_corr_calc = DensityCorrelations( +density_correlations = DensityCorrelations( n_correlations=2, max_angular=angular_cutoff, ) -bispectrum_truncated = dens_corr_calc.compute(density, angular_cutoff=angular_cutoff) +bispectrum_truncated = density_correlations.compute( + density, angular_cutoff=angular_cutoff +) # Move the "l_" and "k_" keys to properties bispectrum_truncated = bispectrum_truncated.keys_to_properties( ["l_1", "l_2", "k_2", "l_3"] ) -print("truncated bispectrum:", bispectrum_truncated) -print("truncated bispectrum first block:", bispectrum_truncated[0]) +print("truncated bispectrum:", bispectrum_truncated.keys) + +# %% + +print("first block:", bispectrum_truncated[0]) + +# %% print( "total number of features:", np.sum([len(block.properties) for block in bispectrum_truncated]), @@ -257,11 +263,11 @@ # # To compute a descriptor that matches the symmetry of a given target property, the # ``selected_keys`` argument can be passed to the ``compute`` method. This was also seen -# in the previous tutorial on :ref:`computing lambda-SOAP `. +# in the previous tutorial on :ref:`computing lambda-SOAP `. # # Following this example, to compute the truncated bispectrum for a polarizability # tensor: -bispectrum_truncated_key_select = dens_corr_calc.compute( +bispectrum_truncated_key_select = density_correlations.compute( density, angular_cutoff=angular_cutoff, selected_keys=mts.Labels( @@ -275,7 +281,9 @@ ["l_1", "l_2", "k_2", "l_3"] ) -print("truncated bispectrum with selected keys:", bispectrum_truncated_key_select) +print("truncated bispectrum with selected keys:", bispectrum_truncated_key_select.keys) + +# %% print( "total number of features:", np.sum([len(block.properties) for block in bispectrum_truncated_key_select]), @@ -291,10 +299,7 @@ # a spherical expansion of decorated atomic point cloud data. # # A key limitation of this approach is an exploding feature size. To reduce the number -# of output blocks, the ``angular_cutoff`` parameter can be used. This can be controlled -# by introducing contractions after each CG combination in the iteration. This will be -# covered in the next tutorial, :ref:`Clebsch-Gordan products -# `. +# of output blocks, the ``angular_cutoff`` parameter can be used. # %% # .. end-body