Python numpy.indices() Examples

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 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 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 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 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 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 _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 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 _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 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 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 _format_spec(self, indices, iops, lx=None, li=None, lls=None, VT100=True):
        if li is None: li = self._max_len_index(indices)
        if lx is None: lx = self._max_len_values()
        if lls is None: lls = [self._max_len_names(iop, name) for name, iop in iops.items()]

        if VT100:
            format_spec = ["  \033[1m{{index!s:<{0}}}\033[0;0m".format(li),"{{value!s:>{0}}}".format(lx)]
        else:
            format_spec = ["  {{index!s:<{0}}}".format(li),"{{value!s:>{0}}}".format(lx)]

        for opname, l in zip(iops, lls):
            f = '{{{1}!s:^{0}}}'.format(l, opname)
            format_spec.append(f)
        return format_spec 

def improve_ipopt(x0, prob, *args, **kwargs):
    try:
        import pyipopt
    except ImportError:
        raise Exception("PyIpopt package is not installed.")

    lb = pyipopt.NLP_LOWER_BOUND_INF
    ub = pyipopt.NLP_UPPER_BOUND_INF
    g_L = np.zeros(prob.m)
    for i in range(prob.m):
        if prob.fs[i].relop == '<=':
            g_L[i] = lb
    g_U = np.zeros(prob.m)

    def eval_grad_f(x, user_data = None):
        return 2*prob.f0.P.dot(x) + prob.f0.qarray
    def eval_g(x, user_data = None):
        return np.array([f.eval(x) for f in prob.fs])

    jac_grid = np.indices((prob.m, prob.n))
    jac_r = jac_grid[0].ravel()
    jac_c = jac_grid[1].ravel()
    def eval_jac_g(x, flag, user_data = None):
        if flag:
            return (jac_r, jac_c)
        else:
            return np.vstack([2*f.P.dot(x)+f.qarray for f in prob.fs])

    nlp = pyipopt.create(
        prob.n, lb*np.ones(prob.n), ub*np.ones(prob.n),
        prob.m, g_L, g_U, prob.m*prob.n, 0,
        prob.f0.eval, eval_grad_f,
        eval_g, eval_jac_g
    )
    try:
        x, zl, zu, constraint_multipliers, obj, status = nlp.solve(x0)
    except:
        pass

    return x 

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 map_incoming_flat(self, incoming_indices):
        """Map (flat) incoming indices to an index in the outgoing pipe.

        Parameters
        ----------
        incoming_indices : iterable of int
            One (flat) index on each of the incoming legs.

        Returns
        -------
        outgoing_index : int
            The index in the outgoing leg.
        """
        # need to calculate the `a_j` in the Notes of the doc-string of self.
        if len(incoming_indices) != self.nlegs:
            raise ValueError("wrong len of flat_ind_incoming")
        qind_in = np.empty((1, self.nlegs), dtype=np.intp)
        within_block_out = 0
        stride = 1
        for ax in range(self.nlegs - 1, -1, -1):  # reversed: C order within the block
            leg = self.legs[ax]
            qind, within_block = leg.get_qindex(incoming_indices[ax])
            qind_in[0, ax] = qind
            within_block_out += stride * within_block
            stride *= (leg.slices[qind + 1] - leg.slices[qind])
        j = self._map_incoming_qind(qind_in)[0]
        q_map = self.q_map[j, :]
        assert (q_map[1] - q_map[0] == stride)
        qind_out = q_map[2]  # I_s
        return self.slices[qind_out] + q_map[0] + within_block_out 

