Update LB kernels on periodicity change

src/core/lb/LBNone.hpp

@@ -67,6 +67,8 @@ struct LBNone {
   void on_node_grid_change() const { throw NoLBActive{}; }
   void on_timestep_change() const { throw NoLBActive{}; }
   void on_temperature_change() const { throw NoLBActive{}; }
+  void on_lees_edwards_change() const { throw NoLBActive{}; }
+  void update_collision_model() const { throw NoLBActive{}; }
 };
 
 } // namespace LB

src/core/lb/LBWalberla.cpp

@@ -30,6 +30,8 @@
 #include "system/System.hpp"
 #include "thermostat.hpp"
 
+#include "lees_edwards/lees_edwards.hpp"
+
 #include <walberla_bridge/lattice_boltzmann/LBWalberlaBase.hpp>
 
 #include <utils/Vector.hpp>
@@ -112,6 +114,54 @@ void LBWalberla::sanity_checks(System::System const &system) const {
                              system.get_time_step());
 }
 
+void LBWalberla::on_lees_edwards_change() { update_collision_model(); }
+
+void LBWalberla::set_collision_model(
+    std::unique_ptr<LeesEdwardsPack> &&lees_edwards_object) {
+  lb_fluid->set_collision_model(std::move(lees_edwards_object));
+  lb_fluid->ghost_communication();
+}
+
+void LBWalberla::set_collision_model(double kT, unsigned int seed) {
+  lb_fluid->set_collision_model(kT, seed);
+  lb_fluid->ghost_communication();
+}
+
+void LBWalberla::update_collision_model() {
+  auto const energy_conversion =
+      Utils::int_pow<2>(lb_params->get_agrid() / lb_params->get_tau());
+  auto const kT = lb_fluid->get_kT() * energy_conversion;
+  auto const seed = lb_fluid->get_seed();
+  update_collision_model(*lb_fluid, *lb_params, kT, seed);
+}
+
+void LBWalberla::update_collision_model(LBWalberlaBase &lb,
+                                        LBWalberlaParams &params, double kT,
+                                        unsigned int seed) {
+  auto const &system = ::System::get_system();
+  if (auto le_protocol = system.lees_edwards->get_protocol()) {
+    if (kT != 0.) {
+      throw std::runtime_error(
+          "Lees-Edwards LB doesn't support thermalization");
+    }
+    auto const &le_bc = system.box_geo->lees_edwards_bc();
+    auto lees_edwards_object = std::make_unique<LeesEdwardsPack>(
+        le_bc.shear_direction, le_bc.shear_plane_normal,
+        [&params, le_protocol, &system]() {
+          return get_pos_offset(system.get_sim_time(), *le_protocol) /
+                 params.get_agrid();
+        },
+        [&params, le_protocol, &system]() {
+          return get_shear_velocity(system.get_sim_time(), *le_protocol) *
+                 (params.get_tau() / params.get_agrid());
+        });
+    lb.set_collision_model(std::move(lees_edwards_object));
+  } else {
+    lb.set_collision_model(kT, seed);
+  }
+  lb.ghost_communication();
+}
+
 } // namespace LB
 
 #endif // WALBERLA

src/core/lb/LBWalberla.hpp

@@ -31,6 +31,7 @@
 #include <utility>
 
 class LBWalberlaBase;
+struct LeesEdwardsPack;
 namespace System {
 class System;
 }
@@ -88,6 +89,14 @@ struct LBWalberla {
   }
   void on_timestep_change() const {}
   void on_temperature_change() const {}
+  void on_lees_edwards_change();
+  void
+  set_collision_model(std::unique_ptr<LeesEdwardsPack> &&lees_edwards_pack);
+  void set_collision_model(double kT, unsigned int seed);
+  void update_collision_model();
+  static void update_collision_model(LBWalberlaBase &instance,
+                                     LBWalberlaParams &params, double kT,
+                                     unsigned int seed);
 };
 
 } // namespace LB

src/core/lb/Solver.cpp

@@ -140,6 +140,18 @@ void Solver::on_temperature_change() {
   }
 }
 
