From 78839a2e8789cf70b3a05f65d40a49a4816dba3f Mon Sep 17 00:00:00 2001
From: jluethi <joel.luethi@uzh.ch>
Date: Wed, 24 Apr 2024 17:08:39 +0200
Subject: [PATCH] Add  function & use it in illumination_correction with
 overwrite_input=False

---
 fractal_tasks_core/tasks/_zarr_utils.py       | 48 +++++++++++++++++++
 .../tasks/illumination_correction.py          |  4 ++
 .../test_unit_illumination_correction.py      | 11 +++++
 3 files changed, 63 insertions(+)

diff --git a/fractal_tasks_core/tasks/_zarr_utils.py b/fractal_tasks_core/tasks/_zarr_utils.py
index 484058859..6d9abb850 100644
--- a/fractal_tasks_core/tasks/_zarr_utils.py
+++ b/fractal_tasks_core/tasks/_zarr_utils.py
@@ -1,10 +1,13 @@
 import copy
 import logging
 
+import anndata as ad
 import zarr
 from filelock import FileLock
 
 from fractal_tasks_core.ngff.zarr_utils import load_NgffWellMeta
+from fractal_tasks_core.tables import write_table
+from fractal_tasks_core.tables.v1 import get_tables_list_v1
 
 logger = logging.getLogger(__name__)
 
@@ -165,3 +168,48 @@ def _get_matching_ref_acquisition_path_heuristic(
         f"option {sorted_path_list[0]}."
     )
     return sorted_path_list[0]
+
+
+def _copy_tables_from_zarr_url(
+    origin_zarr_url: str,
+    target_zarr_url: str,
+    table_type: str = None,
+    overwrite: bool = True,
+) -> None:
+    """
+    Copies all ROI tables from one Zarr into a new Zarr
+
+    Args:
+        origin_zarr_url: url of the OME-Zarr image that contains tables.
+            e.g. /path/to/my_plate.zarr/B/03/0
+        target_zarr_url: url of the new OME-Zarr image where tables are copied
+            to. e.g. /path/to/my_plate.zarr/B/03/0_illum_corr
+        table_type: Filter for specific table types that should be copied.
+        overwrite: Whether existing tables of the same name in the
+            target_zarr_url should be overwritten.
+    """
+    table_list = get_tables_list_v1(
+        zarr_url=origin_zarr_url, table_type=table_type
+    )
+
+    if table_list:
+        logger.info(
+            f"Copying the tables {table_list} from {origin_zarr_url} to "
+            f"{target_zarr_url}."
+        )
+        new_image_group = zarr.group(target_zarr_url)
+
+        for table in table_list:
+            logger.info(f"Copying table: {table}")
+            # Get the relevant metadata of the Zarr table & add it
+            table_url = f"{origin_zarr_url}/tables/{table}"
+            old_table_group = zarr.open_group(table_url, mode="r")
+            # Write the Zarr table
+            curr_table = ad.read_zarr(table_url)
+            write_table(
+                new_image_group,
+                table,
+                curr_table,
+                table_attrs=old_table_group.attrs.asdict(),
+                overwrite=overwrite,
+            )
diff --git a/fractal_tasks_core/tasks/illumination_correction.py b/fractal_tasks_core/tasks/illumination_correction.py
index 37a2b9fef..06e056181 100644
--- a/fractal_tasks_core/tasks/illumination_correction.py
+++ b/fractal_tasks_core/tasks/illumination_correction.py
@@ -34,6 +34,7 @@
     convert_ROI_table_to_indices,
 )
 from fractal_tasks_core.tasks._zarr_utils import _copy_hcs_ome_zarr_metadata
+from fractal_tasks_core.tasks._zarr_utils import _copy_tables_from_zarr_url
 
 logger = logging.getLogger(__name__)
 
@@ -228,6 +229,9 @@ def illumination_correction(
             dimension_separator="/",
         )
         _copy_hcs_ome_zarr_metadata(zarr_url, zarr_url_new)
+        # Copy ROI tables from the old zarr_url to keep ROI tables and other
+        # tables available in the new Zarr
+        _copy_tables_from_zarr_url(zarr_url, zarr_url_new)
 
     # Iterate over FOV ROIs
     num_ROIs = len(list_indices)
diff --git a/tests/tasks/test_unit_illumination_correction.py b/tests/tasks/test_unit_illumination_correction.py
index f739c18b4..6fcd4f8a4 100644
--- a/tests/tasks/test_unit_illumination_correction.py
+++ b/tests/tasks/test_unit_illumination_correction.py
@@ -7,6 +7,7 @@
 import dask.array as da
 import numpy as np
 import pytest
+from devtools import debug
 from pytest import LogCaptureFixture
 from pytest import MonkeyPatch
 
@@ -15,6 +16,7 @@
 from fractal_tasks_core.roi import (
     convert_ROI_table_to_indices,
 )
+from fractal_tasks_core.tables.v1 import get_tables_list_v1
 from fractal_tasks_core.tasks._registration_utils import (
     _split_well_path_image_path,
 )
@@ -68,6 +70,9 @@ def test_illumination_correction(
     )
     num_FOVs = len(list_indices)
 
+    # Get existing tables before illumination correction
+    tables = get_tables_list_v1(zarr_url)
+
     # Prepared expected number of calls
     expected_tot_calls_correct = num_channels * num_FOVs
 
@@ -136,3 +141,9 @@ def patched_correct(*args, **kwargs):
         assert well_paths == ["0"]
     else:
         assert well_paths == ["0", "0" + suffix]
+
+    # Assert that the image has the same tables after illumination correction
+    # as before
+    debug(tables)
+    new_tables = get_tables_list_v1(new_zarr_url)
+    assert tables == new_tables