diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/simple.py | 23 |
1 files changed, 20 insertions, 3 deletions
diff --git a/src/simple.py b/src/simple.py index 5ffe690..e0ee6a7 100644 --- a/src/simple.py +++ b/src/simple.py @@ -174,7 +174,7 @@ class SIMPLE: def avg_error(self): return np.absolute(self.b).sum() - def plot(self, normalize=False, density=False): + def plot(self, normalize=True, density=False, save_path=''): if self.patch: self.patch.remove() if self.colorbar: @@ -182,7 +182,7 @@ class SIMPLE: if self.plt: self.plt.remove() - self.patch = axes.add_patch(Rectangle((0, 0), *reversed(self.bfs_size))) + self.patch = axes.add_patch(Rectangle((0, 0), *reversed(self.bfs_size), color='gray')) axes.set_title('Velocity field (normalized)') plt.suptitle(f'Avg mass source per grid point = {self.avg_error()}') plt.xlabel('X') @@ -207,4 +207,21 @@ class SIMPLE: ) self.colorbar = plt.colorbar(label='Pressure') - plt.pause(0.0001) + if len(save_path): + plt.savefig(save_path, dpi=300) + else: + plt.pause(0.0001) + + def save(self, path): + with open(path, 'wb') as file: + np.save(file, self.u) + np.save(file, self.v) + np.save(file, self.p) + np.save(file, self.b) + + def load(self, path): + with open(path, 'rb') as file: + self.u = np.load(file) + self.v = np.load(file) + self.p = np.load(file) + self.n = np.load(file) |