Merge branch 'pressure_tensor' into 'master'

Added second order moment getter

See merge request pycodegen/lbmpy!83

lbmpy/methods/conservedquantitycomputation.py

@@ -84,13 +84,15 @@ class AbstractConservedQuantityComputation(abc.ABC):
 class DensityVelocityComputation(AbstractConservedQuantityComputation):
     def __init__(self, stencil, compressible, force_model=None,
                  zeroth_order_moment_symbol=sp.Symbol("rho"),
-                 first_order_moment_symbols=sp.symbols("u_:3")):
+                 first_order_moment_symbols=sp.symbols("u_:3"),
+                 second_order_moment_symbols=sp.symbols("p_:9")):
         dim = len(stencil[0])
         self._stencil = stencil
         self._compressible = compressible
         self._forceModel = force_model
         self._symbolOrder0 = zeroth_order_moment_symbol
         self._symbolsOrder1 = first_order_moment_symbols[:dim]
+        self._symbolsOrder2 = second_order_moment_symbols[:(dim * dim)]
 
     @property
     def conserved_quantities(self):
@@ -166,7 +168,8 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
         dim = len(self._stencil[0])
 
         ac = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
-                                                             self._symbolOrder0, self._symbolsOrder1)
+                                                             self._symbolOrder0, self._symbolsOrder1,
+                                                             self._symbolsOrder2)
 
         if self._compressible:
             ac = divide_first_order_moments_by_rho(ac, dim)
@@ -209,6 +212,7 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
                                                                                   self._symbolsOrder1)
             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:
@@ -222,6 +226,13 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
                 moment1_output_symbol = moment1_output_symbol.center_vector
             main_assignments.extend([Assignment(lhs, sum(d[i] * pdf for d, pdf in zip(self._stencil, pdfs)))
                                      for i, lhs in enumerate(moment1_output_symbol)])
+        if 'moment2' in output_quantity_names_to_symbols:
+            moment2_output_symbol = output_quantity_names_to_symbols['moment2']
+            if isinstance(moment2_output_symbol, Field):
+                moment2_output_symbol = moment2_output_symbol.center_vector
+            for i, p in enumerate(moment2_output_symbol):
+                main_assignments.append(Assignment(p, eqs[self._symbolsOrder2[i]]))
+                del eqs[self._symbolsOrder2[i]]
 
         ac = ac.copy(main_assignments, [Assignment(a, b) for a, b in eqs.items()])
         return ac.new_without_unused_subexpressions()
@@ -233,8 +244,8 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
 # -----------------------------------------  Helper functions ----------------------------------------------------------
 
 
-def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs,
-                                                    symbolic_zeroth_moment, symbolic_first_moments):
+def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs, symbolic_zeroth_moment,
+                                                    symbolic_first_moments, symbolic_second_moments=None):
     r"""
     Returns an equation system that computes the zeroth and first order moments with the least amount of operations
 
@@ -244,12 +255,14 @@ def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs,
 
         \rho = \sum_{d \in S} f_d
         u_j = \sum_{d \in S} f_d u_jd
+        p_j = \sum_{d \in S} {d \in S} f_d u_jd
 
     Args:
         stencil: called :math:`S` above
         symbolic_pdfs: called :math:`f` above
         symbolic_zeroth_moment:  called :math:`\rho` above
         symbolic_first_moments: called :math:`u` above
+        symbolic_second_moments: called :math:`p` above
     """
     def filter_out_plus_terms(expr):
         result = 0
@@ -267,6 +280,12 @@ def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs,
         for i in range(dim):
             u[i] += f * int(offset[i])
 
+    p = [0] * dim * dim
+    for f, offset in zip(symbolic_pdfs, stencil):
+        for i in range(dim):
+            for j in range(dim):
+                p[dim * i + j] += f * int(offset[i]) * int(offset[j])
+
     plus_terms = [set(filter_out_plus_terms(u_i).args) for u_i in u]
     for i in range(dim):
         rhs = plus_terms[i]
@@ -284,6 +303,8 @@ def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs,
     equations = []
     equations += [Assignment(symbolic_zeroth_moment, pdf_sum)]
     equations += [Assignment(u_i_sym, u_i) for u_i_sym, u_i in zip(symbolic_first_moments, u)]
+    if symbolic_second_moments:
+        equations += [Assignment(symbolic_second_moments[i], p[i]) for i in range(dim**2)]
 
     return AssignmentCollection(equations, subexpressions)
 

