Python numpy.broadcast_arrays函数代码示例

def refractions(n1, n2, ray_dirs, normals):
	"""Generates directions of rays refracted according to Snells's law (in its vector
	form, [2]
	
	Arguments: 
	n1, n2 - respectively the refractive indices of the medium the unrefracted ray
		travels in and of the medium the ray is entering.
	ray_dirs, normals - each a row of 3-component vectors (as an array) with the
		direction of incoming rays and corresponding normals at the points of
		incidence with the refracting surface.
	
	Returns:
	refracted - a boolean array stating which of the incoming rays has not
		undergone total internal reflection.
	refr_dirs - new ray directions as the result of refraction, for the non-TIR
		rays in the input bundle.
	"""
	# Broadcast all necessary arrays to the larger size required:
	n = N.broadcast_arrays(n2/n1, ray_dirs[0])[0]
	normals = N.broadcast_arrays(normals, ray_dirs)[0]
	cos1 = (normals*ray_dirs).sum(axis=0)
	refracted = cos1**2 >= 1 - n**2
	
	# Throw away totally-reflected rays.
	cos1 = cos1[refracted]
	ray_dirs = ray_dirs[:,refracted]
	normals = normals[:,refracted]
	n = n[refracted]
	
	refr_dirs = (ray_dirs - cos1*normals)/n
	cos2 = N.sqrt(1 - 1./n**2*(1 - cos1**2))
	refr_dirs += normals*cos2*N.where(cos1 < 0, -1, 1)
	
	return refracted, refr_dirs

    def draw(self,animating = False):
        cols, rows = self.size
        minx, maxx = self.xlimits
        miny, maxy = self.ylimits

        cwidth, cheight = self.cell_size

        x = map(lambda i: minx + cwidth*i, range(cols+1))
        y = map(lambda i: miny + cheight*i, range(rows+1))

	if(not animating):
        	f = plt.figure(figsize=self.figsize)
	else:
		plt.clf()
        hlines = np.column_stack(np.broadcast_arrays(x[0], y, x[-1], y))
        vlines = np.column_stack(np.broadcast_arrays(x, y[0], x, y[-1]))
        lines = np.concatenate([hlines, vlines]).reshape(-1, 2, 2)
        line_collection = LineCollection(lines, color="black", linewidths=0.5)
        ax = plt.gca()
        ax.add_collection(line_collection)
        ax.set_xlim(x[0]-1, x[-1]+1)
        ax.set_ylim(y[0]-1, y[-1]+1)
        plt.gca().set_aspect('equal', adjustable='box')
        plt.axis('off')
        self._draw_obstacles(plt.gca())
        self._draw_start_goal(plt.gca())

        return plt.gca()

    def draw(self):
        cols, rows = self.size
        minx, maxx = self.xlimits
        miny, maxy = self.ylimits

        width, height = self.cell_dimensions

        x = map(lambda i: minx + width*i, range(cols+1))
        y = map(lambda i: miny + height*i, range(rows+1))

        f = plt.figure(figsize=self.figsize)

        hlines = np.column_stack(np.broadcast_arrays(x[0], y, x[-1], y))
        vlines = np.column_stack(np.broadcast_arrays(x, y[0], x, y[-1]))
        lines = np.concatenate([hlines, vlines]).reshape(-1, 2, 2)
        line_collection = LineCollection(lines, color="black", linewidths=0.5)
        ax = plt.gca()
        ax.add_collection(line_collection)
        ax.set_xlim(x[0]-1, x[-1]+1)
        ax.set_ylim(y[0]-1, y[-1]+1)
        plt.gca().set_aspect('equal', adjustable='box')
        plt.axis('off')
        self.draw_obstacles(plt.gca())

        return plt.gca()

 def half_edge_align(p, pts, polys):
     poly = align_polys(p, polys)
     mid   = pts[poly].mean(1)
     left  = pts[poly[:,[0,2]]].mean(1)
     right = pts[poly[:,[0,1]]].mean(1)
     s1 = np.array(np.broadcast_arrays(pts[p], mid, left)).swapaxes(0,1)
     s2 = np.array(np.broadcast_arrays(pts[p], mid, right)).swapaxes(0,1)
     return np.vstack([s1, s2])

 def test_simple(self):
     orig = np.ma.masked_array([[1], [2], [3]], mask=[[1], [0], [1]])
     result = BroadcastArray(orig, {1: 2}, (2,)).masked_array()
     expected, _ = np.broadcast_arrays(orig.data, result)
     expected_mask, _ = np.broadcast_arrays(orig.mask, result)
     expected = np.ma.masked_array(expected, mask=expected_mask)
     assert_array_equal(result.mask, expected.mask)
     assert_array_equal(result.data, expected.data)

    def axes_to_table(axes):
        """Fill the observation group axes into a table.

        Define one row for each possible combination of the
        observation group axis bins. Each row will represent
        an observation group.

        Parameters
        ----------
        axes : `~gammapy.data.ObservationGroupAxis`
            List of observation group axes.

        Returns
        -------
        table : `~astropy.table.Table`
            Table containing the observation group definitions.
        """
        # define table column data
        column_data_min = []
        column_data_max = []
        # loop over observation axes
        for i_axis in range(len(axes)):
            if axes[i_axis].fmt == 'values':
                column_data_min.append(axes[i_axis].bins)
                column_data_max.append(axes[i_axis].bins)
            elif axes[i_axis].fmt == 'edges':
                column_data_min.append(axes[i_axis].bins[:-1])
                column_data_max.append(axes[i_axis].bins[1:])

        # define grids of column data
        ndim = len(axes)
        s0 = (1,) * ndim
        expanding_arrays = [x.reshape(s0[:i] + (-1,) + s0[i + 1::])
                            for i, x in enumerate(column_data_min)]
        column_data_expanded_min = np.broadcast_arrays(*expanding_arrays)
        expanding_arrays = [x.reshape(s0[:i] + (-1,) + s0[i + 1::])
                            for i, x in enumerate(column_data_max)]
        column_data_expanded_max = np.broadcast_arrays(*expanding_arrays)

        # recover units
        for i_dim in range(ndim):
            column_data_expanded_min[i_dim] = _recover_units(column_data_expanded_min[i_dim],
                                                             column_data_min[i_dim])
            column_data_expanded_max[i_dim] = _recover_units(column_data_expanded_max[i_dim],
                                                             column_data_max[i_dim])

        # Make the table
        table = Table()
        for i_axis in range(len(axes)):
            if axes[i_axis].fmt == 'values':
                table[axes[i_axis].name] = column_data_expanded_min[i_axis].flatten()
            elif axes[i_axis].fmt == 'edges':
                table[axes[i_axis].name + "_MIN"] = column_data_expanded_min[i_axis].flatten()
                table[axes[i_axis].name + "_MAX"] = column_data_expanded_max[i_axis].flatten()

        ObservationGroups._add_group_id(table, axes)

        return table

