refactor: cleanup

1 files changed, 22 insertions, 26 deletions

diff --git a/src/main.py b/src/main.py
index fffb570..4df1526 100644
--- a/src/main.py
+++ b/src/main.py
@@ -8,7 +8,7 @@ figure, axes = plt.subplots()
 cb = None
 
 
-Re = 400
+Re = 200
 nu = 1 / Re
 domain_size = (1, 2)
 step = 0.05
@@ -17,7 +17,7 @@ M = int(domain_size[1] / step)
 
 
 shape = (N, M)
-alpha = 1  # Coefficient
+alpha = 0.8  # Pressure under-relaxation coefficient
 
 t_m = 1
 h_c = 0.3
@@ -26,13 +26,12 @@ l_c = 0.6
 # 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.random.rand(shape[0], shape[1])
+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)
@@ -45,7 +44,7 @@ def u_boundary(t, y):
     return 6 * f(t) * (y - h_c) * (1 - y) / (1 - h_c)**2
 
 def assert_rule2(value):
-    assert(value > -0.01) # > 0
+    assert value > -0.01, f'Coefficient must be positive: {value}' # > 0
 
 # Loop
 error = 1
@@ -74,7 +73,7 @@ while error > precision:
         for j in range(int(l_c / step)):
             u_star[i][j] = 0
             v_star[i][j] = 0
-            p[i][j] = 0
+            p_star[i][j] = 0
 
     # x-momentum
     for i in range(1, N - 1):
@@ -99,7 +98,7 @@ while error > precision:
             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])
+            + d_e[i][j] * (p_star[i][j + 1] - p_star[i][j])
 
     # y-momentum
     for i in range(1, N - 1):
@@ -124,10 +123,10 @@ while error > precision:
             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])
+            + d_n[i][j] * (p_star[i][j] - p_star[i + 1][j])
 
     # Pressure correction
-    p_c = np.zeros(shape=shape, dtype=float)
+    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
@@ -141,29 +140,30 @@ while error > precision:
             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
+            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
 
     # 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]
+        p[i][0] = p[i][1]
+        p[i][M - 1] = p[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]
+        p[0][j] = p[1][j]
+        p[N - 1][j] = p[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])
+            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_new[i][j] = 0
-            v_new[i][j] = 0
-            p_new[i][j] = 0
+            u[i][j] = 0
+            v[i][j] = 0
+            p[i][j] = 0
 
     # Continuity residual as error measure
     error = 0
@@ -171,10 +171,6 @@ while error > precision:
         for j in range(M):
             error += abs(b[i][j])
 
-    u = u_new
-    v = v_new
-    p = p_new
-
     # Plotting
     print(error)
     x, y = np.meshgrid(
@@ -191,7 +187,7 @@ while error > precision:
         factor = np.sqrt(u ** 2 + v ** 2)
         u_normalized = u / factor
         v_normalized = v / factor
-        plt.quiver(x, y, u_normalized, v_normalized, scale=30)
+        plt.quiver(x, y, u_normalized, v_normalized, scale=45)
         plt.pause(0.0001)