Python numpy.digitize函数代码示例

def group(angle, wind, bsp, speedbins, anglebins, fct=np.median):
    '''Group data in bins according to wind angle and wind speed.

    Parameters
    ----------
    angle : np.ndarry
        Wind angles in degrees
    wind : np.ndarray
        wind speed in kn
    bsp : np.ndarray
        Boat speed in kn
    speedbins : ndarray
        bin boundaries for speed binning
    anglebins : ndarray
        bin boundaries for angle binning.
        Make sure that 180. is included in last bin and not on the boundary.
    fct : function
        Given all bsp values in one (speedbin,anglebin) select on value to
        be used. Common examples are np.median or np.mean

    Returns
    -------
    polar : ndarray([len(speedbins)+1, len(anglebins)])
        This contains the data array with one speed for each (speedbin, anglebin)
    '''
    if (angle.shape != wind.shape) or (angle.shape != bsp.shape):
        raise ValueError('angle, wind and bsp must have same number of elements')

    digspeed = np.digitize(wind, speedbins)
    digangle = np.digitize(np.abs(angle), anglebins)
    polar = np.zeros([len(speedbins)+1, len(anglebins)])
    for i in np.arange(1, len(speedbins)+1):
        for j in np.arange(1, len(anglebins)):
            polar[i, j] = fct(bsp[(digspeed == i) & (digangle == j)])
    return polar

def pos2Grid(x,y,data,xbins=None,ybins=None):
    '''Make a pixellated grid image from a 1d array
        of positions x,y,d. No smoothing, just binning.'''

    if (xbins == None):
	xbins = np.arange(x.min(),x.max()+1)
    if (ybins == None):
	ybins = np.arange(y.min(),y.max()+1)
            
    xd = np.digitize(x,xbins)
    xd -= 1 
	
    yd = np.digitize(y,ybins)
    yd -= 1
	
    (w,) = np.where((xd != 0) & (yd != 0))

    xd = xd[w]
    yd = yd[w]
    data = data[w]
	
    xi,yi = np.array(np.meshgrid(xbins,ybins,indexing='ij'))
    zi = xi*0
    zi[xd,yd] = data
    return zi,xd,yd

def means2idxarrays(g1, g2, i_bins, c_bins, difference):
    '''take two arrays of values and return the initial values
    and differences as numpy digitised arrays'''

    if difference == "relative":
        # calculate difference between mean values for group1 and group2
        # g1 and g2 always the same length
        change = [g2[x] - g1[x] for x in range(0, len(g1))]
        initial = g1

    elif difference == "logfold":
        change = [np.log2((g2[x] + 1.0) / (g1[x] + 1.0))
                  for x in range(0, len(g1))]
        initial = [np.log2(g1[x] + 1.0) for x in range(0, len(g1))]

    elif difference == "abs_logfold":
        change = [abs(np.log2((g2[x] + 1.0) / (g1[x] + 1.0)))
                  for x in range(0, len(g1))]
        initial = [max(np.log2(g1[x] + 1.0), np.log2(g2[x] + 1.0))
                   for x in range(0, len(g1))]

    # return arrays of len(change) with the index position in c_bins
    # corresponding to the bin in which the value of change falls
    change_idx = np.digitize(change, c_bins, right=True)
    initial_idx = np.digitize(initial, i_bins, right=True)

    return(change_idx, initial_idx)

def scatter_density(x, y, xlabel=None, ylabel=None, title=None, xlims=None,
                    ylims=None, filename=None):
    plt.figure()
    plt.grid()
    hist, xedges, yedges = np.histogram2d(x, y)
    xidx = np.clip(np.digitize(x, xedges), 0, hist.shape[0] - 1)
    yidx = np.clip(np.digitize(y, yedges), 0, hist.shape[1] - 1)
    c = hist[xidx, yidx]
    print "starting to plot the scatter plot"
    plt.scatter(x, y, c=c)

    if xlabel:
        plt.xlabel(xlabel)
    if ylabel:
        plt.ylabel(ylabel)
    if title:
        plt.title(title)
    if xlims:
        plt.xlim(xlims)
    if ylims:
        plt.ylim(ylims)

    if filename:
        plt.savefig(filename)
    else:
        plt.show()

def vertical_length_distribution(src_alt, simplex_alt, simplex_lengths, 
        alt_bins, norm=True):
    """ given input altitudes and lengths in km, create vertical
        profiles of source counts and total length.
        
        Returns alt_bins, bin_total_src, bin_total_length
        
        If norm==True, divide the counts by the bin width, returning
        km, counts/km and km/km. Otherwise just return km, counts and km.
        """
        
    # Not sure why we're not using histogram here, so that's a TODO
    # d_alt = 0.5
    d_alt = alt_bins[1:]-alt_bins[:-1]
    # alt_bins = np.arange(0.0,max_alt+d_alt, d_alt)
    bin_total_length = np.zeros(alt_bins.shape[0]-1, dtype=float)
    bin_total_src = np.zeros(alt_bins.shape[0]-1, dtype=float)
    # bin_total_length_sq = np.zeros(alt_bins.shape[0]-1, dtype=float)
    tri_bin_idx = np.digitize(simplex_alt, alt_bins)
    src_bin_idx = np.digitize(src_alt,alt_bins)
    tri_bin_idx[tri_bin_idx>(bin_total_length.shape[0]-1)]=bin_total_length.shape[0]-1
    src_bin_idx[src_bin_idx>(bin_total_src.shape[0]-1)]=bin_total_src.shape[0]-1

    for idx in src_bin_idx:
        bin_total_src[idx] += 1

    for lw,idx in zip(simplex_lengths,tri_bin_idx):
        bin_total_length[idx]+=lw
        # bin_total_length_sq[idx] += lw*lw
    # bin_total_length[tri_bin_idx] += length_weighted
    if norm==True:
        return alt_bins, bin_total_src/d_alt, bin_total_length/d_alt
    else:
        return alt_bins, bin_total_src, bin_total_length

    def _bin_descriptors(self, siftgeo, pca, grid, dimensions, duration):
        """ Groups the points in different bins using the gridding specified
        by grid. The returned results is a dictionary that has a key the bin
        number on each of the three dimensions x, y and t.

        """
        W, H = dimensions
        t_init, t_final = duration
        # Create equally spaced bins.
        bins_x = linspace(0, W + 1, grid[0] + 1)
        bins_y = linspace(0, H + 1, grid[1] + 1)
        bins_t = linspace(t_init, t_final + 1, grid[2] + 1)
        bag_xx = defaultdict(list)
        bag_ll = defaultdict(list)
        N = 0
        for ss in siftgeo:
            xx = pca.transform(ss[1])
            N += 1
            id_x = digitize([ss[0]['x']], bins_x)
            id_y = digitize([ss[0]['y']], bins_y)
            id_t = digitize([ss[0]['t']], bins_t)
            bag_xx[(id_x[0], id_y[0], id_t[0])].append(xx)
            bag_ll[(id_x[0], id_y[0], id_t[0])].append([ss[0]['x'] / W, ss[0]['y'] / H, (ss[0]['t'] - t_init) / (t_final + 1 - t_init)])
            assert (1 <= id_x <= grid[0] and
                    1 <= id_y <= grid[1] and
                    1 <= id_t <= grid[2])
        return bag_xx, bag_ll

def take2D(histogram, x, y, bins_x, bins_y):
    """
    Take the value from a two-dimensional histogram from the bin corresponding to (x, y).

    Parameters:
    -----------
    histogram : The values in the histogram (n,m) (ADW: is this ordering right?)
    x : the x-value to take from the hist
    y : the y-value to take from the hist
    bins_x : the xbin edges, including upper edge (n-dim)
    bins_y : the ybin edges, including upper edge (m-dim)
    """
    histogram = np.array(histogram)
    
    if np.isscalar(x):
        x = [x]
    if np.isscalar(y):
        y = [y]

    bins_x[-1] += 1.e-10 * (bins_x[-1] - bins_x[-2]) # Numerical stability
    bins_y[-1] += 1.e-10 * (bins_y[-1] - bins_y[-2])

    #return np.take(histogram, (histogram.shape[1] * (np.digitize(y, bins_y) - 1)) + (np.digitize(x, bins_x) - 1))

    # Return np.nan for entries which are outside the binning range on either axis
    index = (histogram.shape[1] * (np.digitize(y, bins_y) - 1)) + (np.digitize(x, bins_x) - 1)
    index_clipped = np.clip(index, 0, (histogram.shape[0] * histogram.shape[1]) - 1)
    val = np.take(histogram, index_clipped)

    outlier_x = np.logical_or(x < bins_x[0], x > bins_x[-1])
    outlier_y = np.logical_or(y < bins_y[0], y > bins_y[-1])
    outlier = np.logical_or(outlier_x, outlier_y)
    val[outlier] = np.nan

    return val 

def get_line_histos(results, temp, image, axis=0, bins=None):
    """This function creates an ADU histogram per each pixel in the direction defined by the axis parameter.
    """
    if image is None:
        temp["current_entry"] += 1
        return results, temp

    if bins is None:
        bins = np.arange(-100, 1000, 5)

    for i in range(image.shape[axis]):
        if axis == 0:
            t_histo = np.bincount(np.digitize(image[i, :].flatten(), bins[1:-1]), 
                          minlength=len(bins) - 1)
        elif axis == 1:
            t_histo = np.bincount(np.digitize(image[:, i].flatten(), bins[1:-1]), 
                          minlength=len(bins) - 1)

        if temp["current_entry"] == 0 and i == 0:
            results["histo_counts_line"] = np.empty([image.shape[axis], t_histo.shape[0]], 
                                                dtype=image.dtype)
        if temp["current_entry"] == 0:           
            results["histo_counts_line"][i] = t_histo
        else:
            results["histo_counts_line"][i] += t_histo
    temp["current_entry"] += 1
    return results, temp

def _get_rejrej_array(flat_eff, flat_x, flat_y, x_range=None, y_range=None):
    indices = np.nonzero((flat_eff > 0.005) & np.isfinite(flat_x) & np.isfinite(flat_y))

    used_x = np.log10(flat_x[indices])
    used_y = np.log10(flat_y[indices])
    used_eff = flat_eff[indices]

    if not x_range:
        # allow 1% safety margin on max value
        max_x = _max_noninf(used_x) * 1.0001
        min_x = np.min(used_x)
    else:
        min_x, max_x = x_range

    if not y_range:
        max_y = _max_noninf(used_y) * 1.0001
        min_y = np.min(used_y)
    else:
        min_y, max_y = y_range

    n_out_bins = 100

    x_bin_values = np.linspace(min_x, max_x, n_out_bins)
    x_bins = np.digitize(used_x, bins=x_bin_values) - 1  # no underflow

    y_bin_values = np.linspace(min_y, max_y, n_out_bins)
    y_bins = np.digitize(used_y, bins=y_bin_values) - 1  # no underflow

    make_eff_array = _loop_over_entries  # the other method seems slower

    eff_array = make_eff_array(x_bins, y_bins, used_eff, n_out_bins)

    return eff_array, (min_x, max_x), (min_y, max_y)

    def place(self, sig, bg_x, bg_y, cut_1_range, cut_2_range):
        """
        calculates x,y,z coordinates (rej x, rej y, eff)
        
        NOTE: make sure the eff, rej_x, rej_y arrays are integrated
        """
        assert bg_x.shape == bg_y.shape
        npts_1, npts_2 = bg_x.shape

        c1_bin_bounds = np.linspace(*cut_1_range, num=(npts_1 + 1))
        c1_bin = np.digitize([self._cut_1], c1_bin_bounds) - 1

        c2_bin_bounds = np.linspace(*cut_2_range, num=(npts_2 + 1))
        c2_bin = np.digitize([self._cut_2], c2_bin_bounds) - 1

        if any(b < 0 for b in [c1_bin, c2_bin]):
            raise ValueError("can't put a cut in the underflow bin")

        eff = float(sig[c1_bin, c2_bin] / sig.max())

        def get_rej(bkg_array):
            array_val = bkg_array.max() / bkg_array[c1_bin, c2_bin]
            return float(array_val)

        rej_x, rej_y = [get_rej(ar) for ar in [bg_x, bg_y]]

        self._xyz = rej_x, rej_y, eff
        self._cut_ranges = (cut_1_range, cut_2_range)

    def updateViz(self):
    	if self.gridRadiusViz == 0:
		vals=[]
		for name in self.vizObjectNames:
			r = moose.element(name+self.moosepath)
			d = float(r.getField(self.variable))
                        vals.append(d)
		inds = digitize(vals,self.stepVals)
		for i in range(0,len(self.vizObjects)):
			self.vizObjects[i].r,self.vizObjects[i].g,self.vizObjects[i].b=self.colorMap[inds[i]-1]

	else:
		vals=[]
		vals_2=[]
		for name in self.vizObjectNames:
                    r=mc.pathToId(name+self.moosepath)
                    d=float(mc.getField(r,self.variable))
                    

                    r2=mc.pathToId(name+self.moosepath_2)
                    d2=float(mc.getField(r2,self.variable_2))
				
                    vals.append(d)
                    vals_2.append(d2)
			
		inds = digitize(vals,self.stepVals)
		inds_2 = digitize(vals_2,self.stepVals_2)

		for i in range(0,len(self.vizObjects)):
			self.vizObjects[i].r,self.vizObjects[i].g,self.vizObjects[i].b=self.colorMap[inds[i]-1]
			self.vizObjects[i].radius=self.indRadius[inds_2[i]-1]

	self.updateGL()

