From 1114754da4dc2a330b6843297d647a38cd7cacf4 Mon Sep 17 00:00:00 2001
From: "Edgar P. Burkhart" <git@edgarpierre.fr>
Date: Thu, 14 Apr 2022 12:37:42 +0200
Subject: [PATCH] Compress pickled object

---
 olaflow/config.ini            |  2 +-
 olaflow/processing/animate.py | 18 ++++++------------
 olaflow/processing/pickle.py  |  8 +++++---
 3 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/olaflow/config.ini b/olaflow/config.ini
index 2bac493..3026449 100644
--- a/olaflow/config.ini
+++ b/olaflow/config.ini
@@ -17,6 +17,6 @@ tf = 14300
 x0 = -150
 
 [post]
-pickle = out_post/model.pickle
+out = out_post
 x = -50
 z = 5
diff --git a/olaflow/processing/animate.py b/olaflow/processing/animate.py
index 78da9f9..da24ab6 100644
--- a/olaflow/processing/animate.py
+++ b/olaflow/processing/animate.py
@@ -1,5 +1,6 @@
 import argparse
 import configparser
+import gzip
 import logging
 import multiprocessing as mp
 import pathlib
@@ -23,10 +24,10 @@ log = logging.getLogger("ola_post")
 log.info("Starting sws -> olaFlow converter")
 config = configparser.ConfigParser()
 config.read(args.config)
-out = pathlib.Path(config.get("post", "pickle"))
-out.parent.mkdir(parents=True, exist_ok=True)
+out = pathlib.Path(config.get("post", "out"))
+out.mkdir(parents=True, exist_ok=True)
 
-with out.open("rb") as f:
+with gzip.open(out.joinpath("pickle.gz"), "rb") as f:
     model = pickle.load(f)
 
 x0 = config.getfloat("post", "x")
@@ -115,12 +116,5 @@ def animU(i):
 ani = animation.FuncAnimation(fig, anim, frames=model.t.size, interval=1/24)
 aniU = animation.FuncAnimation(figU, animU, frames=model.t.size, interval=1/24)
 
-a = mp.Process(target=ani.save, args=(out.parent.joinpath("anim.mp4"),), kwargs={"fps": 24})
-aU = mp.Process(target=aniU.save, args=(out.parent.joinpath("animU.mp4"),), kwargs={"fps": 24})
-
-a.start()
-aU.start()
-
-a.join()
-aU.join()
-#plt.show()
+ani.save(out.joinpath("anim.mp4"), fps=24)
+aniU.save(out.joinpath("animU.mp4"), fps=24)
diff --git a/olaflow/processing/pickle.py b/olaflow/processing/pickle.py
index 1dcf79f..a82cb2b 100644
--- a/olaflow/processing/pickle.py
+++ b/olaflow/processing/pickle.py
@@ -1,8 +1,10 @@
 import argparse
+import gzip
 import configparser
 import logging
 import pathlib
 import pickle
+from time import time
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -22,8 +24,8 @@ log = logging.getLogger("ola_post")
 log.info("Starting sws -> olaFlow converter")
 config = configparser.ConfigParser()
 config.read(args.config)
-out = pathlib.Path(config.get("post", "pickle"))
-out.parent.mkdir(parents=True, exist_ok=True)
+out = pathlib.Path(config.get("post", "out"))
+out.mkdir(parents=True, exist_ok=True)
 
 olaflow_root = args.output
 model = OFModel(olaflow_root)
@@ -35,5 +37,5 @@ model.read_field_all("p")
 model.read_field_all("p_rgh")
 model.read_field_all("U")
 
-with out.open("wb") as f:
+with gzip.open(out.joinpath("pickle.gz"), "wb") as f:
     pickle.dump(model, f)