From fc2bd74ed89f12d539aef61a2119292c0a8849d9 Mon Sep 17 00:00:00 2001
From: Philipp Suffa <philipp.suffa@fau.de>
Date: Fri, 27 Sep 2024 12:58:41 +0200
Subject: [PATCH] Implements a getter function for the strain rate tensor

---
 src/lbmpy/__init__.py                  |  2 ++
 src/lbmpy/macroscopic_value_kernels.py | 38 +++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/lbmpy/__init__.py b/src/lbmpy/__init__.py
index 63cc20bf..36454d29 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 bd12a4ea..0f86da9a 100644
--- a/src/lbmpy/macroscopic_value_kernels.py
+++ b/src/lbmpy/macroscopic_value_kernels.py
@@ -1,11 +1,14 @@
 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,
@@ -65,6 +68,39 @@ 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):
+
+    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}.")
+
+    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,
-- 
GitLab