Merge pull request #79 from stochasticHydroTools/dpstokes_open

Error checking to ensure net zero forces/torques for DPStokes when Z is open

Author: rykerfish

solvers/DPStokes/extra/uammd_interface.h

@@ -3,15 +3,18 @@
 */
 #include <memory>
 #include <string>
+#include <vector_types.h>
 namespace uammd_dpstokes {
 // This is in order not to use any UAMMD related includes here.
 // Instead of using uammd::real I have to re define real here.
 #ifndef DOUBLE_PRECISION
 using real = float;
 using real3 = float3;
+using real4 = float4;
 #else
 using real = double;
 using real3 = double3;
+using real4 = double4;
 #endif
 
 // This function returns either 'single' or 'double' according to the UAMMD's
@@ -39,6 +42,7 @@ struct PyParameters {
   // Can be either none, bottom, slit or periodic
   std::string mode;
   bool allowChangingBoxSize = false;
+  bool allowUnsafeForces = false;
 };
 
 class DPStokesUAMMD;

solvers/DPStokes/mobility.h

@@ -12,8 +12,29 @@ in linear time. A. Hashemi et al. J. Chem. Phys. 158, 154101 (2023)
 #include <MobilityInterface/MobilityInterface.h>
 #include <MobilityInterface/random_finite_differences.h>
 #include <cmath>
+#include <thrust/transform_reduce.h>
 #include <vector>
 
+namespace {
+struct vec3_to_vec4 {
+  __device__ uammd_dpstokes::real4 operator()(const uammd_dpstokes::real3 &a) {
+    return {a.x, a.y, a.z, abs(a.x) + abs(a.y) + abs(a.z)};
+  }
+};
+
+uammd_dpstokes::real3 calculate_relative_vec3(const uammd_dpstokes::real3 *vptr,
+                                              uint numberParticles) {
+  uammd_dpstokes::real4 sum_v = thrust::transform_reduce(
+      thrust::cuda::par, vptr, vptr + numberParticles, vec3_to_vec4(),
+      uammd_dpstokes::real4{0, 0, 0, 0},
+      cuda::std::plus<uammd_dpstokes::real4>());
+  uammd_dpstokes::real3 rel_v =
+      uammd_dpstokes::real3{sum_v.x, sum_v.y, sum_v.z} / sum_v.w;
+
+  return rel_v;
+}
+} // namespace
+
 class DPStokes : public libmobility::Mobility {
   using periodicity_mode = libmobility::periodicity_mode;
   using Configuration = libmobility::Configuration;
@@ -54,6 +75,7 @@ class DPStokes : public libmobility::Mobility {
   }
 
   void initialize(Parameters ipar) override {
+    this->par.includeAngular = ipar.includeAngular;
     this->dppar.viscosity = ipar.viscosity;
     this->lanczosTolerance = ipar.tolerance;
     this->dppar.mode = this->wallmode;
@@ -183,6 +205,40 @@ class DPStokes : public libmobility::Mobility {
     device_adapter<real> linear(ilinear, device::cuda);
     device_adapter<real> angular(iangular, device::cuda);
 
+    if (this->wallmode == "nowall" && !this->dppar.allowUnsafeForces) {
+      constexpr real errTol = 1e-6;
+      bool errFlag = false;
+      if (iforces.size() > 0) {
+        const uammd_dpstokes::real3 *fptr =
+            reinterpret_cast<const uammd_dpstokes::real3 *>(forces.data());
+        uammd_dpstokes::real3 rel_f =
+            calculate_relative_vec3(fptr, this->numberParticles);
+        if (fabs(rel_f.x) > errTol || fabs(rel_f.y) > errTol ||
+            fabs(rel_f.z) > errTol) {
+          errFlag = true;
+        }
+      }
+
+      if (this->par.includeAngular && itorques.size() > 0) {
+        const uammd_dpstokes::real3 *tptr =
+            reinterpret_cast<const uammd_dpstokes::real3 *>(torques.data());
+        uammd_dpstokes::real3 rel_t =
+            calculate_relative_vec3(tptr, this->numberParticles);
+        if (fabs(rel_t.x) > errTol || fabs(rel_t.y) > errTol ||
+            fabs(rel_t.z) > errTol) {
+          errFlag = true;
+        }
+      }
+      if (errFlag) {
+        throw std::runtime_error(
+            "[DPStokes] If using an open boundary in z, the mean force and "
+            "torque in each dimension must be zero or the solver will produce "
+            "non-physical results. If you are sure your forces and torques "
+            "satisfy this condition, you can set allowUnsafeForces=True to "
+            "bypass this check.");
+      }
+    }
+
     dpstokes->Mdot(forces.data(), torques.data(), linear.data(), angular.data(),
                    this->getNumberParticles(), this->getIncludeAngular());
   }
@@ -228,12 +284,12 @@ class DPStokes : public libmobility::Mobility {
     this->setPositions(original_pos);
     thrust::transform(thrust::cuda::par, thermal_drift_m.begin(),
                       thermal_drift_m.end(), linear.begin(), linear.begin(),
-                      thrust::plus<real>());
+                      cuda::std::plus<real>());
     if (this->getIncludeAngular()) {
       device_adapter<real> angular(iangular, device::cuda);
       thrust::transform(thrust::cuda::par, thermal_drift_d.begin(),
                         thermal_drift_d.end(), angular.begin(), angular.begin(),
-                        thrust::plus<real>());
+                        cuda::std::plus<real>());
     }
   }
 

solvers/DPStokes/python_wrapper.cu

@@ -23,7 +23,9 @@ zmax : float
 allowChangingBoxSize : bool
     Whether the periodic extents Lx & Ly can be modified during parameter selection. Default: false.
 delta : float
-    The finite difference step size for random finite differences. Specified in units of hydrodynamicRadius. Default is 1e-3.
+    The finite difference step size for random finite differences. Specified in units of hydrodynamicRadius. Default: 1e-3.
+allowUnsafeForces : bool
+    When the Z periodicity is `open` the net force/torque on the domain must be zero. Numerically, this is checked to some tolerance. Enabling this flag will by bypass the check but can lead to unphysical results if used incorrectly. Default: false.
 )pbdoc";
 
 static const char *docstring = R"pbdoc(
