ngio 0.2.8__py3-none-any.whl → 0.3.0a0__py3-none-any.whl

ngio/hcs/plate.py

@@ -1,7 +1,18 @@
 """A module for handling the Plate Collection in an OME-Zarr file."""
 
-from typing import Literal, overload
-
+import asyncio
+import warnings
+from collections.abc import Collection
+from typing import Literal
+
+from ngio.common import (
+    concatenate_image_tables,
+    concatenate_image_tables_as,
+    concatenate_image_tables_as_async,
+    concatenate_image_tables_async,
+    list_image_tables,
+    list_image_tables_async,
+)
 from ngio.images import OmeZarrContainer
 from ngio.ome_zarr_meta import (
     ImageInWellPath,
@@ -15,12 +26,15 @@ from ngio.ome_zarr_meta import (
     path_in_well_validation,
 )
 from ngio.tables import (
+    ConditionTable,
     FeatureTable,
     GenericRoiTable,
     MaskingRoiTable,
     RoiTable,
     Table,
+    TableBackendProtocol,
     TablesContainer,
+    TableType,
     TypedTable,
 )
 from ngio.utils import (
@@ -186,6 +200,28 @@ class OmeZarrWell:
         )
 
 
+def _buil_extras(paths: Collection[str]) -> list[dict[str, str]]:
+    """Build the extras for the images.
+
+    Args:
+        paths (Collection[str]): The paths of the images.
+
+    Returns:
+        list[dict[str, str]]: The extras for the images.
+    """
+    extras = []
+    for path in paths:
+        row, column, path_in_well = path.split("/")
+        extras.append(
+            {
+                "row": row,
+                "column": column,
+                "path_in_well": path_in_well,
+            }
+        )
+    return extras
+
+
 class OmeZarrPlate:
     """A class to handle the Plate Collection in an OME-Zarr file."""
 
@@ -253,6 +289,22 @@ class OmeZarrPlate:
         """Return the wells paths in the plate."""
         return self.meta.wells_paths
 
+    async def images_paths_async(self, acquisition: int | None = None) -> list[str]:
+        """Return the images paths in the plate asynchronously.
+
+        If acquisition is None, return all images paths in the plate.
+        Else, return the images paths in the plate for the given acquisition.
+
+        Args:
+            acquisition (int | None): The acquisition id to filter the images.
+        """
+        wells = await self.get_wells_async()
+        paths = []
+        for well_path, well in wells.items():
+            for img_path in well.paths(acquisition):
+                paths.append(f"{well_path}/{img_path}")
+        return paths
+
     def images_paths(self, acquisition: int | None = None) -> list[str]:
         """Return the images paths in the plate.
 
@@ -262,9 +314,10 @@ class OmeZarrPlate:
         Args:
             acquisition (int | None): The acquisition id to filter the images.
         """
+        wells = self.get_wells()
         images = []
-        for well_path, wells in self.get_wells().items():
-            for img_path in wells.paths(acquisition):
+        for well_path, well in wells.items():
+            for img_path in well.paths(acquisition):
                 images.append(f"{well_path}/{img_path}")
         return images
 
@@ -317,6 +370,37 @@ class OmeZarrPlate:
         group_handler = self._group_handler.derive_handler(well_path)
         return OmeZarrWell(group_handler)
 
+    async def get_wells_async(self) -> dict[str, OmeZarrWell]:
+        """Get all wells in the plate asynchronously.
+
+        This method processes wells in parallel for improved performance
+        when working with a large number of wells.
+
+        Returns:
+            dict[str, OmeZarrWell]: A dictionary of wells, where the key is the well
+                path and the value is the well object.
+        """
+        wells = self._group_handler.get_from_cache("wells")
+        if wells is not None:
+            return wells  # type: ignore[return-value]
+
+        def process_well(well_path):
+            group_handler = self._group_handler.derive_handler(well_path)
+            well = OmeZarrWell(group_handler)
+            return well_path, well
+
+        wells, tasks = {}, []
+        for well_path in self.wells_paths():
+            task = asyncio.to_thread(process_well, well_path)
+            tasks.append(task)
+
+        results = await asyncio.gather(*tasks)
+        for well_path, well in results:
+            wells[well_path] = well
+
+        self._group_handler.add_to_cache("wells", wells)
+        return wells
+
     def get_wells(self) -> dict[str, OmeZarrWell]:
         """Get all wells in the plate.
 
@@ -324,23 +408,86 @@ class OmeZarrPlate:
             dict[str, OmeZarrWell]: A dictionary of wells, where the key is the well
                 path and the value is the well object.
         """
-        wells = {}
-        for well_path in self.wells_paths():
+        wells = self._group_handler.get_from_cache("wells")
+        if wells is not None:
+            return wells  # type: ignore[return-value]
+
+        def process_well(well_path):
             group_handler = self._group_handler.derive_handler(well_path)
             well = OmeZarrWell(group_handler)
+            return well_path, well
+
+        wells = {}
+        for well_path in self.wells_paths():
+            _, well = process_well(well_path)
             wells[well_path] = well
+
+        self._group_handler.add_to_cache("wells", wells)
         return wells
 
+    async def get_images_async(
+        self, acquisition: int | None = None
+    ) -> dict[str, OmeZarrContainer]:
+        """Get all images in the plate asynchronously.
+
+        This method processes images in parallel for improved performance
+        when working with a large number of images.
+
+        Args:
+            acquisition: The acquisition id to filter the images.
+
+        Returns:
+            dict[str, OmeZarrContainer]: A dictionary of images, where the key is the
+                image path and the value is the image object.
+        """
+        images = self._group_handler.get_from_cache("images")
+        if images is not None:
+            return images  # type: ignore[return-value]
+
+        paths = await self.images_paths_async(acquisition=acquisition)
+
+        def process_image(image_path):
+            """Process a single image and return the image path and image object."""
+            img_group_handler = self._group_handler.derive_handler(image_path)
+            image = OmeZarrContainer(img_group_handler)
+            return image_path, image
+
+        images, tasks = {}, []
+        for image_path in paths:
+            task = asyncio.to_thread(process_image, image_path)
+            tasks.append(task)
+
+        results = await asyncio.gather(*tasks)
+
+        for image_path, image in results:
+            images[image_path] = image
+
+        self._group_handler.add_to_cache("images", images)
+        return images
+
     def get_images(self, acquisition: int | None = None) -> dict[str, OmeZarrContainer]:
         """Get all images in the plate.
 
         Args:
             acquisition: The acquisition id to filter the images.
         """
-        images = {}
-        for image_path in self.images_paths(acquisition):
+        images = self._group_handler.get_from_cache("images")
+        if images is not None:
+            return images  # type: ignore[return-value]
+        paths = self.images_paths(acquisition=acquisition)
+
+        def process_image(image_path):
+            """Process a single image and return the image path and image object."""
             img_group_handler = self._group_handler.derive_handler(image_path)
-            images[image_path] = OmeZarrContainer(img_group_handler)
+            image = OmeZarrContainer(img_group_handler)
+            return image_path, image
+
+        images = {}
+        for image_path in paths:
+            _, image = process_image(image_path)
+            images[image_path] = image
+
+        self._group_handler.add_to_cache("images", images)
         return images
 
     def get_image(
@@ -667,91 +814,122 @@ class OmeZarrPlate:
                 raise NgioValidationError("No tables found in the image.")
         return self._tables_container
 
-    @property
-    def list_tables(self) -> list[str]:
+    def list_tables(self, filter_types: str | None = None) -> list[str]:
         """List all tables in the image."""
-        return self.tables_container.list()
+        return self.tables_container.list(filter_types=filter_types)
 
     def list_roi_tables(self) -> list[str]:
         """List all ROI tables in the image."""
         return self.tables_container.list_roi_tables()
 
-    @overload
-    def get_table(self, name: str) -> Table: ...
+    def get_roi_table(self, name: str) -> RoiTable:
+        """Get a ROI table from the image.
+
+        Args:
+            name (str): The name of the table.
+        """
+        table = self.tables_container.get(name=name, strict=True)
+        if not isinstance(table, RoiTable):
+            raise NgioValueError(f"Table {name} is not a ROI table. Got {type(table)}")
+        return table
 
-    @overload
-    def get_table(self, name: str, check_type: None) -> Table: ...
+    def get_masking_roi_table(self, name: str) -> MaskingRoiTable:
+        """Get a masking ROI table from the image.
 
-    @overload
-    def get_table(self, name: str, check_type: Literal["roi_table"]) -> RoiTable: ...
+        Args:
+            name (str): The name of the table.
+        """
+        table = self.tables_container.get(name=name, strict=True)
+        if not isinstance(table, MaskingRoiTable):
+            raise NgioValueError(
+                f"Table {name} is not a masking ROI table. Got {type(table)}"
+            )
+        return table
 
-    @overload
-    def get_table(
-        self, name: str, check_type: Literal["masking_roi_table"]
-    ) -> MaskingRoiTable: ...
+    def get_feature_table(self, name: str) -> FeatureTable:
+        """Get a feature table from the image.
 
-    @overload
-    def get_table(
-        self, name: str, check_type: Literal["feature_table"]
-    ) -> FeatureTable: ...
+        Args:
+            name (str): The name of the table.
+        """
+        table = self.tables_container.get(name=name, strict=True)
+        if not isinstance(table, FeatureTable):
+            raise NgioValueError(
+                f"Table {name} is not a feature table. Got {type(table)}"
+            )
+        return table
 
-    @overload
-    def get_table(
-        self, name: str, check_type: Literal["generic_roi_table"]
-    ) -> GenericRoiTable: ...
+    def get_generic_roi_table(self, name: str) -> GenericRoiTable:
+        """Get a generic ROI table from the image.
+
+        Args:
+            name (str): The name of the table.
+        """
+        table = self.tables_container.get(name=name, strict=True)
+        if not isinstance(table, GenericRoiTable):
+            raise NgioValueError(
+                f"Table {name} is not a generic ROI table. Got {type(table)}"
+            )
+        return table
+
+    def get_condition_table(self, name: str) -> ConditionTable:
+        """Get a condition table from the image.
+
+        Args:
+            name (str): The name of the table.
+        """
+        table = self.tables_container.get(name=name, strict=True)
+        if not isinstance(table, ConditionTable):
+            raise NgioValueError(
+                f"Table {name} is not a condition table. Got {type(table)}"
+            )
+        return table
 
     def get_table(self, name: str, check_type: TypedTable | None = None) -> Table:
         """Get a table from the image.
 
         Args:
             name (str): The name of the table.
-            check_type (TypedTable | None): The type of the table. If None, the
-                type is not checked. If a type is provided, the table must be of that
-                type.
-        """
-        if check_type is None:
-            table = self.tables_container.get(name, strict=False)
-            return table
-
-        table = self.tables_container.get(name, strict=True)
-        match check_type:
-            case "roi_table":
-                if not isinstance(table, RoiTable):
-                    raise NgioValueError(
-                        f"Table '{name}' is not a ROI table. Found type: {table.type()}"
-                    )
-                return table
-            case "masking_roi_table":
-                if not isinstance(table, MaskingRoiTable):
-                    raise NgioValueError(
-                        f"Table '{name}' is not a masking ROI table. "
-                        f"Found type: {table.type()}"
-                    )
-                return table
-
-            case "generic_roi_table":
-                if not isinstance(table, GenericRoiTable):
-                    raise NgioValueError(
-                        f"Table '{name}' is not a generic ROI table. "
-                        f"Found type: {table.type()}"
-                    )
-                return table
-
-            case "feature_table":
-                if not isinstance(table, FeatureTable):
-                    raise NgioValueError(
-                        f"Table '{name}' is not a feature table. "
-                        f"Found type: {table.type()}"
-                    )
-                return table
-            case _:
-                raise NgioValueError(f"Unknown check_type: {check_type}")
+            check_type (TypedTable | None): Deprecated. Please use
+                'get_table_as' instead, or one of the type specific
+                get_*table() methods.
+
+        """
+        if check_type is not None:
+            warnings.warn(
+                "The 'check_type' argument is deprecated, and will be removed in "
+                "ngio=0.3. Use 'get_table_as' instead or one of the "
+                "type specific get_*table() methods.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        return self.tables_container.get(name=name, strict=False)
+
+    def get_table_as(
+        self,
+        name: str,
+        table_cls: type[TableType],
+        backend: str | TableBackendProtocol | None = None,
+    ) -> TableType:
+        """Get a table from the image as a specific type.
+
+        Args:
+            name (str): The name of the table.
+            table_cls (type[TableType]): The type of the table.
+            backend (str | TableBackendProtocol | None): The backend to use. If None,
+                the default backend is used.
+        """
+        return self.tables_container.get_as(
+            name=name,
+            table_cls=table_cls,
+            backend=backend,
+        )
 
     def add_table(
         self,
         name: str,
         table: Table,
-        backend: str | None = None,
+        backend: str | TableBackendProtocol = "anndata_v1",
         overwrite: bool = False,
     ) -> None:
         """Add a table to the image."""
@@ -759,6 +937,186 @@ class OmeZarrPlate:
             name=name, table=table, backend=backend, overwrite=overwrite
         )
 
+    def list_image_tables(
+        self,
+        acquisition: int | None = None,
+        filter_types: str | None = None,
+        mode: Literal["common", "all"] = "common",
+    ) -> list[str]:
+        """List all image tables in the image.
+
+        Args:
+            acquisition (int | None): The acquisition id to filter the images.
+            filter_types (str | None): The type of tables to filter. If None,
+                return all tables. Defaults to None.
+            mode (Literal["common", "all"]): The mode to use for listing the tables.
+                If 'common', return only common tables between all images.
+                If 'all', return all tables. Defaults to 'common'.
+        """
+        images = self.get_images(acquisition=acquisition)
+        return list_image_tables(
+            images=images.values(),
+            filter_types=filter_types,
+            mode=mode,
+        )
+
+    async def list_image_tables_async(
+        self,
+        acquisition: int | None = None,
+        filter_types: str | None = None,
+        mode: Literal["common", "all"] = "common",
+    ) -> list[str]:
+        """List all image tables in the image asynchronously.
+
+        Args:
+            acquisition (int | None): The acquisition id to filter the images.
+            filter_types (str | None): The type of tables to filter. If None,
+                return all tables. Defaults to None.
+            mode (Literal["common", "all"]): The mode to use for listing the tables.
+                If 'common', return only common tables between all images.
+                If 'all', return all tables. Defaults to 'common'.
+        """
+        images = await self.get_images_async(acquisition=acquisition)
+        return await list_image_tables_async(
+            images=images.values(),
+            filter_types=filter_types,
+            mode=mode,
+        )
+
+    def concatenate_image_tables(
+        self,
+        table_name: str,
+        acquisition: int | None = None,
+        strict: bool = True,
+        index_key: str | None = None,
+        mode: Literal["eager", "lazy"] = "eager",
+    ) -> Table:
+        """Concatenate tables from all images in the plate.
+
+        Args:
+            table_name: The name of the table to concatenate.
+            index_key: The key to use for the index of the concatenated table.
+            acquisition: The acquisition id to filter the images.
+            strict: If True, raise an error if the table is not found in the image.
+            index_key: If a string is provided, a new index column will be created
+                new_index_pattern = {row}_{column}_{path_in_well}_{label}
+            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.
+        """
+        images = self.get_images(acquisition=acquisition)
+        extras = _buil_extras(images.keys())
+        return concatenate_image_tables(
+            images=images.values(),
+            extras=extras,
+            table_name=table_name,
+            index_key=index_key,
+            strict=strict,
+            mode=mode,
+        )
+
+    def concatenate_image_tables_as(
+        self,
+        table_name: str,
+        table_cls: type[TableType],
+        acquisition: int | None = None,
+        index_key: str | None = None,
+        strict: bool = True,
+        mode: Literal["eager", "lazy"] = "eager",
+    ) -> TableType:
+        """Concatenate tables from all images in the plate as a specific type.
+
+        Args:
+            table_name: The name of the table to concatenate.
+            table_cls: The type of the table to concatenate.
+            index_key: The key to use for the index of the concatenated table.
+            acquisition: The acquisition id to filter the images.
+            index_key: If a string is provided, a new index column will be created
+                new_index_pattern = {row}_{column}_{path_in_well}_{label}
+            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.
+        """
+        images = self.get_images(acquisition=acquisition)
+        extras = _buil_extras(images.keys())
+        return concatenate_image_tables_as(
+            images=images.values(),
+            extras=extras,
+            table_name=table_name,
+            table_cls=table_cls,
+            index_key=index_key,
+            strict=strict,
+            mode=mode,
+        )
+
+    async def concatenate_image_tables_async(
+        self,
+        table_name: str,
+        acquisition: int | None = None,
+        index_key: str | None = None,
+        strict: bool = True,
+        mode: Literal["eager", "lazy"] = "eager",
+    ) -> Table:
+        """Concatenate tables from all images in the plate asynchronously.
+
+        Args:
+            table_name: The name of the table to concatenate.
+            index_key: The key to use for the index of the concatenated table.
+            acquisition: The acquisition id to filter the images.
+            index_key: If a string is provided, a new index column will be created
+                new_index_pattern = {row}_{column}_{path_in_well}_{label}
+            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.
+        """
+        images = await self.get_images_async(acquisition=acquisition)
+        extras = _buil_extras(images.keys())
+        return await concatenate_image_tables_async(
+            images=images.values(),
+            extras=extras,
+            table_name=table_name,
+            index_key=index_key,
+            strict=strict,
+            mode=mode,
+        )
+
+    async def concatenate_image_tables_as_async(
+        self,
+        table_name: str,
+        table_cls: type[TableType],
+        acquisition: int | None = None,
+        index_key: str | None = None,
+        strict: bool = True,
+        mode: Literal["eager", "lazy"] = "eager",
+    ) -> TableType:
+        """Concatenate tables from all images in the plate as a specific type.
+
+        Args:
+            table_name: The name of the table to concatenate.
+            table_cls: The type of the table to concatenate.
+            index_key: The key to use for the index of the concatenated table.
+            acquisition: The acquisition id to filter the images.
+            index_key: If a string is provided, a new index column will be created
+                new_index_pattern = {row}_{column}_{path_in_well}_{label}
+            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.
+        """
+        images = await self.get_images_async(acquisition=acquisition)
+        extras = _buil_extras(images.keys())
+        return await concatenate_image_tables_as_async(
+            images=images.values(),
+            extras=extras,
+            table_name=table_name,
+            table_cls=table_cls,
+            index_key=index_key,
+            strict=strict,
+            mode=mode,
+        )
+
 
 def open_ome_zarr_plate(
     store: StoreOrGroup,