Merge pull request #581 from fractal-analytics-platform/580-always-us…

…e-write_table-in-tasks-rather-than-directly-calling-anndata

Always use `write_table` wrapper in tasks

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 # Unreleased
 
+* Always use `write_table` in tasks, rather than AnnData `write_elem` (\#581).
 * Testing:
     * Cache Zenodo data, within GitHub actions (\#585).
 

fractal_tasks_core/tasks/apply_registration_to_ROI_tables.py

@@ -22,13 +22,13 @@
 import numpy as np
 import pandas as pd
 import zarr
-from anndata._io.specs import write_elem
 from pydantic.decorator import validate_arguments
 
 from fractal_tasks_core.lib_ngff import load_NgffWellMeta
 from fractal_tasks_core.lib_regions_of_interest import (
     are_ROI_table_columns_valid,
 )
+from fractal_tasks_core.lib_write import write_table
 
 logger = logging.getLogger(__name__)
 
@@ -152,15 +152,14 @@ def apply_registration_to_ROI_tables(
         logger.info(
             f"Write the registered ROI table {new_roi_table} for {acq=}"
         )
-        # TODO rewrite this block through `write_table` (ref PR #499)
-        # Save the shifted ROI tables as a new tables per acquisition
-        group_tables = zarr.group(f"{well_zarr}/{acq}/tables/")
-        write_elem(group_tables, new_roi_table, shifted_rois[acq])
-        # Update list of available tables & its metadata
-        current_tables = group_tables.attrs.asdict().get("tables") or []
-        new_tables = current_tables + [new_roi_table]
-        group_tables.attrs["tables"] = new_tables
-        group_tables[roi_table].attrs["type"] = "ngff:region_table"
+        # Save the shifted ROI table as a new table
+        image_group = zarr.group(f"{well_zarr}/{acq}")
+        write_table(
+            image_group,
+            new_roi_table,
+            shifted_rois[acq],
+            table_attrs=dict(type="ngff:region_table"),
+        )
 
     # TODO: Optionally apply registration to other tables as well?
     # e.g. to well_ROI_table based on FOV_ROI_table

fractal_tasks_core/tasks/apply_registration_to_image.py

@@ -24,7 +24,6 @@
 import dask.array as da
 import numpy as np
 import zarr
-from anndata._io.specs import write_elem
 from pydantic.decorator import validate_arguments
 
 from fractal_tasks_core.lib_ngff import load_NgffImageMeta
@@ -37,6 +36,7 @@
 )
 from fractal_tasks_core.lib_regions_of_interest import is_standard_roi_table
 from fractal_tasks_core.lib_regions_of_interest import load_region
+from fractal_tasks_core.lib_write import write_table
 from fractal_tasks_core.lib_zattrs_utils import get_table_path_dict
 
 logger = logging.getLogger(__name__)
@@ -203,25 +203,25 @@ def apply_registration_to_image(
 
     if table_dict:
         logger.info(f"Processing the tables: {table_dict}")
-        new_tables_group = zarr.group(f"{input_path / new_component}/tables")
-        new_tables_group.attrs["tables"] = list(table_dict.keys())
+        new_image_group = zarr.group(f"{input_path / new_component}")
 
         for table in table_dict.keys():
             logger.info(f"Copying table: {table}")
-            # Write the Zarr table
-            curr_table = ad.read_zarr(table_dict[table])
-            write_elem(new_tables_group, table, curr_table)
             # Get the relevant metadata of the Zarr table & add it
             # See issue #516 for the need for this workaround
             try:
                 old_table_group = zarr.open_group(table_dict[table], mode="r")
             except zarr.errors.GroupNotFoundError:
                 time.sleep(5)
                 old_table_group = zarr.open_group(table_dict[table], mode="r")
-            new_table_group = zarr.open_group(
-                f"{input_path / new_component}/tables/{table}"
+            # Write the Zarr table
+            curr_table = ad.read_zarr(table_dict[table])
+            write_table(
+                new_image_group,
+                table,
+                curr_table,
+                table_attrs=old_table_group.attrs.asdict(),
             )
-            new_table_group.attrs.put(old_table_group.attrs.asdict())
 
     ####################
     # Clean up Zarr file

fractal_tasks_core/tasks/calculate_registration_image_based.py

@@ -22,7 +22,6 @@
 import numpy as np
 import pandas as pd
 import zarr
-from anndata._io.specs import write_elem
 from pydantic.decorator import validate_arguments
 from skimage.registration import phase_cross_correlation
 
@@ -37,6 +36,7 @@
     convert_ROI_table_to_indices,
 )
 from fractal_tasks_core.lib_regions_of_interest import load_region
+from fractal_tasks_core.lib_write import write_table
 
 logger = logging.getLogger(__name__)
 
@@ -252,10 +252,15 @@ def calculate_registration_image_based(
 
     # Write physical shifts to disk (as part of the ROI table)
     logger.info(f"Updating the {roi_table=} with translation columns")
+    image_group = zarr.group(zarr_img_cycle_x)
     new_ROI_table = get_ROI_table_with_translation(ROI_table_x, new_shifts)
-    group_tables = zarr.group(f"{zarr_img_cycle_x}/tables/")
-    write_elem(group_tables, roi_table, new_ROI_table)
-    group_tables[roi_table].attrs["type"] = "ngff:region_table"
+    write_table(
+        image_group,
+        roi_table,
+        new_ROI_table,
+        overwrite=True,
+        table_attrs=dict(type="ngff:region_table"),
+    )
 
     return {}
 

fractal_tasks_core/tasks/copy_ome_zarr.py

@@ -19,7 +19,6 @@
 
 import anndata as ad
 import zarr
-from anndata._io.specs import write_elem
 from pydantic.decorator import validate_arguments
 
 import fractal_tasks_core
@@ -28,6 +27,7 @@
     convert_ROIs_from_3D_to_2D,
 )
 from fractal_tasks_core.lib_write import open_zarr_group_with_overwrite
+from fractal_tasks_core.lib_write import write_table
 
 logger = logging.getLogger(__name__)
 
@@ -170,8 +170,6 @@ def copy_ome_zarr(
                 # Extract pixel sizes, if needed
                 if ROI_table_names:
 
-                    new_tables_group = new_image_group.create_group("tables/")
-                    new_tables_group.attrs["tables"] = ROI_table_names
                     if project_to_2D:
                         path_image = f"{zarrurl_old}/{well_path}/{image_path}"
                         ngff_image_meta = load_NgffImageMeta(path_image)
@@ -188,18 +186,18 @@ def copy_ome_zarr(
                             f"{zarrurl_old=}, convert it to 2D, and "
                             "write it back to the new zarr file."
                         )
-                        ROI_table = ad.read_zarr(
+                        new_ROI_table = ad.read_zarr(
                             f"{zarrurl_old}/{well_path}/{image_path}/"
                             f"tables/{ROI_table_name}"
                         )
                         # Convert 3D ROIs to 2D
                         if project_to_2D:
                             new_ROI_table = convert_ROIs_from_3D_to_2D(
-                                ROI_table, pxl_size_z
+                                new_ROI_table, pxl_size_z
                             )
                         # Write new table
-                        write_elem(
-                            new_tables_group, ROI_table_name, new_ROI_table
+                        write_table(
+                            new_image_group, ROI_table_name, new_ROI_table
                         )
 
     for key in ["plate", "well", "image"]: