Add an optional callback to Lanczos

CMakeLists.txt

@@ -50,14 +50,6 @@ FetchContent_Declare(
   EXCLUDE_FROM_ALL
 )
 FetchContent_MakeAvailable(uammd)
-FetchContent_Declare(
-  lanczos
-  GIT_REPOSITORY https://github.com/RaulPPelaez/LanczosAlgorithm
-  GIT_TAG        v1.0.1
-  EXCLUDE_FROM_ALL
-)
-FetchContent_MakeAvailable(lanczos)
-include_directories(${lanczos_SOURCE_DIR}/include)
 
 set(BLA_VENDOR Generic)
 find_package(BLAS REQUIRED)

include/MobilityInterface/MobilityInterface.h

@@ -7,10 +7,10 @@
 #include "defines.h"
 #include "lanczos.h"
 #include "memory/container.h"
+#include <functional>
 #include <random>
 #include <stdexcept>
 #include <vector>
-
 namespace libmobility {
 
 enum class periodicity_mode {
@@ -28,6 +28,7 @@ struct Parameters {
   real tolerance = 1e-4; // Tolerance for Lanczos fluctuations
   std::uint64_t seed = 0;
   bool includeAngular = false;
+  std::function<void(int, float)> lanczosCallback;
 };
 
 // A list of parameters that cannot be changed by reinitializing a solver and/or
@@ -47,9 +48,10 @@ class Mobility {
   std::uint64_t lanczosSeed;
   real lanczosTolerance;
   std::shared_ptr<LanczosStochasticVelocities> lanczos;
-  std::vector<real> lanczosOutput;
+  thrust::device_vector<real> lanczosOutput;
   bool includeAngular = false;
   std::mt19937 rng;
+  std::function<void(int, real)> lanczosCallback;
 
 protected:
   Mobility() {};
@@ -96,6 +98,7 @@ class Mobility {
     this->initialized = true;
     this->lanczosSeed = this->rng();
     this->lanczosTolerance = par.tolerance;
+    this->lanczosCallback = par.lanczosCallback;
     this->includeAngular = par.includeAngular;
   }
 
@@ -128,8 +131,8 @@ class Mobility {
       throw std::runtime_error(
           "[libMobility] The number of particles is not set. Did you "
           "forget to call setPositions?");
-    device_adapter<real> linear(ilinear, device::cpu);
-    device_adapter<real> angular(iangular, device::cpu);
+    device_adapter<real> linear(ilinear, device::cuda);
+    device_adapter<real> angular(iangular, device::cuda);
     if (linear.empty())
       throw std::runtime_error(
           "[libMobility] This solver requires linear velocities");
@@ -145,8 +148,8 @@ class Mobility {
           this->lanczosTolerance, this->lanczosSeed);
     }
     lanczosOutput.resize(3 * numberElements);
-    std::fill(lanczosOutput.begin(), lanczosOutput.end(), 0);
-    auto dev = linear.dev;
+    thrust::fill(lanczosOutput.begin(), lanczosOutput.end(), 0);
+    auto dev = device::cuda;
     lanczos->sqrtMdotW(
         [this, dev, numberParticles](const real *f, real *mv) {
           // Torques are stored at the end of the force array
@@ -160,14 +163,15 @@ class Mobility {
           device_span<real> s_mv({mv, mv + 3 * N}, dev);
           Mdot(s_f, s_t, s_mv, s_mt);
         },
-        lanczosOutput.data(), numberElements, prefactor);
-    std::transform(lanczosOutput.begin(),
-                   lanczosOutput.begin() + 3 * numberParticles, linear.begin(),
-                   linear.begin(), thrust::plus<real>());
+        lanczosOutput.data().get(), numberElements, lanczosCallback, prefactor);
+    thrust::transform(thrust::cuda::par, lanczosOutput.begin(),
+                      lanczosOutput.begin() + 3 * numberParticles,
+                      linear.begin(), linear.begin(), thrust::plus<real>());
     if (this->includeAngular)
-      std::transform(lanczosOutput.begin() + 3 * numberParticles,
-                     lanczosOutput.end(), angular.begin(), angular.begin(),
-                     thrust::plus<real>());
+      thrust::transform(thrust::cuda::par,
+                        lanczosOutput.begin() + 3 * numberParticles,
+                        lanczosOutput.end(), angular.begin(), angular.begin(),
+                        thrust::plus<real>());
   }
 
   // computes velocities according to the Langevin equation.

include/MobilityInterface/lanczos.h

@@ -3,17 +3,31 @@
  */
 #ifndef LIBMOBILITY_LANCZOS_ADAPTOR_H
 #define LIBMOBILITY_LANCZOS_ADAPTOR_H
-#include "LanczosAlgorithm.h"
+#define CUDA_ENABLED
+#include "lanczos/LanczosAlgorithm.h"
+#include "third_party/saruprng.cuh"
 #include <algorithm>
 #include <cstdint>
 #include <random>
 #include <vector>
 
+namespace detail {
+using real = lanczos::real;
+struct SaruFill {
+  uint seed1, seed2;
+  __device__ real operator()(uint id) {
+    Saru prng(seed1, seed2, id);
+    return prng.gf(real(0), real(1.0)).x;
+  }
+};
+
+} // namespace detail
 // This class uses the LanczosAlgorithm library to compute fluctuations.
 class LanczosStochasticVelocities {
   using real = lanczos::real;
   lanczos::Solver lanczos;
-  std::vector<real> lanczosNoise;
+  // std::vector<real> lanczosNoise;
+  thrust::device_vector<real> lanczosNoise;
   real lanczosTolerance;
   std::mt19937 engine;
 
@@ -27,14 +41,20 @@ class LanczosStochasticVelocities {
   // dW). Where B is an operator that applies the square root of the provided
   // mobility.
   template <class MobilityDot>
-  void sqrtMdotW(MobilityDot dot, real *result, int numberParticles,
-                 real prefactor = 1) {
-    std::normal_distribution<real> dist{0, 1};
-    auto gen = [&]() { return dist(engine); };
+  void sqrtMdotW(MobilityDot idot, real *result, int numberParticles,
+                 std::function<void(int, float)> callback, real prefactor = 1) {
     lanczosNoise.resize(3 * numberParticles);
-    std::generate(lanczosNoise.begin(), lanczosNoise.end(), gen);
-    lanczos.run(dot, result, lanczosNoise.data(), lanczosTolerance,
-                3 * numberParticles);
+    // std::generate(lanczosNoise.begin(), lanczosNoise.end(), gen);
+    uint seed1 = std::uniform_int_distribution<uint>(0, UINT32_MAX)(engine);
+    uint seed2 = std::uniform_int_distribution<uint>(0, UINT32_MAX)(engine);
+    auto cit = thrust::make_counting_iterator<uint>(0);
+    thrust::transform(cit, cit + 3 * numberParticles, lanczosNoise.begin(),
+                      detail::SaruFill{seed1, seed2});
+    std::function<void(real *, real *)> dot = [&](real *f, real *mv) {
+      idot(f, mv);
+    };
+    lanczos.run(dot, result, lanczosNoise.data().get(), lanczosTolerance,
+                3 * numberParticles, callback);
   }
 };
 

include/MobilityInterface/pythonify.h

@@ -7,13 +7,17 @@ python (accompanied by the default documentation of the mobility interface.
 #ifndef MOBILITY_PYTHONIFY_H
 #include "MobilityInterface/MobilityInterface.h"
 #include "memory/python_tensor.h"
+#include <functional>
 #include <nanobind/nanobind.h>
 #include <nanobind/ndarray.h>
+#include <nanobind/stl/optional.h>
 #include <nanobind/stl/pair.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/unique_ptr.h>
 #include <nanobind/stl/vector.h>
+#include <optional>
 #include <stdexcept>
+
 namespace nb = nanobind;
 using namespace nb::literals;
 namespace py = nb;
@@ -159,6 +163,8 @@ includeAngular : bool, optional
 		Whether the solver will produce angular velocities. Needed if torques are given. Default is false.
 tolerance : float, optional
 		Tolerance, used for approximate methods and also for Lanczos (default fluctuation computation). Default is 1e-4.
+lanczos_callback : callable, optional
+		Callback function to be called during the Lanczos process. It should take two arguments: the current iteration (int) and the current error (float). Default is None, which means no callback will be used.
 )pbdoc";
 
 template <class Solver> auto call_sqrtMdotW(Solver &myself, real prefactor) {
@@ -234,12 +240,19 @@ array_like
 
 template <class Solver>
 void call_initialize(Solver &myself, real eta, real a, bool includeAngular,
-                     real tol) {
+                     real tol, std::optional<nb::callable> i_lanczosCallback) {
   libmobility::Parameters par;
   par.viscosity = eta;
   par.hydrodynamicRadius = {a};
   par.tolerance = tol;
   par.includeAngular = includeAngular;
+  std::function<void(int, real)> lanczosCallback;
+  if (i_lanczosCallback.has_value() && bool(i_lanczosCallback.value())) {
+    lanczosCallback = [i_lanczosCallback](int i, real err) {
+      i_lanczosCallback.value()(i, err);
+    };
+  }
+  par.lanczosCallback = lanczosCallback;
   myself.initialize(par);
 }
 
@@ -356,7 +369,7 @@ auto define_module_content(
            "periodicityY"_a, "periodicityZ"_a)
       .def("initialize", call_initialize<MODULENAME>, initialize_docstring,
            "viscosity"_a, "hydrodynamicRadius"_a, "includeAngular"_a = false,
-           "tolerance"_a = 1e-4)
+           "tolerance"_a = 1e-4, "lanczos_callback"_a = nb::none())
       .def("setPositions", call_setPositions<MODULENAME>,
            "The module will compute the mobility according to this set of "
            "positions.",