SD: Test CPU/GPU, FT/FTS, checkpointing, sample

samples/dancing.py

@@ -16,34 +16,86 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
-
+"""
+Stokesian Dynamics simulation of particle sedimentation.
+Reproduce the trajectory in Figure 5b from :cite:`durlofsky87a`.
+"""
 import espressomd
-from espressomd import constraints
+import espressomd.constraints
+import espressomd.observables
+import espressomd.accumulators
 import numpy as np
 import matplotlib.pyplot as plt
 
-espressomd.assert_features("STOKESIAN_DYNAMICS")
+import argparse
+
+parser = argparse.ArgumentParser(epilog=__doc__)
+group = parser.add_mutually_exclusive_group()
+group.add_argument('--cpu', action='store_true', help='Use CPU implementation')
+group.add_argument('--gpu', action='store_true', help='Use GPU implementation')
+group = parser.add_mutually_exclusive_group()
+group.add_argument('--ft', action='store_true', help='Use FT approximation')
+group.add_argument('--fts', action='store_true', help='Use FTS approximation')
+args = parser.parse_args()
+
+required_features = ["STOKESIAN_DYNAMICS"]
+
+if args.gpu:
+    print("Using GPU implementation")
+    required_features.append("CUDA")
+    required_features.append("STOKESIAN_DYNAMICS_GPU")
+    sd_device = "gpu"
+else:
+    print("Using CPU implementation")
+    sd_device = "cpu"
+
+if args.ft:
+    print("Using FT approximation method")
+    sd_method = "ft"
+else:
+    print("Using FTS approximation method")
+    sd_method = "fts"
+
+espressomd.assert_features(required_features)
 
 system = espressomd.System(box_l=[10, 10, 10])
-system.time_step = 1.0
+system.time_step = 1.5
 system.cell_system.skin = 0.4
+system.periodicity = [False, False, False]
 
-system.integrator.set_sd(viscosity=1.0, radii={0: 1.0})
+system.integrator.set_sd(
+    viscosity=1.0, radii={0: 1.0}, device=sd_device,
+    approximation_method=sd_method)
 
 system.part.add(pos=[-5, 0, 0], rotation=[1, 1, 1])
 system.part.add(pos=[0, 0, 0], rotation=[1, 1, 1])
 system.part.add(pos=[7, 0, 0], rotation=[1, 1, 1])
 
-gravity = constraints.Gravity(g=[0, -1, 0])
+gravity = espressomd.constraints.Gravity(g=[0, -1, 0])
 system.constraints.add(gravity)
 
-intsteps = int(13000 / system.time_step)
-pos = np.empty([intsteps, 3 * len(system.part)])
-for i in range(intsteps):
-    system.integrator.run(1)
-    for n, p in enumerate(system.part):
-        pos[i, 3 * n:3 * n + 3] = p.pos
+obs = espressomd.observables.ParticlePositions(ids=system.part[:].id)
+acc = espressomd.accumulators.TimeSeries(obs=obs, delta_N=1)
+system.auto_update_accumulators.add(acc)
+acc.update()
+intsteps = int(10500 / system.time_step)
+system.integrator.run(intsteps)
+
+positions = acc.time_series()
+ref_data = "../testsuite/python/data/dancing.txt"
+data = np.loadtxt(ref_data)
+
+for i in range(3):
+    plt.plot(positions[:, i, 0], positions[:, i, 1], linestyle='solid')
+
+plt.gca().set_prop_cycle(None)
+
+for i in range(0, 6, 2):
+    plt.plot(data[:, i], data[:, i + 1], linestyle='dashed')
 
-for n, p in enumerate(system.part):
-    plt.plot(pos[:, 3 * n], pos[:, 3 * n + 1])
+plt.title("Trajectory of sedimenting spheres\nsolid line: simulation "
+          "({} on {}), dashed line: paper (FTS)"
+          .format(sd_method.upper(), sd_device.upper()))
+plt.xlabel("x")
+plt.ylabel("y")
 plt.show()

testsuite/python/CMakeLists.txt