@@ -41,16 +43,18 @@ MOBILITY_PYTHONIFY_WITH_EXTRA_CODE(
     solver.def(
         "setParameters",
         [](DPStokes &self, real Lx, real Ly, real zmin, real zmax,
-           bool allowChangingBoxSize, real delta) {
+           bool allowChangingBoxSize, real delta, bool allowUnsafeForces) {
           DPStokesParameters params;
           params.Lx = Lx;
           params.Ly = Ly;
           params.zmin = zmin;
           params.zmax = zmax;
           params.allowChangingBoxSize = allowChangingBoxSize;
           params.delta = delta;
+          params.allowUnsafeForces = allowUnsafeForces;
           self.setParametersDPStokes(params);
         },
         "Lx"_a, "Ly"_a, "zmin"_a, "zmax"_a, "allowChangingBoxSize"_a = false,
-        "delta"_a = 1e-3, setparameters_docstring);
+        "delta"_a = 1e-3, "allowUnsafeForces"_a = false,
+        setparameters_docstring);
     , docstring);

tests/test_dpstokes.py

@@ -1,8 +1,9 @@
 import numpy as np
 import libMobility as lm
+import pytest
 
 
-def compute_with_dpstokes(pos, lx, ly, forces=None, torques=None):
+def compute_with_dpstokes(pos, lx, ly, forces=None, torques=None) -> np.ndarray:
     params = {"viscosity": 1 / (6 * np.pi), "hydrodynamicRadius": 1.73}
     dpstokes = lm.DPStokes("periodic", "periodic", "single_wall")
     torques_on = True if torques is not None else False
@@ -11,9 +12,9 @@ def compute_with_dpstokes(pos, lx, ly, forces=None, torques=None):
     dpstokes.setPositions(pos)
     mf_dp, mt_dp = dpstokes.Mdot(forces=forces, torques=torques)
     if torques_on:
-        return mt_dp
+        return np.array(mt_dp)
     else:
-        return mf_dp
+        return np.array(mf_dp)
 
 
 def test_non_square_box():
@@ -34,28 +35,21 @@ def test_non_square_box():
     # This must be identical than doubling the box size in x and repeating the particles
     pos = np.vstack((pos_or, pos_or + np.array([0.5 * L_long, 0.0, 0.0])))
     forces = np.vstack((forces_or, forces_or))