lbmpy_tests/test_shear_flow.py

+"""
+Test shear flow velocity and pressureagainst analytical solutions
+"""
+
+
+import numpy as np
+import pytest
+import sympy as sp
+
+import lbmpy
+from lbmpy.boundaries import UBB
+from lbmpy.boundaries.boundaryhandling import LatticeBoltzmannBoundaryHandling
+from lbmpy.creationfunctions import create_lb_update_rule, create_stream_pull_with_output_kernel
+from lbmpy.macroscopic_value_kernels import macroscopic_values_setter
+from lbmpy.stencils import get_stencil
+
+import pystencils as ps
+
+
+def shear_flow(x, t, nu, v, h, k_max):
+    """
+    Analytical solution for driven shear flow between two plates.
+
+    Parameters
+    ----------
+    x : :obj:`float`
+        Position from the left plane.
+    t : :obj:`float`
+        Time since start of the shearing.
+    nu : :obj:`float`
+        Kinematic viscosity.
+    v : :obj:`float`
+        Shear rate.
+    h : :obj:`float`
+        Distance between the plates.
+    k_max : :obj:`int`
+        Maximum considered wave number.
+
+    Returns
+    -------
+    :obj:`double` : Analytical velocity
+
+    """
+
+    u = x / h - 0.5
+    for k in np.arange(1, k_max + 1):
+        u += 1.0 / (np.pi * k) * np.exp(
+            -4 * np.pi ** 2 * nu * k ** 2 / h ** 2 * t) * np.sin(2 * np.pi / h * k * x)
+    return -v * u
+
+
+rho_0 = 2.2  # density
+eta = 1.6  # kinematic viscosity
+width = 40  # of box
+wall_thickness = 2
+actual_width = width - wall_thickness  # subtract boundary layer from box width
+shear_velocity = 0.2  # scale by width to keep stable
+t_max = 2000
+
+
+@pytest.mark.parametrize('target', ('cpu', 'gpu', 'opencl'))
+@pytest.mark.parametrize('stencil_name', ("D2Q9", "D3Q19",))
+def test_shear_flow(target, stencil_name):
+    # OpenCL and Cuda
+    if target == 'opencl':
+        import pytest
+        pytest.importorskip("pyopencl")
+        import pystencils.opencl.autoinit
+    elif target == 'gpu':
+        import pytest
+        pytest.importorskip("pycuda")
+
+    # LB parameters
+    lb_stencil = get_stencil(stencil_name)
+    dim = len(lb_stencil[0])
+
+    if dim == 2:
+        L = [4, width]
+    elif dim == 3:
+        L = [4, width, 4]
+    else:
+        raise Exception()
+    periodicity = [True, False] + [True] * (dim - 2)
+
+    omega = lbmpy.relaxationrates.relaxation_rate_from_lattice_viscosity(eta)
+
+    # ## Data structures
+    dh = ps.create_data_handling(L, periodicity=periodicity, default_target=target)
+
+    src = dh.add_array('src', values_per_cell=len(lb_stencil))
+    dst = dh.add_array_like('dst', 'src')
+    ρ = dh.add_array('rho', latex_name='\\rho')
+    u = dh.add_array('u', values_per_cell=dh.dim)
+    p = dh.add_array('p', values_per_cell=dh.dim**2)
+
+    # LB Setup
+    collision = create_lb_update_rule(stencil=lb_stencil,
+                                      relaxation_rate=omega,
+                                      method="trt",
+                                      compressible=True,
+                                      kernel_type='collide_only',
+                                      optimization={'symbolic_field': src})
+
+    stream = create_stream_pull_with_output_kernel(collision.method, src, dst, {'velocity': u})
+
+    opts = {'cpu_openmp': False,
+            'cpu_vectorize_info': None,
+            'target': dh.default_target}
+
+    stream_kernel = ps.create_kernel(stream, **opts).compile()
+    collision_kernel = ps.create_kernel(collision, **opts).compile()
+
+    # Boundaries
+    lbbh = LatticeBoltzmannBoundaryHandling(collision.method, dh, src.name, target=dh.default_target)
+
+    # Second moment test setup
+    cqc = collision.method.conserved_quantity_computation
+    getter_eqs = cqc.output_equations_from_pdfs(src.center_vector,
+                                                {'moment2': p})
+
+    kernel_compute_p = ps.create_kernel(getter_eqs, **opts).compile()
+
+    # ## Set up the simulation
+
+    init = macroscopic_values_setter(collision.method, velocity=(0,) * dh.dim,
+                                     pdfs=src.center_vector, density=ρ.center)
+    init_kernel = ps.create_kernel(init, ghost_layers=0).compile()
+
+    vel_vec = sp.Matrix([0.5 * shear_velocity] + [0] * (dim - 1))
+    if dim == 2:
+        lbbh.set_boundary(UBB(velocity=vel_vec), ps.make_slice[:, :wall_thickness])
+        lbbh.set_boundary(UBB(velocity=-vel_vec), ps.make_slice[:, -wall_thickness:])
+    elif dim == 3:
+        lbbh.set_boundary(UBB(velocity=vel_vec), ps.make_slice[:, :wall_thickness, :])
+        lbbh.set_boundary(UBB(velocity=-vel_vec), ps.make_slice[:, -wall_thickness:, :])
+    else:
+        raise Exception()
+
+    for bh in lbbh, :
+        assert len(bh._boundary_object_to_boundary_info) == 2, "Restart kernel to clear boundaries"
+
+    def init():
+        dh.fill(ρ.name, rho_0)
+        dh.fill(u.name, np.nan, ghost_layers=True, inner_ghost_layers=True)
+        dh.fill(u.name, 0)
+
+        dh.run_kernel(init_kernel)
+
+    sync_pdfs = dh.synchronization_function([src.name])
+
+    # Time loop
+    def time_loop(steps):
+        dh.all_to_gpu()
+        for i in range(steps):
+            dh.run_kernel(collision_kernel)
+            sync_pdfs()
+            lbbh()
+            dh.run_kernel(stream_kernel)
+            dh.run_kernel(kernel_compute_p)
+
+            dh.swap(src.name, dst.name)
+
+        if u.name in dh.gpu_arrays:
+            dh.to_cpu(u.name)
+        uu = dh.gather_array(u.name)
+        # average periodic directions
+        if dim == 3:  # dont' swap order
+            uu = np.average(uu, axis=2)
+        uu = np.average(uu, axis=0)
+
+        if p.name in dh.gpu_arrays:
+            dh.to_cpu(p.name)
+        pp = dh.gather_array(p.name)
+        # average periodic directions
+        if dim == 3:  # dont' swap order
+            pp = np.average(pp, axis=2)
+        pp = np.average(pp, axis=0)
+
+        # cut off wall regions
+        uu = uu[wall_thickness:-wall_thickness]
+        pp = pp[wall_thickness:-wall_thickness]
+
+        if(dh.dim == 2):
+            pp = pp.reshape((len(pp), 2, 2))
+        if(dh.dim == 3):
+            pp = pp.reshape((len(pp), 3, 3))
+        return uu, pp
+
+    init()
+    # Simulation
+    profile, pressure_profile = time_loop(t_max)
+
+    expected = shear_flow(x=(np.arange(0, actual_width) + .5),
+                          t=t_max,
+                          nu=eta / rho_0,
+                          v=shear_velocity,
+                          h=actual_width,
+                          k_max=100)
+
+    if(dh.dim == 2):
+        shear_direction = np.array((1, 0), dtype=float)
+        shear_plane_normal = np.array((0, 1), dtype=float)
+    if(dh.dim == 3):
+        shear_direction = np.array((1, 0, 0), dtype=float)
+        shear_plane_normal = np.array((0, 1, 0), dtype=float)
+
+    shear_rate = shear_velocity / actual_width
+    dynamic_viscosity = eta * rho_0
+    correction_factor = eta / (eta - 1. / 6.)
+
+    p_expected = rho_0 * np.identity(dh.dim) / 3.0 + dynamic_viscosity * shear_rate / correction_factor * (
+        np.outer(shear_plane_normal, shear_direction) + np.transpose(np.outer(shear_plane_normal, shear_direction)))
+
+    # Sustract the tensorproduct of the velosity to get the pressure
+    pressure_profile[:, 0, 0] -= rho_0 * profile[:, 0]**2
+    
+    np.testing.assert_allclose(profile[:, 0], expected[1:-1], atol=1E-9)
+    for i in range(actual_width - 2):
+        np.testing.assert_allclose(pressure_profile[i], p_expected, atol=1E-9, rtol=1E-3)