#!/usr/bin/env python # coding: utf-8 # In[1]: # Author: Dr. Christian Kehl get_ipython().run_line_magic('pylab', 'inline') import os import math import xarray as xr import numpy as np from IPython.display import Video, HTML from base64 import b64encode from matplotlib import pyplot as plt from matplotlib import colors from datetime import timedelta from matplotlib.animation import FuncAnimation, writers from DecayLine import DecayLine from CloudFileHelper import * # In[2]: get_ipython().run_line_magic('load_ext', 'autoreload') get_ipython().run_line_magic('autoreload', '2') get_ipython().run_line_magic('load_ext', 'version_information') get_ipython().run_line_magic('version_information', 'numpy, matplotlib, xarray, cartopy') # In[3]: def time_index_value(tx, _ft): # expect ft to be forward-linear ft = _ft if isinstance(_ft, xr.DataArray): ft = ft.data f_dt = ft[1] - ft[0] if type(f_dt) is not np.float64: f_dt = timedelta(f_dt).total_seconds() f_interp = tx / f_dt ti = int(math.floor(f_interp)) return ti def time_partion_value(tx, _ft): # expect ft to be forward-linear ft = _ft if isinstance(_ft, xr.DataArray): ft = ft.data f_dt = ft[1] - ft[0] if type(f_dt) is not np.float64: f_dt = timedelta(f_dt).total_seconds() f_interp = tx / f_dt f_t = f_interp - math.floor(f_interp) return f_t # In[4]: filedir = os.path.abspath(".") pfile_name = "benchmark_doublegyre_noMPI_n1024_fwd.nc" ufile_name = "doublegyreU.nc" vfile_name = "doublegyreV.nc" vdir = os.path.abspath(os.path.join(filedir, "..", "images")) vfile = "parcels_trails.mp4" vfile_path = os.path.join(vdir, vfile) # In[5]: if __name__ == "__main__": if not os.path.exists(vfile_path): Writer = writers['ffmpeg'] ani_writer_nc = Writer(fps=15, metadata=dict(artist='CKehl'), bitrate=3200) get_from_surfdrive("https://surfdrive.surf.nl/files/index.php/s/ui1wS4BNTRpKJOz", os.path.join(filedir, pfile_name)) get_from_surfdrive("https://surfdrive.surf.nl/files/index.php/s/EIWkS3xcMqP8FbQ", os.path.join(filedir, ufile_name)) get_from_surfdrive("https://surfdrive.surf.nl/files/index.php/s/Gm8Sh9VgZfRWDBc", os.path.join(filedir, vfile_name)) zorder = 1 Pn = 96 di = 0 trail_ccode = 'cyan' # ==== Load flow-field data ==== # f_u_nc = xr.open_dataset(os.path.join(filedir, ufile_name), decode_cf=True, engine='netcdf4') fT_nc = f_u_nc['time_counter'] fX_nc = f_u_nc['x'] fY_nc = f_u_nc['y'] fU_nc = f_u_nc['vozocrtx'] max_u_value = max(abs(float(fU_nc.min())), abs(float(fU_nc.max()))) fu_ext = (-max_u_value, +max_u_value) f_v_nc = xr.open_dataset(os.path.join(filedir, vfile_name), decode_cf=True, engine='netcdf4') fV_nc = f_v_nc['vomecrty'] extends = (float(fX_nc.min()), float(fX_nc.max()), float(fY_nc.min()), float(fY_nc.max())) max_v_value = max(abs(float(fV_nc.min())), abs(float(fV_nc.max()))) fv_ext = (-max_v_value, +max_v_value) f_velmag_ext_nc = (0, np.sqrt(max_v_value**2 + max_u_value**2)) total_items = fT_nc.shape[0] data_xarray = xr.open_dataset(os.path.join(filedir, pfile_name)) np.set_printoptions(linewidth=160) ns_per_sec = np.timedelta64(1, 's') # nanoseconds in an sec sec_per_day = 86400.0 time_array = data_xarray['time'].data / ns_per_sec N = data_xarray['lon'].shape[0] tN = data_xarray['lon'].shape[1] indices = np.random.randint(0, N - 1, Pn, dtype=int) time_since_release_nc = (time_array.transpose() - time_array[:, 0]) sim_dt_nc = time_since_release_nc[1, 0] - time_since_release_nc[0, 0] print("dt = {}".format(sim_dt_nc)) pX_nc = data_xarray['lon'] pY_nc = data_xarray['lat'] speed_nc_0 = fU_nc[0, di] ** 2 + fV_nc[0, di] ** 2 speed_nc_0 = np.where(speed_nc_0 > 0, np.sqrt(speed_nc_0), 0) fig_anim_nc = plt.figure() ax_anim_nc = plt.axes() cs_nca_velmag = ax_anim_nc.pcolormesh(fX_nc, fY_nc, speed_nc_0, cmap="Greys", norm=colors.Normalize(vmin=f_velmag_ext_nc[0], vmax=f_velmag_ext_nc[1]), shading='gouraud', zorder=1) cbar_nca_velmag = fig_anim_nc.add_axes([0.0, 0.9, 0.05, 0.07]) plt.colorbar(cs_nca_velmag, cax=cbar_nca_velmag) ax_anim_nc.set_title("Simulation - NetCDF data - t = %5.1f d" % (time_since_release_nc[0, 0] / sec_per_day)) trail_color = list(colors.to_rgba(trail_ccode))[0:3] lines_nc = [] for i in range(0, Pn): lines_nc.append(DecayLine(14, 8, trail_color, zorder=2 + i)) for i in range(0, Pn): lines_nc[i].add_point(pX_nc[indices[i], 0], pY_nc[indices[i], 0]) for l in lines_nc: ax_anim_nc.add_collection(l.get_LineCollection()) def init_nc_animation(): cs_nca_velmag.set_array(speed_nc_0) results = [] results.append(cs_nca_velmag) for l in lines_nc: results.append(l.get_LineCollection()) ax_anim_nc.set_title("Simulation - NetCDF data - t = %5.1f d" % (time_since_release_nc[0, 0] / sec_per_day)) return results def update_flow_only_nc(frames, *args): percent = float(frames) / float(tN) print("\rPlotting progress: [{0:50s}] {1:.1f}%".format('#' * int(percent * 50), percent * 100), end="", flush=True) dt = args[0] tx = float(frames) * dt tx = math.fmod(tx, fT_nc[-1]) ti0 = time_index_value(tx, fT_nc) tt = time_partion_value(tx, fT_nc) ti1 = 0 if ti0 < (len(fT_nc) - 1): ti1 = ti0 + 1 else: ti1 = 0 speed_1 = fU_nc[ti0, di] ** 2 + fV_nc[ti0, di] ** 2 speed_1 = np.where(speed_1 > 0, np.sqrt(speed_1), 0) speed_2 = fU_nc[ti1, di] ** 2 + fV_nc[ti1, di] ** 2 speed_2 = np.where(speed_2 > 0, np.sqrt(speed_2), 0) fs_show = (1.0 - tt) * speed_1 + tt * speed_2 cs_nca_velmag.set_array(fs_show) # == add new lines == # if frames > 0: for pindex in range(0, Pn): lines_nc[pindex].add_point(pX_nc[indices[pindex], frames], pY_nc[indices[pindex], frames]) # == collect results == # results = [] results.append(cs_nca_velmag) for l in lines_nc: results.append(l.get_LineCollection()) ax_anim_nc.set_title("Simulation - NetCDF data - t = %5.1f d" % (tx / sec_per_day)) return results ani = FuncAnimation(fig_anim_nc, update_flow_only_nc, init_func=init_nc_animation, frames=tN, interval=100, fargs=[sim_dt_nc, ]) ani.save(vfile_path, writer=ani_writer_nc, dpi=92) plt.close() # In[6]: Video(vfile_path, embed=True) # In[7]: if os.path.exists(os.path.join(filedir, pfile_name)): os.remove(os.path.join(filedir, pfile_name)) if os.path.exists(os.path.join(filedir, ufile_name)): os.remove(os.path.join(filedir, ufile_name)) if os.path.exists(os.path.join(filedir, vfile_name)): os.remove(os.path.join(filedir, vfile_name)) # In[ ]: