Black+isort, layers, vz in sws_npz

swash/config.ini

@@ -19,9 +19,9 @@ out=out_sws
 mpi=4
 
 [post]
-inp=inp_post
-case=sws_spec_buoy.npz
-compare=sws_spec_buoy_nb.npz
+inp=inp_post_test
+#case=sws_spec_buoy.npz
+#compare=sws_spec_buoy_nb.npz
 out=out_post
 #nperseg=1024
 dt=0.25

swash/processing/animate.py

@@ -3,12 +3,11 @@ import configparser
 import logging
 import pathlib
 
-import matplotlib.pyplot as plt
 import matplotlib.animation as animation
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 
-
 parser = argparse.ArgumentParser(description="Animate swash output")
 parser.add_argument("-v", "--verbose", action="count", default=0)
 args = parser.parse_args()

swash/processing/layers.py

@@ -3,12 +3,11 @@ import configparser
 import logging
 import pathlib
 
-import matplotlib.pyplot as plt
 import matplotlib.animation as animation
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 
-
 parser = argparse.ArgumentParser(description="Animate swash output")
 parser.add_argument("-v", "--verbose", action="count", default=0)
 args = parser.parse_args()
@@ -39,31 +38,55 @@ 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.plot(x, -botl, c="k")
 # ax.fill_between(
 #    x, -botl, -data["botl"] + bathy.hstru, color="k", alpha=0.2
 # )
-lines = ax.plot(x, zk[0].T)
 
-print(velk.shape)
+n = 0
 velk = velk.reshape((6001, 10, 2, 1251))
-vk = np.sqrt((velk[1000] ** 2).sum(axis=1))
-print(vk.shape)
-plt.imshow(vk)
-plt.colorbar()
+vk = np.sqrt((velk[n] ** 2).sum(axis=1))
+# print(vk.shape)
+# plt.imshow(vk)
+# plt.colorbar()
 
-# def animate(i):
-#    for line, z in zip(lines, zk[i]):
-#        line.set_ydata(z)
-#        return lines
-#
-# ani = animation.FuncAnimation(
-#    fig, animate, frames=wl[:, 0].size, interval=20, blit=True
-# )
+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])
+
+
+ani = animation.FuncAnimation(
+    fig, animate, frames=wl[:, 0].size, interval=20, blit=True
+)
 
 plt.show(block=True)

swash/processing/post.py

@@ -5,8 +5,8 @@ import pathlib
 
 import matplotlib.pyplot as plt
 import numpy as np
-import scipy.signal as sgl
 import scipy.fft as fft
+import scipy.signal as sgl
 
 from .read_swash import *
 

swash/processing/sws_npz.py

@@ -1,8 +1,8 @@
 import argparse
 import configparser
 import logging
-from multiprocessing.pool import ThreadPool
 import pathlib
+from multiprocessing.pool import ThreadPool
 
 import numpy as np
 import pandas as pd
@@ -36,6 +36,7 @@ var = {
     "press": rsws.read_scalar,
     "zk": rsws.read_scalar_lay,
     "velk": rsws.read_vector_lay,
+    "vz": rsws.read_scalar_lay,
 }
 
 inp.mkdir(exist_ok=True)
@@ -43,7 +44,8 @@ with ThreadPool() as pool:
     log.info("Converting all data")
     pool.map(
         lambda x: np.save(
-            inp.joinpath(x[0]), x[1](sws_out.joinpath(x[0]).with_suffix(".dat"))
+            inp.joinpath(x[0]),
+            x[1](sws_out.joinpath(x[0]).with_suffix(".dat")),
         ),
         var.items(),
     )