def project(self, mask):
        """Return copy keeping only the indices specified by `mask`.

        Parameters
        ----------
        mask : 1D array(bool)
            Whether to keep of the indices.

        Returns
        -------
        map_qind : 1D array
            Map of qindices, such that ``qind_new = map_qind[qind_old]``,
            and ``map_qind[qind_old] = -1`` for qindices projected out.
        block_masks : 1D array
            The bool mask for each of the *remaining* blocks.
        projected_copy : :class:`LegCharge`
            Copy of self with the qind projected by `mask`.
        """
        mask = np.asarray(mask, dtype=np.bool_)
        cp = self.copy()
        block_masks = [mask[b:e] for b, e in self._slice_start_stop()]
        new_block_lens = [np.sum(bm) for bm in block_masks]
        keep = np.nonzero(new_block_lens)[0]
        block_masks = [block_masks[i] for i in keep]
        cp._set_charges(cp.charges[keep])
        map_qind = -np.ones(self.block_number, np.intp)
        map_qind[keep] = np.arange(len(keep))
        cp._set_block_sizes(np.array(new_block_lens, dtype=np.intp)[keep])
        cp.bunched = self.is_blocked()  # no, it's not `is_bunched`
        return map_qind, block_masks, cp 

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 from_trivial(cls, ind_len, chargeinfo=None, qconj=1):
        """Create trivial (qnumber=0) LegCharge for given len of indices `ind_len`."""
        if chargeinfo is None:
            chargeinfo = ChargeInfo()
            charges = [[]]
        else:
            charges = [[0] * chargeinfo.qnumber]
        return cls(chargeinfo, [0, ind_len], charges, qconj) 

def _advanced_setitem_npc(self, inds, other):
        """Self[inds] = other for non-integer `inds` and :class:`Array` `other`.

        This function is called by self.__setitem__(inds, other)."""
        # suppress warning if we project a pipe
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            map_part2self, permutations, self_part = self._advanced_getitem(inds,
                                                                            calc_map_qind=True,
                                                                            permute=False)
        # permuations are ignored by map_part2self.
        # instead of figuring out permuations in self, apply the *reversed* permutations ot other
        for ax, perm in permutations:
            other = other.permute(inverse_permutation(perm), ax)
        # now test compatibility of self_part with `other`
        if self_part.rank != other.rank:
            raise IndexError("wrong number of indices")
        for pl, ol in zip(self_part.legs, other.legs):
            pl.test_contractible(ol.conj())
        if np.any(self_part.qtotal != other.qtotal):
            raise ValueError("wrong charge for assinging self[inds] = other")
        # note: a block exists in self_part, if and only if its extended version exists in self.
        # by definition, non-existent blocks in `other` are zero.
        # instead of checking which blocks are non-existent,
        # we first set self[inds] completely to zero
        for p_qindices in self_part._qdata:
            qindices, block_mask = map_part2self(p_qindices)
            block = self.get_block(qindices)
            block[block_mask] = 0.  # overwrite data in self
        # now we copy blocks from other
        for o_block, o_qindices in zip(other._data, other._qdata):
            qindices, block_mask = map_part2self(o_qindices)
            block = self.get_block(qindices, insert=True)
            block[block_mask] = o_block  # overwrite data in self
        self.ipurge_zeros(0.)  # remove blocks identically zero 

def _pre_indexing(self, inds):
        """Check if `inds` are valid indices for ``self[inds]`` and replaces Ellipsis by slices.

        Returns
        -------
        only_integer : bool
            Whether all of `inds` are (convertible to) np.intp.
        inds : tuple, len=self.rank
            `inds`, where ``Ellipsis`` is replaced by the correct number of slice(None).
        """
        if type(inds) != tuple:  # for rank 1
            inds = (inds, )
        i = next((i for i, idx in enumerate(inds) if idx is Ellipsis), None)
        # i is index of Ellipsis or None if we don't have one
        if i is None and len(inds) < self.rank:  # need an Ellipsis
            i = len(inds)
            inds = inds + (Ellipsis, )
        if i is not None:
            # replace Ellipsis with slice(None)
            fill = tuple([slice(None)] * (self.rank - len(inds) + 1))
            inds = inds[:i] + fill + inds[i + 1:]
        if len(inds) > self.rank:
            raise IndexError("too many indices for Array")
        # do we have only integer entries in `inds`?
        try:
            only_int = np.array(inds, dtype=np.intp)
            assert (only_int.shape == (len(inds), ))
        except:
            return False, inds
        else:
            return True, inds 

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 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]