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 from matplotlib.ticker import MultipleLocator 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, ) parser.add_argument( "-m", "--max", help="Only compute maximum rather than animation", action="store_true", ) parser.add_argument( "-i", "--initial", help="Only compute initial domain", action="store_true", ) 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 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, idx0 = np.unique(model.x.astype(np.half), return_inverse=True) z0, idz0 = np.unique(model.z.astype(np.half), return_inverse=True) ix0 = np.argsort(x0) iz0 = np.argsort(z0)[::-1] X, Z = np.meshgrid(x0, z0) P = np.full((model.t.size, *X.shape), np.nan) P[:, iz0[idz0], ix0[idx0]] = model.fields["porosity"] AW = np.full((model.t.size, *X.shape), np.nan) AW[:, iz0[idz0], ix0[idx0]] = model.fields["alpha.water"] U = np.full((model.t.size, *X.shape), np.nan) U[:, iz0[idz0], ix0[idx0]] = np.linalg.norm(model.fields["U"], axis=1) fig = plt.figure( figsize=(15 / 2.54, 4 / 2.54), dpi=200, constrained_layout=True ) gs = GridSpec( 3 if not args.initial else 1, 1, figure=fig, height_ratios=[1, 0.1, 0.1] if not args.initial else [1], ) ax = fig.add_subplot(gs[0]) if not args.initial: cax1 = fig.add_subplot(gs[1]) cax2 = fig.add_subplot(gs[2]) aw_m = ax.imshow( AW[0], vmin=0, vmax=1, extent=(x0.min(), x0.max(), z0.min(), z0.max()), cmap="Blues", zorder=1, ) p_m = ax.imshow( P[1], vmin=0, vmax=1, extent=(x0.min(), x0.max(), z0.min(), z0.max()), 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) if not args.initial: 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) ax.xaxis.set_minor_locator(MultipleLocator(5)) ax.yaxis.set_minor_locator(MultipleLocator(5)) figU = plt.figure() gsU = GridSpec( 2 if args.max else 3, 1, figure=figU, height_ratios=[1, 0.05] if args.max else [1, 0.05, 0.05], ) axU = figU.add_subplot(gsU[0]) caxu1 = figU.add_subplot(gsU[1]) if not args.max: caxu2 = figU.add_subplot(gsU[2]) u_m = axU.imshow( U[0], cmap="BuPu", vmin=0, vmax=20, extent=(x0.min(), x0.max(), z0.min(), z0.max()), zorder=1, alpha=np.nan_to_num(AW[0]).clip(0, 1), ) ur_m = axU.imshow( U[0], cmap="YlOrBr", vmin=0, vmax=20, extent=(x0.min(), x0.max(), z0.min(), z0.max()), 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) figU.colorbar(u_m, label=r"$U_w$", cax=caxu1, shrink=0.6, orientation="horizontal") if args.max: aw_m.set_array(AW.max(axis=0)) u_m.set_array(np.nanmax(np.where(AW > 0.5, U, np.nan), axis=0, initial=0)) u_m.set_alpha(1) u_m.set_cmap("hot_r") ur_m.remove() fig.savefig(out.joinpath("max_aw.pdf")) figU.savefig(out.joinpath("max_U.pdf")) elif args.initial: aw_m.set_array(AW[0]) ax.vlines(-20, -15, 15, color="k", lw=1, ls="--", label="Measurements") ax.text(-20, 15, "Measurements", ha="right", va="bottom") fig.savefig(out.joinpath("aw_t0.pdf")) fig.savefig(out.joinpath("aw_t0.jpg"), dpi=200) else: fig.set(figwidth=19.2, figheight=10.8, dpi=100) figU.set(figwidth=19.2, figheight=10.8, dpi=100) figU.colorbar(ur_m, label=r"$U_a$", cax=caxu2, shrink=0.6, orientation="horizontal") tit = ax.text( 0.5, 0.95, f"t={model.t[0]}s", horizontalalignment="center", verticalalignment="top", transform=ax.transAxes, ) titU = axU.text( 0.5, 0.95, f"t={model.t[0]}s", horizontalalignment="center", verticalalignment="top", transform=axU.transAxes, ) def anim(i): tit.set_text(f"t={model.t[i]}s") aw_m.set_array(AW[i]) 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)