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

105 lines
2.7 KiB
Python
Raw Permalink Normal View History

2022-04-20 13:54:22 +02:00
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
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,
)
2022-04-20 13:54:22 +02:00
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
2022-04-20 13:54:22 +02:00
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)
flt = np.where((model.x >= -60) & (model.x <= -20) & (model.z >= 0) & (model.z <= 10))[
0
]
x0, idx0 = np.unique(model.x[flt].astype(np.half), return_inverse=True)
z0, idz0 = np.unique(model.z[flt].astype(np.half), return_inverse=True)
X, Z = np.meshgrid(x0, z0)
P = np.full((model.t.size, *X.shape), np.nan)
P[:, idz0, idx0] = model.fields["porosity"][:, flt]
AW = np.full((model.t.size, *X.shape), np.nan)
AW[:, idz0, idx0] = model.fields["alpha.water"][:, flt]
watl = z0[np.argmax((AW > 0.5)[:, ::-1, :], axis=1)]
U = np.full((model.t.size, 2, *X.shape), np.nan)
UU = np.full((model.t.size, *X.shape), np.nan)
U[..., idz0, idx0] = model.fields["U"][..., flt][:, (0, 2)]
UU[..., idz0, idx0] = np.linalg.norm(model.fields["U"][..., flt], axis=1)
figU = plt.figure(figsize=(19.2, 10.8), dpi=100)
gsU = GridSpec(2, 1, figure=figU, height_ratios=[1, 0.05])
axU = figU.add_subplot(gsU[0])
caxu1 = figU.add_subplot(gsU[1])
# caxu2 = figU.add_subplot(gsU[2])
alp = np.clip(np.nan_to_num(AW), 0, 1)
axU.pcolormesh(X, Z, P[1], vmin=0, vmax=1, cmap="Greys_r")
u_m = axU.quiver(
X,
Z,
*U[0],
UU[0],
alpha=alp[0],
2022-05-09 12:53:36 +02:00
cmap="inferno_r",
clim=(0, np.nanmax(UU)),
2022-04-20 13:54:22 +02:00
)
# (wat_p,) = axU.plot(x0, watl[0])
axU.set(xlabel="x (m)", ylabel="z (m)", aspect="equal", facecolor="#bebebe")
axU.grid(c="k", alpha=0.2)
titU = axU.text(
0.5,
0.95,
f"t={model.t[0]}s",
horizontalalignment="center",
verticalalignment="top",
transform=axU.transAxes,
)
figU.colorbar(u_m, label=r"$U$", cax=caxu1, shrink=0.6, orientation="horizontal")
def animU(i):
titU.set_text(f"t={model.t[i]}s")
u_m.set_UVC(*U[i], UU[i])
u_m.set_alpha(alp[i])
# wat_p.set_ydata(watl[i])
aniU = animation.FuncAnimation(figU, animU, frames=model.t.size, interval=1 / 24)
aniU.save(out.joinpath("animUzoom.mp4"), fps=24)