+void Solver::on_lees_edwards_change() {
+  if (impl->solver) {
+    std::visit([](auto &ptr) { ptr->on_lees_edwards_change(); }, *impl->solver);
+  }
+}
+
+void Solver::update_collision_model() {
+  if (impl->solver) {
+    std::visit([](auto &ptr) { ptr->update_collision_model(); }, *impl->solver);
+  }
+}
+
 bool Solver::is_gpu() const {
   check_solver(impl);
   return std::visit([](auto &ptr) { return ptr->is_gpu(); }, *impl->solver);

src/core/lb/Solver.hpp

@@ -168,6 +168,8 @@ struct Solver : public System::Leaf<Solver> {
   void on_cell_structure_change();
   void on_timestep_change();
   void on_temperature_change();
+  void on_lees_edwards_change();
+  void update_collision_model();
   void veto_boxl_change() const;
 
 private:

src/core/system/System.cpp

@@ -378,6 +378,8 @@ void System::on_observable_calc() {
   clear_particle_node();
 }
 
+void System::on_lees_edwards_change() { lb.on_lees_edwards_change(); }
+
 void System::update_local_geo() {
   *local_geo = LocalBox::make_regular_decomposition(
       box_geo->length(), ::communicator.calc_node_index(),

src/core/system/System.hpp

@@ -253,6 +253,7 @@ class System : public std::enable_shared_from_this<System> {
    *  initialized immediately (P3M etc.).
    */
   void on_observable_calc();
+  void on_lees_edwards_change();
   void veto_boxl_change(bool skip_particle_checks = false) const;
   /**@}*/
 

src/core/unit_tests/lb_particle_coupling_test.cpp

@@ -608,11 +608,13 @@ BOOST_AUTO_TEST_CASE(lb_exceptions) {
     BOOST_CHECK_THROW(lb.veto_kT(0.), NoLBActive);
     BOOST_CHECK_THROW(lb.lebc_sanity_checks(0u, 1u), NoLBActive);
     BOOST_CHECK_THROW(lb.propagate(), NoLBActive);
+    BOOST_CHECK_THROW(lb.update_collision_model(), NoLBActive);
     BOOST_CHECK_THROW(lb.on_cell_structure_change(), NoLBActive);
     BOOST_CHECK_THROW(lb.on_boxl_change(), NoLBActive);
     BOOST_CHECK_THROW(lb.on_node_grid_change(), NoLBActive);
     BOOST_CHECK_THROW(lb.on_timestep_change(), NoLBActive);
     BOOST_CHECK_THROW(lb.on_temperature_change(), NoLBActive);
+    BOOST_CHECK_THROW(lb.on_lees_edwards_change(), NoLBActive);
     BOOST_CHECK_THROW(lb_impl->get_density_at_pos(vec, true), NoLBActive);
     BOOST_CHECK_THROW(lb_impl->get_velocity_at_pos(vec, true), NoLBActive);
     BOOST_CHECK_THROW(lb_impl->add_force_at_pos(vec, vec), NoLBActive);

src/script_interface/lees_edwards/LeesEdwards.hpp

@@ -103,12 +103,13 @@ class LeesEdwards : public AutoParameters<LeesEdwards, System::Leaf> {
           throw std::invalid_argument("Parameters 'shear_direction' and "
                                       "'shear_plane_normal' must differ");
         }
-        auto const &system = get_system();
+        auto &system = get_system();
         system.lb.lebc_sanity_checks(shear_direction, shear_plane_normal);
         // update box geometry and cell structure
         system.box_geo->set_lees_edwards_bc(
             LeesEdwardsBC{0., 0., shear_direction, shear_plane_normal});
         m_lees_edwards->set_protocol(m_protocol->protocol());
+        system.on_lees_edwards_change();
       });
     }
     return {};

src/script_interface/walberla/LBFluid.cpp

@@ -158,7 +158,7 @@ void LBFluid::do_construct(VariantMap const &params) {
   auto const ext_f = get_value<Utils::Vector3d>(params, "ext_force_density");
   m_lb_params = std::make_shared<::LB::LBWalberlaParams>(agrid, tau);
   m_is_active = false;
-  m_seed = get_value<int>(params, "seed");
+  auto const seed = get_value<int>(params, "seed");
   context()->parallel_try_catch([&]() {
     if (tau <= 0.) {
       throw std::domain_error("Parameter 'tau' must be > 0");
@@ -175,7 +175,7 @@ void LBFluid::do_construct(VariantMap const &params) {
     auto const lb_dens = m_conv_dens * dens;
     auto const lb_kT = m_conv_energy * kT;
     auto const lb_ext_f = m_conv_force_dens * ext_f;
-    if (m_seed < 0) {
+    if (seed < 0) {
       throw std::domain_error("Parameter 'seed' must be >= 0");
     }
     if (lb_kT < 0.) {
@@ -188,28 +188,8 @@ void LBFluid::do_construct(VariantMap const &params) {
       throw std::domain_error("Parameter 'kinematic_viscosity' must be >= 0");
     }
     make_instance(params);
-    auto const &system = ::System::get_system();
-    if (auto le_protocol = system.lees_edwards->get_protocol()) {
-      if (lb_kT != 0.) {
-        throw std::runtime_error(
-            "Lees-Edwards LB doesn't support thermalization");
-      }
-      auto const &le_bc = system.box_geo->lees_edwards_bc();
-      auto lees_edwards_object = std::make_unique<LeesEdwardsPack>(
-          le_bc.shear_direction, le_bc.shear_plane_normal,
-          [this, le_protocol, &system]() {
-            return get_pos_offset(system.get_sim_time(), *le_protocol) /
-                   m_lb_params->get_agrid();
-          },
-          [this, le_protocol, &system]() {
-            return get_shear_velocity(system.get_sim_time(), *le_protocol) *
-                   (m_lb_params->get_tau() / m_lb_params->get_agrid());
-          });
-      m_instance->set_collision_model(std::move(lees_edwards_object));
-    } else {
-      m_instance->set_collision_model(lb_kT, m_seed);
-    }
-    m_instance->ghost_communication(); // synchronize ghost layers
+    ::LB::LBWalberla::update_collision_model(*m_instance, *m_lb_params, lb_kT,
+                                             static_cast<unsigned int>(seed));
     m_instance->set_external_force(lb_ext_f);
     for (auto &vtk : m_vtk_writers) {
       vtk->attach_to_lattice(m_instance, get_latice_to_md_units_conversion());

src/script_interface/walberla/LBFluid.hpp

@@ -51,7 +51,6 @@ class LBFluid : public LatticeModel<::LBWalberlaBase, LBVTKHandle> {
   using Base = LatticeModel<::LBWalberlaBase, LBVTKHandle>;
   std::shared_ptr<::LB::LBWalberlaParams> m_lb_params;
   bool m_is_active;
-  int m_seed;
   double m_conv_dist;
   double m_conv_visc;
   double m_conv_dens;
@@ -77,7 +76,8 @@ class LBFluid : public LatticeModel<::LBWalberlaBase, LBVTKHandle> {
          [this]() { return m_instance->get_lattice().get_grid_dimensions(); }},
         {"kT", AutoParameter::read_only,
          [this]() { return m_instance->get_kT() / m_conv_energy; }},
-        {"seed", AutoParameter::read_only, [this]() { return m_seed; }},
+        {"seed", AutoParameter::read_only,
+         [this]() { return static_cast<int>(m_instance->get_seed()); }},
         {"rng_state",
          [this](Variant const &v) {
            auto const rng_state = get_value<int>(v);

src/walberla_bridge/include/walberla_bridge/lattice_boltzmann/LBWalberlaBase.hpp

@@ -257,10 +257,13 @@ class LBWalberlaBase : public LatticeModel {
   /** @brief Get the fluid temperature (if thermalized). */
   virtual double get_kT() const noexcept = 0;
 
+  /** @brief Get the RNG seed (if thermalized). */
+  virtual unsigned int get_seed() const noexcept = 0;
+
   /** @brief Set the RNG counter (if thermalized). */
   [[nodiscard]] virtual std::optional<uint64_t> get_rng_state() const = 0;
 
-  /** @brief Set the rng state of thermalized LBs */
+  /** @brief Set the RNG counter (if thermalized). */
   virtual void set_rng_state(uint64_t counter) = 0;
 
   /** @brief get the velocity field id */

src/walberla_bridge/src/lattice_boltzmann/LBWalberlaImpl.hpp

@@ -257,6 +257,7 @@ class LBWalberlaImpl : public LBWalberlaBase {
   FloatType m_viscosity; /// kinematic viscosity
   FloatType m_density;
   FloatType m_kT;
+  unsigned int m_seed;
 
   // Block data access handles
   BlockDataID m_pdf_field_id;
@@ -401,7 +402,7 @@ class LBWalberlaImpl : public LBWalberlaBase {
   LBWalberlaImpl(std::shared_ptr<LatticeWalberla> lattice, double viscosity,
                  double density)
       : m_viscosity(FloatType_c(viscosity)), m_density(FloatType_c(density)),
-        m_kT(FloatType{0}), m_lattice(std::move(lattice)) {
+        m_kT(FloatType{0}), m_seed(0u), m_lattice(std::move(lattice)) {
 
     auto const &blocks = m_lattice->get_blocks();
     auto const n_ghost_layers = m_lattice->get_ghost_layers();
@@ -592,6 +593,7 @@ class LBWalberlaImpl : public LBWalberlaBase {
     auto const omega_odd = odd_mode_relaxation_rate(omega);
     auto const blocks = get_lattice().get_blocks();
     m_kT = FloatType_c(kT);
+    m_seed = seed;
     auto obj = CollisionModelThermalized(m_last_applied_force_field_id,
                                          m_pdf_field_id, m_kT, omega, omega,
                                          omega_odd, omega, seed, uint32_t{0u});
@@ -1366,6 +1368,10 @@ class LBWalberlaImpl : public LBWalberlaBase {
     return numeric_cast<double>(m_kT);
   }
 
+  [[nodiscard]] unsigned int get_seed() const noexcept override {
+    return m_seed;
+  }
+
   [[nodiscard]] std::optional<uint64_t> get_rng_state() const override {
     auto const cm = std::get_if<CollisionModelThermalized>(&*m_collision_model);
     if (!cm or m_kT == 0.) {

testsuite/python/lb_lees_edwards.py

@@ -191,8 +191,6 @@ def test_velocity_shift_from_particles(self):
             with LBContextManager() as lbf:
                 for profile in self.sample_lb_velocities(lbf):
                     self.check_profile(profile, stencil_D2Q8, '', 'SNWE', tol)
-
-        # order of instantiation doesn't matter
         with LBContextManager() as lbf:
             with LEContextManager('x', 'y', 0):
                 for profile in self.sample_lb_velocities(lbf):
@@ -208,6 +206,13 @@ def test_velocity_shift_from_particles(self):
             with LBContextManager() as lbf:
                 for profile in self.sample_lb_velocities(lbf):
                     self.check_profile(profile, stencil, 'SN', 'WE', tol)
+        with LBContextManager() as lbf:
+            with LEContextManager('x', 'y', le_offset):
+                stencil = {'N~': (8 - le_offset, 0),
+                           'S~': (8 + le_offset, 16),
+                           **stencil_D2Q8}
+                for profile in self.sample_lb_velocities(lbf):
+                    self.check_profile(profile, stencil, 'SN', 'WE', tol)
 
         # TODO: re-enable this check once LB can be sheared in any direction
 #        # East and West are sheared vertically
@@ -253,8 +258,6 @@ def create_impulse(lbf, stencil):
                 create_impulse(lbf, stencil_D2Q8)
                 for profile in self.sample_lb_velocities(lbf):
                     self.check_profile(profile, stencil_D2Q8, '', 'SNWE', tol)
-
-        # order of instantiation doesn't matter
         with LBContextManager() as lbf:
             with LEContextManager('x', 'y', 0):
                 create_impulse(lbf, stencil_D2Q8)
@@ -272,6 +275,14 @@ def create_impulse(lbf, stencil):
                 create_impulse(lbf, stencil_D2Q8)
                 for profile in self.sample_lb_velocities(lbf):
                     self.check_profile(profile, stencil, 'SN', 'WE', tol)
+        with LBContextManager() as lbf:
+            with LEContextManager('x', 'y', le_offset):
+                stencil = {'N~': (8 - le_offset, 1),
+                           'S~': (8 + le_offset, 15),
+                           **stencil_D2Q8}
+                create_impulse(lbf, stencil_D2Q8)
+                for profile in self.sample_lb_velocities(lbf):
+                    self.check_profile(profile, stencil, 'SN', 'WE', tol)
 
         # TODO: re-enable this check once LB can be sheared in any direction
 #        # East and West are sheared vertically
@@ -335,6 +346,10 @@ def test_lebc_mismatch(self):
                     agrid=1., density=1., kinematic_viscosity=1., kT=1., seed=42,
                     tau=system.time_step)
         self.assertIsNone(system.lb)
+        with self.assertRaisesRegex(RuntimeError, "Lees-Edwards LB doesn't support thermalization"):
+            with LBContextManager(kT=1., seed=42) as lbf:
+                LEContextManager('x', 'y', 1.)
+        self.assertIsNone(system.lb)
 
         with self.assertRaisesRegex(ValueError, "Lees-Edwards sweep is implemented for a ghost layer of thickness 1"):
             lattice = espressomd.lb.LatticeWalberla(agrid=1., n_ghost_layers=2)