feat: add research methods

4 files changed, 150 insertions, 89 deletions

diff --git a/src/main.py b/src/main.py
index 551595e..cdbf16e 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,16 +1,13 @@
 from simple import SIMPLE
+from research import Research
 
-problem = SIMPLE((1, 2), (0.7, 0.5), 0.02, 120)
-problem.prep()
+model = SIMPLE((1, 2), (0.7, 0.5), 0.02, 120)
 
-error = 1
-iteration = 0
-while error > 10 ** -7:
-    iteration += 1
-    problem.iterate()
-    error = problem.avg_error()
-    print(f'Iteration {iteration}')
-    print(f'Avg error: {error}')
+research = Research(model, '07x05_002')
 
-    if iteration % 5 == 0 or iteration == 1:
-        problem.plot(True, 2)
+is_complete = research.load()
+
+if is_complete:
+    research.inspect()
+else:
+    research.solve(preview=True, save_plot=True, save_model=True)
diff --git a/src/plotter.py b/src/plotter.py
new file mode 100644
index 0000000..d19d344
--- /dev/null
+++ b/src/plotter.py
@@ -0,0 +1,59 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+
+figure, axes = plt.subplots()
+
+
+class Plotter:
+    def __init__(self):
+        self.plt = None
+        self.colorbar = None
+        self.patch = None
+
+    def clear(self):
+        if self.patch:
+            self.patch.remove()
+        if self.colorbar:
+            self.colorbar.remove()
+        if self.plt:
+            self.plt.remove()
+
+    def plot(self, model, normalize=True, density=False, save_path=''):
+        self.clear()
+
+        self.patch = axes.add_patch(Rectangle((0, 0), *reversed(model.bfs_size), color='gray'))
+        axes.set_title('Velocity field (normalized)')
+        plt.suptitle(f'Avg mass source per grid point = {model.avg_error()}')
+        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
+
+        x, y = np.meshgrid(
+            np.linspace(0, model.domain_size[1], model.shape[1]),
+            np.linspace(0, model.domain_size[0], model.shape[0]),
+        )
+
+        density = density or int((max(model.domain_size) / model.step) / 40)
+
+        self.plt = plt.quiver(
+            x[::density, ::density],
+            y[::density, ::density],
+            u[::density, ::density],
+            v[::density, ::density],
+            model.p[::density, ::density],
+            scale=30,
+            cmap='inferno'
+        )
+        self.colorbar = plt.colorbar(label='Pressure')
+
+    def save(self, path):
+        return plt.savefig(path, dpi=300)
+
+    def show(self):
+        return plt.pause(0.0001)
diff --git a/src/research.py b/src/research.py
new file mode 100644
index 0000000..f0374cf
--- /dev/null
+++ b/src/research.py
@@ -0,0 +1,52 @@
+import os
+from plotter import Plotter
+
+
+class Research:
+    def __init__(self, model, name):
+        self.model = model
+        self.name = name
+        self.plotter = Plotter()
+        self.path = os.path.join('data', name)
+        os.makedirs(self.path, exist_ok=True)
+
+        self.model_path = os.path.join(self.path, 'model.npy')
+        self.solution_path = os.path.join(self.path, 'solution.npy')
+
+    def load(self):
+        if os.path.exists(self.solution_path):
+            self.model.load(self.solution_path)
+            return 1
+
+        if os.path.exists(self.model_path):
+            self.model.load(self.model_path)
+
+        return 0
+
+    def solve(self, precision=10 ** -4, preview=True, save_plot=False, save_model=False):
+        error = 1
+        iteration = 0
+        while error > precision:
+            iteration += 1
+            self.model.iterate()
+            error = self.model.avg_error()
+            print(f'Iteration {iteration}, avg error: {error}')
+
+            if iteration % 10 == 0 or iteration == 1:
+                if preview or save_plot:
+                    self.plotter.plot(self.model)
+                if preview:
+                    self.plotter.show()
+
+            if iteration % 50 == 0:
+                if save_model:
+                    self.model.save(self.model_path)
+
+        self.model.save(self.solution_path)
+        self.inspect()
+
+    def inspect(self):
+        self.plotter.plot(self.model, density=1)
+        while True:
+            self.plotter.show()
+
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