rustalign_aligner/

smith_waterman.rs

//! Smith-Waterman dynamic programming alignment
//!
//! This implements the banded Smith-Waterman algorithm for local alignment,
//! matching the C++ aligner_sw.h implementation.

use super::Alignment;
use rustalign_common::{Nuc, Result as BwtResult, Score, Strand};

/// Parameters for Smith-Waterman alignment
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SwParams {
    /// Match score
    pub match_score: Score,

    /// Mismatch penalty
    pub mismatch_penalty: Score,

    /// Gap open penalty
    pub gap_open: Score,

    /// Gap extend penalty
    pub gap_extend: Score,

    /// Band width for banded alignment (0 = no banding)
    pub band_width: usize,

    /// Minimum score threshold
    pub min_score: Score,
}

impl Default for SwParams {
    fn default() -> Self {
        Self {
            match_score: 2,
            mismatch_penalty: -3,
            gap_open: -5,
            gap_extend: -2,
            band_width: 15,
            min_score: 30,
        }
    }
}

/// DP cell for Smith-Waterman alignment
#[derive(Debug, Clone, Copy, Default)]
struct DpCell {
    /// Score from diagonal (match/mismatch)
    h: Score,
    /// Score from above (insertion)
    e: Score,
    /// Score from left (deletion)
    f: Score,
    /// Best score at this cell
    #[allow(dead_code)]
    best: Score,
}

/// Smith-Waterman aligner with banded DP
pub struct SwAligner {
    /// Alignment parameters
    params: SwParams,
}

impl SwAligner {
    /// Create a new Smith-Waterman aligner
    pub fn new(params: SwParams) -> Self {
        Self { params }
    }

    /// Align a read to a reference sequence
    ///
    /// Uses Smith-Waterman DP with optional banding for optimization.
    /// When band_width > 0, only computes cells within band_width/2 of the main diagonal.
    pub fn align(&self, read: &[Nuc], reference: &[Nuc]) -> BwtResult<Alignment> {
        let read_len = read.len();
        let ref_len = reference.len();

        if read_len == 0 || ref_len == 0 {
            return Ok(Alignment {
                ref_start: 0,
                ref_end: 0,
                score: 0,
                strand: Strand::Forward,
                edits: 0,
                ..Default::default()
            });
        }

        // Use banded alignment if band_width is set
        if self.params.band_width > 0 && self.params.band_width < read_len {
            return self.align_banded(read, reference);
        }

        // Use full DP matrix for correctness or small sequences
        self.align_full(read, reference)
    }

    /// Full (non-banded) Smith-Waterman alignment
    fn align_full(&self, read: &[Nuc], reference: &[Nuc]) -> BwtResult<Alignment> {
        let read_len = read.len();
        let ref_len = reference.len();

        let mut dp_matrix = vec![vec![DpCell::default(); ref_len + 1]; read_len + 1];
        let mut best_score = Score::MIN;
        let mut best_pos = (0, 0);

        // Initialize first row and column
        #[allow(clippy::needless_range_loop)]
        for i in 0..=read_len {
            dp_matrix[i][0].h = 0;
            dp_matrix[i][0].e = Score::MIN / 2;
            dp_matrix[i][0].f = Score::MIN / 2;
        }
        #[allow(clippy::needless_range_loop)]
        for j in 0..=ref_len {
            dp_matrix[0][j].h = 0;
            dp_matrix[0][j].e = Score::MIN / 2;
            dp_matrix[0][j].f = Score::MIN / 2;
        }

        // Fill DP matrix
        for i in 1..=read_len {
            for j in 1..=ref_len {
                // Calculate H score (diagonal)
                let match_score = if read[i - 1] == reference[j - 1] {
                    self.params.match_score
                } else {
                    self.params.mismatch_penalty
                };
                let h_diag = dp_matrix[i - 1][j - 1].h + match_score;

                // Calculate E score (from above/insertion)
                let e_gap = dp_matrix[i - 1][j].e + self.params.gap_extend;
                let h_gap = dp_matrix[i - 1][j].h + self.params.gap_open + self.params.gap_extend;
                let e_new = e_gap.max(h_gap).max(Score::MIN / 2);

                // Calculate F score (from left/deletion)
                let f_gap = dp_matrix[i][j - 1].f + self.params.gap_extend;
                let h_prev = dp_matrix[i][j - 1].h;
                let f_new = f_gap
                    .max(h_prev + self.params.gap_open + self.params.gap_extend)
                    .max(Score::MIN / 2);

                // H cell: max of diagonal, E, F, and 0 (for local alignment)
                let h_new = h_diag.max(e_new).max(f_new).max(0);

                dp_matrix[i][j].h = h_new;
                dp_matrix[i][j].e = e_new;
                dp_matrix[i][j].f = f_new;

                // Track best score
                if h_new > best_score {
                    best_score = h_new;
                    best_pos = (i, j);
                }
            }
        }

        // Count edits from traceback
        let (edits, ref_start) = self.traceback_simple(read, reference, best_pos);

        Ok(Alignment {
            ref_start,
            ref_end: best_pos.1,
            score: best_score,
            strand: Strand::Forward,
            edits,
            ..Default::default()
        })
    }

    /// Banded Smith-Waterman alignment
    ///
    /// Only computes cells within band_width/2 of the main diagonal.
    /// This reduces complexity from O(m*n) to O(band_width * min(m,n)).
    #[allow(clippy::manual_div_ceil)]
    fn align_banded(&self, read: &[Nuc], reference: &[Nuc]) -> BwtResult<Alignment> {
        let read_len = read.len();
        let ref_len = reference.len();
        let band_half = (self.params.band_width + 1) / 2; // Round up

        // Banded DP matrix: only store cells within the band
        // For each row i, we store columns from max(0, i-band_half) to min(ref_len, i+band_half)
        let mut best_score = Score::MIN;
        let mut best_pos = (0, 0);

        // Use a 2D array for the banded region to simplify diagonal access
        // Band size needs to account for the offset
        let band_size = self.params.band_width + 2;
        let mut dp_band = vec![vec![DpCell::default(); band_size]; read_len + 1];

        for i in 0..=read_len {
            // Determine column range for this row
            let col_start = i.saturating_sub(band_half);
            let col_end = (i + band_half + 1).min(ref_len + 1);

            for j in col_start..col_end {
                let band_idx = j.saturating_sub(col_start);
                if band_idx >= band_size {
                    continue;
                }

                // Calculate all needed scores first to avoid borrow conflicts
                let (h_diag, e_above, h_above, f_left, h_left, _nuc_match) = if i == 0 || j == 0 {
                    (Score::MIN, Score::MIN / 2, 0, Score::MIN / 2, 0, false)
                } else {
                    // Calculate diagonal score
                    let prev_col_start = (i - 1).saturating_sub(band_half);
                    let diag_band_idx = (j - 1).saturating_sub(prev_col_start);
                    let h_diag = if diag_band_idx < band_size && i > 0 && j > 0 {
                        let match_score = if read[i - 1] == reference[j - 1] {
                            self.params.match_score
                        } else {
                            self.params.mismatch_penalty
                        };
                        dp_band[i - 1][diag_band_idx].h + match_score
                    } else {
                        Score::MIN
                    };

                    // Calculate score from above
                    let above_band_idx = j.saturating_sub(col_start);
                    let (e_above, h_above) = if i > 0 && above_band_idx < band_size {
                        (
                            dp_band[i - 1][above_band_idx].e,
                            dp_band[i - 1][above_band_idx].h,
                        )
                    } else {
                        (Score::MIN / 2, 0)
                    };

                    // Calculate score from left
                    let left_band_idx = (j - 1).saturating_sub(col_start);
                    let (f_left, h_left) = if left_band_idx < band_size {
                        (dp_band[i][left_band_idx].f, dp_band[i][left_band_idx].h)
                    } else {
                        (Score::MIN / 2, 0)
                    };

                    let nuc_match = read[i - 1] == reference[j - 1];
                    (h_diag, e_above, h_above, f_left, h_left, nuc_match)
                };

                let cell = &mut dp_band[i][band_idx];

                if i == 0 || j == 0 {
                    // Boundary conditions
                    cell.h = 0;
                    cell.e = Score::MIN / 2;
                    cell.f = Score::MIN / 2;
                } else {
                    // Calculate E score (from above/insertion)
                    let e_gap = e_above + self.params.gap_extend;
                    let h_gap = h_above + self.params.gap_open + self.params.gap_extend;
                    cell.e = e_gap.max(h_gap).max(Score::MIN / 2);

                    // Calculate F score (from left/deletion)
                    let f_gap = f_left + self.params.gap_extend;
                    let h_prev = h_left;
                    cell.f = f_gap
                        .max(h_prev + self.params.gap_open + self.params.gap_extend)
                        .max(Score::MIN / 2);

                    // H cell: max of diagonal, E, F, and 0
                    // We already computed h_diag earlier, use it directly
                    cell.h = h_diag.max(cell.e).max(cell.f).max(0);
                }

                // Track best score
                if cell.h > best_score {
                    best_score = cell.h;
                    best_pos = (i, j);
                }
            }
        }

        // For banded alignment, estimate edits based on score
        let matches = (best_score as usize).div_ceil(self.params.match_score.max(1) as usize);
        let edits = best_pos.0.saturating_sub(matches);

        Ok(Alignment {
            ref_start: best_pos.1.saturating_sub(best_pos.0.min(best_pos.1)),
            ref_end: best_pos.1,
            score: best_score,
            strand: Strand::Forward,
            edits,
            ..Default::default()
        })
    }

    /// Simple traceback to count edits
    fn traceback_simple(
        &self,
        read: &[Nuc],
        reference: &[Nuc],
        best_pos: (usize, usize),
    ) -> (usize, usize) {
        let (mut i, mut j) = best_pos;
        let mut edits = 0;
        let mut ref_start = j;

        while i > 0 && j > 0 {
            let _match_score = if read[i - 1] == reference[j - 1] {
                self.params.match_score
            } else {
                self.params.mismatch_penalty
            };

            // For simplicity, just count mismatches along diagonal
            if read[i - 1] != reference[j - 1] {
                edits += 1;
            }

            ref_start = j;
            i -= 1;
            j -= 1;
        }

        (edits, ref_start)
    }

    /// Traceback from best position to reconstruct alignment
    #[allow(dead_code)]
    fn traceback_internal(
        dp_matrix: &[DpCell],
        read: &[Nuc],
        reference: &[Nuc],
        best_pos: (usize, usize),
        band_half: usize,
        params: &SwParams,
    ) -> (usize, usize, usize) {
        let (mut i, mut j) = best_pos;
        let mut edits = 0;
        let mut ref_start = j;

        // Simple traceback - count edits
        #[allow(clippy::implicit_saturating_sub)]
        while i > 0 && j > 0 {
            let row_start = if i > band_half { i - band_half } else { 0 };
            let band_idx = j - row_start;

            if band_idx >= dp_matrix.len() {
                break;
            }

            let cell = dp_matrix[band_idx];

            // Check if we came from diagonal
            let prev_band_idx = if j > row_start && band_idx > 0 {
                band_idx - 1
            } else {
                0
            };

            let match_score = if read[i - 1] == reference[j - 1] {
                params.match_score
            } else {
                params.mismatch_penalty
            };

            let h_diag = if prev_band_idx < dp_matrix.len() {
                dp_matrix[prev_band_idx].h + match_score
            } else {
                Score::MIN
            };

            if cell.h == h_diag && cell.h > 0 {
                // Came from diagonal
                if read[i - 1] != reference[j - 1] {
                    edits += 1;
                }
                i -= 1;
                j -= 1;
                ref_start = j;
            } else if cell.h == cell.e && cell.h > 0 {
                // Came from above (insertion in reference)
                edits += 1;
                i -= 1;
            } else if cell.h == cell.f && cell.h > 0 {
                // Came from left (deletion in reference)
                edits += 1;
                j -= 1;
                ref_start = j;
            } else {
                // Reached boundary
                break;
            }
        }

        (edits, ref_start, best_pos.1)
    }
}

