1
Fork 0
internship/olaflow/processing/animate.py
2022-04-13 15:10:41 +02:00

126 lines
3.4 KiB
Python

import argparse
import configparser
import logging
import multiprocessing as mp
import pathlib
import pickle
import matplotlib.pyplot as plt
import matplotlib.animation as animation
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("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)
with out.open("rb") as f:
model = pickle.load(f)
x0 = config.getfloat("post", "x")
z0 = config.getfloat("post", "z")
i0 = np.argmin(np.abs((model.x - x0) + 1j * (model.z - z0)))
X, Z = np.meshgrid(np.unique(model.x), np.unique(model.z))
C = np.where(
(model.x[:, None, None].astype(np.single) == X[None, :, :].astype(np.single))
& (model.z[:, None, None].astype(np.single) == Z[None, :, :].astype(np.single))
)
P = np.full((model.t.size, *X.shape), np.nan)
P[:, C[1], C[2]] = model.fields["porosity"][:, C[0]]
AW = np.full((model.t.size, *X.shape), np.nan)
AW[:, C[1], C[2]] = model.fields["alpha.water"][:, C[0]]
U = np.full((model.t.size, *X.shape), np.nan)
U[:, C[1], C[2]] = np.linalg.norm(model.fields["U"], axis=1)[:, C[0]]
fig, ax = plt.subplots(figsize=(19.2, 10.8), dpi=100)
tit = ax.text(
0.5,
0.95,
f"t={model.t[0]}s",
horizontalalignment="center",
verticalalignment="top",
transform=ax.transAxes,
)
aw_m = ax.pcolormesh(X, Z, AW[0], vmin=0, vmax=1, cmap="Blues", zorder=1)
ax.pcolormesh(
X,
Z,
P[1],
vmin=0,
vmax=1,
cmap="Greys_r",
alpha=(np.nan_to_num(1 - P[1]) / 2).clip(0, 1),
zorder=1.1,
)
ax.axhline(4.5, ls="-.", lw=1, c="k", alpha=0.2, zorder=1.2)
fig.colorbar(aw_m)
ax.set(xlabel="x (m)", ylabel="z (m)", aspect="equal", facecolor="#bebebe")
ax.grid(c="k", alpha=0.2)
def anim(i):
tit.set_text(f"t={model.t[i]}s")
aw_m.set_array(AW[i])
figU, axU = plt.subplots(figsize=(19.2, 10.8), dpi=100)
u_m = axU.pcolormesh(
X, Z, U[0], cmap="BuPu", vmin=0, vmax=np.nanquantile(U, .99), zorder=1, alpha=np.nan_to_num(AW[0]).clip(0, 1)
)
ur_m = axU.pcolormesh(
X, Z, U[0], cmap="YlOrBr", vmin=0, vmax=np.nanquantile(U, .99), zorder=1, alpha=1-np.nan_to_num(AW[0]).clip(0, 1)
)
# aw_u = axU.contour(X, Z, AW[0], levels=(.5,))
figU.colorbar(u_m)
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,
)
def animU(i):
titU.set_text(f"t={model.t[i]}s")
u_m.set_array(U[i])
u_m.set_alpha(np.nan_to_num(AW[i]).clip(0, 1))
ur_m.set_array(U[i])
ur_m.set_alpha(1-np.nan_to_num(AW[i]).clip(0, 1))
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()