Thermalized Lees-Edwards LB Kernels

- Changed generate_lb_kernels.py to thermalize LE.
- Added kT and seed parameters in all the LE functions in walberla
  bridge.
- Generated new kernels

Co-authored-by: RudolfWeeber <[email protected]>

Author: Yashas17

maintainer/walberla_kernels/generate_lb_kernels.py

@@ -148,22 +148,36 @@ def generate_init_kernels(ctx, method):
 def generate_stream_collide_lees_edwards_kernels(
         ctx, method, data_type, fields):
     precision_prefix = pystencils_espresso.precision_prefix[ctx.double_accuracy]
-    lbm_opt = lbmpy.LBMOptimisation(symbolic_field=fields["pdfs"],
-                                    symbolic_temporary_field=fields["pdfs_tmp"])
-    shear_dir_normal = 1  # y-axis
+    precision_rng = pystencils_espresso.precision_rng[ctx.double_accuracy]
+    block_offsets = tuple(
+        ps.TypedSymbol(f"block_offset_{i}", np.uint32)
+        for i in range(3))
+    optimization = {"cse_global": True,
+                    "double_precision": ctx.double_accuracy}
     le_config = lbmpy.LBMConfig(stencil=stencil,
                                 method=lbmpy.Method.TRT,
-                                relaxation_rate=sp.Symbol("omega_shear"),
                                 streaming_pattern="pull",
+                                relaxation_rate=sp.Symbol("omega_shear"),
                                 force_model=lbmpy.ForceModel.GUO,
                                 force=fields["force"].center_vector,
                                 kernel_type="stream_pull_collide",
                                 **lbm_config_kwargs)
-    le_update_rule_unthermalized = lbmpy.create_lb_update_rule(
+    lbm_opt = lbmpy.LBMOptimisation(symbolic_field=fields["pdfs"],
+                                    symbolic_temporary_field=fields["pdfs_tmp"])
+    shear_dir_normal = 1  # y-axis
+    lb_collision_rule_thermalized = lbmpy.creationfunctions.create_lb_collision_rule(
+        method,
         lbm_config=le_config,
-        lbm_optimisation=lbm_opt)
-    le_collision_rule_unthermalized = lees_edwards.add_lees_edwards_to_collision(
-        config, le_update_rule_unthermalized, fields["pdfs"], stencil,
+        fluctuating={
+            "temperature": kT,
+            "block_offsets": block_offsets,
+            "rng_node": precision_rng
+        },
+        lbm_optimisation=lbm_opt,
+        optimization=optimization
+    )
+    le_lb_collision_rule_thermalized = lees_edwards.add_lees_edwards_to_collision(
+        config, lb_collision_rule_thermalized, fields["pdfs"], stencil,
         shear_dir_normal, True)
     optimization = {"cse_global": True,
                     "double_precision": ctx.double_accuracy}
@@ -174,10 +188,11 @@ def generate_stream_collide_lees_edwards_kernels(
             method,
             le_config,
             data_type,
-            le_collision_rule_unthermalized,
+            le_lb_collision_rule_thermalized,
             stem,
             optimization,
-            params
+            params,
+            block_offset=block_offsets
         )
         ctx.patch_file(stem, get_ext_source(target_suffix),
                        patch_openmp_kernels)

src/core/lb/LBWalberla.cpp

@@ -156,10 +156,6 @@ void LBWalberla::update_collision_model(LBWalberlaBase &lb,
   auto le_protocol = system.lees_edwards->get_protocol();
   if (le_protocol and
       not std::holds_alternative<LeesEdwards::Off>(*le_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,
@@ -171,7 +167,7 @@ void LBWalberla::update_collision_model(LBWalberlaBase &lb,
           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));
+    lb.set_collision_model(std::move(lees_edwards_object), kT, seed);
   } else {
     lb.set_collision_model(kT, seed);
   }

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

@@ -241,7 +241,8 @@ class LBWalberlaBase : public LatticeModel {
 
   /** @brief Configure a thermalized collision model for Lees-Edwards. */
   virtual void
-  set_collision_model(std::unique_ptr<LeesEdwardsPack> &&lees_edwards_pack) = 0;
+  set_collision_model(std::unique_ptr<LeesEdwardsPack> &&lees_edwards_pack,
+                      double kT, unsigned int seed) = 0;
 
   /** @brief Check Lees-Edwards boundary conditions. */
   virtual void check_lebc(unsigned int shear_direction,

src/walberla_bridge/src/lattice_boltzmann/LBWalberlaImpl.hpp

@@ -210,6 +210,7 @@ class LBWalberlaImpl : public LBWalberlaBase {
     }
 
     void operator()(StreamCollisionModelLeesEdwards &cm, IBlock *b) {
+      cm.configure(m_storage, b);
       cm.setV_s(static_cast<decltype(cm.getV_s())>(
           m_lees_edwards_callbacks->get_shear_velocity()));
       cm(b);
@@ -531,9 +532,8 @@ class LBWalberlaImpl : public LBWalberlaBase {
       auto const block_variant = std::variant<IBlock *>(&block);
       std::visit(m_run_stream_collide_sweep, cm_variant, block_variant);
     }
-    if (auto *cm = std::get_if<StreamCollisionModelThermalized>(&cm_variant)) {
-      cm->setTime_step(cm->getTime_step() + 1u);
-    }
+    std::visit([](auto &cm) { cm.setTime_step(cm.getTime_step() + 1u); },
+               cm_variant);
   }
 
   auto has_lees_edwards_bc() const {
@@ -700,8 +700,8 @@ class LBWalberlaImpl : public LBWalberlaBase {
     setup_streaming_communicator();
   }
 
-  void set_collision_model(
-      std::unique_ptr<LeesEdwardsPack> &&lees_edwards_pack) override {
+  void set_collision_model(std::unique_ptr<LeesEdwardsPack> &&lees_edwards_pack,
+                           double kT, unsigned int seed) override {
     assert(m_kT == 0.);
 #if defined(__CUDACC__)
     if constexpr (Architecture == lbmpy::Arch::GPU) {
@@ -712,6 +712,7 @@ class LBWalberlaImpl : public LBWalberlaBase {
     auto const shear_plane_normal = lees_edwards_pack->shear_plane_normal;
     auto const shear_vel = FloatType_c(lees_edwards_pack->get_shear_velocity());
     auto const omega = shear_mode_relaxation_rate();
+    auto const omega_odd = odd_mode_relaxation_rate(omega);
     if (shear_plane_normal != 1u) {
       throw std::domain_error(
           "Lees-Edwards LB only supports shear_plane_normal=\"y\"");
@@ -728,10 +729,13 @@ class LBWalberlaImpl : public LBWalberlaBase {
     }
     auto const &grid_dimensions = lattice.get_grid_dimensions();
     auto const grid_size = FloatType_c(grid_dimensions[shear_plane_normal]);
+    m_kT = FloatType_c(kT);
+    m_seed = seed;
     m_collision_model =
         std::make_shared<CollisionModel>(StreamCollisionModelLeesEdwards(
-            m_last_applied_force_field_id, m_pdf_field_id, grid_size, omega,
-            shear_vel));
+            m_last_applied_force_field_id, m_pdf_field_id, grid_size,
+            zero_centered_to_lb(m_kT), omega, omega, omega_odd, omega, seed,
+            uint32_t{0u}, shear_vel));
     m_lees_edwards_callbacks = std::move(lees_edwards_pack);
     m_run_stream_collide_sweep =
         StreamCollideSweepVisitor(blocks, m_lees_edwards_callbacks);
@@ -1830,23 +1834,25 @@ class LBWalberlaImpl : public LBWalberlaBase {
   }
 
   [[nodiscard]] std::optional<uint64_t> get_rng_state() const override {
-    auto const cm =
-        std::get_if<StreamCollisionModelThermalized>(&*m_collision_model);
-    if (!cm or m_kT == 0.) {
+    if (m_kT == 0.) {
       return std::nullopt;
     }
-    return {static_cast<uint64_t>(cm->getTime_step())};
+    return std::visit(
+        [](auto const &cm) { return static_cast<uint64_t>(cm.getTime_step()); },
+        *m_collision_model);
   }
 
   void set_rng_state(uint64_t counter) override {
-    auto const cm =
-        std::get_if<StreamCollisionModelThermalized>(&*m_collision_model);
-    if (!cm or m_kT == 0.) {
+    if (m_kT == 0.) {
       throw std::runtime_error("This LB instance is unthermalized");
     }
     assert(counter <=
            static_cast<uint32_t>(std::numeric_limits<uint_t>::max()));
-    cm->setTime_step(static_cast<uint32_t>(counter));
+    std::visit(
+        [counter](auto &cm) {
+          cm.setTime_step(static_cast<uint32_t>(counter));
+        },
+        *m_collision_model);
   }
 
   [[nodiscard]] LatticeWalberla const &get_lattice() const noexcept override {