impl Default for SwAligner {
    fn default() -> Self {
        Self::new(SwParams::default())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_params_default() {
        let params = SwParams::default();
        assert_eq!(params.match_score, 2);
        assert_eq!(params.mismatch_penalty, -3);
        assert_eq!(params.band_width, 15);
    }

    #[test]
    fn test_aligner_new() {
        let aligner = SwAligner::new(SwParams::default());
        assert_eq!(aligner.params.match_score, 2);
    }

    #[test]
    fn test_perfect_match() {
        let read = vec![Nuc::A, Nuc::C, Nuc::G, Nuc::T];
        let reference = vec![Nuc::A, Nuc::C, Nuc::G, Nuc::T];
        // Use larger band width for short sequences
        let params = SwParams {
            band_width: 10,
            min_score: 0, // Accept any positive score
            ..SwParams::default()
        };
        let aligner = SwAligner::new(params);

        let aln = aligner.align(&read, &reference).unwrap();
        // For local alignment, we should get at least 8
        assert!(aln.score >= 8);
        assert_eq!(aln.edits, 0);
    }

    #[test]
    fn test_one_mismatch() {
        let read = vec![Nuc::A, Nuc::C, Nuc::G, Nuc::T];
        let reference = vec![Nuc::A, Nuc::C, Nuc::A, Nuc::T]; // G->A mismatch
        let params = SwParams {
            band_width: 10,
            min_score: 0, // Accept any positive score
            ..SwParams::default()
        };
        let aligner = SwAligner::new(params);

        let aln = aligner.align(&read, &reference).unwrap();
        // Should find alignment with at least 3 matches (score >= 6)
        assert!(aln.score >= 3);
    }

    #[test]
    fn test_empty_sequence() {
        let read = vec![];
        let reference = vec![Nuc::A, Nuc::C, Nuc::G, Nuc::T];
        let aligner = SwAligner::new(SwParams::default());

        let aln = aligner.align(&read, &reference).unwrap();
        assert_eq!(aln.score, 0);
        assert_eq!(aln.edits, 0);
    }

    #[test]
    fn test_banded_alignment() {
        let read: Vec<Nuc> = (0..50)
            .map(|i| match i % 4 {
                0 => Nuc::A,
                1 => Nuc::C,
                2 => Nuc::G,
                _ => Nuc::T,
            })
            .collect();
        let mut reference = read.clone();
        // Add a few mutations
        reference[10] = Nuc::T; // was G
        reference[20] = Nuc::A; // was C
        reference[30] = Nuc::G; // was A

        let params = SwParams {
            band_width: 15, // Enable banding
            min_score: 0,
            ..SwParams::default()
        };

        let aligner = SwAligner::new(params);
        let aln = aligner.align(&read, &reference).unwrap();

        // Should find good alignment despite mutations
        assert!(aln.score > 80);
        assert!(aln.edits <= 5);
    }

    #[test]
    fn test_banded_vs_full_consistency() {
        let read: Vec<Nuc> = (0..30)
            .map(|i| match i % 4 {
                0 => Nuc::A,
                1 => Nuc::C,
                2 => Nuc::G,
                _ => Nuc::T,
            })
            .collect();
        let mut reference = read.clone();
        reference[15] = Nuc::T;

        // Banded alignment
        let params_banded = SwParams {
            band_width: 10,
            min_score: 0,
            ..SwParams::default()
        };
        let aligner_banded = SwAligner::new(params_banded);
        let aln_banded = aligner_banded.align(&read, &reference).unwrap();

        // Full alignment (no banding)
        let params_full = SwParams {
            band_width: 0, // Disable banding
            min_score: 0,
            ..SwParams::default()
        };
        let aligner_full = SwAligner::new(params_full);
        let aln_full = aligner_full.align(&read, &reference).unwrap();

        // Both should find similar alignments
        assert_eq!(aln_banded.score, aln_full.score);
    }

    #[test]
    fn test_banded_narrow_band() {
        let read: Vec<Nuc> = (0..100)
            .map(|i| match i % 4 {
                0 => Nuc::A,
                1 => Nuc::C,
                2 => Nuc::G,
                _ => Nuc::T,
            })
            .collect();
        let reference = read.clone();

        // Very narrow band - should still work for perfect match
        let params = SwParams {
            band_width: 5,
            min_score: 0,
            ..SwParams::default()
        };

        let aligner = SwAligner::new(params);
        let aln = aligner.align(&read, &reference).unwrap();

        // Perfect match should get max score
        assert_eq!(aln.score, 200); // 100 * 2
    }

    #[test]
    fn test_banded_disabled() {
        let read: Vec<Nuc> = (0..30)
            .map(|i| match i % 4 {
                0 => Nuc::A,
                1 => Nuc::C,
                2 => Nuc::G,
                _ => Nuc::T,
            })
            .collect();
        let reference = read.clone();

        // band_width = 0 means no banding
        let params = SwParams {
            band_width: 0,
            min_score: 0,
            ..SwParams::default()
        };

        let aligner = SwAligner::new(params);
        let aln = aligner.align(&read, &reference).unwrap();

        assert_eq!(aln.score, 60); // 30 * 2
    }
}