Add bond-atom spherical expansion, feature-gated.

[restructured contribution, commit 3]

python/rascaline/rascaline/__init__.py

@@ -12,6 +12,7 @@
     SortedDistances,
     SphericalExpansion,
     SphericalExpansionByPair,
+    SphericalExpansionForBonds,
 )
 from .log import set_logging_callback  # noqa
 from .profiling import Profiler  # noqa
@@ -29,4 +30,5 @@
     "SortedDistances",
     "SphericalExpansion",
     "SphericalExpansionByPair",
+    "SphericalExpansionForBonds",
 ]

python/rascaline/rascaline/calculators.py

@@ -201,6 +201,57 @@ def __init__(
         super().__init__("spherical_expansion_by_pair", json.dumps(parameters))
 
 
+class SphericalExpansionForBonds(CalculatorBase):
+    """A SOAP-like spherical expansion coefficients for bond-centered environments
+    In other words, the spherical expansion of the neighbor density function centered
+    on the center of a bond,
+    'after' rotating the system so that the bond is aligned with the z axis.
+
+    This is not rotationally invariant, and as such you should use some
+    not-implemented-here matheatical trick
+    similar to what SOAP (the :py:class:`SoapPowerSpectrum` class) uses.
+
+    Most hyperparameters are identical to that of the regulat spherical expansion:
+    :ref:`documentation <spherical-expansion>`.
+
+    the few changes to this are:
+
+    - "cutoff" renamed to "third_cutoff"
+    - "bond_cutoff" which expresses how the pairs of atoms used for the 'bonds' are
+      chosen.
+    - "center_atomS_weight" (caps only used for emphasis): the weight multiplier
+      for the coefficients of the self interactions
+      (where the neighboring atom is one of the pair's atoms).
+    """
+
+    def __init__(
+        self,
+        bond_cutoff,
+        third_cutoff,
+        max_radial,
+        max_angular,
+        atomic_gaussian_width,
+        center_atoms_weight,
+        radial_basis,
+        cutoff_function,
+        radial_scaling=None,
+    ):
+        parameters = {
+            "cutoffs": [bond_cutoff, third_cutoff],
+            "max_radial": max_radial,
+            "max_angular": max_angular,
+            "atomic_gaussian_width": atomic_gaussian_width,
+            "center_atoms_weight": center_atoms_weight,
+            "radial_basis": radial_basis,
+            "cutoff_function": cutoff_function,
+        }
+
+        if radial_scaling is not None:
+            parameters["radial_scaling"] = radial_scaling
+
+        super().__init__("spherical_expansion_for_bonds", json.dumps(parameters))
+
+
 class SoapRadialSpectrum(CalculatorBase):
     """Radial spectrum of Smooth Overlap of Atomic Positions (SOAP).
 

rascaline/src/calculator.rs

@@ -549,6 +549,7 @@ use crate::calculators::{SphericalExpansionByPair, SphericalExpansionParameters}
 use crate::calculators::SphericalExpansion;
 use crate::calculators::{SoapPowerSpectrum, PowerSpectrumParameters};
 use crate::calculators::{SoapRadialSpectrum, RadialSpectrumParameters};
+use crate::calculators::{SphericalExpansionForBonds, SphericalExpansionForBondsParameters};
 use crate::calculators::{LodeSphericalExpansion, LodeSphericalExpansionParameters};
 type CalculatorCreator = fn(&str) -> Result<Box<dyn CalculatorBase>, Error>;
 
@@ -581,6 +582,7 @@ static REGISTERED_CALCULATORS: Lazy<BTreeMap<&'static str, CalculatorCreator>> =
     add_calculator!(map, "spherical_expansion", SphericalExpansion, SphericalExpansionParameters);
     add_calculator!(map, "soap_radial_spectrum", SoapRadialSpectrum, RadialSpectrumParameters);
     add_calculator!(map, "soap_power_spectrum", SoapPowerSpectrum, PowerSpectrumParameters);
+    add_calculator!(map, "spherical_expansion_for_bonds", SphericalExpansionForBonds, SphericalExpansionForBondsParameters);
 
     add_calculator!(map, "lode_spherical_expansion", LodeSphericalExpansion, LodeSphericalExpansionParameters);
     return map;