From b812fc1d37f9a549fd464ed2f2262b6d7d9f2b78 Mon Sep 17 00:00:00 2001
From: David Beyer <dbeyer@icp.uni-stuttgart.de>
Date: Mon, 19 Aug 2024 15:06:47 +0200
Subject: [PATCH] Add function update_collision_model

---
 src/core/lb/LBWalberla.cpp | 44 ++++++++++++++++++++++++++++++++++++++
 src/core/lb/LBWalberla.hpp |  5 +++++
 2 files changed, 49 insertions(+)

diff --git a/src/core/lb/LBWalberla.cpp b/src/core/lb/LBWalberla.cpp
index 77f149ac3c..1de06f11ec 100644
--- a/src/core/lb/LBWalberla.cpp
+++ b/src/core/lb/LBWalberla.cpp
@@ -114,6 +114,7 @@ void LBWalberla::sanity_checks(System::System const &system) const {
                              system.get_time_step());
 }
 
+/*
 void LBWalberla::on_lees_edwards_change() {
   auto &system = System::get_system();
   auto const energy_conversion =
@@ -137,6 +138,49 @@ void LBWalberla::on_lees_edwards_change() {
   lb_fluid->set_collision_model(std::move(lees_edwards_object));
   lb_fluid->ghost_communication(); // synchronize ghost layers
 }
+*/
+
+
+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(); // synchronize ghost layers
+
+}
+
+void LBWalberla::set_collision_model(double kT, unsigned int seed) {
+  lb_fluid->set_collision_model(kT, seed);
+  lb_fluid->ghost_communication(); // synchronize ghost layers
+}
+
+void LBWalberla::update_collision_model() {
+  auto &system = System::get_system();
+  auto const energy_conversion =
+     Utils::int_pow<2>(lb_params->get_agrid() / lb_params->get_tau());
+  auto const lb_kT = lb_fluid->get_kT() * energy_conversion;
+  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) /
+               get_agrid();
+      },
+      [this, le_protocol, &system]() {
+        return get_shear_velocity(system.get_sim_time(), *le_protocol) *
+               (get_tau() / get_agrid());
+      });
+    set_collision_model(std::move(lees_edwards_object));
+  } else {
+    set_collision_model(lb_kT, 42);
+  }
+}
 
 } // namespace LB
 
diff --git a/src/core/lb/LBWalberla.hpp b/src/core/lb/LBWalberla.hpp
index 83a2c952a8..2173c55549 100644
--- a/src/core/lb/LBWalberla.hpp
+++ b/src/core/lb/LBWalberla.hpp
@@ -23,6 +23,8 @@
 
 #ifdef WALBERLA
 
+#include <walberla_bridge/lattice_boltzmann/LBWalberlaBase.hpp>
+
 #include <utils/Vector.hpp>
 
 #include <memory>
@@ -89,6 +91,9 @@ 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();
 };
 
 } // namespace LB