refactor: use iterator for grid

2 files changed, 72 insertions, 76 deletions

diff --git a/src/research.py b/src/research.py
index f0374cf..d660f8e 100644
--- a/src/research.py
+++ b/src/research.py
@@ -37,6 +37,8 @@ class Research:
                     self.plotter.plot(self.model)
                 if preview:
                     self.plotter.show()
+                if save_plot:
+                    self.plotter.save(os.path.join(self.path, f'{iteration}.png'))
 
             if iteration % 50 == 0:
                 if save_model:
diff --git a/src/simple.py b/src/simple.py
index dfd7e0f..0282dd1 100644
--- a/src/simple.py
+++ b/src/simple.py
@@ -14,6 +14,7 @@ class SIMPLE:
 
         self.step = step
         self.shape = tuple(int(x / step) for x in domain_size)
+        self.bfs_node_size = tuple(int(x / step) for x in bfs_size)
 
         # Allocations
         self.u = self.allocate_field()
@@ -45,115 +46,108 @@ class SIMPLE:
         assert value > -0.01, f'WARNING: Value must be positive: {value}'
         return value
 
+    def grid(self):
+        '''Iterator over all grid points, excluding the obstacle'''
+        for i in range(self.shape[0] - 1):
+            for j in range(1, self.shape[1] - 1):
+                if i > self.bfs_node_size[0] or j > self.bfs_node_size[1]:
+                    yield (i, j)
+
     def solve_momentum_equations(self):
         # Momentum along X direction
-        for i in range(1, self.shape[0] - 1):
-            for j in range(1, self.shape[1] - 1):
-                u_E = 0.5 * (self.u[i][j] + self.u[i][j + 1])
-                u_W = 0.5 * (self.u[i][j] + self.u[i][j - 1])
-                v_N = 0.5 * (self.v[i - 1][j] + self.v[i - 1][j + 1])
-                v_S = 0.5 * (self.v[i][j] + self.v[i][j + 1])
+        for i, j in self.grid():
+            u_E = 0.5 * (self.u[i][j] + self.u[i][j + 1])
+            u_W = 0.5 * (self.u[i][j] + self.u[i][j - 1])
+            v_N = 0.5 * (self.v[i - 1][j] + self.v[i - 1][j + 1])
+            v_S = 0.5 * (self.v[i][j] + self.v[i][j + 1])
 
-                a_E = self.assert_positive(-0.5 * u_E * self.step + self.nu)
-                a_W = self.assert_positive(+0.5 * u_W * self.step + self.nu)
-                a_N = self.assert_positive(-0.5 * v_N * self.step + self.nu)
-                a_S = self.assert_positive(+0.5 * v_S * self.step + self.nu)
+            a_E = self.assert_positive(-0.5 * u_E * self.step + self.nu)
+            a_W = self.assert_positive(+0.5 * u_W * self.step + self.nu)
+            a_N = self.assert_positive(-0.5 * v_N * self.step + self.nu)
+            a_S = self.assert_positive(+0.5 * v_S * self.step + self.nu)
 
-                a_e = 0.5 * self.step * (u_E - u_W + v_N - v_S) + 4 * self.nu
-                A_e = -self.step
+            a_e = 0.5 * self.step * (u_E - u_W + v_N - v_S) + 4 * self.nu
+            A_e = -self.step
 
-                self.d_e[i][j] = A_e / a_e
+            self.d_e[i][j] = A_e / a_e
 
-                self.u_star[i][j] = (
-                    a_E * self.u[i][j + 1] +
-                    a_W * self.u[i][j - 1] +
-                    a_N * self.u[i - 1][j] +
-                    a_S * self.u[i + 1][j]
-                ) / a_e + self.d_e[i][j] * (self.p_star[i][j + 1] - self.p_star[i][j])
+            self.u_star[i][j] = (
+                a_E * self.u[i][j + 1] +
+                a_W * self.u[i][j - 1] +
+                a_N * self.u[i - 1][j] +
+                a_S * self.u[i + 1][j]
+            ) / a_e + self.d_e[i][j] * (self.p_star[i][j + 1] - self.p_star[i][j])
 
         # Momentum along Y direction
-        for i in range(1, self.shape[0] - 1):
-            for j in range(1, self.shape[1] - 1):
-                u_E = 0.5 * (self.u[i][j] + self.u[i + 1][j])
-                u_W = 0.5 * (self.u[i][j - 1] + self.u[i + 1][j - 1])
-                v_N = 0.5 * (self.v[i - 1][j] + self.v[i][j])
-                v_S = 0.5 * (self.v[i][j] + self.v[i + 1][j])
+        for i, j in self.grid():
+            u_E = 0.5 * (self.u[i][j] + self.u[i + 1][j])
+            u_W = 0.5 * (self.u[i][j - 1] + self.u[i + 1][j - 1])
+            v_N = 0.5 * (self.v[i - 1][j] + self.v[i][j])
+            v_S = 0.5 * (self.v[i][j] + self.v[i + 1][j])
 
-                a_E = self.assert_positive(-0.5 * u_E * self.step + self.nu)
-                a_W = self.assert_positive(+0.5 * u_W * self.step + self.nu)
-                a_N = self.assert_positive(-0.5 * v_N * self.step + self.nu)
-                a_S = self.assert_positive(+0.5 * v_S * self.step + self.nu)
+            a_E = self.assert_positive(-0.5 * u_E * self.step + self.nu)
+            a_W = self.assert_positive(+0.5 * u_W * self.step + self.nu)
+            a_N = self.assert_positive(-0.5 * v_N * self.step + self.nu)
+            a_S = self.assert_positive(+0.5 * v_S * self.step + self.nu)
 
-                a_n = 0.5 * self.step * (u_E - u_W + v_N - v_S) + 4 * self.nu
-                A_n = -self.step
+            a_n = 0.5 * self.step * (u_E - u_W + v_N - v_S) + 4 * self.nu
+            A_n = -self.step
 
-                self.d_n[i][j] = A_n / a_n
+            self.d_n[i][j] = A_n / a_n
 
-                self.v_star[i][j] = (
-                    a_E * self.v[i][j + 1] +
-                    a_W * self.v[i][j - 1] +
-                    a_N * self.v[i - 1][j] +
-                    a_S * self.v[i + 1][j]
-                ) / a_n + self.d_n[i][j] * (self.p_star[i][j] - self.p_star[i + 1][j])
+            self.v_star[i][j] = (
+                a_E * self.v[i][j + 1] +
+                a_W * self.v[i][j - 1] +
+                a_N * self.v[i - 1][j] +
+                a_S * self.v[i + 1][j]
+            ) / a_n + self.d_n[i][j] * (self.p_star[i][j] - self.p_star[i + 1][j])
 
     def apply_sides_boundary(self):
         for j in range(self.shape[1]):
             self.v_star[0][j] = 0
-            self.u_star[0][j] = 0
-
-            self.v_star[-1][j] = 0
-            self.u_star[-1][j] = 0
+            self.v_star[-2][j] = 0
 
     def apply_bfs_boundary(self):
         '''Apply Backwards Facing Step boundary conditions'''
-        bfs_nodes_width = int(self.bfs_size[0] / self.step)
-        bfs_nodes_height = int(self.bfs_size[1] / self.step)
-        for i in range(bfs_nodes_width + 1):
-            for j in range(bfs_nodes_height + 1):
-                self.u_star[i][j] = 0
-                self.v_star[i][j] = 0
+        for i in range(self.bfs_node_size[0]):
+            self.u_star[i][self.bfs_node_size[1]] = 0
 
     def correct_pressure(self):
         self.p_prime = self.allocate_field()
-        for i in range(1, self.shape[0] - 1):
-            for j in range(1, self.shape[1] - 1):
-                a_E = self.assert_positive(-self.d_e[i][j] * self.step)
-                a_W = self.assert_positive(-self.d_e[i][j-1] * self.step)
-                a_N = self.assert_positive(-self.d_n[i-1][j] * self.step)
-                a_S = self.assert_positive(-self.d_n[i][j] * self.step)
-                a_P = a_E + a_W + a_N + a_S
-
-                self.b[i][j] = self.step * (
-                    -(self.u_star[i][j] - self.u_star[i][j-1])
-                    + (self.v_star[i][j] - self.v_star[i-1][j])
-                )
-
-                self.p_prime[i][j] = (
-                    a_E * self.p_prime[i][j+1] +
-                    a_W * self.p_prime[i][j-1] +
-                    a_N * self.p_prime[i-1][j] +
-                    a_S * self.p_prime[i+1][j] +
-                    self.b[i][j]
-                ) / a_P
+        for i, j in self.grid():
+            a_E = self.assert_positive(-self.d_e[i][j] * self.step)
+            a_W = self.assert_positive(-self.d_e[i][j-1] * self.step)
+            a_N = self.assert_positive(-self.d_n[i-1][j] * self.step)
+            a_S = self.assert_positive(-self.d_n[i][j] * self.step)
+            a_P = a_E + a_W + a_N + a_S
+
+            self.b[i][j] = self.step * (
+                -(self.u_star[i][j] - self.u_star[i][j-1])
+                + (self.v_star[i][j] - self.v_star[i-1][j])
+            )
+
+            self.p_prime[i][j] = (
+                a_E * self.p_prime[i][j+1] +
+                a_W * self.p_prime[i][j-1] +
+                a_N * self.p_prime[i-1][j] +
+                a_S * self.p_prime[i+1][j] +
+                self.b[i][j]
+            ) / a_P
 
         self.p = self.p_star + self.p_prime * self.alpha
         self.p_star = self.p
 
     def correct_velocities(self):
-        for i in range(1, self.shape[0] - 1):
-            for j in range(1, self.shape[1] - 1):
-                self.u[i][j] = self.u_star[i][j] + self.d_e[i][j] * (self.p_prime[i][j + 1] - self.p_prime[i][j])
-                self.v[i][j] = self.v_star[i][j] + self.d_n[i][j] * (self.p_prime[i][j] - self.p_prime[i + 1][j])
+        for i, j in self.grid():
+            self.u[i][j] = self.u_star[i][j] + self.d_e[i][j] * (self.p_prime[i][j + 1] - self.p_prime[i][j])
+            self.v[i][j] = self.v_star[i][j] + self.d_n[i][j] * (self.p_prime[i][j] - self.p_prime[i + 1][j])
 
     def iterate(self):
-        self.apply_sides_boundary()
         self.apply_inflow_boundary()
-        self.apply_bfs_boundary()
 
         self.solve_momentum_equations()
 
         self.apply_sides_boundary()
-        self.apply_inflow_boundary()
         self.apply_bfs_boundary()
 
         self.correct_pressure()
@@ -177,4 +171,4 @@ class SIMPLE:
             self.v = np.load(file)
             self.p = np.load(file)
             self.b = np.load(file)
-            self.p_star = self.allocate_field() # Make sure it's zeros
+            self.p_star = self.p