BiaHub Registration & Stabilization Refactoring

[tayllatheodoro]

BiaHub Registration & Stabilization Refactoring Plan

Issue: #89
Continue work after: #98

Current Problems Identified

1. Code Duplication: Registration (between channels) and stabilization (within channels over time) share similar transformation logic
2. CLI Confusion: Multiple similar commands (estimate-registration, optimize-registration, register, estimate-stabilization, stabilize)
3. Conceptual Overlap: Both involve spatial/temporal alignment with similar mathematical operations
4. Maintenance Burden: Changes need to be replicated across multiple similar codebases

Proposed Architecture

1. Core Abstraction: Unified Transformation Engine

Create a shared Transform system that handles both registration and stabilization.

2. Implementation Strategy

Phase 1: Create Core Transform Engine

• Implement TransformEngine base class with shared algorithms
• Create Transform data structures for parameter representation
• Implement shared estimation, optimization, and application logic
• Add comprehensive tests for core functionality

Phase 2: Refactor Registration

• Move registration logic to use TransformEngine
• Implement registration-specific Estimator, Optimizer, Applicator
• Create new biahub registration CLI commands
• Maintain backward compatibility with deprecation warnings

Phase 3: Refactor Stabilization

• Move stabilization logic to use TransformEngine
• Implement stabilization-specific components
• Create new biahub stabilization CLI commands
• Maintain backward compatibility

Phase 4: Migration & Cleanup

• Add migration guide and examples
• Remove deprecated CLI commands after the appropriate notice period
• Update documentation and tutorials

5. Configuration Unification

Create a unified configuration schema that works for:

• estimate/apply transforms between channels (registration) or within one channel over time (stabilization)
• estimate/apply a transform for all time points or one per time point
• estimate/apply transforms for all positions or one per position
• flexible transform dimension
• flexible use of the methods
• add common transforms manipulation options (inv, multiply, transpose,...) with hpc

6. Benefits of This Refactoring

1. Code Reuse: ~60-70% of transformation logic can be shared
2. Maintainability: Single place to fix bugs and add features
3. Extensibility: Easy to add new transform types (e.g., chromatic aberration correction)
4. Consistency: Unified parameter formats and algorithms across operations
5. Testing: Centralized test suites for core functionality
6. Documentation: Single place to document transformation concepts

7. Backward Compatibility Strategy

• Keep old CLI commands working with deprecation warnings for 2-3 releases
• Provide automatic config migration utilities
• Maintain identical output formats during transition
• Create clear migration documentation with examples

8. Success Metrics

• Lines of code reduction: Target 40% reduction in transform-related code
• Bug reduction: Fewer transform-related issues due to code consolidation
• Development velocity: Faster implementation of new transform features
• User experience: Clearer, more intuitive CLI interface
• Test coverage: Improved test coverage through shared test infrastructure

Implementation Priority

1. High Priority: Core TransformEngine and shared algorithms
2. Medium Priority: Registration refactoring (more complex, used more frequently)
3. Medium Priority: Stabilization refactoring
4. Low Priority: Advanced features (new transform types, GUI tools, etc.)

This refactoring will significantly improve code maintainability while providing a clearer, more powerful interface for users performing image registration and stabilization tasks.

[mattersoflight]

Hi @tayllatheodoro thanks for writing this plan. One important addition to the testing strategy is to include real data in the testing loop. I suggest 5 crops from DynaCell and Neuromast movies (about the size of few cells) that need stabilization and registration, and automating tests. The automation can be via pytest or via a script.

The initial evaluation can be done with human-in-the-loop where the code generates 5 movies before/after correction and plays them for you to quickly see if the code refactor or algorithmic changes are affecting the stabilization and registration.

See an example of similar (not the same) strategy for testing widgets here: https://github.com/mehta-lab/waveorder/blob/main/tests/widget_tests/test_sample_contributions.py#L1