diff options
Diffstat (limited to 'src/plotter.py')
| -rw-r--r-- | src/plotter.py | 23 |
1 files changed, 14 insertions, 9 deletions
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) |