import argparse import gzip from itertools import starmap import logging from multiprocessing import pool import pathlib import pickle import sys from cycler import cycler import matplotlib.pyplot as plt import matplotlib.gridspec as 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", action="append", type=pathlib.Path, help="Post-processing directory", required=True, ) parser.add_argument( "-t", "--timestep", type=float, help="Time-step to compare", ) parser.add_argument( "-f", "--func", type=str, help="Post-process function to compare", default="graphUniform", choices=("graphUniform", "graphUniform2"), ) parser.add_argument( "-y", "--field", type=str, help="Field to compare", default="alpha.water", choices=("alpha.water", "U"), ) args = parser.parse_args() logging.basicConfig(level=max((10, 20 - 10 * args.verbose))) log = logging.getLogger("ola_post") log.info("Plotting comparison of model output") def get_pickle(out): with ( path.open("rb") if (path := out.joinpath("pickle")).exists() else gzip.open(path.with_suffix(".gz"), "rb") ) as f: return pickle.load(f) models = list(map(get_pickle, args.output)) fig, (ax,) = plt.subplots( len(models), figsize=(6, 1.5 * len(models)), dpi=100, constrained_layout=True, squeeze=False, ) if args.timestep is None: match args.field: case "alpha.water": for i, (_ax, _model) in enumerate(zip(ax, models)): _ax.contour( _model.t, _model.post_fields[args.func][f"x_{args.field}"], _model.post_fields[args.func][args.field].T, (0.5,), colors="k", ) case "U": for i, (_ax, _model) in enumerate(zip(ax, models)): _c = _ax.imshow( np.linalg.norm(_model.post_fields[args.func][args.field], axis=2).T, vmin=0, cmap="inferno", extent=( _model.t.min(), _model.t.max(), _model.post_fields[args.func][f"x_{args.field}"].min(), _model.post_fields[args.func][f"x_{args.field}"].max(), ), ) fig.colorbar(_c, label=f"{args.field} (m/s)", ax=_ax) case _: log.error(f"Cannot plot field {args.field} from {args.func}") sys.exit(1) for i, (_ax, _model) in enumerate(zip(ax, models)): _ax.set(xlabel="t (s)", ylabel="z (m)", title=f"Case {i}") _ax.grid(color="k", alpha=0.2) fig.savefig( args.output[0].joinpath( f"diff_{args.func}_{args.field}_{'_'.join([o.name for o in args.output])}.pdf" ) ) else: match args.field: case "alpha.water": for i, (_ax, _model) in enumerate(zip(ax, models)): _ax.tricontour( _model.x, _model.z, _model.fields[args.field][np.where(_model.t == args.timestep)[0]][0], levels=(0.5,), colors="k", ) case _: log.error(f"Cannot plot field {args.field} from {args.func} at timestep") sys.exit(1) for i, (_ax, _model) in enumerate(zip(ax, models)): _ax.set(xlabel="x (m)", ylabel="z (m)", title=f"Case {i}") _ax.grid() fig.savefig( args.output[0].joinpath( f"diff_t{args.timestep}_{'_'.join([o.name for o in args.output])}.pdf" ) )