feat: add research methods

1 files changed, 30 insertions, 77 deletions

diff --git a/src/simple.py b/src/simple.py
index e0ee6a7..dfd7e0f 100644
--- a/src/simple.py
+++ b/src/simple.py
@@ -1,7 +1,4 @@
 import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib.patches import Rectangle
-figure, axes = plt.subplots()
 
 
 class SIMPLE:
@@ -18,21 +15,7 @@ class SIMPLE:
         self.step = step
         self.shape = tuple(int(x / step) for x in domain_size)
 
-        self.x, self.y = np.meshgrid(
-            np.linspace(0, domain_size[1], self.shape[1]),
-            np.linspace(0, domain_size[0], self.shape[0]),
-        )
-
-        self.plt = None
-        self.colorbar = None
-        self.patch = None
-
-    def allocate_field(self, random=False):
-        if random:
-            return np.random.rand(*self.shape)
-        return np.zeros(shape=self.shape, dtype=float)
-
-    def prep(self):
+        # Allocations
         self.u = self.allocate_field()
         self.u_star = self.allocate_field()
 
@@ -47,16 +30,19 @@ class SIMPLE:
         self.d_n = self.allocate_field()
         self.b = self.allocate_field()
 
+    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(self.shape[0]):
             self.u_star[i][0] = 1
             self.v_star[i][0] = 0
 
-    def assert_rule2(self, value):
+    def assert_positive(self, value):
         '''Assert that the value is nearly positive'''
-        if value < -0.1:
-            # TODO: assert
-            print(f'WARNING: Value must be positive: {value}')
+        assert value > -0.01, f'WARNING: Value must be positive: {value}'
         return value
 
     def solve_momentum_equations(self):
@@ -68,10 +54,10 @@ class SIMPLE:
                 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_rule2(-0.5 * u_E * self.step + self.nu)
-                a_W = self.assert_rule2(+0.5 * u_W * self.step + self.nu)
-                a_N = self.assert_rule2(-0.5 * v_N * self.step + self.nu)
-                a_S = self.assert_rule2(+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
@@ -93,10 +79,10 @@ class SIMPLE:
                 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_rule2(-0.5 * u_E * self.step + self.nu)
-                a_W = self.assert_rule2(+0.5 * u_W * self.step + self.nu)
-                a_N = self.assert_rule2(-0.5 * v_N * self.step + self.nu)
-                a_S = self.assert_rule2(+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
@@ -115,13 +101,15 @@ class SIMPLE:
             self.v_star[0][j] = 0
             self.u_star[0][j] = 0
 
-            self.v_star[self.shape[0] - 1][j] = 0
-            self.u_star[self.shape[0] - 1][j] = 0
+            self.v_star[-1][j] = 0
+            self.u_star[-1][j] = 0
 
     def apply_bfs_boundary(self):
         '''Apply Backwards Facing Step boundary conditions'''
-        for i in range(int(self.bfs_size[0] / self.step) + 1):
-            for j in range(int(self.bfs_size[1] / self.step) + 1):
+        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
 
@@ -129,10 +117,10 @@ class SIMPLE:
         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_rule2(-self.d_e[i][j] * self.step)
-                a_W = self.assert_rule2(-self.d_e[i][j-1] * self.step)
-                a_N = self.assert_rule2(-self.d_n[i-1][j] * self.step)
-                a_S = self.assert_rule2(-self.d_n[i][j] * self.step)
+                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 * (
@@ -174,45 +162,8 @@ class SIMPLE:
     def avg_error(self):
         return np.absolute(self.b).sum()
 
-    def plot(self, normalize=True, density=False, save_path=''):
-        if self.patch:
-            self.patch.remove()
-        if self.colorbar:
-            self.colorbar.remove()
-        if self.plt:
-            self.plt.remove()
-
-        self.patch = axes.add_patch(Rectangle((0, 0), *reversed(self.bfs_size), color='gray'))
-        axes.set_title('Velocity field (normalized)')
-        plt.suptitle(f'Avg mass source per grid point = {self.avg_error()}')
-        plt.xlabel('X')
-        plt.ylabel('Y')
-
-        u, v = self.u, self.v
-        if normalize:
-            factor = np.sqrt(u ** 2 + v ** 2)
-            u = u / factor
-            v = v / factor
-
-        density = density or int((max(self.domain_size) / self.step) / 40)
-
-        self.plt = plt.quiver(
-            self.x[::density, ::density],
-            self.y[::density, ::density],
-            u[::density, ::density],
-            v[::density, ::density],
-            self.p[::density, ::density],
-            scale=30,
-            cmap='inferno'
-        )
-        self.colorbar = plt.colorbar(label='Pressure')
-
-        if len(save_path):
-            plt.savefig(save_path, dpi=300)
-        else:
-            plt.pause(0.0001)
-
     def save(self, path):
+        print('SAVE', path)
         with open(path, 'wb') as file:
             np.save(file, self.u)
             np.save(file, self.v)
@@ -220,8 +171,10 @@ class SIMPLE:
             np.save(file, self.b)
 
     def load(self, path):
+        print('LOAD', path)
         with open(path, 'rb') as file:
             self.u = np.load(file)
             self.v = np.load(file)
             self.p = np.load(file)
-            self.n = np.load(file)
+            self.b = np.load(file)
+            self.p_star = self.allocate_field() # Make sure it's zeros