Update 'diff' with U

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()