#!/usr/bin/env python
# coding: utf-8

# ## Compute bedload transport from FVCOM 
# 

# Demonstration using the NetCDF4-Python library to compute bedload transport and bottom velocity (1 meter above bottom) from a triangular grid ocean model (FVCOM) via OPeNDAP.  The results are stored in a new NetCDF4 file. 
# 
# NECOFS (Northeastern Coastal Ocean Forecast System) is run by groups at the University of Massachusetts Dartmouth and the Woods Hole Oceanographic Institution, led by Drs. C. Chen, R. C. Beardsley, G. Cowles and B. Rothschild. Funding is provided to run the model by the NOAA-led Integrated Ocean Observing System and the State of Massachusetts.
# 
# NECOFS is a coupled numerical model that uses nested weather models, a coastal ocean circulation model, and a wave model. The ocean model is a volume-mesh model with horizontal resolution that is finer in complicated regions. It is layered (not depth-averaged) and includes the effects of tides, winds, and varying water densities caused by temperature and salinity changes.
# 
# * Model description: http://fvcom.smast.umassd.edu/research_projects/NECOFS/model_system.html
# * THREDDS server with other forecast and archive products: http://www.smast.umassd.edu:8080/thredds/catalog.html

# In[1]:


import numpy as np
import netCDF4
import datetime as dt


# In[2]:


# Input FVCOM Dataset: DAP Data URL
url = 'http://www.smast.umassd.edu:8080/thredds/dodsC/FVCOM/NECOFS/Forecasts/NECOFS_FVCOM_OCEAN_MASSBAY_FORECAST.nc'
url = 'http://www.smast.umassd.edu:8080/thredds/dodsC/FVCOM/NECOFS/Forecasts/NECOFS_GOM3_FORECAST.nc'
url = 'http://www.smast.umassd.edu:8080/thredds/dodsC/fvcom/archives/necofs_mb'
url = 'http://www.smast.umassd.edu:8080/thredds/dodsC/fvcom/hindcasts/30yr_gom3'

url_out = '/usgs/data2/notebook/fvcom/gom3_bedload_20140805.nc'

# Open DAP
nci = netCDF4.Dataset(url)

# specific times
start = dt.datetime(2014,8,5,10,0,0) 
end = start + dt.timedelta(hours=148)    # 12.42*12 = 149.04


# In[3]:


def z0tocd(z0=3.3546e-04,zr=1.0):
    """ Compute the drag coefficient CD based on roughness height z0 and distance above bottom zr"""
    kappa = 0.4
    cd=(kappa * np.ones_like(zr) / np.log(zr/z0))**2
    return cd


# In[4]:


def cdtoz0(cd=2.5e-3,zr=1.0):
    """ Compute the roughness height z0 based on drag coefficient CD and distance above bottom zr"""
    kappa = 0.4
    z0 = zr / (np.exp(kappa * np.ones_like(cd) / np.sqrt(cd)))
    return z0


# In[5]:


def w100(w=0.1+0j, z0=3.35e-04, zr=1):
    """ Compute the velocity 1 meter above bottom and friction velocity
        from velocity measured at height zr above bottom.

    Keyword arguments
    -----------------
    w : east velocity component  + j*north velocity component (m/s) [complex]
    z0 : roughness height = kb/30 (m) 
    zr : height above bottom for input velocity "w" (m)

    Returns
    -------
    ustar : friction velocity (m/s) [complex]
    w : velocity 1 mab (m/s) [complex]
    
    """
    
    cd = z0tocd(z0,zr)
    ustar = np.sqrt(cd)*w
    kappa = 0.4
    ur = abs(ustar)/kappa*np.log(np.ones_like(zr)/z0)
    wbot = w*ur/(np.abs(w)+1e-16)
    return ustar,wbot


# In[6]:


time_var = nci['time']
istart = netCDF4.date2index(start, time_var, select='nearest')
iend = netCDF4.date2index(end, time_var, select='nearest')
jd = netCDF4.num2date(time_var[istart:iend+1], time_var.units)


# In[7]:


itimes = range(istart,iend+1)


# In[8]:


len(itimes)


# In[9]:


# read connectivity array
nv = nci['nv'][:].T - 1


# In[10]:


# print info on velocity varibble
print(nci['u'])
node = len(nci['h'])
nt, nsig, nele = np.shape(nci['u'])


# In[11]:


# OUTPUT: create NetCDF4 file with deflation on variables


nco = netCDF4.Dataset(url_out, 'w', clobber=True)

# create dimensions
nco.createDimension('nele', nele)
nco.createDimension('node', node)
nco.createDimension('three', 3)
nco.createDimension('time', None)

