ngio 0.2.9__py3-none-any.whl → 0.3.0__py3-none-any.whl

ngio/common/__init__.py

@@ -24,6 +24,15 @@ from ngio.common._slicer import (
     numpy_get_slice,
     numpy_set_slice,
 )
+from ngio.common._table_ops import (
+    concatenate_image_tables,
+    concatenate_image_tables_as,
+    concatenate_image_tables_as_async,
+    concatenate_image_tables_async,
+    conctatenate_tables,
+    list_image_tables,
+    list_image_tables_async,
+)
 from ngio.common._zoom import dask_zoom, numpy_zoom
 
 __all__ = [
@@ -34,6 +43,11 @@ __all__ = [
     "SliceTransform",
     "compute_and_slices",
     "compute_masking_roi",
+    "concatenate_image_tables",
+    "concatenate_image_tables_as",
+    "concatenate_image_tables_as_async",
+    "concatenate_image_tables_async",
+    "conctatenate_tables",
     "consolidate_pyramid",
     "dask_get_slice",
     "dask_set_slice",
@@ -41,6 +55,8 @@ __all__ = [
     "get_masked_pipe",
     "get_pipe",
     "init_empty_pyramid",
+    "list_image_tables",
+    "list_image_tables_async",
     "numpy_get_slice",
     "numpy_set_slice",
     "numpy_zoom",

ngio/common/_array_pipe.py

@@ -1,10 +1,11 @@
 from collections.abc import Collection, Iterable
 from typing import Literal
 
-import dask
-import dask.delayed
+import dask.array as da
 import numpy as np
 import zarr
+from dask.array import Array as DaskArray
+from dask.delayed import Delayed, delayed
 
 from ngio.common._axes_transforms import transform_dask_array, transform_numpy_array
 from ngio.common._common_types import ArrayLike
@@ -55,26 +56,26 @@ def _numpy_get_pipe(
     slices: SliceTransform,
     transformations: tuple[AxesTransformation, ...],
 ) -> np.ndarray:
-    array = numpy_get_slice(array, slices)
-    return transform_numpy_array(array, transformations)
+    _array = numpy_get_slice(array, slices)
+    return transform_numpy_array(_array, transformations)
 
 
 def _delayed_numpy_get_pipe(
     array: zarr.Array,
     slices: SliceTransform,
     transformations: tuple[AxesTransformation, ...],
-) -> dask.delayed:
-    array = dask.delayed(numpy_get_slice)(array, slices)
-    return dask.delayed(transform_numpy_array)(array, transformations)
+) -> Delayed:
+    _array = delayed(numpy_get_slice)(array, slices)
+    return delayed(transform_numpy_array)(_array, transformations)
 
 
 def _dask_get_pipe(
     array: zarr.Array,
     slices: SliceTransform,
     transformations: tuple[AxesTransformation, ...],
-) -> dask.array:
-    array = dask_get_slice(array, slices)
-    return transform_dask_array(array, transformations)
+) -> DaskArray:
+    _array = dask_get_slice(array, slices)
+    return transform_dask_array(_array, transformations)
 
 
 def _numpy_set_pipe(
@@ -89,22 +90,22 @@ def _numpy_set_pipe(
 
 def _dask_set_pipe(
     array: zarr.Array,
-    patch: np.ndarray,
+    patch: DaskArray,
     slices: SliceTransform,
     transformations: tuple[AxesTransformation, ...],
 ) -> None:
-    patch = transform_dask_array(patch, transformations)
-    dask_set_slice(array, patch, slices)
+    _patch = transform_dask_array(patch, transformations)
+    dask_set_slice(array, _patch, slices)
 
 
 def _delayed_numpy_set_pipe(
     array: zarr.Array,
-    patch: np.ndarray,
+    patch: np.ndarray | Delayed,
     slices: SliceTransform,
     transformations: tuple[AxesTransformation, ...],
-) -> dask.delayed:
-    patch = dask.delayed(transform_numpy_array)(patch, transformations)
-    return dask.delayed(numpy_set_slice)(array, patch, slices)
+) -> Delayed:
+    _patch = delayed(transform_numpy_array)(patch, transformations)
+    return delayed(numpy_set_slice)(array, _patch, slices)
 
 
 def get_pipe(
@@ -144,7 +145,7 @@ def set_pipe(
     slices, transformations = _compute_to_disk_transforms(
         dimensions=dimensions, axes_order=axes_order, **slice_kwargs
     )
-    if isinstance(patch, dask.array.Array):
+    if isinstance(patch, DaskArray):
         _dask_set_pipe(
             array=array, patch=patch, slices=slices, transformations=transformations
         )
@@ -152,7 +153,7 @@ def set_pipe(
         _numpy_set_pipe(
             array=array, patch=patch, slices=slices, transformations=transformations
         )
-    elif isinstance(patch, dask.delayed.Delayed):
+    elif isinstance(patch, Delayed):
         _delayed_numpy_set_pipe(
             array=array, patch=patch, slices=slices, transformations=transformations
         )
@@ -193,12 +194,15 @@ def _mask_pipe_common(
         **slice_kwargs,
     )
 
-    if isinstance(array_patch, np.ndarray):
+    if isinstance(array_patch, np.ndarray) and isinstance(label_patch, np.ndarray):
         label_patch = np.broadcast_to(label_patch, array_patch.shape)
-    elif isinstance(array_patch, dask.array.Array):
-        label_patch = dask.array.broadcast_to(label_patch, array_patch.shape)
+    elif isinstance(array_patch, DaskArray) and isinstance(label_patch, DaskArray):
+        label_patch = da.broadcast_to(label_patch, array_patch.shape)
     else:
-        raise NgioValueError(f"Mode {mode} not yet supported for masked array.")
+        raise NgioValueError(
+            "Incompatible types for array and label: "
+            f"{type(array_patch)} and {type(label_patch)}"
+        )
 
     mask = label_patch == label
     return array_patch, mask
@@ -225,7 +229,14 @@ def get_masked_pipe(
         mode=mode,
         **slice_kwargs,
     )
-    array_patch[~mask] = 0
+    if isinstance(array_patch, np.ndarray):
+        array_patch[~mask] = 0
+    elif isinstance(array_patch, DaskArray):
+        array_patch = da.where(mask, array_patch, 0)
+    else:
+        raise NgioValueError(
+            "Mode not yet supported for masked array. Expected a numpy or dask array."
+        )
     return array_patch
 
 
@@ -240,7 +251,7 @@ def set_masked_pipe(
     axes_order: Collection[str] | None = None,
     **slice_kwargs: slice | int | Iterable[int],
 ):
-    if isinstance(patch, dask.array.Array):
+    if isinstance(patch, DaskArray):
         mode = "dask"
     elif isinstance(patch, np.ndarray):
         mode = "numpy"
@@ -259,7 +270,19 @@ def set_masked_pipe(
         mode=mode,
         **slice_kwargs,
     )
-    patch = np.where(mask, patch, array_patch)
+    if isinstance(patch, np.ndarray):
+        assert isinstance(array_patch, np.ndarray)
+        _patch = np.where(mask, patch, array_patch)
+    elif isinstance(patch, DaskArray):
+        _patch = da.where(mask, patch, array_patch)
+    else:
+        raise NgioValueError(
+            "Mode not yet supported for masked array. Expected a numpy or dask array."
+        )
     set_pipe(
-        array, patch, dimensions=dimensions_array, axes_order=axes_order, **slice_kwargs
+        array,
+        _patch,
+        dimensions=dimensions_array,
+        axes_order=axes_order,
+        **slice_kwargs,
     )

ngio/common/_table_ops.py

@@ -0,0 +1,471 @@
+"""Aggregation and filtering operations for tables."""
+
+import asyncio
+from collections import Counter
+from collections.abc import Collection
+from dataclasses import dataclass, field
+from typing import Literal
+
+import pandas as pd
+import polars as pl
+
+from ngio.images._ome_zarr_container import OmeZarrContainer
+from ngio.tables import Table, TableType
+
+
+@dataclass
+class TableWithExtras:
+    """A class to hold a table and its extras."""
+
+    table: Table
+    extras: dict[str, str] = field(default_factory=dict)
+
+
+def _reindex_dataframe(
+    dataframe, index_cols: list[str], index_key: str | None = None
+) -> pd.DataFrame:
+    """Reindex a dataframe using an hash of the index columns."""
+    # Reindex the dataframe
+    old_index = dataframe.index.name
+    if old_index is not None:
+        dataframe = dataframe.reset_index()
+        index_cols.append(old_index)
+    dataframe.index = dataframe[index_cols].astype(str).agg("_".join, axis=1)
+
+    if index_key is None:
+        dataframe.index.name = index_key
+    return dataframe
+
+
+def _add_const_columns(
+    dataframe: pd.DataFrame,
+    new_cols: dict[str, str],
+    index_key: str | None = None,
+) -> pd.DataFrame:
+    for col, value in new_cols.items():
+        dataframe[col] = value
+
+    if index_key is not None:
+        dataframe = _reindex_dataframe(
+            dataframe=dataframe,
+            index_cols=list(new_cols.keys()),
+            index_key=index_key,
+        )
+    return dataframe
+
+
+def _add_const_columns_pl(
+    dataframe: pl.LazyFrame,
+    new_cols: dict[str, str],
+    index_key: str | None = None,
+) -> pl.LazyFrame:
+    dataframe = dataframe.with_columns(
+        [pl.lit(value, dtype=pl.String()).alias(col) for col, value in new_cols.items()]
+    )
+
+    if index_key is not None:
+        dataframe = dataframe.with_columns(
+            [
+                pl.concat_str(
+                    [pl.col(col) for col in new_cols.keys()],
+                    separator="_",
+                ).alias(index_key)
+            ]
+        )
+    return dataframe
+
+
+def _pd_concat(
+    tables: Collection[TableWithExtras], index_key: str | None = None
+) -> pd.DataFrame:
+    """Concatenate tables from different plates into a single table."""
+    if len(tables) == 0:
+        raise ValueError("No tables to concatenate.")
+
+    dataframes = []
+    for table in tables:
+        dataframe = _add_const_columns(
+            dataframe=table.table.dataframe, new_cols=table.extras, index_key=index_key
+        )
+        dataframes.append(dataframe)
+    concatenated_table = pd.concat(dataframes, axis=0)
+    return concatenated_table
+
+
+def _pl_concat(
+    tables: Collection[TableWithExtras], index_key: str | None = None
+) -> pl.LazyFrame:
+    """Concatenate tables from different plates into a single table."""
+    if len(tables) == 0:
+        raise ValueError("No tables to concatenate.")
+
+    dataframes = []
+    for table in tables:
+        polars_ls = _add_const_columns_pl(
+            dataframe=table.table.lazy_frame,
+            new_cols=table.extras,
+            index_key=index_key,
+        )
+        dataframes.append(polars_ls)
+
+    concatenated_table = pl.concat(dataframes, how="vertical")
+    return concatenated_table
+
+
+def conctatenate_tables(
+    tables: Collection[TableWithExtras],
+    mode: Literal["eager", "lazy"] = "eager",
+    index_key: str | None = None,
+    table_cls: type[TableType] | None = None,
+) -> Table:
+    """Concatenate tables from different plates into a single table."""
+    if len(tables) == 0:
+        raise ValueError("No tables to concatenate.")
+
+    table0 = next(iter(tables)).table
+
+    if mode == "lazy":
+        concatenated_table = _pl_concat(tables=tables, index_key=index_key)
+    elif mode == "eager":
+        concatenated_table = _pd_concat(tables=tables, index_key=index_key)
+    else:
+        raise ValueError(f"Unknown mode: {mode}. Use 'eager' or 'lazy'.")
+
+    meta = table0.meta
+    meta.index_key = index_key
+    meta.index_type = "str"
+
+    if table_cls is not None:
+        return table_cls.from_table_data(
+            table_data=concatenated_table,
+            meta=meta,
+        )
+    return table0.from_table_data(
+        table_data=concatenated_table,
+        meta=meta,
+    )
+
+
+def _check_images_and_extras(
+    images: Collection[OmeZarrContainer],
+    extras: Collection[dict[str, str]],
+) -> None:
+    """Check if the images and extras are valid."""
+    if len(images) == 0:
+        raise ValueError("No images to concatenate.")
+
+    if len(images) != len(extras):
+        raise ValueError("The number of images and extras must be the same.")
+
+
+def _concatenate_image_tables(
+    images: Collection[OmeZarrContainer],
+    extras: Collection[dict[str, str]],
+    table_name: str,
+    table_cls: type[TableType] | None = None,
+    index_key: str | None = None,
+    strict: bool = True,
+    mode: Literal["eager", "lazy"] = "eager",
+) -> Table:
+    """Concatenate tables from different images into a single table."""
+    _check_images_and_extras(images=images, extras=extras)
+
+    tables = []
+    for image, extra in zip(images, extras, strict=True):
+        if not strict and table_name not in image.list_tables():
+            continue
+        table = image.get_table(table_name)
+        tables.append(TableWithExtras(table=table, extras=extra))
+
+    return conctatenate_tables(
+        tables=tables,
+        mode=mode,
+        index_key=index_key,
+        table_cls=table_cls,
+    )
+
+
+def concatenate_image_tables(
+    images: Collection[OmeZarrContainer],
+    extras: Collection[dict[str, str]],
+    table_name: str,
+    index_key: str | None = None,
+    strict: bool = True,
+    mode: Literal["eager", "lazy"] = "eager",
+) -> Table:
+    """Concatenate tables from different images into a single table.
+
+    Args:
+        images: A collection of images.
+        extras: A collection of extras dictionaries for each image.
+            this will be added as columns to the table, and will be
+            concatenated with the table index to create a new index.
+        table_name: The name of the table to concatenate.
+        index_key: The key to use for the index of the concatenated table.
+        strict: If True, raise an error if the table is not found in the image.
+        mode: The mode to use for concatenation. Can be 'eager' or 'lazy'.
+            if 'eager', the table will be loaded into memory.
+            if 'lazy', the table will be loaded as a lazy frame.
+    """
+    return _concatenate_image_tables(
+        images=images,
+        extras=extras,
+        table_name=table_name,
+        table_cls=None,
+        index_key=index_key,
+        strict=strict,
+        mode=mode,
+    )
+
+
+def concatenate_image_tables_as(
+    images: Collection[OmeZarrContainer],
+    extras: Collection[dict[str, str]],
+    table_name: str,
+    table_cls: type[TableType],
+    index_key: str | None = None,
+    strict: bool = True,
+    mode: Literal["eager", "lazy"] = "eager",
+) -> TableType:
+    """Concatenate tables from different images into a single table.
+
+    Args:
+        images: A collection of images.
+        extras: A collection of extras dictionaries for each image.
+            this will be added as columns to the table, and will be
+            concatenated with the table index to create a new index.
+        table_name: The name of the table to concatenate.
+        table_cls: The output will be casted to this class, if the new table_cls is
+            compatible with the table_cls of the input tables.
+        index_key: The key to use for the index of the concatenated table.
+        strict: If True, raise an error if the table is not found in the image.
+        mode: The mode to use for concatenation. Can be 'eager' or 'lazy'.
+            if 'eager', the table will be loaded into memory.
+            if 'lazy', the table will be loaded as a lazy frame.
+    """
+    table = _concatenate_image_tables(
+        images=images,
+        extras=extras,
+        table_name=table_name,
+        table_cls=table_cls,
+        index_key=index_key,
+        strict=strict,
+        mode=mode,
+    )
+    if not isinstance(table, table_cls):
+        raise ValueError(f"Table is not of type {table_cls}. Got {type(table)}")
+    return table
+
+
+async def _concatenate_image_tables_async(
+    images: Collection[OmeZarrContainer],
+    extras: Collection[dict[str, str]],
+    table_name: str,
+    table_cls: type[TableType] | None = None,
+    index_key: str | None = None,
+    strict: bool = True,
+    mode: Literal["eager", "lazy"] = "eager",
+) -> Table:
+    """Concatenate tables from different images into a single table."""
+    _check_images_and_extras(images=images, extras=extras)
+
+    def process_image(
+        image: OmeZarrContainer,
+        table_name: str,
+        extra: dict[str, str],
+        mode: Literal["eager", "lazy"] = "eager",
+        strict: bool = True,
+    ) -> TableWithExtras | None:
+        """Process a single image and return the table."""
+        if not strict and table_name not in image.list_tables():
+            return None
+        _table = image.get_table(table_name)
+        if mode == "lazy":
+            # make sure the table is loaded lazily
+            # It the backend is not lazy, this will be
+            # loaded eagerly
+            _ = _table.lazy_frame
+        elif mode == "eager":
+            # make sure the table is loaded eagerly
+            _ = _table.dataframe
+        table = TableWithExtras(
+            table=_table,
+            extras=extra,
+        )
+        return table
+
+    tasks = []
+    for image, extra in zip(images, extras, strict=True):
+        task = asyncio.to_thread(
+            process_image,
+            image=image,
+            table_name=table_name,
+            extra=extra,
+            strict=strict,
+        )
+        tasks.append(task)
+    tables = await asyncio.gather(*tasks)
+    tables = [table for table in tables if table is not None]
+    return conctatenate_tables(
+        tables=tables,
+        mode=mode,
+        index_key=index_key,
+        table_cls=table_cls,
+    )
+
+
+async def concatenate_image_tables_async(
+    images: Collection[OmeZarrContainer],
+    extras: Collection[dict[str, str]],
+    table_name: str,
+    index_key: str | None = None,
+    strict: bool = True,
+    mode: Literal["eager", "lazy"] = "eager",
+) -> Table:
+    """Concatenate tables from different images into a single table.
+
+    Args:
+        images: A collection of images.
+        extras: A collection of extras dictionaries for each image.
+            this will be added as columns to the table, and will be
+            concatenated with the table index to create a new index.
+        table_name: The name of the table to concatenate.
+        index_key: The key to use for the index of the concatenated table.
+        strict: If True, raise an error if the table is not found in the image.
+        mode: The mode to use for concatenation. Can be 'eager' or 'lazy'.
+            if 'eager', the table will be loaded into memory.
+            if 'lazy', the table will be loaded as a lazy frame.
+    """
+    return await _concatenate_image_tables_async(
+        images=images,
+        extras=extras,
+        table_name=table_name,
+        table_cls=None,
+        index_key=index_key,
+        strict=strict,
+        mode=mode,
+    )
+
+
+async def concatenate_image_tables_as_async(
+    images: Collection[OmeZarrContainer],
+    extras: Collection[dict[str, str]],
+    table_name: str,
+    table_cls: type[TableType],
+    index_key: str | None = None,
+    strict: bool = True,
+    mode: Literal["eager", "lazy"] = "eager",
+) -> TableType:
+    """Concatenate tables from different images into a single table.
+
+    Args:
+        images: A collection of images.
+        extras: A collection of extras dictionaries for each image.
+            this will be added as columns to the table, and will be
+            concatenated with the table index to create a new index.
+        table_name: The name of the table to concatenate.
+        table_cls: The output will be casted to this class, if the new table_cls is
+            compatible with the table_cls of the input tables.
+        index_key: The key to use for the index of the concatenated table.
+        strict: If True, raise an error if the table is not found in the image.
+        mode: The mode to use for concatenation. Can be 'eager' or 'lazy'.
+            if 'eager', the table will be loaded into memory.
+            if 'lazy', the table will be loaded as a lazy frame.
+    """
+    table = await _concatenate_image_tables_async(
+        images=images,
+        extras=extras,
+        table_name=table_name,
+        table_cls=table_cls,
+        index_key=index_key,
+        strict=strict,
+        mode=mode,
+    )
+    if not isinstance(table, table_cls):
+        raise ValueError(f"Table is not of type {table_cls}. Got {type(table)}")
+    return table
+
+
+def _tables_names_coalesce(
+    tables_names: list[list[str]],
+    mode: Literal["common", "all"] = "common",
+) -> list[str]:
+    num_images = len(tables_names)
+    if num_images == 0:
+        raise ValueError("No images to concatenate.")
+
+    names = [name for _table_names in tables_names for name in _table_names]
+    names_counts = Counter(names)
+
+    if mode == "common":
+        # Get the names that are present in all images
+        common_names = [
+            name for name, count in names_counts.items() if count == num_images
+        ]
+        return common_names
+    elif mode == "all":
+        # Get all names
+        return list(names_counts.keys())
+    else:
+        raise ValueError(f"Unknown mode: {mode}. Use 'common' or 'all'.")
+
+
+def list_image_tables(
+    images: Collection[OmeZarrContainer],
+    filter_types: str | None = None,
+    mode: Literal["common", "all"] = "common",
+) -> list[str]:
+    """List all table names in the images.
+
+    Args:
+        images: A collection of images.
+        filter_types (str | None): The type of tables to filter. If None,
+            return all tables. Defaults to None.
+        mode (Literal["common", "all"]): Whether to return only common tables
+            between all images or all tables. Defaults to "common".
+    """
+    tables_names = []
+    for image in images:
+        tables = image.list_tables(filter_types=filter_types)
+        tables_names.append(tables)
+
+    return _tables_names_coalesce(
+        tables_names=tables_names,
+        mode=mode,
+    )
+
+
+async def list_image_tables_async(
+    images: Collection[OmeZarrContainer],
+    filter_types: str | None = None,
+    mode: Literal["common", "all"] = "common",
+) -> list[str]:
+    """List all image tables in the image asynchronously.
+
+    Args:
+        images: A collection of images.
+        filter_types (str | None): The type of tables to filter. If None,
+            return all tables. Defaults to None.
+        mode (Literal["common", "all"]): Whether to return only common tables
+            between all images or all tables. Defaults to "common".
+    """
+    images_ids = []
+
+    # key table name, value list of paths
+    def process_image(
+        image: OmeZarrContainer, filter_types: str | None = None
+    ) -> list[str]:
+        tables = image.list_tables(filter_types=filter_types)
+        return tables
+
+    tasks = []
+    for i, image in enumerate(images):
+        images_ids.append(i)
+        task = asyncio.to_thread(process_image, image, filter_types=filter_types)
+        tasks.append(task)
+
+    tables_names = await asyncio.gather(*tasks)
+    return _tables_names_coalesce(
+        tables_names=tables_names,
+        mode=mode,
+    )

ngio/hcs/__init__.py

@@ -1,6 +1,6 @@
 """OME-Zarr HCS objects models."""
 
-from ngio.hcs.plate import (
+from ngio.hcs._plate import (
     OmeZarrPlate,
     OmeZarrWell,
     create_empty_plate,