Pull in basics from symlbm_playground

It's time to extract the generator-part of my GPU LBM playground and turn it
into a nice reusable library. The goal is to produce a framework that can be
used to generate collision and streaming programs from symbolic descriptions.

i.e. it should be possible to select a LB model, the desired boundary
conditions as well as a data structure / streaming model and use this
information to automatically generate matching OpenCL / CUDA / C++
programs.

6 files changed, 137 insertions, 0 deletions

diff --git a/boltzgen/lbm/__init__.py b/boltzgen/lbm/__init__.py
new file mode 100644
index 0000000..f80feaa
--- /dev/null
+++ b/boltzgen/lbm/__init__.py
@@ -0,0 +1,60 @@
+from sympy import *
+from sympy.codegen.ast import Assignment
+
+import utility.optimizations as optimizations
+from lbm.model.characteristics import weights, c_s
+
+
+def assign(names, definitions):
+    return list(map(lambda x: Assignment(*x), zip(names, definitions)))
+
+class LBM:
+    def __init__(self, descriptor):
+        self.descriptor = descriptor
+        self.f_next = symarray('f_next', descriptor.q)
+        self.f_curr = symarray('f_curr', descriptor.q)
+
+        if not hasattr(descriptor, 'w'):
+            self.descriptor.w = weights(descriptor.d, descriptor.c)
+
+        if not hasattr(descriptor, 'c_s'):
+            self.descriptor.c_s = c_s(descriptor.d, descriptor.c, self.descriptor.w)
+
+    def moments(self, optimize = True):
+        rho = symbols('rho')
+        u   = Matrix(symarray('u', self.descriptor.d))
+
+        exprs = [ Assignment(rho, sum(self.f_curr)) ]
+
+        for i, u_i in enumerate(u):
+            exprs.append(
+                Assignment(u_i, sum([ (c_j*self.f_curr[j])[i] for j, c_j in enumerate(self.descriptor.c) ]) / sum(self.f_curr)))
+
+        if optimize:
+            return cse(exprs, optimizations=optimizations.custom, symbols=numbered_symbols(prefix='m'))
+        else:
+            return ([], exprs)
+
+    def equilibrium(self):
+        rho = symbols('rho')
+        u   = Matrix(symarray('u', self.descriptor.d))
+
+        f_eq = []
+
+        for i, c_i in enumerate(self.descriptor.c):
+            f_eq_i = self.descriptor.w[i] * rho * ( 1
+                                                  + c_i.dot(u)    /    self.descriptor.c_s**2
+                                                  + c_i.dot(u)**2 / (2*self.descriptor.c_s**4)
+                                                  - u.dot(u)      / (2*self.descriptor.c_s**2) )
+            f_eq.append(f_eq_i)
+
+        return f_eq
+
+    def bgk(self, tau, f_eq, optimize = True):
+        exprs = [ self.f_curr[i] + 1/tau * (f_eq_i - self.f_curr[i]) for i, f_eq_i in enumerate(f_eq) ]
+
+        if optimize:
+            subexprs, f = cse(exprs, optimizations=optimizations.custom)
+            return (subexprs, assign(self.f_next, f))
+        else:
+            return ([], assign(self.f_next, exprs))
diff --git a/boltzgen/lbm/model/D2Q9.py b/boltzgen/lbm/model/D2Q9.py
new file mode 100644
index 0000000..e3ac9de
--- /dev/null
+++ b/boltzgen/lbm/model/D2Q9.py
@@ -0,0 +1,7 @@
+from sympy import Matrix
+from itertools import product
+
+d = 2
+q = 9
+
+c = [ Matrix(x) for x in product([-1,0,1], repeat=d) ]
diff --git a/boltzgen/lbm/model/D3Q19.py b/boltzgen/lbm/model/D3Q19.py
new file mode 100644
index 0000000..e9e6eec
--- /dev/null
+++ b/boltzgen/lbm/model/D3Q19.py
@@ -0,0 +1,18 @@
+from sympy import Matrix, Rational, sqrt
+
+d = 3
+q = 19
+
+c = [ Matrix(x) for x in [
+    ( 0, 1, 1), (-1, 0, 1), ( 0, 0, 1), ( 1, 0, 1), ( 0, -1, 1),
+    (-1, 1, 0), ( 0, 1, 0), ( 1, 1, 0), (-1, 0, 0), ( 0, 0, 0), ( 1, 0, 0), (-1,-1, 0), ( 0, -1, 0), ( 1, -1, 0),
+    ( 0, 1,-1), (-1, 0,-1), ( 0, 0,-1), ( 1, 0,-1), ( 0, -1,-1)
+]]
+
+w = [Rational(*x) for x in [
+    (1,36), (1,36), (1,18), (1,36), (1,36),
+    (1,36), (1,18), (1,36), (1,18), (1,3), (1,18), (1,36), (1,18), (1,36),
+    (1,36), (1,36), (1,18), (1,36), (1,36)
+]]
+
+c_s = sqrt(Rational(1,3))
diff --git a/boltzgen/lbm/model/D3Q27.py b/boltzgen/lbm/model/D3Q27.py
new file mode 100644
index 0000000..6fb1f80
--- /dev/null
+++ b/boltzgen/lbm/model/D3Q27.py
@@ -0,0 +1,7 @@
+from sympy import Matrix
+from itertools import product
+
+d = 3
+q = 27
+
+c = [ Matrix(x) for x in product([-1,0,1], repeat=d) ]
diff --git a/boltzgen/lbm/model/D3Q7.py b/boltzgen/lbm/model/D3Q7.py
new file mode 100644
index 0000000..04e16a3
--- /dev/null
+++ b/boltzgen/lbm/model/D3Q7.py
@@ -0,0 +1,18 @@
+from sympy import *
+
+q = 7
+d = 3
+
+c = [ Matrix(x) for x in [
+    ( 0, 0, 1),
+    ( 0, 1, 0), (-1, 0, 0), ( 0, 0, 0), ( 1, 0, 0), ( 0, -1, 0),
+    ( 0, 0,-1)
+]]
+
+w = [Rational(*x) for x in [
+    (1,8),
+    (1,8), (1,8), (1,4), (1,8), (1,8),
+    (1,8)
+]]
+
+c_s = sqrt(Rational(1,4))
diff --git a/boltzgen/lbm/model/characteristics.py b/boltzgen/lbm/model/characteristics.py
new file mode 100644
index 0000000..b68afeb
--- /dev/null
+++ b/boltzgen/lbm/model/characteristics.py
@@ -0,0 +1,27 @@
+from sympy import *
+
+# copy of `sympy.integrals.quadrature.gauss_hermite` sans evaluation
+def gauss_hermite(n):
+    x = Dummy("x")
+    p  = hermite_poly(n, x, polys=True)
+    p1 = hermite_poly(n-1, x, polys=True)
+    xi = []
+    w  = []
+    for r in p.real_roots():
+        xi.append(r)
+        w.append(((2**(n-1) * factorial(n) * sqrt(pi))/(n**2 * p1.subs(x, r)**2)))
+    return xi, w
+
+# determine weights of a d-dimensional LBM model on velocity set c
+# (only works for velocity sets that result into NSE-recovering LB models when
+#  plugged into Gauss-Hermite quadrature without any additional arguments
+#  i.e. D2Q9 and D3Q27 but not D3Q19)
+def weights(d, c):
+    _, omegas = gauss_hermite(3)
+    return list(map(lambda c_i: Mul(*[ omegas[1+c_i[iDim]] for iDim in range(0,d) ]) / pi**(d/2), c))
+
+# determine lattice speed of sound using directions and their weights
+def c_s(d, c, w):
+    speeds = set([ sqrt(sum([ w[i] * c_i[j]**2 for i, c_i in enumerate(c) ])) for j in range(0,d) ])
+    assert len(speeds) == 1 # verify isotropy
+    return speeds.pop()