# create variables
timeo = nco.createVariable('time', 'f4',  ('time'))
ho = nco.createVariable('h', 'f4',  ('node'))
nvo = nco.createVariable('nv', 'i4',  ('three', 'nele'))
lonco = nco.createVariable('lonc', 'f4',  ( 'nele'))
latco = nco.createVariable('latc', 'f4',  ( 'nele'))
lono = nco.createVariable('lon', 'f4',  ( 'node'))
lato = nco.createVariable('lat', 'f4',  ( 'node'))

ubot = nco.createVariable('ubot', 'f4',  ('time', 'nele'))
vbot = nco.createVariable('vbot', 'f4',  ('time', 'nele'))
ubedload = nco.createVariable('ubedload', 'f4',  ('time', 'nele'))
vbedload = nco.createVariable('vbedload', 'f4',  ('time', 'nele'))

# write variable attributes
timeo.units=nci['time'].units
ho.units=nci['h'].units
lono.units=nci['lon'].units
lato.units=nci['lat'].units
lonco.units=nci['lonc'].units
latco.units=nci['latc'].units
ubot.units=nci['u'].units
vbot.units=nci['v'].units
ubot.standard_name = 'eastward_component_of_bottom_velocity'
vbot.standard_name = 'northward_component_of_bottom_velocity'

ubedload.units='kg m-1 s-1'
vbedload.units='kg m-1 s-1'
ubedload.standard_name = 'eastward_component_of_bedload_transport'
vbedload.standard_name = 'northward_component_of_bedload_transport'
# write data with no time dimension
lonco[:]=nci['lonc'][:]
latco[:]=nci['latc'][:]
lono[:]=nci['lon'][:]
lato[:]=nci['lat'][:]
nvo[:]=nci['nv'][:]
ho[:]=nci['h'][:]


# In[12]:


# specify bottom layer, but handle case where there is just 1 layer in input file
if np.shape(nci['siglay'])[0]==1:
    ilayer = 0
else:
    ilayer = -1


# In[13]:


# neither z0 or cd is saved in this FVCOM output, so just use canonical bottom roughness kb=0.5 cm
kb=0.005
z0=kb/30.
rho = 1025.  # density plays a small role in stress, so just specify as constant here


# In[14]:


def mpm_bedload(w, z0, zr, d50, rho, rho_s):
    g = 9.81         # gravity
    theta_c = 0.047   # shields parameter

    # compute bottom friction velocity and velocity at 1 mab
    ustar, wbot = w100(w, z0, zr)    

    # compute bottom stress from friction velocity
    cd = z0tocd(z0,zr)

    bstr = cd * rho * ustar * np.abs(ustar) 

    theta_sf = np.abs(bstr) / ((s - 1.0) * g * d50)

    theta = (theta_sf - theta_c)
    theta[theta<0.0] = 0.0
    phi = 8.0*(theta**1.5)

    q = phi * np.sqrt((s-1.0) * g * d50**3)*rho_s

    return q, wbot


# In[15]:


# pick a sand grain size for bedload transport calculation
d50 = 200.0e-06  # 200 micron sand
rho_s = 2650. # density of sediment 
s = rho_s/rho
g = 9.81
theta_c = 0.047


# In[16]:


# test
w = np.array([0, 0.5]) + 1j*np.array([.5, .5])
zr = 1.0
q, wbot = mpm_bedload(w, z0, zr, d50, rho, rho_s)
print(q, wbot)


# In[17]:


from tqdm import tqdm_notebook as tqdm

# loop through time, writing each 2D or 3D field to output file
k=0
for itime in tqdm(itimes):
#    print('{} percent done'.format(float(k)/len(itimes)*100.))
    # get current at layer [0 = surface, -1 = bottom]
    zbed = -1.0
    zr = (nci['siglay'][ilayer, :] - zbed) * (nci['h'][:]+nci['zeta'][itime,:])
    u = nci['u'][itime, ilayer, :]
    v = nci['v'][itime, ilayer, :]

    # average nodes to get bottom layer thicknesses at faces (velocity points)
    zr_face = np.mean(zr[nv],axis=1)
    
    # create complex velocity from components
    w = u + 1j*v
  
    q, wbot = mpm_bedload(w, z0, zr_face, d50, rho, rho_s)

    # write bottom velocity and stress components to output file
    ubot[k,:]=wbot.real 
    vbot[k,:]=wbot.imag 
    ubedload[k,:] = q * np.real(w) / np.abs(w)   # bedload in x direction
    vbedload[k,:] = q * np.imag(w) / np.abs(w)   # bedload in y direction
    timeo[k] = nci['time'][itime]
    k += 1


# In[18]:


nci.close()
nco.close()


# To calculate net bedload transport, use NCO:  
# ```
# ncra -O gom3_bedload_20150805.nc gom3_bedload_20150805_mean.nc
# ```
#     

# In[19]:


get_ipython().system('ls *.nc')


# In[20]:


get_ipython().system('ls /usgs/data2/notebook/fvcom')


# In[ ]: