feat: add initial python implementation

1 files changed, 164 insertions, 0 deletions

diff --git a/src/main.py b/src/main.py
new file mode 100644
index 0000000..69a612a
--- /dev/null
+++ b/src/main.py
@@ -0,0 +1,164 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+PI = np.pi
+
+np.set_printoptions(precision=2, floatmode="maxprec", suppress=True)
+figure, axes = plt.subplots()
+
+
+Re = 100
+nu = 1 / Re
+domain_size = (1, 2)
+step = 0.05
+N = int(domain_size[1] / step)
+M = int(domain_size[0] / step)
+shape = (N, M)
+alpha = 0.8  # Coefficient
+
+t_m = 1
+h_c = 0.5
+
+# Staggered vars
+u = np.zeros(shape=shape, dtype=float)
+u_star = np.zeros(shape=shape, dtype=float)
+u_new = np.zeros(shape=shape, dtype=float)
+
+v = np.zeros(shape=shape, dtype=float)
+v_star = np.zeros(shape=shape, dtype=float)
+v_new = np.zeros(shape=shape, dtype=float)
+
+p = np.zeros(shape=shape, dtype=float)
+p_star = np.zeros(shape=shape, dtype=float)
+p_new = np.zeros(shape=shape, dtype=float)
+
+d_e = np.zeros(shape=shape, dtype=float)
+d_n = np.zeros(shape=shape, dtype=float)
+b = np.zeros(shape=shape, dtype=float)
+
+
+def u_boundary(t, y):
+    def f(t):
+        return 1 if t_m < t else 0.5 * (np.sin(0.5 * PI * (2 * t / t_m - 1)) + 1)
+    return 6 * f(t) * (y - h_c) * (1 - y) / (1 - h_c)**2
+
+
+# Loop
+error = 1
+precision = 10 ** -7
+
+t = 0.15
+
+iteration = 0
+while error > precision:
+    iteration += 1
+    # Inflow boundary condition
+    for i in range(N):
+        y = i * step
+        u_star[i][0] = u_boundary(t, y)
+        v_star[i][0] = 0
+
+    # 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
+
+    # 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
+
+            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[i][j + 1] - p[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
+
+            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[i][j] - p[i + 1][j])
+
+    # Pressure correction
+    p_c = 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
+            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_c[i][j] = (a_E * p_c[i][j+1] + a_W * p_c[i][j-1] + a_N * p_c[i-1][j] + a_S * p_c[i+1][j] + b[i][j]) / a_P
+    p_new = p + p_c * alpha
+
+    # Pressure boundaries
+    for i in range(N - 1):
+        p_new[i][0] = p_new[i][1]
+        p_new[i][M - 1] = p_new[i][M - 2]
+    for j in range(M - 1):
+        p_new[0][j] = p_new[1][j]
+        p_new[N - 1][j] = p_new[N - 2][j]
+
+    # Velocity correction
+    for i in range(1, N - 1):
+        for j in range(1, M - 1):
+            u_new[i][j] = u_star[i][j] + alpha * d_e[i][j] * (p_c[i + 1][j] - p_c[i][j])
+            v_new[i][j] = v_star[i][j] + alpha * d_n[i][j] * (p_c[i][j] - p_c[i + 1][j])
+
+    # Continuity residual as error measure
+    error = 0
+    for i in range(N):
+        for j in range(M):
+            error += abs(b[i][j])
+
+    u = u_new
+    v = v_new
+    p = p_new
+
+    # Plotting
+    print(u)
+    print(v)
+    print(p)
+    print(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:
+        axes.pcolormesh(x, y, p, cmap='hot')
+        plt.quiver(x, y, u, v, scale=1)
+        plt.pause(0.0001)
+
+
+plt.show()