Support for non square boxes in DPStokes

CMakeLists.txt

@@ -46,7 +46,7 @@ endif()
 FetchContent_Declare(
   uammd
   GIT_REPOSITORY https://github.com/RaulPPelaez/uammd/
-  GIT_TAG        v2.11.0
+  GIT_TAG        v3.0.0
   EXCLUDE_FROM_ALL
 )
 FetchContent_MakeAvailable(uammd)

solvers/DPStokes/extra/poly_fits.h

@@ -37,9 +37,17 @@ double polyEval(std::vector<double> polyCoeffs, double x) {
 }
 
 // Coefficients for the polynomial fit of c^{-1} from [1]
-std::vector<double> cbetam_inv = {
+// the magic numbers come from the original DPStokes repo
+// https://github.com/stochasticHydroTools/DoublyPeriodicStokes/blob/main/source/cpu/python/GridAndKernelConfig.py
+std::vector<double> cbeta_monopole_inv = {
     4131643418.193291,  -10471683395.26777, 11833009228.6429,
     -7851132955.882548, 3388121732.651829,  -994285251.2185925,
     201183449.7086889,  -27776767.88241613, 2515647.646492857,
     -136305.2970161326, 3445.959503226691};
+
+std::vector<double> cbeta_dipole_inv = {
+    39129605.95992675,  -141095001.9380851, 226461386.9399433,
+    -213225975.7007014, 130589579.4920038,  -54470021.75408537,
+    15723274.21737542,  -3119398.842468358, 411336.473729197,
+    -33239.91568304499, 1311.988715157221};
 } // namespace dpstokes_polys

solvers/DPStokes/extra/uammd_interface.h

@@ -8,8 +8,10 @@ namespace uammd_dpstokes {
 // Instead of using uammd::real I have to re define real here.
 #ifndef DOUBLE_PRECISION
 using real = float;
+using real3 = float3;
 #else
 using real = double;
+using real3 = double3;
 #endif
 
 // This function returns either 'single' or 'double' according to the UAMMD's
@@ -23,7 +25,6 @@ struct PyParameters {
   int nx = -1;
   int ny = -1;
   int nz = -1;
-  real dt = 0;
   real viscosity;
   real Lx;
   real Ly;
@@ -32,8 +33,8 @@ struct PyParameters {
   real tolerance = 1e-5;
   real w, w_d;
   real hydrodynamicRadius = -1;
-  real beta = -1;
-  real beta_d = -1;
+  real3 beta = {-1.0, -1.0, -1.0};
+  real3 beta_d = {-1.0, -1.0, -1.0};
   real alpha = -1;
   real alpha_d = -1;
   // Can be either none, bottom, slit or periodic

solvers/DPStokes/extra/uammd_wrapper.cu

@@ -53,10 +53,13 @@ auto createFCMParameters(PyParameters pypar) {
   par.tolerance = pypar.tolerance;
   par.box = uammd::Box({pypar.Lx, pypar.Ly, pypar.zmax - pypar.zmin});
   par.cells = {pypar.nx, pypar.ny, pypar.nz};
-  par.kernel = std::make_shared<FCM_BM>(pypar.w, pypar.alpha, pypar.beta,
-                                        pypar.Lx / pypar.nx);
+  par.kernel =
+      std::make_shared<FCM_BM>(pypar.w, pypar.alpha,
+                               pypar.beta.x, // TODO beta parameter may need to
+                                             // be adjusted for non-square?
+                               pypar.Lx / pypar.nx);
   par.kernelTorque = std::make_shared<FCM_BM>(
-      pypar.w_d, pypar.alpha_d, pypar.beta_d, pypar.Lx / pypar.nx);
+      pypar.w_d, pypar.alpha_d, pypar.beta_d.x, pypar.Lx / pypar.nx);
   return par;
 }
 
@@ -76,14 +79,12 @@ auto createDPStokesParameters(PyParameters pypar) {
   par.nx = pypar.nx;
   par.ny = pypar.ny;
   par.nz = pypar.nz;
-  par.dt = pypar.dt;
   par.viscosity = pypar.viscosity;
   par.Lx = pypar.Lx;
   par.Ly = pypar.Ly;
   par.H = pypar.zmax - pypar.zmin;
   par.w = pypar.w;
   par.w_d = pypar.w_d;
-  par.hydrodynamicRadius = pypar.hydrodynamicRadius;
   par.beta = pypar.beta;
   par.beta_d = pypar.beta_d;
   par.alpha = pypar.alpha;
@@ -98,8 +99,7 @@ private:
   auto computeHydrodynamicDisplacements(const auto *d_pos,
                                         const uammd::real4 *d_force,
                                         const uammd::real4 *d_torques,
-                                        int numberParticles, real dt,
-                                        real dtTorque, cudaStream_t st) {
+                                        int numberParticles, cudaStream_t st) {
     if (fcm) {
       return fcm->computeHydrodynamicDisplacements(
           (uammd::real4 *)(d_pos), (uammd::real4 *)(d_force),
@@ -192,8 +192,7 @@ public:
                       Real3ToReal4SubstractOriginZ(zOrigin));
     auto mob = this->computeHydrodynamicDisplacements(
         stored_positions.data().get(), force4.data().get(),
-        (includeAngular) ? torque4.data().get() : nullptr, numberParticles, 0.0,
-        0.0, st);
+        (includeAngular) ? torque4.data().get() : nullptr, numberParticles, st);
     // always copy linear translation but only angular if needed
     thrust::copy(thrust::cuda::par.on(st), mob.first.begin(), mob.first.end(),
                  (uammd::real3 *)h_MF);

solvers/DPStokes/mobility.h

@@ -50,10 +50,6 @@ class DPStokes : public libmobility::Mobility {
 
   void setParametersDPStokes(DPStokesParameters i_dppar) {
     this->dppar = i_dppar;
-    if (this->dppar.Lx != this->dppar.Ly)
-      throw std::runtime_error("[DPStokes] Only square periodic boxes (Lx = "
-                               "Ly) are currently supported.\n");
-
     dpstokes = std::make_shared<uammd_dpstokes::DPStokesGlue>();
   }
 
@@ -71,33 +67,76 @@ class DPStokes : public libmobility::Mobility {
     if (ipar.includeAngular) {
       this->dppar.w = 6;
       this->dppar.w_d = 6;
-      this->dppar.beta = 1.327 * this->dppar.w;
-      this->dppar.beta_d = 2.217 * this->dppar.w;
+      this->dppar.beta = {real(1.327) * this->dppar.w,
+                          real(1.327) * this->dppar.w, -1.0};
+      this->dppar.beta_d = {real(2.217) * this->dppar.w_d,
+                            real(2.217) * this->dppar.w_d, -1.0};
       h = this->dppar.hydrodynamicRadius / 1.731;
       this->dppar.alpha_d = this->dppar.w_d * 0.5;
     } else {
+      // w=4
       this->dppar.w = 4;
-      this->dppar.beta = 1.785 * this->dppar.w;
+      this->dppar.beta = {real(1.785) * this->dppar.w,
+                          real(1.785) * this->dppar.w, -1.0};
       h = this->dppar.hydrodynamicRadius / 1.205;
+
+      // w=6
+      // this->dppar.w = 6;
+      // this->dppar.beta = {real(1.714) * this->dppar.w, real(1.714) *
+      // this->dppar.w, -1.0};
+      // h = this->dppar.hydrodynamicRadius / 1.554;
     }
     this->dppar.alpha = this->dppar.w * 0.5;
     this->dppar.tolerance = 1e-6;
 
-    int N = floor(this->dppar.Lx / h);
-    N += N % 2;
+    int Nx = floor(this->dppar.Lx / h);
+    Nx += Nx % 2;
+    int Ny = floor(this->dppar.Ly / h);
+    Ny += Ny % 2;
 
-    this->dppar.nx = N;
-    this->dppar.ny = N;
+    this->dppar.nx = Nx;
+    this->dppar.ny = Ny;
 
-    // note: this part is only configured for square boxes
     if (this->dppar.allowChangingBoxSize) { // adjust box size to suit h
-      this->dppar.Lx = N * h;
-      this->dppar.Ly = N * h;
+      this->dppar.Lx = Nx * h;
+      this->dppar.Ly = Ny * h;
     } else { // adjust h so that L/h is an integer
-      h = this->dppar.Lx / N;
-      double arg = this->dppar.hydrodynamicRadius / (this->dppar.w * h);
-      this->dppar.beta =
-          dpstokes_polys::polyEval(dpstokes_polys::cbetam_inv, arg);
+      real h_x = this->dppar.Lx / Nx;
+      real h_y = this->dppar.Ly / Ny;
+      h = min(h_x, h_y);
+      double arg = this->dppar.hydrodynamicRadius / (this->dppar.w * h_x);
+      real beta_x =
+          dpstokes_polys::polyEval(dpstokes_polys::cbeta_monopole_inv, arg);
+      arg = this->dppar.hydrodynamicRadius / (this->dppar.w * h_y);
+      real beta_y =
+          dpstokes_polys::polyEval(dpstokes_polys::cbeta_monopole_inv, arg);
+
+      if (beta_x < 4.0 || beta_x > 18.0 || beta_y < 4.0 || beta_y > 18.0) {
+        throw std::runtime_error(
+            "[DPStokes] Could not find (h,beta) within interp range. This "
+            "means the particle radius and grid spacing are incompatible- try "
+            "a square domain.");
+      }
+
+      this->dppar.beta = {beta_x, beta_y, min(beta_x, beta_y)};
+
+      if (ipar.includeAngular) {
+        arg = this->dppar.hydrodynamicRadius / (this->dppar.w_d * h_x);
+        real beta_xd =
+            dpstokes_polys::polyEval(dpstokes_polys::cbeta_dipole_inv, arg);
+        arg = this->dppar.hydrodynamicRadius / (this->dppar.w_d * h_y);
+        real beta_yd =
+            dpstokes_polys::polyEval(dpstokes_polys::cbeta_dipole_inv, arg);
+        if (beta_xd < 4.0 || beta_xd > 18.0 || beta_yd < 4.0 ||
+            beta_yd > 18.0) {
+          throw std::runtime_error(
+              "[DPStokes] Could not find (h,beta) within interp range. This "
+              "means the particle radius and grid spacing are incompatible- "
+              "try "
+              "a square domain.");
+        }
+        this->dppar.beta_d = {beta_xd, beta_yd, min(beta_xd, beta_yd)};
+      }
     }
 
     // Add a buffer of 1.5*w*h/2 when there is an open boundary
@@ -113,7 +152,8 @@ class DPStokes : public libmobility::Mobility {
     // sets chebyshev node spacing at its coarsest (in the middle) to be h
     real nz_actual = M_PI / (asin(h / H)) + 1;
 
-    // pick nearby N such that 2(Nz-1) has two factors of 2 and is FFT friendly
+    // pick nearby N such that 2(Nz-1) has two factors of 2 and is FFT
+    // friendly
     this->dppar.nz = floor(nz_actual);
     this->dppar.nz += (int)ceil(nz_actual) % 2;
 

tests/ref/readme.txt

@@ -2,6 +2,9 @@ Some of this data is generated from the original DoublyPeriodicStokes repository
 
 The reason we've regenerated data is that default parameter selection has been changed to fix the grid parameters to the values that best preserve the hydrodynamic radius. When data has been regenerated, consistency was checked by first matching the original DoublyPeriodicStokes repository by using their same parameters, then switching to the new parameter selection and regenerating.
 
+8/21/2025: Regenerated DPStokes data for angular tests since the procedure for selecting beta for torques was updated.
+           Agreed with previous data at tol=5e-2.
+
 Abbreviations: 
 bw- bottom wall
 sc- slit channel

tests/test_dpstokes.py

@@ -0,0 +1,128 @@
+import numpy as np
+import libMobility as lm
+
+
+def compute_with_dpstokes(pos, lx, ly, forces=None, torques=None):
+    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
+    dpstokes.setParameters(Lx=lx, Ly=ly, zmin=-8.0, zmax=8.0)
+    dpstokes.initialize(**params, includeAngular=torques_on)
+    dpstokes.setPositions(pos)
+    mf_dp, mt_dp = dpstokes.Mdot(forces=forces, torques=torques)
+    if torques_on:
+        return mt_dp
+    else:
+        return mf_dp
+
+
+def test_non_square_box():
+    nP = 10
+    pos_or = np.random.uniform(-8, 8, (nP, 3)).astype(np.float32)
+    forces_or = np.random.uniform(-1, 1, (nP, 3)).astype(np.float32)
+    forces_or -= np.mean(forces_or, axis=0)  # Center the forces around zero
+
+    L_fact = 2.0
+    L_short = 16.0
+    L_long = L_short * L_fact
+
+    mf_dp_cube = compute_with_dpstokes(pos_or, forces=forces_or, lx=L_short, ly=L_short)
+    mt_dp_cube = compute_with_dpstokes(
+        pos_or, torques=forces_or, lx=L_short, ly=L_short
+    )
+
+    # 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), :
+    ]
+
+    # 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), :
+    ]
+
+    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)
+
+    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)
+
+
+def test_isotropy():
+    L0 = 16.0
+    L_fact = 2.0
+
+    f_x = np.array([1.0, 0.0, 0.0])
+    f_y = np.array([0.0, 1.0, 0.0])
+
+    n_runs = 5
+    for _ in range(n_runs):
+        pos = np.random.uniform(-8, 8, (3)).astype(np.float32)
+
+        mf_x = compute_with_dpstokes(pos, forces=f_x, lx=L0 * L_fact, ly=L0)
+        mf_y = compute_with_dpstokes(pos, forces=f_y, lx=L0, ly=L0 * L_fact)
+
+        mt_x = compute_with_dpstokes(pos, torques=f_x, lx=L0 * L_fact, ly=L0)
+        mt_y = compute_with_dpstokes(pos, torques=f_y, lx=L0, ly=L0 * L_fact)
+
+        assert np.allclose(mf_x[0], mf_y[1], rtol=1e-3, atol=1e-2)
+        assert np.allclose(mf_x[1], mf_y[0], rtol=1e-3, atol=1e-2)
+        assert np.allclose(mf_x[2], mf_y[2], rtol=1e-3, atol=1e-2)
+
+        assert np.allclose(mt_x[0], mt_y[1], rtol=1e-3, atol=1e-2)
+        assert np.allclose(mt_x[1], mt_y[0], rtol=1e-3, atol=1e-2)
+        assert np.allclose(mt_x[2], mt_y[2], rtol=1e-3, atol=1e-2)
+
+
+def test_non_square_matching_rpy():
+
+    a = np.random.uniform(0.5, 2.0)
+    eta = 1 / (6 * np.pi * a)
+    L_min = 100 * a  # need a fairly large domain to neglect periodic effects
+    Lx = np.random.uniform(L_min, 3.0 * L_min)
+    Ly = np.random.uniform(L_min, 3.0 * L_min)
+    max_height = 10.0 * a
+    min_height = 1.5 * a  # need a buffer for regularized kernels to match
+
+    solver = lm.DPStokes("periodic", "periodic", "single_wall")
+    solver.setParameters(Lx=Lx, Ly=Ly, zmin=0.0, zmax=max_height)
+    solver.initialize(hydrodynamicRadius=a, viscosity=eta, includeAngular=True)
+
+    rpy = lm.NBody("open", "open", "single_wall")
+    rpy.setParameters(wallHeight=0.0)
+    rpy.initialize(hydrodynamicRadius=a, viscosity=eta, includeAngular=True)
+
+    F = np.eye(6)
+
+    heights = np.linspace(min_height, max_height, 20)
+
+    results_dpstokes = np.zeros((2 * len(heights), 3))
+    results_rpy = np.zeros((2 * len(heights), 3))
+
+    for i, h in enumerate(heights):
+        pos = np.array([0.0, 0.0, h])
+
+        rpy.setPositions(pos)
+        solver.setPositions(pos)
+        for j in range(6):
+            f_j = F[0:3, j].copy()
+            t_j = F[3:6, j].copy()
+            mf_j, mt_j = solver.Mdot(forces=f_j, torques=t_j)
+            results_dpstokes[i] += mf_j
+            results_dpstokes[i + 3] += mt_j
+
+            mf_j, mt_j = rpy.Mdot(forces=f_j, torques=t_j)
+            results_rpy[i] += mf_j
+            results_rpy[i + 3] += mt_j
+
+    assert np.allclose(results_dpstokes, results_rpy, rtol=1e-3, atol=1e-2)

tests/test_initialization.py

@@ -32,14 +32,6 @@ def test_invalid_throws(Solver):
         solver = Solver("periodicasdas", "periodic", "open")
 
 
-def test_dpstokes_invalid_box():
-    with pytest.raises(RuntimeError):
-        Lx, Ly = 10, 20
-        solver = DPStokes("periodic", "periodic", "single_wall")
-        params = {"Lx": Lx, "Ly": Ly, "zmin": 0, "zmax": 19.2}
-        solver.setParameters(**params)
-
-
 @pytest.mark.parametrize(
     ("NBatch", "NperBatch"),
     [