Beads registration: core module, graph matching, and optimize_matches by tayllatheodoro · Pull Request #211 · czbiohub-sf/biahub

tayllatheodoro · 2026-03-10T00:59:02Z

Summary

Add biahub/core module with Transform class (immutable homogeneous matrix operations, ANTs/scipy backends) and Graph/GraphMatcher (graph-based point matching for bead registration).
Add biahub/registration/beads.py with full beads-based registration pipeline: peak detection, graph matching (Hungarian/descriptor), iterative transform refinement with QC scoring, and optimize_matches for grid search over matching parameters.
Add biahub/registration/ants.py with ANTs intensity-based registration pipeline (preprocessing, registration, postprocessing with crop offset handling).
Add biahub/registration/utils.py with shared utilities: transform validation/interpolation, approximate transform computation, I/O, and volume manipulation helpers.
Improve documentation: module-level docstrings, function docstrings (NumPy-style), and descriptive section comments in debug scripts.
Replace all print() with click.echo() in library code for consistent logging.
Update settings with new BeadsMatchSettings fields (QC, filtering, Hungarian/descriptor matching configs).

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Add module-level docstrings to core and registration modules describing pipeline overview and conventions. Improve function docstrings across beads.py, ants.py, utils.py, and graph_matching.py. Replace all print() calls with click.echo() in library code (scripts unchanged). Add section comments to debug_beads_registration.py. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Copilot

Pull request overview

This PR introduces a comprehensive refactor and extension of the registration pipeline in biahub. It adds a new biahub/core module with a reusable Transform class and Graph/GraphMatcher for graph-based point matching. The existing registration logic scattered across biahub/estimate_registration.py is consolidated into dedicated biahub/registration/beads.py, biahub/registration/ants.py, and biahub/registration/utils.py modules.

Changes:

New biahub/core/transform.py and biahub/core/graph_matching.py providing immutable Transform and graph-based matching classes.
New biahub/registration/beads.py with the full beads pipeline: peak detection, matching, iterative refinement (estimate, optimize_transform), and parameter grid search (optimize_matches).
Settings refactored: t_reference moved from BeadsMatchSettings to AffineTransformSettings; new FilterMatchesSettings and QCBeadsRegistrationSettings added; stabilization_type extended with "affine".

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 13 comments.

Show a summary per file

File	Description
`biahub/core/transform.py`	New `Transform` class wrapping homogeneous matrices with ANTs/scipy backends
`biahub/core/graph_matching.py`	New `Graph` and `GraphMatcher` classes for point cloud matching
`biahub/core/__init__.py`	Package init for new core module
`biahub/registration/beads.py`	Full beads-based registration pipeline including `optimize_transform` and `estimate`
`biahub/registration/utils.py`	Shared utilities moved/refactored from `estimate_registration.py`
`biahub/registration/ants.py`	New ANTs-based registration module replacing `_optimize_registration`
`biahub/settings.py`	Added `FilterMatchesSettings`, `QCBeadsRegistrationSettings`; moved `t_reference` to `AffineTransformSettings`
`biahub/estimate_stabilization.py`	Replaced `estimate_xyz_stabilization_with_beads` with `estimate_tczyx` from new module
`biahub/estimate_registration.py`	Replaced inline beads/ANTs logic with new module imports
`biahub/stabilize.py`	Replaced `create_empty_hcs_zarr` with `create_empty_plate` from iohub
`settings/example_estimate_registration_settings_beads.yml`	Updated to new settings structure
`settings/example_estimate_stabilization_settings_xyz_beads.yml`	Updated to new settings structure
`scripts/debug_beads_registration.py`	Refactored debug script using new module APIs
`scripts/debug_beads_stabilization.py`	Refactored debug script using new module APIs

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

biahub/registration/beads.py

Copilot · 2026-03-10T01:04:02Z

biahub/registration/beads.py

+            ref_voxel_size=ref_voxel_size,
+            mov_voxel_size=mov_voxel_size,
+        )
+        click.echo("Computed approx transform: ", approx_transform)


click.echo() does not accept multiple positional arguments the way print() does. The call click.echo("Computed approx transform: ", approx_transform) passes approx_transform as the second positional argument (file), silently writing the string to stdout and the transform value being interpreted as a file object. This will raise a TypeError at runtime. The string and value should be combined in a single argument, e.g. click.echo(f"Computed approx transform: {approx_transform}").

Suggested change

click.echo("Computed approx transform: ", approx_transform)

click.echo(f"Computed approx transform: {approx_transform}")

Copilot · 2026-03-10T01:04:03Z

biahub/registration/beads.py

+
+        if user_transform is not None and current_iterations == 0:
+            click.echo("Optimizing user transform:")
+            user_transform = Transform(matrix=np.asarray(user_transform))


The user_transform parameter is typed as Transform in the signature (line 799), but on line 1082 it is wrapped again in Transform(matrix=np.asarray(user_transform)). Since the input is already a Transform object, calling np.asarray() on it would not return a matrix array — it would just create an array containing the Transform object itself. This will raise a ValueError inside Transform.__init__ because the resulting array shape won't be (3,3) or (4,4). The Transform(...) wrapping should be removed since user_transform is already a Transform instance.

Suggested change

user_transform = Transform(matrix=np.asarray(user_transform))

Copilot · 2026-03-10T01:04:03Z

biahub/registration/utils.py

+    axs[0].plot(z_transforms)
+    axs[0].set_title("Z-Translation")
+    axs[1].plot(x_transforms)
+    axs[1].set_title("X-Translation")
+    axs[2].plot(y_transforms)
+    axs[2].set_title("Y-Translation")


The plot_translations function plots X-Translation in the middle panel (axs[1]) and Y-Translation in the bottom panel (axs[2]), but the variable names and plot order are mismatched: y_transforms feeds axs[1] which is labeled "X-Translation", and x_transforms feeds axs[2] which is labeled "Y-Translation". Either the plot titles or the variable assignment is inverted — since the matrix indexing correctly assigns transforms_zyx[:, 1, 3] to y_transforms and transforms_zyx[:, 2, 3] to x_transforms, the titles should be swapped so that index 1 → Y-Translation and index 2 → X-Translation.

Copilot · 2026-03-10T01:04:03Z

biahub/registration/beads.py

+    mov_peaks_optimized, ref_peaks_optimized = peaks_from_beads(
+        mov=mov_reg_optimized,
+        ref=ref,
+        mov_peaks_settings=beads_match_settings.source_peaks_settings,
+        ref_peaks_settings=beads_match_settings.target_peaks_settings,
+        verbose=debug,
+    )
+
+    quality_score_optimized = overlap_score(
+        mov_peaks=mov_peaks_optimized,
+        ref_peaks=ref_peaks_optimized,
+        radius=beads_match_settings.qc_settings.score_centroid_mask_radius,
+        verbose=debug,
+    )


The peaks_from_beads function returns None implicitly (via bare return) when fewer than 2 peaks are detected (line 622). However, the callers in optimize_transform (line 969-975) and optimize_matches (lines 201-209) call the function expecting a tuple and then immediately unpack both results. When peaks_from_beads returns None, the unpacking mov_peaks_optimized, ref_peaks_optimized = peaks_from_beads(...) will raise a TypeError. The null-check guard is only present in optimize_matches, not in optimize_transform for the second call (lines 969-975).

Copilot · 2026-03-10T01:04:04Z

biahub/registration/ants.py