@@ -56,7 +56,7 @@ endfunction(PYTHON_TEST)
 # Separate features with hyphens, use a period to add an optional flag.
 foreach(
   TEST_COMBINATION
-  lb.cpu-p3m.cpu-lj-therm.lb;lb.gpu-p3m.cpu-lj-therm.lb;ek.gpu;lb.off-therm.npt-int.npt-minimization;lb.off-therm.langevin-int.nvt;lb.off-therm.dpd-int.sd;lb.off-therm.bd-int.bd
+  lb.cpu-p3m.cpu-lj-therm.lb;lb.gpu-p3m.cpu-lj-therm.lb;ek.gpu;lb.off-therm.npt-int.npt-minimization;lb.off-therm.langevin-int.nvt;lb.off-therm.dpd-int.sd;lb.off-therm.bd-int.bd;lb.off-therm.sdm-int.sdm.cpu;lb.off-therm.sdm-int.sdm.gpu
 )
   if(${TEST_COMBINATION} MATCHES "\\.gpu")
     set(TEST_LABELS "gpu")
@@ -200,10 +200,11 @@ python_test(FILE gpu_availability.py MAX_NUM_PROC 1 LABELS gpu)
 python_test(FILE features.py MAX_NUM_PROC 1)
 python_test(FILE galilei.py MAX_NUM_PROC 32)
 python_test(FILE time_series.py MAX_NUM_PROC 1)
-python_test(FILE linear_momentum.py)
+python_test(FILE linear_momentum.py MAX_NUM_PROC 4)
 python_test(FILE linear_momentum_lb.py MAX_NUM_PROC 2 LABELS gpu)
 python_test(FILE mmm1d.py MAX_NUM_PROC 2 LABELS gpu)
-python_test(FILE stokesian_dynamics.py)
+python_test(FILE stokesian_dynamics_cpu.py MAX_NUM_PROC 2)
+python_test(FILE stokesian_dynamics_gpu.py MAX_NUM_PROC 4 LABELS gpu)
 python_test(FILE elc.py MAX_NUM_PROC 2)
 python_test(FILE elc_vs_analytic.py MAX_NUM_PROC 2)
 python_test(FILE rotation.py MAX_NUM_PROC 1)
@@ -225,6 +226,9 @@ add_custom_target(
     ${CMAKE_CURRENT_BINARY_DIR}
   COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_SOURCE_DIR}/ek_common.py
           ${CMAKE_CURRENT_BINARY_DIR}
+  COMMAND
+    ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_SOURCE_DIR}/stokesian_dynamics.py
+    ${CMAKE_CURRENT_BINARY_DIR}
   COMMAND
     ${CMAKE_COMMAND} -E copy
     ${CMAKE_CURRENT_SOURCE_DIR}/ek_eof_one_species_base.py

testsuite/python/save_checkpoint.py

@@ -153,6 +153,9 @@
         system.thermostat.set_npt(kT=1.0, gamma0=2.0, gammav=0.1, seed=42)
     elif 'THERM.DPD' in modes and has_features('DPD'):
         system.thermostat.set_dpd(kT=1.0, seed=42)
