Source code for jina.types.ndarray.sparse.scipy

from typing import List

import numpy as np

from . import BaseSparseNdArray
from ....proto import jina_pb2

if False:
    import scipy.sparse

__all__ = ['SparseNdArray']


[docs]class SparseNdArray(BaseSparseNdArray):
    """
    Scipy powered sparse ndarray.

    .. warning::
        scipy only supports ndim=2

    .. seealso::
        https://docs.scipy.org/doc/scipy/reference/sparse.html

    :param proto: the protobuf message, when not given then create a new one via :meth:`get_null_proto`
    :param sp_format: the sparse format of the scipy matrix. one of 'coo', 'bsr', 'csc', 'csr'.
    """

    def __init__(
        self,
        proto: 'jina_pb2.SparseNdArrayProto' = None,
        sp_format: str = 'coo',
        *args,
        **kwargs,
    ):
        """Set constructor method."""
        import scipy.sparse

        super().__init__(proto, *args, **kwargs)
        support_fmt = {'coo', 'bsr', 'csc', 'csr'}
        if sp_format in support_fmt:
            self.spmat_fn = getattr(scipy.sparse, f'{sp_format}_matrix')
        else:
            raise ValueError(
                f'{sp_format} sparse matrix is not supported, please choose one of those: {support_fmt}'
            )

[docs]    def sparse_constructor(
        self, indices: 'np.ndarray', values: 'np.ndarray', shape: List[int]
    ) -> 'scipy.sparse.spmatrix':
        """
        Sparse NdArray constructor for scipy.sparse.spmatrix.

        :param indices: the indices of the sparse array
        :param values: the values of the sparse array
        :param shape: the shape of the sparse array
        :return: SparseTensor
        """
        if indices.shape[-1] != 2:
            raise ValueError(
                f'scipy backend only supports ndim=2 sparse matrix, given {indices.shape}'
            )
        return self.spmat_fn((values, indices.T), shape=shape)

[docs]    def sparse_parser(self, value: 'scipy.sparse.spmatrix'):
        """
        Parse a scipy.sparse.spmatrix to indices, values and shape.

        :param value: the scipy.sparse.spmatrix.
        :return: a Dict with three entries {'indices': ..., 'values':..., 'shape':...}
        """
        v = value.tocoo()
        return {
            'indices': np.stack([v.row, v.col], axis=1),
            'values': v.data,
            'shape': v.shape,
        }