From 9371d6c80a9b9b112153f00124630a020ecbdf77 Mon Sep 17 00:00:00 2001 From: Ryker Fish Date: Thu, 20 Nov 2025 17:02:50 -0700 Subject: [PATCH 1/6] deterministic div M test working for forces --- tests/test_thermal_drift.py | 104 +++++++++++++++++++++++++++++++++++- 1 file changed, 102 insertions(+), 2 deletions(-) diff --git a/tests/test_thermal_drift.py b/tests/test_thermal_drift.py index 16dcdfa..b8ce112 100644 --- a/tests/test_thermal_drift.py +++ b/tests/test_thermal_drift.py @@ -21,9 +21,75 @@ def average(function, num_averages): return result_m / num_averages, result_d -def divM_rfd(solver, positions): +def average_with_error(function, num_averages, soln, N): + # Track cumulative (running) error: error of the running average after each new sample + avg_err_m = np.zeros(num_averages) + rfd_m = np.zeros(3 * N) + avg_err_d = np.zeros(num_averages) + rfd_d = np.zeros(3 * N) + + for i in range(num_averages): + new_result_m, new_result_d = function() + rfd_m += new_result_m + running_m = rfd_m / float(i + 1) + avg_err_m[i] = np.linalg.norm(running_m - soln) + if new_result_d is not None: + rfd_d += new_result_d + running_d = rfd_d / float(i + 1) + avg_err_d[i] = np.linalg.norm(running_d - soln) / np.linalg.norm(soln) + + rfd_m /= float(num_averages) + rfd_d /= float(num_averages) + + import matplotlib.pyplot as plt + + plt.figure() + n = np.arange(num_averages) + plt.loglog(n, avg_err_m, marker="o", linestyle="-") + plt.loglog(n, 0.1 * 1 / np.sqrt(n), linestyle="--") + plt.legend(["Measured error", "O(1/sqrt(N))"]) + # plt.plot( + # np.arange(num_averages), + # avg_err_d, + # linestyle="--", + # color="red", + # ) + plt.xlabel("Sample index") + plt.ylabel("Error (||running_avg - sol||) or relative for dipole") + plt.title("Cumulative running error of the averaged estimate (linear)") + plt.grid(True) + plt.tight_layout() + plt.savefig("avg_err_m.png", dpi=150) + plt.close() + return rfd_m, rfd_d + + +def deterministic_div_m(solver, positions, delta): + N = np.size(positions) // 3 + print("Number of particles:", N) + + positions = positions.flatten() + div_M = np.zeros(3 * N, dtype=positions.dtype) + for i in range(3 * N): + for j in range(3 * N): + pos_plus = positions.copy() + pos_minus = positions.copy() + pos_plus[j] += 0.5 * delta + pos_minus[j] -= 0.5 * delta + solver.setPositions(pos_plus) + e_j = np.zeros(3 * N, dtype=positions.dtype) + e_j[j] = 1.0 + M_plus, _ = solver.Mdot(e_j) + solver.setPositions(pos_minus) + M_minus, _ = solver.Mdot(e_j) + div_M[i] += M_plus[i] - M_minus[i] + + div_M /= delta + return div_M + + +def divM_rfd(solver, positions, delta): # RFD works by approximating \partial_q \dot M = 1/\delta \langle M(q+\delta/2 W)W - M(q-\delta/2 W)W \rangle - delta = 1e-3 def rfd_func(): W = np.random.normal(size=positions.shape).astype(positions.dtype) @@ -42,6 +108,40 @@ def rfd_func(): return tdrift_m, tdrift_d +@pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) +def test_deterministic_divM_matches_rfd(Solver, periodicity): + solver = Solver(*periodicity) + parameters = get_sane_params(Solver.__name__, periodicity[2]) + solver.setParameters(**parameters) + precision = np.float32 if solver.precision == "float" else np.float64 + solver.initialize( + viscosity=1.0, + hydrodynamicRadius=1.0, + includeAngular=False, + ) + numberParticles = 10 + positions = generate_positions_in_box(parameters, numberParticles).astype(precision) + solver.setPositions(positions) + + delta = 1e-3 + det_div_m = deterministic_div_m(solver, positions, delta) + # rfd_dM, _ = divM_rfd(solver, positions, delta) + # rfd_m, rfd_d = average(solver.divM, 10000) + m, d = average_with_error(solver.divM, 10000, det_div_m, numberParticles) + + breakpoint() + print("-----------------------") + print("Deterministic divM:", det_div_m) + print("RFD divM:", rfd_m) + + assert np.allclose( + det_div_m, + rfd_m, + atol=1e-3, + rtol=1e-3, + ), f"Deterministic divM does not match RFD divM: {np.max(np.abs(det_div_m - rfd_m))}" + + @pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) def test_divM_does_not_change_positions(Solver, periodicity): # Compute the Mdot with some random forces, then compute divM, then compute Mdot again From 58c6fb961f00407c7f63e2ebdba728785f8d179b Mon Sep 17 00:00:00 2001 From: Ryker Fish Date: Thu, 20 Nov 2025 18:29:01 -0700 Subject: [PATCH 2/6] deterministic divM test working for dipole as well --- tests/test_thermal_drift.py | 99 ++++++++++++++++++++----------------- 1 file changed, 53 insertions(+), 46 deletions(-) diff --git a/tests/test_thermal_drift.py b/tests/test_thermal_drift.py index b8ce112..9f71abd 100644 --- a/tests/test_thermal_drift.py +++ b/tests/test_thermal_drift.py @@ -21,22 +21,25 @@ def average(function, num_averages): return result_m / num_averages, result_d -def average_with_error(function, num_averages, soln, N): - # Track cumulative (running) error: error of the running average after each new sample +def average_with_error(function, num_averages, soln, N, includeAngular): + # track error of the running average of RFDs after each new sample avg_err_m = np.zeros(num_averages) rfd_m = np.zeros(3 * N) avg_err_d = np.zeros(num_averages) rfd_d = np.zeros(3 * N) + soln_m = soln[0 : 3 * N] + soln_d = soln[3 * N :] if includeAngular else None + for i in range(num_averages): new_result_m, new_result_d = function() rfd_m += new_result_m running_m = rfd_m / float(i + 1) - avg_err_m[i] = np.linalg.norm(running_m - soln) + avg_err_m[i] = np.linalg.norm(running_m - soln_m) if new_result_d is not None: rfd_d += new_result_d running_d = rfd_d / float(i + 1) - avg_err_d[i] = np.linalg.norm(running_d - soln) / np.linalg.norm(soln) + avg_err_d[i] = np.linalg.norm((running_d - soln_d)) rfd_m /= float(num_averages) rfd_d /= float(num_averages) @@ -46,17 +49,17 @@ def average_with_error(function, num_averages, soln, N): plt.figure() n = np.arange(num_averages) plt.loglog(n, avg_err_m, marker="o", linestyle="-") - plt.loglog(n, 0.1 * 1 / np.sqrt(n), linestyle="--") - plt.legend(["Measured error", "O(1/sqrt(N))"]) - # plt.plot( - # np.arange(num_averages), - # avg_err_d, - # linestyle="--", - # color="red", - # ) - plt.xlabel("Sample index") - plt.ylabel("Error (||running_avg - sol||) or relative for dipole") - plt.title("Cumulative running error of the averaged estimate (linear)") + plt.loglog(n, avg_err_d, marker="o", linestyle="-") + plt.loglog(n, 0.1 * 1 / np.sqrt(n + 1), linestyle="--", color="black") + if includeAngular: + plt.legend( + ["Measured error (linear)", "Measured error (dipole)", "O(1/sqrt(N))"] + ) + else: + plt.legend(["Measured error", "O(1/sqrt(N))"]) + plt.xlabel("RFD count") + plt.ylabel("Error (||running_avg - sol||)") + plt.title("Error of running average of RFDs compared to deterministic divM") plt.grid(True) plt.tight_layout() plt.savefig("avg_err_m.png", dpi=150) @@ -64,25 +67,33 @@ def average_with_error(function, num_averages, soln, N): return rfd_m, rfd_d -def deterministic_div_m(solver, positions, delta): +def deterministic_div_m(solver, positions, includeAngular): + delta = 1e-3 N = np.size(positions) // 3 - print("Number of particles:", N) + size = 6 * N if includeAngular else 3 * N positions = positions.flatten() - div_M = np.zeros(3 * N, dtype=positions.dtype) - for i in range(3 * N): - for j in range(3 * N): - pos_plus = positions.copy() - pos_minus = positions.copy() - pos_plus[j] += 0.5 * delta - pos_minus[j] -= 0.5 * delta - solver.setPositions(pos_plus) - e_j = np.zeros(3 * N, dtype=positions.dtype) - e_j[j] = 1.0 - M_plus, _ = solver.Mdot(e_j) - solver.setPositions(pos_minus) - M_minus, _ = solver.Mdot(e_j) - div_M[i] += M_plus[i] - M_minus[i] + div_M = np.zeros(size, dtype=positions.dtype) + for j in range(3 * N): + pos_plus = positions.copy() + pos_minus = positions.copy() + pos_plus[j] += 0.5 * delta + pos_minus[j] -= 0.5 * delta + + e_j = np.zeros(size, dtype=positions.dtype) + e_j[j] = 1.0 + f = e_j[0 : 3 * N] if includeAngular else e_j + t = e_j[3 * N :] if includeAngular else None + + solver.setPositions(pos_plus) + Mf_plus, Mt_plus = solver.Mdot(forces=f, torques=t) + + solver.setPositions(pos_minus) + Mf_minus, Mt_minus = solver.Mdot(forces=f, torques=t) + + div_M[0 : 3 * N] += Mf_plus - Mf_minus + if includeAngular: + div_M[3 * N :] += Mt_plus - Mt_minus div_M /= delta return div_M @@ -111,35 +122,31 @@ def rfd_func(): @pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) def test_deterministic_divM_matches_rfd(Solver, periodicity): solver = Solver(*periodicity) + includeAngular = True parameters = get_sane_params(Solver.__name__, periodicity[2]) solver.setParameters(**parameters) precision = np.float32 if solver.precision == "float" else np.float64 solver.initialize( viscosity=1.0, hydrodynamicRadius=1.0, - includeAngular=False, + includeAngular=includeAngular, ) numberParticles = 10 positions = generate_positions_in_box(parameters, numberParticles).astype(precision) solver.setPositions(positions) - delta = 1e-3 - det_div_m = deterministic_div_m(solver, positions, delta) + det_div_m = deterministic_div_m(solver, positions, includeAngular) # rfd_dM, _ = divM_rfd(solver, positions, delta) # rfd_m, rfd_d = average(solver.divM, 10000) - m, d = average_with_error(solver.divM, 10000, det_div_m, numberParticles) - - breakpoint() - print("-----------------------") - print("Deterministic divM:", det_div_m) - print("RFD divM:", rfd_m) + rfd_m, rfd_d = average_with_error( + solver.divM, 10000, det_div_m, numberParticles, includeAngular + ) - assert np.allclose( - det_div_m, - rfd_m, - atol=1e-3, - rtol=1e-3, - ), f"Deterministic divM does not match RFD divM: {np.max(np.abs(det_div_m - rfd_m))}" + assert np.allclose(det_div_m[0 : 3 * numberParticles], rfd_m, atol=1e-3, rtol=1e-3) + if includeAngular: + assert np.allclose( + det_div_m[3 * numberParticles :], rfd_d, atol=1e-3, rtol=1e-3 + ) @pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) From a189007835cb0115f2c55ed8a7e7f7f9be8d31b0 Mon Sep 17 00:00:00 2001 From: Ryker Fish Date: Fri, 21 Nov 2025 15:20:17 -0700 Subject: [PATCH 3/6] feat: add RFD delta to interface for NBody and DPStokes --- .../random_finite_differences.h | 3 +- solvers/DPStokes/extra/uammd_interface.h | 1 + solvers/DPStokes/mobility.h | 3 +- solvers/DPStokes/python_wrapper.cu | 32 +++++++++++-------- solvers/NBody/mobility.h | 5 ++- solvers/NBody/python_wrapper.cu | 8 +++-- tests/utils.py | 5 +-- 7 files changed, 34 insertions(+), 23 deletions(-) diff --git a/include/MobilityInterface/random_finite_differences.h b/include/MobilityInterface/random_finite_differences.h index 22b580a..4722661 100644 --- a/include/MobilityInterface/random_finite_differences.h +++ b/include/MobilityInterface/random_finite_differences.h @@ -31,8 +31,7 @@ template void random_finite_differences(Mdot mdot, device_span positions, device_span ilinear, device_span iangular, uint seed, - real prefactor = 1) { - constexpr real delta = 1e-3; + const real delta, real prefactor = 1) { const uint N = ilinear.size() / 3; using device_vector = thrust::device_vector>; diff --git a/solvers/DPStokes/extra/uammd_interface.h b/solvers/DPStokes/extra/uammd_interface.h index c7e0057..1139320 100644 --- a/solvers/DPStokes/extra/uammd_interface.h +++ b/solvers/DPStokes/extra/uammd_interface.h @@ -31,6 +31,7 @@ struct PyParameters { real zmin, zmax; // Tolerance will be ignored in DP mode, TP will use only tolerance and nxy/nz real tolerance = 1e-5; + real delta = 1e-3; // RFD step size real w, w_d; real hydrodynamicRadius = -1; real3 beta = {-1.0, -1.0, -1.0}; diff --git a/solvers/DPStokes/mobility.h b/solvers/DPStokes/mobility.h index 7da91fc..0e96ac1 100644 --- a/solvers/DPStokes/mobility.h +++ b/solvers/DPStokes/mobility.h @@ -223,7 +223,8 @@ class DPStokes : public libmobility::Mobility { device_vector thermal_drift_m(ilinear.size(), 0); device_vector thermal_drift_d(iangular.size(), 0); libmobility::random_finite_differences(mdot, original_pos, thermal_drift_m, - thermal_drift_d, seed, prefactor); + thermal_drift_d, seed, + this->dppar.delta, prefactor); this->setPositions(original_pos); thrust::transform(thrust::cuda::par, thermal_drift_m.begin(), thermal_drift_m.end(), linear.begin(), linear.begin(), diff --git a/solvers/DPStokes/python_wrapper.cu b/solvers/DPStokes/python_wrapper.cu index 4cf4e59..5c14028 100644 --- a/solvers/DPStokes/python_wrapper.cu +++ b/solvers/DPStokes/python_wrapper.cu @@ -22,6 +22,8 @@ zmax : float The maximum value of the z coordinate. This is the position of the top wall if the Z periodicity is `two_walls`. 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. Default is 1e-3, units of length. )pbdoc"; static const char *docstring = R"pbdoc( @@ -35,18 +37,20 @@ The periodicity must be set to `periodic` in the X and Y directions. The Z perio )pbdoc"; MOBILITY_PYTHONIFY_WITH_EXTRA_CODE( - DPStokes, solver.def( - "setParameters", - [](DPStokes &self, real Lx, real Ly, real zmin, real zmax, - bool allowChangingBoxSize) { - DPStokesParameters params; - params.Lx = Lx; - params.Ly = Ly; - params.zmin = zmin; - params.zmax = zmax; - params.allowChangingBoxSize = allowChangingBoxSize; - self.setParametersDPStokes(params); - }, - "Lx"_a, "Ly"_a, "zmin"_a, "zmax"_a, - "allowChangingBoxSize"_a = false, setparameters_docstring); + DPStokes, + solver.def( + "setParameters", + [](DPStokes &self, real Lx, real Ly, real zmin, real zmax, + bool allowChangingBoxSize, real delta) { + DPStokesParameters params; + params.Lx = Lx; + params.Ly = Ly; + params.zmin = zmin; + params.zmax = zmax; + params.allowChangingBoxSize = allowChangingBoxSize; + params.delta = delta; + self.setParametersDPStokes(params); + }, + "Lx"_a, "Ly"_a, "zmin"_a, "zmax"_a, "allowChangingBoxSize"_a = false, + "delta"_a = 1e-3, setparameters_docstring); , docstring); diff --git a/solvers/NBody/mobility.h b/solvers/NBody/mobility.h index 6ead265..70588fd 100644 --- a/solvers/NBody/mobility.h +++ b/solvers/NBody/mobility.h @@ -28,6 +28,7 @@ class NBody : public libmobility::Mobility { nbody_rpy::algorithm algorithm = nbody_rpy::algorithm::advise; real wallHeight; // location of the wall in z + real delta; // finite difference step size for random finite differences // Batched functionality configuration int Nbatch = -1; @@ -53,6 +54,7 @@ class NBody : public libmobility::Mobility { int Nbatch = -1; int NperBatch = -1; std::optional wallHeight = std::nullopt; + real delta = 1e-3; }; /** * @brief Sets the parameters for the N-body computation @@ -196,7 +198,8 @@ class NBody : public libmobility::Mobility { device_vector thermal_drift_d(iangular.size(), 0); libmobility::random_finite_differences(mdot, original_pos, thermal_drift_m, - thermal_drift_d, seed, prefactor); + thermal_drift_d, seed, this->delta, + prefactor); device_adapter linear(ilinear, device::cuda); this->setPositions(original_pos); thrust::transform(thrust::cuda::par, thermal_drift_m.begin(), diff --git a/solvers/NBody/python_wrapper.cu b/solvers/NBody/python_wrapper.cu index add597d..c01f25a 100644 --- a/solvers/NBody/python_wrapper.cu +++ b/solvers/NBody/python_wrapper.cu @@ -17,6 +17,8 @@ static const char *docstringSetParameters = R"pbdoc( The number of particles per batch. If -1 (default), the number of particles per batch is automatically determined. wallHeight : float The height of the wall. Only valid if periodicityZ is single_wall. + delta : float + The finite difference step size for random finite differences. Default is 1e-3, units of length. )pbdoc"; static const char *docstring = R"pbdoc( @@ -50,10 +52,10 @@ MOBILITY_PYTHONIFY_WITH_EXTRA_CODE( solver.def( "setParameters", [](NBody &myself, std::string algo, int NBatch, int NperBatch, - std::optional wallHeight) { + std::optional wallHeight, real delta) { myself.setParametersNBody({nbody_rpy::string2NBodyAlgorithm(algo), - NBatch, NperBatch, wallHeight}); + NBatch, NperBatch, wallHeight, delta}); }, docstringSetParameters, "algorithm"_a = "advise", "Nbatch"_a = -1, - "NperBatch"_a = -1, "wallHeight"_a = std::nullopt); + "NperBatch"_a = -1, "wallHeight"_a = std::nullopt, "delta"_a = 1e-3); , docstring); diff --git a/tests/utils.py b/tests/utils.py index 5e95380..5b414f0 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -4,14 +4,15 @@ sane_parameters = { "PSE": {"psi": 1.0, "Lx": 32, "Ly": 32, "Lz": 32, "shearStrain": 0.0}, - "NBody": {"algorithm": "advise"}, - "NBody_wall": {"algorithm": "advise", "wallHeight": 0.0}, + "NBody": {"algorithm": "advise", "delta": 1e-3}, + "NBody_wall": {"algorithm": "advise", "wallHeight": 0.0, "delta": 1e-3}, "DPStokes": { "Lx": 16, "Ly": 16, "zmin": -6, "zmax": 6, "allowChangingBoxSize": False, + "delta": 1e-3, }, "SelfMobility": {"parameter": 5.0}, } From fd6dec2b7eb95419c091529d66c61a7c1ff98498 Mon Sep 17 00:00:00 2001 From: Ryker Fish Date: Sat, 22 Nov 2025 16:19:28 -0700 Subject: [PATCH 4/6] set delta in units of hydrodynamicRadius since it has units of length. the RFD step size should be small with respect to the particle, so delta=1e-3*a gives better convergence for small particle radii instead of simply using delta=1e-3 always --- solvers/DPStokes/mobility.h | 6 +++--- solvers/DPStokes/python_wrapper.cu | 2 +- solvers/NBody/mobility.h | 9 +++++---- solvers/NBody/python_wrapper.cu | 28 ++++++++++++++-------------- 4 files changed, 23 insertions(+), 22 deletions(-) diff --git a/solvers/DPStokes/mobility.h b/solvers/DPStokes/mobility.h index 0e96ac1..876179e 100644 --- a/solvers/DPStokes/mobility.h +++ b/solvers/DPStokes/mobility.h @@ -222,9 +222,9 @@ class DPStokes : public libmobility::Mobility { uint seed = rng(); device_vector thermal_drift_m(ilinear.size(), 0); device_vector thermal_drift_d(iangular.size(), 0); - libmobility::random_finite_differences(mdot, original_pos, thermal_drift_m, - thermal_drift_d, seed, - this->dppar.delta, prefactor); + libmobility::random_finite_differences( + mdot, original_pos, thermal_drift_m, thermal_drift_d, seed, + this->dppar.delta * this->dppar.hydrodynamicRadius, prefactor); this->setPositions(original_pos); thrust::transform(thrust::cuda::par, thermal_drift_m.begin(), thermal_drift_m.end(), linear.begin(), linear.begin(), diff --git a/solvers/DPStokes/python_wrapper.cu b/solvers/DPStokes/python_wrapper.cu index 5c14028..402e118 100644 --- a/solvers/DPStokes/python_wrapper.cu +++ b/solvers/DPStokes/python_wrapper.cu @@ -23,7 +23,7 @@ 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. Default is 1e-3, units of length. + The finite difference step size for random finite differences. Specified in units of hydrodynamicRadius. Default is 1e-3. )pbdoc"; static const char *docstring = R"pbdoc( diff --git a/solvers/NBody/mobility.h b/solvers/NBody/mobility.h index 70588fd..13cdeb2 100644 --- a/solvers/NBody/mobility.h +++ b/solvers/NBody/mobility.h @@ -54,7 +54,7 @@ class NBody : public libmobility::Mobility { int Nbatch = -1; int NperBatch = -1; std::optional wallHeight = std::nullopt; - real delta = 1e-3; + real delta = 1e-3; // in units of particle radius }; /** * @brief Sets the parameters for the N-body computation @@ -96,6 +96,7 @@ class NBody : public libmobility::Mobility { "you want to use a wall, set periodicityZ to single_wall in the " "configuration."); } + this->delta = par.delta; } void initialize(Parameters ipar) override { @@ -197,9 +198,9 @@ class NBody : public libmobility::Mobility { device_vector thermal_drift_m(ilinear.size(), 0); device_vector thermal_drift_d(iangular.size(), 0); - libmobility::random_finite_differences(mdot, original_pos, thermal_drift_m, - thermal_drift_d, seed, this->delta, - prefactor); + libmobility::random_finite_differences( + mdot, original_pos, thermal_drift_m, thermal_drift_d, seed, + this->delta * this->hydrodynamicRadius, prefactor); device_adapter linear(ilinear, device::cuda); this->setPositions(original_pos); thrust::transform(thrust::cuda::par, thermal_drift_m.begin(), diff --git a/solvers/NBody/python_wrapper.cu b/solvers/NBody/python_wrapper.cu index c01f25a..7d6c31f 100644 --- a/solvers/NBody/python_wrapper.cu +++ b/solvers/NBody/python_wrapper.cu @@ -5,21 +5,21 @@ #include static const char *docstringSetParameters = R"pbdoc( - Set the parameters for the NBody solver. +Set the parameters for the NBody solver. - Parameters - ---------- - algorithm : str - The algorithm to use. Options are "naive", "fast", "block" and "advise". Default is "advise". - NBatch : int - The number of batches to use. If -1 (default), the number of batches is automatically determined. - NperBatch : int - The number of particles per batch. If -1 (default), the number of particles per batch is automatically determined. - wallHeight : float - The height of the wall. Only valid if periodicityZ is single_wall. - delta : float - The finite difference step size for random finite differences. Default is 1e-3, units of length. - )pbdoc"; +Parameters +---------- +algorithm : str + The algorithm to use. Options are "naive", "fast", "block" and "advise". Default is "advise". +NBatch : int + The number of batches to use. If -1 (default), the number of batches is automatically determined. +NperBatch : int + The number of particles per batch. If -1 (default), the number of particles per batch is automatically determined. +wallHeight : float + The height of the wall. Only valid if periodicityZ is single_wall. +delta : float + The finite difference step size for random finite differences. Specified in units of hydrodynamicRadius. Default is 1e-3. +)pbdoc"; static const char *docstring = R"pbdoc( This module computes hydrodynamic interactions using an :math:`O(N^2)` algorithm. From 7a56048e4e266609246fe18dd072c1a16c7ebd23 Mon Sep 17 00:00:00 2001 From: Ryker Fish Date: Sun, 23 Nov 2025 15:46:39 -0700 Subject: [PATCH 5/6] add early exit condition when averaging RFDs. set threshold relatively high for NBody since it can sometimes take a long time to converge --- tests/test_thermal_drift.py | 72 ++++++++++++++++++++++++++----------- 1 file changed, 52 insertions(+), 20 deletions(-) diff --git a/tests/test_thermal_drift.py b/tests/test_thermal_drift.py index 9f71abd..f443474 100644 --- a/tests/test_thermal_drift.py +++ b/tests/test_thermal_drift.py @@ -21,8 +21,8 @@ def average(function, num_averages): return result_m / num_averages, result_d -def average_with_error(function, num_averages, soln, N, includeAngular): - # track error of the running average of RFDs after each new sample +def average_until_error_tolerance(function, num_averages, soln, N, includeAngular, tol): + # track error of the running average of RFDs after each new sample. exit early if the error is below the tolerance avg_err_m = np.zeros(num_averages) rfd_m = np.zeros(3 * N) avg_err_d = np.zeros(num_averages) @@ -31,6 +31,10 @@ def average_with_error(function, num_averages, soln, N, includeAngular): soln_m = soln[0 : 3 * N] soln_d = soln[3 * N :] if includeAngular else None + err_check_iter = 0 + check_every = 5000 + exit_early = False + last_iter = -1 for i in range(num_averages): new_result_m, new_result_d = function() rfd_m += new_result_m @@ -41,16 +45,41 @@ def average_with_error(function, num_averages, soln, N, includeAngular): running_d = rfd_d / float(i + 1) avg_err_d[i] = np.linalg.norm((running_d - soln_d)) + if err_check_iter == check_every: + err_m = np.allclose(running_m, soln_m, atol=tol, rtol=tol) + errs = [err_m] + if includeAngular: + err_d = np.allclose(running_d, soln_d, atol=tol, rtol=tol) + errs.append(err_d) + if all(errs): + last_iter = i + exit_early = True + break + err_check_iter = 0 + err_check_iter += 1 + + if exit_early: + num_averages = last_iter + 1 + rfd_m /= float(num_averages) rfd_d /= float(num_averages) + avg_err_m = avg_err_m[:num_averages] + avg_err_d = avg_err_d[:num_averages] + + # plot_error(avg_err_m, avg_err_d, num_averages, includeAngular, delta, a) + + return rfd_m, rfd_d + + +def plot_error(avg_err_m, avg_err_d, num_averages, includeAngular): import matplotlib.pyplot as plt plt.figure() - n = np.arange(num_averages) + n = np.arange(1, num_averages + 1) plt.loglog(n, avg_err_m, marker="o", linestyle="-") plt.loglog(n, avg_err_d, marker="o", linestyle="-") - plt.loglog(n, 0.1 * 1 / np.sqrt(n + 1), linestyle="--", color="black") + plt.loglog(n, 0.1 * 1 / np.sqrt(n), linestyle="--", color="black") if includeAngular: plt.legend( ["Measured error (linear)", "Measured error (dipole)", "O(1/sqrt(N))"] @@ -62,13 +91,11 @@ def average_with_error(function, num_averages, soln, N, includeAngular): plt.title("Error of running average of RFDs compared to deterministic divM") plt.grid(True) plt.tight_layout() - plt.savefig("avg_err_m.png", dpi=150) + plt.savefig("rfd_error.png", dpi=150) plt.close() - return rfd_m, rfd_d -def deterministic_div_m(solver, positions, includeAngular): - delta = 1e-3 +def deterministic_div_m(solver, positions, includeAngular, delta): N = np.size(positions) // 3 size = 6 * N if includeAngular else 3 * N @@ -120,33 +147,38 @@ def rfd_func(): @pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) -def test_deterministic_divM_matches_rfd(Solver, periodicity): +@pytest.mark.parametrize("includeAngular", [True, False]) +@pytest.mark.parametrize("a", [0.4, 1.0, 2.5]) +def test_deterministic_divM_matches_rfd(Solver, periodicity, a, includeAngular): + if Solver.__name__ == "PSE" and includeAngular: + pytest.skip("PSE does not support torques") + N_iter = 100001 if Solver.__name__ == "NBody" else 25001 + tol = 1e-2 if Solver.__name__ == "DPStokes" else 1e-3 solver = Solver(*periodicity) - includeAngular = True parameters = get_sane_params(Solver.__name__, periodicity[2]) + delta = 1e-3 + if "delta" in parameters: + parameters["delta"] = delta solver.setParameters(**parameters) precision = np.float32 if solver.precision == "float" else np.float64 solver.initialize( viscosity=1.0, - hydrodynamicRadius=1.0, + hydrodynamicRadius=a, includeAngular=includeAngular, ) numberParticles = 10 positions = generate_positions_in_box(parameters, numberParticles).astype(precision) solver.setPositions(positions) - det_div_m = deterministic_div_m(solver, positions, includeAngular) - # rfd_dM, _ = divM_rfd(solver, positions, delta) - # rfd_m, rfd_d = average(solver.divM, 10000) - rfd_m, rfd_d = average_with_error( - solver.divM, 10000, det_div_m, numberParticles, includeAngular + # use delta*a in deterministic since libMobility solvers expect delta to be in units of a + det_div_m = deterministic_div_m(solver, positions, includeAngular, delta * a) + rfd_m, rfd_d = average_until_error_tolerance( + solver.divM, N_iter, det_div_m, numberParticles, includeAngular, tol ) - assert np.allclose(det_div_m[0 : 3 * numberParticles], rfd_m, atol=1e-3, rtol=1e-3) + assert np.allclose(det_div_m[0 : 3 * numberParticles], rfd_m, atol=tol, rtol=tol) if includeAngular: - assert np.allclose( - det_div_m[3 * numberParticles :], rfd_d, atol=1e-3, rtol=1e-3 - ) + assert np.allclose(det_div_m[3 * numberParticles :], rfd_d, atol=tol, rtol=tol) @pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) From eb3090cc474ff221bd219356ded8a4c13c7d5f8f Mon Sep 17 00:00:00 2001 From: Ryker Fish Date: Sun, 23 Nov 2025 19:18:13 -0700 Subject: [PATCH 6/6] reoreganize test files --- tests/test_thermal_drift.py | 301 +++++++++++------------------------- 1 file changed, 93 insertions(+), 208 deletions(-) diff --git a/tests/test_thermal_drift.py b/tests/test_thermal_drift.py index f443474..bde3e27 100644 --- a/tests/test_thermal_drift.py +++ b/tests/test_thermal_drift.py @@ -1,24 +1,105 @@ import pytest import numpy as np -from libMobility import SelfMobility, PSE, NBody, DPStokes from utils import ( solver_configs_all, - solver_configs_torques, get_sane_params, generate_positions_in_box, ) -def average(function, num_averages): - result_m, result_d = function() - for i in range(num_averages - 1): - new_result_m, new_result_d = function() - result_m += new_result_m - if result_d is not None: - result_d += new_result_d - if result_d is not None: - result_d /= num_averages - return result_m / num_averages, result_d +@pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) +@pytest.mark.parametrize("includeAngular", [True, False]) +@pytest.mark.parametrize("a", [0.4, 1.0, 2.5]) +def test_deterministic_divM_matches_rfd(Solver, periodicity, a, includeAngular): + if Solver.__name__ == "PSE" and includeAngular: + pytest.skip("PSE does not support torques") + N_iter = 250001 if Solver.__name__ == "NBody" else 25001 + tol = 1e-2 if Solver.__name__ == "DPStokes" else 1e-3 + solver = Solver(*periodicity) + parameters = get_sane_params(Solver.__name__, periodicity[2]) + delta = 1e-3 + if "delta" in parameters: + parameters["delta"] = delta + solver.setParameters(**parameters) + precision = np.float32 if solver.precision == "float" else np.float64 + solver.initialize( + viscosity=1.0, + hydrodynamicRadius=a, + includeAngular=includeAngular, + ) + numberParticles = 10 + positions = generate_positions_in_box(parameters, numberParticles).astype(precision) + solver.setPositions(positions) + + # use delta*a in deterministic since libMobility solvers expect delta to be in units of a + det_div_m = deterministic_div_m(solver, positions, includeAngular, delta * a) + rfd_m, rfd_d = average_until_error_tolerance( + solver.divM, N_iter, det_div_m, numberParticles, includeAngular, tol + ) + + assert np.allclose(det_div_m[0 : 3 * numberParticles], rfd_m, atol=tol, rtol=tol) + if includeAngular: + assert np.allclose(det_div_m[3 * numberParticles :], rfd_d, atol=tol, rtol=tol) + + +@pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) +def test_divM_does_not_change_positions(Solver, periodicity): + # Compute the Mdot with some random forces, then compute divM, then compute Mdot again + # Check that Mdot has not changed + includeAngular = False + precision = np.float32 if Solver.precision == "float" else np.float64 + solver = Solver(*periodicity) + parameters = get_sane_params(Solver.__name__, periodicity[2]) + solver.setParameters(**parameters) + solver.initialize( + viscosity=1.0, + hydrodynamicRadius=1.0, + includeAngular=includeAngular, + ) + positions = generate_positions_in_box(parameters, 10).astype(precision) + solver.setPositions(positions) + forces = np.random.normal(size=positions.shape).astype(precision) + mf, _ = solver.Mdot(forces) + solver.divM() + mf2, _ = solver.Mdot(forces) + assert np.allclose( + mf, + mf2, + atol=1e-7, + rtol=1e-7, + ), f"Mdot has changed: {np.max(np.abs(mf - mf2))}" + + +@pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) +@pytest.mark.parametrize("hydrodynamicRadius", [1.0, 0.95, 1.12]) +@pytest.mark.parametrize("numberParticles", [1, 2, 3, 10]) +@pytest.mark.parametrize("includeAngular", [False, True]) +def test_divM_returns_different_numbers( + Solver, periodicity, hydrodynamicRadius, numberParticles, includeAngular +): + if Solver.__name__ == "PSE" and includeAngular: + pytest.skip("PSE does not support torques") + + precision = np.float32 if Solver.precision == "float" else np.float64 + solver = Solver(*periodicity) + parameters = get_sane_params(Solver.__name__, periodicity[2]) + solver.setParameters(**parameters) + solver.initialize( + viscosity=1.0, + hydrodynamicRadius=hydrodynamicRadius, + includeAngular=includeAngular, + ) + positions = np.asarray( + generate_positions_in_box(parameters, numberParticles).astype(precision) * 0.8 + ) + solver.setPositions(positions) + rfd1, _ = solver.divM() + if np.all(rfd1 == 0): + pytest.skip("RFD is zero, skipping test") + rfd2, _ = solver.divM() + assert np.any( + np.abs(rfd1 - rfd2) > 1e-5 + ), f"RFD is not different: {np.max(np.abs(rfd1 - rfd2))}" def average_until_error_tolerance(function, num_averages, soln, N, includeAngular, tol): @@ -124,199 +205,3 @@ def deterministic_div_m(solver, positions, includeAngular, delta): div_M /= delta return div_M - - -def divM_rfd(solver, positions, delta): - # RFD works by approximating \partial_q \dot M = 1/\delta \langle M(q+\delta/2 W)W - M(q-\delta/2 W)W \rangle - - def rfd_func(): - W = np.random.normal(size=positions.shape).astype(positions.dtype) - solver.setPositions(positions + delta / 2 * W) - _tdriftp_m, _tdriftp_d = solver.Mdot(W) - solver.setPositions(positions - delta / 2 * W) - _tdriftm_m, _tdriftm_d = solver.Mdot(W) - _tdrift_m = (_tdriftp_m - _tdriftm_m) / delta - _tdrift_d = ( - (_tdriftp_d - _tdriftm_d) / delta if _tdriftm_d is not None else None - ) - return _tdrift_m, _tdrift_d - - solver.setPositions(positions) - tdrift_m, tdrift_d = average(lambda: average(rfd_func, 400), 100) - return tdrift_m, tdrift_d - - -@pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) -@pytest.mark.parametrize("includeAngular", [True, False]) -@pytest.mark.parametrize("a", [0.4, 1.0, 2.5]) -def test_deterministic_divM_matches_rfd(Solver, periodicity, a, includeAngular): - if Solver.__name__ == "PSE" and includeAngular: - pytest.skip("PSE does not support torques") - N_iter = 100001 if Solver.__name__ == "NBody" else 25001 - tol = 1e-2 if Solver.__name__ == "DPStokes" else 1e-3 - solver = Solver(*periodicity) - parameters = get_sane_params(Solver.__name__, periodicity[2]) - delta = 1e-3 - if "delta" in parameters: - parameters["delta"] = delta - solver.setParameters(**parameters) - precision = np.float32 if solver.precision == "float" else np.float64 - solver.initialize( - viscosity=1.0, - hydrodynamicRadius=a, - includeAngular=includeAngular, - ) - numberParticles = 10 - positions = generate_positions_in_box(parameters, numberParticles).astype(precision) - solver.setPositions(positions) - - # use delta*a in deterministic since libMobility solvers expect delta to be in units of a - det_div_m = deterministic_div_m(solver, positions, includeAngular, delta * a) - rfd_m, rfd_d = average_until_error_tolerance( - solver.divM, N_iter, det_div_m, numberParticles, includeAngular, tol - ) - - assert np.allclose(det_div_m[0 : 3 * numberParticles], rfd_m, atol=tol, rtol=tol) - if includeAngular: - assert np.allclose(det_div_m[3 * numberParticles :], rfd_d, atol=tol, rtol=tol) - - -@pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) -def test_divM_does_not_change_positions(Solver, periodicity): - # Compute the Mdot with some random forces, then compute divM, then compute Mdot again - # Check that Mdot has not changed - includeAngular = False - precision = np.float32 if Solver.precision == "float" else np.float64 - solver = Solver(*periodicity) - parameters = get_sane_params(Solver.__name__, periodicity[2]) - solver.setParameters(**parameters) - solver.initialize( - viscosity=1.0, - hydrodynamicRadius=1.0, - includeAngular=includeAngular, - ) - positions = generate_positions_in_box(parameters, 10).astype(precision) - solver.setPositions(positions) - forces = np.random.normal(size=positions.shape).astype(precision) - mf, _ = solver.Mdot(forces) - solver.divM() - mf2, _ = solver.Mdot(forces) - assert np.allclose( - mf, - mf2, - atol=1e-7, - rtol=1e-7, - ), f"Mdot has changed: {np.max(np.abs(mf - mf2))}" - - -@pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) -@pytest.mark.parametrize("hydrodynamicRadius", [1.0, 0.95, 1.12]) -@pytest.mark.parametrize("numberParticles", [1, 2, 3, 10]) -@pytest.mark.parametrize("includeAngular", [False, True]) -def test_divM_is_zero( - Solver, periodicity, hydrodynamicRadius, numberParticles, includeAngular -): - if not np.all(np.array(periodicity) == "open") and not np.all( - np.array(periodicity) == "periodic" - ): - pytest.skip("Only periodic and open boundary conditions have zero divergence") - if Solver.__name__ == "PSE" and includeAngular: - pytest.skip("PSE does not support torques") - precision = np.float32 if Solver.precision == "float" else np.float64 - solver = Solver(*periodicity) - parameters = get_sane_params(Solver.__name__, periodicity[2]) - solver.setParameters(**parameters) - solver.initialize( - viscosity=1.0, - hydrodynamicRadius=hydrodynamicRadius, - includeAngular=includeAngular, - ) - positions = generate_positions_in_box(parameters, numberParticles).astype(precision) - solver.setPositions(positions) - thermal_drift_m, thermal_drift_d = average(solver.divM, 3000) - assert np.allclose( - np.abs(thermal_drift_m), - 0.0, - atol=1e-5, - rtol=0, - ), f"Linear RFD drift is not zero: {np.max(np.abs(thermal_drift_m))}" - if thermal_drift_d is not None: - assert np.allclose( - np.abs(thermal_drift_d), - 0.0, - atol=1e-5, - rtol=0, - ), f"Dipolar RFD drift is not zero: {np.max(np.abs(thermal_drift_d))}" - - -@pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) -@pytest.mark.parametrize("hydrodynamicRadius", [1.0, 0.95, 1.12]) -@pytest.mark.parametrize("numberParticles", [1, 2, 3, 10]) -@pytest.mark.parametrize("includeAngular", [False, True]) -def test_divM_returns_different_numbers( - Solver, periodicity, hydrodynamicRadius, numberParticles, includeAngular -): - if Solver.__name__ == "PSE" and includeAngular: - pytest.skip("PSE does not support torques") - - precision = np.float32 if Solver.precision == "float" else np.float64 - solver = Solver(*periodicity) - parameters = get_sane_params(Solver.__name__, periodicity[2]) - solver.setParameters(**parameters) - solver.initialize( - viscosity=1.0, - hydrodynamicRadius=hydrodynamicRadius, - includeAngular=includeAngular, - ) - positions = np.asarray( - generate_positions_in_box(parameters, numberParticles).astype(precision) * 0.8 - ) - solver.setPositions(positions) - rfd1, _ = solver.divM() - if np.all(rfd1 == 0): - pytest.skip("RFD is zero, skipping test") - rfd2, _ = solver.divM() - assert np.any( - np.abs(rfd1 - rfd2) > 1e-5 - ), f"RFD is not different: {np.max(np.abs(rfd1 - rfd2))}" - - -@pytest.mark.parametrize(("Solver", "periodicity"), solver_configs_all) -@pytest.mark.parametrize("hydrodynamicRadius", [0.95]) -@pytest.mark.parametrize("numberParticles", [1, 10]) -@pytest.mark.parametrize("includeAngular", [False, True]) -def test_divM_matches_rfd( - Solver, periodicity, hydrodynamicRadius, numberParticles, includeAngular -): - if Solver.__name__ == "PSE" and includeAngular: - pytest.skip("PSE does not support torques") - precision = np.float32 if Solver.precision == "float" else np.float64 - solver = Solver(*periodicity) - parameters = get_sane_params(Solver.__name__, periodicity[2]) - solver.setParameters(**parameters) - solver.initialize( - viscosity=1.0, - hydrodynamicRadius=hydrodynamicRadius, - includeAngular=includeAngular, - ) - positions = np.asarray( - generate_positions_in_box(parameters, numberParticles).astype(precision) - ) - solver.setPositions(positions) - reference_m, reference_d = divM_rfd(solver, positions) - - solver.setPositions(positions) - rfd_m, rfd_d = average(lambda: average(solver.divM, 400), 100) - assert np.allclose( - reference_m, - rfd_m, - atol=1e-3, - rtol=1e-3, - ), f"Linear RFD does not match: {np.max(np.abs(reference_m - rfd_m))}" - if reference_d is not None and rfd_d is not None: - assert np.allclose( - reference_d, - rfd_d, - atol=1e-3, - rtol=1e-3, - ), f"Dipole RFD does not match: {np.max(np.abs(reference_d - rfd_d))}"