Field tested assignment collections in 2D

lbmweights/lattice.py

@@ -260,29 +260,30 @@ class Lattice:
         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 * (
-                    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 * (
-                        -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 not self._debug:
+            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 self.eq_order >= 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 * \
+                                   (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)))
+
+                symbolic_description.append(
+                    ps.Assignment(dst[0, 0](i), omega * feq_formula + (1 - omega) * dst[0, 0](i))
                 )
-            if self.eq_order >= 4:
-                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)))
-
-            symbolic_description.append(
-                ps.Assignment(dst[0, 0](i), omega * feq_formula + (1 - omega) * dst[0, 0](i))
-            )
 
         return ps.AssignmentCollection(symbolic_description)
 
@@ -291,9 +292,9 @@ class Lattice:
     def _assignment_collection3d(self, ps):
         c_s_sq, weights, velocities = self.velocity_set()
         Q = self.q
-        c_ix = velocities[0]
-        c_iy = velocities[1]
-        c_iz = velocities[2]
+        c_ix = [vel[0] for vel in velocities]
+        c_iy = [vel[1] for vel in velocities]
+        c_iz = [vel[2] for vel in velocities]
         src, dst = ps.fields("src({Q}), dst({Q}): double[2D]".format(Q=Q))
         density, u_x, u_y, u_z, omega = sp.symbols("density, u_x, u_y, u_z, omega")
 
@@ -586,6 +587,9 @@ class Lattice:
             if c_s_sq.imag != 0:
                 raise ValueError("Imaginary root returned. This should not happen")
             c_s_sq = float(c_s_sq.real)
+        if c_s_sq < 0:
+            raise ValueError("Negative Value for c_s_sq")
+
         self._c_s_sq = c_s_sq
         self._reduced_weights = weights
         return weights