Simplified MBTR calculation.

lauri-codes · lauri-codes · commit 7cc189bbb03e · 2024-11-19T22:40:04.000+02:00
diff --git a/dscribe/descriptors/acsf.py b/dscribe/descriptors/acsf.py
@@ -169,14 +169,16 @@ def is_static():
 
         return output
 
-    def create_single(self, system, centers=None):
+    def create_single(self, system, centers=None, return_descriptor=True, return_derivatives=False):
         """Creates the descriptor for the given system.
 
         Args:
             system (:class:`ase.Atoms` | :class:`.System`): Input system.
             centers (iterable): Indices of the atoms around which the ACSF
                 will be returned. If no centers defined, ACSF will be created
                 for all atoms in the system.
+            return_descriptor: Whether to return the descriptor
+            return_derivatives: Whether to return the derivatives
 
         Returns:
             np.ndarray: The ACSF output for the given system and centers. The
diff --git a/dscribe/descriptors/coulombmatrix.py b/dscribe/descriptors/coulombmatrix.py
@@ -116,10 +116,12 @@ def create(self, system, n_jobs=1, only_physical_cores=False, verbose=False):
 
         return output
 
-    def create_single(self, system):
+    def create_single(self, system, return_descriptor=True, return_derivatives=False):
         """
         Args:
             system (:class:`ase.Atoms`): Input system.
+            return_descriptor: Whether to return the descriptor
+            return_derivatives: Whether to return the derivatives
 
         Returns:
             ndarray: The zero padded matrix as a flattened 1D array.
@@ -140,6 +142,8 @@ def create_single(self, system):
             atomic_numbers,
             cell,
             pbc,
+            return_descriptor,
+            return_derivatives,
         )
 
         return out_des
diff --git a/dscribe/descriptors/descriptor.py b/dscribe/descriptors/descriptor.py
@@ -307,7 +307,7 @@ def create_multiple(arguments, func, is_sparse, index, verbose):
             n_samples = len(arguments)
 
             for i_sample, i_arg in enumerate(arguments):
-                i_out = func(*i_arg, True, False)
+                i_out = func(*i_arg, return_descriptor=True, return_derivatives=False)
                 i_out = self.format_array(i_out)
 
                 # If the shape varies, just add result into a list
diff --git a/dscribe/descriptors/descriptormatrix.py b/dscribe/descriptors/descriptormatrix.py
@@ -101,10 +101,12 @@ def get_matrix(self, system):
             np.ndarray: The final two-dimensional matrix for this descriptor.
         """
 
-    def create_single(self, system):
+    def create_single(self, system, return_descriptor=True, return_derivatives=False):
         """
         Args:
             system (:class:`ase.Atoms` | :class:`.System`): Input system.
+            return_descriptor: Whether to return the descriptor
+            return_derivatives: Whether to return the derivatives
 
         Returns:
             ndarray: The zero padded matrix either as a 1D array.
