Python numpy.indices() Examples

def get_to_common_frame(obs, frame_wcs):
    """ Matches an `Observation`'s coordinates to a `Frame`'s wcs

    Parameters
    ----------
    obs: `Observation`
        An observation instance for which we want to know the coordinates in the frame of `frame_wcs`
    frame_wcs: `WCS`
        a wcs that gives the mapping between pixel coordinates and sky coordinates for a given frame
    Returns
    -------
        coord_obs_frame: `tuple`
            Coordinates of the observations's pixels in the frame of the provided wcs
    """
    c, ny, nx = obs.frame.shape
    #Positions of the observation's pixels
    y, x = np.indices((ny, nx))
    y = y.flatten()
    x = x.flatten()

    # Positions of Observation's pixel in the frame of the wcs
    Y,X = convert_coordinates((y,x), obs.frame.wcs, frame_wcs)

    coord = (Y,X)
    return coord 

def coefficients_to_polynomial(coefficients, variables):
  """Converts array of lists of coefficients to a polynomial."""
  coefficients = np.asarray(coefficients)
  shape = coefficients.shape

  indices = list(zip(*np.indices(shape).reshape([len(shape), -1])))
  monomials = []
  for power in indices:
    coeffs = coefficients.item(power)
    if (number.is_integer_or_rational(coeffs)
        or isinstance(coeffs, sympy.Symbol)):
      coeffs = [coeffs]
    elif not isinstance(coeffs, list):
      raise ValueError('Unrecognized coeffs={} type={}'
                       .format(coeffs, type(coeffs)))
    for coeff in coeffs:
      monomials.append(monomial(coeff, variables, power))
  random.shuffle(monomials)
  return ops.Add(*monomials) 

def trim(coefficients):
  """Makes non-zero entry in the final slice along each axis."""
  coefficients = np.asarray(coefficients)
  non_zero = np.not_equal(coefficients, 0)
  ndim = coefficients.ndim

  for axis in range(ndim):
    length = coefficients.shape[axis]
    axis_complement = list(range(0, axis)) + list(range(axis + 1, ndim))
    non_zero_along_axis = np.any(non_zero, axis=tuple(axis_complement))

    slice_to = 0
    for index in range(length - 1, -1, -1):
      if non_zero_along_axis[index]:
        slice_to = index + 1
        break

    if slice_to < length:
      coefficients = coefficients.take(axis=axis, indices=list(range(slice_to)))

  return coefficients 

def position(self, lat_idx):
        """return 'space' position of one or multiple sites.

        Parameters
        ----------
        lat_idx : ndarray, ``(... , dim+1)``
            Lattice indices.

        Returns
        -------
        pos : ndarray, ``(..., dim)``
            The position of the lattice sites specified by `lat_idx` in real-space.
        """
        idx = self._asvalid_latidx(lat_idx)
        res = np.take(self.unit_cell_positions, idx[..., -1], axis=0)
        for i in range(self.dim):
            res += idx[..., i, np.newaxis] * self.basis[i]
        return res 

def mps2lat_idx(self, i):
        """Translate MPS index `i` to lattice indices ``(x_0, ..., x_{dim-1}, u)``.

        Parameters
        ----------
        i : int | array_like of int
            MPS index/indices.

        Returns
        -------
        lat_idx : array
            First dimensions like `i`, last dimension has len `dim`+1 and contains the lattice
            indices ``(x_0, ..., x_{dim-1}, u)`` corresponding to `i`.
            For `i` accross the MPS unit cell and "infinite" `bc_MPS`, we shift `x_0` accordingly.
        """
        if self.bc_MPS == 'infinite':
            # allow `i` outsit of MPS unit cell for bc_MPS infinite
            i0 = i
            i = np.mod(i, self.N_sites)
            if np.any(i0 != i):
                lat = self.order[i].copy()
                lat[..., 0] += (i0 - i) * self.N_rings // self.N_sites
                # N_sites_per_ring might not be set for IrregularLattice
                return lat
        return self.order[i].copy() 

def mps_lat_idx_fix_u(self, u=None):
        """Similar as :meth:`mps_idx_fix_u`, but return also the corresponding lattice indices.

        Parameters
        ----------
        u : None | int
            Selects a site of the unit cell. ``None`` (default) means all sites.

        Returns
        -------
        mps_idx : array
            MPS indices `i` for which ``self.site(i) is self.unit_cell[u]``.
        lat_idx : 2D array
            The row `j` contains the lattice index (without `u`) corresponding to ``mps_idx[j]``.
        """
        mps_idx = self.mps_idx_fix_u(u)
        return mps_idx, self.order[mps_idx, :-1] 

def idrop_labels(self, old_labels=None):
        """Remove leg labels from self; in place.

        Parameters
        ----------
        old_labels : list of str|int
            The leg labels/indices for which the label should be removed.
            By default (None), remove all labels.
        """
        if old_labels is None:
            self._labels = [None] * self.rank
            return self
        old_inds = self.get_leg_indices(old_labels)
        labels = self._labels[:]
        for i in old_inds:
            labels[i] = None
        self._labels = labels
        return self

    # string output =========================================================== 

def as_completely_blocked(self):
        """Gives a version of self which is completely blocked by charges.

        Functions like :func:`svd` or :func:`eigh` require a complete blocking by charges.
        This can be achieved by encapsulating each leg which is not completely blocked into a
        :class:`LegPipe` (containing only that single leg). The LegPipe will then contain all
        necessary information to revert the blocking.

        Returns
        -------
        encapsulated_axes : list of int
            The leg indices which have been encapsulated into Pipes.
        blocked_self : :class:`Array`
            Self (if ``len(encapsulated_axes) = 0``) or a copy of self,
            which is completely blocked.
        """
        enc_axes = [a for a, l in enumerate(self.legs) if not l.is_blocked()]
        if len(enc_axes) == 0:
            return enc_axes, self
        qconj = [self.legs[a].qconj for a in enc_axes]
        return enc_axes, self.combine_legs([[a] for a in enc_axes], qconj=qconj) 

def _iter_common_sorted(a, b):
    """Yield indices ``i, j`` for which ``a[i] == b[j]``.

    *Assumes* that ``a[i_start:i_stop]`` and ``b[j_start:j_stop]`` are strictly ascending.
    Given that, it is equivalent to (but faster than)
    ``[(i, j) for j, i in itertools.product(range(len(a)), range(len(b)) if a[i] == b[j]]``
    """
    l_a = len(a)
    l_b = len(b)
    i, j = 0, 0
    res = []
    while i < l_a and j < l_b:
        if a[i] < b[j]:
            i += 1
        elif b[j] < a[i]:
            j += 1
        else:
            res.append((i, j))
            i += 1
            j += 1
    return res 

def bunch(self):
        """Return a copy with bunched self.charges: form blocks for contiguous equal charges.

        Returns
        -------
        idx : 1D array
            ``idx[:-1]`` are the indices of the old qind which are kept,
            ``idx[-1] = old_block_number``.
        cp : :class:`LegCharge`
            A new LegCharge with the same charges at given indices of the leg,
            but (possibly) shorter ``self.charges`` and ``self.slices``.

        See also
        --------
        sort : sorts by charges, thus enforcing complete blocking in combination with bunch.
        """
        if self.bunched:  # nothing to do
            return np.arange(self.block_number + 1, dtype=np.intp), self
        cp = self.copy()
        idx = _find_row_differences(self.charges)
        cp._set_charges(cp.charges[idx[:-1]])  # avanced indexing -> copy
        cp._set_slices(cp.slices[idx])
        cp.bunched = True
        return idx, cp 

def _find_row_differences(qflat):
    """Return indices where the rows of the 2D array `qflat` change.

    Parameters
    ----------
    qflat : 2D array
        The rows of this array are compared.

    Returns
    -------
    diffs: 1D array
        The indices where rows change, including the first and last. Equivalent to:
        ``[0]+[i for i in range(1, len(qflat)) if np.any(qflat[i-1] != qflat[i])] + [len(qflat)]``
    """
    if qflat.shape[1] == 0:
        return np.array([0, qflat.shape[0]], dtype=np.intp)
    diff = np.ones(qflat.shape[0] + 1, dtype=np.bool_)
    diff[1:-1] = np.any(qflat[1:] != qflat[:-1], axis=1)
    return np.nonzero(diff)[0]  # get the indices of True-values 

def collate_fn_base(
        self, attributes_and_types: Dict[str, type], batch: Union[List[int], np.ndarray]
    ) -> Tuple[torch.Tensor, ...]:
        """Given indices and attributes to batch, returns a full batch of ``Torch.Tensor``
        """
        indices = np.asarray(batch)
        data_numpy = [
            getattr(self, attr)[indices].astype(dtype)
            if isinstance(getattr(self, attr), np.ndarray)
            else getattr(self, attr)[indices].toarray().astype(dtype)
            for attr, dtype in attributes_and_types.items()
        ]

        data_torch = tuple(torch.from_numpy(d) for d in data_numpy)
        return data_torch

    #############################
    #                           #
    #      GENE FILTERING       #
    #                           #
    ############################# 

