7 files changed, 77 insertions, 83 deletions
diff --git a/boltzgen/lbm/model/BGK.py b/boltzgen/lbm/model/BGK.py
new file mode 100644
index 0000000..f5a34e2
--- /dev/null
+++ b/boltzgen/lbm/model/BGK.py
@@ -0,0 +1,76 @@
+from sympy import *
+from sympy.codegen.ast import Assignment
+
+from boltzgen.utility import assign
+import boltzgen.utility.optimizations as optimizations
+
+from boltzgen.lbm.lattice.characteristics import weights, c_s
+
+class BGK:
+    def __init__(self, descriptor, tau, optimize = True):
+        self.descriptor = descriptor
+        self.tau        = tau
+        self.optimize   = optimize
+
+        if self.tau <= 0.5:
+            raise Exception('Relaxation time must be larger than 0.5')
+
+        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 = None):
+        if optimize is None:
+            optimize = self.optimize
+
+        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, resolve_moments = False):
+        rho = symbols('rho')
+        u   = Matrix(symarray('u', self.descriptor.d))
+
+        if resolve_moments:
+            moments = self.moments(optimize = False)[1]
+            rho = moments[0].rhs
+            for i, m in enumerate(moments[1:]):
+                u[i] = m.rhs
+
+        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 collision(self, f_eq, optimize = None):
+        if optimize is None:
+            optimize = self.optimize
+
+        exprs = [ self.f_curr[i] + 1/self.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
deleted file mode 100644
index e3ac9de..0000000
--- a/boltzgen/lbm/model/D2Q9.py
+++ /dev/null
@@ -1,7 +0,0 @@
-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
deleted file mode 100644
index e9e6eec..0000000
--- a/boltzgen/lbm/model/D3Q19.py
+++ /dev/null
@@ -1,18 +0,0 @@
-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
deleted file mode 100644
index 6fb1f80..0000000
--- a/boltzgen/lbm/model/D3Q27.py
+++ /dev/null
@@ -1,7 +0,0 @@
-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
deleted file mode 100644
index 04e16a3..0000000
--- a/boltzgen/lbm/model/D3Q7.py
+++ /dev/null
@@ -1,18 +0,0 @@
-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/__init__.py b/boltzgen/lbm/model/__init__.py
index c956759..fb8f12b 100644
--- a/boltzgen/lbm/model/__init__.py
+++ b/boltzgen/lbm/model/__init__.py
@@ -1,6 +1 @@
-from . import D2Q9
-from . import D3Q7
-from . import D3Q19
-from . import D3Q27
-
-from . import characteristics
+from .BGK import BGK
diff --git a/boltzgen/lbm/model/characteristics.py b/boltzgen/lbm/model/characteristics.py
deleted file mode 100644
index b68afeb..0000000
--- a/boltzgen/lbm/model/characteristics.py
+++ /dev/null
@@ -1,27 +0,0 @@
-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()