DataRegistration Class

class DataRegistration:
    """Register 2D or 3D MERFISH data across rounds.

    Parameters
    ----------
    datastore : qi2labDataStore
        Initialized qi2labDataStore object
    decon_fiducial: bool, default True
        Deconvolve ALL fiducial rounds. False = only deconvolve round 1 for downstream
        stitching.
    decon_readout: bool, default False
        Deconvolve readout images before saving unwarped readout data for
        decode-time registration.
    overwrite_outputs: bool, default False
        Overwrite existing local transforms and preprocessing outputs.
    perform_deformable_registration: bool, default True
        Estimate and save a residual SOFIMA deformable flow field after affine
        fiducial registration.
    num_gpus: int, default 1
        Number of GPUs to use for registration.
    crop_yx_decon: int, default 2048
        Crop size for deconvolution applied to both y and x dimensions.
    ufish_model: str or pathlib.Path or None, default None
        U-FISH model to use for feature prediction. If omitted or ``None``, use
        the package default ``simfish`` weights. Known aliases include
        ``simfish``, ``merfish``, ``seqfish``,
        ``deepspot``, and ``exseq``. A local ``.onnx``/``.pth`` path or
        HuggingFace weights filename can also be supplied.
    verbose : int, default 1
        Progress verbosity. Set to 0 to suppress routine progress prints.
    """

    def __init__(
        self,
        datastore: qi2labDataStore,
        decon_fiducial: bool = True,
        decon_readout: bool = False,
        overwrite_outputs: bool = False,
        perform_deformable_registration: bool = True,
        num_gpus: int = 1,
        crop_yx_decon: int = 2048,
        ufish_model: str | Path | None = None,
        global_registration: bool = False,
        registration_diagnostics: bool = False,
        sofima_config: SofimaRegistrationConfig | None = None,
        global_registration_config: GlobalRegistrationConfig | None = None,
        verbose: int = 1,
    ) -> None:
        """
        Initialize the object.

        Parameters
        ----------
        datastore : qi2labDataStore
            Datastore containing corrected images, PSFs, metadata, and output
            groups.
        decon_fiducial : bool
            If True, deconvolve fiducial images before registration.
        decon_readout : bool
            If True, deconvolve readout images before U-FISH prediction and
            warping.
        overwrite_outputs : bool
            If True, regenerate local transforms, optional decon images, and
            feature predictor outputs even when they already exist.
        perform_deformable_registration : bool
            If True, estimate a SOFIMA flow field after affine registration.
            The saved field uses channel order X, Y, Z; spatial order Z, Y, X;
            and a patch-centered map origin. Pixel decoding composes this field
            with affine and chromatic transforms when readout images are loaded.
        num_gpus : int
            Number of GPUs available for deconvolution and readout processing.
        crop_yx_decon : int
            Initial lateral RLGC crop size. GPU memory fallback can reduce this
            value and cache the successful size.
        ufish_model : str | Path | None
            U-FISH model alias, weights filename, local path, or None for the
            default model.
        global_registration : bool, default=False
            If True, run global tile registration and fused fiducial OME-Zarr
            generation after local tile preprocessing.
        registration_diagnostics : bool, default=False
            If True, print detailed registration diagnostics.
        sofima_config : SofimaRegistrationConfig or None, default=None
            Explicit SOFIMA deformable registration parameters.
        global_registration_config : GlobalRegistrationConfig or None, default=None
            Explicit multiview-stitcher global registration parameters.
        verbose : int
            Verbosity level for progress messages.
        """
        self._datastore = datastore
        self._decon_fiducial = decon_fiducial
        self._tile_ids = self._datastore.tile_ids
        self._round_ids = self._datastore.round_ids
        self._bit_ids = self._datastore.bit_ids
        self._psfs = self._datastore.channel_psfs
        self._num_gpus = num_gpus
        self._crop_yx_decon = crop_yx_decon
        self._perform_deformable_registration = perform_deformable_registration
        self._tile_id: str | None = None
        self._overwrite_outputs = overwrite_outputs
        self._decon_readout = decon_readout
        self._ufish_model = ufish_model
        self._global_registration = global_registration
        self._registration_diagnostics = bool(registration_diagnostics)
        self._sofima_config = sofima_config or SofimaRegistrationConfig()
        self._global_registration_config = (
            global_registration_config or GlobalRegistrationConfig()
        )
        self._verbose = verbose

    # -----------------------------------
    # property access for class variables
    # -----------------------------------
    @property
    def datastore(self) -> qi2labDataStore:
        """Return the qi2labDataStore object.

        Returns
        -------
        qi2labDataStore
            qi2labDataStore object
        """
        return self._datastore

    @property
    def tile_id(self) -> str | None:
        """Get the current tile id.

        Returns
        -------
        tile_id: int or str
            Tile id
        """
        return self._tile_id

    @tile_id.setter
    def tile_id(self, value: int | str) -> None:
        """Set the tile id.

        Parameters
        ----------
        value : int or str
            Tile id
        """
        if isinstance(value, int):
            if value < 0 or value >= self._datastore.num_tiles:
                print("Set value index >=0 and <" + str(self._datastore.num_tiles))
                return None
            else:
                self._tile_id = self._datastore.tile_ids[value]
        elif isinstance(value, str):
            if value not in self._datastore.tile_ids:
                print("set valid tile id")
                return None
            else:
                self._tile_id = value

    @property
    def perform_deformable_registration(self) -> bool:
        """Get the deformable registration flag.

        Returns
        -------
        bool
            True when SOFIMA deformable flow-field registration is enabled.
        """
        return self._perform_deformable_registration

    @perform_deformable_registration.setter
    def perform_deformable_registration(self, value: bool) -> None:
        """Set the deformable registration flag.

        Parameters
        ----------
        value : bool
            True to enable SOFIMA deformable flow-field registration.
        """
        self._perform_deformable_registration = value

    @property
    def overwrite_outputs(self) -> bool:
        """Get the overwrite_outputs flag.

        Returns
        -------
        bool
            Overwrite existing local transforms and preprocessing outputs.
        """
        return self._overwrite_outputs

    @overwrite_outputs.setter
    def overwrite_outputs(self, value: bool) -> None:
        """Set the overwrite_outputs flag.

        Parameters
        ----------
        value : bool
            Overwrite existing local transforms and preprocessing outputs.
        """
        self._overwrite_outputs = value

    def _entity_root(
        self,
        tile_id: str,
        round_id: str | None = None,
        bit_id: str | None = None,
    ) -> Path:
        """
        Return the native tile/round or tile/bit image directory.

        Parameters
        ----------
        tile_id : str
            Stored tile identifier; None selects the current tile.
        round_id : str | None
            Stored fiducial round identifier.
        bit_id : str | None
            Stored readout bit identifier.

        Returns
        -------
        Path
            Native tile/round or tile/bit directory containing the image stores.
        """
        if (round_id is None and bit_id is None) or (
            round_id is not None and bit_id is not None
        ):
            raise ValueError("Provide either round_id or bit_id, but not both.")

        return self._datastore.local_image_path(
            tile_id, "corrected_data", round=round_id, bit=bit_id
        ).parent

    def _has_valid_deconvolved_fiducial_image(
        self,
        tile_id: str | None = None,
        round_id: str | None = None,
    ) -> bool:
        """
        Has valid native-frame deconvolved fiducial image.

        Parameters
        ----------
        tile_id : str | None
            Stored tile identifier; None selects the current tile.
        round_id : str | None
            Stored fiducial round identifier.

        Returns
        -------
        bool
            True when ``decon_data`` exists and matches the corrected image shape.
        """
        if round_id is None:
            raise ValueError("round_id is required for fiducial images.")

        tile_id = self._tile_id if tile_id is None else tile_id
        entity_root = self._entity_root(tile_id=tile_id, round_id=round_id)
        corrected_shape = self._datastore.image_shape(
            entity_root / Path("corrected_data")
        )
        decon_shape = self._datastore.image_shape(entity_root / Path("decon_data"))
        return corrected_shape is not None and decon_shape == corrected_shape

    def _has_valid_deconvolved_readout_image(
        self,
        tile_id: str | None = None,
        bit_id: str | None = None,
    ) -> bool:
        """
        Has valid native-frame deconvolved readout image.

        Parameters
        ----------
        tile_id : str | None
            Tile identifier.
        bit_id : str | None
            Bit identifier.

        Returns
        -------
        bool
            True when ``decon_data`` exists and matches the corrected image shape.
        """
        tile_id = self._tile_id if tile_id is None else tile_id
        if bit_id is None:
            raise ValueError("bit_id is required for readout images.")

        entity_root = self._entity_root(tile_id=tile_id, bit_id=bit_id)
        corrected_shape = self._datastore.image_shape(
            entity_root / Path("corrected_data")
        )
        readout_shape = self._datastore.image_shape(entity_root / Path("decon_data"))
        return corrected_shape is not None and readout_shape == corrected_shape

    def _has_valid_feature_predictor_outputs(
        self,
        tile_id: str | None = None,
        bit_id: str | None = None,
    ) -> bool:
        """
        Has valid feature predictor outputs.

        Parameters
        ----------
        tile_id : str | None
            Stored tile identifier; None selects the current tile.
        bit_id : str | None
            Stored readout bit identifier.

        Returns
        -------
        bool
            True when predictor and corrected image shapes match and spot localizations exist.
        """
        tile_id = self._tile_id if tile_id is None else tile_id
        if bit_id is None:
            raise ValueError("bit_id is required for feature predictor outputs.")

        entity_root = self._entity_root(tile_id=tile_id, bit_id=bit_id)
        corrected_shape = self._datastore.image_shape(
            entity_root / Path("corrected_data")
        )
        feature_shape = self._datastore.image_shape(
            entity_root / Path(f"{self._datastore.feature_predictor_folder_name}_data")
        )
        spots_path = self._datastore.local_feature_predictor_spots_path(tile_id, bit_id)
        return (
            corrected_shape is not None
            and feature_shape == corrected_shape
            and spots_path.exists()
        )

    def _is_tile_complete(
        self,
        tile_id: str,
        *,
        process_readouts: bool = True,
    ) -> bool:
        """
        Is tile complete.

        Parameters
        ----------
        tile_id : str
            Tile identifier.
        process_readouts : bool, default=True
            If False, assess only fiducial deconvolution and local transforms.

        Returns
        -------
        bool
            True when every requested image and local transform is already present.
        """
        if self._decon_fiducial:
            for round_id in self._round_ids:
                if not self._has_valid_deconvolved_fiducial_image(
                    tile_id=tile_id,
                    round_id=round_id,
                ):
                    return False

        for round_id in self._round_ids:
            if (
                self._datastore.load_local_round_transform_zyx_um(
                    tile=tile_id,
                    round=round_id,
                )
                is None
            ):
                return False

        if self._perform_deformable_registration:
            for round_id in self._round_ids[1:]:
                field = self._datastore.load_local_sofima_flow_field(
                    tile=tile_id, round=round_id, return_future=None
                )
                if field is None or field[0] is None:
                    return False

        if not process_readouts:
            return True

        for bit_id in self._bit_ids:
            if self._decon_readout and not self._has_valid_deconvolved_readout_image(
                tile_id=tile_id, bit_id=bit_id
            ):
                return False
            if not self._has_valid_feature_predictor_outputs(
                tile_id=tile_id, bit_id=bit_id
            ):
                return False

        return True

    def register_all_tiles(self, *, process_readouts: bool = True) -> None:
        """
        Register all tiles.

        Parameters
        ----------
        process_readouts : bool, default=True
            If True, also regenerate readout deconvolution and U-FISH outputs.
            Set False to recompute only the complete fiducial registration
            chain: local affine transforms, optional SOFIMA fields, global
            transforms, and global fiducial fusion.
        """
        tile_ids = list(self._datastore.tile_ids)
        start_idx = 0
        if not self._overwrite_outputs:
            for idx, tile_id in enumerate(tile_ids):
                if self._is_tile_complete(
                    tile_id,
                    process_readouts=process_readouts,
                ):
                    start_idx = idx + 1
                    continue
                start_idx = idx
                break
            else:
                start_idx = len(tile_ids)

            if start_idx >= len(tile_ids):
                if self._verbose >= 1:
                    print(
                        time_stamp(),
                        "All tiles already have complete preprocessing outputs.",
                    )
            elif start_idx > 0:
                if self._verbose >= 1:
                    print(
                        time_stamp(),
                        f"Resuming local registration at tile id: {tile_ids[start_idx]}.",
                    )

        for tile_id in tile_ids[start_idx:]:
            self.tile_id = tile_id
            self._generate_registrations()
            if process_readouts:
                self._apply_registration_to_bits()

        if self._global_registration:
            self.global_register()

    def register_one_tile(self, tile_id: int | str) -> None:
        """Register one tile.

        Parameters
        ----------
        tile_id : int or str
            Tile id
        """
        self.tile_id = tile_id
        if tile_id not in self._tile_ids and tile_id not in range(len(self._tile_ids)):
            return
        self._generate_registrations()
        self._apply_registration_to_bits()

    def apply_registration_to_one_tile(self, tile_id: int | str) -> None:
        """
        Apply existing local registrations to readout bits for one tile.

        This uses the local transforms already stored in the datastore. It does
        not estimate or overwrite fiducial registrations.

        Parameters
        ----------
        tile_id : int or str
            Tile identifier.
        """
        self.tile_id = tile_id
        if tile_id not in self._tile_ids and tile_id not in range(len(self._tile_ids)):
            return
        self._apply_registration_to_bits()

    def load_global_fiducial_views(
        self,
        *,
        use_stored_global_transforms: bool = False,
    ) -> list[Any]:
        """Load native fiducial views with stage or stored global transforms.

        Parameters
        ----------
        use_stored_global_transforms : bool, default False
            Attach saved global transforms instead of only stage metadata.

        Returns
        -------
        list
            Multiview-stitcher images in datastore tile order. Pixels remain lazy.
        """
        from multiview_stitcher import msi_utils, param_utils

        stage_transform_key = "stage_metadata"
        global_transform_key = "global_registered"
        voxel_zyx_um = round_spacing_um(self._datastore.voxel_size_zyx_um)
        if (
            voxel_zyx_um.shape != (3,)
            or not np.all(np.isfinite(voxel_zyx_um))
            or np.any(voxel_zyx_um <= 0)
        ):
            raise ValueError(
                "Global registration requires three positive finite voxel sizes in Z, Y, X order."
            )
        scale = {
            "z": float(voxel_zyx_um[0]),
            "y": float(voxel_zyx_um[1]),
            "x": float(voxel_zyx_um[2]),
        }
        reference_round_id = self._round_ids[0]
        msims = []

        for tile_index, tile_id in enumerate(self._tile_ids):
            if self._verbose >= 1:
                print(
                    time_stamp(),
                    "Loading global registration input "
                    f"tile={tile_id} ({tile_index + 1}/{len(self._tile_ids)})",
                )
            tile_position_zyx_um, affine_zyx_px = (
                self._datastore.load_local_stage_position_zyx_um(
                    tile_id, reference_round_id
                )
            )
            tile_grid_positions = {
                "z": float(np.round(tile_position_zyx_um[0], 2)),
                "y": float(np.round(tile_position_zyx_um[1], 2)),
                "x": float(np.round(tile_position_zyx_um[2], 2)),
            }
            input_path = _local_fiducial_path(
                datastore=self._datastore,
                tile_id=tile_id,
                round_id=reference_round_id,
            )
            sim = _read_fiducial_sim(
                input_path=input_path,
                scale=scale,
                translation=tile_grid_positions,
                affine_zyx_px=affine_zyx_px,
                transform_key=stage_transform_key,
            )
            msim = msi_utils.get_msim_from_sim(sim, scale_factors=[])

            if use_stored_global_transforms:
                affine_zyx_um, _origin_zyx_um, _spacing_zyx_um = (
                    self._datastore.load_global_coord_xforms_um(tile=tile_id)
                )
                if affine_zyx_um is None:
                    raise RuntimeError(
                        "Stored global transform is required for direct global "
                        f"fusion, but none was found for tile {tile_id!r}."
                    )
                msi_utils.set_affine_transform(
                    msim,
                    param_utils.affine_to_xaffine(
                        np.asarray(affine_zyx_um, dtype=np.float32)
                    ),
                    transform_key=global_transform_key,
                    base_transform_key=stage_transform_key,
                )

            msims.append(msim)
            if self._verbose >= 1:
                print(
                    time_stamp(),
                    f"Loaded global registration input tile={tile_id}",
                )
            gc.collect()

        return msims

    def fuse_global_fiducial_views(
        self,
        msims: list[Any],
        *,
        create_max_proj_tiff: bool = True,
    ) -> None:
        """Write registered views on the downsampled fiducial segmentation grid.

        Parameters
        ----------
        msims : list
            Multiview-stitcher images carrying global_registered transforms.
        create_max_proj_tiff : bool, default True
            Also write the maximum-Z OME-TIFF used by Cellpose.
        """
        from multiview_stitcher import fusion, misc_utils, msi_utils
        from multiview_stitcher import spatial_image_utils as si_utils
        from yaozarrs import open_group

        global_transform_key = "global_registered"
        output_zarr_path = self._datastore.fused_image_path()
        # downsample Y/X toward Z spacing for segmentation and ROI coordinates
        voxel_zyx_um = round_spacing_um(self._datastore.voxel_size_zyx_um)
        output_spacing = {
            "z": voxel_zyx_um[0],
            "y": round_pixel_size_um(
                voxel_zyx_um[1] * np.round(voxel_zyx_um[0] / voxel_zyx_um[1], 1)
            ),
            "x": round_pixel_size_um(
                voxel_zyx_um[2] * np.round(voxel_zyx_um[0] / voxel_zyx_um[2], 1)
            ),
        }

        if self._verbose >= 1:
            print(time_stamp(), "Starting global fiducial fusion.")
        fusion_start_time = timeit.default_timer()
        fused_msim = fusion.fuse(
            images=msims,
            transform_key=global_transform_key,
            output_spacing=output_spacing,
            output_zarr_url=str(output_zarr_path),
            **_direct_zarr_fusion_kwargs(misc_utils=misc_utils),
        )
        if self._verbose >= 1:
            print(
                time_stamp(),
                "Finished fusion.fuse direct-to-zarr "
                f"elapsed_s={timeit.default_timer() - fusion_start_time:.2f}",
            )

        metadata_start_time = timeit.default_timer()
        affine = msi_utils.get_transform_from_msim(
            fused_msim,
            transform_key=global_transform_key,
        ).data.squeeze()
        fused_scale0 = msi_utils.get_sim_from_msim(fused_msim)
        origin = si_utils.get_origin_from_sim(fused_scale0, asarray=True)
        spacing = round_spacing_um(
            si_utils.get_spacing_from_sim(fused_scale0, asarray=True)
        )

        qi2labDataStore.save_image_metadata(
            image_path=output_zarr_path,
            extra_attributes={
                "affine_zyx_um": np.asarray(affine, dtype=np.float32).tolist(),
                "origin_zyx_um": np.asarray(origin, dtype=np.float32).tolist(),
                "spacing_zyx_um": round_spacing_um(spacing).tolist(),
            },
            merge=True,
        )
        if self._verbose >= 1:
            print(
                time_stamp(),
                "Finished fused metadata write "
                f"elapsed_s={timeit.default_timer() - metadata_start_time:.2f}",
            )

        del fused_msim
        gc.collect()

        if create_max_proj_tiff:
            projection_start_time = timeit.default_timer()
            cellpose_path = (
                self._datastore.datastore_path / Path("segmentation") / Path("cellpose")
            )
            cellpose_path.mkdir(exist_ok=True)
            filename_path = cellpose_path / Path("fiducial_max_projection.ome.tiff")
            # The TIFF tile iterator requires zarr-python slicing and chunk
            # metadata; yaozarrs opens the store and supplies that adapter.
            fused_array = open_group(output_zarr_path)["0"].to_zarr_python()
            _write_zarr_max_projection_tiff(
                array=fused_array,
                filename_path=filename_path,
                spacing_zyx_um=np.asarray(spacing, dtype=np.float32),
            )
            del fused_array
            if self._verbose >= 1:
                print(
                    time_stamp(),
                    "Finished fused max-projection TIFF "
                    f"elapsed_s={timeit.default_timer() - projection_start_time:.2f}",
                )

        datastore_state = self._datastore.datastore_state.copy()
        datastore_state.update({"GlobalRegistered": True, "Fused": True})
        self._datastore.datastore_state = datastore_state

        if self._verbose >= 1:
            print(
                time_stamp(),
                "Finished global fiducial fusion "
                f"elapsed_s={timeit.default_timer() - fusion_start_time:.2f}",
            )

    def global_register(
        self,
        create_max_proj_tiff: bool = True,
    ) -> None:
        """
        Globally register first-round fiducial tiles and write fused OME-Zarr.

        The method reads each tile's first fiducial round, combines stage
        metadata with multiview-stitcher global registration, saves per-tile
        global transforms back to the datastore, and fuses the transformed
        views directly into the datastore as OME-Zarr v0.5.

        Parameters
        ----------
        create_max_proj_tiff : bool, default=True
            If True, write ``segmentation/cellpose/fiducial_max_projection.ome.tiff``
            from the downsampled fused fiducial OME-Zarr, retaining its spacing.

        Returns
        -------
        None
            Global transforms, fused fiducial OME-Zarr, datastore state, and
            optional max projection are written to the datastore.

        Notes
        -----
        For one tile, save its stage origin and an identity registration
        correction without a pairwise fit. Use ``fuse_global_registered`` to also produce its mosaic.
        """
        from dask import config as dask_config
        from dask.diagnostics import ProgressBar
        from multiview_stitcher import (
            msi_utils,
            param_utils,
            registration,
        )
        from multiview_stitcher import spatial_image_utils as si_utils

        if not self._tile_ids:
            raise ValueError("Global registration requires at least one tile.")
        if len(self._tile_ids) == 1:
            origin, _camera_to_stage = self._datastore.load_local_stage_position_zyx_um(
                self._tile_ids[0], self._round_ids[0]
            )
            self._datastore.save_global_coord_xforms_um(
                affine_zyx_um=np.eye(4, dtype=np.float32),
                origin_zyx_um=origin,
                spacing_zyx_um=round_spacing_um(self._datastore.voxel_size_zyx_um),
                tile=self._tile_ids[0],
            )
            state = self._datastore.datastore_state.copy()
            state["GlobalRegistered"] = True
            self._datastore.datastore_state = state
            if self._verbose >= 1:
                print(
                    time_stamp(),
                    "Skipping global registration because datastore has one tile.",
                )
            return

        if self._verbose >= 1:
            print(time_stamp(), "Starting global fiducial registration.")

        msims = self.load_global_fiducial_views(
            use_stored_global_transforms=False,
        )

        registration_config = self._global_registration_config
        stage_transform_key = "stage_metadata"
        global_transform_key = "global_registered"
        registration_binning = dict(
            zip(
                "zyx",
                (int(value) for value in registration_config.registration_binning_zyx),
                strict=True,
            )
        )
        if self._verbose >= 1:
            print(
                time_stamp(),
                "Running multiview-stitcher global registration "
                f"tiles={len(msims)} registration_binning={registration_binning} "
                "pre_registration_pruning_method=None "
                "post_registration_do_quality_filter="
                f"{registration_config.post_registration_do_quality_filter} "
                "n_parallel_pairwise_regs="
                f"{registration_config.n_parallel_pairwise_regs}",
            )
        registration_start_time = timeit.default_timer()
        progress = ProgressBar() if self._verbose >= 1 else nullcontext()
        with (
            dask_config.set(scheduler=registration_config.dask_scheduler),
            progress,
        ):
            global_transforms = registration.register(
                msims,
                reg_channel_index=int(registration_config.reg_channel_index),
                transform_key=stage_transform_key,
                new_transform_key=global_transform_key,
                pre_registration_pruning_method=None,
                registration_binning=registration_binning,
                post_registration_do_quality_filter=bool(
                    registration_config.post_registration_do_quality_filter
                ),
                groupwise_resolution_kwargs={
                    "reference_view": int(registration_config.reference_view),
                    "transform": registration_config.transform,
                },
                n_parallel_pairwise_regs=int(
                    registration_config.n_parallel_pairwise_regs
                ),
            )
        if self._verbose >= 1:
            print(
                time_stamp(),
                "Finished multiview-stitcher global registration "
                f"elapsed_s={timeit.default_timer() - registration_start_time:.2f}",
            )

        for tile_idx, (msim, transform) in enumerate(
            zip(msims, global_transforms, strict=True)
        ):
            affine = np.asarray(transform).squeeze()
            if registration_config.affine_round_decimals is not None:
                affine = np.round(
                    affine,
                    int(registration_config.affine_round_decimals),
                )
            msi_utils.set_affine_transform(
                msim,
                param_utils.affine_to_xaffine(affine),
                transform_key=global_transform_key,
                base_transform_key=stage_transform_key,
            )
            sim = msi_utils.get_sim_from_msim(msim)
            origin = si_utils.get_origin_from_sim(sim, asarray=True)
            spacing = round_spacing_um(si_utils.get_spacing_from_sim(sim, asarray=True))
            self._datastore.save_global_coord_xforms_um(
                affine_zyx_um=affine,
                origin_zyx_um=origin,
                spacing_zyx_um=spacing,
                tile=tile_idx,
            )

        self.fuse_global_fiducial_views(
            msims=msims,
            create_max_proj_tiff=create_max_proj_tiff,
        )

        if self._verbose >= 1:
            print(time_stamp(), "Finished global fiducial registration.")

    def fuse_global_registered(
        self,
        create_max_proj_tiff: bool = True,
    ) -> None:
        """
        Fuse fiducials, registering globally first when transforms are absent.

        Existing per-tile global transforms are reused when every tile has one.
        Before normal preprocessing, or for a partially registered datastore,
        global registration is run from the first-round fiducials before the
        fused OME-Zarr is created.

        Parameters
        ----------
        create_max_proj_tiff : bool, default=True
            If True, write ``segmentation/cellpose/fiducial_max_projection.ome.tiff``
            from the downsampled fused fiducial OME-Zarr, retaining its spacing.

        Returns
        -------
        None
            Fused fiducial OME-Zarr, datastore state, and optional max
            projection are written to the datastore.
        """
        if not self._global_transforms_available():
            if self._verbose >= 1:
                print(
                    time_stamp(),
                    "Stored global transforms are incomplete; running global "
                    "registration before fusion.",
                )
            self.global_register(create_max_proj_tiff=create_max_proj_tiff)
            if len(self._tile_ids) != 1:
                return

        if self._verbose >= 1:
            print(
                time_stamp(),
                "Starting global fiducial fusion from stored transforms.",
            )

        msims = self.load_global_fiducial_views(
            use_stored_global_transforms=True,
        )
        self.fuse_global_fiducial_views(
            msims=msims,
            create_max_proj_tiff=create_max_proj_tiff,
        )

        if self._verbose >= 1:
            print(
                time_stamp(),
                "Finished global fiducial fusion from stored transforms.",
            )

    def _global_transforms_available(self) -> bool:
        """
        Check whether every tile has a stored global affine transform.

        Returns
        -------
        bool
            ``True`` only when every tile has a non-null global affine.
        """
        for tile_id in self._tile_ids:
            affine_zyx_um, _origin_zyx_um, _spacing_zyx_um = (
                self._datastore.load_global_coord_xforms_um(tile=tile_id)
            )
            if affine_zyx_um is None:
                return False
        return True

    def _generate_registrations(self) -> None:
        """
        Deconvolve fiducials, register rounds, and save local transforms.

        The first fiducial round is loaded, optionally deconvolved, saved, and
        stored as the local reference for the tile. Remaining fiducial rounds
        are split across GPU worker processes. Each worker loads its assigned
        rounds, optionally deconvolves them, registers them to round 1, saves the
        physical-space transform, and estimates SOFIMA residuals when requested.

        Returns
        -------
        None
            Results are written to the datastore.
        """
        if self._num_gpus == 0:
            raise RuntimeError(
                "No GPUs detected. Cannot run _generate_registrations()."
            )

        start_time = timeit.default_timer()
        reference_round_id = self._round_ids[0]

        import cupy as cp

        from merfish3danalysis.utils.rlgc import chunked_rlgc, clear_rlgc_caches

        cp.cuda.Device(0).use()
        reference_image = _load_deconvolve_fiducial_round(
            dr=self,
            round_id=reference_round_id,
            gpu_id=0,
            chunked_rlgc=chunked_rlgc,
        )
        self._datastore.save_local_round_transform_zyx_um(
            np.eye(4, dtype=np.float32),
            tile=self._tile_id,
            round=reference_round_id,
        )
        if self._verbose >= 1:
            print(
                time_stamp(),
                f"Finished fiducial tile id: {self._tile_id}; round id: {reference_round_id}.",
            )

        clear_rlgc_caches(clear_memory_pool=True)

        moving_rounds = list(self._round_ids[1:])
        if not moving_rounds:
            del reference_image
            gc.collect()
            return

        result_queue = mp.Queue()
        num_workers = min(self._num_gpus, len(moving_rounds))
        chunk_size = (len(moving_rounds) + num_workers - 1) // num_workers
        processes = []

        for gpu_id in range(num_workers):
            start = gpu_id * chunk_size
            end = min(start + chunk_size, len(moving_rounds))
            if start >= end:
                break

            subset = moving_rounds[start:end]
            process = mp.Process(
                target=_process_fiducial_rounds_on_gpu,
                args=(self, subset, reference_image, gpu_id, result_queue),
            )
            process.start()
            processes.append(process)

        completed_rounds = set()
        errors = []
        while len(completed_rounds) + len(errors) < len(moving_rounds):
            try:
                status, round_id, payload = result_queue.get(timeout=5)
            except queue.Empty:
                if any(process.exitcode not in (None, 0) for process in processes):
                    break
                if all(process.exitcode is not None for process in processes):
                    break
                continue
            if status == "result":
                completed_rounds.add(round_id)
                if payload is not None:
                    requested_crop_yx = min(
                        int(self._crop_yx_decon),
                        int(max(reference_image.shape[-2:])),
                    )
                    if int(payload) < requested_crop_yx:
                        self._crop_yx_decon = int(payload)
            else:
                errors.append(payload)

        for process in processes:
            process.join()
            if process.exitcode not in (0, None):
                errors.append(
                    f"Fiducial worker pid={process.pid} failed with "
                    f"exitcode={process.exitcode}."
                )

        missing_rounds = [
            round_id for round_id in moving_rounds if round_id not in completed_rounds
        ]
        if missing_rounds:
            errors.append(f"Missing processed fiducial rounds: {missing_rounds}")
        if errors:
            raise RuntimeError("Fiducial registration failed:\n" + "\n".join(errors))

        del reference_image
        clear_rlgc_caches(clear_memory_pool=True)
        gc.collect()

        if self._perform_deformable_registration:
            sofima_queue = mp.Queue()
            sofima_processes = []

            for gpu_id in range(num_workers):
                start = gpu_id * chunk_size
                end = min(start + chunk_size, len(moving_rounds))
                if start >= end:
                    break

                subset = moving_rounds[start:end]
                process = mp.Process(
                    target=_process_sofima_rounds_on_gpu,
                    args=(self, subset, gpu_id, sofima_queue),
                )
                process.start()
                sofima_processes.append(process)

            completed_sofima_rounds = set()
            sofima_errors = []
            while len(completed_sofima_rounds) + len(sofima_errors) < len(
                moving_rounds
            ):
                try:
                    status, round_id, payload = sofima_queue.get(timeout=5)
                except queue.Empty:
                    if any(
                        process.exitcode not in (None, 0)
                        for process in sofima_processes
                    ):
                        break
                    if all(
                        process.exitcode is not None for process in sofima_processes
                    ):
                        break
                    continue
                if status == "result":
                    completed_sofima_rounds.add(round_id)
                else:
                    sofima_errors.append(payload)

            for process in sofima_processes:
                process.join()
                if process.exitcode not in (0, None):
                    sofima_errors.append(
                        f"SOFIMA worker pid={process.pid} failed with "
                        f"exitcode={process.exitcode}."
                    )

            missing_sofima_rounds = [
                round_id
                for round_id in moving_rounds
                if round_id not in completed_sofima_rounds
            ]
            if missing_sofima_rounds:
                sofima_errors.append(
                    f"Missing processed SOFIMA rounds: {missing_sofima_rounds}"
                )
            if sofima_errors:
                raise RuntimeError(
                    "SOFIMA registration failed:\n" + "\n".join(sofima_errors)
                )

        _registration_diag(
            "fiducial_registration_all_done "
            f"tile={self._tile_id} rounds={len(self._round_ids)} "
            f"elapsed_s={timeit.default_timer() - start_time:.2f}",
            enabled=self._registration_diagnostics,
        )
        gc.collect()

    def _apply_registration_to_bits(self) -> None:
        """
        Deconvolve readout bits, run U-FISH, and save unwarped outputs.

        Returns
        -------
        None
            Results are written to the datastore.
        """
        # 1) How many GPUs do we have?
        if self._num_gpus == 0:
            raise RuntimeError(
                "No GPUs detected. Cannot run _apply_registration_to_bits()."
            )

        # 2) Grab all bit IDs and split into `num_gpus` chunks
        all_bits = list(self._bit_ids)
        chunk_size = (
            len(all_bits) + self._num_gpus - 1
        ) // self._num_gpus  # ceiling division

        # 3) Launch one process per GPU (only as many as needed)
        processes = []
        for gpu_id in range(self._num_gpus):
            start = gpu_id * chunk_size
            end = min(start + chunk_size, len(all_bits))
            if start >= end:
                break  # no more bits to assign

            subset = all_bits[start:end]
            p = mp.Process(target=_apply_bits_on_gpu, args=(self, subset, gpu_id))
            p.start()
            processes.append(p)

        # 4) Wait for all GPU-workers to finish
        errors = []
        for p in processes:
            p.join()
            if p.exitcode not in (0, None):
                errors.append(
                    f"Readout worker pid={p.pid} failed with exitcode={p.exitcode}."
                )
        if errors:
            raise RuntimeError("Readout preprocessing failed:\n" + "\n".join(errors))