Source code for delayedarray.Transpose

from typing import Callable, Optional, Tuple, Sequence, Any
from numpy import dtype, transpose
import numpy
import copy

from .DelayedOp import DelayedOp
from .SparseNdarray import SparseNdarray
from .extract_dense_array import extract_dense_array
from .extract_sparse_array import extract_sparse_array
from .create_dask_array import create_dask_array
from .chunk_grid import chunk_grid
from .is_sparse import is_sparse
from .is_masked import is_masked


__author__ = "ltla"
__copyright__ = "ltla"
__license__ = "MIT"




class Transpose(DelayedOp):
    """Delayed transposition, based on Bioconductor's ``DelayedArray::DelayedAperm`` class.

    This will create a matrix transpose in the 2-dimensional case; for a high-dimensional array, it will permute the
    dimensions.

    This class is intended for developers to construct new :py:class:`~delayedarray.DelayedArray.DelayedArray`
    instances. In general, end users should not be interacting with ``Transpose`` objects directly.
    """



    def __init__(self, seed, perm: Optional[Tuple[int, ...]]):
        """
        Args:
            seed:
                Any object that satisfies the seed contract,
                see :py:class:`~delayedarray.DelayedArray.DelayedArray` for details.

            perm:
                Tuple of length equal to the dimensionality of ``seed``,
                containing the permutation of dimensions.  If None, the
                dimension ordering is assumed to be reversed.
        """

        curshape = seed.shape
        ndim = len(curshape)
        if perm is not None:
            if len(perm) != ndim:
                raise ValueError(
                    "Dimensionality of 'seed' and 'perm' should be the same."
                )
        else:
            perm = (*range(ndim - 1, -1, -1),)

        final_shape = []
        for x in perm:
            final_shape.append(curshape[x])

        self._seed = seed
        self._perm = perm
        self._shape = (*final_shape,)


    @property
    def shape(self) -> Tuple[int, ...]:
        """
        Returns:
            Tuple of integers specifying the extent of each dimension of the
            transposed object.
        """
        return self._shape

    @property
    def dtype(self) -> dtype:
        """
        Returns:
            NumPy type for the transposed contents, same as ``seed``.
        """
        return self._seed.dtype

    @property
    def seed(self):
        """
        Returns:
            The seed object.
        """
        return self._seed

    @property
    def perm(self) -> Tuple[int, ...]:
        """
        Returns:
            Permutation of dimensions in the transposition.
        """
        return self._perm



def _simplify_transpose(x: Transpose) -> Any:
    seed = x.seed
    if not type(seed) is Transpose:
        # Don't use isinstance, we don't want to collapse for Transpose
        # subclasses that might be doing god knows what.
        return x

    new_perm = []
    noop = True
    for i, p in enumerate(x.perm):
        new_p = seed.perm[p]
        if new_p != i:
            noop = False
        new_perm.append(new_p)
    if noop:
        return seed.seed

    new_x = copy.copy(x)
    new_x._seed = seed.seed
    new_x._perm = (*new_perm,)
    return new_x


def _extract_array(x: Transpose, subset: Tuple[Sequence[int], ...], f: Callable):
    permsub = [None] * len(subset)
    for i, j in enumerate(x._perm):
        permsub[j] = subset[i]

    target = f(x._seed, (*permsub,))
    return transpose(target, axes=x._perm)




@extract_dense_array.register
def extract_dense_array_Transpose(x: Transpose, subset: Tuple[Sequence[int], ...]) -> numpy.ndarray:
    """See :py:meth:`~delayedarray.extract_dense_array.extract_dense_array`."""
    return _extract_array(x, subset, extract_dense_array)





@extract_sparse_array.register
def extract_sparse_array_Transpose(x: Transpose, subset: Tuple[Sequence[int], ...]) -> SparseNdarray:
    """See :py:meth:`~delayedarray.extract_sparse_array.extract_sparse_array`."""
    return _extract_array(x, subset, extract_sparse_array)





@create_dask_array.register
def create_dask_array_Transpose(x: Transpose):
    """See :py:meth:`~delayedarray.create_dask_array.create_dask_array`."""
    target = create_dask_array(x._seed)
    return transpose(target, axes=x._perm)





@chunk_grid.register
def chunk_grid_Transpose(x: Transpose):
    """See :py:meth:`~delayedarray.chunk_grid.chunk_grid`."""
    chunks = chunk_grid(x._seed)
    return chunks.transpose(x._perm)





@is_sparse.register
def is_sparse_Transpose(x: Transpose):
    """See :py:meth:`~delayedarray.is_sparse.is_sparse`."""
    return is_sparse(x._seed)





@is_masked.register
def is_masked_Transpose(x: Transpose):
    """See :py:meth:`~delayedarray.is_masked.is_masked`."""
    return is_masked(x._seed)