From 71a897cf58e91b66a74af8c1d93bf26e3e2a48f4 Mon Sep 17 00:00:00 2001
From: "Edgar P. Burkhart" <git@edgarpierre.fr>
Date: Fri, 6 May 2022 11:14:05 +0200
Subject: [PATCH] Add options to diff

---
 olaflow/processing/diff.py | 62 ++++++++++++++++++++++++++++++--------
 1 file changed, 49 insertions(+), 13 deletions(-)

diff --git a/olaflow/processing/diff.py b/olaflow/processing/diff.py
index 6cf0025..6f41763 100644
--- a/olaflow/processing/diff.py
+++ b/olaflow/processing/diff.py
@@ -24,6 +24,20 @@ parser.add_argument(
     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",
+)
+
 args = parser.parse_args()
 
 logging.basicConfig(level=max((10, 20 - 10 * args.verbose)))
@@ -43,19 +57,41 @@ def get_pickle(out):
 
 models = list(map(get_pickle, args.output))
 
-fig, ax = plt.subplots(len(models), constrained_layout=True)
-for i, (_ax, _model) in enumerate(zip(ax, models)):
-    _ax.contour(
-        _model.t,
-        _model.post_fields["graphUniform"]["x_alpha.water"],
-        _model.post_fields["graphUniform"]["alpha.water"].T,
-        (0.5,),
-        colors="k",
+if args.timestep is None:
+    fig, ax = plt.subplots(len(models), constrained_layout=True)
+    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",
+        )
+
+        _ax.set(xlabel="t (s)", ylabel="z (m)", title=f"Case {i}")
+        _ax.grid()
+
+    fig.savefig(
+        args.output[0].joinpath(
+            f"diff_{args.func}_{'_'.join([o.name for o in args.output])}.pdf"
+        )
     )
+else:
+    fig, ax = plt.subplots(len(models), constrained_layout=True)
+    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",
+        )
 
-    _ax.set(xlabel="t (s)", ylabel="z (m)", title=f"Case {i}")
-    _ax.grid()
+        _ax.set(xlabel="x (m)", ylabel="z (m)", title=f"Case {i}")
+        _ax.grid()
 
-fig.savefig(
-    args.output[0].joinpath(f"diff_{'_'.join([o.name for o in args.output])}.pdf")
-)
+    fig.savefig(
+        args.output[0].joinpath(
+            f"diff_t{args.timestep}_{'_'.join([o.name for o in args.output])}.pdf"
+        )
+    )