From 402ad87091b2f9c3a2ca02fa578a749337fde1ca Mon Sep 17 00:00:00 2001 From: "Edgar P. Burkhart" Date: Mon, 9 May 2022 11:23:09 +0200 Subject: [PATCH] Update 'diff' with U --- olaflow/processing/diff.py | 84 +++++++++++++++++++++++++++++--------- 1 file changed, 64 insertions(+), 20 deletions(-) diff --git a/olaflow/processing/diff.py b/olaflow/processing/diff.py index c8788b1..c9a2454 100644 --- a/olaflow/processing/diff.py +++ b/olaflow/processing/diff.py @@ -5,6 +5,7 @@ import logging from multiprocessing import pool import pathlib import pickle +import sys from cycler import cycler import matplotlib.pyplot as plt @@ -36,6 +37,15 @@ parser.add_argument( 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() @@ -57,35 +67,69 @@ def get_pickle(out): models = list(map(get_pickle, args.output)) -fig, ax = plt.subplots(len(models), figsize=(6, 1.5 * len(models)), dpi=100, constrained_layout=True) -if args.timestep is None: - for i, (_ax, _model) in enumerate(zip(ax, models)): - _ax.contour( - _model.t, - _model.post_fields[args.func]["x_alpha.water"], - _model.post_fields[args.func]["alpha.water"].T, - (0.5,), - colors="k", - ) +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() + _ax.grid(color="k", alpha=0.2) fig.savefig( args.output[0].joinpath( - f"diff_{args.func}_{'_'.join([o.name for o in args.output])}.pdf" + f"diff_{args.func}_{args.field}_{'_'.join([o.name for o in args.output])}.pdf" ) ) else: - for i, (_ax, _model) in enumerate(zip(ax, models)): - _ax.tricontour( - _model.x, - _model.z, - _model.fields["alpha.water"][np.where(_model.t == args.timestep)[0]][0], - levels=(0.5,), - colors="k", - ) + 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()