Ejemplo n.º 1
0
    def _major_slice(self, idx, copy=False):
        """Index along the major axis where idx is a slice object.
        """

        if idx == slice(None):
            return self.copy() if copy else self

        M, N = self._swap(*self.shape)
        start, stop, step = idx.indices(M)
        M = len(range(start, stop, step))
        new_shape = self._swap(M, N)
        if M == 0 or self.nnz == 0:
            return self.__class__(new_shape)

        if step == 1:
            indptr, indices, data = _index._get_csr_submatrix_major_axis(
                self.indptr, self.indices, self.data, start, stop)
        else:
            rows = cupy.arange(start,
                               start + M * step,
                               step,
                               dtype=self.indptr.dtype)
            indptr, indices, data = _index._csr_row_index(
                rows, self.indptr, self.indices, self.data)

        return self.__class__((data, indices, indptr),
                              shape=new_shape,
                              copy=False)
Ejemplo n.º 2
0
    def _major_index_fancy(self, idx):
        """Index along the major axis where idx is an array of ints.
        """
        _, N = self._swap(*self.shape)
        M = idx.size
        new_shape = self._swap(M, N)
        if self.nnz == 0 or M == 0:
            return self.__class__(new_shape)

        return self.__class__(
            _index._csr_row_index(self.data, self.indices, self.indptr, idx),
            shape=new_shape, copy=False)
Ejemplo n.º 3
0
    def _major_index_fancy(self, idx):
        """Index along the major axis where idx is an array of ints.
        """
        _, N = self._swap(*self.shape)
        M = len(idx)
        new_shape = self._swap(M, N)
        if M == 0:
            return self.__class__(new_shape)

        row_nnz = cupy.diff(self.indptr)
        idx_dtype = self.indices.dtype
        res_indptr = cupy.zeros(M + 1, dtype=idx_dtype)
        cupy.cumsum(row_nnz[idx], out=res_indptr[1:])

        res_indices, res_data = _index._csr_row_index(idx, self.indptr,
                                                      self.indices, self.data,
                                                      res_indptr)

        return self.__class__((res_data, res_indices, res_indptr),
                              shape=new_shape,
                              copy=False)
Ejemplo n.º 4
0
    def _major_slice(self, idx, copy=False):
        """Index along the major axis where idx is a slice object.
        """

        if idx == slice(None):
            return self.copy() if copy else self

        M, N = self._swap(*self.shape)
        start, stop, step = idx.indices(M)
        M = len(range(start, stop, step))
        new_shape = self._swap(M, N)
        if M == 0:
            return self.__class__(new_shape)

        row_nnz = cupy.diff(self.indptr)
        idx_dtype = self.indices.dtype
        res_indptr = cupy.zeros(M + 1, dtype=idx_dtype)

        cupy.cumsum(row_nnz[idx], out=res_indptr[1:])

        if step == 1:
            idx_start = self.indptr[start]
            idx_stop = self.indptr[stop]
            res_indices = cupy.array(self.indices[idx_start:idx_stop],
                                     copy=copy)
            res_data = cupy.array(self.data[idx_start:idx_stop], copy=copy)
        else:
            rows = cupy.arange(start,
                               start + (res_indptr.size - 1) * step,
                               step,
                               dtype=res_indptr.dtype)
            res_indices, res_data = _index._csr_row_index(
                rows, self.indptr, self.indices, self.data, res_indptr)

        return self.__class__((res_data, res_indices, res_indptr),
                              shape=new_shape,
                              copy=False)