feat: minor improvements

3 files changed, 18 insertions, 13 deletions

diff --git a/src/main.py b/src/main.py
index 485755e..b74e12c 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,9 +1,9 @@
 from simple import SIMPLE
 from research import Research
 
-model = SIMPLE((100, 200), (50, 50), 0.001, 100)
+model = SIMPLE((50, 100), (40, 40), 0.001, 700)
 
-research = Research(model, f'{model.p.shape[0]}x{model.p.shape[1]}-{model.Re}')
+research = Research(model, f'{model.p.shape[0] * model.step}x{model.p.shape[1] * model.step}_{model.bfs_shape[0] * model.step}x{model.bfs_shape[1] * model.step}_Re{model.Re}/{model.step}')
 
 is_complete = research.load()
 
diff --git a/src/plotter.py b/src/plotter.py
index c322de9..b493c3e 100644
--- a/src/plotter.py
+++ b/src/plotter.py
@@ -18,7 +18,10 @@ class Plotter:
         if self.colorbar:
             self.colorbar.remove()
         if self.plt:
-            self.plt.remove()
+            try:
+                self.plt.remove()
+            except:
+                pass
 
     def plot(self, model, normalize=True, density=False, save_path='', streamplot=False):
         self.clear()
@@ -51,20 +54,22 @@ class Plotter:
             u = u / factor
             v = v / factor
 
-        # density = density or int((max(model.domain_size) / model.step) / 40)
+        density = density or int(min(model.p.shape) / 40)
 
-        plt.contourf(x, y, model.p)
-        # self.patch = axes.add_patch(Rectangle((0, 0), *reversed(model.bfs_size), color='gray'))
+        plt.contourf(x, y, model.p, cmap='inferno')
+        self.colorbar = plt.colorbar(label='Pressure')
+
+        self.patch = axes.add_patch(Rectangle((0, 0), *reversed(list(x * model.step for x in model.bfs_shape)), color='gray'))
 
         plotter = plt.streamplot if streamplot else plt.quiver
         self.plt = plotter(
-            x,
-            y,
-            u,
-            v,
+            x[::density, ::density],
+            y[::density, ::density],
+            u[::density, ::density],
+            v[::density, ::density],
             color='black'
         )
-        self.colorbar = plt.colorbar(label='Pressure')
+
 
     def save(self, path):
         return plt.savefig(path, dpi=300)
diff --git a/src/research.py b/src/research.py
index 8c46877..e085b07 100644
--- a/src/research.py
+++ b/src/research.py
@@ -32,7 +32,7 @@ class Research:
             error = self.model.avg_error()
             print(f'Iteration {iteration}, avg error: {error}')
 
-            if iteration % 10 == 0 or iteration == 1:
+            if iteration % 50 == 0 or iteration == 1:
                 if preview or save_plot:
                     self.plotter.plot(self.model, normalize=True)
                 if preview:
@@ -52,7 +52,7 @@ class Research:
         self.inspect()
 
     def inspect(self):
-        self.plotter.plot(self.model, streamplot=True)
+        self.plotter.plot(self.model, streamplot=True, density=1)
         while True:
             self.plotter.show()