Update kernels and walberla bridge headers

Lees Edwards collison model now takes kT and seed parameters for
thermalization.

Author: Yashas17

maintainer/walberla_kernels/generate_lb_kernels.py

@@ -135,10 +135,8 @@ 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]
+    shear_dir_normal = 1  # y-axis
     block_offsets = tuple(
         ps.TypedSymbol(f"block_offset_{i}", np.uint32)
         for i in range(3))
@@ -150,6 +148,8 @@ def generate_stream_collide_lees_edwards_kernels(
                                 force=fields["force"].center_vector,
                                 kernel_type="stream_pull_collide",
                                 **lbm_config_kwargs)
+    lbm_opt = lbmpy.LBMOptimisation(symbolic_field=fields["pdfs"],
+                                    symbolic_temporary_field=fields["pdfs_tmp"])
     optimization = {"cse_global": True,
                     "double_precision": ctx.double_accuracy}
     le_collision_rule_thermalized = lbmpy.creationfunctions.create_lb_collision_rule(
@@ -160,6 +160,7 @@ def generate_stream_collide_lees_edwards_kernels(
             "block_offsets": block_offsets,
             "rng_node": precision_rng
         },
+        lbm_optimisation=lbm_opt,
         optimization=optimization
     )
     le_collision_rule_thermalized = lees_edwards.add_lees_edwards_to_collision(
@@ -175,7 +176,8 @@ def generate_stream_collide_lees_edwards_kernels(
             le_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

@@ -171,7 +171,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,7 @@ 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,7 @@ 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 {
@@ -701,8 +700,7 @@ class LBWalberlaImpl : public LBWalberlaBase {
   }
 
   void set_collision_model(
-      std::unique_ptr<LeesEdwardsPack> &&lees_edwards_pack) override {
-    assert(m_kT == 0.);
+      std::unique_ptr<LeesEdwardsPack> &&lees_edwards_pack, double kT, unsigned int seed) override {
 #if defined(__CUDACC__)
     if constexpr (Architecture == lbmpy::Arch::GPU) {
       throw std::runtime_error("Lees-Edwards LB doesn't support GPU yet");
@@ -712,6 +710,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 +727,12 @@ 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);
@@ -1771,23 +1772,19 @@ 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 {