src/lbmpy/updatekernels.py

@@ -120,6 +120,26 @@ def create_stream_pull_with_output_kernel(lb_method, src_field, dst_field=None,
                             simplification_hints=output_eq_collection.simplification_hints)
 
 
+def create_copy_kernel(stencil, src_field, dst_field, accessor=StreamPullTwoFieldsAccessor()):
+    """Creates a copy kernel, which can be used to transfer information from src to dst field.
+
+    Args:
+        stencil: lattice Boltzmann stencil which is used in the form of a tuple of tuples
+        src_field: field used for reading pdf values
+        dst_field: field used for writing pdf values
+        accessor: instance of PdfFieldAccessor, defining where to read and write values
+                  to create e.g. a fused stream-collide kernel See 'fieldaccess.PdfFieldAccessor'
+
+    Returns:
+        AssignmentCollection of a copy update rule
+    """
+
+    temporary_symbols = sp.symbols(f'copied:{stencil.Q}')
+    subexpressions = [Assignment(tmp, acc) for tmp, acc in zip(temporary_symbols, accessor.write(src_field, stencil))]
+    main_assignments = [Assignment(acc, tmp) for acc, tmp in zip(accessor.write(dst_field, stencil), temporary_symbols)]
+    return AssignmentCollection(main_assignments, subexpressions=subexpressions)
+
+
 # ---------------------------------- Pdf array creation for various layouts --------------------------------------------
 
 

tests/advanced_streaming/test_communication.py

@@ -5,36 +5,52 @@ import numpy as np
 from lbmpy.stencils import LBStencil
 from pystencils.slicing import get_slice_before_ghost_layer, get_ghost_region_slice
 from lbmpy.creationfunctions import create_lb_update_rule, LBMConfig, LBMOptimisation
-from lbmpy.advanced_streaming.communication import get_communication_slices, _fix_length_one_slices, \
-    LBMPeriodicityHandling
+from lbmpy.advanced_streaming.communication import (
+    get_communication_slices,
+    _fix_length_one_slices,
+    LBMPeriodicityHandling,
+    periodic_pdf_gpu_copy_kernel,
+)
 from lbmpy.advanced_streaming.utility import streaming_patterns, Timestep
 from lbmpy.enums import Stencil
 
 import pytest
 
 
-@pytest.mark.parametrize('stencil', [Stencil.D2Q9, Stencil.D3Q15, Stencil.D3Q19, Stencil.D3Q27])
-@pytest.mark.parametrize('streaming_pattern', streaming_patterns)
-@pytest.mark.parametrize('timestep', [Timestep.EVEN, Timestep.ODD])
+@pytest.mark.parametrize(
+    "stencil", [Stencil.D2Q9, Stencil.D3Q15, Stencil.D3Q19, Stencil.D3Q27]
+)
+@pytest.mark.parametrize("streaming_pattern", streaming_patterns)
+@pytest.mark.parametrize("timestep", [Timestep.EVEN, Timestep.ODD])
 def test_slices_not_empty(stencil, streaming_pattern, timestep):
     stencil = LBStencil(stencil)
     arr = np.zeros((4,) * stencil.D + (stencil.Q,))
-    slices = get_communication_slices(stencil, streaming_pattern=streaming_pattern, prev_timestep=timestep,
-                                      ghost_layers=1)
+    slices = get_communication_slices(
+        stencil,
+        streaming_pattern=streaming_pattern,
+        prev_timestep=timestep,
+        ghost_layers=1,
+    )
     for _, slices_list in slices.items():
         for src, dst in slices_list:
             assert all(s != 0 for s in arr[src].shape)
             assert all(s != 0 for s in arr[dst].shape)
 
 
-@pytest.mark.parametrize('stencil', [Stencil.D2Q9, Stencil.D3Q15, Stencil.D3Q19, Stencil.D3Q27])
-@pytest.mark.parametrize('streaming_pattern', streaming_patterns)
-@pytest.mark.parametrize('timestep', [Timestep.EVEN, Timestep.ODD])
+@pytest.mark.parametrize(
+    "stencil", [Stencil.D2Q9, Stencil.D3Q15, Stencil.D3Q19, Stencil.D3Q27]
+)
+@pytest.mark.parametrize("streaming_pattern", streaming_patterns)
+@pytest.mark.parametrize("timestep", [Timestep.EVEN, Timestep.ODD])
 def test_src_dst_same_shape(stencil, streaming_pattern, timestep):
     stencil = LBStencil(stencil)
     arr = np.zeros((4,) * stencil.D + (stencil.Q,))
