Add max graph to animate

olaflow/processing/animate.py

@@ -22,6 +22,12 @@ parser.add_argument(
     help="Output directory for pickled data",
     required=True,
 )
+parser.add_argument(
+    "-m",
+    "--max",
+    help="Only compute maximum rather than animation",
+    action="store_true",
+)
 args = parser.parse_args()
 
 logging.basicConfig(level=max((10, 20 - 10 * args.verbose)))
@@ -41,37 +47,38 @@ with (
 x0, idx0 = np.unique(model.x.astype(np.half), return_inverse=True)
 z0, idz0 = np.unique(model.z.astype(np.half), return_inverse=True)
 
+ix0 = np.argsort(x0)
+iz0 = np.argsort(z0)[::-1]
+
 X, Z = np.meshgrid(x0, z0)
 
 P = np.full((model.t.size, *X.shape), np.nan)
-P[:, idz0, idx0] = model.fields["porosity"]
+P[:, iz0[idz0], ix0[idx0]] = model.fields["porosity"]
 
 AW = np.full((model.t.size, *X.shape), np.nan)
-AW[:, idz0, idx0] = model.fields["alpha.water"]
+AW[:, iz0[idz0], ix0[idx0]] = model.fields["alpha.water"]
 
 U = np.full((model.t.size, *X.shape), np.nan)
-U[:, idz0, idx0] = np.linalg.norm(model.fields["U"], axis=1)
+U[:, iz0[idz0], ix0[idx0]] = np.linalg.norm(model.fields["U"], axis=1)
 
 fig = plt.figure(figsize=(19.2, 10.8), dpi=100)
 gs = GridSpec(3, 1, figure=fig, height_ratios=[1, 0.05, 0.05])
 ax = fig.add_subplot(gs[0])
 cax1 = fig.add_subplot(gs[1])
 cax2 = fig.add_subplot(gs[2])
-tit = ax.text(
+aw_m = ax.imshow(
-    0.5,
+    AW[0],
-    0.95,
+    vmin=0,
-    f"t={model.t[0]}s",
+    vmax=1,
-    horizontalalignment="center",
+    extent=(x0.min(), x0.max(), z0.min(), z0.max()),
-    verticalalignment="top",
+    cmap="Blues",
-    transform=ax.transAxes,
+    zorder=1,
 )
-aw_m = ax.pcolormesh(X, Z, AW[0], vmin=0, vmax=1, cmap="Blues", zorder=1)
+p_m = ax.imshow(
-p_m = ax.pcolormesh(
-    X,
-    Z,
     P[1],
     vmin=0,
     vmax=1,
+    extent=(x0.min(), x0.max(), z0.min(), z0.max()),
     cmap="Greys_r",
     alpha=(np.nan_to_num(1 - P[1]) / 2).clip(0, 1),
     zorder=1.1,
@@ -85,62 +92,84 @@ ax.set(xlabel="x (m)", ylabel="z (m)", aspect="equal", facecolor="#000000")
 ax.grid(c="k", alpha=0.2)
 
 
-def anim(i):
-    tit.set_text(f"t={model.t[i]}s")
-    aw_m.set_array(AW[i])
-
-
 figU = plt.figure(figsize=(19.2, 10.8), dpi=100)
-gsU = GridSpec(3, 1, figure=figU, height_ratios=[1, 0.05, 0.05])
+gsU = GridSpec(
+    2 if args.max else 3,
+    1,
+    figure=figU,
+    height_ratios=[1, 0.05] if args.max else [1, 0.05, 0.05],
+)
 axU = figU.add_subplot(gsU[0])
 caxu1 = figU.add_subplot(gsU[1])
-caxu2 = figU.add_subplot(gsU[2])
+if not args.max:
-u_m = axU.pcolormesh(
+    caxu2 = figU.add_subplot(gsU[2])
-    X,
+u_m = axU.imshow(
-    Z,
     U[0],
     cmap="BuPu",
     vmin=0,
-    vmax=np.nanmax(U),
+    vmax=np.nanmax(np.where(AW > 0.5, U, np.nan), initial=0),
+    extent=(x0.min(), x0.max(), z0.min(), z0.max()),
     zorder=1,
     alpha=np.nan_to_num(AW[0]).clip(0, 1),
 )
-ur_m = axU.pcolormesh(
+ur_m = axU.imshow(
-    X,
-    Z,
     U[0],
     cmap="YlOrBr",
     vmin=0,
-    vmax=np.nanmax(U),
+    vmax=np.nanmax(np.where(AW > 0.5, U, np.nan), initial=0),
+    extent=(x0.min(), x0.max(), z0.min(), z0.max()),
     zorder=1,
     alpha=1 - np.nan_to_num(AW[0]).clip(0, 1),
 )
 # aw_u = axU.contour(X, Z, AW[0], levels=(.5,))
 axU.set(xlabel="x (m)", ylabel="z (m)", aspect="equal", facecolor="#bebebe")
 axU.grid(c="k", alpha=0.2)
-titU = axU.text(
-    0.5,
-    0.95,
-    f"t={model.t[0]}s",
-    horizontalalignment="center",
-    verticalalignment="top",
-    transform=axU.transAxes,
-)
-
 figU.colorbar(u_m, label=r"$U_w$", cax=caxu1, shrink=0.6, orientation="horizontal")
-figU.colorbar(ur_m, label=r"$U_a$", cax=caxu2, shrink=0.6, orientation="horizontal")
 
 
-def animU(i):
+if args.max:
-    titU.set_text(f"t={model.t[i]}s")
+    aw_m.set_array(AW.max(axis=0))
-    u_m.set_array(U[i])
-    u_m.set_alpha(np.nan_to_num(AW[i]).clip(0, 1))
-    ur_m.set_array(U[i])
-    ur_m.set_alpha(1 - np.nan_to_num(AW[i]).clip(0, 1))
 
+    u_m.set_array(np.nanmax(np.where(AW > 0.5, U, np.nan), axis=0, initial=0))
+    u_m.set_alpha(1)
+    u_m.set_cmap("hot_r")
+    ur_m.remove()
 
-ani = animation.FuncAnimation(fig, anim, frames=model.t.size, interval=1 / 24)
+    fig.savefig(out.joinpath("max_aw.pdf"))
-aniU = animation.FuncAnimation(figU, animU, frames=model.t.size, interval=1 / 24)
+    figU.savefig(out.joinpath("max_U.pdf"))
+else:
+    figU.colorbar(ur_m, label=r"$U_a$", cax=caxu2, shrink=0.6, orientation="horizontal")
 
-ani.save(out.joinpath("anim.mp4"), fps=24)
+    tit = ax.text(
-aniU.save(out.joinpath("animU.mp4"), fps=24)
+        0.5,
+        0.95,
+        f"t={model.t[0]}s",
+        horizontalalignment="center",
+        verticalalignment="top",
+        transform=ax.transAxes,
+    )
+    titU = axU.text(
+        0.5,
+        0.95,
+        f"t={model.t[0]}s",
+        horizontalalignment="center",
+        verticalalignment="top",
+        transform=axU.transAxes,
+    )
+
+    def anim(i):
+        tit.set_text(f"t={model.t[i]}s")
+        aw_m.set_array(AW[i])
+
+    def animU(i):
+        titU.set_text(f"t={model.t[i]}s")
+        u_m.set_array(U[i])
+        u_m.set_alpha(np.nan_to_num(AW[i]).clip(0, 1))
+        ur_m.set_array(U[i])
+        ur_m.set_alpha(1 - np.nan_to_num(AW[i]).clip(0, 1))
+
+    ani = animation.FuncAnimation(fig, anim, frames=model.t.size, interval=1 / 24)
+    aniU = animation.FuncAnimation(figU, animU, frames=model.t.size, interval=1 / 24)
+
+    ani.save(out.joinpath("anim.mp4"), fps=24)
+    aniU.save(out.joinpath("animU.mp4"), fps=24)