feat: improve plotting and output

1 files changed, 27 insertions, 10 deletions

diff --git a/src/main.py b/src/main.py
index 4df1526..5904fbb 100644
--- a/src/main.py
+++ b/src/main.py
@@ -4,8 +4,8 @@ import matplotlib.pyplot as plt
 PI = np.pi
 
 np.set_printoptions(precision=2, floatmode="maxprec", suppress=True)
-figure, axes = plt.subplots()
 cb = None
+plot = None
 
 
 Re = 200
@@ -166,28 +166,45 @@ while error > precision:
             p[i][j] = 0
 
     # Continuity residual as error measure
-    error = 0
+    new_error = 0
     for i in range(N):
         for j in range(M):
-            error += abs(b[i][j])
+            new_error += abs(b[i][j])
+    new_error /= N * M
+    if abs(new_error - error) > 10 ** -12:
+        error = new_error
+    else:
+        print('Error decrease too small, you probably wont do better')
+        break
 
     # Plotting
-    print(error)
+    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 % 100 == 0:
-        pcolormesh = axes.pcolormesh(x, y, p, cmap='PuBuGn')
-        if (cb):
+    if iteration % 5 == 0 or iteration == 1:
+        print(f'Plotting, iteration: {iteration}')
+        if cb:
             cb.remove()
-        cb = plt.colorbar(pcolormesh)
-
+        if plot:
+            plot.remove()
         factor = np.sqrt(u ** 2 + v ** 2)
         u_normalized = u / factor
         v_normalized = v / factor
-        plt.quiver(x, y, u_normalized, v_normalized, scale=45)
+
+        density = 2
+        plot = plt.quiver(
+            x[::density, ::density],
+            y[::density, ::density],
+            u_normalized[::density, ::density],
+            v_normalized[::density, ::density],
+            p[::density, ::density],
+            scale=30,
+            cmap='inferno'
+        )
+        cb = plt.colorbar()
         plt.pause(0.0001)