def _initInterp(self):
"""
Initialise the interpolant
Finds the horizontal indices of the slice points and
constructs the 3D mask array
"""
# Find the cell index of each point along the slice
self.Tri = GridSearch(self.xp, self.yp, self.cells)
self.cellind = self.Tri(self.xslice, self.yslice)
klayer, Nkmax = self.get_klayer()
# Construct the 3D coordinate arrays
self.xslice = np.repeat(self.xslice.reshape((1, self.Npt)),
self.Nkmax,
axis=0)
self.yslice = np.repeat(self.yslice.reshape((1, self.Npt)),
self.Nkmax,
axis=0)
self.distslice = np.repeat(self.distslice.reshape((1, self.Npt)),
self.Nkmax,
axis=0)
self.zslice = np.repeat(-self.z_r[klayer].reshape((self.Nkmax, 1)),
self.Npt,
axis=1)
# Construct the mask array
self.calc_mask()
# Get the bathymetry along the slice
self.hslice = -self.dv[self.cellind]
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 __call__(self,X,Y,Z,data,update=True):
if update:
# The Update the cell index using GridSearch class
if not self.__dict__.has_key('cellind'):
GridSearch.__call__(self,X,Y)
else:
if np.sum(np.abs(X-self.xpt))>0+1e-8:
self.updatexy(X,Y)
# Return the nearest data point (... for now)
if self.method == 'nearest':
dataout = data[self.cellind]
#if self.method == 'linear':
# dataout = self.lininterp(X,Y,Z,data,kind)
# Mask bogey points
dataout[self.cellind==-1]=0.0
return dataout
def __call__(self,X,Y,Z,data,update=True):
if update:
# The Update the cell index using GridSearch class
if 'cellind' not in self.__dict__:
GridSearch.__call__(self,X,Y)
else:
if np.sum(np.abs(X-self.xpt))>0+1e-8:
self.updatexy(X,Y)
# Return the nearest data point (... for now)
if self.method == 'nearest':
dataout = data[self.cellind]
#if self.method == 'linear':
# dataout = self.lininterp(X,Y,Z,data,kind)
# Mask bogey points
dataout[self.cellind==-1]=0.0
return dataout
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,x,y,z,cells,nfaces,mask,method='nearest',grdfile=None):
self.method=method
# Initialise the trisearch array
GridSearch.__init__(self,x,y,cells,nfaces=nfaces,force_inside=True)
if self.method == 'linear':
Grid.__init__(self,grdfile)
self.datatmp = np.zeros(mask.shape,dtype=np.double)
self.z = np.sort(z)
self.z[-1]=10.0 # Set the surface layer to large
self.Nkmax = z.size-1
self.mask3d = mask
self.maskindex = -1*np.ones(self.mask3d.shape,dtype=np.int32)
rr=0
for ii in range(self.mask3d.shape[0]):
for jj in range(self.mask3d.shape[1]):
if self.mask3d[ii,jj]:
self.maskindex[ii,jj]=rr
rr+=1
def __init__(self,x,y,z,cells,nfaces,mask,method='nearest',grdfile=None):
self.method=method
# Initialise the trisearch array
GridSearch.__init__(self,x,y,cells,nfaces=nfaces,force_inside=True)
if self.method == 'linear':
Grid.__init__(self,grdfile)
self.datatmp = np.zeros(mask.shape,dtype=np.double)
self.z = np.sort(z)
self.z[-1]=10.0 # Set the surface layer to large
self.Nkmax = z.size-1
self.mask3d = mask
self.maskindex = -1*np.ones(self.mask3d.shape,dtype=np.int32)
rr=0
for ii in range(self.mask3d.shape[0]):
for jj in range(self.mask3d.shape[1]):
if self.mask3d[ii,jj]:
self.maskindex[ii,jj]=rr
rr+=1
def __call__(self,X,Y,Z,data,update=True):
if update:
# The Update the cell index using TriSearch class
if not self.__dict__.has_key('cellind'):
#print ' Finding initial particle index...'
GridSearch.__call__(self,X,Y)
else:
if np.sum(np.abs(X-self.xpt))>0+1e-8:
#print ' updating location index...'
self.updatexy(X,Y)
# Find the k-index
kind=self.z.searchsorted(Z)
kind = self.Nkmax - kind
kind[kind>=self.Nkmax-1] = self.Nkmax-1
ind = self.maskindex[kind,self.cellind]
#maskpts = self.mask3d[kind,self.cellind]
#if np.sum(maskpts==False)>10:
# pdb.set_trace()
#
#ind[maskpts == False] = 0
# Return the nearest data point (... for now)
if self.method == 'nearest':
dataout = data[ind]
if self.method == 'linear':
dataout = self.lininterp(X,Y,Z,data,kind)
# Mask bogey points
#dataout[maskpts==False]=0.0
dataout[ind==-1]=0 # these are the masked points
dataout[self.cellind==-1]=0.0
return dataout
def __call__(self,X,Y,Z,data,update=True):
if update:
# The Update the cell index using TriSearch class
if 'cellind' not in self.__dict__:
#print ' Finding initial particle index...'
GridSearch.__call__(self,X,Y)
else:
if np.sum(np.abs(X-self.xpt))>0+1e-8:
#print ' updating location index...'
self.updatexy(X,Y)
# Find the k-index
kind=self.z.searchsorted(Z)
kind = self.Nkmax - kind
kind[kind>=self.Nkmax-1] = self.Nkmax-1
ind = self.maskindex[kind,self.cellind]
#maskpts = self.mask3d[kind,self.cellind]
#if np.sum(maskpts==False)>10:
# pdb.set_trace()
#
#ind[maskpts == False] = 0
# Return the nearest data point (... for now)
if self.method == 'nearest':
dataout = data[ind]
if self.method == 'linear':
dataout = self.lininterp(X,Y,Z,data,kind)
# Mask bogey points
#dataout[maskpts==False]=0.0
dataout[ind==-1]=0 # these are the masked points
dataout[self.cellind==-1]=0.0
return dataout
class SliceEdge(Slice):
"""
Slice suntans edge-based data at all edges near a line
Used for e.g. flux calculations along a profile
"""
edgemethod=1
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 update_xy(self,xpt,ypt):
"""
Updates the x and y coordinate info in the object
"""
if xpt == None or ypt == None:
self._getXYgraphically()
else:
self.xpt=xpt
self.ypt=ypt
self._getSliceCoords(kind='linear')
# List of the edge indices
self.j,self.nodelist =\
self.get_edgeindices(self.xslice,self.yslice,method=self.edgemethod)
self.nslice = len(self.j)
# Update the x and y axis of the slice
self.xslice=self.xp[self.nodelist]
self.yslice=self.yp[self.nodelist]
self.zslice = -self.z_r
self._getDistCoords()
self.edgexy()
# The x and y arrys need to be resized
self.xslice = 0.5*(self.xslice[1:]+self.xslice[0:-1])
self.yslice = 0.5*(self.yslice[1:]+self.yslice[0:-1])
self.distslice = 0.5*(self.distslice[1:]+self.distslice[0:-1])
# Get the mask
self.calc_mask()
# Calculate the area
self.area = self.calc_area()
# Calculate the normnal
self.ne1, self.ne2, self.enormal = self.calc_normal(self.nodelist,self.j)
# Get the bathymetry along the slice
de = self.get_edgevar(self.dv)
self.hslice = -de[self.j]
def loadData(self,variable=None,setunits=True,method='mean'):
"""
Load the specified suntans variable data as a vector
Overloaded method for edge slicing - it is quicker to load time step by
time step in a loop.
method: edge interpolation method - 'mean', 'max'
"""
nc = self.nc
if variable is None:
variable=self.variable
if setunits:
try:
self.long_name = nc.variables[variable].long_name
self.units= nc.variables[variable].units
except:
self.long_name = ''
self.units=''
j=self.j
# Check if cell-centered variable
is3D=True
isCell=False
if self.hasVar(variable):
if self.hasDim(variable,self.griddims['Ne']):
isCell=False
elif self.hasDim(variable,self.griddims['Nc']):
isCell=True
# Check if 3D
if self.hasDim(variable,self.griddims['Nk']) or\
self.hasDim(variable,'Nkw'): # 3D
is3D=True
else:
is3D=False
else:
isCell=True
if isCell:
nc1 = self.grad[j,0].copy()
nc2 = self.grad[j,1].copy()
# check for edges (use logical indexing)
ind1 = nc1==-1
nc1[ind1]=nc2[ind1]
ind2 = nc2==-1
nc2[ind2]=nc1[ind2]
klayer,Nkmax = self.get_klayer()
if self.hasDim(variable,'Nkw'): # vertical velocity
Nkmax +=1
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,:]
# For loop where the data is extracted
nt = len(self.tstep)
ne = len(self.j)
if is3D==True:
self.data = np.zeros((nt,Nkmax,ne))
else:
self.data = np.zeros((nt,ne))
for ii,tt in enumerate(self.tstep):
#tmp=nc.variables[variable][tt,:,:]
tmp = ncload(nc,variable,tt)
# Return the mean for cell-based variables
if isCell:
if method == 'mean':
self.data[ii,...] = 0.5*(tmp[...,nc1]+tmp[...,nc2])
elif method == 'max':
tmp2 = np.dstack((tmp[...,nc1], tmp[...,nc2]))
self.data[ii,...] =tmp2.max(axis=-1)
else:
self.data[ii,...]=tmp[...,self.j]
# Mask 3D data
if is3D:
maskval=0
self.data[ii,self.maskslice]=maskval
#fillval = 999999.0
#self.mask = self.data==fillval
#self.data[self.mask]=0.
self.data[self.data==self._FillValue]=0.
self.data = self.data.squeeze()
return self.data
def edgexy(self):
"""
Nx2 vectors outlining each cell in the edge slice
"""
def closePoly(xp,node,k):
return np.array([ [xp[node],\
xp[node+1],xp[node+1], xp[node],xp[node]],\
[-self.z_w[k],-self.z_w[k],-self.z_w[k+1],-self.z_w[k+1],-self.z_w[k]],\
]).T
self.xye = [closePoly(self.distslice,jj,kk) for kk in range(self.Nkmax) \
for jj in range(len(self.j)) ]
def calc_normal(self,nodelist,j):
"""
Calculate the edge normal
"""
# Calculate the unit normal along the edge
P1 = GPoint(self.xp[nodelist][0:-1],self.yp[nodelist][0:-1])
P2 = GPoint(self.xp[nodelist][1:],self.yp[nodelist][1:])
L = Line(P1,P2)
ne1,ne2 = L.unitnormal()
# Compute the unique normal of the dot product
enormal = np.round(self.n1[j]*ne1 +\
self.n2[j]*ne2)
return ne1,ne2,enormal
def mean(self,phi,axis='time'):
"""
Calculate the mean of the sliced data along an axis
axis: time, depth, area
time : returns the time mean. size= (Nk, Nj)
depth: returns the time and spatial mean. Size = (Nk)
area: returns the area mean. Size = (Nt)
"""
if axis=='time':
return np.mean(phi,axis=0)
elif axis=='area':
area_norm = self.area / self.area.sum()
return np.sum( np.sum(phi*area_norm,axis=-1),axis=-1)
elif axis=='depth':
dx = self.df[self.j]
dx_norm = dx / dx.sum()
return np.sum( self.mean(phi,axis='time')*dx_norm,axis=-1)
def plot(self,z,titlestr=None,**kwargs):
"""
Pcolor plot of the slice
"""
if self.clim is None:
self.clim=[]
self.clim.append(np.min(z))
self.clim.append(np.max(z))
# Set the xy limits
xlims=[self.distslice.min(),self.distslice.max()]
ylims=[-self.z_w.max(),-self.z_w.min()]
self.fig,self.ax,self.patches,self.cb=unsurf(self.xye,z.ravel(),xlim=xlims,ylim=ylims,\
clim=self.clim,**kwargs)
self.ax.set_aspect('auto')
def plotedges(self,color='m',**kwargs):
"""
plot for testing
"""
self.plotmesh()
#plt.plot(self.edgeline.xy,'r')
for ee in self.j:
plt.plot([self.xp[self.edges[ee,0]],self.xp[self.edges[ee,1]]],\
[self.yp[self.edges[ee,0]],self.yp[self.edges[ee,1]]],color=color,\
**kwargs)
def calc_area(self,eta=None):
"""
Calculate thee cross-sectional area of each face
"""
if eta is None:
eta = np.zeros((self.nslice,)) # Assumes the free-surface is zero
dzf = self.getdzf(eta)
area = dzf * self.df[self.j]
area[self.maskslice]=0
return area
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 get_width(self):
"""
Calculate the width of each edge as a 2d array
Missing cells are masked
"""
df = self.df[self.j]
width = np.ones((self.Nkmax,1)) * df[np.newaxis,:]
width[self.maskslice]=0
return width
def calc_mask(self):
""" Construct the mask array"""
klayer,Nkmax=self.get_klayer()
self.maskslice = np.zeros((Nkmax,len(self.j)),dtype=np.bool)
for k,kk in enumerate(klayer):
for ii,j in enumerate(self.j):
if kk >= self.Nke[j]:
self.maskslice[k,ii]=True
def get_edgeindices(self,xpt,ypt,method=1, abortedge=False):
"""
Return the indices of the edges (in order) along the line
method - method for line finding algorithm
0 - closest point to line
1 - closest point without doing a u-turn
abortedge - Set true to abort when slice hits a boundary
"""
# Load the line as a shapely object
#edgeline = asLineString([self.xslice,self.yslice])
Npt = xpt.shape[0]
xyline = [(xpt[ii],ypt[ii]) for ii in range(Npt)]
self.edgeline = LineString(xyline)
# Find the nearest grid Node to the start and end of the line
xy_1 = np.vstack((xpt[0],ypt[0])).T
node0 = self.grd.findnearest(xy_1)
xy_2 = np.vstack((xpt[-1],ypt[-1])).T
endnode = self.grd.findnearest(xy_2)
# This is the list containing all edge nodes
nodelist = [node0[0]]
def connecting_nodes(node,nodelist):
""" finds the nodes connecting to the node"""
edges = self.grd.pnt2edges(node)
cnodes = []
for ee in edges:
for nn in self.grd.edges[ee]:
if nn not in nodelist:
cnodes.append(nn)
return cnodes
def min_dist(nodes,line):
"""Returns the index of the node with the minimum distance
to the line"""
# Convert all nodes to a point object
points = [Point((self.xp[nn],self.yp[nn])) for nn in nodes]
# Calculate the distance
dist = [line.distance(pp) for pp in points]
for ii,dd in enumerate(dist):
if dd == min(dist):
return nodes[ii]
def min_dist_line(cnode,nodes,line):
"""Returns the index of the node with the minimum distance
to the line"""
# Convert all nodes to a point object
points = [Point((0.5*(self.xp[nn]+self.xp[cnode]),\
0.5*(self.yp[nn]+self.yp[cnode]))) for nn in nodes]
#lines = [LineString([(self.xp[cnode],self.yp[cnode]),\
# (self.xp[nn],self.yp[nn])]) for nn in nodes]
# Calculate the distance
dist = [line.distance(pp) for pp in points]
for ii,dd in enumerate(dist):
if dd == min(dist):
return nodes[ii]
def min_dist_angle(cnode,nodes,line):
"""Returns the index of the node with the minimum distance
to the line"""
# Convert all nodes to a point object
points = [Point((0.5*(self.xp[nn]+self.xp[cnode]),\
0.5*(self.yp[nn]+self.yp[cnode]))) for nn in nodes]
# Calculate the distance
dist = [line.distance(pp) for pp in points]
dist = np.array(dist)
# Calculate the angle along the line of the new coordinate
def calc_ang(x1,x2,y1,y2):
return np.arctan2( (y2-y1),(x2-x1) )
angle1 = [calc_ang(self.xp[cnode],self.xp[nn],\
self.yp[cnode],self.yp[nn]) for nn in nodes]
# Calculate the heading of the line near the two points
def calc_heading(P1,P2,L):
d1 = L.project(P1)
d2 = L.project(P2)
if d1 <= d2:
P3 = L.interpolate(d1)
P4 = L.interpolate(d2)
else:
P3 = L.interpolate(d2)
P4 = L.interpolate(d1)
return calc_ang(P3.xy[0][0],P4.xy[0][0],P3.xy[1][0],P4.xy[1][0])
P1 = Point((self.xp[cnode],self.yp[cnode]))
angle2 = [calc_heading(P1,Point( (self.xp[nn],self.yp[nn]) ),line) \
for nn in nodes]
angdiff = np.array(angle2) - np.array(angle1)
# Use the minimum distance unless the point is a u-turn
rank = np.argsort(dist)
for nn in range(dist.shape[0]):
if np.abs(angdiff[rank[nn]]) <= np.pi/2:
return nodes[rank[nn]]
# if they all u-turn return the min dist
return nodes[rank[0]]
# Loop through and find all of the closest points to the line
MAXITER=10000
for ii in range(MAXITER):
cnodes = connecting_nodes(nodelist[-1],nodelist)
#if method==0:
# newnode = min_dist(cnodes,self.edgeline)
if method==0:
newnode = min_dist_line(nodelist[-1],cnodes,self.edgeline)
elif method==1:
newnode = min_dist_angle(nodelist[-1],cnodes,self.edgeline)
#print 'Found new node: %d...'%newnode
if newnode is None:
break
if ii>1 and abortedge:
if self.mark[self.grd.find_edge([newnode,nodelist[-1]])] not in [0,5]:
print 'Warning: reached a boundary cell. Aborting edge finding routine'
break
nodelist.append(newnode)
if newnode == endnode:
#print 'Reached end node.'
break
# Return the list of edges connecting all of the nodes
return [self.grd.find_edge([nodelist[ii],nodelist[ii+1]]) for ii in\
range(len(nodelist)-1)], nodelist
def __init__(self,x,y,cells,nfaces,method='nearest',grdfile=None):
self.method=method
# Initialise the trisearch array
GridSearch.__init__(self,x,y,cells,nfaces=nfaces,force_inside=True)
class SliceEdge(Slice):
"""
Slice suntans edge-based data at all edges near a line
Used for e.g. flux calculations along a profile
"""
edgemethod = 1
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 update_xy(self, xpt, ypt):
"""
Updates the x and y coordinate info in the object
"""
if xpt == None or ypt == None:
self._getXYgraphically()
else:
self.xpt = xpt
self.ypt = ypt
self._getSliceCoords(kind='linear')
# List of the edge indices
self.j,self.nodelist =\
self.get_edgeindices(self.xslice,self.yslice,method=self.edgemethod)
self.nslice = len(self.j)
# Update the x and y axis of the slice
self.xslice = self.xp[self.nodelist]
self.yslice = self.yp[self.nodelist]
self.zslice = -self.z_r
self._getDistCoords()
self.edgexy()
# The x and y arrys need to be resized
self.xslice = 0.5 * (self.xslice[1:] + self.xslice[0:-1])
self.yslice = 0.5 * (self.yslice[1:] + self.yslice[0:-1])
self.distslice = 0.5 * (self.distslice[1:] + self.distslice[0:-1])
# Get the mask
self.calc_mask()
# Calculate the area
self.area = self.calc_area()
# Calculate the normnal
self.ne1, self.ne2, self.enormal = self.calc_normal(
self.nodelist, self.j)
# Get the bathymetry along the slice
de = self.get_edgevar(self.dv)
self.hslice = -de[self.j]
def loadData(self, variable=None, setunits=True, method='mean'):
"""
Load the specified suntans variable data as a vector
Overloaded method for edge slicing - it is quicker to load time step by
time step in a loop.
method: edge interpolation method - 'mean', 'max'
"""
nc = self.nc
if variable is None:
variable = self.variable
if setunits:
try:
self.long_name = nc.variables[variable].long_name
self.units = nc.variables[variable].units
except:
self.long_name = ''
self.units = ''
j = self.j
# Check if cell-centered variable
is3D = True
isCell = False
if self.hasVar(variable):
if self.hasDim(variable, self.griddims['Ne']):
isCell = False
elif self.hasDim(variable, self.griddims['Nc']):
isCell = True
# Check if 3D
if self.hasDim(variable,self.griddims['Nk']) or\
self.hasDim(variable,'Nkw'): # 3D
is3D = True
else:
is3D = False
else:
isCell = True
if isCell:
nc1 = self.grad[j, 0].copy()
nc2 = self.grad[j, 1].copy()
# check for edges (use logical indexing)
ind1 = nc1 == -1
nc1[ind1] = nc2[ind1]
ind2 = nc2 == -1
nc2[ind2] = nc1[ind2]
klayer, Nkmax = self.get_klayer()
if self.hasDim(variable, 'Nkw'): # vertical velocity
Nkmax += 1
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, :]
# For loop where the data is extracted
nt = len(self.tstep)
ne = len(self.j)
if is3D == True:
self.data = np.zeros((nt, Nkmax, ne))
else:
self.data = np.zeros((nt, ne))
for ii, tt in enumerate(self.tstep):
#tmp=nc.variables[variable][tt,:,:]
tmp = ncload(nc, variable, tt)
# Return the mean for cell-based variables
if isCell:
if method == 'mean':
self.data[ii, ...] = 0.5 * (tmp[..., nc1] + tmp[..., nc2])
elif method == 'max':
tmp2 = np.dstack((tmp[..., nc1], tmp[..., nc2]))
self.data[ii, ...] = tmp2.max(axis=-1)
else:
self.data[ii, ...] = tmp[..., self.j]
# Mask 3D data
if is3D:
maskval = 0
self.data[ii, self.maskslice] = maskval
#fillval = 999999.0
#self.mask = self.data==fillval
#self.data[self.mask]=0.
self.data[self.data == self._FillValue] = 0.
self.data = self.data.squeeze()
return self.data
def edgexy(self):
"""
Nx2 vectors outlining each cell in the edge slice
"""
def closePoly(xp, node, k):
return np.array([ [xp[node],\
xp[node+1],xp[node+1], xp[node],xp[node]],\
[-self.z_w[k],-self.z_w[k],-self.z_w[k+1],-self.z_w[k+1],-self.z_w[k]],\
]).T
self.xye = [closePoly(self.distslice,jj,kk) for kk in range(self.Nkmax) \
for jj in range(len(self.j)) ]
def calc_normal(self, nodelist, j):
"""
Calculate the edge normal
"""
# Calculate the unit normal along the edge
P1 = GPoint(self.xp[nodelist][0:-1], self.yp[nodelist][0:-1])
P2 = GPoint(self.xp[nodelist][1:], self.yp[nodelist][1:])
L = Line(P1, P2)
ne1, ne2 = L.unitnormal()
# Compute the unique normal of the dot product
enormal = np.round(self.n1[j]*ne1 +\
self.n2[j]*ne2)
return ne1, ne2, enormal
def mean(self, phi, axis='time'):
"""
Calculate the mean of the sliced data along an axis
axis: time, depth, area
time : returns the time mean. size= (Nk, Nj)
depth: returns the time and spatial mean. Size = (Nk)
area: returns the area mean. Size = (Nt)
"""
if axis == 'time':
return np.mean(phi, axis=0)
elif axis == 'area':
area_norm = self.area / self.area.sum()
return np.sum(np.sum(phi * area_norm, axis=-1), axis=-1)
elif axis == 'depth':
dx = self.df[self.j]
dx_norm = dx / dx.sum()
return np.sum(self.mean(phi, axis='time') * dx_norm, axis=-1)
def plot(self, z, titlestr=None, **kwargs):
"""
Pcolor plot of the slice
"""
if self.clim is None:
self.clim = []
self.clim.append(np.min(z))
self.clim.append(np.max(z))
# Set the xy limits
xlims = [self.distslice.min(), self.distslice.max()]
ylims = [-self.z_w.max(), -self.z_w.min()]
self.fig,self.ax,self.patches,self.cb=unsurf(self.xye,z.ravel(),xlim=xlims,ylim=ylims,\
clim=self.clim,**kwargs)
self.ax.set_aspect('auto')
def plotedges(self, color='m', **kwargs):
"""
plot for testing
"""
self.plotmesh()
#plt.plot(self.edgeline.xy,'r')
for ee in self.j:
plt.plot([self.xp[self.edges[ee,0]],self.xp[self.edges[ee,1]]],\
[self.yp[self.edges[ee,0]],self.yp[self.edges[ee,1]]],color=color,\
**kwargs)
def calc_area(self, eta=None):
"""
Calculate thee cross-sectional area of each face
"""
if eta is None:
eta = np.zeros((self.nslice, )) # Assumes the free-surface is zero
dzf = self.getdzf(eta)
area = dzf * self.df[self.j]
area[self.maskslice] = 0
return area
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 get_width(self):
"""
Calculate the width of each edge as a 2d array
Missing cells are masked
"""
df = self.df[self.j]
width = np.ones((self.Nkmax, 1)) * df[np.newaxis, :]
width[self.maskslice] = 0
return width
def calc_mask(self):
""" Construct the mask array"""
klayer, Nkmax = self.get_klayer()
self.maskslice = np.zeros((Nkmax, len(self.j)), dtype=np.bool)
for k, kk in enumerate(klayer):
for ii, j in enumerate(self.j):
if kk >= self.Nke[j]:
self.maskslice[k, ii] = True
def get_edgeindices(self, xpt, ypt, method=1, abortedge=False):
"""
Return the indices of the edges (in order) along the line
method - method for line finding algorithm
0 - closest point to line
1 - closest point without doing a u-turn
abortedge - Set true to abort when slice hits a boundary
"""
# Load the line as a shapely object
#edgeline = asLineString([self.xslice,self.yslice])
Npt = xpt.shape[0]
xyline = [(xpt[ii], ypt[ii]) for ii in range(Npt)]
self.edgeline = LineString(xyline)
# Find the nearest grid Node to the start and end of the line
xy_1 = np.vstack((xpt[0], ypt[0])).T
node0 = self.grd.findnearest(xy_1)
xy_2 = np.vstack((xpt[-1], ypt[-1])).T
endnode = self.grd.findnearest(xy_2)
# This is the list containing all edge nodes
nodelist = [node0[0]]
def connecting_nodes(node, nodelist):
""" finds the nodes connecting to the node"""
edges = self.grd.pnt2edges(node)
cnodes = []
for ee in edges:
for nn in self.grd.edges[ee]:
if nn not in nodelist:
cnodes.append(nn)
return cnodes
def min_dist(nodes, line):
"""Returns the index of the node with the minimum distance
to the line"""
# Convert all nodes to a point object
points = [Point((self.xp[nn], self.yp[nn])) for nn in nodes]
# Calculate the distance
dist = [line.distance(pp) for pp in points]
for ii, dd in enumerate(dist):
if dd == min(dist):
return nodes[ii]
def min_dist_line(cnode, nodes, line):
"""Returns the index of the node with the minimum distance
to the line"""
# Convert all nodes to a point object
points = [Point((0.5*(self.xp[nn]+self.xp[cnode]),\
0.5*(self.yp[nn]+self.yp[cnode]))) for nn in nodes]
#lines = [LineString([(self.xp[cnode],self.yp[cnode]),\
# (self.xp[nn],self.yp[nn])]) for nn in nodes]
# Calculate the distance
dist = [line.distance(pp) for pp in points]
for ii, dd in enumerate(dist):
if dd == min(dist):
return nodes[ii]
def min_dist_angle(cnode, nodes, line):
"""Returns the index of the node with the minimum distance
to the line"""
# Convert all nodes to a point object
points = [Point((0.5*(self.xp[nn]+self.xp[cnode]),\
0.5*(self.yp[nn]+self.yp[cnode]))) for nn in nodes]
# Calculate the distance
dist = [line.distance(pp) for pp in points]
dist = np.array(dist)
# Calculate the angle along the line of the new coordinate
def calc_ang(x1, x2, y1, y2):
return np.arctan2((y2 - y1), (x2 - x1))
angle1 = [calc_ang(self.xp[cnode],self.xp[nn],\
self.yp[cnode],self.yp[nn]) for nn in nodes]
# Calculate the heading of the line near the two points
def calc_heading(P1, P2, L):
d1 = L.project(P1)
d2 = L.project(P2)
if d1 <= d2:
P3 = L.interpolate(d1)
P4 = L.interpolate(d2)
else:
P3 = L.interpolate(d2)
P4 = L.interpolate(d1)
return calc_ang(P3.xy[0][0], P4.xy[0][0], P3.xy[1][0],
P4.xy[1][0])
P1 = Point((self.xp[cnode], self.yp[cnode]))
angle2 = [calc_heading(P1,Point( (self.xp[nn],self.yp[nn]) ),line) \
for nn in nodes]
angdiff = np.array(angle2) - np.array(angle1)
# Use the minimum distance unless the point is a u-turn
rank = np.argsort(dist)
for nn in range(dist.shape[0]):
if np.abs(angdiff[rank[nn]]) <= np.pi / 2:
return nodes[rank[nn]]
# if they all u-turn return the min dist
return nodes[rank[0]]
# Loop through and find all of the closest points to the line
MAXITER = 10000
for ii in range(MAXITER):
cnodes = connecting_nodes(nodelist[-1], nodelist)
#if method==0:
# newnode = min_dist(cnodes,self.edgeline)
if method == 0:
newnode = min_dist_line(nodelist[-1], cnodes, self.edgeline)
elif method == 1:
newnode = min_dist_angle(nodelist[-1], cnodes, self.edgeline)
#print 'Found new node: %d...'%newnode
if newnode is None:
break
if ii > 1 and abortedge:
if self.mark[self.grd.find_edge([newnode,
nodelist[-1]])] not in [0, 5]:
print(
'Warning: reached a boundary cell. Aborting edge finding routine'
)
break
nodelist.append(newnode)
if newnode == endnode:
#print 'Reached end node.'
break
# Return the list of edges connecting all of the nodes
return [self.grd.find_edge([nodelist[ii],nodelist[ii+1]]) for ii in\
range(len(nodelist)-1)], nodelist
def GridParticles(grdfile,dx,dy,nz,xypoly=None,splitvec=1):
"""
Returns the locations of particles on a regular grid inside of suntans grid
Inputs:
grdfile - netcdf filename containing the suntans grid
dx,dy - resolution in x and y component respectively.
nz - number of particles in the vertical. Particles are arranged in sigma layers.
xypoly - [optional] coordinates of an additional bounding polygon [Nx2 array]
"""
# Load the suntans grid
sun = Grid(grdfile)
# Load a trisearch object
tri = GridSearch(sun.xp,sun.yp,sun.cells,nfaces=sun.nfaces,verbose=False)
if xypoly == None:
xlims = [sun.xlims[0],sun.xlims[1]]
ylims = [sun.ylims[0],sun.ylims[1]]
else:
xlims = [xypoly[:,0].min(),xypoly[:,0].max()]
ylims = [xypoly[:,1].min(),xypoly[:,1].max()]
# Construct a 2D mesh of particles
x = np.arange(xlims[0],xlims[1],dx)
y = np.arange(ylims[0],ylims[1],dy)
X,Y = np.meshgrid(x,y)
X=X.ravel()
Y=Y.ravel()
# Check which particles are inside the grid
cellind = tri(X,Y)
mask = cellind!=-1
# Check which particles are also inside of the polygon
if not xypoly == None:
inpoly = inpolygon(np.vstack((X,Y)).T,xypoly)
mask = operator.and_(mask,inpoly)
xout = X[mask]
yout = Y[mask]
nx = xout.shape[0]
# Construct the 3D mesh
xout = np.repeat(xout.reshape((nx,1)),nz,axis=1)
yout = np.repeat(yout.reshape((nx,1)),nz,axis=1)
zout = np.linspace(0.05,0.95,nz)
zout = np.repeat(zout.reshape((nz,1)),nx,axis=1)
zout *= -sun.dv[cellind[mask]]
zout = zout.T
xout = xout.ravel()
yout = yout.ravel()
zout = zout.ravel()
# Rearrange the vectors to avoid clustering (this helps even the MPI workload)
#xout_split=[]
#yout_split=[]
#zout_split=[]
#for start in range(splitvec):
# xout_split.append(xout[start::splitvec])
# yout_split.append(yout[start::splitvec])
# zout_split.append(zout[start::splitvec])
#xout = np.hstack(xout_split)
#yout = np.hstack(yout_split)
#zout = np.hstack(zout_split)
return xout, yout, zout
def __init__(self,x,y,cells,nfaces,method='nearest',grdfile=None):
self.method=method
# Initialise the trisearch array
GridSearch.__init__(self,x,y,cells,nfaces=nfaces,force_inside=True)
