Python pylab 模块，array() 实例源码

def joint_density(X, Y, bounds=None):
    """
Plots joint distribution of variables.
Inherited from method in src/graphics.py module in project 
git://github.com/aflaxman/pymc-example-tfr-hdi.git
    """
    if bounds:
        X_min, X_max, Y_min, Y_max = bounds
    else:
        X_min = X.min()
        X_max = X.max()
        Y_min = Y.min()
        Y_max = Y.max()

    pylab.plot(X, Y, linestyle='none', marker='o', color='green', mec='green', alpha=.2, zorder=-99)

    gkde = scipy.stats.gaussian_kde([X, Y])
    x,y = pylab.mgrid[X_min:X_max:(X_max-X_min)/25.,Y_min:Y_max:(Y_max-Y_min)/25.]
    z = pylab.array(gkde.evaluate([x.flatten(), y.flatten()])).reshape(x.shape)
    pylab.contour(x, y, z, linewidths=2)

    pylab.axis([X_min, X_max, Y_min, Y_max])

def drawROI(self, ax=None, value=None, pixel=None):
        if not ax: ax = plt.gca()
        roi_map = np.array(healpy.UNSEEN*np.ones(healpy.nside2npix(self.nside)))

        if value is None:
            roi_map[self.roi.pixels] = 1
            roi_map[self.roi.pixels_annulus] = 0
            roi_map[self.roi.pixels_target] = 2
        elif value is not None and pixel is None:
            roi_map[self.pixels] = value
        elif value is not None and pixel is not None:
            roi_map[pixel] = value
        else:
            logger.warning('Unable to parse input')
        #im = healpy.gnomview(roi_map,**self.gnom_kwargs)
        im = drawHealpixMap(roi_map,self.glon,self.glat,self.radius,coord=self.coord)
        return im

def drawStellarDensity(self,ax=None):
        if not ax: ax = plt.gca()
        # Stellar Catalog
        self._create_catalog()
        catalog = self.catalog
        #catalog=ugali.observation.catalog.Catalog(self.config,roi=self.roi)
        pix = ang2pix(self.nside, catalog.lon, catalog.lat)
        counts = collections.Counter(pix)
        pixels, number = numpy.array(sorted(counts.items())).T
        star_map = healpy.UNSEEN * numpy.ones(healpy.nside2npix(self.nside))
        star_map[pixels] = number
        star_map = numpy.where(star_map == 0, healpy.UNSEEN, star_map)

        #im = healpy.gnomview(star_map,**self.gnom_kwargs)
        #healpy.graticule(dpar=1,dmer=1,color='0.5',verbose=False)
        #pylab.close()

        im = drawHealpixMap(star_map,self.glon,self.glat,self.radius,coord=self.coord)
        #im = ax.imshow(im,origin='bottom')
        try:    ax.cax.colorbar(im)
        except: pylab.colorbar(im,ax=ax)
        ax.annotate("Stars",**self.label_kwargs)
        return im

def set_zidx(self):
        names = [n.upper() for n in self.obj.array.dtype.names]
        mod = np.array(self.config['scan']['distance_modulus_array'])
        if 'ZIDX_MAX' in names:
            self.zidx = self.obj['ZIDX_MAX'] 
        elif 'DISTANCE_MODULUS' in names:
            dist_mod = self.obj['DISTANCE_MODULUS']
            self.zidx = np.abs(mod - dist_mod).argmin()
        elif 'MODULUS' in names:
            dist_mod = self.obj['MODULUS']
            self.zidx = np.abs(mod - dist_mod).argmin()
        elif 'DISTANCE' in names:
            dist_mod = mod2dist(self.obj['DISTANCE'])
            self.zidx = np.argmax((mod - dist_mod) > 0)
        else:
            msg = "Failed to parse distance index"
            raise Exception(msg)

