Olaflow processing for article

olaflow/processing/snap.py

@@ -0,0 +1,79 @@
+import argparse
+import gzip
+import logging
+import multiprocessing as mp
+import pathlib
+import pickle
+
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+from matplotlib.gridspec import GridSpec
+from matplotlib.ticker import MultipleLocator
+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(
+    "-o",
+    "--output",
+    type=pathlib.Path,
+    help="Output directory for pickled data",
+    required=True,
+)
+parser.add_argument(
+    "-m",
+    "--max",
+    help="Only compute maximum rather than animation",
+    action="store_true",
+)
+parser.add_argument(
+    "-i",
+    "--initial",
+    help="Only compute initial domain",
+    action="store_true",
+)
+args = parser.parse_args()
+
+logging.basicConfig(level=max((10, 20 - 10 * args.verbose)))
+log = logging.getLogger("ola_post")
+
+log.info("Animating olaFlow output")
+out = args.output
+out.mkdir(parents=True, exist_ok=True)
+
+with (
+    path.open("rb")
+    if (path := out.joinpath("pickle")).exists()
+    else gzip.open(path.with_suffix(".gz"), "rb")
+) as f:
+    model = pickle.load(f)
+
+x0, idx0 = np.unique(model.x.astype(np.half), return_inverse=True)
+z0, idz0 = np.unique(model.z.astype(np.half), return_inverse=True)
+
+ix0 = np.argsort(x0)
+iz0 = np.argsort(z0)[::-1]
+
+X, Z = np.meshgrid(x0, z0)
+
+P = np.full((model.t.size, *X.shape), np.nan)
+P[:, iz0[idz0], ix0[idx0]] = model.fields["porosity"]
+
+AW = np.full((model.t.size, *X.shape), np.nan)
+AW[:, iz0[idz0], ix0[idx0]] = model.fields["alpha.water"]
+
+#U = np.full((model.t.size, *X.shape), np.nan)
+#U[:, iz0[idz0], ix0[idx0]] = np.linalg.norm(model.fields["U"], axis=1)
+
+i0 = np.argmin(np.abs(model.t[:, None] - np.asarray((102, 118, 144.5, 176.5))[None, :]), axis=0) 
+
+fig, ax_ = plt.subplots(
+    2, 2, figsize=(15 / 2.54, 4 / 2.54), dpi=200, constrained_layout=True
+)
+for ax, i in zip(ax_.flatten(), i0):
+    ax.imshow(AW[i], cmap="Blues", vmin=0, vmax=1)
+
+fig.savefig(out.joinpath("snap.pdf"))