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()