ngio 0.1.6__py3-none-any.whl → 0.2.0a2__py3-none-any.whl

ngio/tables/backends/_abstract_backend.py

@@ -0,0 +1,71 @@
+from abc import ABC, abstractmethod
+from collections.abc import Collection
+from typing import Literal
+
+from anndata import AnnData
+from pandas import DataFrame
+
+from ngio.utils import ZarrGroupHandler
+
+
+class AbstractTableBackend(ABC):
+    """Abstract class for table backends."""
+
+    def __init__(
+        self,
+        group_handler: ZarrGroupHandler,
+        index_key: str | None = None,
+        index_type: Literal["int", "str"] = "int",
+    ):
+        """Initialize the handler.
+
+        Args:
+            group_handler (ZarrGroupHandler): An object to handle the Zarr group
+                containing the table data.
+            index_key (str): The column name to use as the index of the DataFrame.
+            index_type (str): The type of the index column in the DataFrame.
+        """
+        self._group_handler = group_handler
+        self._index_key = index_key
+        self._index_type = index_type
+
+    @staticmethod
+    @abstractmethod
+    def backend_name() -> str:
+        """The name of the backend."""
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def implements_anndata() -> bool:
+        """Whether the handler implements the anndata protocol."""
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def implements_dataframe() -> bool:
+        """Whether the handler implements the dataframe protocol."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_columns(self) -> list[str]:
+        """List all labels in the group."""
+        raise NotImplementedError
+
+    def load_as_anndata(self, columns: Collection[str] | None = None) -> AnnData:
+        """Load the metadata in the store."""
+        raise NotImplementedError
+
+    def load_as_dataframe(self, columns: Collection[str] | None = None) -> DataFrame:
+        """List all labels in the group."""
+        raise NotImplementedError
+
+    def write_from_dataframe(
+        self, table: DataFrame, metadata: dict | None = None
+    ) -> None:
+        """Consolidate the metadata in the store."""
+        raise NotImplementedError
+
+    def write_from_anndata(self, table: AnnData, metadata: dict | None = None) -> None:
+        """Consolidate the metadata in the store."""
+        raise NotImplementedError

ngio/tables/backends/_anndata_utils.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+import pandas.api.types as ptypes
+import zarr
+from anndata import AnnData
+from anndata._io.specs import read_elem
+from anndata._io.utils import _read_legacy_raw
+from anndata._io.zarr import read_dataframe
+from anndata.compat import _clean_uns
+from anndata.experimental import read_dispatched
+
+from ngio.tables._validators import validate_index_dtype, validate_index_key
+from ngio.utils import (
+    NgioTableValidationError,
+    StoreOrGroup,
+    open_group_wrapper,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Collection
+
+
+def custom_read_zarr(
+    store: StoreOrGroup, elem_to_read: Collection[str] | None = None
+) -> AnnData:
+    """Read from a hierarchical Zarr array store.
+
+    # Implementation originally from https://github.com/scverse/anndata/blob/main/src/anndata/_io/zarr.py
+    # Original implementation would not work with remote storages so we had to copy it
+    # here and slightly modified it to work with remote storages.
+
+    Args:
+        store (StoreOrGroup): A store or group to read the AnnData from.
+        elem_to_read (Collection[str] | None): The elements to read from the store.
+    """
+    group, _ = open_group_wrapper(store=store, mode="r")
+
+    if elem_to_read is None:
+        elem_to_read = [
+            "X",
+            "obs",
+            "var",
+            "uns",
+            "obsm",
+            "varm",
+            "obsp",
+            "varp",
+            "layers",
+        ]
+
+    # Read with handling for backwards compat
+    def callback(func: Callable, elem_name: str, elem: Any, iospec: Any) -> Any:
+        if iospec.encoding_type == "anndata" or elem_name.endswith("/"):
+            ad_kwargs = {}
+            # Some of these elem fail on https
+            # So we only include the ones that are strictly necessary
+            # for fractal tables
+            # This fails on some https
+            # base_elem += list(elem.keys())
+            for k in elem_to_read:
+                v = elem.get(k)
+                if v is not None and not k.startswith("raw."):
+                    ad_kwargs[k] = read_dispatched(v, callback)  # type: ignore
+            return AnnData(**ad_kwargs)
+
+        elif elem_name.startswith("/raw."):
+            return None
+        elif elem_name in {"/obs", "/var"}:
+            return read_dataframe(elem)
+        elif elem_name == "/raw":
+            # Backwards compat
+            return _read_legacy_raw(group, func(elem), read_dataframe, func)
+        return func(elem)
+
+    adata = read_dispatched(group, callback=callback)  # type: ignore
+
+    # Backwards compat (should figure out which version)
+    if "raw.X" in group:
+        raw = AnnData(**_read_legacy_raw(group, adata.raw, read_dataframe, read_elem))  # type: ignore
+        raw.obs_names = adata.obs_names  # type: ignore
+        adata.raw = raw  # type: ignore
+
+    # Backwards compat for <0.7
+    if isinstance(group["obs"], zarr.Array):
+        _clean_uns(adata)
+
+    if not isinstance(adata, AnnData):
+        raise ValueError(f"Expected an AnnData object, but got {type(adata)}")
+    return adata
+
+
+def _check_for_mixed_types(series: pd.Series) -> None:
+    """Check if the column has mixed types."""
+    if series.apply(type).nunique() > 1:  # type: ignore
+        raise NgioTableValidationError(
+            f"Column {series.name} has mixed types: "
+            f"{series.apply(type).unique()}. "  # type: ignore
+            "Type of all elements must be the same."
+        )
+
+
+def _check_for_supported_types(series: pd.Series) -> Literal["str", "int", "numeric"]:
+    """Check if the column has supported types."""
+    if ptypes.is_string_dtype(series):
+        return "str"
+    if ptypes.is_integer_dtype(series):
+        return "int"
+    if ptypes.is_numeric_dtype(series):
+        return "numeric"
+    raise NgioTableValidationError(
+        f"Column {series.name} has unsupported type: {series.dtype}."
+        " Supported types are string and numerics."
+    )
+
+
+def dataframe_to_anndata(
+    dataframe: pd.DataFrame,
+    index_key: str | None = None,
+    overwrite: bool = False,
+) -> ad.AnnData:
+    """Convert a table DataFrame to an AnnData object.
+
+    Args:
+        dataframe (pd.DataFrame): A pandas DataFrame representing a fractal table.
+        index_key (str): The column name to use as the index of the DataFrame.
+            Default is None.
+        overwrite (bool): Whether to overwrite the index if a different index is found.
+            Default is False.
+    """
+    # Check if the index_key is present in the data frame + optional validations
+    dataframe = validate_index_key(dataframe, index_key, overwrite=overwrite)
+    dataframe = validate_index_dtype(dataframe, index_type="str")
+
+    str_columns, int_columns, num_columns = [], [], []
+    for c_name in dataframe.columns:
+        column_df = dataframe[c_name]
+        _check_for_mixed_types(column_df)  # Mixed types are not allowed in the table
+        c_type = _check_for_supported_types(
+            column_df
+        )  # Only string and numeric types are allowed
+
+        if c_type == "str":
+            str_columns.append(c_name)
+
+        elif c_type == "int":
+            int_columns.append(c_name)
+
+        elif c_type == "numeric":
+            num_columns.append(c_name)
+
+    # Converting all observations to string
+    obs_df = dataframe[str_columns + int_columns]
+    obs_df.index = dataframe.index
+
+    x_df = dataframe[num_columns]
+
+    if x_df.dtypes.nunique() > 1:
+        x_df = x_df.astype("float64")
+
+    if x_df.empty:
+        # If there are no numeric columns, create an empty array
+        # to avoid AnnData failing to create the object
+        x_df = np.zeros((len(obs_df), 0), dtype="float64")
+
+    return ad.AnnData(X=x_df, obs=obs_df)
+
+
+def anndata_to_dataframe(
+    anndata: ad.AnnData,
+    index_key: str | None = "label",
+    index_type: str = "int",
+    overwrite: bool = False,
+) -> pd.DataFrame:
+    """Convert a AnnData object representing a fractal table to a pandas DataFrame.
+
+    Args:
+        anndata (ad.AnnData): An AnnData object representing a fractal table.
+        index_key (str): The column name to use as the index of the DataFrame.
+            Default is 'label'.
+        index_type (str): The type of the index column in the DataFrame.
+            Either 'str' or 'int'. Default is 'int'.
+        overwrite (bool): Whether to overwrite the index if a different index is found.
+            Default is False.
+    """
+    dataframe = anndata.to_df()
+    dataframe[anndata.obs_keys()] = anndata.obs
+    dataframe = validate_index_key(dataframe, index_key, overwrite=overwrite)
+    dataframe = validate_index_dtype(dataframe, index_type)
+    return dataframe

ngio/tables/backends/_anndata_v1.py

@@ -0,0 +1,75 @@
+from collections.abc import Collection
+from pathlib import Path
+
+from anndata import AnnData
+from pandas import DataFrame
+
+from ngio.tables.backends._abstract_backend import AbstractTableBackend
+from ngio.tables.backends._anndata_utils import (
+    anndata_to_dataframe,
+    custom_read_zarr,
+    dataframe_to_anndata,
+)
+
+
+class AnnDataBackend(AbstractTableBackend):
+    """A class to load and write tables from/to an AnnData object."""
+
+    @staticmethod
+    def backend_name() -> str:
+        """The name of the backend."""
+        return "anndata_v1"
+
+    @staticmethod
+    def implements_anndata() -> bool:
+        """Whether the handler implements the anndata protocol."""
+        return True
+
+    @staticmethod
+    def implements_dataframe() -> bool:
+        """Whether the handler implements the dataframe protocol."""
+        return True
+
+    def load_columns(self) -> list[str]:
+        """List all labels in the group."""
+        return list(self.load_as_dataframe().columns)
+
+    def load_as_anndata(self, columns: Collection[str] | None = None) -> AnnData:
+        """Load the metadata in the store."""
+        anndata = custom_read_zarr(self._group_handler._group)
+        if columns is not None:
+            raise NotImplementedError(
+                "Selecting columns is not implemented for AnnData."
+            )
+        return anndata
+
+    def load_as_dataframe(self, columns: Collection[str] | None = None) -> DataFrame:
+        """List all labels in the group."""
+        dataframe = anndata_to_dataframe(
+            self.load_as_anndata(),
+            index_key=self._index_key,
+            index_type=self._index_type,
+        )
+        if columns is not None:
+            dataframe = dataframe[columns]
+        return dataframe
+
+    def write_from_dataframe(
+        self, table: DataFrame, metadata: dict | None = None
+    ) -> None:
+        """Consolidate the metadata in the store."""
+        anndata = dataframe_to_anndata(table, index_key=self._index_key)
+        self.write_from_anndata(anndata, metadata)
+
+    def write_from_anndata(self, table: AnnData, metadata: dict | None = None) -> None:
+        """Consolidate the metadata in the store."""
+        store = self._group_handler.store
+        if not isinstance(store, str | Path):
+            raise ValueError(
+                "To write an AnnData object the store must be a local path/str."
+            )
+
+        store = Path(store) / self._group_handler.group.path
+        table.write_zarr(store)
+        if metadata is not None:
+            self._group_handler.write_attrs(metadata)

ngio/tables/backends/_json_v1.py

@@ -0,0 +1,56 @@
+from collections.abc import Collection
+
+import pandas as pd
+from pandas import DataFrame
+
+from ngio.tables.backends._abstract_backend import AbstractTableBackend
+from ngio.utils import NgioFileNotFoundError
+
+
+class JsonTableBackend(AbstractTableBackend):
+    """A class to load and write small tables in the zarr group .attrs (json)."""
+
+    @staticmethod
+    def backend_name() -> str:
+        """The name of the backend."""
+        return "json_v1"
+
+    @staticmethod
+    def implements_anndata() -> bool:
+        """Whether the handler implements the anndata protocol."""
+        return False
+
+    @staticmethod
+    def implements_dataframe() -> bool:
+        """Whether the handler implements the dataframe protocol."""
+        return True
+
+    def load_columns(self) -> list[str]:
+        """List all labels in the group."""
+        return list(self.load_as_dataframe().columns)
+
+    def _get_table_group(self):
+        try:
+            table_group = self._group_handler.get_group(path="table")
+        except NgioFileNotFoundError:
+            table_group = self._group_handler.group.create_group("table")
+        return table_group
+
+    def load_as_dataframe(self, columns: Collection[str] | None = None) -> DataFrame:
+        """List all labels in the group."""
+        table_group = self._get_table_group()
+        table_dict = dict(table_group.attrs)
+        data_frame = pd.DataFrame.from_dict(table_dict)
+        if columns is not None:
+            data_frame = data_frame[columns]
+        return data_frame
+
+    def write_from_dataframe(
+        self, table: DataFrame, metadata: dict | None = None
+    ) -> None:
+        """Consolidate the metadata in the store."""
+        table_group = self._get_table_group()
+        table_group.attrs.clear()
+        table_group.attrs.update(table.to_dict())
+        if metadata is not None:
+            self._group_handler.write_attrs(metadata)

ngio/tables/backends/_table_backends.py

@@ -0,0 +1,102 @@
+"""Protocol for table backends handlers."""
+
+from collections.abc import Collection
+from typing import Literal, Protocol
+
+from anndata import AnnData
+from pandas import DataFrame
+
+from ngio.tables.backends._anndata_v1 import AnnDataBackend
+from ngio.tables.backends._json_v1 import JsonTableBackend
+from ngio.utils import NgioValueError, ZarrGroupHandler
+
+
+class TableBackendProtocol(Protocol):
+    def __init__(
+        self,
+        group_handler: ZarrGroupHandler,
+        index_key: str | None = None,
+        index_type: Literal["int", "str"] = "int",
+    ): ...
+
+    @staticmethod
+    def backend_name() -> str: ...
+
+    @staticmethod
+    def implements_anndata() -> bool: ...
+
+    @staticmethod
+    def implements_dataframe() -> bool: ...
+
+    def load_columns(self) -> list[str]: ...
+
+    def load_as_anndata(self, columns: Collection[str] | None = None) -> AnnData: ...
+
+    def load_as_dataframe(
+        self, columns: Collection[str] | None = None
+    ) -> DataFrame: ...
+
+    def write_from_dataframe(
+        self, table: DataFrame, metadata: dict | None = None
+    ) -> None: ...
+
+    def write_from_anndata(
+        self, table: AnnData, metadata: dict | None = None
+    ) -> None: ...
+
+
+class ImplementedTableBackends:
+    """A class to manage the available table backends."""
+
+    _instance = None
+    _implemented_backends: dict[str, type[TableBackendProtocol]]
+
+    def __new__(cls):
+        """Create a new instance of the class if it does not exist."""
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._implemented_backends = {}
+        return cls._instance
+
+    @property
+    def available_backends(self) -> list[str]:
+        """Return the available table backends."""
+        return list(self._implemented_backends.keys())
+
+    def get_backend(
+        self,
+        backend_name: str | None,
+        group_handler: ZarrGroupHandler,
+        index_key: str | None = None,
+        index_type: Literal["int", "str"] = "int",
+    ) -> TableBackendProtocol:
+        """Try to get a handler for the given store based on the metadata version."""
+        if backend_name is None:
+            # Default to anndata since it is currently
+            # the only backend in use.
+            backend_name = "anndata_v1"
+
+        if backend_name not in self._implemented_backends:
+            raise NgioValueError(f"Table backend {backend_name} not implemented.")
+        handler = self._implemented_backends[backend_name](
+            group_handler=group_handler, index_key=index_key, index_type=index_type
+        )
+        return handler
+
+    def add_backend(
+        self,
+        table_beckend: type[TableBackendProtocol],
+        overwrite: bool = False,
+    ):
+        """Register a new handler."""
+        backend_name = table_beckend.backend_name()
+        if backend_name in self._implemented_backends and not overwrite:
+            raise NgioValueError(
+                f"Table backend {backend_name} already implemented. "
+                "Use the `overwrite=True` parameter to overwrite it."
+            )
+        self._implemented_backends[backend_name] = table_beckend
+
+
+ImplementedTableBackends().add_backend(AnnDataBackend)
+ImplementedTableBackends().add_backend(JsonTableBackend)