diff --git a/olaflow/processing/animate.py b/olaflow/processing/animate.py index 32fb8c7..e1a79a4 100644 --- a/olaflow/processing/animate.py +++ b/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( - 0.5, - 0.95, - f"t={model.t[0]}s", - horizontalalignment="center", - verticalalignment="top", - transform=ax.transAxes, +aw_m = ax.imshow( + AW[0], + vmin=0, + vmax=1, + extent=(x0.min(), x0.max(), z0.min(), z0.max()), + cmap="Blues", + zorder=1, ) -aw_m = ax.pcolormesh(X, Z, AW[0], vmin=0, vmax=1, cmap="Blues", zorder=1) -p_m = ax.pcolormesh( - X, - Z, +p_m = ax.imshow( 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]) -u_m = axU.pcolormesh( - X, - Z, +if not args.max: + caxu2 = figU.add_subplot(gsU[2]) +u_m = axU.imshow( 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( - X, - Z, +ur_m = axU.imshow( 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): - 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)) +if args.max: + aw_m.set_array(AW.max(axis=0)) + 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) -aniU = animation.FuncAnimation(figU, animU, frames=model.t.size, interval=1 / 24) + fig.savefig(out.joinpath("max_aw.pdf")) + 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) -aniU.save(out.joinpath("animU.mp4"), fps=24) + tit = ax.text( + 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)