Finished 2D assignment collection

lbmweights/lattice.py

@@ -225,48 +225,63 @@ class Lattice:
             self._eq_order = 2
         return self._eq_order
 
+    @property
+    def necessary_ghost_layers(self):
+        velocities = self.velocities
+        return np.array(velocities).max()
+
     @property
     @ps_installed
     def assignment_collection(self, ps):
         if self._dimension == 3:
             return self._assignment_collection3d()
-        c_s_sq, weights, velocities = self.velocity_set()
+        symbolic_description = []
+        c_s_sq = self.c_s_sq
+        weights = self.weights
+        velocities = self.velocities
         Q = self.q
-        c_ix = velocities[0]
-        c_iy = velocities[1]
+        weights = [complex(x).real for x in weights]
+        c_ix = [vel[0] for vel in velocities]
+        c_iy = [vel[1] for vel in velocities]
+
         src, dst = ps.fields("src({Q}), 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
+        for i in range(Q):
+            symbolic_description.append(
+                ps.Assignment(dst[0, 0](i), src[c_ix[i], c_iy[i]](i), )
+            )
 
-        symbolic_description = [ps.Assignment(density, density_formula),
-                                ps.Assignment(u_x, vel_x_formula),
-                                ps.Assignment(u_y, vel_y_formula), ]
+        density_formula = sum([dst[0, 0](i) for i in range(Q)])
+        vel_x_formula = (sum([dst[0, 0](i) * c_ix[i] for i in range(Q)])) / density
+        vel_y_formula = (sum([dst[0, 0](i) * c_iy[i] for i in range(Q)])) / density
+
+        symbolic_description.append(ps.Assignment(density, density_formula))
+        symbolic_description.append(ps.Assignment(u_x, vel_x_formula))
+        symbolic_description.append(ps.Assignment(u_y, vel_y_formula))
 
         for i in range(Q):
-            feq_formula = weights[i] * density[0, 0] * (
+            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 self.eq_order >= 3:
-                feq_formula += weights[i] * density[0, 0] * (
-                    -1 / (2 * c_s_sq**2) * (u_x * c_ix[i] + u_y * c_iy[i]) * (u_x ** 2 + u_y ** 2)
-                    + 4 / 3 * (1 / (2 * c_s_sq)) ** 3 * (u_x * c_ix[i] + u_y * c_iy[i]) ** 3
+                feq_formula += weights[i] * density * (
+                        -1 / (2 * c_s_sq ** 2) * (u_x * c_ix[i] + u_y * c_iy[i]) * (u_x ** 2 + u_y ** 2)
+                        + 4 / 3 * (1 / (2 * c_s_sq)) ** 3 * (u_x * c_ix[i] + u_y * c_iy[i]) ** 3
                 )
             if self.eq_order >= 4:
-                feq_formula += weights[i] * density[0, 0] * \
+                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)))
+                                       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)))
 
             symbolic_description.append(
-                ps.Assignment(dst[-1 * c_iy[i], c_ix[i]](i), omega * feq_formula + (1 - omega) * src[0, 0](i))
+                ps.Assignment(dst[0, 0](i), omega * feq_formula + (1 - omega) * dst[0, 0](i))
             )
 
         return ps.AssignmentCollection(symbolic_description)