2D Assignment Collection

cli.py

@@ -30,12 +30,14 @@ def benchmark(x=10):
 )
 @click.option("--seed", type=int, help="Random number to be used as seed")
 def lbmweights(dimension, order, shells, seed):
-    #shell_list = [int(x) for x in re.findall(u"\d+", shells)]
-    #lattice = Lattice(
-    #    dimension=dimension, order=order, shell_list=shell_list, seed=seed
-    #)
+    """Usage:
+    Test functionality of this module from the command line.
 
+    lbmweights --dimension=2 --order=4 --shells="1,2,4"
+    """
     lattice = Lattice.from_name("D2Q9")
+    lattice.weights
+    lattice.assignment_collection
 
 
 

lbmweights/lattice.py

@@ -18,6 +18,7 @@ from .utils.utils import (
     lattice_sum,
     approximate_ratio,
     timer,
+    ps_installed,
     get_random_vector,
     velocities_for_shell,
 )
@@ -182,6 +183,61 @@ class Lattice:
             return self._c_s_sq
         return self.velocity_set()[0]
 
+    @property
+    @ps_installed
+    def assignment_collection(self):
+        print("okay")
+        return
+
+        if self._dimension == 3:
+            return self._assignment_collection3d()
+        Q = len(self.weights)
+        weights = self.weights
+        c_s_sq = float(self.c_s_sq)
+        c_ix = self.velocities[0]
+        c_iy = self.velocities[1]
+        src, dst = ps.fields("src, dst({Q}): double[2D]".format(Q=Q))
+        density, u_x, u_y, omega = sp.symbols("density, u_x, u_y, omega")
+
+        density_formula = sum([src[0, 0](i) for i in range(Q)])
+        vel_x_formula = (sum([src[0, 0](i) * c_ix[i] for i in range(Q)])) / density
+        vel_y_formula = (sum([src[0, 0](i) * c_iy[i] for i in range(Q)])) / density
+
+        symbolic_description = [ps.Assignment(density, density_formula),
+                                ps.Assignment(u_x, vel_x_formula),
+                                ps.Assignment(u_y, vel_y_formula), ]
+
+        for i in range(Q):
+            feq_formula = weights[i] * density * (
+                    1
+                    + (u_x * c_ix[i] + u_y * c_iy[i]) / c_s_sq
+                    + (u_x * c_ix[i] + u_y * c_iy[i]) ** 2 / (2 * c_s_sq ** 2)
+                    - (u_x ** 2 + u_y ** 2) / (2 * c_s_sq)
+            )
+            if len(c_ix) >= 17:
+                feq_formula += weights[i] * density * (
+                        -2 * (1 / (2 * c_s_sq)) ** 2 * (
+                            c_ix[i] * u_x ** 3 + c_ix[i] * u_x * u_y ** 2 + c_iy[i] * u_y * u_x ** 2 + c_iy[
+                        i] * u_y ** 3)
+                        + 4 / 3 * (1 / (2 * c_s_sq)) ** 3 * (u_x * c_ix[i] + u_y * c_iy[i]) ** 3
+                )
+            if len(c_ix) >= 37:
+                feq_formula += weights[i] * density * \
+                               (1 / (6 * c_s_sq ** 2) * (
+                                       (u_x * c_ix[i] + u_y * c_iy[i]) ** 4 / (4 * c_s_sq ** 2) - (
+                                           3 * (u_x ** 2 + u_y ** 2) *
+                                           (u_x * c_ix[i] + u_y * c_iy[i]) ** 2) / (2 * c_s_sq) + 3 * (
+                                                   u_x ** 4 + u_y ** 4)))
+
+        for i in range(Q):
+            symbolic_description.append(
+                ps.Assignment(dst[-1 * c_iy[i], c_ix[i]](i), omega * feq_formula + (1 - omega) * src[0, 0](i))
+            )
+        return ps.AssignmentCollection(symbolic_description)
+
+    def _assignment_collection3d(self):
+        pass
+
     def shell_from_type(self, type):
         for shell in self._shells:
             if shell.type == type:

lbmweights/utils/utils.py

@@ -53,6 +53,22 @@ def timer(func, *args, **kwargs):
     return wrapper
 
 
+def ps_installed(func, *args, **kwargs):
+    """
+    Wrapper that imports Pystencils or warns the user if it can't be installed.
+    """
+
+    def wrapper(*args, **kwargs):
+        try:
+            ps = __import__("pystencils")
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError("Pystencils not found. Please install it to get AssignmentCollections")
+        ret = func(*args, **kwargs)
+        return ret
+
+    return wrapper
+
+
 def analyse_tensor_dimension(tensor_rank):
     """