def _reduce_points_and_bounds(points, lower_and_upper_bounds=None):
    """
    Reduce the dimensionality of arrays of coordinate points (and optionally
    bounds).

    Dimensions over which all values are the same are reduced to size 1, using
    :func:`_collapse_degenerate_points_and_bounds`.
    All size-1 dimensions are then removed.
    If the bounds arrays are also passed in, then all three arrays must have
    the same shape or be capable of being broadcast to match.

    Args:

    * points (array-like):
        Coordinate point values.

    Kwargs:

    * lower_and_upper_bounds (pair of array-like, or None):
        Corresponding bounds values (lower, upper), if any.

    Returns:
        dims (iterable of ints), points(array), bounds(array)

        * 'dims' is the mapping from the result array dimensions to the
            original dimensions.  However, when 'array' is scalar, 'dims' will
            be None (rather than an empty tuple).
        * 'points' and 'bounds' are the reduced arrays.
            If no bounds were passed, None is returned.

    """
    orig_points_dtype = np.asarray(points).dtype
    bounds = None
    if lower_and_upper_bounds is not None:
        lower_bounds, upper_bounds = np.broadcast_arrays(
            *lower_and_upper_bounds)
        orig_bounds_dtype = lower_bounds.dtype
        bounds = np.vstack((lower_bounds, upper_bounds)).T

    # Attempt to broadcast points to match bounds to handle scalars.
    if bounds is not None and points.shape != bounds.shape[:-1]:
        points, _ = np.broadcast_arrays(points, bounds[..., 0])

    points, bounds = _collapse_degenerate_points_and_bounds(points, bounds)

    used_dims = tuple(i_dim for i_dim in range(points.ndim)
                      if points.shape[i_dim] > 1)
    reshape_inds = tuple([points.shape[dim] for dim in used_dims])
    points = points.reshape(reshape_inds)
    points = points.astype(orig_points_dtype)
    if bounds is not None:
        bounds = bounds.reshape(reshape_inds + (2,))
        bounds = bounds.astype(orig_bounds_dtype)

    if not used_dims:
        used_dims = None

    return used_dims, points, bounds

