Restructure LBM model / lattice distinction

9 files changed, 85 insertions, 84 deletions

diff --git a/boltzgen/lbm/__init__.py b/boltzgen/lbm/__init__.py
index 3f608cb..276c0bb 100644
--- a/boltzgen/lbm/__init__.py
+++ b/boltzgen/lbm/__init__.py
@@ -1,78 +1,2 @@
-from sympy import *
-from sympy.codegen.ast import Assignment
-
-import boltzgen.utility.optimizations as optimizations
-from boltzgen.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, 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 bgk(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))
+from . import model
+from . import lattice
diff --git a/boltzgen/lbm/model/D2Q9.py b/boltzgen/lbm/lattice/D2Q9.py
index e3ac9de..e3ac9de 100644
--- a/boltzgen/lbm/model/D2Q9.py
+++ b/boltzgen/lbm/lattice/D2Q9.py
diff --git a/boltzgen/lbm/model/D3Q19.py b/boltzgen/lbm/lattice/D3Q19.py
index e9e6eec..e9e6eec 100644
--- a/boltzgen/lbm/model/D3Q19.py
+++ b/boltzgen/lbm/lattice/D3Q19.py
diff --git a/boltzgen/lbm/model/D3Q27.py b/boltzgen/lbm/lattice/D3Q27.py
index 6fb1f80..6fb1f80 100644
--- a/boltzgen/lbm/model/D3Q27.py
+++ b/boltzgen/lbm/lattice/D3Q27.py
diff --git a/boltzgen/lbm/model/D3Q7.py b/boltzgen/lbm/lattice/D3Q7.py
index 04e16a3..04e16a3 100644
--- a/boltzgen/lbm/model/D3Q7.py
+++ b/boltzgen/lbm/lattice/D3Q7.py
diff --git a/boltzgen/lbm/lattice/__init__.py b/boltzgen/lbm/lattice/__init__.py
new file mode 100644
index 0000000..c956759
--- /dev/null
+++ b/boltzgen/lbm/lattice/__init__.py
@@ -0,0 +1,6 @@
+from . import D2Q9
+from . import D3Q7
+from . import D3Q19
+from . import D3Q27
+
+from . import characteristics
diff --git a/boltzgen/lbm/model/characteristics.py b/boltzgen/lbm/lattice/characteristics.py
index b68afeb..b68afeb 100644
--- a/boltzgen/lbm/model/characteristics.py
+++ b/boltzgen/lbm/lattice/characteristics.py
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/__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