diff --git a/olaflow/processing/animate.py b/olaflow/processing/animate.py index adcd43b..78da9f9 100644 --- a/olaflow/processing/animate.py +++ b/olaflow/processing/animate.py @@ -1,6 +1,7 @@ import argparse import configparser import logging +import multiprocessing as mp import pathlib import pickle @@ -41,9 +42,13 @@ C = np.where( P = np.full((model.t.size, *X.shape), np.nan) P[:, C[1], C[2]] = model.fields["porosity"][:, C[0]] + AW = np.full((model.t.size, *X.shape), np.nan) AW[:, C[1], C[2]] = model.fields["alpha.water"][:, C[0]] +U = np.full((model.t.size, *X.shape), np.nan) +U[:, C[1], C[2]] = np.linalg.norm(model.fields["U"], axis=1)[:, C[0]] + fig, ax = plt.subplots(figsize=(19.2, 10.8), dpi=100) tit = ax.text( 0.5, @@ -61,20 +66,61 @@ ax.pcolormesh( vmin=0, vmax=1, cmap="Greys_r", - alpha=np.nan_to_num(1 - P[1])/2, + alpha=(np.nan_to_num(1 - P[1]) / 2).clip(0, 1), zorder=1.1, ) ax.axhline(4.5, ls="-.", lw=1, c="k", alpha=0.2, zorder=1.2) -def anim(i): - tit.set_text(f"t={model.t[i]}s") - aw_m.set_array(AW[i]) - return (aw_m,) - - fig.colorbar(aw_m) ax.set(xlabel="x (m)", ylabel="z (m)", aspect="equal", facecolor="#bebebe") ax.grid(c="k", alpha=0.2) -ani = animation.FuncAnimation(fig, anim, frames=model.t.size) -ani.save(out.parent.joinpath("anim.mp4"), fps=24) + + +def anim(i): + tit.set_text(f"t={model.t[i]}s") + aw_m.set_array(AW[i]) + + + +figU, axU = plt.subplots(figsize=(19.2, 10.8), dpi=100) +u_m = axU.pcolormesh( + X, Z, U[0], cmap="BuPu", vmin=0, vmax=np.nanquantile(U, .99), zorder=1, alpha=np.nan_to_num(AW[0]).clip(0, 1) +) +ur_m = axU.pcolormesh( + X, Z, U[0], cmap="YlOrBr", vmin=0, vmax=np.nanquantile(U, .99), zorder=1, alpha=1-np.nan_to_num(AW[0]).clip(0, 1) +) +# aw_u = axU.contour(X, Z, AW[0], levels=(.5,)) +figU.colorbar(u_m) +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, +) + + +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) + +a = mp.Process(target=ani.save, args=(out.parent.joinpath("anim.mp4"),), kwargs={"fps": 24}) +aU = mp.Process(target=aniU.save, args=(out.parent.joinpath("animU.mp4"),), kwargs={"fps": 24}) + +a.start() +aU.start() + +a.join() +aU.join() +#plt.show()