Split descriptors and symbolic formulation

7 files changed, 64 insertions, 91 deletions

diff --git a/implosion.py b/implosion.py
index 05f3644..2ba831f 100644
--- a/implosion.py
+++ b/implosion.py
@@ -5,7 +5,8 @@ import matplotlib
 import matplotlib.pyplot as plt
 matplotlib.use('AGG')
 
-from lbm import Lattice, Geometry
+from simulation         import Lattice, Geometry
+from symbolic.generator import LBM
 
 import symbolic.D2Q9 as D2Q9
 
@@ -58,12 +59,14 @@ moments = []
 
 print("Initializing simulation...\n")
 
+lbm = LBM(D2Q9)
+
 lattice = Lattice(
     descriptor = D2Q9,
     geometry   = Geometry(1024, 1024),
 
-    moments = D2Q9.moments(optimize = False),
-    collide = D2Q9.bgk(tau = 0.8),
+    moments = lbm.moments(optimize = False),
+    collide = lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.8),
 
     pop_eq_src   = pop_eq,
     boundary_src = boundary)
diff --git a/ldc_2d.py b/ldc_2d.py
index 0aa1e71..979f7be 100644
--- a/ldc_2d.py
+++ b/ldc_2d.py
@@ -5,7 +5,8 @@ import matplotlib
 import matplotlib.pyplot as plt
 matplotlib.use('AGG')
 
-from lbm import Lattice, Geometry
+from simulation         import Lattice, Geometry
+from symbolic.generator import LBM
 
 import symbolic.D2Q9 as D2Q9
 
@@ -50,12 +51,14 @@ moments = []
 
 print("Initializing simulation...\n")
 
+lbm = LBM(D2Q9)
+
 lattice = Lattice(
     descriptor = D2Q9,
     geometry   = Geometry(256, 256),
 
-    moments = D2Q9.moments(optimize = False),
-    collide = D2Q9.bgk(tau = 0.52),
+    moments = lbm.moments(optimize = False),
+    collide = lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.52),
 
     boundary_src = boundary)
 
diff --git a/ldc_3d.py b/ldc_3d.py
index 62f820e..c008b06 100644
--- a/ldc_3d.py
+++ b/ldc_3d.py
@@ -5,7 +5,8 @@ import matplotlib
 import matplotlib.pyplot as plt
 matplotlib.use('AGG')
 
-from lbm import Lattice, Geometry
+from simulation         import Lattice, Geometry
+from symbolic.generator import LBM
 
 import symbolic.D3Q19 as D3Q19
 
@@ -68,12 +69,14 @@ moments = []
 
 print("Initializing simulation...\n")
 
+lbm = LBM(D3Q19)
+
 lattice = Lattice(
     descriptor = D3Q19,
     geometry   = Geometry(128, 128, 128),
 
-    moments = D3Q19.moments(optimize = False),
-    collide = D3Q19.bgk(tau = 0.52),
+    moments = lbm.moments(optimize = False),
+    collide = lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.52),
 
     boundary_src = boundary)
 
diff --git a/lbm.py b/simulation.py
index ef05e4f..ef05e4f 100644
--- a/lbm.py
+++ b/simulation.py
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)