Add options to diff

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,12 +57,13 @@ 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)):
+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["graphUniform"]["x_alpha.water"],
-        _model.post_fields["graphUniform"]["alpha.water"].T,
+            _model.post_fields[args.func]["x_alpha.water"],
+            _model.post_fields[args.func]["alpha.water"].T,
             (0.5,),
             colors="k",
         )
@@ -56,6 +71,27 @@ for i, (_ax, _model) in enumerate(zip(ax, models)):
         _ax.set(xlabel="t (s)", 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_{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="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"
+        )
+    )