Python pylab.pcolor函数代码示例

 def test2_3x4_to_5x7_cart(self):
     # Test non-periodic grid returning double grid resolution
     roESMP = CdmsRegrid(self.fromGrid3x4, self.toGrid5x7,
                         dtype = self.data3x4.dtype,
                         srcGridMask = self.data3x4.mask,
                         regridTool = 'ESMP',
                         periodicity = 0,
                         regridMethod = 'Conserve', 
                         coordSys = 'cart')
     diag = {'srcAreas':0, 'dstAreas':0, 
             'srcAreaFractions':0, 'dstAreaFractions':0}
     ESMP5x7 = roESMP(self.data3x4, diag = diag)
     dstMass = (ESMP5x7 * diag['dstAreas']).sum()
     srcMass = (self.data3x4 * diag['srcAreas'] \
                             * diag['srcAreaFractions']).sum()
     if False:
         pylab.figure(1)
         pylab.pcolor(self.data3x4)
         pylab.colorbar()
         pylab.title('original self.data3x4')
         pylab.figure(2)
         pylab.pcolor(ESMP5x7)
         pylab.colorbar()
         pylab.title('interpolated ESMP5x7')
     self.assertLess(abs(srcMass - dstMass), self.eps)
     self.assertEqual(self.data3x4[0,0], ESMP5x7[0,0])
     self.assertEqual(1.0, ESMP5x7[0,0])
     self.assertEqual(0.25, ESMP5x7[1,1])
     self.assertEqual(0.0, ESMP5x7[2,2])

def plot(nxG, nyG, iBeg, iEnd, jBeg, jEnd, data, title=''):
    """
    Plot distributed array
    @param nxG number of global cells in x
    @param nyG number of global cells in y
    @param iBeg global starting index in x
    @param iEnd global ending index in x
    @param jBeg global starting index in y
    @param jEnd global ending index in y
    @param data local array
    @param title plot title
    """
    sz = MPI.COMM_WORLD.Get_size()
    rk = MPI.COMM_WORLD.Get_rank()
    iBegs = MPI.COMM_WORLD.gather(iBeg, root=0)
    iEnds = MPI.COMM_WORLD.gather(iEnd, root=0)
    jBegs = MPI.COMM_WORLD.gather(jBeg, root=0)
    jEnds = MPI.COMM_WORLD.gather(jEnd, root=0)
    arrays = MPI.COMM_WORLD.gather(numpy.array(data), root=0)
    if rk == 0:
        bigArray = numpy.zeros((nxG, nyG), data.dtype)
        for pe in range(sz):
            bigArray[iBegs[pe]:iEnds[pe], jBegs[pe]:jEnds[pe]] = arrays[pe]
        from matplotlib import pylab
        pylab.pcolor(bigArray.transpose())
        # add the decomp domains
        for pe in range(sz):
            pylab.plot([iBegs[pe], iBegs[pe]], [0, nyG - 1], 'w--')
            pylab.plot([0, nxG - 1], [jBegs[pe], jBegs[pe]], 'w--')
        pylab.title(title)
        pylab.show()

 def test2_varRegrid(self):
     print
     print 'test2_varRegrid'
     srcF = cdms2.open(sys.prefix + \
                           '/sample_data/so_Omon_ACCESS1-0_historical_r1i1p1_185001-185412_2timesteps.nc')
     so = srcF('so')[0, 0, ...]
     clt = cdms2.open(sys.prefix + '/sample_data/clt.nc')('clt')
     diag = {'srcAreas': None, 'dstAreas': None,
             'srcAreaFractions': None, 'dstAreaFractions': None}
     soInterp = so.regrid(clt.getGrid(), 
                          regridTool = 'esmf', 
                          regridMethod='conserve', 
                    	     diag = diag)
     if self.pe == 0:
         totSrcArea = diag['srcAreas'].sum()
         totDstArea = diag['dstAreas'].sum()
         totSrcFrac = diag['srcAreaFractions'].sum()
         self.assertEqual(numpy.isnan(totSrcFrac).sum(), 0)
         self.assertLess(abs(totSrcArea - 4*pi)/(4*pi), 0.02)
         self.assertLess(abs(totDstArea - 4*pi)/(4*pi), 0.01)
         soMass = (so*diag['srcAreas']).sum()
         inMass = (soInterp*diag['dstAreas']).sum()
         print soMass, inMass
         diff = abs(soMass - inMass)/soMass
         self.assertLess(diff, 7.e-7)
     if False:
         pylab.subplot(1, 2, 1)
         pylab.pcolor(so, vmin = 20, vmax = 40)
         pylab.colorbar()
         pylab.title('so')
         pylab.subplot(1, 2, 2)
         pylab.pcolor(soInterp, vmin = 20, vmax = 40)
         pylab.colorbar()
         pylab.title('soInterp')

def density_cloud_by_tags(df, columns, silent=False):
    """Create density cloud of data for a given tag or group of tags
    For example:
        columns='DayOfWeek' --> Plots for Mon, Tue, Wed, Thur, ...
        columns='Weekday' --> Plots of weekends vs weekday
        columns=['Season','Weekday'] 
            --> Plots of Summer, Spring, Winter, Fall Weekdays and Weekends
    """
    figures = []
    if columns == 'hr' or 'hr' in columns:
        raise ValueError("Columns cannot contain hr tag")
        
    # Create a profile for day of week
    maxY = df['USAGE'].max()
    for label, data in df.groupby(columns):    
        
        # Find mean
        mean = data.groupby('hr')['USAGE'].agg('mean')
        # Add in any missing hours
        for h in set(range(24)) - set(data['hr']):
            mean = mean.set_value(h, None)
            
    
        # Create a density cloud of the MW
        X = np.zeros([24, 100]) # Hours by resolution
        Y = np.zeros([24, 100])
        C = np.zeros([24, 100])    
        for hr, data2 in data.groupby('hr'):        
            freq = []
            step = 1
            rng = range(0,51,step)[1:]
            freq += rng
            bins = np.percentile(data2['USAGE'], rng)
            
            rng = range(50,101,step)[1:]
            freq += [100 - a for a in rng]
            bins = np.hstack([bins, np.percentile(data2['USAGE'], rng)])
            freq = np.array(freq)
               
            X[hr,:] = np.ones(len(bins))*hr
            Y[hr,:] = bins
            C[hr,:] = freq
        
        plt.figure()
        #plt.xkcd()
        plt.pcolor(X, Y, C, cmap=plt.cm.YlOrRd)
        plt.plot(X[:,1], mean, color='k', label='Mean')
        plt.colorbar().set_label('Probability Higher/Lower than Median')    
        plt.legend(loc='upper left')
        plt.xlabel('Hour of Day')
        plt.ylabel('Usage (kWh)')
        plt.ylim([0, maxY])
        plt.xlim([0,23])
        plt.title('Typical usage on %s' % str(label))
        plt.grid(axis='y')
        figures.append(plt.gcf())
        if not silent:
            plt.show()
        
    return figures

    def XXtest5_regrid(self):
        srcF = cdms2.open(sys.prefix + \
                              '/sample_data/so_Omon_ACCESS1-0_historical_r1i1p1_185001-185412_2timesteps.nc')
        so = srcF('so')[0, 0, ...]
        clt = cdms2.open(sys.prefix + '/sample_data/clt.nc')('clt')
        dstData = so.regrid(clt.getGrid(), 
                            regridTool = 'esmf', 
                            regridMethod='conserve')

        if self.pe == 0:
            dstDataMask = (dstData == so.missing_value)
            dstDataFltd = dstData * (1 - dstDataMask)
            zeroValCnt = (dstData == 0).sum()
            if so.missing_value > 0:
                dstDataMin = dstData.min()
                dstDataMax = dstDataFltd.max()
            else:
                dstDataMin = dstDataFltd.min()
                dstDataMax = dstData.max()
                zeroValCnt = (dstData == 0).sum()
            print 'Number of zero valued cells', zeroValCnt
            print 'min/max value of dstData: %f %f' % (dstDataMin, dstDataMax)                   
            self.assertLess(dstDataMax, so.max())
            if False:
                pylab.figure(1)
                pylab.pcolor(so, vmin=20, vmax=40)
                pylab.colorbar()
                pylab.title('so')
                pylab.figure(2)
                pylab.pcolor(dstData, vmin=20, vmax=40)
                pylab.colorbar()
                pylab.title('dstData')

    def plot_expansion(time, A_array, E_array, eig, c, file_to_save='exp.png',figuresize=(10,10)):
        from matplotlib.colors import LogNorm
        plt.figure(figsize=figuresize) # in inches!

        plt.subplot(1,2,1)
        plt.plot(A_array[0], time, label = r'$A_x$')
        plt.plot(E_array[0], time, label = r'$E_x$')
        plt.plot(A_array[1], time, label = r'$A_y$', ls = '--')
        plt.plot(E_array[1], time, label = r'$E_y$', ls = '--')
        plt.ylim(0, max(time))
        plt.ylabel(r'$t, s$', fontsize=26)
        plt.legend(prop={'size':16})

        plt.subplot(1,2,2)
        X,Y = np.meshgrid(eig.real, time)
        plt.pcolor(X, Y, c, norm=LogNorm(vmin=1.0e-3, vmax=1.0))
        plt.xlim(min(eig.real), max(eig.real))
        plt.ylim(0.0, max(time))
        plt.colorbar()
        plt.xlabel(r'$E, eV$', fontsize=26)
        plt.xlim(-0.3,0.3)
        if file_to_save:
            plt.savefig(file_to_save)
            plt.close()
        return None

 def plot_pcolor(self, X, Y, Z):
     Xm, Ym = np.meshgrid(X, Y)
     try:
         plt.pcolor(Xm, Ym, Z)
     except:
         plt.pcolor(Z)
     plt.colorbar()
     return None

    def test3(self):
        """
        2D + level
        """
        print 'This is a known failure for now. ESMF are looking into the error'
        u = cdms2.open(cdat_info.get_prefix() + '/sample_data/clt.nc')('u')[0, :,...]
        uCart = u.regrid( u.getGrid(),
                          regridTool='esmf', regridMethod='linear',
                          coordSys = 'cart',
                          periodicity = 1)
        uDegr = u.regrid( u.getGrid(),
                          regridTool='esmf', regridMethod='linear',
                          coordSys = 'deg',
                          periodicity = 1)
        n = reduce(lambda x,y: x*y, uCart.shape)
        mask = (u == u.missing_value)
        if PLOT:
            import matplotlib.pylab as pl
            fig = pl.figure()
            fig.add_subplot(2,1,1)
            pl.pcolor(u.getLongitude()[:], u.getLatitude()[:], uCart[1,...],
                      vmin = 0)
            pl.colorbar()
            pl.title('Cartiesian')
            fig.add_subplot(2,1,2)
            pl.pcolor(u.getLongitude()[:], u.getLatitude()[:], uDegr[1,...],
                      vmin = 0)
            pl.colorbar()
            pl.title('Degrees')
            pl.show()
        if PRINT:
            print
            print 'Level 0, u.min() = %11.3f, uInterp.min() = %11.3f' % \
                     (u[0,...].min(), uInterp[0,...].min())
            print 'Level 1, u.min() = %11.3f, uInterp.min() = %11.3f' % \
                     (u[1,...].min(), uInterp[1,...].min())

            print '\nFor indices 52,59 and 60, 68'
            print 'm means missing'
            print 'd means uInterp < u'
            print '. means uInterp == u'
            for i in range(52, 59):
                string = ""
                for j in range(60, 68):
                    if uInterp.mask[1,i,j] == True:
                        string = string+"m"
                    elif (uInterp[1,i,j] < u[1,i,j]):
                        string = string+"d"
                    else:
                        string = string + "."
                print string
            print

        diff = abs(numpy.sum(u*(1-mask) - uCart)/float(n))
        self.assertLess(diff, 1.e-3)
        diff = abs(numpy.sum(u*(1-mask) - uDegr)/float(n))
        self.assertLess(diff, 1.e-3)

def corr(data,labels,**kwargs):
    data=np.transpose(data)
    corrs=np.corrcoef(data)
    
    labelsDict=dict((i,labels[i]) for i in range(len(labels)))
    if 'makeGraph' in kwargs.keys():
        if kwargs['makeGraph']==True:
            fig,ax=plt.subplots()
    #         plt.pcolor(corrs)
            plt.pcolor(corrs>=kwargs['Tresh'])
            plt.xticks([i for i in range(44)],rotation=45)
            
            ax.set_xticklabels(labels)
            ax.set_yticklabels(labels)
            plt.tick_params(axis='both', which='both', labelsize=7)
    #         plt.imshow(corrs>=kwargs['Tresh'],interpolation=None)
    #         plt.colorbar()
            plt.show()
        
    if 'undGraph' in kwargs:
        plt.figure()
        if kwargs['undGraph']==True:
            gcorrs=np.copy(corrs)
            if 'Tresh' in kwargs:
                idx=np.where(corrs<=kwargs['Tresh'])
                gcorrs[idx]=0
                gcorrs=gcorrs-np.identity(gcorrs.shape[0])
                
            
            G=nx.from_numpy_matrix(np.triu(gcorrs))
            for node in nx.nodes(G):
                edges=np.sum([ 1 for i in nx.all_neighbors(G, node)])
                if edges==0:
                    G.remove_node(node)
                    labelsDict.pop(node)

            G=nx.relabel_nodes(G,labelsDict)
            
            pos=nx.spring_layout(G,iterations=200)
            
#             pos=nx.shell_layout(G)
            nx.draw_networkx(G,pos,font_size=9)
#             nx.draw_spring(G)
#             nx.draw(G,pos,font_size=9)
            plt.show()
            
            
    if 'ret' in kwargs.keys():    
        if kwargs['ret']==True:
            corrs2=np.triu(corrs-np.diagflat(np.diag(corrs)))
            i,j=np.where(np.abs(corrs2)>=kwargs['Tresh'])
    #         print corrs2[i,j]
    #         print i
    #         print j
            feats=set(list(i)+list(j))
    #         print feats
            return feats

def plot_contour(x, y, z, file_name):
    xx, yy = meshgrid(y, x)
    plt.figure()
    plt.pcolor(xx, yy, z, vmax=abs(z).max(), vmin=abs(z).min())
    cbar = plt.colorbar()
    cbar.ax.set_ylabel('Scalar Modifier', rotation=270)
    plt.xlabel('Calendar Years')
    plt.ylabel('Age')
    with open(str(file_name), 'wb') as f:
        plt.savefig(f, format='png')
    plt.close()

 def test2_tasWithLevelsUsingLibCF(self):
     test = 'Time test using regrid2 gsRegrid'
     r2 = regrid2.GenericRegrid(self.gGrid, self.fGrid, 
                                dtype = self.gtas.dtype,
                                regridMethod='linear', 
                                regridTool='libcf')
     r2.computeWeights()
     fShape = list(self.gtas.shape[:-2]) + list(self.fGrid[0].shape)
     ftas = np.ma.masked_array(np.ones(fShape, self.gtas.dtype)*self.gtas.missing_value,
                               mask = np.zeros(fShape))
     ftas.missing_value = self.gtas.missing_value
     r2.apply(self.gtas, ftas)
     if PLOT:
         vmin = self.gtas.min()
         vmax = self.gtas.max()
         pylab.subplot(1,4,1)
         pylab.pcolor(self.gtas[-1,...], vmin=vmin, vmax=vmax)
         pylab.colorbar()
         pylab.title('self.gtas[-1,...]')
         pylab.subplot(1,4,2)
         pylab.pcolor(ftasOld[-1,...], vmin=vmin, vmax=vmax)
         pylab.colorbar()
         pylab.title('ftasOld[-1,...]')
         pylab.subplot(1,4,3)
         pylab.pcolor(ftas[0,...], vmin=vmin, vmax=vmax)
         pylab.colorbar()
         pylab.title('ftas[-1,...]')
         pylab.subplot(1,4,4)
         pylab.pcolor(ftas[-1,...] - ftasOld[-1,...], vmin=-1, vmax=1)
         pylab.colorbar()
         pylab.title('ftas[-1,...] - ftasOld[-1,...]')
     
     self.assertEqual(ftas[0,...].shape, self.fclt[0,...].shape)
     self.assertEqual(ftas.shape[0], self.gtas.shape[0])
     self.assertGreater(abs(ftas[ftas>0].mean()-self.gtas.mean()), 1)

    def test_2d_esmf_conserv(self):
        print 'running test_2d_esmf_conserv...'
        f = cdms2.open(sys.prefix + \
                           '/sample_data/so_Omon_ACCESS1-0_historical_r1i1p1_185001-185412_2timesteps.nc')
        so = f('so')[0, 0, :, :]
        clt = cdms2.open(sys.prefix + '/sample_data/clt.nc')('clt')[0, :, :]
        tic = time.time()
        soInterp = so.regrid(clt.getGrid(), regridTool='ESMF', regridMethod='Conservative')
        soInterpInterp = soInterp.regrid(so.getGrid(), regridTool='ESMF', 
                                         regridMethod='Conservative')
        toc = time.time()
        print 'time to interpolate (ESMF conservative) forward/backward: ', toc - tic
        ntot = reduce(operator.mul, so.shape)
        avgdiff = numpy.sum(so - soInterpInterp) / float(ntot)
        print 'avgdiff = ', avgdiff

        if PLOT:
            pylab.figure(2)
            pylab.subplot(2, 2, 1)
            pylab.pcolor(so, vmin=20.0, vmax=40.0)
            pylab.colorbar()
            pylab.title('ESMF conserve regrid: so')
            pylab.subplot(2, 2, 2)
            pylab.pcolor(soInterp, vmin=20.0, vmax=40.0)
            pylab.colorbar()
            pylab.title('ESMF conserve regrid: soInterp')
            pylab.subplot(2, 2, 3)
            pylab.pcolor(soInterpInterp, vmin=20.0, vmax=40.0)
            pylab.colorbar()
            pylab.title('ESMF conserve regrid: soInterpInterp')
            pylab.subplot(2, 2, 4)
            pylab.pcolor(so - soInterpInterp, vmin=-0.5, vmax=0.5)
            pylab.colorbar()
            pylab.title('ESMF conserve regrid: error')

    def PlotFieldQuad(self):
        if not os.path.isdir(self.outdir):
            os.makedirs(self.outdir)

        fig_path = self.outdir

        from matplotlib import pylab as plt

        plt.close('all')
        plt.figure()
        quad = self.quad.copy()
        plt.plot(quad[0,:]/1e-15,quad[1,:],'g--',label='$\mathscr{S}$')
        plt.plot(quad[0,:]/1e-15,quad[2,:],'r--',label='$\mathscr{I}$')
        xlabel = 'time ($fs$)'
        ylabel = '$\mathscr{S}$, $\mathscr{I}$ ($a.u.$)'
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.legend(frameon=False)
        plt.draw()
        plt.savefig(os.path.join(fig_path,'si_vs_t.png'),dpi=240)
        plt.close()
        plt.figure()
        plt.plot(quad[0,:]/1e-15,quad[3,:],'g--',label='$\partial_t \mathscr{S}$')
        plt.plot(quad[0,:]/1e-15,quad[4,:],'r--',label='$\partial_t \mathscr{I}$')
        xlabel = 'time ($fs$)'
        ylabel = '$\partial_t$ $\mathscr{S}$, $\mathscr{I}$ ($fs^{-1}$)'
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.legend(frameon=False)
        plt.draw()
        plt.savefig(os.path.join(fig_path,'dtsi_vs_t.png'),dpi=240)
        plt.close()

        plt.figure()
        t = np.linspace(self.quad[0,0],self.quad[0,-1],len(self.quad[0,:]))
        x = self._x.copy()
        #print t.shape, x.shape,self.u.shape
        T,X = np.meshgrid(t,x)
        #print T[0::5,:].shape,X[0::5,:].shape,self.u[0::5,:].transpose().shape
        u = self.u.transpose()
        plt.pcolor(X[::10,:],T[::10,:],u[::10,:],shading='interp')
        xlabel = '$x$ ($\mu m$)'
        ylabel = 'time ($fs$)'
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.colorbar()
        plt.draw()
        plt.savefig(os.path.join(fig_path,'u.png'),dpi=240)
        plt.close()
        plt.close('all')

    def test3(self):
        """
        2D + level
        """
        print "This is a known failure for now. ESMF are looking into the error"
        u = cdms2.open(sys.prefix + "/sample_data/clt.nc")("u")[0, :, ...]
        uCart = u.regrid(u.getGrid(), regridTool="esmf", regridMethod="linear", coordSys="cart", periodicity=1)
        uDegr = u.regrid(u.getGrid(), regridTool="esmf", regridMethod="linear", coordSys="deg", periodicity=1)
        n = reduce(lambda x, y: x * y, uCart.shape)
        mask = u == u.missing_value
        if PLOT:
            import matplotlib.pylab as pl

            fig = pl.figure()
            fig.add_subplot(2, 1, 1)
            pl.pcolor(u.getLongitude()[:], u.getLatitude()[:], uCart[1, ...], vmin=0)
            pl.colorbar()
            pl.title("Cartiesian")
            fig.add_subplot(2, 1, 2)
            pl.pcolor(u.getLongitude()[:], u.getLatitude()[:], uDegr[1, ...], vmin=0)
            pl.colorbar()
            pl.title("Degrees")
            pl.show()
        if PRINT:
            print
            print "Level 0, u.min() = %11.3f, uInterp.min() = %11.3f" % (u[0, ...].min(), uInterp[0, ...].min())
            print "Level 1, u.min() = %11.3f, uInterp.min() = %11.3f" % (u[1, ...].min(), uInterp[1, ...].min())

            print "\nFor indices 52,59 and 60, 68"
            print "m means missing"
            print "d means uInterp < u"
            print ". means uInterp == u"
            for i in range(52, 59):
                string = ""
                for j in range(60, 68):
                    if uInterp.mask[1, i, j] == True:
                        string = string + "m"
                    elif uInterp[1, i, j] < u[1, i, j]:
                        string = string + "d"
                    else:
                        string = string + "."
                print string
            print

        diff = abs(numpy.sum(u * (1 - mask) - uCart) / float(n))
        self.assertLess(diff, 1.0e-3)
        diff = abs(numpy.sum(u * (1 - mask) - uDegr) / float(n))
        self.assertLess(diff, 1.0e-3)

 def test1_regrid(self):
     clt = cdms2.open(cdat_info.get_sampledata_path() + '/clt.nc')('clt')[0,...]
     ta = cdms2.open(cdat_info.get_sampledata_path() + '/ta_ncep_87-6-88-4.nc')('ta')[0, 0,...]
     diag = {}
     cltInterp = clt.regrid( ta.getGrid(), 
                             regridTool = 'libcf', 
                             mkCyclic = True, verbose = True,
                             diag = diag )
     print cltInterp.sum()
     n = reduce(lambda x,y: x*y, cltInterp.shape)
     self.assertLess(abs(cltInterp.sum() - 696921.0)/n, 0.3)
     if PLOT:
         pylab.pcolor(ta.getLongitude()[:], ta.getLatitude()[:], cltInterp)
         pylab.colorbar()
         pylab.title('cltInterp')
     self.assertEqual(True, True)