diff --git a/src/lbmpy/__init__.py b/src/lbmpy/__init__.py
index 63cc20bfef8619d909ecce9e55c20141a4e1925a..36454d29781e556c49d8e06581be060c8fa49fa4 100644
--- a/src/lbmpy/__init__.py
+++ b/src/lbmpy/__init__.py
@@ -12,6 +12,7 @@ from .lbstep import LatticeBoltzmannStep
from .macroscopic_value_kernels import (
pdf_initialization_assignments,
macroscopic_values_getter,
+ strain_rate_tensor_getter,
compile_macroscopic_values_getter,
compile_macroscopic_values_setter,
create_advanced_velocity_setter_collision_rule,
@@ -42,6 +43,7 @@ __all__ = [
"LatticeBoltzmannStep",
"pdf_initialization_assignments",
"macroscopic_values_getter",
+ "strain_rate_tensor_getter",
"compile_macroscopic_values_getter",
"compile_macroscopic_values_setter",
"create_advanced_velocity_setter_collision_rule",
diff --git a/src/lbmpy/macroscopic_value_kernels.py b/src/lbmpy/macroscopic_value_kernels.py
index bd12a4ea97dda40234b4b7d2cdbcf6a52f54a825..ae99bbb36d4e8ed5eb55bfd1685f0d73b79b1bc8 100644
--- a/src/lbmpy/macroscopic_value_kernels.py
+++ b/src/lbmpy/macroscopic_value_kernels.py
@@ -1,32 +1,39 @@
import functools
+import sympy as sp
from copy import deepcopy
from lbmpy.simplificationfactory import create_simplification_strategy
-from pystencils import create_kernel, CreateKernelConfig
+from pystencils import create_kernel, CreateKernelConfig, Assignment
from pystencils.field import Field, get_layout_of_array
from pystencils.enums import Target
from lbmpy.advanced_streaming.utility import get_accessor, Timestep
+from lbmpy.relaxationrates import get_shear_relaxation_rate
+from lbmpy.utils import second_order_moment_tensor
-def pdf_initialization_assignments(lb_method, density, velocity, pdfs,
- streaming_pattern='pull', previous_timestep=Timestep.BOTH,
- set_pre_collision_pdfs=False):
- """Assignments to initialize the pdf field with equilibrium"""
+def get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, pre_collision_pdfs):
if isinstance(pdfs, Field):
accessor = get_accessor(streaming_pattern, previous_timestep)
- if set_pre_collision_pdfs:
+ if pre_collision_pdfs:
field_accesses = accessor.read(pdfs, lb_method.stencil)
else:
field_accesses = accessor.write(pdfs, lb_method.stencil)
- elif streaming_pattern == 'pull' and not set_pre_collision_pdfs:
+ elif streaming_pattern == 'pull' and not pre_collision_pdfs:
field_accesses = pdfs
else:
raise ValueError("Invalid value of pdfs: A PDF field reference is required to derive "
+ f"initialization assignments for streaming pattern {streaming_pattern}.")
+ return field_accesses
+
+
+def pdf_initialization_assignments(lb_method, density, velocity, pdfs,
+ streaming_pattern='pull', previous_timestep=Timestep.BOTH,
+ set_pre_collision_pdfs=False):
+ """Assignments to initialize the pdf field with equilibrium"""
+ field_accesses = get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, set_pre_collision_pdfs)
if isinstance(density, Field):
density = density.center
-
if isinstance(velocity, Field):
velocity = velocity.center_vector
@@ -41,17 +48,9 @@ def pdf_initialization_assignments(lb_method, density, velocity, pdfs,
def macroscopic_values_getter(lb_method, density, velocity, pdfs,
streaming_pattern='pull', previous_timestep=Timestep.BOTH,
use_pre_collision_pdfs=False):
- if isinstance(pdfs, Field):
- accessor = get_accessor(streaming_pattern, previous_timestep)
- if use_pre_collision_pdfs:
- field_accesses = accessor.read(pdfs, lb_method.stencil)
- else:
- field_accesses = accessor.write(pdfs, lb_method.stencil)
- elif streaming_pattern == 'pull' and not use_pre_collision_pdfs:
- field_accesses = pdfs
- else:
- raise ValueError("Invalid value of pdfs: A PDF field reference is required to derive "
- + f"getter assignments for streaming pattern {streaming_pattern}.")
+
+ field_accesses = get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, use_pre_collision_pdfs)
+
cqc = lb_method.conserved_quantity_computation
assert not (velocity is None and density is None)
output_spec = {}
@@ -65,6 +64,27 @@ def macroscopic_values_getter(lb_method, density, velocity, pdfs,
macroscopic_values_setter = pdf_initialization_assignments
+def strain_rate_tensor_getter(lb_method, strain_rate_tensor, pdfs, streaming_pattern='pull',
+ previous_timestep=Timestep.BOTH, use_pre_collision_pdfs=False):
+
+ field_accesses = get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, use_pre_collision_pdfs)
+
+ if isinstance(strain_rate_tensor, Field):
+ strain_rate_tensor = strain_rate_tensor.center_vector
+
+ omega_s = get_shear_relaxation_rate(lb_method)
+ equilibrium = lb_method.equilibrium_distribution
+ rho = equilibrium.density if equilibrium.compressible else equilibrium.background_density
+
+ f_neq = sp.Matrix([field_accesses[i] for i in range(lb_method.stencil.Q)]) - lb_method.get_equilibrium_terms()
+ pi = second_order_moment_tensor(f_neq, lb_method.stencil)
+ strain_rate_tensor_equ = - 1.5 * (omega_s / rho) * pi
+
+ assignments = [Assignment(strain_rate_tensor[i * lb_method.stencil.D + j], strain_rate_tensor_equ[i, j])
+ for i in range(lb_method.stencil.D) for j in range(lb_method.stencil.D)]
+ return assignments
+
+
def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None,
ghost_layers=1, iteration_slice=None,
field_layout='numpy', target=Target.CPU,
diff --git a/tests/test_macroscopic_value_kernels.py b/tests/test_macroscopic_value_kernels.py
index 97fc0bd1bcb28153cfd11a6abbdefcf232eab941..eb5b8d330c70ada432a38c696508bdbf6eba173e 100644
--- a/tests/test_macroscopic_value_kernels.py
+++ b/tests/test_macroscopic_value_kernels.py
@@ -1,10 +1,11 @@
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