feat: finish working bfs example

4 files changed, 103 insertions, 133 deletions

diff --git a/src/main.py b/src/main.py
index 27620b7..485755e 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,13 +1,13 @@
 from simple import SIMPLE
 from research import Research
 
-model = SIMPLE((30, 30), (0, 0), 0.02, 40)
+model = SIMPLE((100, 200), (50, 50), 0.001, 100)
 
-research = Research(model, 'testing')
+research = Research(model, f'{model.p.shape[0]}x{model.p.shape[1]}-{model.Re}')
 
 is_complete = research.load()
 
 if is_complete:
     research.inspect()
 else:
-    research.solve(preview=True)
+    research.solve(preview=True, save_model=True, save_plot=True)
diff --git a/src/plotter.py b/src/plotter.py
index 45d95da..c322de9 100644
--- a/src/plotter.py
+++ b/src/plotter.py
@@ -20,7 +20,7 @@ class Plotter:
         if self.plt:
             self.plt.remove()
 
-    def plot(self, model, normalize=True, density=False, save_path=''):
+    def plot(self, model, normalize=True, density=False, save_path='', streamplot=False):
         self.clear()
 
         axes.set_title('Velocity field (normalized)')
@@ -28,38 +28,41 @@ class Plotter:
         plt.xlabel('X')
         plt.ylabel('Y')
 
-        u, v = model.u, model.v
-        if normalize:
-            factor = np.sqrt(u ** 2 + v ** 2)
-            u = u / factor
-            v = v / factor
+        shape = model.p.shape
+
+        u = np.zeros(shape, dtype=float)
+        v = np.zeros(shape, dtype=float)
+
+        for i in range(shape[0]):
+            for j in range(shape[1]):
+                u[i][j] = 0.5 * (model.u[i][j] + model.u[i][j + 1])
+                v[i][j] = 0.5 * (model.v[i][j] + model.v[i + 1][j])
 
+        assert not v[0, :].any()
+        assert not v[-1, :].any()
 
-        shape = (model.p.shape[0] + 1, model.p.shape[1] + 1)
         x, y = np.meshgrid(
             np.linspace(0, shape[1] * model.step, shape[1]),
             np.linspace(0, shape[0] * model.step, shape[0]),
         )
 
-        u = copy(model.u)
-        u.resize(shape)
-        v = copy(model.v)
-        v.resize(shape)
-        p = copy(model.p)
-        p.resize(shape)
-
-        print(shape, u.shape, v.shape)
+        if normalize:
+            factor = np.sqrt(u ** 2 + v ** 2)
+            u = u / factor
+            v = v / factor
 
         # density = density or int((max(model.domain_size) / model.step) / 40)
 
-        plt.contourf(x, y, p)
+        plt.contourf(x, y, model.p)
         # self.patch = axes.add_patch(Rectangle((0, 0), *reversed(model.bfs_size), color='gray'))
