3 files changed, 46 insertions, 82 deletions
diff --git a/symbolic/D2Q9.py b/symbolic/D2Q9.py
index 8d42245..22f7ed5 100644
--- a/symbolic/D2Q9.py
+++ b/symbolic/D2Q9.py
@@ -1,5 +1,4 @@
 from sympy import *
-from sympy.codegen.ast import Assignment
 
 q = 9
 d = 2
@@ -8,43 +7,3 @@ c = [ Matrix(x) for x in [(-1, 1), ( 0, 1), ( 1, 1), (-1, 0), ( 0, 0), ( 1, 0),
 w = [ Rational(*x) for x in [(1,36), (1,9), (1,36), (1,9), (4,9), (1,9), (1,36), (1,9), (1,36)] ]
 
 c_s = sqrt(Rational(1,3))
-
-f_next = symarray('f_next', q)
-f_curr = symarray('f_curr', q)
-
-def moments(f = f_curr, optimize = True):
-    rho = symbols('rho')
-    u   = Matrix(symarray('u', d))
-
-    exprs = [ Assignment(rho, sum(f)) ]
-
-    for i, u_i in enumerate(u):
-        exprs.append(Assignment(u_i, sum([ (c_j*f[j])[i] for j, c_j in enumerate(c) ]) / sum(f)))
-
-    if optimize:
-        return cse(exprs, optimizations='basic', symbols=numbered_symbols(prefix='m'))
-    else:
-        return ([], exprs)
-
-def equilibrium():
-    rho = symbols('rho')
-    u   = Matrix(symarray('u', d))
-
-    f_eq = []
-
-    for i, c_i in enumerate(c):
-        f_eq_i = w[i] * rho * (  1
-                               + c_i.dot(u)    /    c_s**2
-                               + c_i.dot(u)**2 / (2*c_s**4)
-                               - u.dot(u)      / (2*c_s**2) )
-        f_eq.append(f_eq_i)
-
-    return f_eq
-
-def bgk(tau, f_eq = equilibrium(), optimize = True):
-    exprs = [ Assignment(f_next[i], f_curr[i] + 1/tau * ( f_eq_i - f_curr[i] )) for i, f_eq_i in enumerate(f_eq) ]
-
-    if optimize:
-        return cse(exprs, optimizations='basic')
-    else:
-        return ([], exprs)
diff --git a/symbolic/D3Q19.py b/symbolic/D3Q19.py
index 789b083..4e84908 100644
--- a/symbolic/D3Q19.py
+++ b/symbolic/D3Q19.py
@@ -1,5 +1,4 @@
 from sympy import *
-from sympy.codegen.ast import Assignment
 
 q = 19
 d = 3
@@ -17,43 +16,3 @@ w = [Rational(*x) for x in [
 ]]
 
 c_s = sqrt(Rational(1,3))
-
-f_next = symarray('f_next', q)
-f_curr = symarray('f_curr', q)
-
-def moments(f = f_curr, optimize = True):
-    rho = symbols('rho')
-    u   = Matrix(symarray('u', d))
-
-    exprs = [ Assignment(rho, sum(f)) ]
-
-    for i, u_i in enumerate(u):
-        exprs.append(Assignment(u_i, sum([ (c_j*f[j])[i] for j, c_j in enumerate(c) ]) / sum(f)))
-
-    if optimize:
-        return cse(exprs, optimizations='basic', symbols=numbered_symbols(prefix='m'))
-    else:
-        return ([], exprs)
-
-def equilibrium():
-    rho = symbols('rho')
-    u   = Matrix(symarray('u', d))
-
-    f_eq = []
-
-    for i, c_i in enumerate(c):
-        f_eq_i = w[i] * rho * (  1
-                               + c_i.dot(u)    /    c_s**2
-                               + c_i.dot(u)**2 / (2*c_s**4)
-                               - u.dot(u)      / (2*c_s**2) )
-        f_eq.append(f_eq_i)
-
-    return f_eq
-
-def bgk(tau, f_eq = equilibrium(), optimize = True):
-    exprs = [ Assignment(f_next[i], f_curr[i] + 1/tau * ( f_eq_i - f_curr[i] )) for i, f_eq_i in enumerate(f_eq) ]
-
-    if optimize:
-        return cse(exprs, optimizations='basic')
-    else:
-        return ([], exprs)
diff --git a/symbolic/generator.py b/symbolic/generator.py
new file mode 100644
index 0000000..e57154c
--- /dev/null
+++ b/symbolic/generator.py
@@ -0,0 +1,46 @@
+from sympy import *
+from sympy.codegen.ast import Assignment
+
+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)
+
+    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='basic', 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 = [ Assignment(self.f_next[i], self.f_curr[i] + 1/tau * (f_eq_i - self.f_curr[i])) for i, f_eq_i in enumerate(f_eq) ]
+
+        if optimize:
+            return cse(exprs, optimizations='basic')
+        else:
+            return ([], exprs)