def subsample_cells(self, size: Union[int, float] = 1.0):
        """Wrapper around ``update_cells`` allowing for automatic (based on sum of counts) subsampling.

        If size is a:

            * (0,1) float: subsample 100*``size`` % of the cells
            * int: subsample ``size`` cells
        """
        new_n_cells = (
            int(size * self.nb_cells)
            if not isinstance(size, (int, np.integer))
            else size
        )
        indices = np.argsort(np.squeeze(np.asarray(self.X.sum(axis=1))))[::-1][
            :new_n_cells
        ]
        self.update_cells(indices) 

def filter_cell_types(self, cell_types: Union[List[str], List[int], np.ndarray]):
        """Performs in-place filtering of cells by keeping cell types in ``cell_types``.

        Parameters
        ----------
        cell_types
            numpy array of type np.int (indices) or np.str (cell-types names)
        """
        cell_types = np.asarray(cell_types)
        if isinstance(cell_types[0], str):
            labels_to_keep = self.cell_types_to_labels(cell_types)
        elif isinstance(cell_types[0], (int, np.integer)):
            labels_to_keep = cell_types
        else:
            raise ValueError(
                "Wrong dtype for cell_types. Should be either str or int (labels)."
            )

        subset_cells = self._get_cells_filter_mask_by_attribute(
            attribute_name="labels",
            attribute_values_to_keep=labels_to_keep,
            return_data=False,
        )

        self.update_cells(subset_cells) 

def _indices(self, slice_index=None):
        # get a int-array containing all indices in the first axis.
        if slice_index is None:
            slice_index = self._current_slice_
        #try:
        indices = np.indices(self._realshape_, dtype=int)
        indices = indices[(slice(None),)+slice_index]
        indices = np.rollaxis(indices, 0, indices.ndim).reshape(-1,self._realndim_)
            #print indices_
            #if not np.all(indices==indices__):
            #    import ipdb; ipdb.set_trace()
        #except:
        #    indices = np.indices(self._realshape_, dtype=int)
        #    indices = indices[(slice(None),)+slice_index]
        #    indices = np.rollaxis(indices, 0, indices.ndim)
        return indices 

def __str__(self, indices=None, iops=None, lx=None, li=None, lls=None, only_name=False, VT100=True):
        filter_ = self._current_slice_
        vals = self.flat
        if indices is None: indices = self._indices(filter_)
        if iops is None:
            ravi = self._raveled_index(filter_)
            iops = OrderedDict([name, iop.properties_for(ravi)] for name, iop in self._index_operations.items())
        if lls is None: lls = [self._max_len_names(iop, name) for name, iop in iops.items()]

        format_spec = '  |  '.join(self._format_spec(indices, iops, lx, li, lls, VT100))

        to_print = []

        if not only_name: to_print.append(format_spec.format(index=__index_name__, value=self.hierarchy_name(), **dict((name, name) for name in iops)))
        else: to_print.append(format_spec.format(index='-'*li, value=self.hierarchy_name(), **dict((name, '-'*l) for name, l in zip(iops, lls))))

        for i in range(self.size):
            to_print.append(format_spec.format(index=indices[i], value="{1:.{0}f}".format(__precision__, vals[i]), **dict((name, ' '.join(map(str, iops[name][i]))) for name in iops)))
        return '\n'.join(to_print) 

def perform_nms(math_regions):

    # convert from x1,y1,x2,y2 to x,y,w,h
    math_regions[:, 2] = math_regions[:, 2] - math_regions[:, 0]
    math_regions[:, 3] = math_regions[:, 3] - math_regions[:, 1]

    scores = math_regions[:, 4]
    math_regions = np.delete(math_regions, 4, 1)

    math_regions = math_regions.tolist()
    scores = scores.tolist()

    indices = NMSBoxes(math_regions, scores, 0.2, 0.5)

    indices = [item for sublist in indices for item in sublist]
    math_regions = [math_regions[i] for i in indices]

    math_regions = np.array(math_regions)

    # restore to x1,y1,x2,y2
    math_regions[:, 2] = math_regions[:, 2] + math_regions[:, 0]
    math_regions[:, 3] = math_regions[:, 3] + math_regions[:, 1]

    return math_regions.tolist() 

def perform_nms(math_regions):

    # convert from x1,y1,x2,y2 to x,y,w,h
    math_regions[:, 2] = math_regions[:, 2] - math_regions[:, 0]
    math_regions[:, 3] = math_regions[:, 3] - math_regions[:, 1]

    scores = math_regions[:, 4]
    math_regions = np.delete(math_regions, 4, 1)

    math_regions = math_regions.tolist()
    scores = scores.tolist()

    indices = NMSBoxes(math_regions, scores, 0.2, 0.5)

    indices = [item for sublist in indices for item in sublist]
    math_regions = [math_regions[i] for i in indices]

    math_regions = np.array(math_regions)

    # restore to x1,y1,x2,y2
    math_regions[:, 2] = math_regions[:, 2] + math_regions[:, 0]
    math_regions[:, 3] = math_regions[:, 3] + math_regions[:, 1]

    return math_regions.tolist() 

def centroid(component):
    """Determine centroid of model

    Parameters
    ----------
    component: `scarlet.Component` or `scarlet.ComponentTree`
        Component to analyze
    """
    model = component.get_model()
    indices = np.indices(model.shape)
    centroid = np.array([np.sum(ind * model) for ind in indices]) / model.sum()
    return centroid + component.bbox.origin 

def moments(data):
	"""Returns (height, x, y, width_x, width_y)
	the gaussian parameters of a 2D distribution by calculating its
	moments """
	total = data.sum()
	X, Y = np.indices(data.shape)
	x = (X*data).sum()/total
	y = (Y*data).sum()/total
	col = data[:, int(y)]
	width_x = np.sqrt(abs((np.arange(col.size)-y)**2*col).sum()/col.sum())
	row = data[int(x), :]
	width_y = np.sqrt(abs((np.arange(row.size)-x)**2*row).sum()/row.sum())
	height = data.max()
	return height, x, y, width_x, width_y 

def fitgaussian(data):
	"""Returns (height, x, y, width_x, width_y)
	the gaussian parameters of a 2D distribution found by a fit"""
	params = moments(data)
	errorfunction = lambda p: np.ravel(gaussian(*p)(*np.indices(data.shape)) - data)
	p, success = scipy.optimize.leastsq(errorfunction, params)
	return p 

def moments(data):
	"""Returns (height, x, y, width_x, width_y)
	the gaussian parameters of a 2D distribution by calculating its
	moments """
	total = data.sum()
	X, Y = np.indices(data.shape)
	x = (X*data).sum()/total
	y = (Y*data).sum()/total
	col = data[:, int(y)]
	width_x = np.sqrt(abs((np.arange(col.size)-y)**2*col).sum()/col.sum())
	row = data[int(x), :]
	width_y = np.sqrt(abs((np.arange(row.size)-x)**2*row).sum()/row.sum())
	height = data.max()
	return height, x, y, width_x, width_y 

def fitgaussian(data):
	"""Returns (height, x, y, width_x, width_y)
	the gaussian parameters of a 2D distribution found by a fit"""
	params = moments(data)
	errorfunction = lambda p: np.ravel(gaussian(*p)(*np.indices(data.shape)) - data)
	p, success = scipy.optimize.leastsq(errorfunction, params)
	return p 

def generate_dataset(k=2):
    shape = (4, 5, 6)
    rank = len(shape)
    init = [gen_rand(s, k) for s in shape]
    hidden = [gen_rand(s, k) for s in shape]
    x = parafac(hidden)
    x_indices = np.array([a.ravel() for a in np.indices(shape)]).T
    x_vals = x.ravel()
    return shape, rank, k, init, x, x_indices, x_vals 

def expand_coefficients(coefficients, entropy, length=None):
  """Expands coefficients to multiple terms that sum to each coefficient.

  Args:
    coefficients: Array, such that `coefficients[i, j, ..., k]` is the
        coefficient of x**i * y**j * ... * z**k.
    entropy: Float >= 0; the entropy to use for generating extra randomness.
    length: Number of terms that appear, e.g., 2x + 3 has two terms. If `None`
        then a suitable length will be picked depending on the entropy
        requested.

  Returns:
    Numpy object array with the same shape as `coefficients`, containing lists.
  """
  coefficients = np.asarray(coefficients)
  shape = coefficients.shape

  expanded_coefficients = np.empty(shape, dtype=np.object)

  min_length = np.count_nonzero(coefficients) + 2
  if length is None:
    max_length = min_length + int(math.ceil(entropy) / 2)
    length = random.randint(min_length, max_length)
  if length < min_length:
    length = min_length

  is_zero_flat = np.reshape(coefficients, [-1]) == 0
  counts = expanded_coefficient_counts(length, is_zero=is_zero_flat)
  coeffs_entropy = entropy * np.random.dirichlet(np.maximum(1e-9, counts - 1))
  counts = np.reshape(counts, shape)
  coeffs_entropy = np.reshape(coeffs_entropy, shape)

  indices = list(zip(*np.indices(shape).reshape([len(shape), -1])))
  for power in indices:
    coeffs = integers_with_sum(
        value=coefficients.item(power),
        count=counts.item(power),
        entropy=coeffs_entropy.item(power))
    expanded_coefficients.itemset(power, coeffs)

  return expanded_coefficients 