def queryDb(db='Pawsey', limit=None, saveToCSV=None):
    """
    Return the query result as a numpy array

    saveToCSV:  csv file name (string)
    """
    if (not DB.has_key(db)):
        raise Exception("Unknown db name {0}".format(db))
    # don't want to use group by to relief the database
    if (limit is None or len(limit) == 0):
        sl = ""
    else:
        sl = " {0}".format(limit)
    hsql = "select file_size from ngas_files {0}".format(sl)
    try:
        t = dbpass
    except NameError:
        dbpass = getpass.getpass('%s DB password: ' % db)
        print "Connecting to DB"
        dbconn = dbdrv.connect(database="ngas", user="ngas_ro",
                             password=dbpass, host=DB[db])

    cur = dbconn.cursor()
    print "Executing query '{0}'".format(hsql)
    cur.execute(hsql)
    r = cur.fetchall()
    dbconn.close()
    del(dbconn)
    res = pylab.array(r)
    ll = len(res)
    res = res[:,0].reshape(ll)
    if (not (saveToCSV is None)):
        f = open(saveToCSV, 'wb')
        for li in res:
            #print "--{0}:{1}".format(int(li), type(int(li)))
            f.write("{0}\n".format(li))

    return res # get the first (only) column from all rows

def queryCSV(csvfile):
    with open(csvfile) as f:
        print "loading records from {0}".format(csvfile)
        re = f.readlines()
        return pylab.array(re)

def drawHealpixMap(map, lon, lat, size=1.0, xsize=501, coord='GC', **kwargs):
    """
    Draw local projection of healpix map.
    """
    ax = plt.gca()
    x = np.linspace(-size,size,xsize)
    y = np.linspace(-size,size,xsize)
    xx, yy = np.meshgrid(x,y)

    coord = coord.upper()

    if coord == 'GC':
        #Assumes map and (lon,lat) are Galactic, but plotting celestial
        llon, llat = image2sphere(*gal2cel(lon,lat),x=xx.flat,y=yy.flat)
        pix = ang2pix(healpy.get_nside(map),*cel2gal(llon,llat))
    elif coord == 'CG':
        #Assumes map and (lon,lat) are celestial, but plotting Galactic
        llon, llat = image2sphere(*cel2gal(lon,lat),x=xx.flat,y=yy.flat)
        pix = ang2pix(healpy.get_nside(map),*gal2cel(llon,llat))
    else:
        #Assumes plotting the native coordinates
        llon, llat = image2sphere(lon,lat,xx.flat,yy.flat)
        pix = ang2pix(healpy.get_nside(map),llon,llat)

    values = map[pix].reshape(xx.shape)
    zz = np.ma.array(values,mask=(values==healpy.UNSEEN),fill_value=np.nan)

    return drawProjImage(xx,yy,zz,coord=coord,**kwargs)

def drawIsochrone(isochrone, **kwargs):
    ax = plt.gca()
    logger.debug(str(isochrone))
    if kwargs.pop('cookie',None):
        # Broad cookie cutter
        defaults = dict(alpha=0.5, color='0.5', zorder=0, 
                        linewidth=15, linestyle='-')
    else:
        # Thin lines
        defaults = dict(color='k', linestyle='-')
    kwargs = dict(defaults.items()+kwargs.items())

    isos = isochrone.isochrones if hasattr(isochrone,'isochrones') else [isochrone]
    for iso in isos:
        iso = copy.deepcopy(iso)
        logger.debug(iso.filename)
        iso.hb_spread = False
        mass_init,mass_pdf,mass_act,mag_1,mag_2 = iso.sample(mass_steps=1e3)
        mag = mag_1 + isochrone.distance_modulus
        color = mag_1 - mag_2

        # Find discontinuities in the color magnitude distributions
        dmag = np.fabs(mag[1:]-mag[:-1])
        dcolor = np.fabs(color[1:]-color[:-1])
        idx = np.where( (dmag>1.0) | (dcolor>0.25))[0]
        # +1 to map from difference array to original array
        mags = np.split(mag,idx+1)
        colors = np.split(color,idx+1)

        for i,(c,m) in enumerate(zip(colors,mags)):
            msg = '%-4i (%g,%g) -- (%g,%g)'%(i,m[0],c[0],m[-1],c[-1])
            logger.debug(msg)
            if i > 0:
                kwargs['label'] = None
            ax.plot(c,m,**kwargs)
    return ax

