def sunhis2initial(hisfile,icfile):
"""
Main function
"""
print '#####\nCreating initial condition file (%s) from history file (%s)...'%(icfile,hisfile)
# Load the history file
sunhis = Spatial(hisfile, tstep=-1, klayer=[-99])
# Load the initial condition object
timeic = datetime.strftime(sunhis.time[-1],'%Y%m%d.%H%M%S')
sunic = InitialCond(hisfile,timeic)
# Load the last time step into each variable
sunic.h = sunhis.loadData(variable='eta').reshape((1,sunic.h.shape))
sunic.T = sunhis.loadData(variable='temp').reshape((1,)+sunic.T.shape)
sunic.S = sunhis.loadData(variable='salt').reshape((1,)+sunic.S.shape)
sunic.uc = sunhis.loadData(variable='uc').reshape((1,)+sunic.uc.shape)
sunic.vc = sunhis.loadData(variable='vc').reshape((1,)+sunic.vc.shape)
print sunic.h.shape
print sunic.T.shape
# Load the age
if sunhis.hasVar('agec'):
sunic.agec = sunhis.loadData(variable='agec').reshape((1,)+sunic.agec.shape)
sunic.agealpha = sunhis.loadData(variable='agealpha').reshape((1,)+sunic.agealpha.shape)
# Write to the output file
sunic.writeNC(icfile)
def _dumpic(self):
"""
Dump initial condition plots
For each variable (uc,vc,T,S,eta):
- surface plot
- seabed plot
"""
varnames = ['uc','vc','temp','salt','eta']
sun = Spatial(self.suntanspath+'/'+self.icfile,klayer=[0])
# Plot the surface variables
for vv in varnames:
h=plt.figure()
sun.variable=vv
sun.loadData()
sun.clim = [sun.data.min(),sun.data.max()]
sun.plot()
outfile = '%s/IC_%s_surface.png'%(self.plotdir,vv)
sun.savefig(outfile)
del h
# Plot the seabed variables
sun.klayer=[-1]
for vv in varnames:
h=plt.figure()
sun.variable=vv
sun.loadData()
sun.clim = [sun.data.min(),sun.data.max()]
sun.plot()
outfile = '%s/IC_%s_seabed.png'%(self.plotdir,vv)
sun.savefig(outfile)
del h
def sunhis2initial(hisfile, icfile):
"""
Main function
"""
print '#####\nCreating initial condition file (%s) from history file (%s)...' % (
icfile, hisfile)
# Load the history file
sunhis = Spatial(hisfile, tstep=-1, klayer=[-99])
# Load the initial condition object
timeic = datetime.strftime(sunhis.time[-1], '%Y%m%d.%H%M%S')
sunic = InitialCond(hisfile, timeic)
# Load the last time step into each variable
sunic.h = sunhis.loadData(variable='eta').reshape((1, sunic.h.shape))
sunic.T = sunhis.loadData(variable='temp').reshape((1, ) + sunic.T.shape)
sunic.S = sunhis.loadData(variable='salt').reshape((1, ) + sunic.S.shape)
sunic.uc = sunhis.loadData(variable='uc').reshape((1, ) + sunic.uc.shape)
sunic.vc = sunhis.loadData(variable='vc').reshape((1, ) + sunic.vc.shape)
print sunic.h.shape
print sunic.T.shape
# Load the age
if sunhis.hasVar('agec'):
sunic.agec = sunhis.loadData(variable='agec').reshape((1, ) +
sunic.agec.shape)
sunic.agealpha = sunhis.loadData(
variable='agealpha').reshape((1, ) + sunic.agealpha.shape)
# Write to the output file
sunic.writeNC(icfile)
def _dumpic(self):
"""
Dump initial condition plots
For each variable (uc,vc,T,S,eta):
- surface plot
- seabed plot
"""
varnames = ['uc', 'vc', 'temp', 'salt', 'eta']
sun = Spatial(self.suntanspath + '/' + self.icfile, klayer=[0])
# Plot the surface variables
for vv in varnames:
h = plt.figure()
sun.variable = vv
sun.loadData()
sun.clim = [sun.data.min(), sun.data.max()]
sun.plot()
outfile = '%s/IC_%s_surface.png' % (self.plotdir, vv)
sun.savefig(outfile)
del h
# Plot the seabed variables
sun.klayer = [-1]
for vv in varnames:
h = plt.figure()
sun.variable = vv
sun.loadData()
sun.clim = [sun.data.min(), sun.data.max()]
sun.plot()
outfile = '%s/IC_%s_seabed.png' % (self.plotdir, vv)
sun.savefig(outfile)
del h
def loadData(self, variable=None):
"""
Overloaded loadData function - updates the unstructured grid object
"""
Spatial.loadData(self, variable=variable)
if self.is3D:
self.data=np.ravel(self.data[self.mask3D])
else:
self.data=np.ravel(self.data)
self.ug.cell_data.scalars = self.data
self.ug.cell_data.scalars.name = 'suntans_scalar'
self.ug.modified()
def loadData(self):
"""
Overloaded loadData function - updates the unstructured grid object
"""
Spatial.loadData(self)
if self.is3D:
self.data=np.ravel(self.data[self.mask3D])
else:
self.data=np.ravel(self.data)
self.ug.cell_data.scalars = self.data
self.ug.cell_data.scalars.name = 'suntans_scalar'
self.ug.modified()
def river_discharge(ncfile,shpfiles):
"""
Calculates the river flux of all type-2 boundaries located with each polygon
specified with a shpfile polygon
"""
# Load the spatial object
sun = Spatial(ncfile,klayer=[-99])
# Indentify the river cells
ind = np.argwhere(sun.mark==2).ravel()
#sun.j = ind
# Calculate the mask region
if type(shpfiles) != type([]):
shpfiles = [shpfiles]
masks = []
polynames = []
for shpfile in shpfiles:
mask, maskpoly = maskShpPoly(sun.xe[ind],sun.ye[ind],shpfile)
masks.append(mask)
polynames.append(os.path.splitext(os.path.basename(shpfile))[0])
# Create a dictionary with output info
data={}
for poly in polynames:
data.update({poly:{'Q_r':np.zeros((sun.Nt,)),'time':sun.time}})
print 'Loading the data...'
#U = sun.loadData(variable='U_F')
sun.tstep = range(sun.Nt)
U = np.zeros((sun.Nt,sun.Nkmax,ind.shape[0]))
for tt in range(sun.Nt):
sun.tstep = tt
tmp = sun.loadData(variable='U_F')
U[tt,...] = tmp[:,ind]
print '\t%d of %d...'%(tt,sun.Nt)
for mask,poly in zip(masks,polynames):
tmp_dA = np.sum( np.sum(U*mask,axis=-1), axis=-1)
data[poly]['Q_r']=tmp_dA
return data
def area_integrate(ncfile,varnames,shpfiles):
"""
Area integrate a suntans variable for all time in the domain
specified with a shpfile polygon
"""
# Use numexpr to try and speed things up
import numexpr as ne
# Load the spatial object
sun = Spatial(ncfile,klayer=[-99])
Ac = sun.Ac
# Calculate the mask region
if type(shpfiles) != type([]):
shpfiles = [shpfiles]
masks = []
polynames = []
for shpfile in shpfiles:
mask, maskpoly = maskShpPoly(sun.xv,sun.yv,shpfile)
masks.append(mask)
polynames.append(os.path.splitext(os.path.basename(shpfile))[0])
# Create a dictionary with output info
data={}
for poly in polynames:
data.update({poly:{'V':np.zeros((sun.Nt,)),'time':sun.time}})
for varname in varnames:
data[poly].update({varname:np.zeros((sun.Nt,))})
sun.tstep = range(sun.Nt)
for varname in varnames:
print 'Area integrating varibles: %s ...'%(varname)
tmp = sun.loadData(variable=varname)
for mask,poly in zip(masks,polynames):
tmp_dA = ne.evaluate("sum(tmp*Ac*mask,axis=1)")
data[poly][varname]=tmp_dA
return data
def energy_budget(energyfile,polyfile,trange):
"""
# Area-integrate the energy terms
"""
varnames = ['KEz','PEz','uP','uKE','uPE','ueta','W_work','B_flux','diss']
# Load the energy file as a suntans object
sun = Spatial(energyfile)
# Create the mask
mask,maskpoly = maskShpPoly(sun.xv,sun.yv,polyfile)
# Initialise the output dictionary
tstep = range(0,sun.Nt)[trange[0]:trange[1]]
nt = len(tstep)
budget ={}
for vv in varnames:
budget.update({vv:np.zeros((nt,))})
for ii,tt in enumerate(tstep):
print 'Area-integrating step: %d of %d...'%(ii,tstep[-1])
for vv in varnames:
sun.tstep=[tt]
data = sun.loadData(variable=vv)
budget[vv][ii],areatotal = sun.areaint(data,mask)
budget.update({'time':sun.time[tstep]})
# Calculate the time-rate of change of KE and PE
dt = sun.timeraw[1]-sun.timeraw[0]
budget.update({'dKE_dt':np.zeros((nt,))})
budget.update({'dPE_dt':np.zeros((nt,))})
budget['dKE_dt'][1::] = (budget['KEz'][1::]-budget['KEz'][0:-1])/dt
budget['dPE_dt'][1::] = (budget['PEz'][1::]-budget['PEz'][0:-1])/dt
return budget
def suntans2ic(self,hisfile,setUV=False,seth=False):
"""
Uses data from another suntans file as initial conditions
Data needs to be on the same grid
"""
# Load the history file
sunhis = Spatial(hisfile, tstep=-1, klayer=[-99])
# Set the time step to grab from the history file
#...
tstep = sunhis.getTstep(self.time,self.time)
sunhis.tstep = [tstep[0]]
print 'Setting the intial condition with time step: %s\nfrom the file:%s'\
%(datetime.strftime(sunhis.time[tstep[0]],\
'%Y-%m-%d %H-%M-%S'),hisfile)
# Npw grab each variable and save in the IC object
if seth:
self.h = sunhis.loadData(variable='eta').reshape((1,self.h.shape))
if sunhis.hasVar('temp'):
self.T = sunhis.loadData(variable='temp').reshape((1,)+self.T.shape)
if sunhis.hasVar('salt'):
self.S = sunhis.loadData(variable='salt').reshape((1,)+self.S.shape)
if setUV:
self.uc = sunhis.loadData(variable='uc').reshape((1,)+self.uc.shape)
self.vc = sunhis.loadData(variable='vc').reshape((1,)+self.vc.shape)
# Load the age
if sunhis.hasVar('agec'):
self.agec = sunhis.loadData(variable='agec').reshape((1,)+self.agec.shape)
self.agealpha = sunhis.loadData(variable='agealpha').reshape((1,)+self.agealpha.shape)
print 'Done setting initial condition data from file.'
wroms = np.load("calcs/mask_bayshelf_weights_radial.npz")["wroms"]
wromsu = wroms[:, 1:] # don't want to interpolate this
wromsv = wroms[1:, :]
etablend = np.empty_like(gridblend.x_rho)
# Read in SUNTANS output
for File in Files:
# klayer = 0 is surface, tstep=-99 reads in all time steps
# Do u, v, eta
sun = Spatial(File, klayer=2, tstep=-99, variable="eta") # , clim=clim)
# sun = Spatial(File, klayer=0, tstep=-99, variable='eta')#, clim=clim)
datesbay = sun.time # get dates available in file
eta = sun.loadData()
# Loop through times in bay file
for i, date in enumerate(datesbay):
# Find the same time from the shelf model output
# Assumes that there is a time in the shelf model that perfectly aligns with the
# times in the bay model at some point
tindshelf = find(datesshelf == date)
print date
# pdb.set_trace()
if not find(datesshelf == date): # if a shelf time doesn't align, don't use this bay model output
continue
# have to do u and v by time step for some reason
u = Spatial(File, klayer=["surface"], tstep=i, variable="uc").loadData()
def suntans2untrim(ncfile, outfile, tstart, tend, grdfile=None):
"""
Converts a suntans averages netcdf file into untrim format
for use in particle tracking
"""
####
# Step 1: Load the suntans data object
####
sun = Spatial(ncfile, klayer=[-99])
# Calculate some other variables
sun.de = sun.get_edgevar(sun.dv, method='min')
sun.mark[sun.mark == 5] = 0
sun.mark[sun.mark == 3] = 2
sun.facemark = np.zeros((sun.Nc, ), dtype=np.int)
# Update the grad variable from the ascii grid file if supplied
if not grdfile == None:
print('Updating grid with ascii values...')
grd = Grid(grdfile)
sun.grad = grd.grad[:, ::-1]
###
# Step 2: Write the grid variables to a netcdf file
###
nc = Dataset(outfile, 'w', format='NETCDF4_CLASSIC')
# Global variable
nc.Description = 'UnTRIM history file converted from SUNTANS output'
# Write the dimensions
for dd in list(untrim_griddims.keys()):
if dd == 'time':
nc.createDimension(untrim_griddims[dd], 0)
elif dd == 'numsides':
nc.createDimension(untrim_griddims[dd], sun.maxfaces)
else:
nc.createDimension(untrim_griddims[dd], sun[dd])
for dd in other_dims:
nc.createDimension(dd, other_dims[dd])
###
# Step 3: Initialize all of the grid variables
###
def create_nc_var(name, dimensions, attdict,data=None, \
dtype='f8',zlib=False,complevel=0,fill_value=999999.0):
tmp=nc.createVariable(name, dtype, dimensions,\
zlib=zlib,complevel=complevel,fill_value=fill_value)
for aa in list(attdict.keys()):
tmp.setncattr(aa, attdict[aa])
if not data == None:
nc.variables[name][:] = data
# Make sure the masked cells have a value of -1
mask = sun['cells'].mask.copy()
sun['cells'][mask] = FILLVALUE
sun['face'][mask] = FILLVALUE
for vv in list(untrim_gridvars.keys()):
vname = untrim_gridvars[vv]
print('Writing grid variable %s (%s)...' % (vname, vv))
if vv == 'time':
continue
# add dz_min attribute to z_r variable
if vv == 'z_r':
ugrid[vname]['attributes'].update({'dz_min': 1e-5})
#sun[vv][:]=sun[vv][::-1]
sun[vv][:] = sun['z_w'][0:-1][::-1]
# Reverse the order of grad(???)
if vv == 'grad':
sun[vv][:] = sun[vv][:, ::-1]
## Fix one-based indexing
#if vv in ['cells','edges','grad']:
# mask = sun[vv][:]==-1
# tmp = sun[vv][:]+1
# tmp[mask]=-1
# #sun[vv][:]=sun[vv][:]+1
# create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
# data=tmp,dtype=ugrid[vname]['dtype'])
create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
data=sun[vv],dtype=ugrid[vname]['dtype'])
# Initialize the two time variables
vname = untrim_gridvars['time']
create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
dtype=ugrid[vname]['dtype'])
vname = 'Mesh2_data_time_string'
create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
dtype=ugrid[vname]['dtype'])
###
# Step 4: Initialize all of the time-varying variables (but don't write)
###
for vname in varnames:
print('Creating variable %s...' % (vname))
create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
dtype=ugrid[vname]['dtype'],zlib=True,complevel=1,fill_value=999999.)
###
# Step 5: Loop through all of the time steps and write the variables
###
tsteps = sun.getTstep(tstart, tend)
tdays = othertime.DaysSince(sun.time, basetime=datetime(1899, 12, 31))
for ii, tt in enumerate(tsteps):
# Convert the time to the untrim formats
timestr = datetime.strftime(sun.time[tt], '%Y-%m-%d %H:%M:%S')
print('Writing data at time %s (%d of %d)...' %
(timestr, tt, tsteps[-1]))
#Write the time variables
nc.variables['Mesh2_data_time'][ii] = tdays[ii]
nc.variables['Mesh2_data_time_string'][:, ii] = timestr
# Load each variable or calculate it and convert it to the untrim format
sun.tstep = [tt]
###
# Compute a few terms first
eta = sun.loadData(variable='eta')
U = sun.loadData(variable='U_F')
dzz = sun.getdzz(eta)
dzf = sun.getdzf(eta)
vname = 'Mesh2_sea_surface_elevation'
#print '\tVariable: %s...'%vname
nc.variables[vname][:, ii] = eta
vname = 'Mesh2_salinity_3d'
#print '\tVariable: %s...'%vname
tmp3d = sun.loadData(variable='salt')
nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]
vname = 'Mesh2_vertical_diffusivity_3d'
#print '\tVariable: %s...'%vname
tmp3d = sun.loadData(variable='nu_v')
nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]
vname = 'h_flow_avg'
#print '\tVariable: %s...'%vname
nc.variables[vname][:, :, ii] = U.swapaxes(0, 1)[:, ::-1]
vname = 'v_flow_avg'
#print '\tVariable: %s...'%vname
tmp3d = sun.loadData(variable='w') * sun.Ac # m^3/s
nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]
# Need to calculate a few terms for the other variables
vname = 'Mesh2_edge_wet_area'
#print '\tVariable: %s...'%vname
#dzf = sun.loadData(variable='dzf')
tmp3d = dzf * sun.df
nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]
vname = 'Mesh2_face_water_volume'
#print '\tVariable: %s...'%vname
#dzz = sun.loadData(variable='dzz')
tmp3d = dzz * sun.Ac
nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]
vname = 'Mesh2_face_wet_area'
#print '\tVariable: %s...'%vname
tmp3d = np.repeat(sun.Ac[np.newaxis, ...], sun.Nkmax, axis=0)
nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]
# UnTRIM references from bottom to top i.e.
# k = 0 @ bed ; k = Nkmax-1 @ top
vname = 'Mesh2_edge_bottom_layer'
#print '\tVariable: %s...'%vname
#tmp2d = sun.Nkmax-sun.Nke # zero based
tmp2d = sun.Nkmax - sun.Nke + 1 # one based
nc.variables[vname][:, ii] = tmp2d
vname = 'Mesh2_edge_top_layer'
#print '\tVariable: %s...'%vname
etop = sun.loadData(variable='etop')
#tmp2d = sun.Nkmax-etop-1 # zero based
tmp2d = sun.Nkmax - etop # one based
nc.variables[vname][:, ii] = tmp2d
vname = 'Mesh2_face_bottom_layer'
#print '\tVariable: %s...'%vname
#tmp2d = sun.Nkmax-sun.Nk + 1 # zero based
tmp2d = sun.Nkmax - sun.Nk # one based
nc.variables[vname][:, ii] = tmp2d
vname = 'Mesh2_face_top_layer'
#print '\tVariable: %s...'%vname
ctop = sun.loadData(variable='ctop')
#tmp2d = sun.Nkmax-ctop-1 # zero based
tmp2d = sun.Nkmax - ctop # one based
nc.variables[vname][:, ii] = tmp2d
print(72 * '#')
print('\t Finished SUNTANS->UnTRIM conversion')
print(72 * '#')
# close the file
nc.close()
"""
Generates a series of animations
"""
from sunpy import Spatial
import numpy as np
import matplotlib.pyplot as plt
runname='SFBay3D'
ncfile = '%s/%s_0*.nc'%('data',runname)
k = 1 # depth layer
plt.figure()
sun = Spatial(ncfile,klayer=[k],variable='salt')
sun.tstep=np.arange(0,len(sun.time))
sun.loadData()
sun.clim = [28.0,32.0]
sun.animate(vector_overlay=False)
sun.saveanim('plots/%s_salt.mov'%runname)
sun.clim=None
plt.figure()
sun = Spatial(ncfile,klayer=[k],variable='uc')
sun.tstep=np.arange(0,len(sun.time))
sun.loadData()
sun.animate(vector_overlay=False)
sun.saveanim('plots/%s_uc.mov'%runname)
sun.clim=None
plt.figure()
sun = Spatial(ncfile,klayer=[k],variable='eta')
def suntans2untrim(ncfile,outfile,tstart,tend,grdfile=None):
"""
Converts a suntans averages netcdf file into untrim format
for use in particle tracking
"""
####
# Step 1: Load the suntans data object
####
sun = Spatial(ncfile,klayer=[-99])
# Calculate some other variables
sun.de = sun.get_edgevar(sun.dv,method='min')
sun.mark[sun.mark==5]=0
sun.mark[sun.mark==3]=2
sun.facemark = np.zeros((sun.Nc,),dtype=np.int)
# Update the grad variable from the ascii grid file if supplied
if not grdfile == None:
print 'Updating grid with ascii values...'
grd = Grid(grdfile)
sun.grad = grd.grad[:,::-1]
###
# Step 2: Write the grid variables to a netcdf file
###
nc = Dataset(outfile,'w',format='NETCDF4_CLASSIC')
# Global variable
nc.Description = 'UnTRIM history file converted from SUNTANS output'
# Write the dimensions
for dd in untrim_griddims.keys():
if dd == 'time':
nc.createDimension(untrim_griddims[dd],0)
elif dd =='numsides':
nc.createDimension(untrim_griddims[dd],sun.maxfaces)
else:
nc.createDimension(untrim_griddims[dd],sun[dd])
for dd in other_dims:
nc.createDimension(dd,other_dims[dd])
###
# Step 3: Initialize all of the grid variables
###
def create_nc_var(name, dimensions, attdict,data=None, \
dtype='f8',zlib=False,complevel=0,fill_value=999999.0):
tmp=nc.createVariable(name, dtype, dimensions,\
zlib=zlib,complevel=complevel,fill_value=fill_value)
for aa in attdict.keys():
tmp.setncattr(aa,attdict[aa])
if not data==None:
nc.variables[name][:] = data
# Make sure the masked cells have a value of -1
mask = sun['cells'].mask.copy()
sun['cells'][mask]=FILLVALUE
sun['face'][mask]=FILLVALUE
for vv in untrim_gridvars.keys():
vname = untrim_gridvars[vv]
print 'Writing grid variable %s (%s)...'%(vname,vv)
if vv=='time':
continue
# add dz_min attribute to z_r variable
if vv == 'z_r':
ugrid[vname]['attributes'].update({'dz_min':1e-5})
#sun[vv][:]=sun[vv][::-1]
sun[vv][:]=sun['z_w'][0:-1][::-1]
# Reverse the order of grad(???)
if vv=='grad':
sun[vv][:]=sun[vv][:,::-1]
## Fix one-based indexing
#if vv in ['cells','edges','grad']:
# mask = sun[vv][:]==-1
# tmp = sun[vv][:]+1
# tmp[mask]=-1
# #sun[vv][:]=sun[vv][:]+1
# create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
# data=tmp,dtype=ugrid[vname]['dtype'])
create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
data=sun[vv],dtype=ugrid[vname]['dtype'])
# Initialize the two time variables
vname=untrim_gridvars['time']
create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
dtype=ugrid[vname]['dtype'])
vname = 'Mesh2_data_time_string'
create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
dtype=ugrid[vname]['dtype'])
###
# Step 4: Initialize all of the time-varying variables (but don't write)
###
for vname in varnames:
print 'Creating variable %s...'%(vname)
create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
dtype=ugrid[vname]['dtype'],zlib=True,complevel=1,fill_value=999999.)
###
# Step 5: Loop through all of the time steps and write the variables
###
tsteps = sun.getTstep(tstart,tend)
tdays = othertime.DaysSince(sun.time,basetime=datetime(1899,12,31))
for ii, tt in enumerate(tsteps):
# Convert the time to the untrim formats
timestr = datetime.strftime(sun.time[tt],'%Y-%m-%d %H:%M:%S')
print 'Writing data at time %s (%d of %d)...'%(timestr,tt,tsteps[-1])
#Write the time variables
nc.variables['Mesh2_data_time'][ii]=tdays[ii]
nc.variables['Mesh2_data_time_string'][:,ii]=timestr
# Load each variable or calculate it and convert it to the untrim format
sun.tstep=[tt]
###
# Compute a few terms first
eta = sun.loadData(variable='eta' )
U = sun.loadData(variable='U_F' )
dzz = sun.getdzz(eta)
dzf = sun.getdzf(eta)
vname='Mesh2_sea_surface_elevation'
#print '\tVariable: %s...'%vname
nc.variables[vname][:,ii]=eta
vname = 'Mesh2_salinity_3d'
#print '\tVariable: %s...'%vname
tmp3d = sun.loadData(variable='salt' )
nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]
vname = 'Mesh2_vertical_diffusivity_3d'
#print '\tVariable: %s...'%vname
tmp3d = sun.loadData(variable='nu_v' )
nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]
vname = 'h_flow_avg'
#print '\tVariable: %s...'%vname
nc.variables[vname][:,:,ii]=U.swapaxes(0,1)[:,::-1]
vname = 'v_flow_avg'
#print '\tVariable: %s...'%vname
tmp3d = sun.loadData(variable='w' ) * sun.Ac # m^3/s
nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]
# Need to calculate a few terms for the other variables
vname = 'Mesh2_edge_wet_area'
#print '\tVariable: %s...'%vname
#dzf = sun.loadData(variable='dzf')
tmp3d = dzf*sun.df
nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]
vname = 'Mesh2_face_water_volume'
#print '\tVariable: %s...'%vname
#dzz = sun.loadData(variable='dzz')
tmp3d = dzz*sun.Ac
nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]
vname = 'Mesh2_face_wet_area'
#print '\tVariable: %s...'%vname
tmp3d = np.repeat(sun.Ac[np.newaxis,...],sun.Nkmax,axis=0)
nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]
# UnTRIM references from bottom to top i.e.
# k = 0 @ bed ; k = Nkmax-1 @ top
vname = 'Mesh2_edge_bottom_layer'
#print '\tVariable: %s...'%vname
#tmp2d = sun.Nkmax-sun.Nke # zero based
tmp2d = sun.Nkmax-sun.Nke+1 # one based
nc.variables[vname][:,ii]=tmp2d
vname = 'Mesh2_edge_top_layer'
#print '\tVariable: %s...'%vname
etop = sun.loadData(variable='etop')
#tmp2d = sun.Nkmax-etop-1 # zero based
tmp2d = sun.Nkmax-etop # one based
nc.variables[vname][:,ii]=tmp2d
vname = 'Mesh2_face_bottom_layer'
#print '\tVariable: %s...'%vname
#tmp2d = sun.Nkmax-sun.Nk + 1 # zero based
tmp2d = sun.Nkmax-sun.Nk # one based
nc.variables[vname][:,ii]=tmp2d
vname = 'Mesh2_face_top_layer'
#print '\tVariable: %s...'%vname
ctop = sun.loadData(variable='ctop')
#tmp2d = sun.Nkmax-ctop-1 # zero based
tmp2d = sun.Nkmax-ctop # one based
nc.variables[vname][:,ii]=tmp2d
print 72*'#'
print '\t Finished SUNTANS->UnTRIM conversion'
print 72*'#'
# close the file
nc.close()
def volume_integrate(ncfile,varnames,shpfiles,constantdzz=False):
"""
Volume integrate a suntans variable for all time in the domain
specified with a shpfile polygon
"""
# Use numexpr to try and speed things up
import numexpr as ne
# Load the spatial object
sun = Spatial(ncfile,klayer=[-99])
Ac = sun.Ac
# Calculate the mask region
if type(shpfiles) != type([]):
shpfiles = [shpfiles]
masks = []
polynames = []
for shpfile in shpfiles:
mask, maskpoly = maskShpPoly(sun.xv,sun.yv,shpfile)
masks.append(mask)
polynames.append(os.path.splitext(os.path.basename(shpfile))[0])
# Create a dictionary with output info
data={}
for poly in polynames:
data.update({poly:{'V':np.zeros((sun.Nt,)),'time':sun.time}})
for varname in varnames:
data[poly].update({varname:np.zeros((sun.Nt,))})
# Fix dzz for testing
sun.tstep = [0]
dzz = sun.loadData(variable='dzz')
h = ne.evaluate("sum(dzz,axis=0)")
#dzz = np.repeat(sun.dz[:,np.newaxis],sun.Nc,axis=1)
for ii in range(sun.Nt):
sun.tstep = [ii]
print 'Volume integrating for time step: %d of %d...'%(ii,sun.Nt)
# Load the depth and mean age arrays
if not constantdzz:
dzz = sun.loadData(variable='dzz')
# Calculate the total volume
#h = np.sum(dzz,axis=0)
h = ne.evaluate("sum(dzz,axis=0)")
for varname in varnames:
tmp = sun.loadData(variable=varname)
for mask,poly in zip(masks,polynames):
V, A = sun.areaint(h,mask=mask) # Depth*area
data[poly]['V'][ii] = V
# Get the depth-integral
#tmp_dz = sun.depthint(tmp,dz=dzz)
tmp_dz = ne.evaluate("sum(tmp*dzz,axis=0)")
# Calculate the volume-integral
#tmp_dV, A = sun.areaint(tmp_dz,mask=mask)
tmp_dV = ne.evaluate("sum(tmp_dz*Ac*mask)")
data[poly][varname][ii]=tmp_dV/V
return data