import argparse import configparser import logging import pathlib import matplotlib.animation as animation import matplotlib.pyplot as plt import numpy as np <<<<<<< HEAD ======= import pandas as pd >>>>>>> olaflow parser = argparse.ArgumentParser(description="Animate swash output") parser.add_argument("-v", "--verbose", action="count", default=0) args = parser.parse_args() logging.basicConfig(level=max((10, 20 - 10 * args.verbose))) log = logging.getLogger("post") log.info("Starting post-processing") config = configparser.ConfigParser() config.read("config.ini") inp = pathlib.Path(config.get("post", "inp")) root = pathlib.Path(config.get("swash", "out")) out = pathlib.Path(config.get("plot", "out")) out.mkdir(exist_ok=True) def data(var): return np.load(inp.joinpath(f"{var}.npy")) x = data("xp") t = data("tsec") watl = data("watl") botl = data("botl") zk = data("zk") velk = data("velk") vz = data("vz") wl = np.maximum(watl, -botl) # print(x.size, -np.arange(0, 1 * bathy.hstru.size, 1)[::-1].size) fig, ax = plt.subplots() # ax.plot(x, -botl, c="k") # ax.fill_between( # x, -botl, -data["botl"] + bathy.hstru, color="k", alpha=0.2 # ) n = 0 vk = np.sqrt((velk[n] ** 2).sum(axis=1)) # print(vk.shape) # plt.imshow(vk) # plt.colorbar() lines = ax.plot(x, zk[n].T, c="#0066cc") quiv = [] for i in range(10): quiv.append( ax.quiver( x[::50], (zk[n, i, ::50] + zk[n, i + 1, ::50]) / 2, velk[n, i, 0, ::50], vz[n, i, ::50], units="dots", width=2, scale=0.05, ) ) ax.autoscale(True, "w", True) ax.set_ylim(top=15) def animate(k): for i, q in enumerate(quiv): q.set_UVC( velk[k, i, 0, ::50], vz[k, i, ::50], ) for i, l in enumerate(lines): l.set_ydata(zk[k, i]) return *quiv, *lines ani = animation.FuncAnimation( fig, animate, frames=wl[:, 0].size, interval=20, blit=True ) ani.save(out.joinpath("layers.mp4", codec="h265"))