feat: add penalty weights for pick_hyb_probe task

fulcrumgenomics · Sep 19, 2024 · 47062c2 · 47062c2
1 parent 8d332df
commit 47062c2
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 4 deletions.
diff --git a/prymer/primer3/primer3_weights.py b/prymer/primer3/primer3_weights.py
@@ -32,20 +32,21 @@
 @dataclass(frozen=True, init=True, slots=True)
 class Primer3Weights:
     """Holds the weights that Primer3 uses to adjust penalties
-     that originate from the designed primer(s).
+     that originate from the designed oligo(s).
 
     The weights that Primer3 uses when a parameter is less than optimal are labeled with "_lt".
     "_gt" weights are penalties applied when a parameter is greater than optimal.
 
+    Some of these settings depart from the default settings enumerated in the Primer3 manual.
     Please see the Primer3 manual for additional details:
      https://primer3.org/manual.html#globalTags
 
      Example:
          >>> Primer3Weights() #default implementation
-         Primer3Weights(product_size_lt=1, product_size_gt=1, product_tm_lt=0.0, product_tm_gt=0.0, primer_end_stability=0.25, primer_gc_lt=0.25, primer_gc_gt=0.25, primer_self_any=0.1, primer_self_end=0.1, primer_size_lt=0.5, primer_size_gt=0.1, primer_tm_lt=1.0, primer_tm_gt=1.0)
+         Primer3Weights(product_size_lt=1, product_size_gt=1, product_tm_lt=0.0, product_tm_gt=0.0, primer_end_stability=0.25, primer_gc_lt=0.25, primer_gc_gt=0.25, primer_self_any=0.1, primer_self_end=0.1, primer_size_lt=0.5, primer_size_gt=0.1, primer_tm_lt=1.0, primer_tm_gt=1.0, probe_size_lt=0.25, probe_size_gt=0.25, probe_tm_lt=1.0, probe_tm_gt=1.0, probe_gc_lt=0.5, probe_gc_gt=0.5, probe_self_any=1.0, probe_self_end=1.0)
 
          >>> Primer3Weights(product_size_lt=5)
-         Primer3Weights(product_size_lt=5, product_size_gt=1, product_tm_lt=0.0, product_tm_gt=0.0, primer_end_stability=0.25, primer_gc_lt=0.25, primer_gc_gt=0.25, primer_self_any=0.1, primer_self_end=0.1, primer_size_lt=0.5, primer_size_gt=0.1, primer_tm_lt=1.0, primer_tm_gt=1.0)
+         Primer3Weights(product_size_lt=5, product_size_gt=1, product_tm_lt=0.0, product_tm_gt=0.0, primer_end_stability=0.25, primer_gc_lt=0.25, primer_gc_gt=0.25, primer_self_any=0.1, primer_self_end=0.1, primer_size_lt=0.5, primer_size_gt=0.1, primer_tm_lt=1.0, primer_tm_gt=1.0, probe_size_lt=0.25, probe_size_gt=0.25, probe_tm_lt=1.0, probe_tm_gt=1.0, probe_gc_lt=0.5, probe_gc_gt=0.5, probe_self_any=1.0, probe_self_end=1.0)
     """  # noqa: E501
 
     product_size_lt: int = 1
@@ -61,6 +62,15 @@ class Primer3Weights:
     primer_size_gt: float = 0.1
     primer_tm_lt: float = 1.0
     primer_tm_gt: float = 1.0
+    probe_size_lt: float = 0.25
+    probe_size_gt: float = 0.25
+    probe_tm_lt: float = 1.0
+    probe_tm_gt: float = 1.0
+    probe_gc_lt: float = 0.5
+    probe_gc_gt: float = 0.5
+    probe_self_any: float = 1.0
+    probe_self_end: float = 1.0
+    probe_hairpin_th: float = 1.0
 
     def to_input_tags(self) -> dict[Primer3InputTag, Any]:
         """Maps weights to Primer3InputTag to feed directly into Primer3."""
@@ -78,5 +88,14 @@ def to_input_tags(self) -> dict[Primer3InputTag, Any]:
             Primer3InputTag.PRIMER_WT_SIZE_GT: self.primer_size_gt,
             Primer3InputTag.PRIMER_WT_TM_LT: self.primer_tm_lt,
             Primer3InputTag.PRIMER_WT_TM_GT: self.primer_tm_gt,
+            Primer3InputTag.PRIMER_INTERNAL_WT_SIZE_LT: self.probe_size_lt,
+            Primer3InputTag.PRIMER_INTERNAL_WT_SIZE_GT: self.probe_size_gt,
+            Primer3InputTag.PRIMER_INTERNAL_WT_TM_LT: self.probe_tm_lt,
+            Primer3InputTag.PRIMER_INTERNAL_WT_TM_GT: self.probe_tm_gt,
+            Primer3InputTag.PRIMER_INTERNAL_WT_GC_PERCENT_LT: self.probe_gc_lt,
+            Primer3InputTag.PRIMER_INTERNAL_WT_GC_PERCENT_GT: self.probe_gc_gt,
+            Primer3InputTag.PRIMER_INTERNAL_WT_SELF_ANY: self.probe_self_any,
+            Primer3InputTag.PRIMER_INTERNAL_WT_SELF_END: self.probe_self_end,
+            Primer3InputTag.PRIMER_INTERNAL_WT_HAIRPIN_TH: self.probe_hairpin_th
         }
         return mapped_dict
diff --git a/tests/primer3/test_primer3_weights.py b/tests/primer3/test_primer3_weights.py
@@ -19,7 +19,16 @@ def test_primer_weights_valid() -> None:
     assert test_dict[Primer3InputTag.PRIMER_WT_SIZE_GT] == 0.1
     assert test_dict[Primer3InputTag.PRIMER_WT_TM_LT] == 1.0
     assert test_dict[Primer3InputTag.PRIMER_WT_TM_GT] == 1.0
-    assert len((test_dict.values())) == 13
+    assert test_dict[Primer3InputTag.PRIMER_INTERNAL_WT_SIZE_LT] == 0.25
+    assert test_dict[Primer3InputTag.PRIMER_INTERNAL_WT_SIZE_GT] == 0.25
+    assert test_dict[Primer3InputTag.PRIMER_INTERNAL_WT_TM_LT] == 1.0
+    assert test_dict[Primer3InputTag.PRIMER_INTERNAL_WT_TM_GT] == 1.0
+    assert test_dict[Primer3InputTag.PRIMER_INTERNAL_WT_GC_PERCENT_LT] == 0.5
+    assert test_dict[Primer3InputTag.PRIMER_INTERNAL_WT_GC_PERCENT_GT] == 0.5
+    assert test_dict[Primer3InputTag.PRIMER_INTERNAL_WT_SELF_ANY] == 1.0
+    assert test_dict[Primer3InputTag.PRIMER_INTERNAL_WT_SELF_END] == 1.0
+    assert test_dict[Primer3InputTag.PRIMER_INTERNAL_WT_HAIRPIN_TH] == 1.0
+    assert len((test_dict.values())) == 22
 
 
 def test_primer_weights_to_input_tags() -> None: