Make the different functions return a numpy array instead of requirin…

…g it as input (#15)

* Update CMakeLists.txt

* Move examples to a new folder.
Change the location of some installation headers

* Update CI

* Update READMEs

* Update README

* Add option for cuda archs

* Update Python interface

* Add a function to get the number of particles

* Change the interface so that Mdot returns an array

* Update docs

* Port all functions to return results.
Add some tests for the interface
Do not call Mdot if forces is null

* Update env

* Relax python dependency

* Add LICENSE

* Update test

* Fix test

* Make fluctuation test more robust

CMakeLists.txt

@@ -3,15 +3,16 @@ project(libMobility)
 enable_language(CUDA)
 set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_SOURCE_DIR}/cmake/" "$ENV{CONDA_PREFIX}/share/cmake/Modules")
 set(CMAKE_BUILD_TYPE Release)
-
 add_compile_definitions(PUBLIC MAXLOGLEVEL=3)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CXX_STANDARD 14)
 set(CMAKE_CUDA_STANDARD 14)
 set(CMAKE_CUDA_STANDARD_REQUIRED ON)
 set(CMAKE_CUDA_SEPARABLE_COMPILATION OFF)
 # Set CUDA archs so all supported GPUs are covered
-set(CMAKE_CUDA_ARCHITECTURES "all")
+if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+  set(CMAKE_CUDA_ARCHITECTURES "all")
+endif()
 
 #add_compile_definitions(PUBLIC DOUBLE_PRECISION)
 set(UAMMD_INCLUDE include/third_party/uammd/src include/third_party/uammd/src/third_party)

LICENSE

@@ -0,0 +1,7 @@
+Copyright 2024 StochasticHydroTools
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

docs/source/usage.rst

@@ -62,7 +62,6 @@ A libMobility solver is initialized in three steps:
 	  precision = np.float32 if Solver.precision=="float" else np.float64
 	  pos = np.random.rand(3*numberParticles).astype(precision)
 	  force = np.ones(3*numberParticles).astype(precision)
-	  result = np.zeros(3*numberParticles).astype(precision)
 
           # The solver will fail if it is not compatible with the provided periodicity
 	  nb = Solver(periodicityX='open',periodicityY='open',periodicityZ='open')
@@ -72,11 +71,11 @@ A libMobility solver is initialized in three steps:
                         hydrodynamicRadius = 1.0,
                         numberParticles = numberParticles)
 	  nb.setPositions(pos)
-	  nb.Mdot(forces = force, result = result)
+	  result = nb.Mdot(forces = force)
 	  print(f"{numberParticles} particles located at ( X Y Z ): {pos}")
 	  print("Forces:", force)
 	  print("M*F:", result)
 	  # result = prefactor*sqrt(2*temperature*M)*dW
-	  nb.sqrtMdotW(prefactor = 1.0, result = result)
+	  result = nb.sqrtMdotW(prefactor = 1.0)
 	  print("sqrt(2*T*M)*N(0,1):", result)
 	  nb.clean()

environment.yml

@@ -2,7 +2,7 @@ name: libmobility
 channels:
   - conda-forge
 dependencies:
-  - cmake >=3.22
+  - cmake >=3.24
   - cuda-version 12.*
   - gxx
   - cuda-libraries-dev
@@ -11,6 +11,5 @@ dependencies:
   - mkl-devel
   - pybind11
   - numpy
-  - python 3.11.*
   - pytest
   - scipy

include/MobilityInterface/MobilityInterface.h

@@ -107,10 +107,16 @@ namespace libmobility{
 
     //Equivalent to calling Mdot and then stochasticDisplacements, can be faster in some solvers
     virtual void hydrodynamicVelocities(const real* forces, real* result, real prefactor = 1){
-      Mdot(forces, result);
+      if(forces){
+	Mdot(forces, result);
+      }
       sqrtMdotW(result, prefactor);
     }
 
+    int getNumberParticles() const{
+      return this->numberParticles;
+    }
+
     //Clean any memory allocated by the solver
     virtual void clean(){
       lanczos.reset();

include/MobilityInterface/pythonify.h

@@ -1,56 +1,69 @@
 /*Raul P. Pelaez 2021.
-The MOBILITY_PYTHONIFY(className, description) macro creates a pybind11 module from a class (called className) that inherits from libmobility::Mobility. "description" is a string that will be printed when calling help(className) from python (accompanied by the default documentation of the mobility interface.
+The MOBILITY_PYTHONIFY(className, description) macro creates a pybind11 module
+from a class (called className) that inherits from libmobility::Mobility.
+"description" is a string that will be printed when calling help(className) from
+python (accompanied by the default documentation of the mobility interface.
  */
 #include <stdexcept>
 #ifndef MOBILITY_PYTHONIFY_H
-#include"MobilityInterface.h"
-#include<pybind11/pybind11.h>
-#include<pybind11/numpy.h>
+#include "MobilityInterface.h"
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
 namespace py = pybind11;
 using namespace pybind11::literals;
 using pyarray = py::array;
+using pyarray_c =
+    py::array_t<libmobility::real, py::array::c_style | py::array::forcecast>;
 
 #define MOBILITYSTR(s) xMOBILITYSTR(s)
 #define xMOBILITYSTR(s) #s
 
-inline auto string2Periodicity(std::string per){
+inline auto string2Periodicity(std::string per) {
   using libmobility::periodicity_mode;
-  if(per == "open") return periodicity_mode::open;
-  else if(per == "unspecified") return periodicity_mode::unspecified;
-  else if(per == "single_wall") return periodicity_mode::single_wall;
-  else if(per == "two_walls") return periodicity_mode::two_walls;
-  else if(per == "periodic") return periodicity_mode::periodic;
-  else throw std::runtime_error("[libMobility] Invalid periodicity");
+  if (per == "open")
+    return periodicity_mode::open;
+  else if (per == "unspecified")
+    return periodicity_mode::unspecified;
+  else if (per == "single_wall")
+    return periodicity_mode::single_wall;
+  else if (per == "two_walls")
+    return periodicity_mode::two_walls;
+  else if (per == "periodic")
+    return periodicity_mode::periodic;
+  else
+    throw std::runtime_error("[libMobility] Invalid periodicity");
 }
 
-inline auto createConfiguration(std::string perx, std::string pery, std::string perz){
+inline auto createConfiguration(std::string perx, std::string pery,
+                                std::string perz) {
   libmobility::Configuration conf;
   conf.periodicityX = string2Periodicity(perx);
   conf.periodicityY = string2Periodicity(pery);
   conf.periodicityZ = string2Periodicity(perz);
   return conf;
 }
 
-template<typename T>
-void check_dtype(pyarray &arr){
-  if(not py::isinstance<py::array_t<T>>(arr)){
+template <typename T, class Array> void check_dtype(Array &arr) {
+  if (not py::isinstance<py::array_t<T>>(arr)) {
     throw std::runtime_error("Input array must have the correct data type.");
   }
-  if(not py::isinstance<py::array_t<T, py::array::c_style | py::array::forcecast>>(arr)){
-    throw std::runtime_error("The input array is not contiguous and cannot be used as a buffer.");
+  if (not py::isinstance<
+          py::array_t<T, py::array::c_style | py::array::forcecast>>(arr)) {
+    throw std::runtime_error(
+        "The input array is not contiguous and cannot be used as a buffer.");
   }
-
 }
 
-libmobility::real* cast_to_real(pyarray &arr){
+template <class Array> libmobility::real *cast_to_real(Array &arr) {
   check_dtype<libmobility::real>(arr);
-  return static_cast<libmobility::real*>(arr.mutable_data());
+  return static_cast<libmobility::real *>(arr.mutable_data());
 }
 
-const libmobility::real* cast_to_const_real(pyarray &arr){
+template <class Array> const libmobility::real *cast_to_const_real(Array &arr) {
   check_dtype<const libmobility::real>(arr);
-  return static_cast<const libmobility::real*>(arr.data());
+  return static_cast<const libmobility::real *>(arr.data());
 }
+
 const char *constructor_docstring = R"pbdoc(
 Initialize the module with a given set of periodicity conditions.
 
@@ -88,39 +101,90 @@ tolerance : float, optional
 		Tolerance, used for approximate methods and also for Lanczos (default fluctuation computation). Default is 1e-4.
 )pbdoc";
 
-const char *mdot_docstring = R"pbdoc(
-Computes the product of the Mobility matrix with a group of forces, :math:`\boldsymbol{\mathcal{M}}\boldsymbol{F}`.
+template <class Solver>
+auto call_sqrtMdotW(Solver &solver, libmobility::real prefactor) {
+  auto result =
+      py::array_t<libmobility::real>({solver.getNumberParticles() * 3});
+  solver.sqrtMdotW(cast_to_real(result), prefactor);
+  return result.reshape({solver.getNumberParticles(), 3});
+}
+
+const char *sqrtMdotW_docstring = R"pbdoc(
+Computes the stochastic contribution, :math:`\text{prefactor}\sqrt{2T\boldsymbol{\mathcal{M}}}d\boldsymbol{W}`, where :math:`\boldsymbol{\mathcal{M}}` is the mobility matrix and :math:`d\boldsymbol{W}` is a Wiener process.
 
 It is required that :py:mod:`setPositions` has been called before calling this function.
-Both inputs must have precision given by the precision attribute of the module.
-Both inputs must have size 3*N, where N is the number of particles.
-The arrays are ordered as :code:`[f0x, f0y, f0z, f1x, f1y, f1z, ...]`.
 
 Parameters
 ----------
-forces : array_like,
-		Forces acting on the particles.
-result : array_like,
-		Where the result will be stored.
+
+prefactor : float, optional
+		Prefactor to multiply the result by. Default is 1.0.
+
+Returns
+-------
+array_like
+		The resulting fluctuations. Shape is (N, 3), where N is the number of particles.
+
 )pbdoc";
 
-const char *sqrtMdotW_docstring = R"pbdoc(
-Computes the stochastic contribution, :math:`\text{prefactor}\sqrt{2T\boldsymbol{\mathcal{M}}}d\boldsymbol{W}`, where :math:`\boldsymbol{\mathcal{M}}` is the mobility matrix and :math:`d\boldsymbol{W}` is a Wiener process.
+template <class Solver> auto call_mdot(Solver &myself, pyarray_c &forces) {
+  int N = myself.getNumberParticles();
+  if (forces.size() < 3 * N and forces.size() > 0) {
+    throw std::runtime_error("The forces array must have size 3*N.");
+  }
+  auto f = forces.size() ? cast_to_const_real(forces) : nullptr;
+  auto result =
+      py::array_t<libmobility::real>(py::array::ShapeContainer({3 * N}));
+  result.attr("fill")(0);
+  myself.Mdot(f, cast_to_real(result));
+  return result.reshape({N, 3});
+}
+
+const char *mdot_docstring = R"pbdoc(
+Computes the product of the Mobility matrix with a group of forces, :math:`\boldsymbol{\mathcal{M}}\boldsymbol{F}`.
 
 It is required that :py:mod:`setPositions` has been called before calling this function.
-Both inputs must have precision given by the precision attribute of the module.
-Both inputs must have size 3*N, where N is the number of particles.
-The arrays are ordered as :code:`[f0x, f0y, f0z, f1x, f1y, f1z, ...]`.
 
 Parameters
 ----------
-result : array_like,
-		Where the result will be stored. The result will have the same format as the forces array.
-prefactor : float, optional
-		Prefactor to multiply the result by. Default is 1.0.
+forces : array_like,
+		Forces acting on the particles. Must have shape (N, 3), where N is the number of particles.
+
+Returns
+-------
+array_like
+		The result of the product. The result will have the same format as the forces array.
 )pbdoc";
 
+template <class Solver>
+void call_initialize(Solver &myself, libmobility::real T, libmobility::real eta,
+                     libmobility::real a, int N, libmobility::real tol) {
+  libmobility::Parameters par;
+  par.temperature = T;
+  par.viscosity = eta;
+  par.hydrodynamicRadius = {a};
+  par.tolerance = tol;
+  par.numberParticles = N;
+  myself.initialize(par);
+}
+
+template <class Solver> void call_setPositions(Solver &myself, pyarray_c &pos) {
+  myself.setPositions(cast_to_const_real(pos));
+}
 
+template <class Solver>
+auto call_hydrodynamicVelocities(Solver &myself, pyarray_c &forces,
+                                 libmobility::real prefactor) {
+  int N = myself.getNumberParticles();
+  if (forces.size() < 3 * N and forces.size() > 0) {
+    throw std::runtime_error("The forces array must have size 3*N.");
+  }
+  auto f = forces.size() ? cast_to_const_real(forces) : nullptr;
+  auto result = py::array_t<libmobility::real>({3 * N});
+  result.attr("fill")(0);
+  myself.hydrodynamicVelocities(f, cast_to_real(result), prefactor);
+  return result.reshape({N, 3});
+}
 
 const char *hydrodynamicvelocities_docstring = R"pbdoc(
 Computes the hydrodynamic (deterministic and stochastic) velocities.
@@ -136,58 +200,57 @@ Parameters
 ----------
 forces : array_like, optional
 		Forces acting on the particles.
-result : array_like
-		Where the result will be stored.
+
 prefactor : float, optional
 		Prefactor to multiply the result by. Default is 1.0.
-)pbdoc";
 
+Returns
+-------
+array_like
+		The resulting velocities. Shape is (N, 3), where N is the number of particles.
+)pbdoc";
 
-#define xMOBILITY_PYTHONIFY(MODULENAME, EXTRACODE, documentation)	\
-  PYBIND11_MODULE(MODULENAME, m){		      \
-  using real = libmobility::real;		      \
-  using Parameters = libmobility::Parameters;				\
-  using Configuration = libmobility::Configuration;			\
-  auto solver = py::class_<MODULENAME>(m, MOBILITYSTR(MODULENAME), documentation); \
-  solver.def(py::init([](std::string perx, std::string pery, std::string perz){	\
-    return std::unique_ptr<MODULENAME>(new MODULENAME(createConfiguration(perx, pery, perz))); }), \
-    constructor_docstring, "periodicityX"_a, "periodicityY"_a, "periodicityZ"_a). \
-    def("initialize", [](MODULENAME &myself, real T, real eta, real a, int N, real tol){ \
-    Parameters par;							\
-    par.temperature = T;						\
-    par.viscosity = eta;						\
-    par.hydrodynamicRadius = {a};					\
-    par.tolerance = tol;						\
-    par.numberParticles = N;						\
-    myself.initialize(par);						\
-  },									\
-    initialize_docstring,		\
-    "temperature"_a, "viscosity"_a,					\
-    "hydrodynamicRadius"_a,						\
-	  "numberParticles"_a,						\
-	  "tolerance"_a = 1e-4).						\
-    def("setPositions", [](MODULENAME &myself, pyarray pos){myself.setPositions(cast_to_const_real(pos));}, \
-	"The module will compute the mobility according to this set of positions.", \
-	"positions"_a).							\
-    def("Mdot", [](MODULENAME &myself, pyarray forces, pyarray result){\
-      auto f = forces.size()?cast_to_const_real(forces):nullptr;	\
-      myself.Mdot(f, cast_to_real(result));},	\
-      mdot_docstring,				\
-      "forces"_a = pyarray(), "result"_a).	\
-    def("sqrtMdotW", [](MODULENAME &myself, pyarray result, libmobility::real prefactor){ \
-      myself.sqrtMdotW(cast_to_real(result), prefactor);},		\
-      sqrtMdotW_docstring,						\
-      "result"_a = pyarray(), "prefactor"_a = 1.0).			\
-    def("hydrodynamicVelocities", [](MODULENAME &myself, pyarray forces,\
-					       pyarray result, libmobility::real prefactor){ \
-      auto f = forces.size()?cast_to_const_real(forces):nullptr;	\
-      myself.hydrodynamicVelocities(f, cast_to_real(result), prefactor);}, \
-	hydrodynamicvelocities_docstring,				\
-	"forces"_a = pyarray(), "result"_a  = pyarray(), "prefactor"_a = 1). \
-    def("clean", &MODULENAME::clean, "Frees any memory allocated by the module."). \
-    def_property_readonly_static("precision", [](py::object){return MODULENAME::precision;}, R"pbdoc(Compilation precision, a string holding either float or double.)pbdoc"); \
-  EXTRACODE\
+template <class Solver>
+auto call_construct(std::string perx, std::string pery, std::string perz) {
+  return std::unique_ptr<Solver>(
+      new Solver(createConfiguration(perx, pery, perz)));
 }
-#define MOBILITY_PYTHONIFY(MODULENAME, documentationPrelude) xMOBILITY_PYTHONIFY(MODULENAME,; ,documentationPrelude)
-#define MOBILITY_PYTHONIFY_WITH_EXTRA_CODE(MODULENAME, EXTRA, documentationPrelude) xMOBILITY_PYTHONIFY(MODULENAME,  EXTRA, documentationPrelude)
+
+#define xMOBILITY_PYTHONIFY(MODULENAME, EXTRACODE, documentation)                            \
+  PYBIND11_MODULE(MODULENAME, m) {                                                           \
+    using real = libmobility::real;                                                          \
+    using Parameters = libmobility::Parameters;                                              \
+    using Configuration = libmobility::Configuration;                                        \
+    auto solver =                                                                            \
+        py::class_<MODULENAME>(m, MOBILITYSTR(MODULENAME), documentation);                   \
+    solver                                                                                   \
+        .def(py::init(&call_construct<MODULENAME>), constructor_docstring,                   \
+             "periodicityX"_a, "periodicityY"_a, "periodicityZ"_a)                           \
+        .def("initialize", call_initialize<MODULENAME>, initialize_docstring,                \
+             "temperature"_a, "viscosity"_a, "hydrodynamicRadius"_a,                         \
+             "numberParticles"_a, "tolerance"_a = 1e-4)                                      \
+        .def("setPositions", call_setPositions<MODULENAME>,                                  \
+             "The module will compute the mobility according to this set of "                \
+             "positions.",                                                                   \
+             "positions"_a)                                                                  \
+        .def("Mdot", call_mdot<MODULENAME>, mdot_docstring,                                  \
+             "forces"_a = pyarray())                                                         \
+        .def("sqrtMdotW", call_sqrtMdotW<MODULENAME>, sqrtMdotW_docstring,                   \
+             "prefactor"_a = 1.0)                                                            \
+        .def("hydrodynamicVelocities",                                                       \
+             call_hydrodynamicVelocities<MODULENAME>,                                        \
+             hydrodynamicvelocities_docstring, "forces"_a = pyarray_c(),                     \
+             "prefactor"_a = 1)                                                              \
+        .def("clean", &MODULENAME::clean,                                                    \
+             "Frees any memory allocated by the module.")                                    \
+        .def_property_readonly_static(                                                       \
+            "precision", [](py::object) { return MODULENAME::precision; },                   \
+            R"pbdoc(Compilation precision, a string holding either float or double.)pbdoc"); \
+    EXTRACODE                                                                                \
+  }
+#define MOBILITY_PYTHONIFY(MODULENAME, documentationPrelude)                   \
+  xMOBILITY_PYTHONIFY(MODULENAME, ;, documentationPrelude)
+#define MOBILITY_PYTHONIFY_WITH_EXTRA_CODE(MODULENAME, EXTRA,                  \
+                                           documentationPrelude)               \
+  xMOBILITY_PYTHONIFY(MODULENAME, EXTRA, documentationPrelude)
 #endif