Animate U

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()