+            center = (ref_mask.shape[-2] // 2, ref_mask.shape[-1] // 2)
+            radius = int(ref_mask_radius * min(center))
+
+            ref_mask[(x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius**2] = True


The circular mask equation has swapped x and y axes. np.ogrid returns y as the first axis and x as the second, but the condition (x - center[0]) ** 2 + (y - center[1]) ** 2 uses center[0] (the Y center) to compare against x and center[1] (the X center) to compare against y. It should be (y - center[0]) ** 2 + (x - center[1]) ** 2, where center = (shape[-2] // 2, shape[-1] // 2) maps to (y_center, x_center).

Suggested change

ref_mask[(x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius**2] = True

ref_mask[(y - center[0]) ** 2 + (x - center[1]) ** 2 <= radius**2] = True

Copilot · 2026-03-10T01:04:04Z

biahub/settings.py

-    )
+    stabilization_type: Literal["z", "xy", "xyz", "affine"]
+    stabilization_method: Literal[
+        "beads", "phase-cross-corr", "focus-finding", "manual", "ants", "beads"


The StabilizationSettings.stabilization_method Literal type has "beads" duplicated in the allowed values list. This is harmless to Pydantic's validation (duplicates are treated as a single option), but it is a typo/mistake. The second "beads" entry should be removed, or if a different value was intended, it should be corrected.

Suggested change

"beads", "phase-cross-corr", "focus-finding", "manual", "ants", "beads"

"beads", "phase-cross-corr", "focus-finding", "manual", "ants"

Copilot · 2026-03-10T01:04:04Z

settings/example_estimate_stabilization_settings_xyz_beads.yml


 # Settings for affine transformation estimation (applies to all stabilization methods)
 affine_transform_settings:
+  t_reference: "first" # Reference timepoint for transform


The t_reference field was moved from BeadsMatchSettings to AffineTransformSettings. However, the stabilization YAML file example_estimate_stabilization_settings_xyz_beads.yml now places t_reference under affine_transform_settings, but the stabilization code in estimate_stabilization.py calls estimate_tczyx(..., mode="stabilization"), which resolves t_reference from affine_transform_settings.t_reference. This is a behavioral change worth being aware of: in the previous BeadsMatchSettings.t_reference it was strictly scoped to bead matching, whereas the new AffineTransformSettings.t_reference applies to all transform estimation modes including ANTs and manual. This cross-cutting concern may cause unexpected behavior when using non-beads methods with this setting in stabilization mode.

Copilot · 2026-03-10T01:04:05Z

biahub/core/transform.py

+    def to_ants(self):
+        """
+        Convert to ANTs transform.
+
+        Returns
+        -------
+        ants.ANTsTransform
+            ANTs transform object.
+
+        Notes
+        -----
+        Requires ANTsPy to be installed.
+        Works for both 2D and 3D transforms.
+        """
+        try:
+            import ants
+        except ImportError:
+            raise ImportError(
+                "ANTsPy is required for to_ants(). Install with: pip install antspyx"
+            )
+        if self._ndim not in (2, 3):
+            raise ValueError(f"Unsupported ndim: {self._ndim}")
+        T_ants_style = self._matrix[:, :-1].ravel()
+        T_ants_style[-self._ndim :] = self._matrix[: self._ndim, -1]
+        T_ants = ants.new_ants_transform(
+            transform_type="AffineTransform",
+            dimension=self._ndim,
+        )
+        T_ants.set_parameters(T_ants_style)
+        return T_ants


The to_ants() method builds the ANTs-style parameter array as self._matrix[:, :-1].ravel() which creates a flat array of shape (ndim*(ndim+1),). It then sets the last ndim elements to the translation: T_ants_style[-self._ndim:] = self._matrix[:self._ndim, -1]. However, the original convert_transform_to_ants in utils.py (which is the reference implementation tested in test_affine.py) uses T_numpy[:, :-1].ravel() and then sets T_ants_style[-3:] = T_numpy[:3, -1].

For a 4x4 matrix, [:, :-1].ravel() gives a 12-element array (4 rows × 3 columns). Setting [-3:] with the translation column correctly fills the last 3 elements. The logic in the new to_ants() in Transform is equivalent. But there's an inconsistency: the old function uses ants.new_ants_transform(transform_type="AffineTransform") without specifying dimension, while the new one passes dimension=self._ndim. Verify this behaves correctly for 2D (3×3) transforms since the default dimension in ANTs may be 3.

Copilot · 2026-03-10T01:04:05Z

biahub/registration/beads.py

+        if quality_score_optimized == 1:
+            break
+        transform = optimized_transform
+
+        if user_transform is not None and current_iterations == 0:
+            click.echo("Optimizing user transform:")
+            user_transform = Transform(matrix=np.asarray(user_transform))
+            optimized_transform_user, quality_score_optimized_user = optimize_transform(
+                transform=user_transform,
+                mov=mov,
+                ref=ref,
+                beads_match_settings=beads_match_settings,
+                affine_transform_settings=affine_transform_settings,
+                verbose=verbose,
+                debug=debug,
+            )
+
+            if quality_score_optimized < quality_score_optimized_user:
+
+                transform_iter_dict[current_iterations] = {
+                    "transform": optimized_transform_user,
+                    "quality_score": quality_score_optimized_user,
+                }
+                if quality_score_optimized_user == 1:
+                    break
+                transform = optimized_transform_user
+
+        if transform is None:
+            break


In estimate (beads.py), when optimize_transform returns (None, -1) (e.g., not enough matches found), the code sets transform = None (line 1078) and then proceeds to call optimize_transform again with user_transform on lines 1083-1091. On the next loop iteration, optimize_transform will receive transform=None and call transform.to_ants(), causing an AttributeError. The if transform is None: break check on line 1103 runs after the user_transform block, so it's too late to protect the next iteration call to optimize_transform(transform=transform, ...) on line 1063 when transform is set to None by optimized_transform on line 1078.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

ieivanov · 2026-03-14T00:38:05Z

Thanks @tayllatheodoro!

Could you put this PR in the context of #89? Does this PR close any outstanding issues? What work is still outstanding in the overhaul of registration and stabilization? Does this PR address beads registration only, but there is more work needed to apply these transforms to cells and neuromasts? Feel free to take over issue #89 and use it to plan and document your work on this project, similar to how I've done in czbiohub-sf/shrimPy#211

ieivanov · 2026-03-14T00:39:23Z

I'm happy if we merge this PR, I just want to understand how it fits in the larger project.

tayllatheodoro and others added 19 commits March 9, 2026 10:41

Add core module and beads/ants registration from refactor branch

0745690

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

beads stabilization unavailable, not used in pipelines

c7c993c

beads registration

c69924c

use prev trasnform for registration

94f27d5

beads settings with filtering and qc settings

5873c7c

beads stabilization

7fdbe50

updated registration settings

134025e

style

46fecf6

remove applier of transform

99bdc0c

add get approx trasnform function from manual registration

4707799

move helper function to utils

21a88b1

update to iohub create empty plate

c07cdee

flag to compute approx trasnform

8b72982

update settings with flag

b4d6fab

style

ffbdcb2

optimize matches for future improvements

3234376

update documnetation

3fe3df8

update debug beads stabilization

ceca913

tayllatheodoro requested a review from Copilot March 10, 2026 00:59

Copilot started reviewing on behalf of tayllatheodoro March 10, 2026 00:59 View session

Copilot AI reviewed Mar 10, 2026

View reviewed changes

Update biahub/registration/beads.py

7369f4c

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

tayllatheodoro self-assigned this Mar 10, 2026

tayllatheodoro added the enhancement New feature or request label Mar 10, 2026

tayllatheodoro requested a review from mattersoflight March 10, 2026 01:07

mattersoflight approved these changes Mar 13, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Beads registration: core module, graph matching, and optimize_matches#211

Beads registration: core module, graph matching, and optimize_matches#211
tayllatheodoro wants to merge 20 commits intomainfrom
beads_registration

tayllatheodoro commented Mar 10, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI Mar 10, 2026

Uh oh!

Copilot AI Mar 10, 2026

Uh oh!

Copilot AI Mar 10, 2026

Uh oh!

Copilot AI Mar 10, 2026

Uh oh!

Copilot AI Mar 10, 2026

Uh oh!

Copilot AI Mar 10, 2026

Uh oh!

Copilot AI Mar 10, 2026

Uh oh!

Copilot AI Mar 10, 2026

Uh oh!

Copilot AI Mar 10, 2026

Uh oh!

ieivanov commented Mar 14, 2026

Uh oh!

ieivanov commented Mar 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	click.echo("Computed approx transform: ", approx_transform)
	click.echo(f"Computed approx transform: {approx_transform}")

	ref_mask[(x - center[0]) 2 + (y - center[1]) 2 <= radius**2] = True
	ref_mask[(y - center[0]) 2 + (x - center[1]) 2 <= radius**2] = True

	"beads", "phase-cross-corr", "focus-finding", "manual", "ants", "beads"
	"beads", "phase-cross-corr", "focus-finding", "manual", "ants"

Conversation

tayllatheodoro commented Mar 10, 2026

Summary

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

ieivanov commented Mar 14, 2026

Uh oh!

ieivanov commented Mar 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants