rustytorch/crates/specialized/rtx-registration/src/lib.rs

//! # RTX Registration
//!
//! Image registration library for medical imaging.
//!
//! This crate provides tools for:
//! - Spatial transformations (rigid, affine)
//! - Image similarity metrics (MSE, NCC)
//! - Optimization algorithms for registration
//! - Volume resampling with interpolation
//!
//! ## Example
//!
//! ```rust,no_run
//! use rtx_registration::prelude::*;
//! use rtx_medical_io::Volume;
//!
//! // Create transforms
//! let rotation = RigidTransform::rotation_z(0.1);
//! let translation = RigidTransform::translation(5.0, 0.0, 0.0);
//! let combined = rotation.compose(&translation);
//!
//! // Transform a point
//! let point = nalgebra::Point3::new(1.0, 2.0, 3.0);
//! let transformed = combined.transform_point(&point);
//! ```

pub mod error;
pub mod interpolate;
pub mod metric;
pub mod optimizer;
pub mod transform;

pub use error::{RegistrationError, Result};

/// Prelude module with commonly used types.
pub mod prelude {
    pub use crate::error::{RegistrationError, Result};
    pub use crate::interpolate::{
        InterpolationMethod, interpolate, interpolate_nearest, interpolate_trilinear,
        resample_volume, resample_volume_parallel,
    };
    pub use crate::metric::{
        MeanSquaredError, NormalizedCrossCorrelation, SimilarityMetric, compute_gradient,
    };
    pub use crate::optimizer::{
        GradientDescent, OptimizationConfig, OptimizationResult, Optimizer, PowellOptimizer,
    };
    pub use crate::transform::{
        AffineTransform, ComposedTransform, RigidTransform, Transform, centroid, transform_points,
    };
}

/// Register a moving volume to a fixed volume using rigid transformation.
///
/// Returns the optimal rigid transform that aligns the moving volume to the fixed volume.
pub fn register_rigid<M: metric::SimilarityMetric>(
    fixed: &rtx_medical_io::Volume,
    moving: &rtx_medical_io::Volume,
    metric: &M,
    mask: Option<&rtx_medical_io::Volume>,
    config: &optimizer::OptimizationConfig,
) -> Result<transform::RigidTransform> {
    use interpolate::{InterpolationMethod, resample_volume_parallel};
    use optimizer::Optimizer;
    use transform::Transform;

    let optimizer = optimizer::PowellOptimizer::new();

    // Cost function: resample moving with transform, compute metric
    let cost_fn = |params: &[f64]| -> f64 {
        let mut t = transform::RigidTransform::identity();
        t.set_parameters(params);

        let resampled = resample_volume_parallel(
            moving,
            &t,
            fixed.shape(),
            fixed.spacing(),
            InterpolationMethod::Trilinear,
        );

        metric.to_cost(metric.compute(fixed, &resampled, mask))
    };

    let initial_params = vec![0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
    let result = optimizer.optimize(&initial_params, cost_fn, config)?;

    let mut transform = transform::RigidTransform::identity();
    transform.set_parameters(&result.parameters);
    Ok(transform)
}

/// Register a moving volume to a fixed volume using affine transformation.
///
/// Returns the optimal affine transform that aligns the moving volume to the fixed volume.
pub fn register_affine<M: metric::SimilarityMetric>(
    fixed: &rtx_medical_io::Volume,
    moving: &rtx_medical_io::Volume,
    metric: &M,
    mask: Option<&rtx_medical_io::Volume>,
    config: &optimizer::OptimizationConfig,
) -> Result<transform::AffineTransform> {
    use interpolate::{InterpolationMethod, resample_volume_parallel};
    use optimizer::Optimizer;
    use transform::Transform;

    let optimizer = optimizer::PowellOptimizer::new();

    // Cost function
    let cost_fn = |params: &[f64]| -> f64 {
        let mut t = transform::AffineTransform::identity();
        t.set_parameters(params);

        // Check if transform is valid (not degenerate)
        if !t.is_invertible() {
            return f64::MAX;
        }

        let resampled = resample_volume_parallel(
            moving,
            &t,
            fixed.shape(),
            fixed.spacing(),
            InterpolationMethod::Trilinear,
        );

        metric.to_cost(metric.compute(fixed, &resampled, mask))
    };

    // Start with identity transform
    let initial_params = transform::AffineTransform::identity().get_parameters();
    let result = optimizer.optimize(&initial_params, cost_fn, config)?;

    let mut transform = transform::AffineTransform::identity();
    transform.set_parameters(&result.parameters);
    Ok(transform)
}

#[cfg(test)]
mod tests {
    use super::*;
    use nalgebra::Point3;
    use rtx_medical_io::Volume;

    #[test]
    fn test_prelude_imports() {
        use prelude::*;

        let t = RigidTransform::identity();
        let p = Point3::new(1.0, 2.0, 3.0);
        let q = t.transform_point(&p);
        assert!((p - q).norm() < 1e-10);
    }

    #[test]
    fn test_transform_chain() {
        use prelude::*;

        // Create a chain of transforms
        let t1 = RigidTransform::rotation_z(std::f64::consts::PI / 4.0);
        let t2 = RigidTransform::translation(10.0, 0.0, 0.0);
        let composed = t1.compose(&t2);

        // Verify it works
        let p = Point3::new(0.0, 0.0, 0.0);
        let q = composed.transform_point(&p);
        assert!(q.coords.norm() > 0.0);
    }

    #[test]
    fn test_similarity_metrics() {
        use prelude::*;

        let mut vol = Volume::zeros([10, 10, 10]);
        for z in 0..10 {
            for y in 0..10 {
                for x in 0..10 {
                    vol.set(x, y, z, (x + y + z) as f64);
                }
            }
        }

        let mse = MeanSquaredError::new();
        let ncc = NormalizedCrossCorrelation::new();

        // Same image should have MSE=0, NCC=1
        assert!(mse.compute(&vol, &vol, None) < 1e-10);
        assert!((ncc.compute(&vol, &vol, None) - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_interpolation() {
        use prelude::*;

        let mut vol = Volume::zeros([10, 10, 10]);
        vol.set(5, 5, 5, 100.0);

        let p = Point3::new(5.0, 5.0, 5.0);
        let v = interpolate(&vol, &p, InterpolationMethod::Trilinear);
        assert_eq!(v, 100.0);
    }

    #[test]
    fn test_optimizer() {
        use prelude::*;

        let optimizer = GradientDescent::new();
        let config = OptimizationConfig {
            max_iterations: 1000,
            tolerance: 1e-8,
            step_size: 0.1,
            verbose: false,
        };

        let cost_fn = |params: &[f64]| params[0].powi(2) + params[1].powi(2);

        let result = optimizer.optimize(&[5.0, 5.0], cost_fn, &config).unwrap();

        // Check that we're close to the minimum
        assert!(result.cost < 0.1);
    }
}