cellarr-array 0.0.3__py3-none-any.whl → 0.2.0__py3-none-any.whl

cellarr_array/{SparseCellArray.py → core/sparse.py}

@@ -3,14 +3,14 @@ try:
 except ImportError:
     # TODO: This is required for Python <3.10. Remove once Python 3.9 reaches EOL in October 2025
     EllipsisType = type(...)
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import tiledb
 from scipy import sparse
 
-from .CellArray import CellArray
 from .helpers import SliceHelper
+from .base import CellArray
 
 __author__ = "Jayaram Kancherla"
 __copyright__ = "Jayaram Kancherla"
@@ -22,18 +22,71 @@ class SparseCellArray(CellArray):
 
     def __init__(
         self,
-        uri: str,
+        uri: Optional[str] = None,
+        tiledb_array_obj: Optional[tiledb.Array] = None,
         attr: str = "data",
-        mode: str = None,
+        mode: Optional[Literal["r", "w", "d", "m"]] = None,
         config_or_context: Optional[Union[tiledb.Config, tiledb.Ctx]] = None,
         return_sparse: bool = True,
-        sparse_coerce: Union[sparse.csr_matrix, sparse.csc_matrix] = sparse.csr_matrix,
+        sparse_format: Union[sparse.csr_matrix, sparse.csc_matrix] = sparse.csr_matrix,
+        validate: bool = True,
+        **kwargs,
     ):
-        """Initialize SparseCellArray."""
-        super().__init__(uri, attr, mode, config_or_context)
+        """Initialize the object.
+
+        Args:
+            uri:
+                URI to the array.
+                Required if 'tiledb_array_obj' is not provided.
+
+            tiledb_array_obj:
+                Optional, an already opened ``tiledb.Array`` instance.
+                If provided, 'uri' can be None, and 'config_or_context' is ignored.
+
+            attr:
+                Attribute to access.
+                Defaults to "data".
+
+            mode:
+                Open the array object in read 'r', write 'w', modify
+                'm' mode, or delete 'd' mode.
+
+                Defaults to None for automatic mode switching.
+
+                If 'tiledb_array_obj' is provided, this mode should ideally match
+                the mode of the provided array or be None.
+
+            config_or_context:
+                Optional config or context object. Ignored if 'tiledb_array_obj' is provided,
+                as context will be derived from the object.
+
+                Defaults to None.
+
+            return_sparse:
+                Whether to return a sparse representation of the data when object is sliced.
+                Default is to return a dictionary that contains coordinates and values.
+
+            sparse_format:
+                Format to return, defaults to csr_matrix.
+
+            validate:
+                Whether to validate the attributes.
+                Defaults to True.
+
+            kwargs:
+                Additional arguments.
+        """
+        super().__init__(
+            uri=uri,
+            tiledb_array_obj=tiledb_array_obj,
+            attr=attr,
+            mode=mode,
+            config_or_context=config_or_context,
+            validate=validate,
+        )
 
         self.return_sparse = return_sparse
-        self.sparse_coerce = sparse.csr_matrix if sparse_coerce is None else sparse_coerce
+        self.sparse_format = sparse.csr_matrix if sparse_format is None else sparse_format
 
     def _validate_matrix_dims(self, data: sparse.spmatrix) -> Tuple[sparse.coo_matrix, bool]:
         """Validate and adjust matrix dimensions if needed.
@@ -73,7 +126,7 @@ class SparseCellArray(CellArray):
                 shape.append(idx.stop - (idx.start or 0))
             elif isinstance(idx, list):
                 shape.append(len(set(idx)))
-            else:  # single integer
+            else:
                 shape.append(1)
 
         # Always return (n,1) shape for CSR matrix
@@ -87,20 +140,17 @@ class SparseCellArray(CellArray):
         """Convert TileDB result to CSR format or dense array."""
         data = result[self._attr]
 
-        # empty result
         if len(data) == 0:
             print("is emoty")
             if not self.return_sparse:
                 return result
             else:
-                # For COO output, return empty sparse matrix
                 if self.ndim == 1:
-                    matrix = self.sparse_coerce((1, shape[0]))
+                    matrix = self.sparse_format((1, shape[0]))
                     return matrix[:, key[0]]
 
-                return self.sparse_coerce(shape)[key]
+                return self.sparse_format(shape)[key]
 
-        # Get coordinates
         coords = []
         for dim_name in self.dim_names:
             dim_coords = result[dim_name]
@@ -111,11 +161,12 @@ class SparseCellArray(CellArray):
             coords = [np.zeros_like(coords[0]), coords[0]]
             shape = (1, shape[0])
 
-        # Create sparse matrix
         matrix = sparse.coo_matrix((data, tuple(coords)), shape=shape)
-        if self.sparse_coerce in (sparse.csr_matrix, sparse.csr_array):
+
+        sliced = matrix
+        if self.sparse_format in (sparse.csr_matrix, sparse.csr_array):
             sliced = matrix.tocsr()
-        elif self.sparse_coerce in (sparse.csc_matrix, sparse.csc_array):
+        elif self.sparse_format in (sparse.csc_matrix, sparse.csc_array):
             sliced = matrix.tocsc()
 
         if self.ndim == 1:
@@ -147,7 +198,6 @@ class SparseCellArray(CellArray):
         if all(isinstance(idx, slice) for idx in optimized_key):
             return self._direct_slice(tuple(optimized_key))
 
-        # For mixed slice-list queries, adjust slice bounds
         tiledb_key = []
         for idx in key:
             if isinstance(idx, slice):
@@ -186,22 +236,20 @@ class SparseCellArray(CellArray):
         if not sparse.issparse(data):
             raise TypeError("Input must be a scipy sparse matrix.")
 
-        # Validate and adjust dimensions
-        data, is_1d = self._validate_matrix_dims(data)
+        coo_data, is_1d = self._validate_matrix_dims(data)
 
-        # Check bounds
-        end_row = start_row + data.shape[0]
+        end_row = start_row + coo_data.shape[0]
         if end_row > self.shape[0]:
             raise ValueError(
                 f"Write operation would exceed array bounds. End row {end_row} > array rows {self.shape[0]}."
             )
 
-        if not is_1d and data.shape[1] != self.shape[1]:
-            raise ValueError(f"Data columns {data.shape[1]} don't match array columns {self.shape[1]}.")
+        if not is_1d and coo_data.shape[1] != self.shape[1]:
+            raise ValueError(f"Data columns {coo_data.shape[1]} don't match array columns {self.shape[1]}.")
 
-        adjusted_rows = data.row + start_row
+        adjusted_rows = coo_data.row + start_row
         with self.open_array(mode="w") as array:
             if is_1d:
-                array[adjusted_rows] = data.data
+                array[adjusted_rows] = coo_data.data
             else:
-                array[adjusted_rows, data.col] = data.data
+                array[adjusted_rows, coo_data.col] = coo_data.data

cellarr_array/dataloaders/__init__.py

@@ -0,0 +1,3 @@
+from .denseloader import DenseArrayDataset, construct_dense_array_dataloader
+from .iterabledataloader import CellArrayIterableDataset, construct_iterable_dataloader
+from .sparseloader import SparseArrayDataset, construct_sparse_array_dataloader

cellarr_array/dataloaders/denseloader.py

@@ -0,0 +1,198 @@
+from typing import Optional
+from warnings import warn
+
+import numpy as np
+import tiledb
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+from ..core.dense import DenseCellArray
+
+__author__ = "Jayaram Kancherla"
+__copyright__ = "Jayaram Kancherla"
+__license__ = "MIT"
+
+
+class DenseArrayDataset(Dataset):
+    def __init__(
+        self,
+        array_uri: str,
+        attribute_name: str = "data",
+        num_rows: Optional[int] = None,
+        num_columns: Optional[int] = None,
+        cellarr_ctx_config: Optional[dict] = None,
+        transform=None,
+    ):
+        """PyTorch Dataset for dense TileDB arrays accessed via DenseCellArray.
+
+        Args:
+            array_uri:
+                URI of the TileDB dense array.
+
+            attribute_name:
+                Name of the attribute to read from.
+
+            num_rows:
+                Total number of rows in the dataset.
+                If None, will infer from `array.shape[0]`.
+
+            num_columns:
+                The number of columns in the dataset.
+                If None, will attempt to infer `from array.shape[1]`.
+
+            cellarr_ctx_config:
+                Optional TileDB context configuration dict for CellArray.
+
+            transform:
+                Optional transform to be applied on a sample.
+        """
+        self.array_uri = array_uri
+        self.attribute_name = attribute_name
+        self.cellarr_ctx_config = cellarr_ctx_config
+        self.transform = transform
+        self.cell_array_instance = None
+
+        if num_rows is not None and num_columns is not None:
+            self._len = num_rows
+            self.num_columns = num_columns
+        else:
+            # Infer the array shape
+            print(f"Dataset '{array_uri}': num_rows or num_columns not provided. Probing array...")
+            init_ctx_config = tiledb.Config(self.cellarr_ctx_config) if self.cellarr_ctx_config else None
+            try:
+                temp_arr = DenseCellArray(
+                    uri=self.array_uri, attr=self.attribute_name, config_or_context=init_ctx_config
+                )
+                if temp_arr.ndim == 1:
+                    self._len = num_rows if num_rows is not None else temp_arr.shape[0]
+                    self.num_columns = 1
+                elif temp_arr.ndim == 2:
+                    self._len = num_rows if num_rows is not None else temp_arr.shape[0]
+                    self.num_columns = num_columns if num_columns is not None else temp_arr.shape[1]
+                else:
+                    raise ValueError(f"Array ndim {temp_arr.ndim} not supported.")
+
+                print(f"Dataset '{array_uri}': Inferred shape. Rows: {self._len}, Columns: {self.num_columns}")
+
+            except Exception as e:
+                if num_rows is None or num_columns is None:
+                    raise ValueError(
+                        f"num_rows and num_columns must be provided if inferring array shape fails for '{array_uri}'."
+                    ) from e
+                self._len = num_rows
+                self.feature_dim = num_columns
+                warn(
+                    f"Falling back to provided or zero dimensions for '{array_uri}' due to inference error: {e}",
+                    RuntimeWarning,
+                )
+
+        if self.num_columns is None or self.num_columns <= 0 and self._len > 0:  # Check if num_columns is valid
+            raise ValueError(
+                f"num_columns ({self.num_columns}) is invalid or could not be determined for array '{array_uri}'."
+            )
+
+        if self._len == 0:
+            warn(f"Dataset for '{array_uri}' has length 0.", RuntimeWarning)
+
+    def _init_worker_state(self):
+        """Initializes the DenseCellArray instance for the current worker."""
+        if self.cell_array_instance is None:
+            ctx = tiledb.Ctx(self.cellarr_ctx_config) if self.cellarr_ctx_config else None
+            self.cell_array_instance = DenseCellArray(
+                uri=self.array_uri, attr=self.attribute_name, mode="r", config_or_context=ctx
+            )
+
+            # Sanity check: worker's shape against dataset's established dims
+            # if self.cell_array_instance.shape[0] != self._len or \
+            #    (self.cell_array_instance.ndim > 1 and self.cell_array_instance.shape[1] != self.feature_dim) or \
+            #    (self.cell_array_instance.ndim == 1 and self.feature_dim != 1) :
+            #     print(f"Warning: Worker for {self.array_uri} sees shape {self.cell_array_instance.shape} "
+            #           f"but dataset initialized with len={self._len}, feat={self.feature_dim}")
+
+    def __len__(self):
+        return self._len
+
+    def __getitem__(self, idx):
+        if not 0 <= idx < self._len:
+            raise IndexError(f"Index {idx} out of bounds for dataset of length {self._len}.")
+
+        self._init_worker_state()
+
+        if self.cell_array_instance.ndim == 2:
+            item_slice = (slice(idx, idx + 1), slice(None))
+        elif self.cell_array_instance.ndim == 1:
+            item_slice = slice(idx, idx + 1)
+        else:
+            raise ValueError(f"Array ndim {self.cell_array_instance.ndim} not supported in __getitem__.")
+
+        sample_data_np = self.cell_array_instance[item_slice]
+        if sample_data_np.ndim == 2 and sample_data_np.shape[0] == 1:
+            sample_data_np = sample_data_np.squeeze(0)
+        elif sample_data_np.ndim == 1 and sample_data_np.shape[0] == 1 and self.feature_dim == 1:
+            pass
+        elif sample_data_np.ndim == 0 and self.feature_dim == 1:
+            sample_data_np = np.array([sample_data_np])
+
+        if self.transform:
+            sample_data_np = self.transform(sample_data_np)
+
+        return torch.from_numpy(sample_data_np)
+
+
+def construct_dense_array_dataloader(
+    array_uri: str,
+    attribute_name: str = "data",
+    num_rows: Optional[int] = None,
+    num_columns: Optional[int] = None,
+    batch_size: int = 1000,
+    num_workers_dl: int = 2,
+) -> DataLoader:
+    """Construct an instance of `DenseArrayDataset` with PyTorch DataLoader.
+
+    Args:
+        array_uri:
+            URI of the TileDB array.
+
+        attribute_name:
+            Name of the attribute to read from.
+
+        num_rows:
+            The total number of rows in the TileDB array.
+
+        num_columns:
+            The total number of columns in the TileDB array.
+
+        batch_size:
+            Number of random samples per batch generated by the dataset.
+
+        num_workers_dl:
+            Number of worker processes for the DataLoader.
+    """
+    tiledb_ctx_config = {
+        "sm.tile_cache_size": 1000 * 1024**2,  # 1000 MB tile cache per worker
+        "sm.num_reader_threads": 4,
+    }
+
+    dataset = DenseArrayDataset(
+        array_uri=array_uri,
+        attribute_name=attribute_name,
+        num_rows=num_rows,
+        num_columns=num_columns,
+        cellarr_ctx_config=tiledb_ctx_config,
+    )
+
+    if len(dataset) == 0:
+        print("Dataset is empty, cannot create DataLoader.")
+        return
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers_dl,
+        pin_memory=True,
+        prefetch_factor=2,
+        persistent_workers=True if num_workers_dl > 0 else False,
+    )
+
+    return dataloader