1 files changed, 59 insertions, 0 deletions
diff --git a/src/plotter.py b/src/plotter.py
new file mode 100644
index 0000000..d19d344
--- /dev/null
+++ b/src/plotter.py
@@ -0,0 +1,59 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+
+figure, axes = plt.subplots()
+
+
+class Plotter:
+    def __init__(self):
+        self.plt = None
+        self.colorbar = None
+        self.patch = None
+
+    def clear(self):
+        if self.patch:
+            self.patch.remove()
+        if self.colorbar:
+            self.colorbar.remove()
+        if self.plt:
+            self.plt.remove()
+
+    def plot(self, model, normalize=True, density=False, save_path=''):
+        self.clear()
+
+        self.patch = axes.add_patch(Rectangle((0, 0), *reversed(model.bfs_size), color='gray'))
+        axes.set_title('Velocity field (normalized)')
+        plt.suptitle(f'Avg mass source per grid point = {model.avg_error()}')
+        plt.xlabel('X')
+        plt.ylabel('Y')
+
+        u, v = model.u, model.v
+        if normalize:
+            factor = np.sqrt(u ** 2 + v ** 2)
+            u = u / factor
+            v = v / factor
+
+        x, y = np.meshgrid(
+            np.linspace(0, model.domain_size[1], model.shape[1]),
+            np.linspace(0, model.domain_size[0], model.shape[0]),
+        )
+
+        density = density or int((max(model.domain_size) / model.step) / 40)
+
+        self.plt = plt.quiver(
+            x[::density, ::density],
+            y[::density, ::density],
+            u[::density, ::density],
+            v[::density, ::density],
+            model.p[::density, ::density],
+            scale=30,
+            cmap='inferno'
+        )
+        self.colorbar = plt.colorbar(label='Pressure')
+
+    def save(self, path):
+        return plt.savefig(path, dpi=300)
+
+    def show(self):
+        return plt.pause(0.0001)