feat: allow saving progress

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)