+    elif 'THERM.SDM' in modes and (has_features('STOKESIAN_DYNAMICS') or has_features('STOKESIAN_DYNAMICS_GPU')):
+        system.periodicity = [0, 0, 0]
+        system.thermostat.set_sd(kT=1.0, seed=42)
     # set integrator
     if 'INT.NPT' in modes and has_features('NPT'):
         system.integrator.set_isotropic_npt(ext_pressure=2.0, piston=0.01,
@@ -164,6 +167,16 @@
         system.integrator.set_nvt()
     elif 'INT.BD' in modes:
         system.integrator.set_brownian_dynamics()
+    elif 'INT.SDM.CPU' in modes and has_features('STOKESIAN_DYNAMICS'):
+        system.periodicity = [0, 0, 0]
+        system.integrator.set_sd(approximation_method='ft', viscosity=0.5,
+                                 device='cpu', radii={0: 1.5},
+                                 pair_mobility=False, self_mobility=True)
+    elif 'INT.SDM.GPU' in modes and has_features('STOKESIAN_DYNAMICS_GPU') and espressomd.gpu_available():
+        system.periodicity = [0, 0, 0]
+        system.integrator.set_sd(approximation_method='fts', viscosity=2.0,
+                                 device='gpu', radii={0: 1.0},
+                                 pair_mobility=True, self_mobility=False)
     # set minimization
     if 'MINIMIZATION' in modes:
         steepest_descent(system, f_max=1, gamma=10, max_steps=0,

testsuite/python/stokesian_dynamics.py

@@ -17,57 +17,49 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 import espressomd
-from espressomd import constraints
+import espressomd.observables
+import espressomd.accumulators
+import espressomd.constraints
 import numpy as np
 import unittest as ut
-import unittest_decorators as utx
 from tests_common import abspath
 
 s = espressomd.System(box_l=[1.0, 1.0, 1.0])
 
 
-def skipIfMissingFeatureStokesianDynamics():
-    """Specialized unittest skipIf decorator for missing Stokesian Dynamics."""
-    if not espressomd.has_features(["STOKESIAN_DYNAMICS"]) and (not espressomd.has_features(
-            ["STOKESIAN_DYNAMICS_GPU"]) or not espressomd.gpu_available()):
-        return ut.skip("Skipping test: feature STOKESIAN_DYNAMICS unavailable")
-    return utx._id
-
-
-@skipIfMissingFeatureStokesianDynamics()
 class StokesianDynamicsSetupTest(ut.TestCase):
     system = s
+    device = 'none'
 
     def setUp(self):
         self.system.box_l = [10] * 3
 
         self.system.time_step = 1.0
         self.system.cell_system.skin = 0.4
 
-        # unset SD integrator so we can test whether set_sd fails
+        # unset SD integrator so we can test whether set_sd() fails.
         # set_nvt() is the only way to ensure that integ_switch is
         # set to a different value than INTEG_METHOD_SD
         self.system.integrator.set_nvt()
 
-    def test_pbc_checks(self):
-
+    def pbc_checks(self):
         self.system.periodicity = [0, 0, 1]
-        with (self.assertRaises(Exception)):
+        with self.assertRaises(Exception):
             self.system.integrator.set_sd(viscosity=1.0,
-                                          device="cpu",
+                                          device=self.device,
                                           radii={0: 1.0})
 
         self.system.periodicity = [0, 0, 0]
         self.system.integrator.set_sd(viscosity=1.0,
-                                      device="cpu",
+                                      device=self.device,
                                       radii={0: 1.0})
-        with (self.assertRaises(Exception)):
+        with self.assertRaises(Exception):
             self.system.periodicity = [0, 1, 0]
 
 
-@skipIfMissingFeatureStokesianDynamics()
 class StokesianDynamicsTest(ut.TestCase):
     system = s
+    device = 'none'
 
     # Digitized reference data of Figure 5b from
     # Durlofsky et al., J. Fluid Mech. 180, 21 (1987)
@@ -79,61 +71,66 @@ def setUp(self):
         self.system.periodicity = [0, 0, 0]
         self.system.cell_system.skin = 0.4
 
-    def falling_spheres(self, time_step, l_factor, t_factor):
+    def falling_spheres(self, time_step, l_factor, t_factor,
+                        sd_method='fts', sd_short=False):
         self.system.time_step = time_step
         self.system.part.add(pos=[-5 * l_factor, 0, 0], rotation=[1, 1, 1])
         self.system.part.add(pos=[0 * l_factor, 0, 0], rotation=[1, 1, 1])
         self.system.part.add(pos=[7 * l_factor, 0, 0], rotation=[1, 1, 1])
 
-        self.system.integrator.set_sd(viscosity=1.0 / (t_factor * l_factor),
-                                      device="cpu",
-                                      radii={0: 1.0 * l_factor})
+        self.system.integrator.set_sd(
+            viscosity=1.0 / (t_factor * l_factor),
+            device=self.device, radii={0: 1.0 * l_factor},
+            approximation_method=sd_method)
 
-        gravity = constraints.Gravity(
+        gravity = espressomd.constraints.Gravity(
             g=[0, -1.0 * l_factor / (t_factor**2), 0])
         self.system.constraints.add(gravity)
         self.system.time_step = 1.0 * t_factor
 
-        intsteps = int(8000 * t_factor / self.system.time_step)
-        pos = np.empty([intsteps + 1, 3 * len(self.system.part)])
-
-        pos[0, :] = self.system.part[:].pos.flatten()
-        for i in range(intsteps):
-            self.system.integrator.run(1)
-            for n, p in enumerate(self.system.part):
-                pos[i + 1, 3 * n:3 * n + 3] = p.pos
-
-        for i in range(self.data.shape[0]):
-            for n in range(3):
-                x_ref = self.data[i, 2 * n] * l_factor
-                y_ref = self.data[i, 2 * n + 1] * l_factor
-                x = pos[:, 3 * n]
-                y = pos[:, 3 * n + 1]
-
-                if y_ref < -555 * l_factor:
-                    continue
-
-                idx = np.abs(y - y_ref).argmin()
-                dist = np.sqrt((x_ref - x[idx])**2 + (y_ref - y[idx])**2)
-                self.assertLess(dist, 0.5 * l_factor)
-
-    def test_default(self):
-        self.falling_spheres(1.0, 1.0, 1.0)
-
-    def test_rescaled(self):
-        self.falling_spheres(1.0, 4.5, 2.5)
-
-    def test_different_time_step(self):
-        self.falling_spheres(0.7, 1.0, 1.0)
+        obs = espressomd.observables.ParticlePositions(ids=(0, 1, 2))
+        acc = espressomd.accumulators.TimeSeries(obs=obs, delta_N=1)
+        self.system.auto_update_accumulators.add(acc)
+        acc.update()
+
+        if sd_short:
+            y_min = -80
+            intsteps = 1300
+        elif sd_method == 'fts':
+            y_min = -555
+            intsteps = 8000
+        else:
+            y_min = -200
+            intsteps = 3000
+        intsteps = int(intsteps * t_factor / self.system.time_step)
+
+        self.system.integrator.run(intsteps)
+
+        simul = acc.time_series()[:, :, 0:2]
+        paper = self.data.reshape([-1, 3, 2])
+
+        for pid in range(3):
+            dist = []
+            # the simulated trajectory is oversampled by a ratio of
+            # (90/t_factor):1 compared to the published trajectory
+            for desired in paper[:, pid] * l_factor:
+                if desired[1] < y_min * l_factor:
+                    break
+                # find the closest point in the simulated trajectory
+                idx = np.abs(simul[:, pid, 1] - desired[1]).argmin()
+                actual = simul[idx, pid]
+                dist.append(np.linalg.norm(actual - desired))
+            self.assertLess(idx, intsteps, msg='Insufficient sampling')
+            np.testing.assert_allclose(dist, 0, rtol=0, atol=0.5 * l_factor)
 
     def tearDown(self):
         self.system.constraints.clear()
         self.system.part.clear()
 
 
-@skipIfMissingFeatureStokesianDynamics()
 class StokesianDiffusionTest(ut.TestCase):
     system = s
+    device = 'none'
 
     kT = 1e-4
     R = 1.5
@@ -148,12 +145,12 @@ def setUp(self):
 
         self.system.part.add(pos=[0, 0, 0], rotation=[1, 1, 1])
 
+    def check(self):
         self.system.integrator.set_sd(viscosity=self.eta,
-                                      device="cpu",
+                                      device=self.device,
                                       radii={0: self.R})
         self.system.thermostat.set_sd(kT=self.kT, seed=42)
 
-    def test(self):
         intsteps = int(100000 / self.system.time_step)
         pos = np.empty([intsteps + 1, 3])
         orientation = np.empty((intsteps + 1, 3))
@@ -206,7 +203,3 @@ def tearDown(self):
         self.system.constraints.clear()
         self.system.part.clear()
         self.system.thermostat.set_sd(kT=0)
-
-
-if __name__ == '__main__':
-    ut.main()