rna_algos/
mccaskill_algo.rs

1use utils::*;
2
3pub struct FoldSums<T: Hash> {
4  pub sums_external: SumMat,
5  pub sums_rightmost_basepairs_external: SumMat,
6  pub sums_rightmost_basepairs_multibranch: SumMat,
7  pub sums_close: SparseSumMat<T>,
8  pub sums_accessible: SparseSumMat<T>,
9  pub sums_multibranch: SumMat,
10  pub sums_1ormore_basepairs: SumMat,
11}
12
13#[derive(Clone)]
14pub struct FoldScores<T: Hash> {
15  pub hairpin_scores: SparseScoreMat<T>,
16  pub twoloop_scores: ScoreMat4d<T>,
17  pub multibranch_close_scores: SparseScoreMat<T>,
18  pub accessible_scores: SparseScoreMat<T>,
19}
20
21pub type ScoreMat4d<T> = HashMap<PosQuad<T>, Score>;
22pub type SparseScoreMat<T> = HashMap<PosPair<T>, Score>;
23
24impl FoldScoreSets {
25  pub fn new(init_val: Score) -> FoldScoreSets {
26    let init_vals = [init_val; NUM_BASES];
27    let mat_2d = [init_vals; NUM_BASES];
28    let mat_3d = [mat_2d; NUM_BASES];
29    let mat_4d = [mat_3d; NUM_BASES];
30    FoldScoreSets {
31      // The CONTRAfold model.
32      hairpin_scores_len: [init_val; MAX_LOOP_LEN + 1],
33      bulge_scores_len: [init_val; MAX_LOOP_LEN],
34      interior_scores_len: [init_val; MAX_LOOP_LEN - 1],
35      interior_scores_symmetric: [init_val; MAX_INTERIOR_SYMMETRIC],
36      interior_scores_asymmetric: [init_val; MAX_INTERIOR_ASYMMETRIC],
37      stack_scores: mat_4d,
38      terminal_mismatch_scores: mat_4d,
39      dangling_scores_left: mat_3d,
40      dangling_scores_right: mat_3d,
41      helix_close_scores: mat_2d,
42      basepair_scores: mat_2d,
43      interior_scores_explicit: [[init_val; MAX_INTERIOR_EXPLICIT]; MAX_INTERIOR_EXPLICIT],
44      bulge_scores_0x1: [init_val; NUM_BASES],
45      interior_scores_1x1: [[init_val; NUM_BASES]; NUM_BASES],
46      multibranch_score_base: init_val,
47      multibranch_score_basepair: init_val,
48      multibranch_score_unpair: init_val,
49      external_score_basepair: init_val,
50      external_score_unpair: init_val,
51      // The cumulative parameters of the CONTRAfold model.
52      hairpin_scores_len_cumulative: [init_val; MAX_LOOP_LEN + 1],
53      bulge_scores_len_cumulative: [init_val; MAX_LOOP_LEN],
54      interior_scores_len_cumulative: [init_val; MAX_LOOP_LEN - 1],
55      interior_scores_symmetric_cumulative: [init_val; MAX_INTERIOR_SYMMETRIC],
56      interior_scores_asymmetric_cumulative: [init_val; MAX_INTERIOR_ASYMMETRIC],
57    }
58  }
59
60  pub fn accumulate(&mut self) {
61    let mut sum = 0.;
62    for i in 0..self.hairpin_scores_len_cumulative.len() {
63      sum += self.hairpin_scores_len[i];
64      self.hairpin_scores_len_cumulative[i] = sum;
65    }
66    let mut sum = 0.;
67    for i in 0..self.bulge_scores_len_cumulative.len() {
68      sum += self.bulge_scores_len[i];
69      self.bulge_scores_len_cumulative[i] = sum;
70    }
71    let mut sum = 0.;
72    for i in 0..self.interior_scores_len_cumulative.len() {
73      sum += self.interior_scores_len[i];
74      self.interior_scores_len_cumulative[i] = sum;
75    }
76    let mut sum = 0.;
77    for i in 0..self.interior_scores_symmetric_cumulative.len() {
78      sum += self.interior_scores_symmetric[i];
79      self.interior_scores_symmetric_cumulative[i] = sum;
80    }
81    let mut sum = 0.;
82    for i in 0..self.interior_scores_asymmetric_cumulative.len() {
83      sum += self.interior_scores_asymmetric[i];
84      self.interior_scores_asymmetric_cumulative[i] = sum;
85    }
86  }
87
88  pub fn transfer(&mut self) {
89    for (v, &w) in self
90      .hairpin_scores_len
91      .iter_mut()
92      .zip(HAIRPIN_SCORES_LEN_ATLEAST.iter())
93    {
94      *v = w;
95    }
96    for (v, &w) in self
97      .bulge_scores_len
98      .iter_mut()
99      .zip(BULGE_SCORES_LEN_ATLEAST.iter())
100    {
101      *v = w;
102    }
103    for (v, &w) in self
104      .interior_scores_len
105      .iter_mut()
106      .zip(INTERIOR_SCORES_LEN_ATLEAST.iter())
107    {
108      *v = w;
109    }
110    for (v, &w) in self
111      .interior_scores_symmetric
112      .iter_mut()
113      .zip(INTERIOR_SCORES_SYMMETRIC_ATLEAST.iter())
114    {
115      *v = w;
116    }
117    for (v, &w) in self
118      .interior_scores_asymmetric
119      .iter_mut()
120      .zip(INTERIOR_SCORES_ASYMMETRIC_ATLEAST.iter())
121    {
122      *v = w;
123    }
124    for (i, x) in STACK_SCORES_CONTRA.iter().enumerate() {
125      for (j, x) in x.iter().enumerate() {
126        if !has_canonical_basepair(&(i, j)) {
127          continue;
128        }
129        for (k, x) in x.iter().enumerate() {
130          for (l, &x) in x.iter().enumerate() {
131            if !has_canonical_basepair(&(k, l)) {
132              continue;
133            }
134            self.stack_scores[i][j][k][l] = x;
135          }
136        }
137      }
138    }
139    for (i, x) in TERMINAL_MISMATCH_SCORES_CONTRA.iter().enumerate() {
140      for (j, x) in x.iter().enumerate() {
141        if !has_canonical_basepair(&(i, j)) {
142          continue;
143        }
144        for (k, x) in x.iter().enumerate() {
145          for (l, &x) in x.iter().enumerate() {
146            self.terminal_mismatch_scores[i][j][k][l] = x;
147          }
148        }
149      }
150    }
151    for (i, x) in DANGLING_SCORES_LEFT.iter().enumerate() {
152      for (j, x) in x.iter().enumerate() {
153        if !has_canonical_basepair(&(i, j)) {
154          continue;
155        }
156        for (k, &x) in x.iter().enumerate() {
157          self.dangling_scores_left[i][j][k] = x;
158        }
159      }
160    }
161    for (i, x) in DANGLING_SCORES_RIGHT.iter().enumerate() {
162      for (j, x) in x.iter().enumerate() {
163        if !has_canonical_basepair(&(i, j)) {
164          continue;
165        }
166        for (k, &x) in x.iter().enumerate() {
167          self.dangling_scores_right[i][j][k] = x;
168        }
169      }
170    }
171    for (i, x) in HELIX_CLOSE_SCORES.iter().enumerate() {
172      for (j, &x) in x.iter().enumerate() {
173        if !has_canonical_basepair(&(i, j)) {
174          continue;
175        }
176        self.helix_close_scores[i][j] = x;
177      }
178    }
179    for (i, x) in BASEPAIR_SCORES.iter().enumerate() {
180      for (j, &x) in x.iter().enumerate() {
181        if !has_canonical_basepair(&(i, j)) {
182          continue;
183        }
184        self.basepair_scores[i][j] = x;
185      }
186    }
187    for (i, x) in INTERIOR_SCORES_EXPLICIT.iter().enumerate() {
188      for (j, &x) in x.iter().enumerate() {
189        self.interior_scores_explicit[i][j] = x;
190      }
191    }
192    for (x, &y) in self
193      .bulge_scores_0x1
194      .iter_mut()
195      .zip(BULGE_SCORES_0X1.iter())
196    {
197      *x = y;
198    }
199    for (i, x) in INTERIOR_SCORES_1X1_CONTRA.iter().enumerate() {
200      for (j, &x) in x.iter().enumerate() {
201        self.interior_scores_1x1[i][j] = x;
202      }
203    }
204    self.multibranch_score_base = MULTIBRANCH_SCORE_BASE;
205    self.multibranch_score_basepair = MULTIBRANCH_SCORE_BASEPAIR;
206    self.multibranch_score_unpair = MULTIBRANCH_SCORE_UNPAIR;
207    self.external_score_basepair = EXTERNAL_SCORE_BASEPAIR;
208    self.external_score_unpair = EXTERNAL_SCORE_UNPAIR;
209    self.accumulate();
210  }
211}
212
213impl<T: Hash> FoldSums<T> {
214  pub fn new(seq_len: usize) -> FoldSums<T> {
215    let neg_infs = vec![vec![NEG_INFINITY; seq_len]; seq_len];
216    FoldSums {
217      sums_external: vec![vec![0.; seq_len]; seq_len],
218      sums_rightmost_basepairs_external: neg_infs.clone(),
219      sums_rightmost_basepairs_multibranch: neg_infs.clone(),
220      sums_close: SparseSumMat::<T>::default(),
221      sums_accessible: SparseSumMat::<T>::default(),
222      sums_multibranch: neg_infs.clone(),
223      sums_1ormore_basepairs: neg_infs,
224    }
225  }
226}
227
228impl<T: HashIndex> Default for FoldScores<T> {
229  fn default() -> Self {
230    Self::new()
231  }
232}
233
234impl<T: HashIndex> FoldScores<T> {
235  pub fn new() -> FoldScores<T> {
236    let scores = SparseScoreMat::<T>::default();
237    let scores_4d = ScoreMat4d::<T>::default();
238    FoldScores {
239      hairpin_scores: scores.clone(),
240      twoloop_scores: scores_4d,
241      multibranch_close_scores: scores.clone(),
242      accessible_scores: scores,
243    }
244  }
245}
246
247pub fn mccaskill_algo<T>(
248  seq: SeqSlice,
249  uses_contra_model: bool,
250  allows_short_hairpins: bool,
251  fold_score_sets: &FoldScoreSets,
252) -> (SparseProbMat<T>, FoldScores<T>)
253where
254  T: HashIndex,
255{
256  let seq_len = seq.len();
257  let mut fold_scores = FoldScores::<T>::new();
258  let fold_sums = if uses_contra_model {
259    get_fold_sums_contra::<T>(
260      seq,
261      &mut fold_scores,
262      allows_short_hairpins,
263      fold_score_sets,
264    )
265  } else {
266    get_fold_sums::<T>(seq, &mut fold_scores)
267  };
268  let basepair_probs = if uses_contra_model {
269    get_basepair_probs_contra::<T>(
270      &fold_sums,
271      seq_len,
272      &fold_scores,
273      allows_short_hairpins,
274      fold_score_sets,
275    )
276  } else {
277    get_basepair_probs::<T>(&fold_sums, seq_len, &fold_scores)
278  };
279  (basepair_probs, fold_scores)
280}
281
282pub fn get_fold_sums<T>(seq: SeqSlice, fold_scores: &mut FoldScores<T>) -> FoldSums<T>
283where
284  T: HashIndex,
285{
286  let seq_len = seq.len();
287  let uses_sentinel_bases = false;
288  let mut fold_sums = FoldSums::<T>::new(seq_len);
289  let seq_len = T::from_usize(seq.len()).unwrap();
290  for subseq_len in range_inclusive(T::from_usize(MIN_SPAN_HAIRPIN_CLOSE).unwrap(), seq_len) {
291    for i in range_inclusive(T::zero(), seq_len - subseq_len) {
292      let j = i + subseq_len - T::one();
293      let (long_i, long_j) = (i.to_usize().unwrap(), j.to_usize().unwrap());
294      let pos_pair_close = (i, j);
295      let long_pos_pair_close = (long_i, long_j);
296      let basepair_close = (seq[long_i], seq[long_j]);
297      let mut sum = NEG_INFINITY;
298      if long_pos_pair_close.1 - long_pos_pair_close.0 + 1 >= MIN_SPAN_HAIRPIN_CLOSE
299        && has_canonical_basepair(&basepair_close)
300      {
301        let hairpin_score = get_hairpin_score(seq, &long_pos_pair_close);
302        fold_scores
303          .hairpin_scores
304          .insert(pos_pair_close, hairpin_score);
305        logsumexp(&mut sum, hairpin_score);
306        for k in range(i + T::one(), j - T::one()) {
307          let long_k = k.to_usize().unwrap();
308          if long_k - long_i - 1 > MAX_2LOOP_LEN {
309            break;
310          }
311          for l in range(k + T::one(), j).rev() {
312            let long_l = l.to_usize().unwrap();
313            if long_j - long_l - 1 + long_k - long_i - 1 > MAX_2LOOP_LEN {
314              break;
315            }
316            let pos_pair_accessible = (k, l);
317            let long_pos_pair_accessible = (long_k, long_l);
318            if let Some(&x) = fold_sums.sums_close.get(&pos_pair_accessible) {
319              let y = get_2loop_score(seq, &long_pos_pair_close, &long_pos_pair_accessible);
320              fold_scores.twoloop_scores.insert((i, j, k, l), y);
321              let y = x + y;
322              logsumexp(&mut sum, y);
323            }
324          }
325        }
326        let multibranch_close_score = get_multibranch_close_score(seq, &long_pos_pair_close);
327        logsumexp(
328          &mut sum,
329          fold_sums.sums_multibranch[long_i + 1][long_j - 1] + multibranch_close_score,
330        );
331        let accessible_score = get_accessible_score(seq, &long_pos_pair_close, uses_sentinel_bases);
332        if sum > NEG_INFINITY {
333          fold_scores
334            .multibranch_close_scores
335            .insert(pos_pair_close, multibranch_close_score);
336          fold_scores
337            .accessible_scores
338            .insert(pos_pair_close, accessible_score);
339          fold_sums.sums_close.insert(pos_pair_close, sum);
340          let sum = sum + accessible_score;
341          fold_sums.sums_accessible.insert(pos_pair_close, sum);
342        }
343      }
344      sum = NEG_INFINITY;
345      for k in range_inclusive(i + T::one(), j) {
346        let pos_pair_accessible = (i, k);
347        if let Some(&x) = fold_sums.sums_accessible.get(&pos_pair_accessible) {
348          logsumexp(&mut sum, x);
349        }
350      }
351      fold_sums.sums_rightmost_basepairs_external[long_i][long_j] = sum;
352      sum = 0.;
353      for k in long_i..long_j {
354        let x = fold_sums.sums_rightmost_basepairs_external[k][long_j];
355        let y = if long_i == 0 && k == 0 {
356          0.
357        } else {
358          fold_sums.sums_external[long_i][k - 1]
359        };
360        let y = x + y;
361        logsumexp(&mut sum, y);
362      }
363      fold_sums.sums_external[long_i][long_j] = sum;
364      sum = fold_sums.sums_rightmost_basepairs_external[long_i][long_j] + COEFF_NUM_BRANCHES;
365      let mut sum2 = NEG_INFINITY;
366      for k in long_i + 1..long_j {
367        let x = fold_sums.sums_rightmost_basepairs_external[k][long_j] + COEFF_NUM_BRANCHES;
368        logsumexp(&mut sum, x);
369        let y = fold_sums.sums_1ormore_basepairs[long_i][k - 1] + x;
370        logsumexp(&mut sum2, y);
371      }
372      fold_sums.sums_multibranch[long_i][long_j] = sum2;
373      logsumexp(&mut sum, sum2);
374      fold_sums.sums_1ormore_basepairs[long_i][long_j] = sum;
375    }
376  }
377  fold_sums
378}
379
380pub fn get_fold_sums_contra<T>(
381  seq: SeqSlice,
382  fold_scores: &mut FoldScores<T>,
383  allows_short_hairpins: bool,
384  fold_score_sets: &FoldScoreSets,
385) -> FoldSums<T>
386where
387  T: HashIndex,
388{
389  let seq_len = seq.len();
390  let uses_sentinel_bases = false;
391  let mut fold_sums = FoldSums::<T>::new(seq_len);
392  let short_seq_len = T::from_usize(seq.len()).unwrap();
393  for subseq_len in range_inclusive(T::one(), short_seq_len) {
394    for i in range_inclusive(T::zero(), short_seq_len - subseq_len) {
395      let j = i + subseq_len - T::one();
396      let (long_i, long_j) = (i.to_usize().unwrap(), j.to_usize().unwrap());
397      let pos_pair_close = (i, j);
398      let long_pos_pair_close = (long_i, long_j);
399      let basepair_close = (seq[long_i], seq[long_j]);
400      let mut sum = NEG_INFINITY;
401      if has_canonical_basepair(&basepair_close)
402        && (allows_short_hairpins
403          || long_pos_pair_close.1 - long_pos_pair_close.0 + 1 >= MIN_SPAN_HAIRPIN_CLOSE)
404      {
405        if long_j - long_i - 1 <= MAX_LOOP_LEN {
406          let hairpin_score = get_hairpin_score_contra(seq, &long_pos_pair_close, fold_score_sets);
407          fold_scores
408            .hairpin_scores
409            .insert(pos_pair_close, hairpin_score);
410          logsumexp(&mut sum, hairpin_score);
411        }
412        for k in range(i + T::one(), j - T::one()) {
413          let long_k = k.to_usize().unwrap();
414          if long_k - long_i - 1 > MAX_LOOP_LEN {
415            break;
416          }
417          for l in range(k + T::one(), j).rev() {
418            let long_l = l.to_usize().unwrap();
419            if long_j - long_l - 1 + long_k - long_i - 1 > MAX_LOOP_LEN {
420              break;
421            }
422            let pos_pair_accessible = (k, l);
423            let long_pos_pair_accessible = (long_k, long_l);
424            if let Some(&x) = fold_sums.sums_close.get(&pos_pair_accessible) {
425              let y = get_2loop_score_contra(
426                seq,
427                &long_pos_pair_close,
428                &long_pos_pair_accessible,
429                fold_score_sets,
430              );
431              fold_scores.twoloop_scores.insert((i, j, k, l), y);
432              let y = x + y;
433              logsumexp(&mut sum, y);
434            }
435          }
436        }
437        let multibranch_close_score = fold_score_sets.multibranch_score_base
438          + fold_score_sets.multibranch_score_basepair
439          + get_junction_score(
440            seq,
441            &long_pos_pair_close,
442            uses_sentinel_bases,
443            fold_score_sets,
444          );
445        logsumexp(
446          &mut sum,
447          fold_sums.sums_multibranch[long_i + 1][long_j - 1] + multibranch_close_score,
448        );
449        let accessible_score = get_junction_score(
450          seq,
451          &(long_pos_pair_close.1, long_pos_pair_close.0),
452          uses_sentinel_bases,
453          fold_score_sets,
454        ) + fold_score_sets.basepair_scores[basepair_close.0]
455          [basepair_close.1];
456        if sum > NEG_INFINITY {
457          fold_scores
458            .multibranch_close_scores
459            .insert(pos_pair_close, multibranch_close_score);
460          fold_scores
461            .accessible_scores
462            .insert(pos_pair_close, accessible_score);
463          fold_sums.sums_close.insert(pos_pair_close, sum);
464          let sum = sum + accessible_score;
465          fold_sums.sums_accessible.insert(pos_pair_close, sum);
466        }
467      }
468      sum = NEG_INFINITY;
469      let mut sum2 = sum;
470      for k in range_inclusive(i + T::one(), j) {
471        let pos_pair_accessible = (i, k);
472        if let Some(&x) = fold_sums.sums_accessible.get(&pos_pair_accessible) {
473          logsumexp(
474            &mut sum,
475            x + fold_score_sets.external_score_basepair
476              + fold_score_sets.external_score_unpair * (j - k).to_f32().unwrap(),
477          );
478          logsumexp(
479            &mut sum2,
480            x + fold_score_sets.multibranch_score_basepair
481              + fold_score_sets.multibranch_score_unpair * (j - k).to_f32().unwrap(),
482          );
483        }
484      }
485      fold_sums.sums_rightmost_basepairs_external[long_i][long_j] = sum;
486      fold_sums.sums_rightmost_basepairs_multibranch[long_i][long_j] = sum2;
487      sum = fold_score_sets.external_score_unpair * subseq_len.to_f32().unwrap();
488      for k in long_i..long_j {
489        let x = fold_sums.sums_rightmost_basepairs_external[k][long_j];
490        let y = if long_i == 0 && k == 0 {
491          0.
492        } else {
493          fold_sums.sums_external[long_i][k - 1]
494        };
495        let y = x + y;
496        logsumexp(&mut sum, y);
497      }
498      fold_sums.sums_external[long_i][long_j] = sum;
499      sum = fold_sums.sums_rightmost_basepairs_multibranch[long_i][long_j];
500      sum2 = NEG_INFINITY;
501      for k in long_i + 1..long_j {
502        let x = fold_sums.sums_rightmost_basepairs_multibranch[k][long_j];
503        logsumexp(
504          &mut sum,
505          x + fold_score_sets.multibranch_score_unpair * (k - long_i) as Score,
506        );
507        let x = fold_sums.sums_1ormore_basepairs[long_i][k - 1] + x;
508        logsumexp(&mut sum2, x);
509      }
510      fold_sums.sums_multibranch[long_i][long_j] = sum2;
511      logsumexp(&mut sum, sum2);
512      fold_sums.sums_1ormore_basepairs[long_i][long_j] = sum;
513    }
514  }
515  fold_sums
516}
517
518pub fn get_basepair_probs<T>(
519  fold_sums: &FoldSums<T>,
520  seq_len: usize,
521  fold_scores: &FoldScores<T>,
522) -> SparseProbMat<T>
523where
524  T: HashIndex,
525{
526  let global_sum = fold_sums.sums_external[0][seq_len - 1];
527  let mut basepair_probs = SparseProbMat::<T>::default();
528  let mut probs_multibranch = vec![vec![NEG_INFINITY; seq_len]; seq_len];
529  let mut probs_multibranch2 = probs_multibranch.clone();
530  let short_seq_len = T::from_usize(seq_len).unwrap();
531  for subseq_len in range_inclusive(
532    T::from_usize(MIN_SPAN_HAIRPIN_CLOSE).unwrap(),
533    short_seq_len,
534  )
535  .rev()
536  {
537    for i in range_inclusive(T::zero(), short_seq_len - subseq_len) {
538      let j = i + subseq_len - T::one();
539      let (long_i, long_j) = (i.to_usize().unwrap(), j.to_usize().unwrap());
540      let mut sum = NEG_INFINITY;
541      let mut sum2 = sum;
542      for k in range(j + T::one(), short_seq_len) {
543        let long_k = k.to_usize().unwrap();
544        let pos_pair_close = (i, k);
545        if let Some(&x) = fold_sums.sums_close.get(&pos_pair_close) {
546          let basepair_prob = basepair_probs[&pos_pair_close];
547          let multibranch_close_score = fold_scores.multibranch_close_scores[&pos_pair_close];
548          let x = basepair_prob + multibranch_close_score - x;
549          logsumexp(
550            &mut sum,
551            x + fold_sums.sums_1ormore_basepairs[long_j + 1][long_k - 1],
552          );
553          logsumexp(&mut sum2, x);
554        }
555      }
556      probs_multibranch[long_i][long_j] = sum;
557      probs_multibranch2[long_i][long_j] = sum2;
558      let pos_pair_accessible = (i, j);
559      if let Some(&sum_close) = fold_sums.sums_close.get(&pos_pair_accessible) {
560        let sum_accessible = fold_sums.sums_accessible[&pos_pair_accessible];
561        let sum_pair = (
562          if pos_pair_accessible.0 < T::one() {
563            0.
564          } else {
565            fold_sums.sums_external[0][long_i - 1]
566          },
567          if pos_pair_accessible.1 > short_seq_len - T::from_usize(2).unwrap() {
568            0.
569          } else {
570            fold_sums.sums_external[long_j + 1][seq_len - 1]
571          },
572        );
573        let mut sum = sum_pair.0 + sum_accessible + sum_pair.1 - global_sum;
574        for k in range(T::zero(), i).rev() {
575          let long_k = k.to_usize().unwrap();
576          if long_i - long_k - 1 > MAX_2LOOP_LEN {
577            break;
578          }
579          for l in range(j + T::one(), short_seq_len) {
580            let long_l = l.to_usize().unwrap();
581            if long_l - long_j - 1 + long_i - long_k - 1 > MAX_2LOOP_LEN {
582              break;
583            }
584            let pos_pair_close = (k, l);
585            if let Some(&x) = fold_sums.sums_close.get(&pos_pair_close) {
586              logsumexp(
587                &mut sum,
588                basepair_probs[&pos_pair_close] + sum_close - x
589                  + fold_scores.twoloop_scores[&(k, l, i, j)],
590              );
591            }
592          }
593        }
594        let sum_accessible = sum_accessible + COEFF_NUM_BRANCHES;
595        for k in 0..long_i {
596          let x = fold_sums.sums_1ormore_basepairs[k + 1][long_i - 1];
597          logsumexp(&mut sum, sum_accessible + probs_multibranch2[k][long_j] + x);
598          let y = probs_multibranch[k][long_j];
599          logsumexp(&mut sum, sum_accessible + y);
600          logsumexp(&mut sum, sum_accessible + x + y);
601        }
602        if sum > NEG_INFINITY {
603          basepair_probs.insert(pos_pair_accessible, sum);
604        }
605      }
606    }
607  }
608  basepair_probs = basepair_probs.iter().map(|(x, &y)| (*x, expf(y))).collect();
609  basepair_probs
610}
611
612pub fn get_basepair_probs_contra<T>(
613  fold_sums: &FoldSums<T>,
614  seq_len: usize,
615  fold_scores: &FoldScores<T>,
616  allows_short_hairpins: bool,
617  fold_score_sets: &FoldScoreSets,
618) -> SparseProbMat<T>
619where
620  T: HashIndex,
621{
622  let global_sum = fold_sums.sums_external[0][seq_len - 1];
623  let mut basepair_probs = SparseProbMat::<T>::default();
624  let mut probs_multibranch = vec![vec![NEG_INFINITY; seq_len]; seq_len];
625  let mut probs_multibranch2 = probs_multibranch.clone();
626  let short_seq_len = T::from_usize(seq_len).unwrap();
627  for subseq_len in range_inclusive(
628    T::from_usize(if allows_short_hairpins {
629      2
630    } else {
631      MIN_SPAN_HAIRPIN_CLOSE
632    })
633    .unwrap(),
634    short_seq_len,
635  )
636  .rev()
637  {
638    for i in range_inclusive(T::zero(), short_seq_len - subseq_len) {
639      let j = i + subseq_len - T::one();
640      let (long_i, long_j) = (i.to_usize().unwrap(), j.to_usize().unwrap());
641      let mut sum = NEG_INFINITY;
642      let mut sum2 = sum;
643      for k in range(j + T::one(), short_seq_len) {
644        let long_k = k.to_usize().unwrap();
645        let pos_pair_close = (i, k);
646        if let Some(&x) = fold_sums.sums_close.get(&pos_pair_close) {
647          let basepair_prob = basepair_probs[&pos_pair_close];
648          let multibranch_close_score = fold_scores.multibranch_close_scores[&pos_pair_close];
649          let x = basepair_prob + multibranch_close_score - x;
650          logsumexp(
651            &mut sum,
652            x + fold_sums.sums_1ormore_basepairs[long_j + 1][long_k - 1],
653          );
654          logsumexp(
655            &mut sum2,
656            x + fold_score_sets.multibranch_score_unpair * (k - j - T::one()).to_f32().unwrap(),
657          );
658        }
659      }
660      probs_multibranch[long_i][long_j] = sum;
661      probs_multibranch2[long_i][long_j] = sum2;
662      let pos_pair_accessible = (i, j);
663      if let Some(&sum_close) = fold_sums.sums_close.get(&pos_pair_accessible) {
664        let sum_pair = (
665          if pos_pair_accessible.0 < T::one() {
666            0.
667          } else {
668            fold_sums.sums_external[0][long_i - 1]
669          },
670          if pos_pair_accessible.1 > short_seq_len - T::from_usize(2).unwrap() {
671            0.
672          } else {
673            fold_sums.sums_external[long_j + 1][seq_len - 1]
674          },
675        );
676        let mut sum = sum_pair.0
677          + sum_pair.1
678          + fold_sums.sums_accessible[&pos_pair_accessible]
679          + fold_score_sets.external_score_basepair
680          - global_sum;
681        for k in range(T::zero(), i).rev() {
682          let long_k = k.to_usize().unwrap();
683          if long_i - long_k - 1 > MAX_LOOP_LEN {
684            break;
685          }
686          for l in range(j + T::one(), short_seq_len) {
687            let long_l = l.to_usize().unwrap();
688            if long_l - long_j - 1 + long_i - long_k - 1 > MAX_LOOP_LEN {
689              break;
690            }
691            let pos_pair_close = (k, l);
692            if let Some(&x) = fold_sums.sums_close.get(&pos_pair_close) {
693              logsumexp(
694                &mut sum,
695                basepair_probs[&pos_pair_close] + sum_close - x
696                  + fold_scores.twoloop_scores[&(k, l, i, j)],
697              );
698            }
699          }
700        }
701        let sum_accessible = fold_sums.sums_accessible[&pos_pair_accessible]
702          + fold_score_sets.multibranch_score_basepair;
703        for k in 0..long_i {
704          let x = fold_sums.sums_1ormore_basepairs[k + 1][long_i - 1];
705          logsumexp(&mut sum, sum_accessible + probs_multibranch2[k][long_j] + x);
706          let y = probs_multibranch[k][long_j];
707          logsumexp(
708            &mut sum,
709            sum_accessible
710              + y
711              + fold_score_sets.multibranch_score_unpair * (long_i - k - 1) as Score,
712          );
713          logsumexp(&mut sum, sum_accessible + x + y);
714        }
715        if sum > NEG_INFINITY {
716          basepair_probs.insert(pos_pair_accessible, sum);
717        }
718      }
719    }
720  }
721  basepair_probs = basepair_probs.iter().map(|(x, &y)| (*x, expf(y))).collect();
722  basepair_probs
723}