Small fixes (#92)

* Move create_zarr_plate out of converter.run.
* Check if tiles exist in lookup table.
    Note: It can happen that certain positions have no data e.g. auto-focus z-shift in CV acquisitions.
* Call alignment on unaligned tiles.
* Prepare acquisitions to be flexible w.r.t. missing channel axis.
* Lowest z position in ImageXpress acquisitions is 1.
* Update coordinate transformation tests.
* Add test for get axes with only a single channel.
* Update converter tests to create the plate before calling run.
* Add a test with empty chunks.

examples/CellVoyagerStackAcquisitionExample.py

@@ -13,7 +13,7 @@ def main():
     shutil.rmtree("cv-stack.zarr", ignore_errors=True)
 
     # Parse CV plate acquisition.
-    plate = StackAcquisition(
+    plate_acquisition = StackAcquisition(
         acquisition_dir=Path(__file__).parent.parent
         / "resources"
         / "CV8000"
@@ -37,9 +37,12 @@ def main():
         fuse_func=stitching_utils.fuse_mean,
     )
 
+    plate = converter.create_zarr_plate(plate_acquisition)
+
     # Run conversion.
     converter.run(
-        plate_acquisition=plate,
+        plate=plate,
+        plate_acquisition=plate_acquisition,
         well_sub_group="0",
         chunks=(2, 1000, 1000),
         max_layer=2,

examples/ImageXpressSinglePlaneAcquisitionExample.py

@@ -13,7 +13,7 @@ def main():
     shutil.rmtree("md-single-plane.zarr", ignore_errors=True)
 
     # Parse MD plate acquisition.
-    plate = SinglePlaneAcquisition(
+    plate_acquisition = SinglePlaneAcquisition(
         acquisition_dir=Path(__file__).parent.parent / "resources" / "Projection-Mix",
         alignment=TileAlignmentOptions.GRID,
     )
@@ -33,9 +33,12 @@ def main():
         fuse_func=stitching_utils.fuse_mean,
     )
 
+    plate = converter.create_zarr_plate(plate_acquisition)
+
     # Run conversion.
     converter.run(
-        plate_acquisition=plate,
+        plate=plate,
+        plate_acquisition=plate_acquisition,
         well_sub_group="0",
         chunks=(1, 512, 512),
         max_layer=2,

examples/ImageXpressStackAcquisitionExample.py

@@ -12,8 +12,16 @@ def main():
     # Remove existing zarr.
     shutil.rmtree("md-stack.zarr", ignore_errors=True)
 
+    import distributed
+
+    client = distributed.Client(
+        n_workers=1,
+        threads_per_worker=1,
+        processes=False,
+    )
+
     # Parse MD plate acquisition.
-    plate = StackAcquisition(
+    plate_acquisition = StackAcquisition(
         acquisition_dir=Path(__file__).parent.parent / "resources" / "Projection-Mix",
         alignment=TileAlignmentOptions.GRID,
     )
@@ -31,11 +39,15 @@ def main():
         stitching_yx_chunk_size_factor=2,
         warp_func=stitching_utils.translate_tiles_2d,
         fuse_func=stitching_utils.fuse_mean,
+        client=client,
     )
 
+    plate = converter.create_zarr_plate(plate_acquisition)
+
     # Run conversion.
     converter.run(
-        plate_acquisition=plate,
+        plate=plate,
+        plate_acquisition=plate_acquisition,
         well_sub_group="0",
         chunks=(1, 512, 512),
         max_layer=2,

src/faim_hcs/alignment/alignment.py

@@ -13,7 +13,7 @@ class AbstractAlignment(ABC):
     def __init__(self, tiles: list[Tile]) -> None:
         super().__init__()
         self._unaligned_tiles = stitching_utils.shift_to_origin(tiles)
-        self._aligned_tiles = self._align(tiles)
+        self._aligned_tiles = self._align(self._unaligned_tiles)
 
     def _align(self, tiles: list[Tile]) -> list[Tile]:
         raise NotImplementedError()

src/faim_hcs/hcs/acquisition.py

@@ -245,7 +245,10 @@ def get_z_spacing(self) -> Optional[float]:
         raise NotImplementedError()
 
     def get_coordinate_transformations(
-        self, max_layer: int, yx_binning: int
+        self,
+        max_layer: int,
+        yx_binning: int,
+        ndim: int,
     ) -> list[dict[str, Any]]:
         """
         Get the NGFF conform coordinate transformations for the well
@@ -255,6 +258,7 @@ def get_coordinate_transformations(
         ----------
         max_layer : Maximum layer of the resolution pyramid.
         yx_binning : Bin factor of the yx resolution.
+        ndim : Number of dimensions of the data.
 
         Returns
         -------
@@ -268,6 +272,9 @@ def get_coordinate_transformations(
                         {
                             "scale": [
                                 1.0,
+                            ]
+                            * (ndim - 3)
+                            + [
                                 self.get_z_spacing(),
                                 self.get_yx_spacing()[0] * yx_binning * 2**s,
                                 self.get_yx_spacing()[1] * yx_binning * 2**s,
@@ -282,6 +289,9 @@ def get_coordinate_transformations(
                         {
                             "scale": [
                                 1.0,
+                            ]
+                            * (ndim - 2)
+                            + [
                                 self.get_yx_spacing()[0] * yx_binning * 2**s,
                                 self.get_yx_spacing()[1] * yx_binning * 2**s,
                             ],

src/faim_hcs/hcs/converter.py

@@ -77,9 +77,21 @@ def __init__(
             self._cluster_factory = None
             self._client = client
 
-    def _create_zarr_plate(
+    def create_zarr_plate(
         self, plate_acquisition: PlateAcquisition, wells: Optional[list[str]] = None
     ) -> zarr.Group:
+        """
+        Create empty NGFF zarr plate.
+
+        Note: Loads the plate from disk if it already exists.
+
+        Parameters
+        ----------
+        plate_acquisition :
+            A single plate acquisition.
+        wells :
+            List of wells to build. If None, all wells are built.
+        """
         plate_path = join(self._ngff_plate.root_dir, self._ngff_plate.name + ".zarr")
         if not os.path.exists(plate_path):
             os.makedirs(plate_path, exist_ok=False)
@@ -110,6 +122,7 @@ def _create_zarr_plate(
 
     def run(
         self,
+        plate: zarr.Group,
         plate_acquisition: PlateAcquisition,
         wells: list[str] = None,
         well_sub_group: str = "0",
@@ -138,7 +151,6 @@ def run(
             zarr.Group of the plate.
         """
         assert 2 <= len(chunks) <= 3, "Chunks must be 2D or 3D."
-        plate = self._create_zarr_plate(plate_acquisition, wells=wells)
         well_acquisitions = plate_acquisition.get_well_acquisitions(wells)
 
         for well_acquisition in well_acquisitions:
@@ -173,6 +185,7 @@ def _write_metadata(
         coordinate_transformations = well_acquisition.get_coordinate_transformations(
             max_layer=max_layer,
             yx_binning=self._yx_binning,
+            ndim=len(shapes[0]),
         )
         fmt = CurrentFormat()
         dims = len(shapes[0])

src/faim_hcs/hcs/imagexpress/ImageXpressWellAcquisition.py

@@ -34,9 +34,9 @@ def _assemble_tiles(self) -> list[Tile]:
             channel = row["channel"]
             metadata = load_metaseries_tiff_metadata(file)
             if self._z_spacing is None:
-                z = 0
+                z = 1
             else:
-                z = int(metadata["stage-position-z"] / self._z_spacing)
+                z = row["z"] if row["z"] is not None else 1
 
             bgcm = None
             if self._background_correction_matrices is not None:
@@ -78,6 +78,11 @@ def get_z_spacing(self) -> Optional[float]:
 
     def get_axes(self) -> list[str]:
         if "z" in self._files.columns:
-            return ["c", "z", "y", "x"]
+            axes = ["z", "y", "x"]
         else:
-            return ["c", "y", "x"]
+            axes = ["y", "x"]
+
+        if self._files["channel"].nunique() > 1:
+            axes = ["c"] + axes
+
+        return axes

src/faim_hcs/stitching/DaskTileStitcher.py

@@ -81,20 +81,20 @@ def _compute_block_to_tile_map(self):
                 position=tuple(block_position * np.array(self.chunk_shape)),
                 shape=self.chunk_shape,
             )
-            for tile in tiles_lut[
-                (block_bbox.time_start, block_bbox.channel_start, block_bbox.z_start)
-            ]:
-                tile_bbox = BoundingBox5D.from_pos_and_shape(
-                    position=tile.get_position(),
-                    shape=(
-                        1,
-                        1,
-                        1,
+            pos = (block_bbox.time_start, block_bbox.channel_start, block_bbox.z_start)
+            if pos in tiles_lut.keys():
+                for tile in tiles_lut[pos]:
+                    tile_bbox = BoundingBox5D.from_pos_and_shape(
+                        position=tile.get_position(),
+                        shape=(
+                            1,
+                            1,
+                            1,
+                        )
+                        + tile.shape,
                     )
-                    + tile.shape,
-                )
-                if block_bbox.overlaps(tile_bbox):
-                    block_to_tile_map[block_position].append(tile)
+                    if block_bbox.overlaps(tile_bbox):
+                        block_to_tile_map[block_position].append(tile)
 
         return block_to_tile_map
 

tests/hcs/imagexpress/test_ImageXpressWellAcquisition.py

@@ -33,18 +33,7 @@ def test__assemble_tiles(files):
         assert tile.shape == (512, 512)
         assert tile.position.channel in [1, 2, 4]
         assert tile.position.time == 0
-        assert tile.position.z in [
-            3106,
-            3107,
-            3109,
-            3111,
-            3112,
-            3114,
-            3116,
-            3117,
-            3119,
-            3121,
-        ]
+        assert tile.position.z in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
 
 
 def test_get_axes(files):
@@ -66,6 +55,15 @@ def test_get_axes(files):
     axes = ix_well_acquisition.get_axes()
     assert axes == ["c", "y", "x"]
 
+    ix_well_acquisition = ImageXpressWellAcquisition(
+        files=files[files["channel"] == "w1"].drop("z", axis=1),
+        alignment=TileAlignmentOptions.GRID,
+        z_spacing=None,
+    )
+
+    axes = ix_well_acquisition.get_axes()
+    assert axes == ["y", "x"]
+
 
 def test_get_yx_spacing(files):
     ix_well_acquisition = ImageXpressWellAcquisition(

tests/hcs/test_acquisition.py

@@ -242,6 +242,7 @@ def test_get_coordiante_transformations_3d(dummy_well):
     ct = dummy_well.get_coordinate_transformations(
         max_layer=2,
         yx_binning=1,
+        ndim=4,
     )
     assert len(ct) == 3
     assert ct[0] == [
@@ -282,10 +283,7 @@ def test_get_coordiante_transformations_3d(dummy_well):
 def test_get_coordiante_transformations_2d(dummy_well):
     dummy_well.get_z_spacing = lambda: None
     dummy_well.get_yx_spacing = lambda: (2.0, 3.0)
-    ct = dummy_well.get_coordinate_transformations(
-        max_layer=2,
-        yx_binning=1,
-    )
+    ct = dummy_well.get_coordinate_transformations(max_layer=2, yx_binning=1, ndim=3)
     assert len(ct) == 3
     assert ct[0] == [
         {

tests/hcs/test_converter.py

@@ -124,7 +124,7 @@ def plate_acquisition_3d():
 
 def test__create_zarr_plate(tmp_dir, plate_acquisition, hcs_plate):
     converter = ConvertToNGFFPlate(hcs_plate)
-    zarr_plate = converter._create_zarr_plate(plate_acquisition, wells=None)
+    zarr_plate = converter.create_zarr_plate(plate_acquisition, wells=None)
 
     assert exists(join(tmp_dir, "plate_name.zarr"))
     assert zarr_plate.attrs["plate"]["name"] == "plate_name"
@@ -142,7 +142,7 @@ def test__create_zarr_plate(tmp_dir, plate_acquisition, hcs_plate):
     assert not exists(join(tmp_dir, "plate_name.zarr", "E"))
     assert not exists(join(tmp_dir, "plate_name.zarr", "F"))
 
-    zarr_plate_1 = converter._create_zarr_plate(plate_acquisition)
+    zarr_plate_1 = converter.create_zarr_plate(plate_acquisition)
     assert zarr_plate_1 == zarr_plate
 
 
@@ -195,7 +195,7 @@ def test__mean_cast_to():
 
 def test__create_well_group(tmp_dir, plate_acquisition, hcs_plate):
     converter = ConvertToNGFFPlate(hcs_plate)
-    zarr_plate = converter._create_zarr_plate(plate_acquisition)
+    zarr_plate = converter.create_zarr_plate(plate_acquisition)
     well_group = converter._create_well_group(
         plate=zarr_plate,
         well_acquisition=plate_acquisition.get_well_acquisitions()[0],
@@ -262,7 +262,8 @@ def test_run(tmp_dir, plate_acquisition, hcs_plate):
             n_workers=1, threads_per_worker=4, processes=False
         ).get_client(),
     )
-    plate = converter.run(plate_acquisition, max_layer=2)
+    plate = converter.create_zarr_plate(plate_acquisition)
+    plate = converter.run(plate=plate, plate_acquisition=plate_acquisition, max_layer=2)
     plate.attrs["plate"]["wells"] == [
         {"columnIndex": 7, "path": "D/08", "rowIndex": 3},
         {"columnIndex": 2, "path": "E/03", "rowIndex": 4},
@@ -297,7 +298,10 @@ def test_run_selection(tmp_dir, plate_acquisition, hcs_plate):
             n_workers=1, threads_per_worker=4, processes=False
         ).get_client(),
     )
-    plate = converter.run(plate_acquisition, max_layer=2, wells=["D08"])
+    plate = converter.create_zarr_plate(plate_acquisition)
+    plate = converter.run(
+        plate=plate, plate_acquisition=plate_acquisition, max_layer=2, wells=["D08"]
+    )
     plate.attrs["plate"]["wells"] == [{"columnIndex": 7, "path": "D/08", "rowIndex": 3}]
     for well in ["D08"]:
         row, col = well[0], well[1:]

tests/stitching/test_TileStitcher.py

@@ -82,6 +82,34 @@ def test_block_to_tile_map_large_chunks(tiles):
     assert ts._block_to_tile_map[(0, 0, 0, 1, 1)] == [tiles[3]]
 
 
+def test_block_to_tile_map_with_missing_tiles(tiles):
+    """
+    Test mapping of chunk-blocks to tiles if the chunk shape is larger than
+    the tile shape.
+    """
+    ts = DaskTileStitcher(
+        tiles=tiles,
+        chunk_shape=(15, 15),
+        output_shape=(1, 1, 2, 20, 20),
+    )
+
+    assert ts._shape == (1, 1, 2, 20, 20)
+    assert len(ts._block_to_tile_map) == 8
+    assert ts._block_to_tile_map[(0, 0, 0, 0, 0)] == [
+        tiles[0],
+        tiles[1],
+        tiles[2],
+        tiles[3],
+    ]
+    assert ts._block_to_tile_map[(0, 0, 0, 0, 1)] == [tiles[1], tiles[3]]
+    assert ts._block_to_tile_map[(0, 0, 0, 1, 0)] == [tiles[2], tiles[3]]
+    assert ts._block_to_tile_map[(0, 0, 0, 1, 1)] == [tiles[3]]
+    assert ts._block_to_tile_map[(0, 0, 1, 0, 0)] == []
+    assert ts._block_to_tile_map[(0, 0, 1, 0, 1)] == []
+    assert ts._block_to_tile_map[(0, 0, 1, 1, 0)] == []
+    assert ts._block_to_tile_map[(0, 0, 1, 1, 1)] == []
+
+
 def test_block_to_tile_map_small_chunks(tiles):
     """
     Test mapping of chunk-blocks to tiles if the chunk shape is smaller than