""" @author: AJMR Plotting and animation script for NTC sediment tests April 17, 2023 """ import os import pandas as pd import geopandas as gpd import netCDF4 as nc import math import numpy as np import matplotlib as mpl import matplotlib.pyplot as plt from scipy.interpolate import LinearNDInterpolator, interp1d import matplotlib.animation as animation import dfm_tools as dfmt # Geotiff import cartopy.crs as ccrs import contextily as ctx import pathlib as pl from shapely.geometry import Point, MultiPoint import alphashape import rioxarray import xarray as xr import cartopy.crs as ccrs import contextily as ctx from dfm_tools.regulargrid import scatter_to_regulargrid from pathlib import Path # Colormaps import cm_xml_to_matplotlib as cm import pytz #%% Read Model Log runLog = pd.read_excel('C:/Users/arey/files/Projects/Newtown/Model Log NTC.xlsx', 'ModelLog') dataPath = "FlowFM/dflowfm/output/FlowFM_map.nc" histPath = "FlowFM/dflowfm/output/FlowFM_his.nc" #%% Import Model results for animations # modelPlot = [20, 22] # modelPlot = [18, 20, 21, 22] # modelPlot = [18, 21] modelPlot = [146] d3dfm_DataArray = [None] * (max(modelPlot)+1) modelVelX = [None] * (max(modelPlot)+1) modelVelY = [None] * (max(modelPlot)+1) tSteps = [None] * (max(modelPlot)+1) # Time steps to import tStepsImport = range(1000, 7000) for i in modelPlot: # Create variables modelVelX[i] = [None] * (max(tStepsImport)+1) modelVelY[i] = [None] * (max(tStepsImport)+1) file_nc_map = Path(runLog['Run Location'][i]) / dataPath d3dfm_DataArray[i] = dfmt.open_partitioned_dataset(file_nc_map.as_posix()) # Get Var info vars_pd = dfmt.get_ncvarproperties(d3dfm_DataArray[i]) # Get time steps tSteps[i] = d3dfm_DataArray[i].time.values # Loop through time steps and import data for t in tStepsImport: # Import 3D velocity data modelVelXin = d3dfm_DataArray[i].mesh2d_ucx[t, :, :].compute() modelVelYin = d3dfm_DataArray[i].mesh2d_ucy[t, :, :].compute() # Average over layers modelVelX[i][t] = modelVelXin.mean(axis=1) modelVelY[i][t] = modelVelYin.mean(axis=1) print(t) # %% Import water levels from hist modelHistWL = [None] * (max(modelPlot)+1) for i in modelPlot: # Hist file path file_nc_hist = Path(runLog['Run Location'][i]) / histPath # Open hist file dataset data_xr = xr.open_mfdataset(file_nc_hist, preprocess=dfmt.preprocess_hisnc) # Get stations stations_pd = data_xr['stations'].to_dataframe() # Convert to Pandas modelHistWL[i] = data_xr.waterlevel.sel(stations=['West_WL', 'East_WL']).to_pandas() # Put in UTC if needed if i == 104 or i == 105 or i >= 135: modelHistWL[i] = modelHistWL[i].tz_localize( pytz.timezone('EST')).tz_convert(pytz.utc) else: modelHistWL[i] = modelHistWL[i].tz_localize( pytz.timezone('UTC')) #%% Velocity Quiver animations plt.rcParams['animation.ffmpeg_path'] = \ 'C:/Users/arey/files/Local/ffmpeg-2022-02-14-git-59c647bcf3-full_build/bin/ffmpeg.exe' mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckex9vtri0o6619p55sl5qiyv/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw' x, y, arrow_length = 0.93, 0.95, 0.12 # fontprops = fm.FontProperties(size=12) # Set model to plot modelPlot = [84] tStepsPlot = tStepsImport tStepsPlot = range(970, 971) tStepsPlot = range(955, 990) vmin = 0 # Legend min vmax = 0.05 # Legend max scale = 30 # Size of arrows, smaller is bigger #25 qmax = 0.002# Scaling threshold. Values above this are set to 1 amin = 0.1 # Minimum arrow length. A dot is used below this value for m in modelPlot: # Find water level plot range wlPlotMinIDX = wlHistTimeIDX = np.argmin(abs(tSteps[m][tStepsPlot[0]] - modelHistWL[m].index.values)) wlPlotMaxIDX = wlHistTimeIDX = np.argmin(abs(tSteps[m][tStepsPlot[-1]] - modelHistWL[m].index.values)) cmapName = 'turbo' cmap = mpl.cm.turbo # # Zoom limits # xLimits = [305900, 306500] # yLimits = [61400, 62200] # Corner limits xLimits = [306088, 306444] yLimits = [61370, 61741] # Setup Video metadata = dict(title='NTC Velocity', artist='Matplotlib', comment='AJMR June 2, 2023') writer = animation.FFMpegWriter(fps=1, metadata=metadata, codec='h264', bitrate=5000) fig, axes = plt.subplots(figsize=(10, 10)) writer.setup(fig, 'C:/Users/arey/files/Projects/Newtown/Vel_Animation4.mp4') mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw' frameCount = 1 for t in tStepsPlot: #289 # Clear axes for video axes.cla() # Create new figure for stills # fig, axes = plt.subplots(figsize=(12, 12)) # Clear inset WL Plot if frameCount > 1: axin1.cla() # Cropt to zoom limits croppedGridMask = (d3dfm_DataArray[m].mesh2d_face_x >= xLimits[0]) & \ (d3dfm_DataArray[m].mesh2d_face_x <= xLimits[1]) & \ (d3dfm_DataArray[m].mesh2d_face_y >= yLimits[0]) & \ (d3dfm_DataArray[m].mesh2d_face_y <= yLimits[1]) croppedGridMask = croppedGridMask.compute() # Rasterize Ugrid rasterU = dfmt.rasterize_ugrid(d3dfm_DataArray[m]['mesh2d_ucx'].isel( time=t, mesh2d_nLayers=4).where(croppedGridMask, drop=True).to_dataset(), resolution=10) rasterV = dfmt.rasterize_ugrid(d3dfm_DataArray[m]['mesh2d_ucy'].isel( time=t, mesh2d_nLayers=4).where(croppedGridMask, drop=True).to_dataset(), resolution=10) # Using fixed grid-> select desiered time step/layer/variable, convert to dataset, rasterize, then to numpy overGrid_U_Scale = rasterU.to_array().to_numpy().squeeze() overGrid_V_Scale = rasterV.to_array().to_numpy().squeeze() r = np.power(np.add(np.power(overGrid_U_Scale, 2), np.power(overGrid_V_Scale, 2)), 0.5) curPlotMask = np.sqrt((overGrid_U_Scale) ** 2 + (overGrid_V_Scale) ** 2) > qmax overGrid_U_Plot = np.zeros(overGrid_U_Scale.shape) overGrid_V_Plot = np.zeros(overGrid_V_Scale.shape) # Normalize arrows to 1 where not zero overGrid_U_Plot[r != 0] = overGrid_U_Scale[r != 0] / r[r != 0] overGrid_V_Plot[r != 0] = overGrid_V_Scale[r != 0] / r[r != 0] # Scale arrows above qmax overGrid_U_Plot[~curPlotMask] = overGrid_U_Plot[~curPlotMask] * (r[~curPlotMask] / qmax) overGrid_V_Plot[~curPlotMask] = overGrid_V_Plot[~curPlotMask] * (r[~curPlotMask] / qmax) # Calculate velocity magnitude velMag = np.sqrt((overGrid_U_Scale) ** 2 + (overGrid_V_Scale) ** 2) # Using fixed grid qv = axes.quiver(rasterU['x'].values, rasterU['y'].values, overGrid_U_Plot, overGrid_V_Plot, velMag, scale=scale, color='k', headlength=5, headwidth=5, headaxislength=5, width=0.002, minlength=amin) #200 # Set map limits axes.set_xlim(xLimits[0], xLimits[1]) axes.set_ylim(yLimits[0], yLimits[1]) # Add basemap ctx.add_basemap(axes, source=mapbox, crs='EPSG:32118') # Color Limits qv.set_clim([vmin, vmax]) # Set colormap for quiver qv.set_cmap(cmap) # Add title axes.title.set_text('Mid-depth velocity: ' + np.datetime_as_string(tSteps[m][t], unit='m')) # Add inset wl plot axin1 = axes.inset_axes([0.1, 0.1, 0.4, 0.15]) # Find matching time step in model water level history wlHistTimeIDX = np.argmin(abs(tSteps[m][t] - modelHistWL[m].index.values)) # Plot model water level history axin1.plot(modelHistWL[m].index, modelHistWL[m]) # Add current point axin1.scatter(modelHistWL[m].index[wlHistTimeIDX], modelHistWL[m]['East_WL'][wlHistTimeIDX], 10, 'r') axin1.set_xticks([]) axin1.set_yticks([]) # Set time axis limits axin1.set_xlim([modelHistWL[m].index[wlPlotMinIDX], modelHistWL[m].index[wlPlotMaxIDX]]) # Set inset axis y limits axin1.set_ylim([-0.5, 2]) plt.setp(axin1.get_xticklabels(), backgroundcolor="white") plt.setp(axin1.get_yticklabels(), backgroundcolor="white") if frameCount == 1: #i < 1000 fig.subplots_adjust(right=0.8) norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) # Create colorbar axis # cax = plt.axes([0.82, 0.25, 0.04, 0.5]) cax = plt.axes([0.85, 0.15, 0.04, 0.7]) # Add colorbar without designated tick marks cb1 = mpl.colorbar.ColorbarBase(cax, cmap=cmap, norm=norm, orientation='vertical') # Colorbar lavel cb1.set_label('Water velocity [m/s]') # Save video frame writer.grab_frame() # plt.show() print(t) plt.show() # Increase frame count frameCount += 1 # Save stills # fig.savefig('//srv-ott3/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/' + # '06_Models/04_Delft3D/Figures/SedConc/Wide_RevC_SedConc_' + # runLog['Run'][m] + '_Frame_' + str(i) + '.png', # bbox_inches='tight', dpi=300) # Save video writer.finish() #%% Save Sediment Geotiff # Create mask within bounds boundSelect = 2 croppedGridMask = (d3dfm_DataArray[i].mesh2d_face_x >= axLim_Xmin[boundSelect]) &\ (d3dfm_DataArray[i].mesh2d_face_x <= axLim_Xmax[boundSelect]) &\ (d3dfm_DataArray[i].mesh2d_face_y >= axLim_Ymin[boundSelect]) &\ (d3dfm_DataArray[i].mesh2d_face_y <= axLim_Ymax[boundSelect]) croppedGridMask = croppedGridMask.compute() # Step through time steps for sedTstep in range(105, 120): # Rasterize sediment concentration within bounds sedConcTif = dfmt.rasterize_ugrid(d3dfm_DataArray[i]['mesh2d_sedfrac_concentration'].isel( time=sedTstep, mesh2d_nLayers=4, nSedSus=0).where(croppedGridMask, drop=True).to_dataset(), resolution=2) # Add CRS sedConcTif.rio.write_crs("epsg:32118", inplace=True) sedConcTif.rio.to_raster('C:/Users/arey/files/Projects/Newtown/SedFigs2/GeoTiff_RevB_' + 'Sed_Conc4_Time_' + str(sedTstep) + '.tif')