-        # TODO: allow using streamplot
-        self.plt = plt.quiver(
+
+        plotter = plt.streamplot if streamplot else plt.quiver
+        self.plt = plotter(
             x,
             y,
             u,
             v,
+            color='black'
         )
         self.colorbar = plt.colorbar(label='Pressure')
 
diff --git a/src/research.py b/src/research.py
index 2a59c49..8c46877 100644
--- a/src/research.py
+++ b/src/research.py
@@ -34,7 +34,7 @@ class Research:
 
             if iteration % 10 == 0 or iteration == 1:
                 if preview or save_plot:
-                    self.plotter.plot(self.model, normalize=False, density=1)
+                    self.plotter.plot(self.model, normalize=True)
                 if preview:
                     self.plotter.show()
                 if save_plot:
@@ -45,10 +45,14 @@ class Research:
                     self.model.save(self.model_path)
 
         self.model.save(self.solution_path)
+
+        self.plotter.plot(self.model, streamplot=True)
+        self.plotter.save(os.path.join(self.path, f'streamplot.png'))
+
         self.inspect()
 
     def inspect(self):
-        self.plotter.plot(self.model, density=1)
+        self.plotter.plot(self.model, streamplot=True)
         while True:
             self.plotter.show()
 
diff --git a/src/simple.py b/src/simple.py
index b54cd38..0e103ea 100644
--- a/src/simple.py
+++ b/src/simple.py
@@ -27,120 +27,89 @@ class SIMPLE:
         self.d_n = np.zeros(shape=self.v.shape, dtype=float)
         self.b = np.zeros(shape=shape, dtype=float)
 
-    def allocate_field(self, random=False):
-        if random:
-            return np.random.rand(*self.shape)
-        return np.zeros(shape=self.shape, dtype=float)
-
-    def apply_inflow_boundary(self):
-        for i in range(1, self.shape[0] - 1):
-            self.u_star[i][0] = 1
-            self.v_star[i][0] = 0
-
-    def apply_outflow_boundary(self):
-        for i in range(0, self.shape[0]):
-            self.u_star[i][-1] = self.u_star[i][-2];
-
     def assert_positive(self, value):
         '''Assert that the value is nearly positive'''
-        assert value > -1, f'WARNING: Value must be positive: {value}'
+        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(1, 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.u.shape[0] - 1):
             for j in range(1, self.u.shape[1] - 1):
-                u_W = 0.5 * (self.u[i][j] + self.u[i][j - 1])
-                u_E = 0.5 * (self.u[i][j] + self.u[i][j + 1])
+                if i >= self.bfs_shape[0] or j >= self.bfs_shape[1]:
+                    u_W = 0.5 * (self.u[i][j] + self.u[i][j - 1])
+                    u_E = 0.5 * (self.u[i][j] + self.u[i][j + 1])
 
-                v_S = 0.5 * (self.v[i][j - 1] + self.v[i][j])
-                v_N = 0.5 * (self.v[i + 1][j - 1] + self.v[i + 1][j])
+                    v_S = 0.5 * (self.v[i][j - 1] + self.v[i][j])
+                    v_N = 0.5 * (self.v[i + 1][j - 1] + 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_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]) # p - p_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]) # p - p_e
 
         # Momentum along Y direction
         for i in range(1, self.v.shape[0] - 1):
             for j in range(1, self.v.shape[1] - 1):
-                u_W = 0.5 * (self.u[i - 1][j] + self.u[i][j])
-                u_E = 0.5 * (self.u[i - 1][j + 1] + self.u[i][j + 1])
-
-                v_N = 0.5 * (self.v[i][j] + self.v[i + 1][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_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.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 - 1][j] - self.p_star[i][j]) # p - p_n
-
-    def apply_sides_boundary(self):
-        for j in range(self.shape[1]):
-            self.v_star[0][j] = 0
-            self.v_star[-2][j] = 0
-            # WORKSNICE: self.v_star[-3][j] = 0
-
-    def apply_bfs_boundary(self):
-        '''Apply Backwards Facing Step boundary conditions'''
-        for i in range(self.bfs_node_size[0]):
-            self.u_star[i][self.bfs_node_size[1]] = 0
-        for j in range(self.bfs_node_size[1]):
-            self.v_star[self.bfs_node_size[0]][j] = 0
+                if i >= self.bfs_shape[0] or j >= self.bfs_shape[1]:
+                    u_W = 0.5 * (self.u[i - 1][j] + self.u[i][j])
+                    u_E = 0.5 * (self.u[i - 1][j + 1] + self.u[i][j + 1])
+
+                    v_N = 0.5 * (self.v[i][j] + self.v[i + 1][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_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.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 - 1][j] - self.p_star[i][j]) # p - p_n
 
     def correct_pressure(self):
         self.p_prime = np.zeros(shape=self.p.shape, dtype=float)
         for i in range(1, self.p.shape[0] - 1):
             for j in range(1, self.p.shape[1] - 1):
