[Affine] Change to new transformation function

RagnarGrootKoerkamp · Dec 31, 2023 · de316f0 · de316f0
1 parent 941f3cb
commit de316f0
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 25 deletions.
diff --git a/pa-affine-constants/src/lib.rs b/pa-affine-constants/src/lib.rs
@@ -3,3 +3,8 @@ use std::cell::RefCell;
 
 thread_local!(pub static INDEL_COST: RefCell<Cost> = RefCell::new(3));
 thread_local!(pub static SUB_COST: RefCell<Cost> = RefCell::new(2));
+thread_local!(pub static R: RefCell<Cost> = INDEL_COST.with(|indel_cost| {
+    SUB_COST.with(|sub_cost| {
+        RefCell::new(std::cmp::min(*indel_cost.borrow(), *sub_cost.borrow()))
+    })
+}));
diff --git a/pa-heuristic/src/matches/exact.rs b/pa-heuristic/src/matches/exact.rs
@@ -11,7 +11,7 @@ use super::*;
 use crate::prelude::*;
 use smallvec::SmallVec;
 
-use pa_affine_constants::{INDEL_COST, SUB_COST};
+use pa_affine_constants::R;
 
 /// Build a hashset of the seeds in a, and query all kmers in b.
 pub fn hash_a<'a>(a: Seq<'a>, b: Seq<'a>, config: MatchConfig, transform_filter: bool) -> Matches {
@@ -62,10 +62,7 @@ fn hash_to_smallvec(
                     start,
                     end: start + Pos(k, k),
                     match_cost: 0,
-                    seed_potential: min(
-                        INDEL_COST.with(|indel_cost| *indel_cost.borrow()),
-                        SUB_COST.with(|sub_cost| *sub_cost.borrow()),
-                    ) as MatchCost,
+                    seed_potential: R.with(|r| *r.borrow()) as MatchCost,
                     pruned: MatchStatus::Active,
                 });
             }

diff --git a/pa-heuristic/src/matches/qgrams.rs b/pa-heuristic/src/matches/qgrams.rs
@@ -3,7 +3,7 @@ use itertools::izip;
 use super::*;
 use crate::prelude::*;
 
-use pa_affine_constants::{INDEL_COST, SUB_COST};
+use pa_affine_constants::R;
 
 // NOTE: This assumes an alphabet of 'ACGT'.
 pub struct QGrams<'a> {
@@ -104,16 +104,10 @@ impl<'a> QGrams<'a> {
             .map(|i| Seed {
                 start: i as I,
                 end: i as I + k,
-                seed_potential: min(
-                    INDEL_COST.with(|indel_cost| *indel_cost.borrow()),
-                    SUB_COST.with(|sub_cost| *sub_cost.borrow()),
-                ) as MatchCost
-                    * r,
-                seed_cost: min(
-                    INDEL_COST.with(|indel_cost| *indel_cost.borrow()),
-                    SUB_COST.with(|sub_cost| *sub_cost.borrow()),
-                ) as MatchCost
-                    * r,
+                // TODO: Use the `r` parameter instead and make the required changes in `matches.rs`.
+                // This assumes exact matches only.
+                seed_potential: R.with(|r| *r.borrow()) as MatchCost,
+                seed_cost: R.with(|r| *r.borrow()) as MatchCost,
             })
             .collect()
     }

diff --git a/pa-heuristic/src/seeds.rs b/pa-heuristic/src/seeds.rs
@@ -2,7 +2,7 @@ use itertools::Itertools;
 
 use crate::prelude::*;
 
-use pa_affine_constants::{INDEL_COST, SUB_COST};
+use pa_affine_constants::{INDEL_COST, R};
 
 /// Type for the cost of a single match/mutation.
 pub type MatchCost = u8;
@@ -139,11 +139,11 @@ impl Seeds {
     #[inline]
     pub fn transform(&self, pos @ Pos(i, j): Pos) -> Pos {
         let p = self.potential(pos);
-        INDEL_COST.with(|indel_cost| {
-            SUB_COST.with(|sub_cost| {
+        R.with(|r| {
+            INDEL_COST.with(|indel_cost| {
                 Pos(
-                    min(*indel_cost.borrow(), *sub_cost.borrow()) * (i - j) - p,
-                    min(*indel_cost.borrow(), *sub_cost.borrow()) * (j - i) - p,
+                    (i - j) - (*r.borrow()).div_ceil(*indel_cost.borrow()) * p,
+                    (j - i) - (*r.borrow()).div_ceil(*indel_cost.borrow()) * p,
                 )
             })
         })
@@ -154,11 +154,12 @@ impl Seeds {
         if pos == Pos(I::MAX, I::MAX) {
             return pos;
         }
-        let p = -(x + y) / 2;
-        let i = self.start_of_potential[p as usize];
-        let diff = INDEL_COST.with(|indel_cost| {
-            SUB_COST.with(|sub_cost| (x - y) / (min(*indel_cost.borrow(), *sub_cost.borrow()) * 2))
+        let p = R.with(|r| {
+            INDEL_COST
+                .with(|indel_cost| -(x + y) / (2 * (*r.borrow()).div_ceil(*indel_cost.borrow())))
         });
+        let i = self.start_of_potential[p as usize];
+        let diff = (x - y) / 2;
         let j = i - diff;
         debug_assert_eq!(pos, self.transform(Pos(i, j)));
         Pos(i, j)