Refactor registration

[tayllatheodoro]

Summary

This PR refactors the registration and stabilization pipeline from a set of independent scripts into a unified, modular system. The core change is moving from per-FOV SLURM jobs to a flexible dispatch model where jobs can be submitted per FOV, per timepoint, or per (FOV, T) depending on the method.

Motivation

• estimate-registration and estimate-stabilization were duplicating orchestration logic
• register and stabilize were duplicating transform-application logic
• No unified way to compose transforms (e.g. apply z-stabilization + xy-stabilization in one pass)
• No SLURM parallelism at the timepoint level — all T were processed sequentially within a single job

Architecture

New module: biahub/registration/

Each registration method is now a self-contained module exposing a per-job unit function registered in a method registry dict:

Dict	Parallelism	Methods
PARALLEL_FOV_T_METHODS	one SLURM job per (FOV, T)	ants, pcc, beads-independently, match-z-focus
SEQUENTIAL_T_METHODS	one SLURM job per FOV, T sequential	beads-with-propagation
PARALLEL_FOV_METHODS	one SLURM job per FOV, all T together	stackreg

registration/estimator.py

Orchestrator that:

• Reads config and dispatches to the right parallelism mode based on method
• Passes method-specific kwargs from settings to the per-job function
• Handles file I/O (opens OME-Zarr, passes correct data slice to each job)
• Unified for both stabilization (same channel across T) and registration (different channels/positions)
• Saves per-T transforms as {output_dir}/transforms_per_position/{fov_key}/{t}.npy

registration/applier.py

Unified transform application that:

• Applies transforms per-T, per-FOV, or broadcasts a single transform across all FOVs
• Accepts multiple transform_dirpaths and composes them (T1 @ T2 @ ...) before applying — enabling combined z+xy stabilization in a single pass
• Replaces both register.py and stabilize.py application logic

core/transform.py

New Transform class wrapping homogeneous matrices (4×4 for 3D, 3×3 for 2D) used consistently across all methods.

registration/match_z_focus.py

New method for matching Z focus between channels, enabling cross-channel z-alignment as part of the registration pipeline.

What's in this PR (done)

• core/ module with Transform class and GraphMatching class
• registration/ants.py refactored to expose per-job unit function
• registration/phase_cross_correlation.py refactored
• registration/beads.py refactored with estimate_independently and estimate_with_propagation
• registration/stackreg.py refactored
• registration/match_z_focus.py created
• registration/estimator.py created with method registry and SLURM orchestrator
• registration/utils.py with shared transform utilities

Work in progress

• Complete estimator.py: settings kwargs passthrough, file I/O, registration vs stabilization mode
• Complete applier.py: per-T / per-FOV / broadcast modes + transform composition
• Thin estimate_registration.py and register.py to CLI-only (delegate to estimator/applier)
• Ensure output format compatibility across all methods
• Clean up old orchestrator patterns in method files
• Update settings classes and CLI registration in main.py
• Adapt pipeline for 2D images and 3×3 transforms

Test plan

• Run estimate CLI with each method (ants, pcc, beads, stackreg, match-z-focus) on a test dataset
• Verify per-T .npy files are saved correctly for each parallelism mode
• Run apply CLI with single transform dir and with composed transform dirs
• Verify stabilization mode (same channel) and registration mode (different channels) both work
• Verify 2D input handling once implemented

🤖 Generated with Claude Code