-    slices = get_communication_slices(stencil, streaming_pattern=streaming_pattern, prev_timestep=timestep,
-                                      ghost_layers=1)
+    slices = get_communication_slices(
+        stencil,
+        streaming_pattern=streaming_pattern,
+        prev_timestep=timestep,
+        ghost_layers=1,
+    )
     for _, slices_list in slices.items():
         for src, dst in slices_list:
             src_shape = arr[src].shape
@@ -42,12 +58,15 @@ def test_src_dst_same_shape(stencil, streaming_pattern, timestep):
             assert src_shape == dst_shape
 
 
-@pytest.mark.parametrize('stencil', [Stencil.D2Q9, Stencil.D3Q15, Stencil.D3Q19, Stencil.D3Q27])
+@pytest.mark.parametrize(
+    "stencil", [Stencil.D2Q9, Stencil.D3Q15, Stencil.D3Q19, Stencil.D3Q27]
+)
 def test_pull_communication_slices(stencil):
     stencil = LBStencil(stencil)
 
     slices = get_communication_slices(
-        stencil, streaming_pattern='pull', prev_timestep=Timestep.BOTH, ghost_layers=1)
+        stencil, streaming_pattern="pull", prev_timestep=Timestep.BOTH, ghost_layers=1
+    )
     for i, d in enumerate(stencil):
         if i == 0:
             continue
@@ -58,21 +77,115 @@ def test_pull_communication_slices(stencil):
                 dst = s[1][:-1]
                 break
 
-        inner_slice = _fix_length_one_slices(get_slice_before_ghost_layer(d, ghost_layers=1))
+        inner_slice = _fix_length_one_slices(
+            get_slice_before_ghost_layer(d, ghost_layers=1)
+        )
         inv_dir = (-e for e in d)
-        gl_slice = _fix_length_one_slices(get_ghost_region_slice(inv_dir, ghost_layers=1))
+        gl_slice = _fix_length_one_slices(
+            get_ghost_region_slice(inv_dir, ghost_layers=1)
+        )
         assert src == inner_slice
         assert dst == gl_slice
 
 
-@pytest.mark.parametrize('stencil_name', [Stencil.D2Q9, Stencil.D3Q15, Stencil.D3Q19, Stencil.D3Q27])
+@pytest.mark.parametrize("direction", LBStencil(Stencil.D3Q27).stencil_entries)
+@pytest.mark.parametrize("pull", [False, True])
+def test_gpu_comm_kernels(direction: tuple, pull: bool):
+    pytest.importorskip("cupy")
+
+    stencil = LBStencil(Stencil.D3Q27)
+    inv_dir = stencil[stencil.inverse_index(direction)]
+    target = ps.Target.GPU
+
+    domain_size = (4,) * stencil.D
+
+    dh: ps.datahandling.SerialDataHandling = ps.create_data_handling(
+        domain_size,
+        periodicity=(True,) * stencil.D,
+        parallel=False,
+        default_target=target,
+    )
+
+    field = dh.add_array("field", values_per_cell=2)
+
+    if pull:
+        dst_slice = get_ghost_region_slice(inv_dir)
+        src_slice = get_slice_before_ghost_layer(direction)
+    else:
+        dst_slice = get_slice_before_ghost_layer(direction)
+        src_slice = get_ghost_region_slice(inv_dir)
+
+    src_slice += (1,)
+    dst_slice += (1,)
+
+    kernel = periodic_pdf_gpu_copy_kernel(field, src_slice, dst_slice)
+
+    dh.cpu_arrays[field.name][src_slice] = 42.0
+    dh.all_to_gpu()
+
+    dh.run_kernel(kernel)
+
+    dh.all_to_cpu()
+    np.testing.assert_equal(dh.cpu_arrays[field.name][dst_slice], 42.0)
+
+
+@pytest.mark.parametrize("stencil", [Stencil.D2Q9, Stencil.D3Q19])
+@pytest.mark.parametrize("streaming_pattern", streaming_patterns)
+def test_direct_copy_and_kernels_equivalence(stencil: Stencil, streaming_pattern: str):
+    pytest.importorskip("cupy")
+
+    target = ps.Target.GPU
+    stencil = LBStencil(stencil)
+    domain_size = (4,) * stencil.D
+
+    dh: ps.datahandling.SerialDataHandling = ps.create_data_handling(
+        domain_size,
+        periodicity=(True,) * stencil.D,
+        parallel=False,
+        default_target=target,
+    )
+
+    pdfs_a = dh.add_array("pdfs_a", values_per_cell=stencil.Q)
+    pdfs_b = dh.add_array("pdfs_b", values_per_cell=stencil.Q)
+
+    dh.fill(pdfs_a.name, 0.0, ghost_layers=True)
+    dh.fill(pdfs_b.name, 0.0, ghost_layers=True)
+
+    for q in range(stencil.Q):
+        sl = ps.make_slice[:4, :4, q] if stencil.D == 2 else ps.make_slice[:4, :4, :4, q]
+        dh.cpu_arrays[pdfs_a.name][sl] = q
+        dh.cpu_arrays[pdfs_b.name][sl] = q
+
+    dh.all_to_gpu()
+
+    direct_copy = LBMPeriodicityHandling(stencil, dh, pdfs_a.name, streaming_pattern, cupy_direct_copy=True)
+    kernels_copy = LBMPeriodicityHandling(stencil, dh, pdfs_b.name, streaming_pattern, cupy_direct_copy=False)
+
+    direct_copy(Timestep.EVEN)
+    kernels_copy(Timestep.EVEN)
+
+    dh.all_to_cpu()
+
+    np.testing.assert_equal(
+        dh.cpu_arrays[pdfs_a.name],
+        dh.cpu_arrays[pdfs_b.name]
+    )
+
+
+@pytest.mark.parametrize(
+    "stencil_name", [Stencil.D2Q9, Stencil.D3Q15, Stencil.D3Q19, Stencil.D3Q27]
+)
 def test_optimised_and_full_communication_equivalence(stencil_name):
     target = ps.Target.CPU
     stencil = LBStencil(stencil_name)
-    domain_size = (4, ) * stencil.D
+    domain_size = (4,) * stencil.D
 
-    dh = ps.create_data_handling(domain_size, periodicity=(True, ) * stencil.D,
-                                 parallel=False, default_target=target)
+    dh = ps.create_data_handling(
+        domain_size,
+        periodicity=(True,) * stencil.D,
+        parallel=False,
+        default_target=target,
+    )
 
     pdf = dh.add_array("pdf", values_per_cell=len(stencil), dtype=np.int64)
     dh.fill("pdf", 0, ghost_layers=True)
@@ -82,9 +195,9 @@ def test_optimised_and_full_communication_equivalence(stencil_name):
     gl = dh.ghost_layers_of_field("pdf")
 
     num = 0
-    for idx, x in np.ndenumerate(dh.cpu_arrays['pdf']):
-        dh.cpu_arrays['pdf'][idx] = num
-        dh.cpu_arrays['pdf_tmp'][idx] = num
+    for idx, x in np.ndenumerate(dh.cpu_arrays["pdf"]):
+        dh.cpu_arrays["pdf"][idx] = num
+        dh.cpu_arrays["pdf_tmp"][idx] = num
         num += 1
 
     lbm_config = LBMConfig(stencil=stencil, kernel_type="stream_pull_only")
@@ -95,21 +208,27 @@ def test_optimised_and_full_communication_equivalence(stencil_name):
     ast = ps.create_kernel(ac, config=config)
     stream = ast.compile()
 
-    full_communication = dh.synchronization_function(pdf.name, target=dh.default_target, optimization={"openmp": True})
+    full_communication = dh.synchronization_function(
+        pdf.name, target=dh.default_target, optimization={"openmp": True}
+    )
     full_communication()
 
     dh.run_kernel(stream)
     dh.swap("pdf", "pdf_tmp")
-    pdf_full_communication = np.copy(dh.cpu_arrays['pdf'])
+    pdf_full_communication = np.copy(dh.cpu_arrays["pdf"])
 
     num = 0
-    for idx, x in np.ndenumerate(dh.cpu_arrays['pdf']):
-        dh.cpu_arrays['pdf'][idx] = num
-        dh.cpu_arrays['pdf_tmp'][idx] = num
+    for idx, x in np.ndenumerate(dh.cpu_arrays["pdf"]):
+        dh.cpu_arrays["pdf"][idx] = num
+        dh.cpu_arrays["pdf_tmp"][idx] = num
         num += 1
 
-    optimised_communication = LBMPeriodicityHandling(stencil=stencil, data_handling=dh, pdf_field_name=pdf.name,
-                                                     streaming_pattern='pull')
+    optimised_communication = LBMPeriodicityHandling(
+        stencil=stencil,
+        data_handling=dh,
+        pdf_field_name=pdf.name,
+        streaming_pattern="pull",
+    )
     optimised_communication()
     dh.run_kernel(stream)
     dh.swap("pdf", "pdf_tmp")
@@ -119,9 +238,14 @@ def test_optimised_and_full_communication_equivalence(stencil_name):
             for j in range(gl, domain_size[1]):
                 for k in range(gl, domain_size[2]):
                     for f in range(len(stencil)):
-                        assert dh.cpu_arrays['pdf'][i, j, k, f] == pdf_full_communication[i, j, k, f], print(f)
+                        assert (
+                            dh.cpu_arrays["pdf"][i, j, k, f]
+                            == pdf_full_communication[i, j, k, f]
+                        ), print(f)
     else:
         for i in range(gl, domain_size[0]):
             for j in range(gl, domain_size[1]):
                 for f in range(len(stencil)):
-                    assert dh.cpu_arrays['pdf'][i, j, f] == pdf_full_communication[i, j, f]
+                    assert (
+                        dh.cpu_arrays["pdf"][i, j, f] == pdf_full_communication[i, j, f]
+                    )

tests/cumulantmethod/test_flow_around_sphere.py

@@ -58,7 +58,13 @@ def flow_around_sphere(stencil, galilean_correction, fourth_order_correction, L_
             boundary_assignments)
         iter_slice = get_slice_before_ghost_layer((1,) + (0,) * (stencil.D - 1))
         extrapolation_ast = create_kernel(
-            boundary_assignments, config=CreateKernelConfig(iteration_slice=iter_slice, ghost_layers=1, target=target))
+            boundary_assignments,
+            config=CreateKernelConfig(
+                iteration_slice=iter_slice,
+                ghost_layers=1,
+                target=target,
+                skip_independence_check=True)
+            )
         return extrapolation_ast.compile()
 
     dh = create_data_handling(channel_size, periodicity=False, default_layout='fzyx', default_target=target)

tests/full_scenarios/shear_wave/scenario_shear_wave.py

@@ -202,19 +202,19 @@ def create_full_parameter_study():
 
     config = CreateKernelConfig(target=Target.CPU)
 
-    methods = [LBMConfig(method=Method.TRT, relaxation_rates=[omega, 1]),
-               LBMConfig(method=Method.MRT, relaxation_rates=[omega], equilibrium_order=2),
-               LBMConfig(method=Method.MRT, relaxation_rates=[omega], equilibrium_order=3),
-               LBMConfig(method=Method.CUMULANT, relaxation_rates=[omega],  # cumulant
+    methods = [LBMConfig(stencil=LBStencil(Stencil.D3Q27), method=Method.TRT, relaxation_rates=[omega, 1]),
+               LBMConfig(stencil=LBStencil(Stencil.D3Q27), method=Method.MRT, relaxation_rates=[omega],
+                         equilibrium_order=2),
+               LBMConfig(stencil=LBStencil(Stencil.D3Q27), method=Method.MRT, relaxation_rates=[omega],
+                         equilibrium_order=3),
+               LBMConfig(stencil=LBStencil(Stencil.D3Q27), method=Method.CUMULANT, relaxation_rates=[omega],  # cumulant
                          compressible=True, continuous_equilibrium=True, equilibrium_order=3),
-               LBMConfig(method=Method.CUMULANT, relaxation_rates=[omega],  # cumulant with fourth-order correction
+               LBMConfig(stencil=LBStencil(Stencil.D3Q27), method=Method.CUMULANT, relaxation_rates=[omega],  # cumulant with fourth-order correction
                          compressible=True, continuous_equilibrium=True, fourth_order_correction=0.1),
-               LBMConfig(method=Method.TRT_KBC_N4, relaxation_rates=[omega, omega_f], entropic=True,
-                         zero_centered=False,  # entropic order 2
-                         continuous_equilibrium=True, equilibrium_order=2),
-               LBMConfig(method=Method.TRT_KBC_N4, relaxation_rates=[omega, omega_f], entropic=True,
-                         zero_centered=False,  # entropic order 3
-                         continuous_equilibrium=True, equilibrium_order=3),
+               LBMConfig(stencil=LBStencil(Stencil.D3Q27), method=Method.TRT_KBC_N4, relaxation_rates=[omega, omega_f],
+                         entropic=True, zero_centered=False, continuous_equilibrium=True, equilibrium_order=2),
+               LBMConfig(stencil=LBStencil(Stencil.D3Q27), method=Method.TRT_KBC_N4, relaxation_rates=[omega, omega_f],
+                         entropic=True, zero_centered=False, continuous_equilibrium=True, equilibrium_order=3),
 
                # entropic cumulant: not supported for the moment
                # LBMConfig(method=Method.CUMULANT, relaxation_rates=["omega", "omega_f"], entropic=True, zero_centered=False,

tests/phasefield/test_n_phase_boyer_noncoupled.ipynb

 %% Cell type:code id: tags:
 
 ``` python
 import pytest
 pytest.importorskip('cupy')
 ```
 
 %% Output
 
     <module 'cupy' from '/home/markus/.local/lib/python3.11/site-packages/cupy/__init__.py'>
 
 %% Cell type:code id: tags:
 
 ``` python
 from lbmpy.session import *
 from lbmpy.phasefield.n_phase_boyer import *
 from lbmpy.phasefield.kerneleqs import *
 from lbmpy.phasefield.contact_angle_circle_fitting import *
-from scipy.ndimage.filters import gaussian_filter
+from scipy.ndimage import gaussian_filter
 from pystencils.simp import sympy_cse_on_assignment_list
 one = sp.sympify(1)
 
 import pyximport
 pyximport.install(language_level=3)
 from lbmpy.phasefield.simplex_projection import simplex_projection_2d  # NOQA
 ```
 
 %% Cell type:markdown id: tags:
 
 # Simulation arbitrary surface tension case
 
 %% Cell type:code id: tags:
 
 ``` python
 n = 4
 dx, dt = 1, 1
 mobility = 2e-3
 domain_size = (150, 150)
 ε = one * 4
 penalty_factor = 0
 stabilization_factor = 10
 
 κ = (one,  one/2, one/3, one/4)
 sigma_factor = one / 15
 σ = sp.ImmutableDenseMatrix(n, n, lambda i,j: sigma_factor* (κ[i] + κ[j]) if i != j else 0 )
 #σ
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 dh = create_data_handling(domain_size, periodicity=True, default_target=ps.Target.GPU)
 c = dh.add_array('c', values_per_cell=n)
 c_tmp = dh.add_array_like('c_tmp', 'c')
 
 μ = dh.add_array('mu', values_per_cell=n)
 
 cvec = c.center_vector
 μvec = μ.center_vector
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 α, _ = diffusion_coefficients(σ)
 
 f = lambda c: c**2 * ( 1 - c ) **2
 a, b = compute_ab(f)
 
 capital_f = capital_f0(cvec, σ) + correction_g(cvec, σ) + stabilization_factor * stabilization_term(cvec, α)
 
 f_bulk = free_energy_bulk(capital_f, b, ε) + penalty_factor * (one - sum(cvec))
 f_if = free_energy_interfacial(cvec, σ, a, ε)
 f = f_bulk + f_if
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 #f_bulk
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 μ_assignments = mu_kernel(f, cvec, c, μ)
 μ_assignments = [Assignment(a.lhs, a.rhs.doit()) for a in μ_assignments]
 μ_assignments = sympy_cse_on_assignment_list(μ_assignments)
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 discretize = fd.Discretization2ndOrder(dx=dx, dt=dt)
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 def lapl(e):
     return sum(ps.fd.diff(e, d, d) for d in range(dh.dim))
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 rhs = α * μvec
 discretized_rhs = [discretize(fd.expand_diff_full( lapl(mobility * rhs_i) + fd.transient(cvec[i], idx=i), functions=μvec))
                    for i, rhs_i in enumerate(rhs)]
 c_assignments = [Assignment(lhs, rhs)
                  for lhs, rhs in zip(c_tmp.center_vector, discretized_rhs)]
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 #c_assignments
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 μ_sync = dh.synchronization_function(μ.name)
 c_sync = dh.synchronization_function(c.name)
 optimization = {'cpu_openmp': False, 'cpu_vectorize_info': None}
 
 config = ps.CreateKernelConfig(cpu_openmp=False, target=dh.default_target)
 
 μ_kernel = create_kernel(μ_assignments, config=config).compile()
 c_kernel = create_kernel(c_assignments, config=config).compile()
 
 def set_c(slice_obj, values):
     for block in dh.iterate(slice_obj):
         arr = block[c.name]
         arr[..., : ] = values
 
 def smooth():
     for block in dh.iterate(ghost_layers=True):
         c_arr = block[c.name]
         for i in range(n):
             gaussian_filter(c_arr[..., i], sigma=2, output=c_arr[..., i])
 
 def time_loop(steps):
     dh.all_to_gpu()
     for t in range(steps):
         c_sync()
         dh.run_kernel(μ_kernel)
         μ_sync()
         dh.run_kernel(c_kernel)
         dh.swap(c.name, c_tmp.name)
         #simplex_projection_2d(dh.cpu_arrays[c.name])
     dh.all_to_cpu()
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 set_c(make_slice[:, :], [0, 0, 0, 0])
 set_c(make_slice[:, 0.5:], [1, 0, 0, 0])
 set_c(make_slice[:, :0.5], [0, 1, 0, 0])
 set_c(make_slice[0.3:0.7, 0.3:0.7], [0, 0, 1, 0])
 smooth()
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 #dh.load_all('n_phases_state_size200_stab10.npz')
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 plt.phase_plot(dh.gather_array(c.name))
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 neumann_angles_from_surface_tensions(lambda i, j: float(σ[i, j]))
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 import time
 for i in range(10):
     start = time.perf_counter()
     time_loop(1_000)
     end = time.perf_counter()
 
     try:
         print(i, end - start, liquid_lens_neumann_angles(dh.gather_array(c.name)))
     except Exception:
         print(i, end - start, "none found")
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 plt.subplot(1,3,1)
 t = dh.gather_array(c.name, make_slice[25, :]).squeeze()
 plt.plot(t);
 
 plt.subplot(1,3,2)
 plt.phase_plot(dh.gather_array(c.name), linewidth=1)
 
 plt.subplot(1,3,3)
 plt.scalar_field(dh.gather_array(μ.name)[:, :, 2])
 plt.colorbar();
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 assert not np.isnan(dh.max(c.name))
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 t = dh.gather_array(c.name, make_slice[25, 55:90]).squeeze()
 plt.hlines(0.5, 0, 30)
 plt.plot(t);
 ```

tests/phasefield_allen_cahn/test_analytical.py

@@ -44,6 +44,6 @@ def test_analytical():
     assert np.isclose(parameters.density_heavy, 1.0)
     assert np.isclose(parameters.density_light, 0.001207114228456914)
     assert np.isclose(parameters.dynamic_viscosity_heavy, 5.733727652152216e-05)
-    assert np.isclose(parameters.dynamic_viscosity_light, 0.0008630017037694861)
+    assert np.isclose(parameters.dynamic_viscosity_light, 1.0417416358027054e-06)
     assert np.isclose(parameters.gravitational_acceleration, -7.407407407407407e-08)
     assert np.isclose(parameters.surface_tension, 3.149857262258028e-05, rtol=1e-05)

tests/test_boundary_handling.py

 import numpy as np
 import pytest
 
-from lbmpy.boundaries import NoSlip, UBB, SimpleExtrapolationOutflow, ExtrapolationOutflow, \
-    FixedDensity, DiffusionDirichlet, NeumannByCopy, StreamInConstant, FreeSlip
-from lbmpy.boundaries.boundaryhandling import LatticeBoltzmannBoundaryHandling
+from lbmpy.boundaries import (NoSlip, NoSlipLinearBouzidi, QuadraticBounceBack,
+                              UBB, SimpleExtrapolationOutflow, ExtrapolationOutflow, FixedDensity, DiffusionDirichlet,
+                              NeumannByCopy, StreamInConstant, FreeSlip)
+from lbmpy.boundaries.boundaryhandling import LatticeBoltzmannBoundaryHandling, create_lattice_boltzmann_boundary_kernel
 from lbmpy.creationfunctions import create_lb_function, create_lb_method, LBMConfig
 from lbmpy.enums import Stencil, Method
 from lbmpy.geometry import add_box_boundary
 from lbmpy.lbstep import LatticeBoltzmannStep
 from lbmpy.stencils import LBStencil
+import pystencils as ps
 from pystencils import create_data_handling, make_slice, Target, CreateKernelConfig
 from pystencils.slicing import slice_from_direction
 from pystencils.stencil import inverse_direction
@@ -435,3 +437,45 @@ def test_boundary_utility_functions():
     assert stream == StreamInConstant(constant=1.0, name="stream")
     assert not stream == StreamInConstant(constant=1.0, name="test")
     assert not stream == noslip
+
+
+@pytest.mark.parametrize("given_force_vector", [True, False])
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_force_on_boundary(given_force_vector, dtype):
+    stencil = LBStencil(Stencil.D2Q9)
+    pdfs = ps.fields(f"pdfs_src({stencil.Q}): {dtype}[{stencil.D}D]", layout='fzyx')
+
+    method = create_lb_method(lbm_config=LBMConfig(stencil=stencil, method=Method.SRT, relaxation_rate=1.8))
+
+    noslip = NoSlip(name="noslip", calculate_force_on_boundary=True)
+    bouzidi = NoSlipLinearBouzidi(name="bouzidi", calculate_force_on_boundary=True)
+    qq_bounce_Back = QuadraticBounceBack(name="qqBB", relaxation_rate=1.8, calculate_force_on_boundary=True)
+
+    boundary_objects = [noslip, bouzidi, qq_bounce_Back]
+    for boundary in boundary_objects:
+        if given_force_vector:
+            force_vector_type = np.dtype([(f"F_{i}", dtype) for i in range(stencil.D)], align=True)
+            force_vector = ps.Field('forceVector', ps.FieldType.INDEXED, force_vector_type, layout=[0],
+                                    shape=(ps.TypedSymbol("forceVectorSize", "int32"), 1), strides=(1, 1))
+        else:
+            force_vector = None
+
+        index_struct_dtype = _numpy_data_type_for_boundary_object(boundary, stencil.D)
+
+        index_field = ps.Field('indexVector', ps.FieldType.INDEXED, index_struct_dtype, layout=[0],
+                               shape=(ps.TypedSymbol("indexVectorSize", "int32"), 1), strides=(1, 1))
+
+        create_lattice_boltzmann_boundary_kernel(pdfs, index_field, method, boundary, force_vector=force_vector)
+
+
+def _numpy_data_type_for_boundary_object(boundary_object, dim):
+    boundary_index_array_coordinate_names = ["x", "y", "z"]
+    direction_member_name = "dir"
+    default_index_array_dtype = np.int32
+
+    coordinate_names = boundary_index_array_coordinate_names[:dim]
+    return np.dtype(
+        [(name, default_index_array_dtype) for name in coordinate_names]
+        + [(direction_member_name, default_index_array_dtype)]
+        + [(i[0], i[1].numpy_dtype) for i in boundary_object.additional_data],
+        align=True,)

tests/test_macroscopic_value_kernels.py

 import pytest
 import numpy as np
+import sympy as sp
 
 from lbmpy.creationfunctions import create_lb_method, LBMConfig
 from lbmpy.enums import Stencil, ForceModel, Method
 from lbmpy.macroscopic_value_kernels import (
-    compile_macroscopic_values_getter, compile_macroscopic_values_setter)
+    compile_macroscopic_values_getter, compile_macroscopic_values_setter, strain_rate_tensor_getter)
 from lbmpy.advanced_streaming.utility import streaming_patterns, Timestep
 from lbmpy.stencils import LBStencil
 
@@ -83,3 +84,18 @@ def test_set_get_constant_velocity(stencil, force_model, compressible):
         computed_velocity = velocity_output_field[0, 0, 0, :]
 
     np.testing.assert_almost_equal(np.array(ref_velocity[:method.dim]), computed_velocity)
+
+
+@pytest.mark.parametrize('stencil', [Stencil.D2Q9, Stencil.D3Q19])
+@pytest.mark.parametrize("method_enum", [Method.SRT, Method.TRT, Method.MRT, Method.CUMULANT])
+def test_get_strain_rate_tensor(stencil, method_enum):
+    stencil = LBStencil(stencil)
+    lbm_config = LBMConfig(stencil=stencil, method=method_enum, compressible=True)
+    method = create_lb_method(lbm_config=lbm_config)
+
+    pdfs = sp.symbols(f"f_:{stencil.Q}")
+    strain_rate_tensor = sp.symbols(f"strain_rate_tensor_:{method.dim * method.dim}")
+
+    getter_assignments = strain_rate_tensor_getter(method, strain_rate_tensor, pdfs)
+
+    assert len(getter_assignments) == method.dim * method.dim

tests/test_psm.py

+import pytest
+import numpy as np
+
+from pystencils import fields, CreateKernelConfig, Target, create_kernel, get_code_str, create_data_handling
+from lbmpy import LBMConfig, Stencil, Method, LBStencil, create_lb_method, create_lb_collision_rule, LBMOptimisation, \
+    create_lb_update_rule
+from lbmpy.partially_saturated_cells import PSMConfig
+
+
+@pytest.mark.parametrize("stencil", [Stencil.D2Q9, Stencil.D3Q19, Stencil.D3Q27])
+@pytest.mark.parametrize("method_enum", [Method.SRT, Method.MRT, Method.CUMULANT])
+def test_psm_integration(stencil, method_enum):
+    stencil = LBStencil(stencil)
+
+    fraction_field = fields("fraction_field: double[10, 10, 5]", layout=(2, 1, 0))
+    object_vel = fields("object_vel(3): double[10, 10, 5]", layout=(3, 2, 1, 0))
+    psm_config = PSMConfig(fraction_field=fraction_field, object_velocity_field=object_vel)
+
+    lbm_config = LBMConfig(stencil=stencil, method=method_enum, relaxation_rate=1.5, compressible=True,
+                           psm_config=psm_config)
+    config = CreateKernelConfig(target=Target.CPU)#
+
+    collision_rule = create_lb_collision_rule(lbm_config=lbm_config)
+
+    ast = create_kernel(collision_rule, config=config)
+    code_str = get_code_str(ast)
+
+
+def get_data_handling_for_psm(use_psm):
+    domain_size = (10, 10)
+    stencil = LBStencil(Stencil.D2Q9)
+    dh = create_data_handling(domain_size=domain_size, periodicity=(True, False))
+
+    f = dh.add_array('f', values_per_cell=len(stencil))
+    dh.fill(f.name, 0.0, ghost_layers=True)
+    f_tmp = dh.add_array('f_tmp', values_per_cell=len(stencil))
+    dh.fill(f_tmp.name, 0.0, ghost_layers=True)
+
+    psm_config = None
+    if use_psm:
+        fraction_field = dh.add_array('fraction_field', values_per_cell=1)
+        dh.fill(fraction_field.name, 0.0, ghost_layers=True)
+        object_vel = dh.add_array('object_vel', values_per_cell=dh.dim)
+        dh.fill(object_vel.name, 0.0, ghost_layers=True)
+        psm_config = PSMConfig(fraction_field=fraction_field, object_velocity_field=object_vel)
+
+    lbm_config = LBMConfig(stencil=stencil, method=Method.SRT, relaxation_rate=1.5, psm_config=psm_config)
+
+    lbm_optimisation = LBMOptimisation(symbolic_field=f, symbolic_temporary_field=f_tmp)
+    update_rule = create_lb_update_rule(lbm_config=lbm_config, lbm_optimisation=lbm_optimisation)
+    kernel_lb_step_with_psm = create_kernel(update_rule).compile()
+    return (dh, kernel_lb_step_with_psm)
+
+
+def test_lbm_vs_psm():
+
+    psm_dh, psm_kernel = get_data_handling_for_psm(True)
+    lbm_dh, lbm_kernel = get_data_handling_for_psm(False)
+
+    for i in range(20):
+        psm_dh.run_kernel(psm_kernel)
+        psm_dh.swap('f', 'f_tmp')
+
+        lbm_dh.run_kernel(lbm_kernel)
+        lbm_dh.swap('f', 'f_tmp')
+
+    max_vel_error = np.max(np.abs(psm_dh.gather_array('f') - lbm_dh.gather_array('f')))
+    np.testing.assert_allclose(max_vel_error, 0, atol=1e-14)

tests/test_welford.py

+import pytest
+import numpy as np
+import pystencils as ps
+from lbmpy.flow_statistics import welford_assignments
+
+
+@pytest.mark.parametrize('order', [1, 2, 3])
+@pytest.mark.parametrize('dim', [1, 2, 3])
+def test_welford(order, dim):
+
+    target = ps.Target.CPU
+
+    # create data handling and fields
+    dh = ps.create_data_handling((1, 1, 1), periodicity=False, default_target=target)
+
+    vector_field = dh.add_array('vector', values_per_cell=dim)
+    dh.fill(vector_field.name, 0.0, ghost_layers=False)
+
+    mean_field = dh.add_array('mean', values_per_cell=dim)
+    dh.fill(mean_field.name, 0.0, ghost_layers=False)
+
+    if order >= 2:
+        sos = dh.add_array('sos', values_per_cell=dim**2)
+        dh.fill(sos.name, 0.0, ghost_layers=False)
+
+        if order == 3:
+            soc = dh.add_array('soc', values_per_cell=dim**3)
+            dh.fill(soc.name, 0.0, ghost_layers=False)
+        else:
+            soc = None
+    else:
+        sos = None
+        soc = None
+
+    welford = welford_assignments(field=vector_field, mean_field=mean_field,
+                                  sum_of_squares_field=sos, sum_of_cubes_field=soc)
+
+    welford_kernel = ps.create_kernel(welford).compile()
+
+    # set random seed
+    np.random.seed(42)
+    n = int(1e4)
+    x = np.random.normal(size=n * dim).reshape(n, dim)
+
+    analytical_mean = np.zeros(dim)
+    analytical_covariance = np.zeros(dim**2)
+    analytical_third_order_moments = np.zeros(dim**3)
+
+    # calculate analytical mean
+    for i in range(dim):
+        analytical_mean[i] = np.mean(x[:, i])
+
+    # calculate analytical covariances
+    for i in range(dim):
+        for j in range(dim):
+            analytical_covariance[i * dim + j] \
+                = (np.sum((x[:, i] - analytical_mean[i]) * (x[:, j] - analytical_mean[j]))) / n
+
+    # calculate analytical third-order central moments
+    for i in range(dim):
+        for j in range(dim):
+            for k in range(dim):
+                analytical_third_order_moments[(i * dim + j) * dim + k] \
+                    = (np.sum((x[:, i] - analytical_mean[i])
+                              * (x[:, j] - analytical_mean[j])
+                              * (x[:, k] - analytical_mean[k]))) / n
+
+    # Time loop
+    counter = 1
+    for i in range(n):
+        for d in range(dim):
+            new_value = x[i, d]
+            if dim > 1:
+                dh.fill(array_name=vector_field.name, val=new_value, value_idx=d, ghost_layers=False)
+            else:
+                dh.fill(array_name=vector_field.name, val=new_value, ghost_layers=False)
+        dh.run_kernel(welford_kernel, counter=counter)
+        counter += 1
+
+    welford_mean = dh.gather_array(mean_field.name).flatten()
+    np.testing.assert_allclose(welford_mean, analytical_mean)
+
+    if order >= 2:
+        welford_covariance = dh.gather_array(sos.name).flatten() / n
+        np.testing.assert_allclose(welford_covariance, analytical_covariance)
+        if order == 3:
+            welford_third_order_moments = dh.gather_array(soc.name).flatten() / n
+            np.testing.assert_allclose(welford_third_order_moments, analytical_third_order_moments)
+
+
+if __name__ == '__main__':
+    test_welford(1, 1)