204 add separate equal cost range mode that ignores geometric tsm constraints

generic_a_star/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "generic_a_star"
 description = "A generic implementation of the A* algorithm"
-version = "5.0.1"
+version = "6.0.1"
 edition.workspace = true
 rust-version.workspace = true
 license.workspace = true

lib_ts_chainalign/Cargo.toml

@@ -2,15 +2,15 @@
 name = "lib_ts_chainalign"
 description = "A chaining-based sequence-to-sequence aligner that accounts for template switches"
 authors = ["Sebastian Schmidt <sebastian.schmidt@helsinki.fi>"]
-version = "5.0.1"
+version = "6.0.1"
 license.workspace = true
 edition.workspace = true
 rust-version.workspace = true
 repository.workspace = true
 
 [dependencies]
-generic_a_star = { version = "5.0.1", path = "../generic_a_star" }
-lib_tsalign = { version = "5.0.1", path = "../lib_tsalign" }
+generic_a_star = { version = "6.0.1", path = "../generic_a_star" }
+lib_tsalign = { version = "6.0.1", path = "../lib_tsalign" }
 ndarray = { version = "0.17.1", features = ["serde"] }
 num-traits.workspace = true
 serde = { workspace = true, features = ["derive"] }

lib_ts_chainalign/src/chain_align.rs

@@ -18,7 +18,7 @@ use itertools::Itertools;
 use lib_tsalign::a_star_aligner::{
     alignment_geometry::AlignmentRange,
     alignment_result::AlignmentResult,
-    template_switch_distance::{EqualCostRange, TemplateSwitchDirection},
+    template_switch_distance::{TSMUncertaintyRange, TemplateSwitchDirection},
 };
 use log::{debug, info, trace};
 use rustc_hash::FxHashMapSeed;
@@ -409,7 +409,7 @@ fn actually_align<
                         multiplicity,
                         TsAlignAlignmentType::TemplateSwitchEntrance {
                             first_offset: jump,
-                            equal_cost_range: EqualCostRange::new_invalid(),
+                            uncertainty_range: TSMUncertaintyRange::new_invalid(),
                             descendant: ts_kind.descendant.into_tsalign_primary(),
                             ancestor: ts_kind.ancestor.into_tsalign_secondary(),
                             direction: TemplateSwitchDirection::Reverse,

lib_tsalign/Cargo.toml

@@ -2,7 +2,7 @@
 name = "lib_tsalign"
 description = "A sequence-to-sequence aligner that accounts for template switches"
 authors = ["Sebastian Schmidt <sebastian.schmidt@helsinki.fi>"]
-version = "5.0.1"
+version = "6.0.1"
 license.workspace = true
 edition.workspace = true
 rust-version.workspace = true
@@ -26,9 +26,9 @@ thiserror.workspace = true
 num-traits.workspace = true
 serde = { workspace = true, features = ["derive"], optional = true }
 noisy_float = { version = "0.2.0" }
-generic_a_star = { version = "5.0.1", path = "../generic_a_star" }
+generic_a_star = { version = "6.0.1", path = "../generic_a_star" }
 log.workspace = true
 rustc-hash = "2.1.1"
-seed_chain = { version = "5.0.1", path = "../seed_chain" }
+seed_chain = { version = "6.0.1", path = "../seed_chain" }
 extend_map = "0.14.3"
 template-switch-error-free-inners = "0.1.1"

lib_tsalign/src/a_star_aligner.rs

@@ -23,7 +23,7 @@ use crate::{
     a_star_aligner::{
         alignment_result::alignment::Alignment,
         template_switch_distance::{
-            context::DynamicStrategies,
+            TSMUncertaintyRangeExtensionMode, context::DynamicStrategies,
             strategies::allow_ts_14_out_of_range::AdditionalExplicitTSMStartsAndEnds,
         },
     },
@@ -222,6 +222,7 @@ pub fn template_switch_distance_a_star_align<
     memory_limit: Option<usize>,
     force_label_correcting: bool,
     extend_beyond_range: bool,
+    tsm_uncertainty_range_extension_mode: TSMUncertaintyRangeExtensionMode,
     template_switch_count_memory: <Strategies::TemplateSwitchCount as TemplateSwitchCountStrategy>::Memory,
 ) -> AlignmentResult<template_switch_distance::AlignmentType, Strategies::Cost>
 where
@@ -276,7 +277,35 @@ where
         info!("Alignment ranges after extension {range}");
     }
 
