LE: add fully connected option to RegularDecomposition (#4958)

Description of changes:
- fully connected means that all cells of the decomposition at the (shear) normal boundary are neighbors with all cells across that boundary in the shear direction. 
- in this way, all pairs will be visible in the non-bonded loop, when a Lees-Edwards offset is applied across that boundary

src/core/cell_system/CellStructure.cpp

@@ -44,6 +44,7 @@
 #include <cstddef>
 #include <iterator>
 #include <memory>
+#include <optional>
 #include <set>
 #include <stdexcept>
 #include <string>
@@ -248,12 +249,13 @@ void CellStructure::set_atom_decomposition() {
   system.on_cell_structure_change();
 }
 
-void CellStructure::set_regular_decomposition(double range) {
+void CellStructure::set_regular_decomposition(
+    double range, std::optional<std::pair<int, int>> fully_connected_boundary) {
   auto &system = get_system();
   auto &local_geo = *system.local_geo;
   auto const &box_geo = *system.box_geo;
   set_particle_decomposition(std::make_unique<RegularDecomposition>(
-      ::comm_cart, range, box_geo, local_geo));
+      ::comm_cart, range, box_geo, local_geo, fully_connected_boundary));
   m_type = CellStructureType::REGULAR;
   local_geo.set_cell_structure_type(m_type);
   system.on_cell_structure_change();

src/core/cell_system/CellStructure.hpp

@@ -46,6 +46,7 @@
 #include <cassert>
 #include <iterator>
 #include <memory>
+#include <optional>
 #include <set>
 #include <span>
 #include <stdexcept>
@@ -557,7 +558,9 @@ struct CellStructure : public System::Leaf<CellStructure> {
    *
    * @param range Interaction range.
    */
-  void set_regular_decomposition(double range);
+  void set_regular_decomposition(
+      double range,
+      std::optional<std::pair<int, int>> fully_connected_boundary);
 
   /**
    * @brief Set the particle decomposition to @ref HybridDecomposition.

src/core/cell_system/HybridDecomposition.cpp

@@ -37,6 +37,7 @@
 #include <cstddef>
 #include <functional>
 #include <iterator>
+#include <optional>
 #include <set>
 #include <utility>
 
@@ -48,7 +49,7 @@ HybridDecomposition::HybridDecomposition(boost::mpi::communicator comm,
                                          std::set<int> n_square_types)
     : m_comm(std::move(comm)), m_box(box_geo), m_cutoff_regular(cutoff_regular),
       m_regular_decomposition(RegularDecomposition(
-          m_comm, cutoff_regular + skin, m_box, local_box)),
+          m_comm, cutoff_regular + skin, m_box, local_box, std::nullopt)),
       m_n_square(AtomDecomposition(m_comm, m_box)),
       m_n_square_types(std::move(n_square_types)),
       m_get_global_ghost_flags(std::move(get_ghost_flags)) {

src/core/cell_system/RegularDecomposition.cpp

@@ -397,6 +397,27 @@ void RegularDecomposition::init_cell_interactions() {
   auto const &node_grid = ::communicator.node_grid;
   auto const global_halo_offset = hadamard_product(node_pos, cell_grid) - halo;
   auto const global_size = hadamard_product(node_grid, cell_grid);
+  auto const at_boundary = [&global_size](int coord, Utils::Vector3i cell_idx) {
+    return (cell_idx[coord] == 0 or cell_idx[coord] == global_size[coord]);
+  };
+
+  // For the fully connected feature (cells that don't share at least a corner)
+  // only apply if one cell is a ghost cell (i.e. connections across the
+  // periodic boundary.
+  auto const fcb_is_inner_connection = [&global_size, this](Utils::Vector3i a,
+                                                            Utils::Vector3i b) {
+    if (fully_connected_boundary()) {
+      auto const [fc_normal, fc_dir] = *fully_connected_boundary();
+      auto const involves_ghost_cell =
+          (a[fc_normal] == -1 or a[fc_normal] == global_size[fc_normal] or
+           b[fc_normal] == -1 or b[fc_normal] == global_size[fc_normal]);
+      if (not involves_ghost_cell) {
+        // check if cells do not share at least a corner
+        return std::abs((a - b)[fc_dir]) > 1;
+      }
+    }
+    return false;
+  };
 
   /* Translate a node local index (relative to the origin of the local grid)
    * to a global index. */
@@ -418,6 +439,19 @@ void RegularDecomposition::init_cell_interactions() {
     return (global_index - global_halo_offset);
   };
 
+  // sanity checks
+  if (fully_connected_boundary()) {
+    auto const [fc_normal, fc_dir] = *fully_connected_boundary();
+    if (fc_normal == fc_dir) {
+      throw std::domain_error("fully_connected_boundary normal and connection "
+                              "coordinates need to differ.");
+    }
+    if (node_grid[fc_dir] != 1) {
+      throw std::runtime_error(
+          "The MPI nodegrid must be 1 in the fully connected direction.");
+    }
+  }
+
   /* We only consider local cells (e.g. not halo cells), which
    * span the range [(1,1,1), cell_grid) in local coordinates. */
   auto const start = global_index(Utils::Vector3i{1, 1, 1});
@@ -431,20 +465,18 @@ void RegularDecomposition::init_cell_interactions() {
         Utils::Vector3i lower_index = {m - 1, n - 1, o - 1};
         Utils::Vector3i upper_index = {m + 1, n + 1, o + 1};
 
-        //        /* In the fully connected case, we consider all cells
-        //         * in the direction as neighbors, not only the nearest ones.
+        /* In the fully connected case, we consider all cells
+        * in the direction as neighbors, not only the nearest ones.
         //         */
-        //        for (int i = 0; i < 3; i++) {
-        //          if (dd.fully_connected[i]) {
-        //            // Fully connected is only needed at the box surface
-        //            if (i==0 and (n!=start[1] or n!=end[1]-1) and (o!=start[2]
-        //            or o!=end[2]-1)) continue; if (i==1 and (m!=start[0] or
-        //            m!=end[0]-1) and (o!=start[2] or o!=end[2]-1)) continue;
-        //            if (i==2 and (m!=start[0] or m!=end[0]-1) and (n!=start[1]
-        //            or n!=end[1]-1)) continue; lower_index[i] = 0;
-        //            upper_index[i] = global_size[i] - 1;
-        //          }
-        //        }
+        if (fully_connected_boundary()) {
+          auto const [fc_boundary, fc_direction] = *fully_connected_boundary();
+
+          // Fully connected is only needed at the box surface
+          if (at_boundary(fc_boundary, {m, n, o})) {
+            lower_index[fc_direction] = -1;
+            upper_index[fc_direction] = global_size[fc_boundary];
+          }
+        }
 
         /* In non-periodic directions, the halo needs not
          * be considered. */
@@ -466,6 +498,12 @@ void RegularDecomposition::init_cell_interactions() {
         for (int p = lower_index[2]; p <= upper_index[2]; p++)
           for (int q = lower_index[1]; q <= upper_index[1]; q++)
             for (int r = lower_index[0]; r <= upper_index[0]; r++) {
+              if (fully_connected_boundary()) {
+                // Avoid fully connecting the boundary layer and the
+                // next INNER layer
+                if (fcb_is_inner_connection({m, n, o}, {r, q, p}))
+                  continue;
+              }
               neighbors.insert(Utils::Vector3i{r, q, p});
             }
 
@@ -629,11 +667,13 @@ GhostCommunicator RegularDecomposition::prepare_comm() {
   return ghost_comm;
 }
 
-RegularDecomposition::RegularDecomposition(boost::mpi::communicator comm,
-                                           double range,
-                                           BoxGeometry const &box_geo,
-                                           LocalBox const &local_geo)
-    : m_comm(std::move(comm)), m_box(box_geo), m_local_box(local_geo) {
+RegularDecomposition::RegularDecomposition(
+    boost::mpi::communicator comm, double range, BoxGeometry const &box_geo,
+    LocalBox const &local_geo,
+    std::optional<std::pair<int, int>> fully_connected)
+    : m_comm(std::move(comm)), m_box(box_geo), m_local_box(local_geo),
+      m_fully_connected_boundary(std::move(fully_connected)) {
+
   /* set up new regular decomposition cell structure */
   create_cell_grid(range);
 

src/core/cell_system/RegularDecomposition.hpp

@@ -79,6 +79,7 @@ struct RegularDecomposition : public ParticleDecomposition {
   boost::mpi::communicator m_comm;
   BoxGeometry const &m_box;
   LocalBox m_local_box;
+  std::optional<std::pair<int, int>> m_fully_connected_boundary = {};
   std::vector<Cell> cells;
   std::vector<Cell *> m_local_cells;
   std::vector<Cell *> m_ghost_cells;
@@ -87,7 +88,8 @@ struct RegularDecomposition : public ParticleDecomposition {
 
 public:
   RegularDecomposition(boost::mpi::communicator comm, double range,
-                       BoxGeometry const &box_geo, LocalBox const &local_geo);
+                       BoxGeometry const &box_geo, LocalBox const &local_geo,
+                       std::optional<std::pair<int, int>> fully_connected);
 
   GhostCommunicator const &exchange_ghosts_comm() const override {
     return m_exchange_ghosts_comm;
@@ -115,6 +117,8 @@ struct RegularDecomposition : public ParticleDecomposition {
   Utils::Vector3d max_cutoff() const override;
   Utils::Vector3d max_range() const override;
 
+  auto fully_connected_boundary() const { return m_fully_connected_boundary; }
+
   std::optional<BoxGeometry> minimum_image_distance() const override {
     return {m_box};
   }

src/core/system/System.cpp

@@ -138,7 +138,17 @@ void System::set_min_global_cut(double value) {
 
 void System::set_cell_structure_topology(CellStructureType topology) {
   if (topology == CellStructureType::REGULAR) {
-    cell_structure->set_regular_decomposition(get_interaction_range());
+    if (cell_structure->decomposition_type() == CellStructureType::REGULAR) {
+      // get fully connected info from exising regular decomposition
+      auto &old_regular_decomposition =
+          dynamic_cast<RegularDecomposition const &>(
+              std::as_const(*cell_structure).decomposition());
+      cell_structure->set_regular_decomposition(
+          get_interaction_range(),
+          old_regular_decomposition.fully_connected_boundary());
+    } else { // prev. decomposition is not a regular decomposition
+      cell_structure->set_regular_decomposition(get_interaction_range(), {});
+    }
   } else if (topology == CellStructureType::NSQUARE) {
     cell_structure->set_atom_decomposition();
   } else {

src/python/espressomd/cell_system.py

@@ -115,6 +115,12 @@ def set_regular_decomposition(self, **kwargs):
         use_verlet_lists : :obj:`bool`, optional
             Activates or deactivates the usage of Verlet lists.
             Defaults to ``True``.
+        fully_connected_boundary : :obj:`dict`, optional
+            If set, connects all cells on a given boundary along the given direction.
+            Example: ``{"boundary": "z", "direction": "x"}`` connects all
+            cells on the boundary normal to the z-direction along the x-axis.
+            This corresponds to z-axis as shear plane normal and x-axis as
+            shear direction in Lees-Edwards boundary conditions.
 
         """
         self.call_method("initialize", name="regular_decomposition", **kwargs)

src/script_interface/cell_system/CellSystem.cpp

@@ -40,7 +40,9 @@
 #include <boost/variant.hpp>
 
 #include <algorithm>
+#include <cassert>
 #include <iterator>
+#include <optional>
 #include <set>
 #include <sstream>
 #include <stdexcept>
@@ -49,6 +51,25 @@
 #include <utility>
 #include <vector>
 
+static int coord(std::string const &s) {
+  if (s == "x")
+    return 0;
+  if (s == "y")
+    return 1;
+  if (s == "z")
+    return 2;
+  throw std::invalid_argument("Invalid Cartesian coordinate: '" + s + "'");
+}
+
+static std::string coord_letter(int c) {
+  if (c == 0)
+    return "x";
+  if (c == 1)
+    return "y";
+  assert(c == 2);
+  return "z";
+}
+
 namespace ScriptInterface {
 namespace CellSystem {
 
@@ -110,6 +131,20 @@ CellSystem::CellSystem() {
          auto const ns_types = hd.get_n_square_types();
          return Variant{std::vector<int>(ns_types.begin(), ns_types.end())};
        }},
+      {"fully_connected_boundary", AutoParameter::read_only,
+       [this]() {
+         if (get_cell_structure().decomposition_type() !=
+             CellStructureType::REGULAR) {
+           return Variant{none};
+         }
+         auto const rd = get_regular_decomposition();
+         auto const fcb = rd.fully_connected_boundary();
+         if (not fcb)
+           return Variant{none};
+         return Variant{std::unordered_map<std::string, Variant>{
+             {{"boundary", Variant{coord_letter((*fcb).first)}},
+              {"direction", Variant{coord_letter((*fcb).second)}}}}};
+       }},
       {"cutoff_regular", AutoParameter::read_only,
        [this]() {
          if (get_cell_structure().decomposition_type() !=
@@ -261,6 +296,21 @@ void CellSystem::initialize(CellStructureType const &cs_type,
         get_value_or<std::vector<int>>(params, "n_square_types", {});
     auto n_square_types = std::set<int>{ns_types.begin(), ns_types.end()};
     m_cell_structure->set_hybrid_decomposition(cutoff_regular, n_square_types);
+  } else if (cs_type == CellStructureType::REGULAR) {
+    std::optional<std::pair<int, int>> fcb_pair = std::nullopt;
+    if (params.contains("fully_connected_boundary") and
+        not is_none(params.at("fully_connected_boundary"))) {
+      auto const variant =
+          get_value<VariantMap>(params, "fully_connected_boundary");
+      context()->parallel_try_catch([&fcb_pair, &variant]() {
+        fcb_pair = {{coord(boost::get<std::string>(variant.at("boundary"))),
+                     coord(boost::get<std::string>(variant.at("direction")))}};
+      });
+    }
+    context()->parallel_try_catch([this, &fcb_pair]() {
+      m_cell_structure->set_regular_decomposition(
+          get_system().get_interaction_range(), fcb_pair);
+    });
   } else {
     system.set_cell_structure_topology(cs_type);
   }