From 12cdd7cf8e20e397e8b633b415305ff9540e3126 Mon Sep 17 00:00:00 2001
From: Rudolf Weeber <weeber@icp.uni-stuttgart.de>
Date: Thu, 26 Aug 2021 11:36:54 +0200
Subject: [PATCH] Engine LB coupling with n square cell system

---
 src/core/grid_based_algorithms/lb.cpp         |  9 ++--
 .../lb_particle_coupling.cpp                  | 46 +++++++++++--------
 testsuite/python/engine_lb.py                 | 31 ++++++++++++-
 3 files changed, 62 insertions(+), 24 deletions(-)

diff --git a/src/core/grid_based_algorithms/lb.cpp b/src/core/grid_based_algorithms/lb.cpp
index 7cac4ed2a0d..3401746e2d0 100644
--- a/src/core/grid_based_algorithms/lb.cpp
+++ b/src/core/grid_based_algorithms/lb.cpp
@@ -614,9 +614,12 @@ void lb_sanity_checks(const LB_Parameters &lb_parameters) {
   if (lb_parameters.viscosity <= 0.0) {
     runtimeErrorMsg() << "Lattice Boltzmann fluid viscosity not set";
   }
-  if (local_geo.cell_structure_type() !=
-      CellStructureType::CELL_STRUCTURE_REGULAR) {
-    runtimeErrorMsg() << "LB requires regular-decomposition cellsystem";
+  if (local_geo.cell_structure_type() ==
+      CellStructureType::CELL_STRUCTURE_NSQUARE) {
+    if (n_nodes > 1) {
+      runtimeErrorMsg() << "LB only works with regular decomposition, "
+                           "if using more than one MPI node";
+    }
   }
 }
 
diff --git a/src/core/grid_based_algorithms/lb_particle_coupling.cpp b/src/core/grid_based_algorithms/lb_particle_coupling.cpp
index 2a3d6c1d57f..349d51427e1 100644
--- a/src/core/grid_based_algorithms/lb_particle_coupling.cpp
+++ b/src/core/grid_based_algorithms/lb_particle_coupling.cpp
@@ -223,24 +223,8 @@ bool in_local_halo(Vector3d const &pos) {
   return in_local_domain(pos, halo);
 }
 
-#ifdef ENGINE
-void add_swimmer_force(Particle const &p, double time_step) {
-  if (p.swimming().swimming) {
-    // calculate source position
-    const double direction =
-        double(p.swimming().push_pull) * p.swimming().dipole_length;
-    auto const director = p.calc_director();
-    auto const source_position = p.pos() + direction * director;
-
-    if (not in_local_halo(source_position)) {
-      return;
-    }
-
-    add_md_force(source_position, p.swimming().f_swim * director, time_step);
-  }
-}
-#endif
-
+/** @brief Return a vector of positions shifted by +,- box length in each
+ ** coordinate */
 std::vector<Utils::Vector3d> shifted_positions(Utils::Vector3d pos,
                                                const BoxGeometry &box) {
   std::vector<Utils::Vector3d> res;
@@ -261,6 +245,30 @@ std::vector<Utils::Vector3d> shifted_positions(Utils::Vector3d pos,
   return res;
 }
 
+#ifdef ENGINE
+void add_swimmer_force(Particle const &p, double time_step, bool has_ghosts) {
+  if (p.swimming().swimming) {
+    // calculate source position
+    const double direction =
+        double(p.swimming().push_pull) * p.swimming().dipole_length;
+    auto const director = p.calc_director();
+    auto const source_position = p.pos() + direction * director;
+    auto const force = p.swimming().f_swim * director;
+
+    if (in_local_halo(source_position)) {
+      add_md_force(source_position, force, time_step);
+    }
+    if (not has_ghosts) {
+      // couple positions shifted by one box length to add forces to ghost
+      // layers
+      for (auto pos : shifted_positions(source_position, box_geo)) {
+        add_md_force(pos, force, time_step);
+      }
+    }
+  }
+}
+#endif
+
 } // namespace
 
 void lb_lbcoupling_calc_particle_lattice_ia(bool couple_virtual,
@@ -346,7 +354,7 @@ void lb_lbcoupling_calc_particle_lattice_ia(bool couple_virtual,
           }
 
 #ifdef ENGINE
-          add_swimmer_force(p, time_step);
+          add_swimmer_force(p, time_step, has_ghosts);
 #endif
         };
 
diff --git a/testsuite/python/engine_lb.py b/testsuite/python/engine_lb.py
index 188fae847f4..7f565d0efa5 100644
--- a/testsuite/python/engine_lb.py
+++ b/testsuite/python/engine_lb.py
@@ -69,6 +69,14 @@ def add_all_types_of_swimmers(
                                   "dipole_length": 0.8})
 
     def setUp(self):
+        if self.decomposition_type == "regular":
+            self.system.cell_system.set_regular_decomposition()
+        elif self.decomposition_type == "n_square":
+            if any(self.system.cell_system.node_grid > 1):
+                self.skipTest(
+                    "N Square and LB not compatible on more than one core")
+            self.system.cell_system.set_n_square()
+
         self.lbf = self.lb_class(**self.LB_params)
         self.system.actors.add(self.lbf)
         self.system.thermostat.set_lb(LB_fluid=self.lbf, gamma=self.gamma)
@@ -126,16 +134,35 @@ def test_particle_forces(self):
 
 
 @utx.skipIfMissingFeatures(["ENGINE", "ROTATIONAL_INERTIA", "MASS"])
-class SwimmerTestCPU(SwimmerTest, ut.TestCase):
+class SwimmerTestRegularCPU(SwimmerTest, ut.TestCase):
+
+    decomposition_type = "regular"
+    lb_class = espressomd.lb.LBFluid
+    tol = 1e-10
+
+
+@utx.skipIfMissingFeatures(["ENGINE", "ROTATIONAL_INERTIA", "MASS"])
+class SwimmerTestNSquareCPU(SwimmerTest, ut.TestCase):
 
+    decomposition_type = "n_square"
     lb_class = espressomd.lb.LBFluid
     tol = 1e-10
 
 
 @utx.skipIfMissingGPU()
 @utx.skipIfMissingFeatures(["ENGINE", "ROTATIONAL_INERTIA", "MASS"])
-class SwimmerTestGPU(SwimmerTest, ut.TestCase):
+class SwimmerTestRegularGPU(SwimmerTest, ut.TestCase):
+
+    decomposition_type = "regular"
+    lb_class = espressomd.lb.LBFluidGPU
+    tol = 1e-5
+
+
+@utx.skipIfMissingGPU()
+@utx.skipIfMissingFeatures(["ENGINE", "ROTATIONAL_INERTIA", "MASS"])
+class SwimmerTestNSquareGPU(SwimmerTest, ut.TestCase):
 
+    decomposition_type = "n_square"
     lb_class = espressomd.lb.LBFluidGPU
     tol = 1e-5