def main(args=sys.argv[1:]):
    # there are some cases when this script is run on systems without DISPLAY variable being set
    # in such case matplotlib backend has to be explicitly specified
    # we do it here and not in the top of the file, as inteleaving imports with code lines is discouraged
    import matplotlib
    matplotlib.use('Agg')
    from pylab import plt, ylabel, grid, xlabel, array

    parser = argparse.ArgumentParser()
    parser.add_argument("rst_file", help="location of rst file in TRiP98 format", type=str)
    parser.add_argument("output_file", help="location of PNG file to save", type=str)
    parser.add_argument("-s", "--submachine", help="Select submachine to plot.", type=int, default=1)
    parser.add_argument("-f", "--factor", help="Factor for scaling the blobs. Default is 1000.", type=int, default=1000)
    parser.add_argument("-v", "--verbosity", action='count', help="increase output verbosity", default=0)
    parser.add_argument('-V', '--version', action='version', version=pt.__version__)
    args = parser.parse_args(args)

    file = args.rst_file
    sm = args.submachine
    fac = args.factor

    a = pt.Rst()
    a.read(file)

    # convert data in submachine to a nice array
    b = a.machines[sm]

    x = []
    y = []
    z = []
    for _x, _y, _z in b.raster_points:
        x.append(_x)
        y.append(_y)
        z.append(_z)

    title = "Submachine: {:d} / {:d} - Energy: {:.3f} MeV/u".format(sm, len(a.machines), b.energy)
    print(title)
    cc = array(z)

    cc = cc / cc.max() * fac

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.scatter(x, y, c=cc, s=cc, alpha=0.75)
    ylabel("mm")
    xlabel("mm")
    grid(True)
    plt.title(title)
    plt.savefig(args.output_file)
    plt.close()

def allplot(xb,yb,bins=30,fig=1,xlabel='x',ylabel='y'):
    """
Input:
X,Y : objects referring to the variables produced by PyMC that you want
to analyze. Example: X=M.theta, Y=M.slope.

Inherited from Tommy LE BLANC's code at astroplotlib|STSCI.
    """
    #X,Y=xb.trace(),yb.trace()
    X,Y=xb,yb

    #pylab.rcParams.update({'font.size': fontsize})
    fig=pylab.figure(fig)
    pylab.clf()

    gs = pylab.GridSpec(2, 2, width_ratios=[3,1], height_ratios=[1,3], wspace=0.07, hspace=0.07)
    scat=pylab.subplot(gs[2])
    histx=pylab.subplot(gs[0], sharex=scat)
    histy=pylab.subplot(gs[3], sharey=scat)
    #scat=fig.add_subplot(2,2,3)
    #histx=fig.add_subplot(2,2,1, sharex=scat)
    #histy=fig.add_subplot(2,2,4, sharey=scat)

    # Scatter plot
    scat.plot(X, Y,linestyle='none', marker='o', color='green', mec='green',alpha=.2, zorder=-99)

    gkde = scipy.stats.gaussian_kde([X, Y])
    x,y = numpy.mgrid[X.min():X.max():(X.max()-X.min())/25.,Y.min():Y.max():(Y.max()-Y.min())/25.]
    z = numpy.array(gkde.evaluate([x.flatten(), y.flatten()])).reshape(x.shape)
    scat.contour(x, y, z, linewidths=2)
    scat.set_xlabel(xlabel)
    scat.set_ylabel(ylabel)

    # X-axis histogram
    histx.hist(X, bins, histtype='stepfilled')
    pylab.setp(histx.get_xticklabels(), visible=False)  # no X label
    #histx.xaxis.set_major_formatter(pylab.NullFormatter()) # no X label

    # Y-axis histogram
    histy.hist(Y, bins, histtype='stepfilled', orientation='horizontal')
    pylab.setp(histy.get_yticklabels(), visible=False)  # no Y label
    #histy.yaxis.set_major_formatter(pylab.NullFormatter()) # no Y label

    #pylab.minorticks_on()
    #pylab.subplots_adjust(hspace=0.1)
    #pylab.subplots_adjust(wspace=0.1)
    pylab.draw()
    pylab.show()

