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)