320 lines
11 KiB
Python
320 lines
11 KiB
Python
"""
|
||
@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')
|