Preserve complete and ordered DMR segmentation

modkit-core/src/dmr/single_site.rs

@@ -808,6 +808,14 @@ fn collapse_counts(
 }
 
 type ChromToSingleScores = (String, Vec<MkResult<SingleSiteDmrScore>>);
+
+fn sort_chrom_to_site_scores(chrom_to_site_scores: &mut [ChromToSingleScores]) {
+    // The batch maps use hash iteration order, while the writer and stateful
+    // segmenter require the lexical contig order used by SingleSiteBatches.
+    chrom_to_site_scores
+        .sort_unstable_by(|(a_chrom, _), (b_chrom, _)| a_chrom.cmp(b_chrom));
+}
+
 fn process_batch_of_positions(
     batch: DmrBatchOfPositions,
     sample_index: Arc<SingleSiteSampleIndex>,
@@ -816,7 +824,7 @@ fn process_batch_of_positions(
     let (a_lines, b_lines) =
         sample_index.read_bedmethyl_lines_organized_by_position(batch)?;
 
-    let chrom_to_site_scores = a_lines
+    let mut chrom_to_site_scores = a_lines
         .into_iter()
         // intersect a_lines and b_lines on contig/chrom, there should be a
         // filter upstream of this to make sure that this is not ever a miss
@@ -852,9 +860,32 @@ fn process_batch_of_positions(
         })
         .collect::<Vec<ChromToSingleScores>>();
 
+    sort_chrom_to_site_scores(&mut chrom_to_site_scores);
+
     Ok(chrom_to_site_scores)
 }
 
+#[cfg(test)]
+mod chrom_score_order_tests {
+    use super::{sort_chrom_to_site_scores, ChromToSingleScores};
+
+    #[test]
+    fn chrom_scores_are_sorted_lexically() {
+        let mut chrom_to_site_scores = ["cc", "aa", "bb"]
+            .into_iter()
+            .map(|chrom| (chrom.to_string(), Vec::new()))
+            .collect::<Vec<ChromToSingleScores>>();
+
+        sort_chrom_to_site_scores(&mut chrom_to_site_scores);
+
+        let chroms = chrom_to_site_scores
+            .iter()
+            .map(|(chrom, _)| chrom.as_str())
+            .collect::<Vec<_>>();
+        assert_eq!(chroms, vec!["aa", "bb", "cc"]);
+    }
+}
+
 struct Coverages {
     a_coverages: Vec<u64>,
     b_coverages: Vec<u64>,
@@ -1319,7 +1350,7 @@ fn path_to_region_labels(
     path: &[States],
     positions: &[u64],
 ) -> Vec<(u64, u64, States)> {
-    assert_eq!(path.len(), positions.len() - 1);
+    assert_eq!(path.len(), positions.len());
     if path.is_empty() {
         return Vec::new();
     } else {
@@ -1345,3 +1376,75 @@ fn path_to_region_labels(
         agg
     }
 }
+
+#[cfg(test)]
+mod segmentation_path_tests {
+    use super::{path_to_region_labels, States};
+
+    fn assert_every_position_covered_once(
+        path: &[States],
+        positions: &[u64],
+        regions: &[(u64, u64, States)],
+    ) {
+        for (&position, &state) in positions.iter().zip(path) {
+            let covering_regions = regions
+                .iter()
+                .filter(|(start, end, _)| *start <= position && position < *end)
+                .collect::<Vec<_>>();
+            assert_eq!(
+                covering_regions.len(),
+                1,
+                "position {position} in state {state:?} was not covered exactly once"
+            );
+            assert_eq!(covering_regions[0].2, state);
+        }
+    }
+
+    fn check_case(
+        path: &[States],
+        positions: &[u64],
+        expected: &[(u64, u64, States)],
+    ) {
+        let regions = path_to_region_labels(path, positions);
+        assert_eq!(regions, expected);
+        assert_every_position_covered_once(path, positions, &regions);
+    }
+
+    #[test]
+    fn singleton_position_becomes_a_single_base_region() {
+        check_case(&[States::Same], &[10], &[(10, 11, States::Same)]);
+    }
+
+    #[test]
+    fn two_positions_are_both_integrated() {
+        check_case(
+            &[States::Same, States::Same],
+            &[10, 20],
+            &[(10, 21, States::Same)],
+        );
+        check_case(
+            &[States::Same, States::Different],
+            &[10, 20],
+            &[(10, 11, States::Same), (20, 21, States::Different)],
+        );
+    }
+
+    #[test]
+    fn multi_position_transitions_cover_each_position_once() {
+        check_case(
+            &[
+                States::Same,
+                States::Same,
+                States::Different,
+                States::Different,
+                States::Same,
+            ],
+            &[10, 20, 30, 40, 50],
+            &[
+                (10, 21, States::Same),
+                (30, 41, States::Different),
+                (50, 51, States::Same),
+            ],
+        );
+    }
+}

modkit-core/src/hmm.rs

@@ -187,7 +187,7 @@ impl HmmModel {
         assert_eq!(probs.len(), transitions.len());
         let (dp_matrix, pointers) = self.viterbi_forward(&probs, &transitions);
         let path = self.viterbi_decode(&dp_matrix, &pointers);
-        assert_eq!(path.len(), scores.len() - 1);
+        assert_eq!(path.len(), scores.len());
         path
     }
 
@@ -197,21 +197,16 @@ impl HmmModel {
         pointers: &[PointerCell],
     ) -> Vec<States> {
         let final_state = dp_matrix.last().unwrap().argmax();
-        // dbg!(final_state);
         let mut path = vec![final_state];
-        let mut curr_pointer =
-            pointers.last().unwrap().get_value(final_state).unwrap();
-        for pointers in pointers.iter().rev().skip(1) {
-            let pointer = pointers.get_value(curr_pointer);
-            if let Some(pointer) = pointer {
-                path.push(pointer);
-                curr_pointer = pointer;
-            } else {
-                break;
-            }
+        let mut current_state = final_state;
+        // The first pointer cell is the empty start cell, and the second
+        // points back to the un-emitted start state. Decode only the emitted
+        // states, from the final score back through the second score.
+        for pointers in pointers.iter().skip(2).rev() {
+            current_state = pointers.get_value(current_state).unwrap();
+            path.push(current_state);
         }
 
-        path.pop();
         path.reverse();
         path
     }
@@ -413,12 +408,141 @@ impl Projection {
 
 #[cfg(test)]
 mod hmm_tests {
-    use crate::hmm::HmmModel;
+    use crate::hmm::{HmmModel, States};
+
+    fn test_model() -> HmmModel {
+        HmmModel::new(0.1, 0.9, 0.3, -0.1, 0.01, 500, true).unwrap()
+    }
+
+    fn emission_score(model: &HmmModel, p: f64, state: States) -> f64 {
+        let p = if p == 0f64 { 1e-5 } else { p };
+        let (factor, log_probability) = match state {
+            States::Same => (model.same_state_factor, p.ln()),
+            States::Different => {
+                (model.diff_state_factor, (1f64 - p + 1e-5).ln())
+            }
+        };
+        factor * (log_probability - model.significance_factor)
+    }
+
+    fn transition_score(
+        model: &HmmModel,
+        previous: States,
+        current: States,
+        diff_stay: f64,
+    ) -> f64 {
+        match (previous, current) {
+            (States::Same, States::Same) => model.same_to_same,
+            (States::Same, States::Different) => model.same_to_diff,
+            (States::Different, States::Same) => (1f64 - diff_stay).ln(),
+            (States::Different, States::Different) => diff_stay.ln(),
+        }
+    }
+
+    /// Exhaustively score the hidden start state and every emitted state.
+    /// This deliberately does not use the dynamic-programming matrix or its
+    /// back-pointers, so it independently checks both state order and length.
+    fn brute_force_path(
+        model: &HmmModel,
+        scores: &[f64],
+        positions: &[u64],
+    ) -> Vec<States> {
+        assert!(!scores.is_empty());
+        assert_eq!(scores.len(), positions.len());
+
+        let probabilities = scores
+            .iter()
+            .map(|&score| (-score.max(0f64)).exp())
+            .collect::<Vec<_>>();
+        let diff_stays = positions.windows(2).fold(
+            vec![model.dmr_prior],
+            |mut transitions, window| {
+                let gap = (window[1] - window[0]) as f64;
+                transitions.push(if model.linear_proj {
+                    model.projection.linear_project_prob(gap)
+                } else {
+                    model.projection.ln_project_prob(gap)
+                });
+                transitions
+            },
+        );
+
+        let state_count = scores.len() + 1;
+        let mut best: Option<(f64, Vec<States>)> = None;
+        for encoded_path in 0..(1usize << state_count) {
+            let states = (0..state_count)
+                .map(|i| {
+                    if encoded_path & (1usize << i) == 0 {
+                        States::Same
+                    } else {
+                        States::Different
+                    }
+                })
+                .collect::<Vec<_>>();
+            let initial_score = match states[0] {
+                States::Same => model.same_to_same,
+                States::Different => model.same_to_diff,
+            };
+            let total_score = probabilities
+                .iter()
+                .zip(diff_stays.iter())
+                .enumerate()
+                .fold(initial_score, |total, (i, (&p, &diff_stay))| {
+                    total
+                        + transition_score(
+                            model,
+                            states[i],
+                            states[i + 1],
+                            diff_stay,
+                        )
+                        + emission_score(model, p, states[i + 1])
+                });
+
+            if best
+                .as_ref()
+                .map(|(best_score, _)| total_score > *best_score)
+                .unwrap_or(true)
+            {
+                best = Some((total_score, states[1..].to_vec()));
+            }
+        }
+        best.unwrap().1
+    }
 
     #[test]
     fn test_prob_to_factor() {
         let sig_fact = 0.01;
         let fact = HmmModel::prob_to_factor(sig_fact).unwrap();
         dbg!(fact);
     }
+
+    #[test]
+    fn viterbi_path_matches_independent_oracle_for_tiny_sequences() {
+        let model = test_model();
+        let cases = [
+            (vec![0.0], vec![10]),
+            (vec![12.0], vec![10]),
+            (vec![0.0, 12.0], vec![10, 20]),
+            (vec![12.0, 0.0], vec![10, 20]),
+            (vec![0.0, 0.0, 12.0, 12.0, 0.0], vec![10, 20, 30, 40, 50]),
+        ];
+
+        for (scores, positions) in cases {
+            let expected = brute_force_path(&model, &scores, &positions);
+            let actual = model.viterbi_path(&scores, &positions);
+            assert_eq!(actual.len(), scores.len());
+            assert_eq!(actual, expected, "scores: {scores:?}");
+        }
+    }
+
+    #[test]
+    fn viterbi_path_preserves_state_transitions_in_order() {
+        let model = test_model();
+        let scores = vec![0.0, 0.0, 12.0, 12.0, 0.0];
+        let positions = vec![10, 20, 30, 40, 50];
+        let expected = brute_force_path(&model, &scores, &positions);
+        assert!(expected.windows(2).any(|states| states[0] != states[1]));
+
+        assert_eq!(model.viterbi_path(&scores, &positions), expected);
+    }
 }