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 __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