From a9230f485ffa9217e4294163586c18d83a858ab1 Mon Sep 17 00:00:00 2001
From: "Edgar P. Burkhart" <git@edgarpierre.fr>
Date: Thu, 14 Apr 2022 14:43:13 +0200
Subject: [PATCH] Post-processing: U, P plots

---
 olaflow/processing/post.py | 88 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 88 insertions(+)
 create mode 100644 olaflow/processing/post.py

diff --git a/olaflow/processing/post.py b/olaflow/processing/post.py
new file mode 100644
index 0000000..9beb31b
--- /dev/null
+++ b/olaflow/processing/post.py
@@ -0,0 +1,88 @@
+import argparse
+import configparser
+import gzip
+from itertools import starmap
+import logging
+import multiprocessing as mp
+import pathlib
+import pickle
+
+from cycler import cycler
+import matplotlib.pyplot as plt
+import matplotlib.gridspec as gridspec
+import numpy as np
+from scipy import interpolate
+
+from .olaflow import OFModel
+
+parser = argparse.ArgumentParser(description="Post-process olaflow results")
+parser.add_argument("-v", "--verbose", action="count", default=0)
+parser.add_argument("-c", "--config", default="config.ini")
+args = parser.parse_args()
+
+logging.basicConfig(level=max((10, 20 - 10 * args.verbose)))
+log = logging.getLogger("ola_post")
+
+log.info("Animating olaFlow output")
+config = configparser.ConfigParser()
+config.read(args.config)
+out = pathlib.Path(config.get("post", "out"))
+out.mkdir(parents=True, exist_ok=True)
+
+with gzip.open(gzf, "rb") if (
+    gzf := out.joinpath("pickle.gz")
+).exists() else gzf.with_suffix("").open("rb") as f:
+    model = pickle.load(f)
+
+x0_conf = config.getfloat("post", "x")
+x0_val = model.x[np.argmin(np.abs(model.x - x0_conf))]
+# z0 = config.getfloat("post", "z")
+# z0 = np.linspace(-5, 5, 16)
+c0_ = ((model.x == x0_val)[None, :] & (model.fields["alpha.water"] > 0.95)).any(axis=0)
+c0 = model.coords[c0_][:: (c0_.sum() // 8 + 1)]
+
+i0 = np.argmin(
+    np.linalg.norm(model.coords[..., None] - c0.T[None, ...], axis=1),
+    axis=0,
+)
+
+aw = model.fields["alpha.water"][:, i0]
+
+U = np.where(aw > 0.95, np.linalg.norm(model.fields["U"][..., i0], axis=1), np.nan)
+P = np.where(aw > 0.95, model.fields["p"][..., i0], np.nan)
+P_rgh = np.where(aw > 0.95, model.fields["p_rgh"][..., i0], np.nan)
+
+with plt.rc_context(
+    {
+        "axes.prop_cycle": cycler(
+            color=np.linspace(0, 1, i0.size + 1)[:-1].astype("U")
+        ),
+        "axes.grid": True,
+        "axes.xmargin": 0,
+    }
+):
+    fig, ax = plt.subplots(3, constrained_layout=True)
+    ax1, ax2, ax3 = ax
+
+    ha = ax1.plot(model.t, U, lw=1)
+    ax1.set(xlabel="t (s)", ylabel="U (m/s)")
+
+    ax2.plot(model.t, P*1e-3, lw=1)
+    ax2.set(xlabel="t (s)", ylabel="P (kPa)")
+
+    ax3.plot(model.t, P_rgh*1e-3, lw=1)
+    ax3.set(xlabel="t (s)", ylabel="P_rgh (kPa)")
+
+    for a in ax:
+        a.set(ylim=0)
+
+    ax2.legend(
+        ha,
+        list(
+            starmap(lambda x, z: f"x={x:8}m; z={z:8}m", zip(model.x[i0], model.z[i0]))
+        ),
+        bbox_to_anchor=(1.05, .5),
+        loc="center left"
+    )
+
+    plt.show()