def _bandflux(model, band, time_or_phase, zp, zpsys):
    """Support function for bandflux in Source and Model.
    This is necessary to have outside because ``phase`` is used in Source
    and ``time`` is used in Model.
    """

    if zp is not None and zpsys is None:
        raise ValueError('zpsys must be given if zp is not None')

    # broadcast arrays
    if zp is None:
        time_or_phase, band = np.broadcast_arrays(time_or_phase, band)
    else:
        time_or_phase, band, zp, zpsys = \
            np.broadcast_arrays(time_or_phase, band, zp, zpsys)

    # convert all to 1d arrays
    ndim = time_or_phase.ndim # save input ndim for return val
    time_or_phase = np.atleast_1d(time_or_phase)
    band = np.atleast_1d(band)
    if zp is not None:
        zp = np.atleast_1d(zp)
        zpsys = np.atleast_1d(zpsys)

    # initialize output arrays
    bandflux = np.zeros(time_or_phase.shape, dtype=np.float)

    # Loop over unique bands.
    for b in set(band):
        mask = band == b
        b = get_bandpass(b)

        # Raise an exception if bandpass is out of model range.
        if (b.wave[0] < model.minwave() or b.wave[-1] > model.maxwave()):
            raise ValueError(
                'bandpass {0!r:s} [{1:.6g}, .., {2:.6g}] '
                'outside spectral range [{3:.6g}, .., {4:.6g}]'
                .format(b.name, b.wave[0], b.wave[-1], 
                        model.minwave(), model.maxwave()))

        # Get the flux
        f = model._flux(time_or_phase[mask], b.wave)
        fsum = np.sum(f * b.trans * b.wave * b.dwave, axis=1) / HC_ERG_AA

        if zp is not None:
            zpnorm = 10.**(0.4 * zp[mask])
            bandzpsys = zpsys[mask]
            for ms in set(bandzpsys):
                mask2 = bandzpsys == ms
                ms = get_magsystem(ms)
                zpnorm[mask2] = zpnorm[mask2] / ms.zpbandflux(b)
            fsum *= zpnorm

        bandflux[mask] = fsum

    if ndim == 0:
        return bandflux[0]
    return bandflux

def uniform_distribution_overlap(m1, s1, m2, s2):
    '''Find the overlap between two uniform distributions
    
    Compute the integral of two uniform distributions
    centered on m1 and m2
    with widths 2 * s1 and 2 * s2
    and heights 1 / (2 * s1) and 1 / (2 * s2)
    '''
    h1h2 = 1 / (4 * s1 * s2)
    return np.clip((np.min(np.broadcast_arrays(m1 + s1, m2 + s2), axis=0) -
                    np.max(np.broadcast_arrays(m1 - s1, m2 - s2), axis=0)),
                   0, None) / h1h2

def test_broadcast_arrays():
    # Currently arrayfire is missing support for int64
    x2 = numpy.array([[1,2,3]], dtype=numpy.float32)
    y2 = numpy.array([[1],[2],[3]], dtype=numpy.float32)
    x1 = afnumpy.array(x2)
    y1 = afnumpy.array(y2)
    iassert(afnumpy.broadcast_arrays(x1, y1), numpy.broadcast_arrays(x2, y2))
    x1 = afnumpy.array([2])
    y1 = afnumpy.array(2)
    x2 = numpy.array([2])
    y2 = numpy.array(2)
    iassert(afnumpy.broadcast_arrays(x1, y1), numpy.broadcast_arrays(x2, y2))

def bkgsubtract(space, bkg):
    if space.dimension == bkg.dimension:
        bkg.photons = bkg.photons * space.contributions / bkg.contributions
        bkg.photons[bkg.contributions == 0] = 0
        bkg.contributions = space.contributions
        return space - bkg
    else:
        photons = numpy.broadcast_arrays(space.photons, bkg.photons)[1]
        contributions = numpy.broadcast_arrays(space.contributions, bkg.contributions)[1]
        bkg = Space(space.axes)
        bkg.photons = photons
        bkg.contributions = contributions
        return bkgsubtract(space, bkg)

    def __setitem__(self, key, x):
        row, col = self._validate_indices(key)

        if isinstance(row, INT_TYPES) and isinstance(col, INT_TYPES):
            x = np.asarray(x, dtype=self.dtype)
            if x.size != 1:
                raise ValueError('Trying to assign a sequence to an item')
            self._set_intXint(row, col, x.flat[0])
            return

        if isinstance(row, slice):
            row = np.arange(*row.indices(self.shape[0]))[:, None]
        else:
            row = np.atleast_1d(row)

        if isinstance(col, slice):
            col = np.arange(*col.indices(self.shape[1]))[None, :]
            if row.ndim == 1:
                row = row[:, None]
        else:
            col = np.atleast_1d(col)

        i, j = np.broadcast_arrays(row, col)
        if i.shape != j.shape:
            raise IndexError('number of row and column indices differ')

        from .base import isspmatrix
        if isspmatrix(x):
            if i.ndim == 1:
                # Inner indexing, so treat them like row vectors.
                i = i[None]
                j = j[None]
            broadcast_row = x.shape[0] == 1 and i.shape[0] != 1
            broadcast_col = x.shape[1] == 1 and i.shape[1] != 1
            if not ((broadcast_row or x.shape[0] == i.shape[0]) and
                    (broadcast_col or x.shape[1] == i.shape[1])):
                raise ValueError('shape mismatch in assignment')
            if x.size == 0:
                return
            x = x.tocoo(copy=True)
            x.sum_duplicates()
            self._set_arrayXarray_sparse(i, j, x)
        else:
            # Make x and i into the same shape
            x = np.asarray(x, dtype=self.dtype)
            x, _ = np.broadcast_arrays(x, i)
            if x.shape != i.shape:
                raise ValueError("shape mismatch in assignment")
            if x.size == 0:
                return
            self._set_arrayXarray(i, j, x)

    def score(self, outcomes, modelparams, expparams, return_L=False):

        na = np.newaxis
        n_m = modelparams.shape[0]
        n_e = expparams.shape[0]
        n_o = outcomes.shape[0]
        n_p = self.n_modelparams
        
        m = expparams['m'].reshape((1, 1, 1, n_e))
        
        L = self.likelihood(outcomes, modelparams, expparams)[na, ...]
        outcomes = outcomes.reshape((1, n_o, 1, 1))
        
        if not self._il:

            p, A, B = modelparams.T[:, :, np.newaxis]
            p = p.reshape((1, 1, n_m, 1))
            A = A.reshape((1, 1, n_m, 1))
            B = B.reshape((1, 1, n_m, 1))
        
            q = (-1)**(1-outcomes) * np.concatenate(np.broadcast_arrays(
                A * m * (p ** (m-1)), p**m, np.ones_like(p),
            ), axis=0) / L
            
        else:
        
            p_tilde, p_ref, A, B = modelparams.T[:, :, np.newaxis]
            p_C = p_tilde * p_ref
            
            mode = expparams['reference'][np.newaxis, :]
            
            p = np.where(mode, p_ref, p_C)
            
            p = p.reshape((1, 1, n_m, n_e))
            A = A.reshape((1, 1, n_m, 1))
            B = B.reshape((1, 1, n_m, 1))
        
            q = (-1)**(1-outcomes) * np.concatenate(np.broadcast_arrays(
                np.where(mode, 0, A * m * (p_tilde ** (m - 1)) * (p_ref ** m)),
                np.where(mode,
                    A * m * (p_ref ** (m - 1)),
                    A * m * (p_ref ** (m - 1)) * (p_tilde ** m)
                ),
                p**m, np.ones_like(p)
            ), axis=0) / L
        
        if return_L:
            # Need to strip off the extra axis we added for broadcasting to q.
            return q, L[0, ...]
        else:
            return q

    def _index_to_arrays(self, i, j):
        if isinstance(i, np.ndarray) and i.dtype.kind == 'b':
            i = self._boolean_index_to_array(i)
            if len(i)==2:
                if isinstance(j, slice):
                    j = i[1]

                else:
                    raise ValueError('too many indices for array')
            i = i[0]
        if isinstance(j, np.ndarray) and j.dtype.kind == 'b':
            j = self._boolean_index_to_array(j)[0]

        i_slice = isinstance(i, slice)
        if i_slice:
            i = self._slicetoarange(i, self.shape[0])[:,None]
        else:
            i = np.atleast_1d(i)

        if isinstance(j, slice):
            j = self._slicetoarange(j, self.shape[1])[None,:]
            if i.ndim == 1:
                i = i[:,None]
            elif not i_slice:
                raise IndexError('index returns 3-dim structure')
        elif isscalarlike(j):
            # row vector special case
            j = np.atleast_1d(j)
            if i.ndim == 1:
                i, j = np.broadcast_arrays(i, j)
                i = i[:, None]
                j = j[:, None]
                return i, j
        else:
            j = np.atleast_1d(j)
            if i_slice and j.ndim>1:
                raise IndexError('index returns 3-dim structure')

        i, j = np.broadcast_arrays(i, j)

        if i.ndim == 1:
            # return column vectors for 1-D indexing
            i = i[None,:]
            j = j[None,:]
        elif i.ndim > 2:
            raise IndexError("Index dimension must be <= 2")

        return i, j

    def noise_variance(self, bl_indices, f_indices, nt_per_day, ndays=None):
        """Calculate the instrumental noise variance.

        Parameters
        ----------
        bl_indices : array_like
            Indices of baselines to calculate.
        f_indices : array_like
            Indices of frequencies to calculate. Must be broadcastable against
            `bl_indices`.
        nt_per_day : integer
            The number of time samples in one sidereal day.
        ndays : integer
            The number of sidereal days observed.

        Returns
        -------
        noise_var : np.ndarray
            The noise variance.
        """

        ndays = self.ndays if not ndays else ndays # Set to value if not set.
        t_int = ndays * units.t_sidereal / nt_per_day
        # bw = 1.0e6 * (self.freq_upper - self.freq_lower) / self.num_freq
        bw = np.abs(self.frequencies[1] - self.frequencies[0]) * 1e6

        # Broadcast arrays against each other
        bl_indices, f_indices = np.broadcast_arrays(bl_indices, f_indices)

        return 2.0*self.tsys(f_indices)**2 / (t_int * bw * self.redundancy[bl_indices]) # 2.0 for two pol