-                a_E = 0 if j == self.p.shape[1] - 1 else self.assert_positive(-self.d_e[i][j] * self.step)
-                a_W = 0 if j == 1 else self.assert_positive(-self.d_e[i][j+1] * self.step)
-                a_N = 0 if i == self.p.shape[0] - 1 else self.assert_positive(-self.d_n[i+1][j] * self.step)
-                a_S = 0 if i == 1 else 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+1][j] - self.v_star[i][j])
-                )
-
-                self.p_prime[i][j] = (
-                    (a_E * self.p_prime[i][j+1] if a_E > 0 else 0) +
-                    (a_W * self.p_prime[i][j-1] if a_W > 0 else 0) +
-                    (a_N * self.p_prime[i+1][j] if a_N > 0 else 0) +
-                    (a_S * self.p_prime[i-1][j] if a_S > 0 else 0) +
-                    self.b[i][j]
-                ) / a_P
+                if i >= self.bfs_shape[0] or j >= self.bfs_shape[1]:
+                    a_E = 0 if j == self.p.shape[1] - 1 else self.assert_positive(-self.d_e[i][j] * self.step)
+                    a_W = 0 if j == 1 else self.assert_positive(-self.d_e[i][j+1] * self.step)
+                    a_N = 0 if i == self.p.shape[0] - 1 else self.assert_positive(-self.d_n[i+1][j] * self.step)
+                    a_S = 0 if i == 1 else 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+1][j] - self.v_star[i][j])
+                    )
+
+                    self.p_prime[i][j] = (
+                        (a_E * self.p_prime[i][j+1] if a_E > 0 else 0) +
+                        (a_W * self.p_prime[i][j-1] if a_W > 0 else 0) +
+                        (a_N * self.p_prime[i+1][j] if a_N > 0 else 0) +
+                        (a_S * self.p_prime[i-1][j] if a_S > 0 else 0) +
+                        self.b[i][j]
+                    ) / a_P
 
         self.p = self.p_star + self.p_prime * self.alpha
         self.p_star = self.p
@@ -158,37 +127,31 @@ class SIMPLE:
         self.solve_momentum_equations()
 
         # Boundary
-        self.u_star[:, 0] = -self.u_star[:, 1]
+        self.u_star[:, 0] = 2 - self.u_star[:, 1]
         self.v_star[:, 0] = 0
 
-        self.u_star[:, -1] = -self.u_star[:, -2]
-        self.v_star[:, -1] = 0
+        self.v_star[-2, :] = -self.v_star[-1, :]
+        self.v_star[1, :] = -self.v_star[0, :]
 
-        self.u_star[-1, :] = 1
-        self.v_star[-1, :] = -self.v_star[-2, :]
+        self.v_star[self.bfs_shape[0], :self.bfs_shape[1]] = self.v_star[self.bfs_shape[0] - 1, :self.bfs_shape[1]]
+        self.u_star[:self.bfs_shape[0], self.bfs_shape[1]] = self.u_star[:self.bfs_shape[0], self.bfs_shape[1] - 1]
 
-        self.u_star[0, :] = 0
-        self.v_star[0, :] = -self.v_star[1, :]
-
-        self.p_star[0, 0] = 0
+        self.p_star[:self.bfs_shape[0], :self.bfs_shape[1]] = 0
 
         self.correct_pressure()
         self.correct_velocities()
 
         # Boundary enforce
-        self.u[:, 0] = -self.u[:, 1]
+        self.u[:, 0] = 2 - self.u[:, 1]
         self.v[:, 0] = 0
 
-        self.u[:, -1] = -self.u[:, -2]
-        self.v[:, -1] = 0
-
-        self.u[-1, :] = 1
-        self.v[-1, :] = -self.v[-2, :]
+        self.v[-2, :] = -self.v[-1, :]
+        self.v[1, :] = -self.v[0, :]
 
-        self.u[0, :] = 0
-        self.v[0, :] = -self.v[1, :]
+        self.v[self.bfs_shape[0], :self.bfs_shape[1]] = self.v[self.bfs_shape[0] - 1, :self.bfs_shape[1]]
+        self.u[:self.bfs_shape[0], self.bfs_shape[1]] = self.u[:self.bfs_shape[0], self.bfs_shape[1] - 1]
 
-        self.p[0, 0] = 0
+        self.p[:self.bfs_shape[0], :self.bfs_shape[1]] = 0
 
     def avg_error(self):
         return np.absolute(self.b).sum()