feat: rewrite with class and find awesome params

2 files changed, 214 insertions, 192 deletions

diff --git a/src/main.py b/src/main.py
index ed363a0..01f0c19 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,195 +1,11 @@
-import numpy as np
-import matplotlib.pyplot as plt
+from simple import SIMPLE
 
-PI = np.pi
+problem = SIMPLE((1, 2), (0.7, 0.5), 0.02, 120)
 
-np.set_printoptions(precision=2, floatmode="maxprec", suppress=True)
-cb = None
-plot = None
+problem.prep()
 
-
-Re = 170
-nu = 1 / Re
-domain_size = (1, 2)
-step = 0.05
-N = int(domain_size[0] / step)
-M = int(domain_size[1] / step)
-shape = (N, M)
-
-alpha = 0.8  # Pressure under-relaxation coefficient
-
-t_m = 1
-h_c = 0.8
-l_c = 0.5
-
-# Staggered vars
-u = np.zeros(shape=shape, dtype=float)
-u_star = np.zeros(shape=shape, dtype=float)
-
-v = np.zeros(shape=shape, dtype=float)
-v_star = np.zeros(shape=shape, dtype=float)
-
-p = np.zeros(shape=shape, dtype=float)
-p_star = np.random.rand(shape[0], shape[1])
-
-d_e = np.zeros(shape=shape, dtype=float)
-d_n = np.zeros(shape=shape, dtype=float)
-b = np.zeros(shape=shape, dtype=float)
-
-
-def assert_rule2(value):
-    assert value > -0.01, f'Coefficient must be positive: {value}' # > 0
-
-# Loop
-error = 1
-precision = 10 ** -7
-
-t = 0.8
-
-iteration = 0
-while error > precision:
-    iteration += 1
-    # Inflow boundary condition
-    for i in range(N):
-        y = i * step
-        u_star[i][0] = 1
-        v_star[i][0] = 0
-
-    # x-momentum
-    for i in range(1, N - 1):
-        for j in range(1, M - 1):
-            u_E = 0.5 * (u[i][j] + u[i][j + 1])
-            u_W = 0.5 * (u[i][j] + u[i][j - 1])
-            v_N = 0.5 * (v[i - 1][j] + v[i - 1][j + 1])
-            v_S = 0.5 * (v[i][j] + v[i][j + 1])
-
-            a_E = -0.5 * u_E * step + nu
-            a_W = +0.5 * u_W * step + nu
-            a_N = -0.5 * v_N * step + nu
-            a_S = +0.5 * v_S * step + nu
-            assert_rule2(a_E)
-            assert_rule2(a_W)
-            assert_rule2(a_N)
-            assert_rule2(a_S)
-
-            a_e = 0.5 * step * (u_E - u_W + v_N - v_S) + 4 * nu
-            A_e = -step
-
-            d_e[i][j] = A_e / a_e
-
-            u_star[i][j] = (a_E * u[i][j + 1] + a_W * u[i][j - 1] + a_N * u[i - 1][j] + a_S * u[i + 1][j]) / a_e
-            + d_e[i][j] * (p_star[i][j + 1] - p_star[i][j])
-
-    # y-momentum
-    for i in range(1, N - 1):
-        for j in range(1, M - 1):
-            u_E = 0.5 * (u[i][j] + u[i + 1][j])
-            u_W = 0.5 * (u[i][j - 1] + u[i + 1][j - 1])
-            v_N = 0.5 * (v[i - 1][j] + v[i][j])
-            v_S = 0.5 * (v[i][j] + v[i + 1][j])
-
-            a_E = -0.5 * u_E * step + nu
-            a_W = +0.5 * u_W * step + nu
-            a_N = -0.5 * v_N * step + nu
-            a_S = +0.5 * v_S * step + nu
-            assert_rule2(a_E)
-            assert_rule2(a_W)
-            assert_rule2(a_N)
-            assert_rule2(a_S)
-
-            a_n = 0.5 * step * (u_E - u_W + v_N - v_S) + 4 * nu
-            A_n = -step
-
-            d_n[i][j] = A_n / a_n
-
-            v_star[i][j] = (a_E * v[i][j + 1] + a_W * v[i][j - 1] + a_N * v[i - 1][j] + a_S * v[i + 1][j]) / a_n
-            + d_n[i][j] * (p_star[i][j] - p_star[i + 1][j])
-
-    # Sides boundary conditions
-    for j in range(M):
-        v_star[0][j] = 0
-        v_star[N - 1][j] = 0
-        u_star[0][j] = 0
-        u_star[N - 1][j] = 0
-
-    # Backwards-facing step boundary conditions (same as sides)
-    for i in range(int(h_c / step)):
-        for j in range(int(l_c / step)):
-            u_star[i][j] = 0
-            v_star[i][j] = 0
-
-    # Pressure correction
-    p_prime = np.zeros(shape=shape, dtype=float)
-    for i in range(1, N - 1):
-        for j in range(1, M - 1):
-            a_E = -d_e[i][j] * step
-            a_W = -d_e[i][j-1] * step
-            a_N = -d_n[i-1][j] * step
-            a_S = -d_n[i][j] * step
-            assert_rule2(a_E)
-            assert_rule2(a_W)
-            assert_rule2(a_N)
-            assert_rule2(a_S)
-            a_P = a_E + a_W + a_N + a_S
-            b[i][j] = step * (-(u_star[i][j] - u_star[i][j-1]) + (v_star[i][j] - v_star[i-1][j]))
-
-            p_prime[i][j] = (a_E * p_prime[i][j+1] + a_W * p_prime[i][j-1] + a_N * p_prime[i-1][j] + a_S * p_prime[i+1][j] + b[i][j]) / a_P
-    p = p_star + p_prime * alpha
-    p_star = p
-
-    # Velocity correction
-    for i in range(1, N - 1):
-        for j in range(1, M - 1):
-            u[i][j] = u_star[i][j] + d_e[i][j] * (p_prime[i][j + 1] - p_prime[i][j])
-            v[i][j] = v_star[i][j] + d_n[i][j] * (p_prime[i][j] - p_prime[i + 1][j])
-
-    # Backwards-facing step boundary conditions enforce
-    for i in range(int(h_c / step)):
-        for j in range(int(l_c / step)):
-            u[i][j] = 0
-            v[i][j] = 0
-
-    # Continuity residual as error measure
-    new_error = 0
-    for i in range(N):
-        for j in range(M):
-            new_error += abs(b[i][j])
-    new_error /= N * M
-    if abs(new_error - error) > 10 ** -20:
-        error = new_error
-    else:
-        print('Error decrease too small, you probably wont do better')
-        break
-
-    # Plotting
-    print(new_error)
-    x, y = np.meshgrid(
-        np.linspace(0, domain_size[1], shape[1]),
-        np.linspace(0, domain_size[0], shape[0]),
-    )
-
-    if iteration % 5 == 0 or iteration == 1:
-        print(f'Plotting, iteration: {iteration}')
-        if cb:
-            cb.remove()
-        if plot:
-            plot.remove()
-        factor = np.sqrt(u ** 2 + v ** 2)
-        u_normalized = u / factor
-        v_normalized = v / factor
-
-        density = 1
-        plot = plt.quiver(
-            x[::density, ::density],
-            y[::density, ::density],
-            u_normalized[::density, ::density],
-            v_normalized[::density, ::density],
-            p[::density, ::density],
-            scale=30,
-            cmap='plasma'
-        )
-        cb = plt.colorbar()
-        plt.pause(0.0001)
-
-
-plt.show()
+for i in range(3000):
+    print(f'Iteration {i}')
+    problem.iterate()
+    if i % 5 == 0 or i == 1:
+        problem.plot(True, 2)
diff --git a/src/simple.py b/src/simple.py
new file mode 100644
index 0000000..d524659
--- /dev/null
+++ b/src/simple.py
@@ -0,0 +1,206 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+figure, axes = plt.subplots()
+
+
+class SIMPLE:
+    def __init__(self, domain_size, bfs_size, step, Re, alpha=0.8):
+        np.set_printoptions(precision=2, floatmode="maxprec", suppress=True)
+
+        self.Re = Re
+        self.nu = 1 / Re
+        self.alpha = alpha
+
+        self.domain_size = domain_size
+        self.bfs_size = bfs_size
+
+        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):
+        self.u = self.allocate_field()
+        self.u_star = self.allocate_field()
+
+        self.v = self.allocate_field()
+        self.v_star = self.allocate_field()
+
+        self.p = self.allocate_field()
+        self.p_star = self.allocate_field(True)
+        self.p_prime = self.allocate_field()
+
+        self.d_e = self.allocate_field()
+        self.d_n = self.allocate_field()
+        self.b = self.allocate_field()
+
+    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):
+        '''Assert that the value is nearly positive'''
+        if value < -0.1:
+            # TODO: assert
+            print(f'WARNING: Value must be positive: {value}')
+        return value
+
+    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])
+
+                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 = 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.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])
+
+                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_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][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.u_star[self.shape[0] - 1][j] = 0
+            self.v_star[self.shape[0] - 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):
+                self.u_star[i][j] = 0
+                self.v_star[i][j] = 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_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_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])
+
+    def apply_outflow_boundary(self):
+        for i in range(self.shape[0]):
+            self.u[i][self.shape[1] - 1] = self.u[i][self.shape[1] - 2]
+            self.v[i][self.shape[1] - 1] = self.v[i][self.shape[1] - 2]
+
+    def iterate(self):
+        self.apply_inflow_boundary()
+
+        self.solve_momentum_equations()
+
+        self.apply_sides_boundary()
+        self.apply_bfs_boundary()
+
+        self.correct_pressure()
+        self.correct_velocities()
+
+        # self.apply_outflow_boundary()
+
+        self.apply_bfs_boundary()  # Is it needed?
+
+    def plot(self, normalize=False, density=1):
+        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)))
+
+        u, v = self.u, self.v
+        if normalize:
+            factor = np.sqrt(u ** 2 + v ** 2)
+            u = u / factor
+            v = v / factor
+
+        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()
+        plt.pause(0.0001)