Support base-cost per ts_kind in tschainalign.

lib_ts_chainalign/src/alignment/ts_kind.rs

@@ -23,18 +23,22 @@ pub enum TsDescendant {
 }
 
 impl TsKind {
+    /// Ancestor is Seq1, descendant is Seq1.
     pub const TS11: Self = TsKind {
         ancestor: TsAncestor::Seq1,
         descendant: TsDescendant::Seq1,
     };
+    /// Ancestor is Seq1, descendant is Seq2.
     pub const TS12: Self = TsKind {
         ancestor: TsAncestor::Seq1,
         descendant: TsDescendant::Seq2,
     };
+    /// Ancestor is Seq2, descendant is Seq1.
     pub const TS21: Self = TsKind {
         ancestor: TsAncestor::Seq2,
         descendant: TsDescendant::Seq1,
     };
+    /// Ancestor is Seq2, descendant is Seq2.
     pub const TS22: Self = TsKind {
         ancestor: TsAncestor::Seq2,
         descendant: TsDescendant::Seq2,

lib_ts_chainalign/src/chaining_cost_function.rs

@@ -286,12 +286,18 @@ impl<Cost: AStarCost> ChainingCostFunction<Cost> {
                         let lower_bound = chaining_lower_bounds.jump_12_lower_bound(gap);
                         jump_12[[from_index, to_index]] = lower_bound.max(
                             max_exact_cost_function_cost
-                                + chaining_lower_bounds.alignment_costs().ts_base_cost
+                                + chaining_lower_bounds
+                                    .alignment_costs()
+                                    .ts_base_cost
+                                    .get(ts_kind)
                                 + Cost::from_usize(1),
                         );
                         if lower_bound
                             <= max_exact_cost_function_cost
-                                + chaining_lower_bounds.alignment_costs().ts_base_cost
+                                + chaining_lower_bounds
+                                    .alignment_costs()
+                                    .ts_base_cost
+                                    .get(ts_kind)
                         {
                             eligible_anchors.push(to_anchor);
                         }
@@ -330,7 +336,10 @@ impl<Cost: AStarCost> ChainingCostFunction<Cost> {
                             align_end,
                             ts_kind,
                             max_exact_cost_function_cost
-                                + chaining_lower_bounds.alignment_costs().ts_base_cost,
+                                + chaining_lower_bounds
+                                    .alignment_costs()
+                                    .ts_base_cost
+                                    .get(ts_kind),
                             &mut additional_secondary_targets_output,
                         );
                     let after_align = Instant::now();
@@ -344,7 +353,10 @@ impl<Cost: AStarCost> ChainingCostFunction<Cost> {
                         jump_12[[from_index, to_index]] = jump_12[[from_index, to_index]].min(cost);
                         if cost
                             <= max_exact_cost_function_cost
-                                + chaining_lower_bounds.alignment_costs().ts_base_cost
+                                + chaining_lower_bounds
+                                    .alignment_costs()
+                                    .ts_base_cost
+                                    .get(ts_kind)
                         {
                             jump_12.set_exact(from_index, to_index);
                         }
@@ -442,7 +454,11 @@ impl<Cost: AStarCost> ChainingCostFunction<Cost> {
                 start,
                 sequences.secondary_end(ts_kind),
                 ts_kind,
-                max_exact_cost_function_cost + chaining_lower_bounds.alignment_costs().ts_base_cost,
+                max_exact_cost_function_cost
+                    + chaining_lower_bounds
+                        .alignment_costs()
+                        .ts_base_cost
+                        .get(ts_kind),
                 &mut additional_secondary_targets_output,
             );
             additional_secondary_targets_output.sort_unstable();
@@ -453,7 +469,10 @@ impl<Cost: AStarCost> ChainingCostFunction<Cost> {
                 jump_12[[primary_start_anchor_index, to_index]] = cost;
                 if cost
                     <= max_exact_cost_function_cost
-                        + chaining_lower_bounds.alignment_costs().ts_base_cost
+                        + chaining_lower_bounds
+                            .alignment_costs()
+                            .ts_base_cost
+                            .get(ts_kind)
                 {
                     jump_12.set_exact(primary_start_anchor_index, to_index);
                 }
@@ -466,7 +485,10 @@ impl<Cost: AStarCost> ChainingCostFunction<Cost> {
                     .jump_12_lower_bound(gap)
                     .max(
                         max_exact_cost_function_cost
-                            + chaining_lower_bounds.alignment_costs().ts_base_cost
+                            + chaining_lower_bounds
+                                .alignment_costs()
+                                .ts_base_cost
+                                .get(ts_kind)
                             + Cost::from_usize(1),
                     )
                     .min(jump_12[[primary_start_anchor_index, index]]);

lib_ts_chainalign/src/chaining_cost_function/tests.rs

@@ -28,7 +28,7 @@ fn create_lower_bounds(max_match_run: u32) -> ChainingLowerBounds<U32Cost> {
         AlignmentCosts::new(
             GapAffineCosts::new(2u8.into(), 3u8.into(), 1u8.into()),
             GapAffineCosts::new(4u8.into(), 6u8.into(), 2u8.into()),
-            2u8.into(),
+            U32Cost::from(2u8).into(),
             TsLimits::new_unlimited(),
         ),
     )

lib_ts_chainalign/src/chaining_lower_bounds/ts_jump.rs

@@ -36,7 +36,7 @@ impl<Cost: AStarCost> TsJumpLowerBounds<Cost> {
             for secondary_descendant_gap in 0..=max_n - primary_descendant_gap {
                 let lower_bound = primary_lower_bounds
                     .variable_gap2_lower_bound(primary_descendant_gap)
-                    + cost_table.ts_base_cost
+                    + cost_table.ts_base_cost.min()
                     + secondary_lower_bounds.variable_gap2_lower_bound(secondary_descendant_gap);
                 lower_bounds_12[[primary_descendant_gap + secondary_descendant_gap]] =
                     lower_bounds_12[[primary_descendant_gap + secondary_descendant_gap]]

lib_ts_chainalign/src/chaining_lower_bounds/ts_jump/tests.rs

@@ -18,7 +18,7 @@ fn test_max_match_run_0() {
             gap_open: U32Cost::from(6u8),
             gap_extend: U32Cost::from(2u8),
         },
-        ts_base_cost: U32Cost::from(2u8),
+        ts_base_cost: U32Cost::from(2u8).into(),
         ts_limits: TsLimits {
             inter_jump_12: -100..100,
             intra_jump_12: -100..100,
@@ -33,7 +33,7 @@ fn test_max_match_run_0() {
 
     let expected_lower_bounds_12 = [2, 4, 6];
     let expected_lower_bounds_34 =
-        expected_lower_bounds_12.map(|cost| cost - cost_table.ts_base_cost.as_primitive());
+        expected_lower_bounds_12.map(|cost| cost - cost_table.ts_base_cost.min().as_primitive());
 
     for descendant_gap in 0..=max_n {
         assert_eq!(
@@ -64,7 +64,7 @@ fn test_max_match_run_1() {
             gap_open: U32Cost::from(6u8),
             gap_extend: U32Cost::from(2u8),
         },
-        ts_base_cost: U32Cost::from(2u8),
+        ts_base_cost: U32Cost::from(2u8).into(),
         ts_limits: TsLimits {
             inter_jump_12: -100..100,
             intra_jump_12: -100..100,
@@ -79,7 +79,7 @@ fn test_max_match_run_1() {
 
     let expected_lower_bounds_12 = [2, 2, 2, 4, 4, 6, 6, 8, 8];
     let expected_lower_bounds_34 =
-        expected_lower_bounds_12.map(|cost| cost - cost_table.ts_base_cost.as_primitive());
+        expected_lower_bounds_12.map(|cost| cost - cost_table.ts_base_cost.min().as_primitive());
 
     for descendant_gap in 0..=max_n {
         assert_eq!(
@@ -110,7 +110,7 @@ fn test_max_match_run_2() {
             gap_open: U32Cost::from(6u8),
             gap_extend: U32Cost::from(2u8),
         },
-        ts_base_cost: U32Cost::from(2u8),
+        ts_base_cost: U32Cost::from(2u8).into(),
         ts_limits: TsLimits {
             inter_jump_12: -100..100,
             intra_jump_12: -100..100,
@@ -125,7 +125,7 @@ fn test_max_match_run_2() {
 
     let expected_lower_bounds_12 = [2, 2, 2, 2, 2, 4, 4, 4, 6, 6];
     let expected_lower_bounds_34 =
-        expected_lower_bounds_12.map(|cost| cost - cost_table.ts_base_cost.as_primitive());
+        expected_lower_bounds_12.map(|cost| cost - cost_table.ts_base_cost.min().as_primitive());
 
     for descendant_gap in 0..=max_n {
         assert_eq!(

lib_ts_chainalign/src/costs.rs

@@ -1,8 +1,10 @@
 use std::ops::Range;
 
-use num_traits::Zero;
+use num_traits::{Zero, bounds::UpperBounded};
 use serde::{Deserialize, Serialize};
 
+use crate::alignment::ts_kind::TsKind;
+
 mod compat;
 
 #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
@@ -30,6 +32,13 @@ pub struct TsLimits {
     pub ancestor_gap: Range<isize>,
 }
 
+/// The base cost of a template switch.
+/// This is applied whenever a template switch is started.
+#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
+pub struct TsBaseCost<Cost> {
+    cost_by_kind: [Cost; 4],
+}
+
 /// The cost function for alignments.
 ///
 /// For convenience, it implements [`TryFrom<TemplateSwitchConfig<Alphabet, Cost>>`](std::convert::TryFrom).
@@ -41,8 +50,8 @@ pub struct AlignmentCosts<Cost> {
     /// Costs for secondary alignment, i.e. for the 23-alignment of a template switch.
     pub secondary_costs: GapAffineCosts<Cost>,
     /// The base cost of a template switch.
-    /// This is applied whenevera template switch is started.
-    pub ts_base_cost: Cost,
+    /// This is applied whenever a template switch is started.
+    pub ts_base_cost: TsBaseCost<Cost>,
     /// Limits on the geometry of a template switch.
     pub ts_limits: TsLimits,
 }
@@ -75,11 +84,60 @@ impl TsLimits {
     }
 }
 
+impl<Cost> TsBaseCost<Cost> {
+    pub fn new(cost_by_kind: [Cost; 4]) -> Self {
+        Self { cost_by_kind }
+    }
+
+    pub fn has_zero_cost(&self) -> bool
+    where
+        Cost: Zero,
+    {
+        self.cost_by_kind.iter().any(|cost| cost.is_zero())
+    }
+
+    pub fn min(&self) -> Cost
+    where
+        Cost: Copy + Ord,
+    {
+        *self.cost_by_kind.iter().min().unwrap()
+    }
+
+    pub fn get(&self, ts_kind: TsKind) -> Cost
+    where
+        Cost: Copy,
+    {
+        self.cost_by_kind[ts_kind.index()]
+    }
+}
+
+impl<Cost: UpperBounded + Eq> FromIterator<(TsKind, Cost)> for TsBaseCost<Cost> {
+    fn from_iter<T: IntoIterator<Item = (TsKind, Cost)>>(iter: T) -> Self {
+        let mut cost_by_kind = [None, None, None, None];
+
+        for (ts_kind, cost) in iter {
+            assert!(
+                cost_by_kind[ts_kind.index()].is_none(),
+                "Duplicate TsKind {ts_kind} in iterator",
+            );
+            cost_by_kind[ts_kind.index()] = Some(cost);
+        }
+
+        Self::new(cost_by_kind.map(|cost| cost.expect("Missing TsKind")))
+    }
+}
+
+impl<Cost: Copy> From<Cost> for TsBaseCost<Cost> {
+    fn from(value: Cost) -> Self {
+        Self::new([value; 4])
+    }
+}
+
 impl<Cost> AlignmentCosts<Cost> {
     pub fn new(
         primary_costs: GapAffineCosts<Cost>,
         secondary_costs: GapAffineCosts<Cost>,
-        ts_base_cost: Cost,
+        ts_base_cost: TsBaseCost<Cost>,
         ts_limits: TsLimits,
     ) -> Self {
         Self {
@@ -95,6 +153,6 @@ impl<Cost: Zero> AlignmentCosts<Cost> {
     pub fn has_zero_cost(&self) -> bool {
         self.primary_costs.has_zero_cost()
             || self.secondary_costs.has_zero_cost()
-            || self.ts_base_cost.is_zero()
+            || self.ts_base_cost.has_zero_cost()
     }
 }

lib_ts_chainalign/src/costs/compat.rs

@@ -7,7 +7,10 @@ use lib_tsalign::{
     config::TemplateSwitchConfig, costs::gap_affine::GapAffineAlignmentCostTable,
 };
 
-use crate::costs::{AlignmentCosts, GapAffineCosts, TsLimits};
+use crate::{
+    alignment::ts_kind::TsKind,
+    costs::{AlignmentCosts, GapAffineCosts, TsBaseCost, TsLimits},
+};
 
 impl<AlphabetType: Alphabet, Cost: AStarCost> From<TemplateSwitchConfig<AlphabetType, Cost>>
     for AlignmentCosts<Cost>
@@ -16,10 +19,12 @@ impl<AlphabetType: Alphabet, Cost: AStarCost> From<TemplateSwitchConfig<Alphabet
         let value = &value;
         assert!(value.left_flank_length == 0 && value.right_flank_length == 0);
 
-        let ts_base_cost = value.base_cost.qqr;
-        assert_eq!(ts_base_cost, value.base_cost.qrr);
-        assert_eq!(ts_base_cost, value.base_cost.rqr);
-        assert_eq!(ts_base_cost, value.base_cost.rrr);
+        let ts_base_cost = TsBaseCost::from_iter([
+            (TsKind::TS11, value.base_cost.rrr),
+            (TsKind::TS12, value.base_cost.qrr), // Note that the ancestor is the secondary and the descendant is the primary, hence these are flipped.
+            (TsKind::TS21, value.base_cost.rqr),
+            (TsKind::TS22, value.base_cost.qqr),
+        ]);
 
         Self {
             primary_costs: (&value.primary_edit_costs).into(),

lib_ts_chainalign/src/exact_chaining/ts_12_jump/algo.rs

@@ -88,7 +88,7 @@ impl<Cost: AStarCost> AStarContext for Context<'_, '_, '_, Cost> {
             },
             predecessor: None,
             predecessor_alignment_type: None,
-            cost: self.costs.ts_base_cost,
+            cost: self.costs.ts_base_cost.get(self.end.ts_kind().unwrap()),
             match_run: 0,
         }
     }