Adapted Forcing Momentum Test Case

lbmpy_tests/test_force.py

 from pystencils.session import *
 from pystencils import Target
 from lbmpy.session import *
-from lbmpy.macroscopic_value_kernels import macroscopic_values_setter
+from lbmpy.macroscopic_value_kernels import macroscopic_values_setter, macroscopic_values_getter
 import lbmpy.forcemodels
 from lbmpy.moments import is_bulk_moment
 from lbmpy.relaxationrates import get_bulk_relaxation_rate, get_shear_relaxation_rate
@@ -45,17 +45,17 @@ def test_total_momentum(method, force_model, omega):
                                       compressible=True,
                                       force_model=force_model,
                                       force=F,
-                                      kernel_type='collide_only',
+                                      streaming_pattern='pull',
                                       optimization={'symbolic_field': src})
 
-    stream = create_stream_pull_with_output_kernel(collision.method, src, dst,
-                                                   {'density': ρ, 'velocity': u})
+    # stream = create_stream_pull_with_output_kernel(collision.method, src, dst,
+    #                                                {'density': ρ, 'velocity': u})
 
     opts = {'cpu_openmp': True,
             'cpu_vectorize_info': None,
             'target': dh.default_target}
 
-    stream_kernel = ps.create_kernel(stream, **opts).compile()
+    # stream_kernel = ps.create_kernel(stream, **opts).compile()
     collision_kernel = ps.create_kernel(collision, **opts).compile()
 
     def init():
@@ -63,24 +63,28 @@ def test_total_momentum(method, force_model, omega):
         dh.fill(u.name, 0)
 
         setter = macroscopic_values_setter(collision.method, velocity=(0,) * dh.dim,
-                                           pdfs=src.center_vector, density=ρ.center)
-        kernel = ps.create_kernel(setter, ghost_layers=0).compile()
+                                           pdfs=src, density=ρ.center,
+                                           set_pre_collision_pdfs=True)
+        kernel = ps.create_kernel(setter).compile()
         dh.run_kernel(kernel)
 
     sync_pdfs = dh.synchronization_function([src.name])
 
+    getter = macroscopic_values_getter(collision.method, ρ.center, u.center_vector, src, use_pre_collision_pdfs=True)
+    getter_kernel = ps.create_kernel(getter).compile()
+
     def time_loop(steps):
         dh.all_to_gpu()
-        for i in range(steps):
+        for _ in range(steps):
             dh.run_kernel(collision_kernel)
-            sync_pdfs()
-            dh.run_kernel(stream_kernel)
             dh.swap(src.name, dst.name)
+            sync_pdfs()
         dh.all_to_cpu()
 
     t = 20
     init()
     time_loop(t)
+    dh.run_kernel(getter_kernel)
     total = np.sum(dh.gather_array(u.name), axis=(0, 1))
     assert np.allclose(total / np.prod(L) / F / t, 1)