Added function get_central_moments to moments.py

lbmpy/moments.py

@@ -417,6 +417,36 @@ def shift_matrix(moments, stencil):
     return sp.Matrix(shift_matrix_list)
 
 
+def get_central_moments(stencil, u=tuple(sp.symbols("u_0 u_1 u_2"))):
+    """
+    Returns the central moments for D2Q9 or D3Q27
+    """
+    x, y, z = MOMENT_SYMBOLS
+    if len(stencil) == 9:
+        central_moments = (sp.sympify(1), x - u[0], y - u[1],
+                           (x - u[0]) ** 2, (y - u[1]) ** 2, (x - u[0]) * (y - u[1]),
+                           (x - u[0]) ** 2 * (y - u[1]), (x - u[0]) * (y - u[1]) ** 2,
+                           (x - u[0]) ** 2 * (y - u[1]) ** 2)
+        return central_moments
+    elif len(stencil) == 27:
+        central_moments = (sp.sympify(1),
+                           x - u[0], y - u[1], z - u[2],
+                           (x - u[0]) ** 2, (y - u[1]) ** 2, (z - u[2]) ** 2, (x - u[0]) * (y - u[1]),
+                           (x - u[0]) * (z - u[2]), (y - u[1]) * (z - u[2]),
+                           (x - u[0]) ** 2 * (y - u[1]), (x - u[0]) ** 2 * (z - u[2]), (x - u[0]) * (y - u[1]) ** 2,
+                           (x - u[0]) * (z - u[2]) ** 2, (y - u[1]) ** 2 * (z - u[2]), (y - u[1]) * (z - u[2]) ** 2,
+                           (x - u[0]) * (y - u[1]) * (z - u[2]),
+                           (x - u[0]) ** 2 * (y - u[1]) ** 2, (x - u[0]) ** 2 * (z - u[2]) ** 2,
+                           (y - u[1]) ** 2 * (z - u[2]) ** 2, (x - u[0]) ** 2 * (y - u[1]) * (z - u[2]),
+                           (x - u[0]) * (y - u[1]) ** 2 * (z - u[2]), (x - u[0]) * (y - u[1]) * (z - u[2]) ** 2,
+                           (x - u[0]) ** 2 * (y - u[1]) ** 2 * (z - u[2]),
+                           (x - u[0]) ** 2 * (y - u[1]) * (z - u[2]) ** 2,
+                           (x - u[0]) * (y - u[1]) ** 2 * (z - u[2]) ** 2,
+                           (x - u[0]) ** 2 * (y - u[1]) ** 2 * (z - u[2]) ** 2)
+        return central_moments
+
+
+
 def gram_schmidt(moments, stencil, weights=None):
     r"""
     Computes orthogonal set of moments using the method by Gram-Schmidt