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