def plotDistance(self):
        filename = self.config.mergefile
        logger.debug("Opening %s..."%filename)
        f = pyfits.open(filename)
        pixels,values = f[1].data['PIXEL'],2*f[1].data['LOG_LIKELIHOOD']
        if values.ndim == 1: values = values.reshape(-1,1)
        distances = f[2].data['DISTANCE_MODULUS']
        if distances.ndim == 1: distances = distances.reshape(-1,1)
        ts_map = healpy.UNSEEN * numpy.ones(healpy.nside2npix(self.nside))

        ndim = len(distances)
        nrows = int(numpy.sqrt(ndim))
        ncols = ndim // nrows + (ndim%nrows > 0)

        fig = pylab.figure()
        axes  = AxesGrid(fig, 111, nrows_ncols = (nrows, ncols),axes_pad=0,
                         label_mode='1', cbar_mode='single',cbar_pad=0,cbar_size='5%',
                         share_all=True,add_all=False)

        images = []
        for i,val in enumerate(values.T):
            ts_map[pixels] = val

            im = healpy.gnomview(ts_map,**self.gnom_kwargs)
            pylab.close()
            images.append(im)
        data = numpy.array(images); mask = (data == healpy.UNSEEN)
        images = numpy.ma.array(data=data,mask=mask)
        vmin = numpy.ma.min(images)
        vmax = numpy.ma.max(images)

        for i,val in enumerate(values.T):
            ax = axes[i]
            im = ax.imshow(images[i],origin='bottom',vmin=vmin,vmax=vmax)
            ax.cax.colorbar(im)

            #ax.annotate(r"$\mu = %g$"%distances[i],**self.label_kwargs)
            ax.annotate(r"$d = %.0f$ kpc"%mod2dist(distances[i]),**self.label_kwargs)
            ax.axis["left"].major_ticklabels.set_visible(False) 
            ax.axis["bottom"].major_ticklabels.set_visible(False) 
            fig.add_axes(ax)
            fig.add_axes(ax.cax)
        return fig,axes

def draw_slices(hist, func=np.sum, **kwargs):
    """ Draw horizontal and vertical slices through histogram """
    from mpl_toolkits.axes_grid1 import make_axes_locatable
    kwargs.setdefault('ls','-')
    ax = plt.gca()

    data = hist

    # Slices
    vslice = func(data,axis=0)
    hslice = func(data,axis=1)

    npix = np.array(data.shape)
    #xlim,ylim = plt.array(zip([0,0],npix-1))
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    #extent = ax.get_extent()
    #xlim =extent[:2]
    #ylim = extent[2:]

    # Bin centers
    xbin = np.linspace(xlim[0],xlim[1],len(vslice))#+0.5 
    ybin = np.linspace(ylim[0],ylim[1],len(hslice))#+0.5
    divider = make_axes_locatable(ax)

    #gh2 = pywcsgrid2.GridHelperSimple(wcs=self.header, axis_nums=[2, 1])
    hax = divider.append_axes("right", size=1.2, pad=0.05,sharey=ax,
                              axes_class=axes_divider.LocatableAxes)
    hax.axis["left"].toggle(label=False, ticklabels=False)
    #hax.plot(hslice, plt.arange(*ylim)+0.5,'-') # Bin center
    hax.plot(hslice, ybin, **kwargs) # Bin center
    hax.xaxis.set_major_locator(MaxNLocator(4,prune='both'))
    hax.set_ylim(*ylim)

    #gh1 = pywcsgrid2.GridHelperSimple(wcs=self.header, axis_nums=[0, 2])
    vax = divider.append_axes("top", size=1.2, pad=0.05, sharex=ax,
                              axes_class=axes_divider.LocatableAxes)
    vax.axis["bottom"].toggle(label=False, ticklabels=False)
    vax.plot(xbin, vslice, **kwargs) 
    vax.yaxis.set_major_locator(MaxNLocator(4,prune='lower'))
    vax.set_xlim(*xlim)

    return vax,hax

