Update colormap and limits for animate & diff

olaflow/processing/animate.py

@@ -107,7 +107,7 @@ u_m = axU.imshow(
     U[0],
     cmap="BuPu",
     vmin=0,
-    vmax=np.nanmax(np.where(AW > 0.5, U, np.nan), initial=0),
+    vmax=20,
     extent=(x0.min(), x0.max(), z0.min(), z0.max()),
     zorder=1,
     alpha=np.nan_to_num(AW[0]).clip(0, 1),
@@ -116,7 +116,7 @@ ur_m = axU.imshow(
     U[0],
     cmap="YlOrBr",
     vmin=0,
-    vmax=np.nanmax(np.where(AW > 0.5, U, np.nan), initial=0),
+    vmax=20,
     extent=(x0.min(), x0.max(), z0.min(), z0.max()),
     zorder=1,
     alpha=1 - np.nan_to_num(AW[0]).clip(0, 1),

olaflow/processing/diff.py

@@ -68,13 +68,14 @@ def get_pickle(out):
 models = list(map(get_pickle, args.output))
 
 
-fig, (ax,) = plt.subplots(
+fig, ax_ = plt.subplots(
     len(models),
     figsize=(6, 1.5 * len(models)),
     dpi=100,
     constrained_layout=True,
     squeeze=False,
 )
+ax = ax_[:, 0]
 
 if args.timestep is None:
     match args.field:
@@ -90,9 +91,10 @@ if args.timestep is None:
         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,
+                    np.where(_model.post_fields[args.func]["alpha.water"] > 0.5, np.linalg.norm(_model.post_fields[args.func][args.field], axis=2), np.nan).T[::-1],
                     vmin=0,
-                    cmap="inferno",
+                    vmax=20,
+                    cmap="inferno_r",
                     extent=(
                         _model.t.min(),
                         _model.t.max(),