Fix: momentum density calculation

lbmpy/forcemodels.py

@@ -113,6 +113,9 @@ class Simple:
     def macroscopic_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
+    def macroscopic_momentum_density_shift(self, density):
+        return default_momentum_density_shift(self._force)
+
 
 class Luo:
     r"""Force model by Luo :cite:`luo1993lattice`.
@@ -135,6 +138,9 @@ class Luo:
     def macroscopic_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
+    def macroscopic_momentum_density_shift(self, density):
+        return default_momentum_density_shift(self._force)
+
 
 class Guo:
     r"""
@@ -155,6 +161,9 @@ class Guo:
     def macroscopic_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
+    def macroscopic_momentum_density_shift(self, density):
+        return default_momentum_density_shift(self._force)
+
     def equilibrium_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
@@ -187,6 +196,9 @@ class Schiller:
     def macroscopic_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
+    def macroscopic_momentum_density_shift(self, density):
+        return default_momentum_density_shift(self._force)
+
 
 class Buick:
     r"""
@@ -208,6 +220,9 @@ class Buick:
     def macroscopic_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
+    def macroscopic_momentum_density_shift(self, density):
+        return default_momentum_density_shift(self._force)
+
     def equilibrium_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
@@ -231,9 +246,16 @@ class EDM:
     def macroscopic_velocity_shift(self, density):
         return default_velocity_shift(density, self._force)
 
+    def macroscopic_momentum_density_shift(self, density):
+        return default_momentum_density_shift(self._force)
+
 
 # --------------------------------  Helper functions  ------------------------------------------------------------------
 
 
 def default_velocity_shift(density, force):
     return [f_i / (2 * density) for f_i in force]
+
+
+def default_momentum_density_shift(force):
+    return [f_i / 2 for f_i in force]

lbmpy/methods/conservedquantitycomputation.py

@@ -207,8 +207,8 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
             mom_density_eq_coll = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
                                                                                   self._symbolOrder0,
                                                                                   self._symbolsOrder1)
-            mom_density_eq_coll = apply_force_model_shift('macroscopic_velocity_shift', dim, mom_density_eq_coll,
-                                                          self._forceModel, self._compressible)
+            mom_density_eq_coll = apply_force_model_shift('macroscopic_momentum_density_shift', dim, 
+                                                          mom_density_eq_coll, self._forceModel, self._compressible)
             for sym, val in zip(momentum_density_output_symbols, mom_density_eq_coll.main_assignments[1:]):
                 main_assignments.append(Assignment(sym, val.rhs))
         if 'moment0' in output_quantity_names_to_symbols:

lbmpy_tests/test_force.py

@@ -155,3 +155,40 @@ def test_modes(stencil, force_model):
     elif force_model == "simple":
         # All other moments should be zero
         assert list(force_moments[dim+1:]) == [0] * (len(stencil)-(dim+1))
+
+
+@pytest.mark.parametrize("force_model", force_models)
+def test_momentum_density_shift(force_model):
+    target = 'cpu'
+
+    stencil = get_stencil('D2Q9')
+    domain_size = (4, 4)
+    dh = ps.create_data_handling(domain_size=domain_size, default_target=target)
+
+    rho = dh.add_array('rho', values_per_cell=1)
+    dh.fill('rho', 0.0, ghost_layers=True)
+
+    velField = dh.add_array('velField', values_per_cell=dh.dim)
+    dh.fill('velField', 0.0, ghost_layers=True)
+
+    momentum_density = dh.add_array('momentum_density', values_per_cell=dh.dim)
+    dh.fill('momentum_density', 0.0, ghost_layers=True)
+
+    src = dh.add_array('src', values_per_cell=len(stencil))
+    dh.fill('src', 0.0, ghost_layers=True)
+
+    method = create_lb_method(method="srt", compressible=True, force_model=force_model, force=[1, 2])
+
+    momentum_density_symbols = sp.symbols("md_:2")
+    cqc = method.conserved_quantity_computation
+
+    momentum_density_getter = cqc.output_equations_from_pdfs(src.center_vector,
+                                                             {'density': rho.center,
+                                                              'momentum_density': momentum_density.center_vector})
+
+    momentum_density_ast = ps.create_kernel(momentum_density_getter, target=dh.default_target)
+    momentum_density_kernel = momentum_density_ast.compile()
+
+    dh.run_kernel(momentum_density_kernel)
+    assert np.sum(dh.gather_array(momentum_density.name)[:, :, 0]) == np.prod(domain_size) / 2
+    assert np.sum(dh.gather_array(momentum_density.name)[:, :, 1]) == np.prod(domain_size)