def differentiate(coefficients, axis):
  """Differentiate coefficients (corresponding to polynomial) along axis."""
  coefficients = np.asarray(coefficients)
  indices = list(range(1, coefficients.shape[axis]))
  coefficients = coefficients.take(axis=axis, indices=indices)

  broadcast_shape = np.ones(coefficients.ndim, dtype=np.int32)
  broadcast_shape[axis] = len(indices)
  broadcast = np.asarray(indices).reshape(broadcast_shape)

  result = broadcast * coefficients
  return trim(result) 

def argwhere(a):
    """
    Find the indices of array elements that are non-zero, grouped by element.

    Parameters
    ----------
    a : array_like
        Input data.

    Returns
    -------
    index_array : ndarray
        Indices of elements that are non-zero. Indices are grouped by element.

    See Also
    --------
    where, nonzero

    Notes
    -----
    ``np.argwhere(a)`` is the same as ``np.transpose(np.nonzero(a))``.

    The output of ``argwhere`` is not suitable for indexing arrays.
    For this purpose use ``nonzero(a)`` instead.

    Examples
    --------
    >>> x = np.arange(6).reshape(2,3)
    >>> x
    array([[0, 1, 2],
           [3, 4, 5]])
    >>> np.argwhere(x>1)
    array([[0, 2],
           [1, 0],
           [1, 1],
           [1, 2]])

    """
    return transpose(nonzero(a)) 

def flatnonzero(a):
    """
    Return indices that are non-zero in the flattened version of a.

    This is equivalent to np.nonzero(np.ravel(a))[0].

    Parameters
    ----------
    a : array_like
        Input data.

    Returns
    -------
    res : ndarray
        Output array, containing the indices of the elements of `a.ravel()`
        that are non-zero.

    See Also
    --------
    nonzero : Return the indices of the non-zero elements of the input array.
    ravel : Return a 1-D array containing the elements of the input array.

    Examples
    --------
    >>> x = np.arange(-2, 3)
    >>> x
    array([-2, -1,  0,  1,  2])
    >>> np.flatnonzero(x)
    array([0, 1, 3, 4])

    Use the indices of the non-zero elements as an index array to extract
    these elements:

    >>> x.ravel()[np.flatnonzero(x)]
    array([-2, -1,  1,  2])

    """
    return np.nonzero(np.ravel(a))[0] 

def test_take(self):
        tgt = [2, 3, 5]
        indices = [1, 2, 4]
        a = [1, 2, 3, 4, 5]

        out = np.take(a, indices)
        assert_equal(out, tgt) 

def test_results(self):
        a = np.arange(1*2*3*4).reshape(1, 2, 3, 4).copy()
        aind = np.indices(a.shape)
        assert_(a.flags['OWNDATA'])
        for (i, j) in self.tgtshape:
            # positive axis, positive start
            res = np.rollaxis(a, axis=i, start=j)
            i0, i1, i2, i3 = aind[np.array(res.shape) - 1]
            assert_(np.all(res[i0, i1, i2, i3] == a))
            assert_(res.shape == self.tgtshape[(i, j)], str((i,j)))
            assert_(not res.flags['OWNDATA'])

            # negative axis, positive start
            ip = i + 1
            res = np.rollaxis(a, axis=-ip, start=j)
            i0, i1, i2, i3 = aind[np.array(res.shape) - 1]
            assert_(np.all(res[i0, i1, i2, i3] == a))
            assert_(res.shape == self.tgtshape[(4 - ip, j)])
            assert_(not res.flags['OWNDATA'])

            # positive axis, negative start
            jp = j + 1 if j < 4 else j
            res = np.rollaxis(a, axis=i, start=-jp)
            i0, i1, i2, i3 = aind[np.array(res.shape) - 1]
            assert_(np.all(res[i0, i1, i2, i3] == a))
            assert_(res.shape == self.tgtshape[(i, 4 - jp)])
            assert_(not res.flags['OWNDATA'])

            # negative axis, negative start
            ip = i + 1
            jp = j + 1 if j < 4 else j
            res = np.rollaxis(a, axis=-ip, start=-jp)
            i0, i1, i2, i3 = aind[np.array(res.shape) - 1]
            assert_(np.all(res[i0, i1, i2, i3] == a))
            assert_(res.shape == self.tgtshape[(4 - ip, 4 - jp)])
            assert_(not res.flags['OWNDATA'])