In [1]:
# Author: Dr. Christian Kehl
%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 *
Populating the interactive namespace from numpy and matplotlib
In [2]:
%load_ext autoreload
%autoreload 2
%load_ext version_information
%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)
    
Out[6]: