Fix bug in self mobility (#30)

* Fix bug in self mobility where sqrtMdotW would overwrite the results of Mdot in hydrodynamicVelocities

* Fill linear with 0 in sqrtMdotW

* Small change

include/MobilityInterface/pythonify.h

@@ -133,8 +133,8 @@ template <class Solver>
 auto call_sqrtMdotW(Solver &solver, libmobility::real prefactor) {
   int N = solver.getNumberParticles();
   auto linear = py::array_t<libmobility::real>({N * 3});
+  linear.attr("fill")(0);
   auto angular = py::array_t<libmobility::real>();
-
   if (solver.getNeedsTorque()) {
     angular.resize({3 * N});
     angular.attr("fill")(0);
@@ -143,7 +143,8 @@ auto call_sqrtMdotW(Solver &solver, libmobility::real prefactor) {
   } else {
     solver.sqrtMdotW(cast_to_real(linear), nullptr, prefactor);
   }
-  return std::make_pair(linear.reshape({N, 3}), angular);
+  linear = linear.reshape({N, 3});
+  return std::make_pair(linear, angular);
 }
 
 const char *sqrtMdotW_docstring = R"pbdoc(

solvers/SelfMobility/mobility.h

@@ -7,11 +7,9 @@
 #ifndef MOBILITY_SELFMOBILITY_H
 #define MOBILITY_SELFMOBILITY_H
 #include<MobilityInterface/MobilityInterface.h>
-#include <cstdint>
 #include <random>
 #include<vector>
 #include<cmath>
-#include<type_traits>
 
 class SelfMobility: public libmobility::Mobility{
   using periodicity_mode = libmobility::periodicity_mode;
@@ -56,12 +54,12 @@ class SelfMobility: public libmobility::Mobility{
   void Mdot(const real* forces, const real* torques, real* linear, real* angular) override{
     if(forces){
       for(int i = 0; i<3*numberParticles; i++){
-        linear[i] = forces[i]*linearMobility;
+        linear[i] += forces[i]*linearMobility;
       }
     }
     if(torques){
       for(int i = 0; i<3*numberParticles; i++){
-        angular[i] = torques[i]*angularMobility;
+        angular[i] += torques[i]*angularMobility;
       }
     }
   }
@@ -71,12 +69,12 @@ class SelfMobility: public libmobility::Mobility{
     std::normal_distribution<real> d{0,1};
     for(int i = 0; i<3*numberParticles; i++){
       real dW = d(rng);
-      linear[i] = prefactor*sqrt(2*temperature*linearMobility)*dW;
+      linear[i] += prefactor*sqrt(2*temperature*linearMobility)*dW;
     }
     if(angular){
       for(int i = 0; i<3*numberParticles; i++){
         real dW = d(rng);
-        angular[i] = prefactor*sqrt(2*temperature*angularMobility)*dW;
+        angular[i] += prefactor*sqrt(2*temperature*angularMobility)*dW;
       }
     }
   }

tests/test_interface.py

@@ -74,6 +74,7 @@ def test_returns_mf(Solver, periodicity):
     mf, _ = solver.Mdot(forces)
     assert mf.shape == (numberParticles, 3)
 
+
 @pytest.mark.parametrize(
     ("Solver", "periodicity"),
     [
@@ -95,7 +96,7 @@ def test_returns_mf_mt(Solver, periodicity):
         viscosity=1.0,
         hydrodynamicRadius=hydrodynamicRadius,
         numberParticles=numberParticles,
-        needsTorque=True
+        needsTorque=True,
     )
 
     # Set precision to be the same as compiled precision
@@ -179,3 +180,44 @@ def test_returns_hydrodisp(Solver, periodicity):
     assert sqrtmw.shape == (numberParticles, 3)
     sqrtmw, _ = solver.hydrodynamicVelocities(forces)
     assert sqrtmw.shape == (numberParticles, 3)
+
+
+@pytest.mark.parametrize(
+    ("Solver", "periodicity"),
+    [
+        (SelfMobility, ("open", "open", "open")),
+        (PSE, ("periodic", "periodic", "periodic")),
+        (NBody, ("open", "open", "open")),
+        (DPStokes, ("periodic", "periodic", "open")),
+        (DPStokes, ("periodic", "periodic", "single_wall")),
+        (DPStokes, ("periodic", "periodic", "two_walls")),
+    ],
+)
+@pytest.mark.parametrize("needsTorque", [True, False])
+def test_hydrodisp_equivalent(Solver, periodicity, needsTorque):
+    #  Check that calling Mdot is equivalent to calling hydrodynamicVelocities with temperature = 0
+    hydrodynamicRadius = 1.0
+    solver = Solver(*periodicity)
+    parameters = sane_parameters[Solver.__name__]
+    solver.setParameters(**parameters)
+    numberParticles = 1
+    solver.initialize(
+        temperature=0.0,
+        viscosity=1.0,
+        hydrodynamicRadius=hydrodynamicRadius,
+        numberParticles=numberParticles,
+        needsTorque=needsTorque,
+    )
+
+    # Set precision to be the same as compiled precision
+    precision = np.float32 if Solver.precision == "float" else np.float64
+    positions = np.random.rand(numberParticles, 3).astype(precision)
+    forces = np.random.rand(numberParticles, 3).astype(precision)
+    torques = np.random.rand(numberParticles, 3).astype(precision)
+    solver.setPositions(positions)
+    args = (forces, torques) if needsTorque else (forces,)
+    mf, mt = solver.Mdot(*args)
+    bothmf, bothmt = solver.hydrodynamicVelocities(*args)
+    assert np.allclose(mf, bothmf, atol=1e-6)
+    if needsTorque:
+        assert np.allclose(mt, bothmt, atol=1e-6)