Source code for compressed_lists.numpy_list

from __future__ import annotations

from typing import Any, Dict, List, Optional, Sequence, Union
from warnings import warn

import biocutils as ut
import numpy as np

from .base import CompressedList
from .partition import Partitioning
from .split_generic import _generic_register_helper, splitAsCompressedList

__author__ = "Jayaram Kancherla"
__copyright__ = "Jayaram Kancherla"
__license__ = "MIT"




class CompressedNumpyList(CompressedList):
    """CompressedList implementation for lists of NumPy vectors."""



    def __init__(
        self,
        unlist_data: np.ndarray,
        partitioning: Partitioning,
        element_metadata: Optional[dict] = None,
        metadata: Optional[Union[Dict[str, Any], ut.NamedList]] = None,
        **kwargs,
    ):
        """Initialize a CompressedNumpyList.

        Args:
            unlist_data:
                List of NumPy vectors.

            partitioning:
                Partitioning object defining element boundaries.

            element_metadata:
                Optional metadata for elements.

            metadata:
                Optional general metadata.

            kwargs:
                Additional arguments.
        """

        if not isinstance(unlist_data, np.ndarray):
            if len(unlist_data) == 0:
                unlist_data = np.asarray([])
            else:
                try:
                    warn("trying to concatenate/coerce 'unlist_data' to a `np.ndarray`..")
                    unlist_data = np.concatenate(unlist_data)
                except Exception as e:
                    raise TypeError("'unlist_data' must be an `np.ndarray`, provided ", type(unlist_data)) from e

        super().__init__(
            unlist_data, partitioning, element_type=np.ndarray, element_metadata=element_metadata, metadata=metadata
        )




    @classmethod
    def from_list(
        cls,
        lst: List[np.ndarray],
        names: Optional[Union[ut.Names, Sequence[str]]] = None,
        metadata: Optional[dict] = None,
    ) -> CompressedNumpyList:
        """
        Create a `CompressedNumpyList` from a list of NumPy vectors.

        Args:
            lst:
                List of NumPy vectors.

            names:
                Optional names for list elements.

            metadata:
                Optional metadata.

        Returns:
            A new `CompressedNumpyList`.
        """
        partitioning = Partitioning.from_list(lst, names)

        if len(lst) == 0:
            unlist_data = np.array([])
        else:
            unlist_data = np.concatenate(lst)

        return cls(unlist_data, partitioning, metadata=metadata)




@splitAsCompressedList.register
def _(
    data: np.ndarray,
    groups_or_partitions: Union[list, Partitioning],
    names: Optional[Union[ut.Names, Sequence[str]]] = None,
    metadata: Optional[dict] = None,
) -> CompressedNumpyList:
    """Handle NumPy arrays."""

    partitioned_data, groups_or_partitions = _generic_register_helper(
        data=data, groups_or_partitions=groups_or_partitions, names=names
    )

    if not isinstance(partitioned_data, np.ndarray) and len(partitioned_data) != 0:
        partitioned_data = ut.combine_sequences(*partitioned_data)

    return CompressedNumpyList(unlist_data=partitioned_data, partitioning=groups_or_partitions, metadata=metadata)