From 3b0f59f388659f40649a6b30f8d95ed265ec5c94 Mon Sep 17 00:00:00 2001
From: Frederik Hennig <frederik.hennig@fau.de>
Date: Thu, 23 Jun 2022 13:27:04 +0200
Subject: [PATCH] Fix: Broken He model population space forcing with
 zero-centered storage

---
 lbmpy/forcemodels.py      |  9 ++++++---
 lbmpy_tests/test_force.py | 11 +++++++++--
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/lbmpy/forcemodels.py b/lbmpy/forcemodels.py
index 78aa9969..dd3296fd 100644
--- a/lbmpy/forcemodels.py
+++ b/lbmpy/forcemodels.py
@@ -91,7 +91,9 @@ import sympy as sp
 
 from pystencils.sympyextensions import scalar_product
 
-from lbmpy.maxwellian_equilibrium import get_equilibrium_values_of_maxwell_boltzmann_function
+from lbmpy.maxwellian_equilibrium import (
+    discrete_maxwellian_equilibrium, get_equilibrium_values_of_maxwell_boltzmann_function
+)
 from lbmpy.moments import (
     MOMENT_SYMBOLS, exponent_tuple_sort_key, exponents_to_polynomial_representations,
     extract_monomials, moment_sort_key,
@@ -304,14 +306,15 @@ class He(AbstractForceModel):
         super(He, self).__init__(force)
 
     def forcing_terms(self, lb_method):
-        rho = lb_method.zeroth_order_equilibrium_moment_symbol
         u = sp.Matrix(lb_method.first_order_equilibrium_moment_symbols)
         force = sp.Matrix(self.symbolic_force_vector)
 
         cs_sq = sp.Rational(1, 3)  # squared speed of sound
         # eq. 6.31 in the LB book by Krüger et al. shows that the equilibrium terms are devided by rho.
         # This is done here with subs({rho: 1}) to be consistent with compressible and incompressible force models
-        eq_terms = lb_method.get_equilibrium_terms().subs({rho: 1})
+        cqc = lb_method.conserved_quantity_computation
+        eq_terms = discrete_maxwellian_equilibrium(lb_method.stencil, rho=sp.Integer(1),
+                                                   u=cqc.velocity_symbols, c_s_sq=sp.Rational(1, 3))
 
         result = []
         for direction, eq in zip(lb_method.stencil, eq_terms):
diff --git a/lbmpy_tests/test_force.py b/lbmpy_tests/test_force.py
index 828f974a..f6f7b5a8 100644
--- a/lbmpy_tests/test_force.py
+++ b/lbmpy_tests/test_force.py
@@ -43,8 +43,10 @@ def test_total_momentum(method_enum, zero_centered, force_model, omega):
 
     collision_kernel = ps.create_kernel(collision, config=config).compile()
 
+    fluid_density = 1.1
+
     def init():
-        dh.fill(ρ.name, 1)
+        dh.fill(ρ.name, fluid_density)
         dh.fill(u.name, 0)
 
         setter = macroscopic_values_setter(collision.method, velocity=(0,) * dh.dim,
@@ -70,8 +72,13 @@ def test_total_momentum(method_enum, zero_centered, force_model, omega):
     init()
     time_loop(t)
     dh.run_kernel(getter_kernel)
+
+    #   Check that density did not change
+    assert_allclose(dh.gather_array(ρ.name), fluid_density)
+
+    #   Check for correct velocity
     total = np.sum(dh.gather_array(u.name), axis=(0, 1))
-    assert_allclose(total / np.prod(L) / F / t, 1)
+    assert_allclose(total / np.prod(L) / F * fluid_density / t, 1)
 
 
 @pytest.mark.parametrize("force_model", force_models)
-- 
GitLab