-    mf_dp_lx = compute_with_dpstokes(pos, forces=forces, lx=L_long, ly=L_short)[
-        : len(pos_or), :
-    ]
-    mt_dp_lx = compute_with_dpstokes(pos, torques=forces, lx=L_long, ly=L_short)[
-        : len(pos_or), :
-    ]
+    mf_dp_lx = compute_with_dpstokes(pos, forces=forces, lx=L_long, ly=L_short)
+    mt_dp_lx = compute_with_dpstokes(pos, torques=forces, lx=L_long, ly=L_short)
 
     # And the same for doubling the box size in y
     pos = np.vstack((pos_or, pos_or + np.array([0.0, 0.5 * L_long, 0.0])))
     forces = np.vstack((forces_or, forces_or))
-    mf_dp_ly = compute_with_dpstokes(pos, forces=forces, lx=L_short, ly=L_long)[
-        : len(pos_or), :
-    ]
-    mt_dp_ly = compute_with_dpstokes(pos, torques=forces, lx=L_short, ly=L_long)[
-        : len(pos_or), :
-    ]
+    mf_dp_ly = compute_with_dpstokes(pos, forces=forces, lx=L_short, ly=L_long)
+    mt_dp_ly = compute_with_dpstokes(pos, torques=forces, lx=L_short, ly=L_long)
 
-    assert np.allclose(mf_dp_lx, mf_dp_cube, rtol=1e-3, atol=1e-2)
-    assert np.allclose(mf_dp_ly, mf_dp_cube, rtol=1e-3, atol=1e-2)
+    slice = np.s_[0 : len(pos_or), :]
+    assert np.allclose(mf_dp_lx[slice], mf_dp_cube, rtol=1e-3, atol=1e-2)
+    assert np.allclose(mf_dp_ly[slice], mf_dp_cube, rtol=1e-3, atol=1e-2)
 
-    assert np.allclose(mt_dp_lx, mt_dp_cube, rtol=1e-3, atol=1e-2)
-    assert np.allclose(mt_dp_ly, mt_dp_cube, rtol=1e-3, atol=1e-2)
+    assert np.allclose(mt_dp_lx[slice], mt_dp_cube, rtol=1e-3, atol=1e-2)
+    assert np.allclose(mt_dp_ly[slice], mt_dp_cube, rtol=1e-3, atol=1e-2)
 
 
 def test_isotropy():
@@ -126,3 +120,29 @@ def test_non_square_matching_rpy():
             results_rpy[i + 3] += mt_j
 
     assert np.allclose(results_dpstokes, results_rpy, rtol=1e-3, atol=1e-2)
+
+
+def test_dpstokes_open_errors():
+    solver = lm.DPStokes("periodic", "periodic", "open")
+    solver.setParameters(Lx=10.0, Ly=10.0, zmin=0.0, zmax=10.0)
+    solver.initialize(hydrodynamicRadius=1.0, viscosity=1.0, includeAngular=True)
+
+    pos = np.random.uniform(0.0, 10.0, (10, 3))
+    solver.setPositions(pos)
+
+    forces = np.random.uniform(-1.0, 1.0, (10, 3))
+    with pytest.raises(RuntimeError, match="DPStokes"):
+        solver.Mdot(forces=forces)
+    with pytest.raises(RuntimeError, match="DPStokes"):
+        solver.Mdot(torques=forces)
+
+    forces -= np.mean(forces, axis=0)
+    solver.Mdot(forces=forces, torques=forces)
+
+    # check that we can bypass the error if we allow unsafe forces
+    forces = np.random.uniform(-1.0, 1.0, (10, 3))
+    solver = lm.DPStokes("periodic", "periodic", "open")
+    solver.setParameters(Lx=10.0, Ly=10.0, zmin=0.0, zmax=10.0, allowUnsafeForces=True)
+    solver.initialize(hydrodynamicRadius=1.0, viscosity=1.0, includeAngular=True)
+    solver.setPositions(pos)
+    solver.Mdot(forces=forces, torques=forces)

tests/utils.py

@@ -108,6 +108,9 @@ def get_sane_params(solverName, geom=None):
         params = sane_parameters["NBody_wall"].copy()
     else:
         params = sane_parameters[solverName].copy()
+
+    if solverName == "DPStokes" and geom == "open":
+        params["allowUnsafeForces"] = True
     return params