feat: rewrite with class and find awesome params

1 files changed, 206 insertions, 0 deletions

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)