feat: improve plotting and parameters

1 files changed, 27 insertions, 12 deletions

diff --git a/src/main.py b/src/main.py
index 69a612a..5e0d682 100644
--- a/src/main.py
+++ b/src/main.py
@@ -5,19 +5,23 @@ PI = np.pi
 
 np.set_printoptions(precision=2, floatmode="maxprec", suppress=True)
 figure, axes = plt.subplots()
+cb = None
 
 
 Re = 100
 nu = 1 / Re
 domain_size = (1, 2)
 step = 0.05
-N = int(domain_size[1] / step)
-M = int(domain_size[0] / step)
+N = int(domain_size[0] / step)
+M = int(domain_size[1] / step)
+
+
 shape = (N, M)
-alpha = 0.8  # Coefficient
+alpha = 1  # Coefficient
 
 t_m = 1
-h_c = 0.5
+h_c = 0.3
+l_c = 0.6
 
 # Staggered vars
 u = np.zeros(shape=shape, dtype=float)
@@ -47,14 +51,14 @@ def u_boundary(t, y):
 error = 1
 precision = 10 ** -7
 
-t = 0.15
+t = 0.35
 
 iteration = 0
 while error > precision:
     iteration += 1
     # Inflow boundary condition
     for i in range(N):
-        y = i * step
+        y = domain_size[0] - i * step
         u_star[i][0] = u_boundary(t, y)
         v_star[i][0] = 0
 
@@ -107,6 +111,13 @@ while error > precision:
             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])
 
+    # 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
+            p_new[i][j] = 100
+
     # Pressure correction
     p_c = np.zeros(shape=shape, dtype=float)
     for i in range(1, N - 1):
@@ -146,18 +157,22 @@ while error > precision:
     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)
+    if iteration % 100 == 0:
+        pcolormesh = axes.pcolormesh(x, y, p, cmap='PuBuGn')
+        if (cb):
+            cb.remove()
+        cb = plt.colorbar(pcolormesh)
+
+        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.pause(0.0001)