Ejemplo n.º 1
0
    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]
Ejemplo n.º 2
0
    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)
Ejemplo n.º 3
0
 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
Ejemplo n.º 4
0
    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
Ejemplo n.º 5
0
    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)
Ejemplo n.º 6
0
    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
Ejemplo n.º 7
0
    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
Ejemplo n.º 8
0
    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
Ejemplo n.º 9
0
    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
Ejemplo n.º 10
0
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
Ejemplo n.º 11
0
 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)
Ejemplo n.º 12
0
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
Ejemplo n.º 13
0
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
Ejemplo n.º 14
0
 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)