This repository has been archived on 2025-05-05. You can view files and clone it, but you cannot make any changes to it's state, such as pushing and creating new issues, pull requests or comments.
internship/olaflow/processing/animate.py

147 lines
4 KiB
Python

import argparse
import configparser
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("-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 = config.getfloat("post", "x")
z0 = config.getfloat("post", "z")
i0 = np.argmin(np.abs((model.x - x0) + 1j * (model.z - z0)))
x0, idx0 = np.unique(model.x.astype(np.half), return_inverse=True)
z0, idz0 = np.unique(model.z.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"]
AW = np.full((model.t.size, *X.shape), np.nan)
AW[:, idz0, idx0] = model.fields["alpha.water"]
U = np.full((model.t.size, *X.shape), np.nan)
U[:, idz0, idx0] = np.linalg.norm(model.fields["U"], axis=1)
fig = plt.figure(figsize=(19.2, 10.8), dpi=100)
gs = GridSpec(3, 1, figure=fig, height_ratios=[1, .05, .05])
ax = fig.add_subplot(gs[0])
cax1 = fig.add_subplot(gs[1])
cax2 = fig.add_subplot(gs[2])
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)
p_m = 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, label=r"$\alpha_w$", cax=cax1, shrink=0.6, orientation="horizontal")
fig.colorbar(p_m, label=r"Porosity", cax=cax2, shrink=0.6, orientation="horizontal")
ax.set(xlabel="x (m)", ylabel="z (m)", aspect="equal", facecolor="#000000")
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 = plt.figure(figsize=(19.2, 10.8), dpi=100)
gsU = GridSpec(3, 1, figure=figU, height_ratios=[1, .05, .05])
axU = figU.add_subplot(gsU[0])
caxu1 = figU.add_subplot(gsU[1])
caxu2 = figU.add_subplot(gsU[2])
u_m = axU.pcolormesh(
X,
Z,
U[0],
cmap="BuPu",
vmin=0,
vmax=np.nanquantile(U, 0.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, 0.99),
zorder=1,
alpha=1 - np.nan_to_num(AW[0]).clip(0, 1),
)
# aw_u = axU.contour(X, Z, AW[0], levels=(.5,))
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_w$", cax=caxu1, shrink=0.6, orientation="horizontal")
figU.colorbar(ur_m, label=r"$U_a$", cax=caxu2, shrink=0.6, orientation="horizontal")
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)
ani.save(out.joinpath("anim.mp4"), fps=24)
aniU.save(out.joinpath("animU.mp4"), fps=24)