Draft:Thermalize Lees Edwards LBM kernel

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)
@@ -206,7 +221,6 @@ def generate_stream_collide_kernels(ctx, method, data_type):
         },
         optimization=optimization
     )
-
     for params, target_suffix in paramlist(parameters, ("GPU", "CPU", "AVX")):
         stem = f"StreamCollideSweepThermalized{precision_prefix}{target_suffix}"  # nopep8
         pystencils_espresso.generate_stream_collision_sweep(

maintainer/walberla_kernels/lees_edwards.py

@@ -45,6 +45,7 @@ def velocity_offset_eqs(config, method, pdfs,
     populations in the boundary layer. Returns an AssignmentCollection
     with one Assignment per stencil direction.
     """
+    print("Lees Edwareds velocity offsets applied to PDF components")
     dim = len(stencil[0])
     default_dtype = config.data_type.default_factory()
 
@@ -84,10 +85,12 @@ def velocity_offset_eqs(config, method, pdfs,
     delta_pdf_eqs = macroscopic_values_setter(
         method, sp.Symbol("dens"), [
             sp.Symbol("v_0"), sp.Symbol("v_1"), sp.Symbol("v_2")], pdfs)
-
+    print(f"{delta_pdf_eqs=}")
     # Replace the assignments of (rho,u) by (rho, u+v) - (rho,u)
     ma = []
+
     for a, c in zip(delta_pdf_eqs.main_assignments, method.stencil):
+        print(f"Direction {c}: {a}")
         # Determine direction of the stencil component in the
         # shear_dir_normal
         if c[shear_dir_normal] == 1:
@@ -99,7 +102,7 @@ def velocity_offset_eqs(config, method, pdfs,
         else:
             up = False
             down = False
-
+        print(f"{up=} {down=} {layer_prefactor=}")
         # Replace (rho,u) by (rho,u+v) in boundary layers
         rhs = sp.simplify(
             a.rhs -
@@ -113,9 +116,10 @@ def velocity_offset_eqs(config, method, pdfs,
         rhs = rhs.replace(points_up, up)
         rhs = rhs.replace(points_down, down)
         new_a = Assignment(a.lhs, rhs)
+        print(f"velocity offset: {new_a}")
+        print()
 
         ma.append(new_a)
-        print(c, ma[-1])
     # Plug in modified assignments
     delta_pdf_eqs.main_assignments = ma
     return delta_pdf_eqs.main_assignments
@@ -132,10 +136,26 @@ def add_lees_edwards_to_collision(
         stencil,
         combined_kernel)
 
-    ma = []
-    for i, a in enumerate(collision.main_assignments):
-        # Add Lees-Edwards-shift to collision main assignments
-        new_a = Assignment(a.lhs, a.rhs + offset[i].rhs)
-        ma.append(new_a)
-    collision.main_assignments = ma
+    print("Lees edwards collision assigments")
+
+    keys = [sp.Symbol(f"d_{i}") for i in range(19)]
+
+    for i, key in enumerate(keys):
+        match_found = False
+        for j, a in enumerate(collision.main_assignments):
+            if a.lhs == key:
+                collision.main_assignments[j] = Assignment(
+                    a.lhs, a.rhs + offset[i].rhs)
+                match_found = True
+                print(f"m:{a.lhs}, o:{offset[i].lhs}")
+                break
+        if match_found:
+            continue
+        for j, a in enumerate(collision.subexpressions):
+            if hasattr(a, 'lhs'):
+                if a.lhs == key:
+                    collision.subexpressions[j] = Assignment(
+                        a.lhs, a.rhs + offset[i].rhs)
+                    print(f"s:{a.lhs}, o:{offset[i].lhs}")
+                    break
     return collision

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 {