diff --git a/dscribe/descriptors/ewaldsummatrix.py b/dscribe/descriptors/ewaldsummatrix.py
@@ -177,7 +177,7 @@ def create(
 
         return output
 
-    def create_single(self, system, accuracy=1e-5, w=1, r_cut=None, g_cut=None, a=None):
+    def create_single(self, system, accuracy=1e-5, w=1, r_cut=None, g_cut=None, a=None, return_descriptor=True, return_derivaties=False):
         """
         Args:
             system (:class:`ase.Atoms` | :class:`.System`): Input system.
@@ -198,6 +198,8 @@ def create_single(self, system, accuracy=1e-5, w=1, r_cut=None, g_cut=None, a=No
                 Gaussians. If not provided, a default value of :math:`\\alpha =
                 \\sqrt{\\pi}\\left(\\frac{N}{V^2}\\right)^{1/6}` is used.
                 Corresponds to the standard deviation of the Gaussians.
+            return_descriptor: Whether to return the descriptor
+            return_derivatives: Whether to return the derivatives
         """
         self.q = system.get_atomic_numbers()
         self.q_squared = self.q**2
diff --git a/dscribe/descriptors/lmbtr.py b/dscribe/descriptors/lmbtr.py
@@ -374,6 +374,8 @@ def create_single(
         self,
         system,
         centers=None,
+        return_descriptor=True,
+        return_derivaties=False
     ):
         """Return the local many-body tensor representation for the given
         system and centers.
@@ -386,6 +388,8 @@ def create_single(
                 If cartesian positions are provided, new atoms are added at that
                 position. If no centers are provided, all atoms in the system
                 will be used as centers.
+            return_descriptor: Whether to return the descriptor
+            return_derivatives: Whether to return the derivatives
 
         Returns:
             1D ndarray: The local many-body tensor representations of given
diff --git a/dscribe/descriptors/mbtr.py b/dscribe/descriptors/mbtr.py
@@ -337,11 +337,13 @@ def create(self, system, n_jobs=1, only_physical_cores=False, verbose=False):
 
         return output
 
-    def create_single(self, system, return_descriptor, return_derivatives):
+    def create_single(self, system, return_descriptor=True, return_derivatives=False):
         """Return the many-body tensor representation for the given system.
 
         Args:
             system (:class:`ase.Atoms`): Input system.
+            return_descriptor: Whether to return the descriptor
+            return_derivatives: Whether to return the derivatives
 
         Returns:
             np.ndarray | sparse.COO: The return type is
diff --git a/dscribe/descriptors/soap.py b/dscribe/descriptors/soap.py
@@ -486,7 +486,7 @@ def is_static():
 
         return output
 
-    def create_single(self, system, centers=None):
+    def create_single(self, system, centers=None, return_descriptor=True, return_derivatives=False):
         """Return the SOAP output for the given system and given centers.
 
         Args:
@@ -495,6 +495,8 @@ def create_single(self, system, centers=None):
                 specified, the SOAP spectrum will be created for these points.
                 If no centers are defined, the SOAP output will be created
                 for all atoms in the system.
+            return_descriptor: Whether to return the descriptor
+            return_derivatives: Whether to return the derivatives
 
         Returns:
             np.ndarray | sparse.COO: The SOAP output for the
diff --git a/dscribe/ext/coulombmatrix.cpp b/dscribe/ext/coulombmatrix.cpp
@@ -39,18 +39,16 @@ CoulombMatrix::CoulombMatrix(
 
 void CoulombMatrix::create(
     py::array_t<double> out, 
-    py::array_t<double> positions,
-    py::array_t<int> atomic_numbers,
-    py::array_t<double> cell,
+    System &system,
     CellList cell_list,
     bool return_descriptor,
     bool return_derivatives
 )
 {
     // Calculate all pairwise distances.
     auto out_mu = out.mutable_unchecked<1>();
-    auto atomic_numbers_u = atomic_numbers.unchecked<1>();
-    auto positions_u = positions.unchecked<2>();
+    auto atomic_numbers_u = system.atomic_numbers.unchecked<1>();
+    auto positions_u = system.positions.unchecked<2>();
     int n_atoms = atomic_numbers_u.shape(0);
     MatrixXd matrix = distancesEigen(positions_u);
 
diff --git a/dscribe/ext/coulombmatrix.h b/dscribe/ext/coulombmatrix.h
@@ -46,9 +46,7 @@ class CoulombMatrix: public DescriptorGlobal {
          */
         void create(
             py::array_t<double> out, 
-            py::array_t<double> positions,
-            py::array_t<int> atomic_numbers,
-            py::array_t<double> cell,
+            System &system,
             CellList cell_list,
             bool return_descriptor,
             bool return_derivatives
diff --git a/dscribe/ext/descriptorglobal.cpp b/dscribe/ext/descriptorglobal.cpp
@@ -39,26 +39,22 @@ void DescriptorGlobal::create(
     // Extend system if periodicity is requested.
     auto pbc_u = pbc.unchecked<1>();
     bool is_periodic = this->periodic && (pbc_u(0) || pbc_u(1) || pbc_u(2));
-    if (is_periodic) {
-        System system_extended = extend_system(positions, atomic_numbers, cell, pbc, this->cutoff);
-        positions = system_extended.positions;
-        atomic_numbers = system_extended.atomic_numbers;
-    }
-    this->create(out, positions, atomic_numbers, cell, return_descriptor, return_derivatives);
+    System system = is_periodic
+        ? extend_system(positions, atomic_numbers, cell, pbc, this->cutoff)
+        : System(positions, atomic_numbers, cell);
+    this->create(out, system, return_descriptor, return_derivatives);
 }
 
 void DescriptorGlobal::create(
     py::array_t<double> out, 
-    py::array_t<double> positions,
-    py::array_t<int> atomic_numbers,
-    py::array_t<double> cell,
+    System &system,
     bool return_descriptor,
     bool return_derivatives
 )
 {
     // Calculate neighbours with a cell list
-    CellList cell_list(positions, this->cutoff);
-    this->create(out, positions, atomic_numbers, cell, cell_list, return_descriptor, return_derivatives);
+    CellList cell_list(system.positions, this->cutoff);
+    this->create(out, system, cell_list, return_descriptor, return_derivatives);
 }
 
 void DescriptorGlobal::derivatives_numerical(
@@ -72,7 +68,6 @@ void DescriptorGlobal::derivatives_numerical(
     bool return_descriptor
 )
 {
-    int n_copies = 1;
     int n_atoms = atomic_numbers.size();
     int n_features = this->get_number_of_features();
     auto derivatives_mu = derivatives.mutable_unchecked<3>();
@@ -81,20 +76,18 @@ void DescriptorGlobal::derivatives_numerical(
 
     // Extend the system if it is periodic
     bool is_periodic = this->periodic && (pbc_u(0) || pbc_u(1) || pbc_u(2));
-    if (is_periodic) {
-        System system_extension = extend_system(positions, atomic_numbers, cell, pbc, this->cutoff);
-        n_copies = system_extension.atomic_numbers.size()/atomic_numbers.size();
-        positions = system_extension.positions;
-        atomic_numbers = system_extension.atomic_numbers;
-    }
-    auto positions_mu = positions.mutable_unchecked<2>();
+    System system = is_periodic
+      ? extend_system(positions, atomic_numbers, cell, pbc, this->cutoff)
+      : System(positions, atomic_numbers, cell);
+    int n_copies = system.atomic_numbers.size() / atomic_numbers.size();
+    auto positions_mu = system.positions.mutable_unchecked<2>();
 
     // Pre-calculate cell list for atoms
-    CellList cell_list_atoms(positions, this->cutoff);
+    CellList cell_list_atoms(system.positions, this->cutoff);
 
     // Calculate the desciptor value if requested
     if (return_descriptor) {
-        this->create(descriptor, positions, atomic_numbers, cell, cell_list_atoms, true, false);
+        this->create(descriptor, system, cell_list_atoms, true, false);
     }
 
     // Central finite difference with error O(h^2)
@@ -141,7 +134,7 @@ void DescriptorGlobal::derivatives_numerical(
                 auto d_mu = d.mutable_unchecked<1>();
 
                 // Calculate descriptor value
-                this->create(d, positions, atomic_numbers, cell, cell_list_atoms, true, false);
+                this->create(d, system, cell_list_atoms, true, false);
 
                 // Add value to final derivative array
                 double coeff = coefficients[i_stencil];
diff --git a/dscribe/ext/descriptorglobal.h b/dscribe/ext/descriptorglobal.h
@@ -20,6 +20,7 @@ limitations under the License.
 #include <string>
 #include "descriptor.h"
 #include "celllist.h"
+#include "geometry.h"
 
 namespace py = pybind11;
 using namespace std;
@@ -58,9 +59,7 @@ class DescriptorGlobal : public Descriptor {
          */
         void create(
             py::array_t<double> out, 
-            py::array_t<double> positions,
-            py::array_t<int> atomic_numbers,
-            py::array_t<double> cell,
+            System &system,
             bool return_descriptor,
             bool return_derivatives
         );
@@ -75,9 +74,7 @@ class DescriptorGlobal : public Descriptor {
          */
         virtual void create(
             py::array_t<double> out, 
-            py::array_t<double> positions,
-            py::array_t<int> atomic_numbers,
-            py::array_t<double> cell,
+            System &system,
             CellList cell_list,
             bool return_descriptor,
             bool return_derivatives
diff --git a/dscribe/ext/geometry.cpp b/dscribe/ext/geometry.cpp
@@ -114,35 +114,6 @@ System::System(
     this->indices = indices;
 }
 
-System::System(
-    py::array_t<double> positions,
-    py::array_t<int> atomic_numbers,
-    py::array_t<double> cell,
-    py::array_t<int> indices
-)
-    : positions(positions)
-    , atomic_numbers(atomic_numbers)
-    , cell(cell)
-    , indices(indices)
-{
-    // Create the default set of interactive atoms, which encompasses the whole
-    // system
-    unordered_set<int> interactive_atoms = unordered_set<int>();
-    int n_atoms = atomic_numbers.size();
-    for (int i = 0; i < n_atoms; ++i) {
-        interactive_atoms.insert(i);
-    }
-    this->interactive_atoms = interactive_atoms;
-
-    // Create the default cell indices
-    py::array_t<int> cell_indices({n_atoms});
-    auto cell_indices_mu = cell_indices.mutable_unchecked<1>();
-    for (int i = 0; i < n_atoms; ++i) {
-        cell_indices_mu(i) = 0;
-    }
-    this->cell_indices = cell_indices;
-}
-
 System::System(
     py::array_t<double> positions,
     py::array_t<int> atomic_numbers,
@@ -223,15 +194,17 @@ System extend_system(
         }
     }
 
-    // Calculate the extended system positions.
+    // Calculate the extended system properties
     int n_rep = (2*n_copies_axis[0]+1)*(2*n_copies_axis[1]+1)*(2*n_copies_axis[2]+1);
     int n_atoms = atomic_numbers.size();
     py::array_t<double> ext_pos({n_atoms*n_rep, 3});
     py::array_t<int> ext_atomic_numbers({n_atoms*n_rep});
     py::array_t<int> ext_indices({n_atoms*n_rep});
+    py::array_t<int> cell_indices({n_atoms*n_rep});
     auto ext_pos_mu = ext_pos.mutable_unchecked<2>();
     auto ext_atomic_numbers_mu = ext_atomic_numbers.mutable_unchecked<1>();
     auto ext_indices_mu = ext_indices.mutable_unchecked<1>();
+    auto cell_indices_mu = cell_indices.mutable_unchecked<1>();
     int i_copy = 0;
     int a_limit = multipliers[0].size();
     int b_limit = multipliers[1].size();
@@ -257,13 +230,20 @@ System extend_system(
                     for (int m=0; m < 3; ++m) {
                         ext_pos_mu(index, m) = positions_u(l, m) + addition[m];
                     }
+                    cell_indices_mu(index) = i_copy;
                 }
                 ++i_copy;
             }
         }
     }
 
-    return System(ext_pos, ext_atomic_numbers, cell, ext_indices);
+    // Store original set of interactive atoms
+    unordered_set<int> interactive_atoms = unordered_set<int>();
+    for (int i = 0; i < n_atoms; ++i) {
+        interactive_atoms.insert(i);
+    }
+
+    return System(ext_pos, ext_atomic_numbers, cell, ext_indices, cell_indices, interactive_atoms);
 }
 
 py::array_t<double> distancesNumpy(py::detail::unchecked_reference<double, 2> &positions_u)
diff --git a/dscribe/ext/geometry.h b/dscribe/ext/geometry.h
@@ -33,12 +33,6 @@ class System {
             py::array_t<double> cell,
             bool extra = true
         );
-        System(
-            py::array_t<double> positions,
-            py::array_t<int> atomic_numbers,
-            py::array_t<double> cell,
-            py::array_t<int> indices
-        );
         System(
             py::array_t<double> positions,
             py::array_t<int> atomic_numbers,
diff --git a/dscribe/ext/mbtr.cpp b/dscribe/ext/mbtr.cpp
diff --git a/dscribe/ext/mbtr.h b/dscribe/ext/mbtr.h
diff --git a/dscribe/ext/mbtr2.cpp b/dscribe/ext/mbtr2.cpp
diff --git a/dscribe/ext/mbtr2.h b/dscribe/ext/mbtr2.h
diff --git a/tests/conftest.py b/tests/conftest.py
diff --git a/tests/test_mbtr.py b/tests/test_mbtr.py