def sim_make_residual_images(rmcal,binX=32,binY=32):
	xBins = np.arange(0,nX+1,binX)
	yBins = np.arange(0,nY+1,binY)
	median_a_offset = 0
	dmag = []
	for i,obj in enumerate(rmcal):
		mag,err = rmcal.get_object_phot(obj)
		dmag.append(obj.refMag - (mag - median_a_offset))
	dmag = np.concatenate(dmag)
	xy = np.hstack( [ [rmcal.objs[i].xpos,rmcal.objs[i].ypos] 
	                             for i in range(rmcal.num_objects()) ] )
	# XXX hack that last index in a_indices is ccdNum
	ccds = np.concatenate( [ rmcal.objs[i].a_indices[-1]
	                             for i in range(rmcal.num_objects()) ] )
	ffmaps = []
	for ccdNum in range(4):
		ffmap = [[[] for xi in xBins] for yi in yBins]
		ii = np.where(ccds==ccdNum)[0]
		for xi,yi,dm in zip(np.digitize(xy[0,ii],xBins),
		                    np.digitize(xy[1,ii],yBins),
		                    dmag[ii]):
			ffmap[yi][xi].append(dm)
		for xi in range(len(xBins)):
			for yi in range(len(yBins)):
				if len(ffmap[yi][xi])==0:
					ffmap[yi][xi] = np.nan
				else:
					ffmap[yi][xi] = np.median(ffmap[yi][xi])
		ffmaps.append(np.array(ffmap))
	return np.array(ffmaps)

def Mars_Year_np(j2k_np, jday_vals, year_vals, year_length, return_length=False):
    jday_vals = np.array(jday_vals)

    year_vals = np.array(year_vals)

    year_length = np.array(year_length)

    if j2k_np < jday_vals[0]:
        return np.floor(1+(j2k_np-jday_vals[0])/year_length[0])
    elif j2k_np >= jday_vals[-1]:
        return np.floor(1+(j2k_np-jday_vals[-1])/year_length[-1])
    else:
        try:
            v=np.clip(np.digitize(j2k_np,jday_vals),1,jday_vals.size)-1
            y = year_vals[v]
            l = year_length[v]
        except:
            v=np.clip(np.digitize([j2k_np],jday_vals),1,jday_vals.size)-1
            y = year_vals[v][0]
            l = year_length[v][0]

    if return_length:
        return (y*1.0,l)
    else:
        return y*1.0

def plot_2Dhist_medians(x, y, z, xlabel=None, ylabel=None, cblabel=None, ranges=[[-0.007, 0.002],[-0.014, 0.005]], vmin=0.0, vmax=10.0,
                filename=None):
    
    xedges = np.linspace(ranges[0][0], ranges[0][1], 51) # these numbers chosen to get 50 bins in final plot
    yedges = np.linspace(ranges[1][0], ranges[1][1], 51)

    xbins = np.digitize(x, xedges) # values falling below min(xedges) assigned 0; values above max(xedges) assigned 51
    ybins = np.digitize(y, yedges)
    
    medians = np.zeros((50,50))
    for i in range(50):
        for j in range(50):
            medians[i,j] = np.nanmedian(z[(xbins == i+1) * (ybins == j+1)])

    fig, ax = plt.subplots(figsize=(6.5, 5))
    plt.gcf().subplots_adjust(bottom=0.15)

    plt.imshow(medians.T, origin='lower', aspect='auto',
               interpolation='nearest', cmap=plt.cm.viridis, vmin=vmin, vmax=vmax,
               extent=(ranges[0][0], ranges[0][1], ranges[1][0], ranges[1][1]))

    if xlabel:
        plt.xlabel(xlabel)
    if ylabel:
        plt.ylabel(ylabel)
    cb = plt.colorbar()
    if cblabel:
        cb.set_label(cblabel)

    plt.draw()
    plt.tight_layout()

    if filename:
        plt.savefig(filename)

    def __getitem__(self, key):
        """
        Implements slicing or indexing of the Histogram
        """
        if key is (): return self # May no longer be necessary
        if isinstance(key, tuple) and len(key) > self.ndims:
            raise Exception("Slice must match number of key dimensions.")

        centers = [(float(l)+r)/2 for (l,r) in zip(self.edges, self.edges[1:])]
        if isinstance(key, slice):
            start, stop = key.start, key.stop
            if [start, stop] == [None,None]: return self
            start_idx, stop_idx = None,None
            if start is not None:
                start_idx = np.digitize([start], centers, right=True)[0]
            if stop is not None:
                stop_idx = np.digitize([stop], centers, right=True)[0]

            slice_end = stop_idx+1 if stop_idx is not None else None
            slice_values = self.values[start_idx:stop_idx]
            slice_edges =  self.edges[start_idx: slice_end]

            extents = (min(slice_edges), self.extents[1],
                       max(slice_edges), self.extents[3])
            return self.clone((slice_values, slice_edges), extents=extents)
        else:
            if not (self.edges.min() <= key < self.edges.max()):
                raise Exception("Key value %s is out of the histogram bounds" % key)
            idx = np.digitize([key], self.edges)[0]
            return self.values[idx-1 if idx>0 else idx]