-    info!("Extending template switches");
-    result.compute_ts_equal_cost_ranges(reference, query, &range, config);
+    match tsm_uncertainty_range_extension_mode {
+        TSMUncertaintyRangeExtensionMode::None => {
+            debug!("TSM uncertainty range extension is disabled.")
+        }
+
+        TSMUncertaintyRangeExtensionMode::EqualCost => {
+            info!("Extending TSM uncertainty ranges without increasing cost");
+            result.compute_ts_uncertainty_ranges(
+                reference,
+                query,
+                &range,
+                config,
+                tsm_uncertainty_range_extension_mode,
+            );
+        }
+
+        TSMUncertaintyRangeExtensionMode::EqualCostIgnoreGeometry => {
+            info!(
+                "Extending TSM uncertainty ranges without increasing cost but ignoring cost increases due to TSM geometry",
+            );
+            result.compute_ts_uncertainty_ranges(
+                reference,
+                query,
+                &range,
+                config,
+                tsm_uncertainty_range_extension_mode,
+            );
+        }
+    }
+
     result
 }

lib_tsalign/src/a_star_aligner/alignment_result.rs

@@ -4,13 +4,14 @@ use a_star_sequences::SequencePair;
 use alignment::Alignment;
 use compact_genome::interface::{alphabet::Alphabet, sequence::GenomeSequence};
 use generic_a_star::{AStarResult, cost::AStarCost};
-use log::{trace, warn};
+use log::{debug, trace, warn};
 use noisy_float::types::{R64, r64};
 use num_traits::{Float, Zero};
 
 use crate::{
     a_star_aligner::{
-        alignment_geometry::AlignmentCoordinates, template_switch_distance::AlignmentType,
+        alignment_geometry::AlignmentCoordinates,
+        template_switch_distance::{AlignmentType, TSMUncertaintyRangeExtensionMode},
     },
     config::TemplateSwitchConfig,
 };
@@ -286,7 +287,7 @@ impl<Cost: AStarCost + From<u64>>
         }
 
         // Compute cost before extending.
-        let mut current_cost = alignment.compute_cost(
+        let mut current_cost = alignment.compute_complete_cost(
             reference,
             query,
             range.reference_offset(),
@@ -325,7 +326,7 @@ impl<Cost: AStarCost + From<u64>>
             }
 
             // Compute cost.
-            let new_cost = alignment.compute_cost(
+            let new_cost = alignment.compute_complete_cost(
                 reference,
                 query,
                 new_range.reference_offset(),
@@ -380,7 +381,7 @@ impl<Cost: AStarCost + From<u64>>
             }
 
             // Compute cost.
-            let new_cost = alignment.compute_cost(
+            let new_cost = alignment.compute_complete_cost(
                 reference,
                 query,
                 new_range.reference_offset(),
@@ -413,7 +414,7 @@ impl<Cost: AStarCost + From<u64>>
     }
 
     #[allow(clippy::too_many_lines)]
-    pub fn compute_ts_equal_cost_ranges<
+    pub fn compute_ts_uncertainty_ranges<
         AlphabetType: Alphabet,
         SubsequenceType: GenomeSequence<AlphabetType, SubsequenceType> + ?Sized,
     >(
@@ -422,6 +423,7 @@ impl<Cost: AStarCost + From<u64>>
         query: &SubsequenceType,
         range: &AlignmentRange,
         config: &TemplateSwitchConfig<AlphabetType, Cost>,
+        mode: TSMUncertaintyRangeExtensionMode,
     ) {
         let Self::WithTarget {
             alignment,
@@ -431,26 +433,35 @@ impl<Cost: AStarCost + From<u64>>
             trace!("There is no alignment, therefore we cannot postprocess it.");
             return;
         };
+        if mode == TSMUncertaintyRangeExtensionMode::None {
+            debug!("TSM uncertainty range extension is disabled");
+            return;
+        }
         if config.left_flank_length > 0 || config.right_flank_length > 0 {
             warn!("TS extension does not support flanks");
             return;
         }
 
         let reference_offset = range.reference_offset();
         let query_offset = range.query_offset();
+        let ignore_geometry_cost = match mode {
+            TSMUncertaintyRangeExtensionMode::None => unreachable!(),
+            TSMUncertaintyRangeExtensionMode::EqualCost => false,
+            TSMUncertaintyRangeExtensionMode::EqualCostIgnoreGeometry => true,
+        };
 
         for i in 0..alignment.inner_mut().len() {
             let alignment_type = alignment.inner_mut()[i].1;
 
             if let super::template_switch_distance::AlignmentType::TemplateSwitchEntrance {
-                mut equal_cost_range,
+                mut uncertainty_range,
                 ..
             } = alignment_type
             {
-                equal_cost_range.min_start = 0;
-                equal_cost_range.max_start = 0;
-                equal_cost_range.min_end = 0;
-                equal_cost_range.max_end = 0;
+                uncertainty_range.min_start = 0;
+                uncertainty_range.max_start = 0;
+                uncertainty_range.min_end = 0;
+                uncertainty_range.max_end = 0;
 
                 // Keep track of the current cost, which might get less while extending but never increase
                 let mut current_cost = alignment.compute_cost(
@@ -459,8 +470,11 @@ impl<Cost: AStarCost + From<u64>>
                     reference_offset,
                     query_offset,
                     config,
+                    ignore_geometry_cost,
                 );
-                debug_assert_eq!(current_cost, (statistics.cost.round().raw() as u64).into());
+                if !ignore_geometry_cost {
+                    debug_assert_eq!(current_cost, (statistics.cost.round().raw() as u64).into());
+                }
 
                 // Move start to the left
                 {
@@ -480,6 +494,7 @@ impl<Cost: AStarCost + From<u64>>
                             reference_offset,
                             query_offset,
                             config,
+                            ignore_geometry_cost,
                         );
                         if new_cost > current_cost {
                             trace!(
@@ -488,7 +503,7 @@ impl<Cost: AStarCost + From<u64>>
                             break;
                         }
                         current_cost = new_cost;
-                        equal_cost_range.min_start -= 1;
+                        uncertainty_range.min_start -= 1;
                     }
                 }
 
@@ -510,6 +525,7 @@ impl<Cost: AStarCost + From<u64>>
                             reference_offset,
                             query_offset,
                             config,
+                            ignore_geometry_cost,
                         );
                         if new_cost > current_cost {
                             trace!(
@@ -518,7 +534,7 @@ impl<Cost: AStarCost + From<u64>>
                             break;
                         }
                         current_cost = new_cost;
-                        equal_cost_range.max_start += 1;
+                        uncertainty_range.max_start += 1;
                     }
                 }
 
@@ -538,6 +554,7 @@ impl<Cost: AStarCost + From<u64>>
                             reference_offset,
                             query_offset,
                             config,
+                            ignore_geometry_cost,
                         );
                         if new_cost > current_cost {
                             trace!(
@@ -546,7 +563,7 @@ impl<Cost: AStarCost + From<u64>>
                             break;
                         }
                         current_cost = new_cost;
-                        equal_cost_range.min_end -= 1;
+                        uncertainty_range.min_end -= 1;
                     }
                 }
 
@@ -566,6 +583,7 @@ impl<Cost: AStarCost + From<u64>>
                             reference_offset,
                             query_offset,
                             config,
+                            ignore_geometry_cost,
                         );
                         if new_cost > current_cost {
                             trace!(
@@ -574,18 +592,18 @@ impl<Cost: AStarCost + From<u64>>
                             break;
                         }
                         current_cost = new_cost;
-                        equal_cost_range.max_end += 1;
+                        uncertainty_range.max_end += 1;
                     }
                 }
 
                 let super::template_switch_distance::AlignmentType::TemplateSwitchEntrance {
-                    equal_cost_range: alignment_equal_cost_range,
+                    uncertainty_range: alignment_uncertainty_range,
                     ..
                 } = &mut alignment.inner_mut()[i].1
                 else {
                     unreachable!()
                 };
-                *alignment_equal_cost_range = equal_cost_range;
+                *alignment_uncertainty_range = uncertainty_range;
             }
         }
     }