Atomic/local stress

bindings/bind_py_models.cc

@@ -82,6 +82,18 @@ namespace rascal {
             SparsePoints, i.e. the basis used by the sparse method.
             The representation of the atomic structures computed with Calculator
             should have already been computed.)");
+    kernel.def(
+        "compute_local",
+        py::overload_cast<const Calculator &, const StructureManagers &,
+                          const SparsePoints &>(
+            &SparseKernel::template compute_local<Calculator, StructureManagers,
+                                            SparsePoints>),
+        py::call_guard<py::gil_scoped_release>(),
+        R"(Compute the sparse kernel between the representation of a set of
+            atomic structures, i.e. StructureManagerCollections, and a set of
+            SparsePoints, i.e. the basis used by the sparse method.
+            The representation of the atomic structures computed with Calculator
+            should have already been computed.)");
     kernel.def("compute",
                py::overload_cast<const SparsePoints &>(
                    &SparseKernel::template compute<SparsePoints>),
@@ -193,6 +205,43 @@ namespace rascal {
           return neg_stress_global;
         },
         py::call_guard<py::gil_scoped_release>());
+
+    mod.def(
+        "compute_sparse_kernel_local_neg_stress",
+        [](const Calculator & calculator, SparseKernel & kernel,
+           ManagerCollection & managers, SparsePoints & sparse_points,
+           math::Vector_t & weights) {
+          std::string neg_stress_name = compute_sparse_kernel_local_neg_stress(
+              calculator, kernel, managers, sparse_points, weights);
+          size_t nb_centers_over_managers{0};
+          for (const auto & manager : managers) {
+
+            nb_centers_over_managers+=manager->size();
+            //for (auto center : manager) {
+            //  nb_centers_over_managers++;
+            //}
+          }
+          math::Matrix_t local_neg_stress{nb_centers_over_managers, 9};
+
+          size_t i_center{0};
+          size_t i_manager_nb_center{0};
+          for (const auto & manager : managers) {
+            i_manager_nb_center=manager->size();
+            // i_manager_nb_center = 0;
+            // for (auto center : manager) {
+            //   i_manager_nb_center++;
+            // }
+            auto && neg_stress{
+                *manager
+                    ->template get_property<Property<double, 1, Manager_t, 9>>(
+                        neg_stress_name, true, true, true)};
+            local_neg_stress.block(i_center, 0, i_manager_nb_center, 9) =
+                Eigen::Map<const math::Matrix_t>(neg_stress.view().data(), 1, 9);
+            i_center += i_manager_nb_center;
+          }
+          return local_neg_stress;
+        },
+        py::call_guard<py::gil_scoped_release>());
   }
 
   /**

bindings/rascal/lib/__init__.py

@@ -4,4 +4,5 @@
     kernels,
     compute_sparse_kernel_gradients,
     compute_sparse_kernel_neg_stress,
+    compute_sparse_kernel_local_neg_stress
 )

bindings/rascal/models/krr.py

@@ -8,7 +8,11 @@
     train_gap_model         Train a GAP model given a kernel matrix and sparse points
 """
 from ..utils import BaseIO
-from ..lib import compute_sparse_kernel_gradients, compute_sparse_kernel_neg_stress
+from ..lib import (
+    compute_sparse_kernel_gradients,
+    compute_sparse_kernel_neg_stress,
+    compute_sparse_kernel_local_neg_stress
+)
 
 import scipy
 import numpy as np
@@ -352,7 +356,7 @@ def predict_forces(self, managers, KNM=None):
 
         return -gradients
 
-    def predict_stress(self, managers, KNM=None):
+    def predict_stress(self, managers, KNM=None, local_stress=False):
         """Predict gradients w.r.t cell parameters, e.g. stress, associated with the atomic structures in managers.
         The stress is returned using the Voigt order: xx, yy, zz, yz, xz, xy.
 
@@ -376,14 +380,26 @@ def predict_stress(self, managers, KNM=None):
 
         if KNM is None:
             rep = self.kernel._representation
-            neg_stress = compute_sparse_kernel_neg_stress(
-                rep,
-                self.kernel._kernel,
-                managers.managers,
-                self.X_train._sparse_points,
-                self.weights.reshape((1, -1)),
-            )
+            if local_stress:
+                neg_stress = compute_sparse_kernel_local_neg_stress(
+                    rep,
+                    self.kernel._kernel,
+                    managers.managers,
+                    self.X_train._sparse_points,
+                    self.weights.reshape((1, -1)),
+                )
+            else:
+                neg_stress = compute_sparse_kernel_neg_stress(
+                    rep,
+                    self.kernel._kernel,
+                    managers.managers,
+                    self.X_train._sparse_points,
+                    self.weights.reshape((1, -1)),
+                )
         else:
+            if local_stress:
+                raise ValueError("Option local_stress does not work with KNM."
+                                 " Please set KNM to None.")
             if 6 * len(managers) != KNM.shape[0]:
                 raise ValueError(
                     "KNM size mismatch {}!={}".format(6 * len(managers), KNM.shape[0])

src/rascal/models/sparse_kernel_predict.hh

@@ -374,5 +374,115 @@ namespace rascal {
     }  // manager
     return neg_stress_name;
   }
+
+
+  /**
+   * Compute the atomic gradients of a structure w.r.t atomic cell paramaters
+   * using a sparse GPR model. Only SOAP-GAP model is implemented at the moment
+   * (see
+   * @ref SparseKernelImpl<internal::SparseKernelType::GAP>::compute_derivative)
+   *
+   * The point of this function is to provide a faster prediction routine
+   * compared to computing the kernel elements and then multiplying them with
+   * the weights of the model.
+   *
+   * The gradients are attached to the input manager in a Property of
+   * Order 1 with the name
+   * '"kernel: "+kernel_type+" ; "+representation_grad_name+" negative stress;
+   * weight_hash:"+weight_hash' gradients [2*n_managers, 3]
+   *
+   * @tparam StructureManagers should be an iterable over shared pointer
+   *          of structure managers like ManagerCollection
+   * @param sparse_points a SparsePoints* class
+   * @param managers a ManagerCollection or similar collection of
+   * structure managers
+   * @param weights regression weights of the sparse GPR model
+   * @return name used to register the gradients in the managers
+   */
+  template <class Calculator, class StructureManagers, class SparsePoints>
+  std::string compute_sparse_kernel_local_neg_stress(const Calculator & calculator,
+                                               SparseKernel & kernel,
+                                               StructureManagers & managers,
+                                               SparsePoints & sparse_points,
+                                               math::Vector_t & weights) {
+    using Manager_t = typename StructureManagers::Manager_t;
+    using Property_t = typename Calculator::template Property_t<Manager_t>;
+    using PropertyGradient_t =
+        typename Calculator::template PropertyGradient_t<Manager_t>;
+    auto && representation_name{calculator.get_name()};
+    const auto representation_grad_name{calculator.get_gradient_name()};
+    size_t n_centers{0};
+    // find the total number of gradients
+    for (const auto & manager : managers) {
+      n_centers += manager->size();
+    }
+
+    internal::Hash<math::Vector_t, double> hasher{};
+    auto kernel_type_str = kernel.parameters.at("name").get<std::string>();
+    std::string weight_hash = std::to_string(hasher(weights));
+    std::string pair_grad_atom_i_r_j_name =
+        std::string("kernel: ") + kernel_type_str + std::string(" ; ") +
+        representation_grad_name +
+        std::string(" partial gradients; weight_hash:") + weight_hash;
+    std::string neg_stress_name =
+        std::string("kernel: ") + kernel_type_str + std::string(" ; ") +
+        representation_grad_name +
+        std::string(" negative local stress; weight_hash:") + weight_hash;
+
+    for (const auto & manager : managers) {
+      if (kernel_type_str == "GAP") {
+        const auto zeta = kernel.parameters.at("zeta").get<size_t>();
+
+        compute_partial_gradients_gap<Property_t, PropertyGradient_t>(
+            manager, sparse_points, weights, zeta, representation_name,
+            representation_grad_name, pair_grad_atom_i_r_j_name);
+      }
+
+      //auto && expansions_coefficients{*manager->template get_property<prop_t>(
+      //    this->get_name(), true, true, ExcludeGhosts)};
+      auto && neg_stress{
+          *manager->template get_property<Property<double, 1, Manager_t, 9>>(
+              neg_stress_name, true, true, true)};
+      if (neg_stress.is_updated()) {
+        continue;
+      }
+
+      neg_stress.resize();
+      neg_stress.setZero();
+
+      auto && pair_grad_atom_i_r_j{*manager->template get_property<
+          Property<double, 2, Manager_t, 1, ThreeD>>(pair_grad_atom_i_r_j_name,
+                                                     true)};
+
+      auto manager_root = extract_underlying_manager<0>(manager);
+      json structure_copy = manager_root->get_atomic_structure();
+      auto atomic_structure =
+          structure_copy.template get<AtomicStructure<ThreeD>>();
+      float volume = atomic_structure.get_volume();
+      for (auto center : manager) {
+        //auto && local_neg_stress = neg_stress[center];
+        //local_neg_stress.setZero();
+
+        Eigen::Vector3d r_i = center.get_position();
+        // accumulate partial gradients onto gradients
+        for (auto neigh : center.pairs_with_self_pair()) {
+          auto && local_neg_stress = neg_stress[neigh.get_atom_j()];
+          //local_neg_stress.setZero();
+          Eigen::Vector3d r_ji = r_i - neigh.get_position();
+          for (int a_der{0}; a_der < ThreeD; a_der++) {
+            for (int b_der{0}; b_der < ThreeD; b_der++) {
+              local_neg_stress(a_der*3 + b_der) +=
+                  r_ji(a_der) * pair_grad_atom_i_r_j[neigh](b_der)/volume;
+            }
+          }
+        }
+        // local_neg_stress /= volume;
+      }
+      neg_stress.set_updated_status(true);
+    }  // manager
+    return neg_stress_name;
+  }
+
+
 }  // namespace rascal
 #endif  // SRC_RASCAL_MODELS_SPARSE_KERNEL_PREDICT_HH_