1
Fork 0
internship/olaflow/processing/post.py

90 lines
2.4 KiB
Python

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 (
path.open("rb")
if (path := out.joinpath("pickle")).exists()
else gzip.open(path.with_suffix(".gz"), "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, 0.5),
loc="center left",
)
fig.savefig(out.joinpath("fig.pdf"))