Option to force all TSM descendants to be the same.

lib_tsalign/src/a_star_aligner.rs

@@ -6,7 +6,7 @@ use std::{
 use alignment_geometry::AlignmentRange;
 use alignment_result::{AlignmentResult, IAlignmentType};
 use compact_genome::interface::{alphabet::Alphabet, sequence::GenomeSequence};
-use generic_a_star::{AStar, AStarContext, AStarNode, AStarResult, cost::AStarCost};
+use generic_a_star::{AStar, AStarContext, AStarNode, AStarResult, cost::AStarCost, reset::Reset};
 use log::{debug, info};
 use template_switch_distance::{
     context::Memory,
@@ -19,7 +19,13 @@ use template_switch_distance::{
 };
 use traitsequence::interface::Sequence;
 
-use crate::{a_star_aligner::alignment_result::alignment::Alignment, config};
+use crate::{
+    a_star_aligner::{
+        alignment_result::alignment::Alignment,
+        template_switch_distance::context::DynamicStrategies,
+    },
+    config,
+};
 
 pub mod alignment_geometry;
 pub mod alignment_result;
@@ -186,6 +192,7 @@ pub fn template_switch_distance_a_star_align<
         Strategies::Alphabet,
         <Strategies as AlignmentStrategySelector>::Cost,
     >,
+    mut dynamic_strategies: DynamicStrategies,
     cost_limit: Option<Strategies::Cost>,
     memory_limit: Option<usize>,
     force_label_correcting: bool,
@@ -209,6 +216,7 @@ where
         shortcut: (),
         primary_match: (),
     };
+    dynamic_strategies.reset();
 
     info!("Calling aligner...");
     let mut result = a_star_align(
@@ -220,6 +228,7 @@ where
             range.clone(),
             config.clone(),
             memory,
+            dynamic_strategies,
             cost_limit,
             memory_limit,
             force_label_correcting,

lib_tsalign/src/a_star_aligner/configurable_a_star_align.rs

@@ -9,14 +9,21 @@ use traitsequence::interface::Sequence;
 
 use crate::{
     a_star_aligner::{
-        template_switch_distance::strategies::{
-            AlignmentStrategySelection,
-            primary_range::NoPrunePrimaryRangeStrategy,
-            secondary_deletion::AllowSecondaryDeletionStrategy,
-            shortcut::NoShortcutStrategy,
-            template_switch_min_length::{
-                PreprocessedLookaheadTemplateSwitchMinLengthStrategy,
-                PreprocessedTemplateSwitchMinLengthStrategy,
+        template_switch_distance::{
+            context::DynamicStrategies,
+            strategies::{
+                AlignmentStrategySelection,
+                descendant::{
+                    AnyTemplateSwitchDescendantStrategy, OnlyEqualTemplateSwitchDescendantStrategy,
+                    TemplateSwitchDescendantStrategy,
+                },
+                primary_range::NoPrunePrimaryRangeStrategy,
+                secondary_deletion::AllowSecondaryDeletionStrategy,
+                shortcut::NoShortcutStrategy,
+                template_switch_min_length::{
+                    PreprocessedLookaheadTemplateSwitchMinLengthStrategy,
+                    PreprocessedTemplateSwitchMinLengthStrategy,
+                },
             },
         },
         template_switch_distance_a_star_align,
@@ -86,6 +93,18 @@ pub enum TotalLengthStrategySelector {
     Maximise,
 }
 
+/// Select the descendant strategy.
+#[derive(Debug, Default)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+#[cfg_attr(feature = "serde", serde(rename_all = "snake_case"))]
+pub enum DescendantStrategySelector {
+    /// The TSM descendant can be any of the two input sequences for every TSM.
+    #[default]
+    AllowAny,
+    /// All TSMs must have the same descendant, but it can be either of the two input sequences.
+    AllowOnlyAllEqual,
+}
+
 /// Just used in this file to bundle the query parameters to make the code more readable
 struct QueryData<'a> {
     reference_name: &'a str,
@@ -107,6 +126,7 @@ pub struct Aligner<AlphabetType: Alphabet = DnaAlphabetOrN> {
     min_length_strategy: MinLengthStrategySelector,
     chaining_strategy: ChainingStrategySelector,
     total_length_strategy: TotalLengthStrategySelector,
+    descendant_strategy: DescendantStrategySelector,
     no_ts: bool,
 }
 
@@ -126,6 +146,7 @@ impl<AlphabetType: Alphabet> Default for Aligner<AlphabetType> {
             min_length_strategy: MinLengthStrategySelector::default(),
             chaining_strategy: ChainingStrategySelector::default(),
             total_length_strategy: TotalLengthStrategySelector::default(),
+            descendant_strategy: DescendantStrategySelector::default(),
             no_ts: false,
         }
     }
@@ -170,6 +191,14 @@ impl<AlphabetType: Alphabet> Aligner<AlphabetType> {
         self
     }
 
+    pub fn set_descendant_strategy(
+        &mut self,
+        descendant_strategy: DescendantStrategySelector,
+    ) -> &mut Self {
+        self.descendant_strategy = descendant_strategy;
+        self
+    }
+
     pub fn set_no_ts(&mut self, no_ts: bool) -> &mut Self {
         self.no_ts = no_ts;
         self
@@ -259,9 +288,39 @@ impl<AlphabetType: Alphabet> Aligner<AlphabetType> {
         data: QueryData,
     ) -> AlignmentResult<AlignmentType, U64Cost> {
         if self.no_ts {
-            self.align_call::<ML, CH, TL, MaxTemplateSwitchCountStrategy>(data, 0)
+            self.align_select_descendant_strategy::<ML, CH, TL, MaxTemplateSwitchCountStrategy>(
+                data, 0,
+            )
         } else {
-            self.align_call::<ML, CH, TL, NoTemplateSwitchCountStrategy>(data, ())
+            self.align_select_descendant_strategy::<ML, CH, TL, NoTemplateSwitchCountStrategy>(
+                data,
+                (),
+            )
+        }
+    }
+
+    fn align_select_descendant_strategy<
+        ML: TemplateSwitchMinLengthStrategy<U64Cost>,
+        CH: ChainingStrategy<U64Cost>,
+        TL: TemplateSwitchTotalLengthStrategy,
+        TC: TemplateSwitchCountStrategy,
+    >(
+        &self,
+        data: QueryData,
+        count_strategy_memory: TC::Memory,
+    ) -> AlignmentResult<AlignmentType, U64Cost> {
+        match self.descendant_strategy {
+            DescendantStrategySelector::AllowAny => self
+                .align_call::<ML, CH, TL, TC, AnyTemplateSwitchDescendantStrategy>(
+                    data,
+                    count_strategy_memory,
+                ),
+            DescendantStrategySelector::AllowOnlyAllEqual => {
+                self.align_call::<ML, CH, TL, TC, OnlyEqualTemplateSwitchDescendantStrategy>(
+                    data,
+                    count_strategy_memory,
+                )
+            }
         }
     }
 
@@ -270,6 +329,7 @@ impl<AlphabetType: Alphabet> Aligner<AlphabetType> {
         CH: ChainingStrategy<U64Cost>,
         TL: TemplateSwitchTotalLengthStrategy,
         TC: TemplateSwitchCountStrategy,
+        DS: TemplateSwitchDescendantStrategy,
     >(
         &self,
         data: QueryData,
@@ -292,6 +352,7 @@ impl<AlphabetType: Alphabet> Aligner<AlphabetType> {
                 AllowPrimaryMatchStrategy,
                 NoPrunePrimaryRangeStrategy,
                 TL,
+                DS,
             >,
             _,
         >(
@@ -302,6 +363,7 @@ impl<AlphabetType: Alphabet> Aligner<AlphabetType> {
             data.ranges
                 .unwrap_or_else(|| AlignmentRange::new_complete(reference.len(), query.len())),
             &self.costs,
+            DynamicStrategies {},
             cost_limit,
             data.memory_limit,
             false,

lib_tsalign/src/a_star_aligner/template_switch_distance.rs

@@ -23,7 +23,10 @@ pub use identifier::{
     GapType, Identifier, TemplateSwitchDirection, TemplateSwitchPrimary, TemplateSwitchSecondary,
 };
 
-use crate::config::BaseCost;
+use crate::{
+    a_star_aligner::template_switch_distance::strategies::descendant::TemplateSwitchDescendantStrategy,
+    config::BaseCost,
+};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Node<Strategies: AlignmentStrategySelector> {
@@ -251,6 +254,15 @@ impl<Strategies: AlignmentStrategySelector> Node<Strategies> {
                     unreachable!("This closure is only called on template switch entrances.")
                 };
 
+                // Check descendant restriction.
+                if !self
+                    .strategies
+                    .descendant_strategy
+                    .is_descendant_allowed(*template_switch_primary)
+                {
+                    return None;
+                }
+
                 let base_cost = base_cost.get(
                     *template_switch_primary,
                     *template_switch_secondary,

lib_tsalign/src/a_star_aligner/template_switch_distance/context.rs

@@ -40,6 +40,7 @@ pub struct Context<
 
     pub a_star_buffers: AStarBuffers<Box<Node<Strategies>>>,
     pub memory: Memory<Strategies>,
+    pub dynamic_strategies: DynamicStrategies,
 
     cost_limit: Option<Strategies::Cost>,
     memory_limit: Option<usize>,
@@ -60,6 +61,8 @@ pub struct Memory<Strategies: AlignmentStrategySelector> {
 >>::Memory,
 }
 
+pub struct DynamicStrategies {}
+
 impl<
     'reference,
     'query,
@@ -76,6 +79,7 @@ impl<
         range: AlignmentRange,
         config: TemplateSwitchConfig<Strategies::Alphabet, Strategies::Cost>,
         memory: Memory<Strategies>,
+        dynamic_strategies: DynamicStrategies,
         cost_limit: Option<Strategies::Cost>,
         memory_limit: Option<usize>,
         force_label_correcting: bool,
@@ -89,6 +93,7 @@ impl<
             config,
             a_star_buffers: Default::default(),
             memory,
+            dynamic_strategies,
             cost_limit,
             memory_limit,
             force_label_correcting,
@@ -780,6 +785,7 @@ impl<
 {
     fn reset(&mut self) {
         self.memory.reset();
+        self.dynamic_strategies.reset();
     }
 }
 
@@ -789,6 +795,10 @@ impl<Strategies: AlignmentStrategySelector> Reset for Memory<Strategies> {
     }
 }
 
+impl Reset for DynamicStrategies {
+    fn reset(&mut self) {}
+}
+
 impl<
     SubsequenceType: GenomeSequence<Strategies::Alphabet, SubsequenceType> + ?Sized,
     Strategies: AlignmentStrategySelector,

lib_tsalign/src/a_star_aligner/template_switch_distance/lower_bounds/template_switch.rs

@@ -18,11 +18,12 @@ use crate::{
         alignment_result::IAlignmentType,
         template_switch_distance::{
             Context, Identifier, Node,
-            context::Memory,
+            context::{DynamicStrategies, Memory},
             identifier::GapType,
             strategies::{
                 AlignmentStrategySelection, chaining::NoChainingStrategy,
-                node_ord::CostOnlyNodeOrdStrategy, primary_match::AllowPrimaryMatchStrategy,
+                descendant::AnyTemplateSwitchDescendantStrategy, node_ord::CostOnlyNodeOrdStrategy,
+                primary_match::AllowPrimaryMatchStrategy,
                 primary_range::NoPrunePrimaryRangeStrategy,
                 secondary_deletion::ForbidSecondaryDeletionStrategy, shortcut::NoShortcutStrategy,
                 template_switch_count::MaxTemplateSwitchCountStrategy,
@@ -62,6 +63,7 @@ type TSLBAlignmentStrategies<AlphabetType, Cost> = AlignmentStrategySelection<
     AllowPrimaryMatchStrategy,
     NoPrunePrimaryRangeStrategy,
     NoTemplateSwitchTotalLengthStrategy,
+    AnyTemplateSwitchDescendantStrategy,
 >;
 
 impl<Cost: AStarCost> TemplateSwitchLowerBoundMatrix<Cost> {
@@ -106,6 +108,7 @@ impl<Cost: AStarCost> TemplateSwitchLowerBoundMatrix<Cost> {
                     shortcut: (),
                     primary_match: (),
                 },
+                DynamicStrategies {},
                 None,
                 None,
                 false,

lib_tsalign/src/a_star_aligner/template_switch_distance/lower_bounds/template_switch_alignment.rs

@@ -19,10 +19,11 @@ use crate::{
         alignment_result::IAlignmentType,
         template_switch_distance::{
             Context, Identifier,
-            context::Memory,
+            context::{DynamicStrategies, Memory},
             strategies::{
                 AlignmentStrategySelection,
                 chaining::NoChainingStrategy,
+                descendant::AnyTemplateSwitchDescendantStrategy,
                 node_ord::CostOnlyNodeOrdStrategy,
                 primary_match::{
                     MaxConsecutivePrimaryMatchMemory, MaxConsecutivePrimaryMatchStrategy,
@@ -60,6 +61,7 @@ type TSALBAlignmentStrategies<AlphabetType, Cost> = AlignmentStrategySelection<
     MaxConsecutivePrimaryMatchStrategy,
     NoPrunePrimaryRangeStrategy,
     NoTemplateSwitchTotalLengthStrategy,
+    AnyTemplateSwitchDescendantStrategy,
 >;
 
 impl<Cost: AStarCost> TemplateSwitchAlignmentLowerBoundMatrix<Cost> {
@@ -113,6 +115,7 @@ impl<Cost: AStarCost> TemplateSwitchAlignmentLowerBoundMatrix<Cost> {
                             .min_substitution_cost(),
                     },
                 },
+                DynamicStrategies {},
                 None,
                 None,
                 false,