def drawKernelHist(ax, kernel):
    ext = kernel.extension
    theta = kernel.theta
    lon, lat = kernel.lon, kernel.lat
    xmin,xmax = -5*ext,5*ext
    ymin,ymax = -5*ext,5*ext,
    x = np.linspace(xmin,xmax,100)+kernel.lon
    y = np.linspace(ymin,ymax,100)+kernel.lat

    xx,yy = np.meshgrid(x,y)
    zz = kernel.pdf(xx,yy)
    im = ax.imshow(zz)#,extent=[xmin,xmax,ymin,ymax])
    hax,vax = draw_slices(ax,zz,color='k')

    mc_lon,mc_lat = kernel.sample(1e5)
    hist,xedges,yedges = np.histogram2d(mc_lon,mc_lat,bins=[len(x),len(y)],
                                        range=[[x.min(),x.max()],[y.min(),y.max()]])
    xbins,ybins = np.arange(hist.shape[0])+0.5,np.arange(hist.shape[1])+0.5

    vzz = zz.sum(axis=0)
    hzz = zz.sum(axis=1)
    vmc = hist.sum(axis=0)
    hmc = hist.sum(axis=1)

    vscale = vzz.max()/vmc.max()
    hscale = hzz.max()/hmc.max()

    kwargs = dict(marker='.',ls='',color='r')
    hax.errorbar(hmc*hscale, ybins, xerr=np.sqrt(hmc)*hscale,**kwargs)
    vax.errorbar(xbins, vmc*vscale,yerr=np.sqrt(vmc)*vscale,**kwargs) 

    ax.set_ylim(0,len(y))
    ax.set_xlim(0,len(x))

    #try: ax.cax.colorbar(im)
    #except: pylab.colorbar(im)

    #a0 = np.array([0.,0.])
    #a1 =kernel.a*np.array([np.sin(np.deg2rad(theta)),-np.cos(np.deg2rad(theta))])
    #ax.plot([a0[0],a1[0]],[a0[1],a1[1]],'-ob')
    # 
    #b0 = np.array([0.,0.])
    #b1 =kernel.b*np.array([np.cos(np.radians(theta)),np.sin(np.radians(theta))])
    #ax.plot([b0[0],b1[0]],[b0[1],b1[1]],'-or')    

    label_kwargs = dict(xy=(0.05,0.05),xycoords='axes fraction', xytext=(0, 0), 
                        textcoords='offset points',ha='left', va='bottom',size=10,
                        bbox={'boxstyle':"round",'fc':'1'}, zorder=10)
    norm = zz.sum() * (x[1]-x[0])**2
    ax.annotate("Sum = %.2f"%norm,**label_kwargs)

    #ax.set_xlabel(r'$\Delta$ LON (deg)')
    #ax.set_ylabel(r'$\Delta$ LAT (deg)')

###################################################

def drawKernel(kernel, contour=False, coords='C', **kwargs):
    ax = plt.gca()

    if 'colors' not in kwargs:
        kwargs.setdefault('cmap',matplotlib.cm.jet)
    kwargs.setdefault('origin','lower')

    ext   = kernel.extension
    theta = kernel.theta

    xmin,xmax = -kernel.edge,kernel.edge
    ymin,ymax = -kernel.edge,kernel.edge

    if coords[-1] == 'G':
        lon, lat = kernel.lon, kernel.lat
    elif coords[-1] == 'C':
        lon,lat = gal2cel(kernel.lon, kernel.lat)
    else:
        msg = 'Unrecognized coordinate: %s'%coords
        raise Exception(msg)

    x = np.linspace(xmin,xmax,500)+lon
    y = np.linspace(ymin,ymax,500)+lat
    xx,yy = np.meshgrid(x,y)
    extent = [x[0],x[-1],y[0],y[-1]]
    kwargs.setdefault('extent',extent)

    if coords[-1] == 'C': xx,yy = cel2gal(xx,yy)

    zz = kernel.pdf(xx.flat,yy.flat).reshape(xx.shape)
    zmax = zz.max()

    if contour:
        levels = kwargs.pop('levels',10)
        #levels = np.logspace(np.log10(zmax)-1,np.log10(zmax),7)
        ret = ax.contour(zz,levels,**kwargs)
    else:
        val = np.ma.array(zz,mask=zz<zz.max()/100.)
        ret = ax.imshow(val,**kwargs)

    return ret


###################################################