def __init__(self,ncfile,**kwargs):
"""
Initialise the suntides class
See sunpy.Spatial class for list of kwargs
"""
self.__dict__.update(kwargs)
Spatial.__init__(self,ncfile,**kwargs)
if self.hasVar('eta_amp'):
print 'Loading existing harmonic data...'
self._loadVars()
else:
# Get the tidal fruequencies
if self.frqnames == None:
# This returns the default frequencies from the uspectra class
self.frq,self.frqnames = uspectra.getTideFreq(Fin=None)
else:
self.frq,self.frqnames = uspectra.getTideFreq(Fin=self.frqnames)
self.Ntide = len(self.frqnames)
self.reftime = datetime(self.baseyear,1,1)
self.Nt = len(self.time)
def __init__(self,infile,**kwargs):
self.__dict__.update(kwargs)
Spatial.__init__(self,infile,**kwargs)
#
if self.is3D:
self.klayer=np.arange(self.kstart,self.Nkmax)
self.Nkmax-=self.kstart
self.Nk -= self.kstart
self.data = np.zeros((self.Nc,self.Nkmax))
self.data = np.ravel(self.data)
if self.maxfaces==3:
self.initTvtk3D()
else:
self.initMixedTvtk3D()
else:
# Initialize the 2D object
self.data = np.zeros((self.Nc,))
self.returnPoints()
if self.maxfaces==3:
self.ug = self.initTvtk2D()
else:
self.ug = self.initMixedTvtk2D()
if self.offscreen:
print 'Using offscreen rendering.'
mlab.options.offscreen=True
def __init__(self,ncfile,**kwargs):
Spatial.__init__(self,ncfile,gridvars=untrim_gridvars,griddims=untrim_griddims,**kwargs)
# Make sure the number of faces array is correct
self.nfaces = np.sum(self.cells.mask==False,axis=1)
self.xy = self.cellxy()
def __init__(self,infile,**kwargs):
self.__dict__.update(kwargs)
Spatial.__init__(self,infile,**kwargs)
#
if self.is3D:
self.klayer=np.arange(self.kstart,self.Nkmax)
self.Nkmax-=self.kstart
self.Nk -= self.kstart
self.data = np.zeros((self.Nc,self.Nkmax))
self.data = np.ravel(self.data)
self.initTvtk3D()
else:
# Initialize the 2D object
self.data = np.zeros((self.Nc,))
self.returnPoints()
if self.maxfaces==3:
self.initTvtk2D()
else:
self.initMixedTvtk2D()
if self.offscreen:
print 'Using offscreen rendering.'
mlab.options.offscreen=True
def __init__(self, ncfile, **kwargs):
Spatial.__init__(self,
ncfile,
gridvars=untrim_gridvars,
griddims=untrim_griddims,
**kwargs)
# Make sure the number of faces array is correct
self.nfaces = np.sum(self.cells.mask == False, axis=1)
self.xy = self.cellxy()
def __init__(self,ncfile,**kwargs):
"""
Initialise the suntides class
See sunpy.Spatial class for list of kwargs
"""
self.__dict__.update(kwargs)
Spatial.__init__(self,ncfile,**kwargs)
self.Nt = len(self.time)
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 __init__(self,ncfile,xpt=None,ypt=None,Npt=100,klayer=[-99],**kwargs):
self.Npt=Npt
Spatial.__init__(self,ncfile,klayer=klayer,**kwargs)
# Load the grid as a hybridgrid
self.grd = GridSearch(self.xp,self.yp,self.cells,nfaces=self.nfaces,\
edges=self.edges,mark=self.mark,grad=self.grad,neigh=self.neigh,\
xv=self.xv,yv=self.yv)
# Find the edge indices along the line
self.update_xy(xpt,ypt)
def __init__(self,ncfile,xpt=None,ypt=None,Npt=100,klayer=[-99],**kwargs):
self.Npt=Npt
Spatial.__init__(self,ncfile,klayer=klayer,**kwargs)
# Load the grid as a hybridgrid
self.grd = GridSearch(self.xp,self.yp,self.cells,nfaces=self.nfaces,\
edges=self.edges,mark=self.mark,grad=self.grad,neigh=self.neigh,\
xv=self.xv,yv=self.yv)
# Find the edge indices along the line
self.update_xy(xpt,ypt)
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 __init__(self,ncfile,xpt=None,ypt=None,Npt=100):
Spatial.__init__(self,ncfile,klayer=[-99])
# Calculate the horizontal coordinates of the slice
self.Npt = Npt
if xpt == None or ypt == None:
self._getXYgraphically()
else:
self.xpt=xpt
self.ypt=ypt
self._getSliceCoords()
# Initialise the slice interpolation object
self._initInterp()
def __init__(self,ncfile,xpt=None,ypt=None,Npt=100):
Spatial.__init__(self,ncfile,klayer=[-99])
# Calculate the horizontal coordinates of the slice
self.Npt = Npt
if xpt == None or ypt == None:
self._getXYgraphically()
else:
self.xpt=xpt
self.ypt=ypt
self._getSliceCoords()
# Initialise the slice interpolation object
self._initInterp()
def on_open_file(self, event):
file_choices = "SUNTANS NetCDF (*.nc)|*.nc*|UnTRIM NetCDF (*.nc)|*.nc*|All Files (*.*)|*.*"
dlg = wx.FileDialog(
self,
message="Open SUNTANS file...",
defaultDir=os.getcwd(),
defaultFile="",
wildcard=file_choices,
style= wx.FD_MULTIPLE)
if dlg.ShowModal() == wx.ID_OK:
self.plot_type='hydro'
path = dlg.GetPaths()
# Initialise the class
if dlg.GetFilterIndex() == 0 or dlg.GetFilterIndex() > 1: #SUNTANS
self.flash_status_message("Opening SUNTANS file: %s" % path)
Spatial.__init__(self,path)
startvar='dv'
if dlg.GetFilterIndex()==1: #UnTRIM
self.flash_status_message("Opening UnTRIMS file: %s" % path)
#Spatial.__init__(self,path,gridvars=untrim_gridvars,griddims=untrim_griddims)
UNTRIMSpatial.__init__(self,path)
startvar='Mesh2_face_depth'
# Populate the drop down menus
vnames = self.listCoordVars()
self.variable_list.SetItems(vnames)
# Update the time drop down list
if self.__dict__.has_key('time'):
self.timestr = [datetime.strftime(tt,'%d-%b-%Y %H:%M:%S') for tt in self.time]
else:
# Assume that it is a harmonic-type file
self.timestr = self.nc.Constituent_Names.split()
self.time_list.SetItems(self.timestr)
# Draw the depth
if startvar in vnames:
self.variable=startvar
self.loadData()
self.create_figure()
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 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 on_open_file(self, event):
file_choices = "SUNTANS NetCDF (*.nc)|*.nc*|UnTRIM NetCDF (*.nc)|*.nc*|All Files (*.*)|*.*"
dlg = wx.FileDialog(
self,
message="Open SUNTANS file...",
defaultDir=os.getcwd(),
defaultFile="",
wildcard=file_choices,
style= wx.FD_MULTIPLE)
if dlg.ShowModal() == wx.ID_OK:
self.plot_type='hydro'
path = dlg.GetPaths()
# Initialise the class
if dlg.GetFilterIndex() == 0 or dlg.GetFilterIndex() > 1: #SUNTANS
self.flash_status_message("Opening SUNTANS file: %s" % path)
Spatial.__init__(self,path)
startvar='dv'
if dlg.GetFilterIndex()==1: #UnTRIM
self.flash_status_message("Opening UnTRIMS file: %s" % path)
#Spatial.__init__(self,path,gridvars=untrim_gridvars,griddims=untrim_griddims)
UNTRIMSpatial.__init__(self,path)
startvar='Mesh2_face_depth'
# Populate the drop down menus
vnames = self.listCoordVars()
self.variable_list.SetItems(vnames)
# Update the time drop down list
if self.__dict__.has_key('time'):
self.timestr = [datetime.strftime(tt,'%d-%b-%Y %H:%M:%S') for tt in self.time]
else:
# Assume that it is a harmonic-type file
self.timestr = self.nc.Constituent_Names.split()
self.time_list.SetItems(self.timestr)
# Draw the depth
if startvar in vnames:
self.variable=startvar
self.loadData()
self.create_figure()
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 __init__(self,ncfile,**kwargs):
"""
Initialize the 3-D grid, etc
"""
self.__dict__.update(kwargs)
if self.is3D:
Spatial.__init__(self,ncfile,klayer=[-99],**kwargs)
# Initialise the 3-D grid
self.init3Dgrid()
else: # surface layer only
print '%s\nRunning in 2D mode...\n%s'%(24*'#',24*'#')
Spatial.__init__(self,ncfile,klayer=['surface'],**kwargs)
self.nActive = self.Nc
#self.klayer=np.arange(0,self.Nkmax)
# Step 2) Initialise the interpolation function
if self.is3D:
if self.interp_method in ('idw','nearest'):
self.UVWinterp = interp3D(self.xv3d,self.yv3d,self.zv3d,method=self.interp_method)
# Interpolation function to find free surface and seabed
self.Hinterp = interp3D(self.xv,self.yv,0*self.xv,method='nearest')
elif self.interp_method == 'mesh':
if self.interp_meshmethod == 'nearest':
self.UVWinterp = \
interp3Dmesh(self.xp,self.yp,-self.z_w,self.cells,\
self.nfaces,self.mask3D,method='nearest')
elif self.interp_meshmethod == 'linear':
self.UVWinterp =\
interp3Dmesh(self.xp,self.yp,-self.z_w,self.cells,self.nfaces,\
self.mask3D,method='linear',grdfile=self.ncfile)
else:
# Surface layer use nearest point only for now
self.UVWinterp = interp2Dmesh(self.xp,self.yp,self.cells,\
self.nfaces,method='nearest')
def __init__(self,ncfile,**kwargs):
"""
Initialize the 3-D grid, etc
"""
self.__dict__.update(kwargs)
if self.is3D:
Spatial.__init__(self,ncfile,klayer=[-99],**kwargs)
# Initialise the 3-D grid
self.init3Dgrid()
else: # surface layer only
print '%s\nRunning in 2D mode...\n%s'%(24*'#',24*'#')
Spatial.__init__(self,ncfile,klayer=['surface'],**kwargs)
self.nActive = self.Nc
#self.klayer=np.arange(0,self.Nkmax)
# Step 2) Initialise the interpolation function
if self.is3D:
if self.interp_method in ('idw','nearest'):
self.UVWinterp = interp3D(self.xv3d,self.yv3d,self.zv3d,method=self.interp_method)
# Interpolation function to find free surface and seabed
self.Hinterp = interp3D(self.xv,self.yv,0*self.xv,method='nearest')
elif self.interp_method == 'mesh':
if self.interp_meshmethod == 'nearest':
self.UVWinterp = \
interp3Dmesh(self.xp,self.yp,-self.z_w,self.cells,\
self.nfaces,self.mask3D,method='nearest')
elif self.interp_meshmethod == 'linear':
self.UVWinterp =\
interp3Dmesh(self.xp,self.yp,-self.z_w,self.cells,self.nfaces,\
self.mask3D,method='linear',grdfile=self.ncfile)
else:
# Surface layer use nearest point only for now
self.UVWinterp = interp2Dmesh(self.xp,self.yp,self.cells,\
self.nfaces,method='nearest')
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 ncload(nc,variable,tt):
if variable=='agemean':
ac = nc.variables['agec'][tt,klayer,:]
aa = nc.variables['agealpha'][tt,klayer,:]
tmp = aa/ac
tmp[ac<1e-12]=0.
return tmp/86400.
if variable=='area':
eta = nc.variables['eta'][tt,:]
dzf = self.getdzf(eta)
dzf = Spatial.getdzf(self,eta)
return self.df*dzf
else:
if self.hasDim(variable,self.griddims['Nk']): # 3D
return nc.variables[variable][tt,klayer,:]
else:
return nc.variables[variable][tt,:]
def ncload(nc,variable,tt):
if variable=='agemean':
ac = nc.variables['agec'][tt,klayer,:]
aa = nc.variables['agealpha'][tt,klayer,:]
tmp = aa/ac
tmp[ac<1e-12]=0.
return tmp/86400.
if variable=='area':
eta = nc.variables['eta'][tt,:]
dzf = self.getdzf(eta)
dzf = Spatial.getdzf(self,eta)
return self.df*dzf
else:
if self.hasDim(variable,self.griddims['Nk']): # 3D
return nc.variables[variable][tt,klayer,:]
else:
return nc.variables[variable][tt,:]
def surface(self,clim=None,**kwargs):
"""
Surface plot of the scalar in the 'data' attribute with mayavi
Works on the 2D and 3D data
"""
if clim==None:
clim = [self.data.min(), self.data.max()]
# Create a new scene if there isn't one
if not self.__dict__.has_key('fig'):
self.newscene()
src = mlab.pipeline.add_dataset(self.ug)
self.h=mlab.pipeline.surface(src,vmin=clim[0],vmax=clim[1],**kwargs)
# Add a colorbar if the isn't one
if not self.__dict__.has_key('cb'):
self.colorbar()
# Add a title if there isn't one
if not self.__dict__.has_key('title'):
self.title=mlab.title(Spatial.genTitle(self),height=0.95,size=0.15)
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()
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.'
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
"""
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 getdzf(self,eta):
""" Get the cell thickness along each edge of the slice"""
dzf = Spatial.getdzf(self,eta,j=self.j)
dzf[self.maskslice]=0
return dzf
def getdzf(self,eta):
""" Get the cell thickness along each edge of the slice"""
dzf = Spatial.getdzf(self,eta,j=self.j)
dzf[self.maskslice]=0
return dzf
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 __init__(self,ncfile,**kwargs):
"""
Calculates the energy variables from suntans output
"""
# Initialize the spatial class
Spatial.__init__(self,ncfile,klayer=[-99])
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
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
import numpy as np import matplotlib.pyplot as plt from datetime import timedelta from scipy.integrate import cumtrapz import pdb ### # Inputs ncfile = 'data/Heatflux_AVG.0' #cellindex=range(0,9) cellindex=[4] t0 = 0 ### sun = Spatial(ncfile,klayer=[-99]) sun.tstep = range(t0,sun.Nt) time = sun.time[sun.tstep] # Constants dt = sun.globalatts['dt']*sun.globalatts['ntaverage'] RHO0 = 1000.0 Cp = 4186.0 fac = (RHO0*Cp) area = sun.Ac[cellindex] sumarea = np.sum(area) depth = sun.dv[cellindex] volume = area*depth # Cell volume
wsun = np.load("calcs/mask_bayshelf_weights_radial.npz")["wsun"]
wsunu = wsun[:, 1:] # don't want to interpolate this
wsunv = wsun[1:, :]
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
"""
Generates a series of animations from the suntans output
"""
from sunpy import Spatial
import numpy as np
import matplotlib.pyplot as plt
ncfile = 'rundata/Estuary_SUNTANS_00*nc'
plt.figure()
sun = Spatial(ncfile,klayer=[1],variable='salt')
sun.tstep=np.arange(0,len(sun.time))
sun.loadData()
sun.animate(vector_overlay=False)
sun.saveanim('plots/salt.mov')
plt.figure()
sun = Spatial(ncfile,klayer=[1],variable='temp')
sun.tstep=np.arange(0,len(sun.time))
sun.loadData()
sun.animate(vector_overlay=False)
sun.saveanim('plots/temp.mov')
plt.figure()
sun = Spatial(ncfile,variable='tau_y')
sun.tstep=np.arange(0,len(sun.time))
sun.loadData()
sun.animate(vector_overlay=False)
sun.saveanim('plots/tau_y.mov')