Fix/pip likelihood crash (#42)

* Set probabilities at root to f64::min if zero

Hack to avoid running into underflow when branch lengths grow to infinity and
site likelihoods become 0. If at the root the individual probability of a single site
is 0.0, set it to f64::MIN_POSITIVE instead to limit the likelihood to a numerical
value. Needs to be tested for stability/numerical correctness.

* Add test for dataset with -Inf likelihood

* Test for repeated blen optimisation

Added a test for repeated branch length opimisation. The test is just to make sure it runs and does not
crash with -Inf likelihood.

* Minor refactoring for meaningful names

* Checking for x < f64::min instead of == 0.0

Replace direct float comparison to 0.0 by comparison to f64::MIN_POSITIVE
instead.

* Fix/input alphabet mismatch (#43)

* Allow to specify alphabet for sequences

* Made alphabets lazy static, remove clones

Made dna and protein alphabets lazy static, made functions returning their copies public,
added derive Copy, removed unnecessary alphabet clones

* Add missing test data

Added missing test data file that is compatible with the dna alphabet but is
actually a very trimmed down proptein alignment.

* Fix log message

* Remove tree file from test

* Change if x.abs() < f64::MIN to x.is_subnormal

Changed check for single column probability from x.abs < f64::MIN_POSITIVE
to a more meaningful x.is_subnormal(). This produces the same result as limiting
the likelihood to be at most f64::MIN / msa_length, but reduces the number of
ln() calls.

* Change comment for pip likelihood correction

Accidentally working on #42 as well.

* Check column probs against 0.0 and is_NaN

* Add test for too small column probs

* Replace repeatedly computed const with lazy_static

* Remove is_nan condition which could cover up bugs

* Removing stray comments

Author: junniest

Diff for: phylo/data/p105.msa.fa

@@ -0,0 +1,18 @@
+>284813
+--------------------------------------------QRLIL-----------
+---------------------------------
+>284811
+-MSDL--E----HESVPK--IPNESVWPDIVYLPDFKPTFPKWQ----------------
+---------------------------------
+>284593
+MSSELPGE----HESIPK--IPSEAVWPDIVYLPDFKPSFPQWR----------------
+---------------------------------
+>237561
+-MVET--E----HESIPQ--MPNEEIWPDVNYLPDFKSSFPQWK----------------
+----------NDDRDHNNYNEDNIGIDKHQNM-
+>284591
+-MVDT--E----HESISP--MPTEETWPGVTALPDYKPTFPQWS-----VGTNMKSGYNA
+AAGPGLMMST-----------------------
+>284812
+-------ERSNC-DRISETGLPNEEVWPGVTLLQDYKSTFPRWK----------------
+--------------------------------F

Diff for: phylo/data/p105.newick

@@ -0,0 +1 @@
+((284811:0.0000000000000002220446049250313,(284593:0.09695146373489569,(237561:0.3249811214954128,(284812:0.2932969786389556,(284813:4183.274915114282,284591:0.16527960508898415):0.16533876806006598):0.04817830194688745):0.10198380019982661):0.041663718143987165):0);

Diff for: phylo/data/p226.msa.fa

@@ -0,0 +1,12 @@
+>5207
+---A
+>284811
+----
+>284593
+----
+>284591
+G-G-
+>284592
+----
+>284812
+-V-G

Diff for: phylo/src/alignment/alignment_builder.rs

@@ -94,7 +94,7 @@ impl<'a> AlignmentBuilder<'a> {
                 .collect();
             Record::with_attrs(rec.id(), rec.desc(), &seq)
         });
-        self.seqs = Sequences::with_alphabet(new_seqs.collect(), self.seqs.alphabet().clone());
+        self.seqs = Sequences::with_alphabet(new_seqs.collect(), *self.seqs.alphabet());
     }
 
     fn stack_maps(msa_len: usize, map_x: &Mapping, map_y: &Mapping) -> Mapping {

Diff for: phylo/src/alignment/mod.rs

@@ -122,10 +122,7 @@ impl Alignment {
             records.push(Record::with_attrs(rec.id(), rec.desc(), &aligned_seq));
         }
 
-        Ok(Sequences::with_alphabet(
-            records,
-            self.seqs.alphabet.clone(),
-        ))
+        Ok(Sequences::with_alphabet(records, self.seqs.alphabet))
     }
 
     fn compile_leaf_map(&self, root: &NodeIdx, tree: &Tree) -> Result<LeafMapping> {

Diff for: phylo/src/alignment/sequences.rs

@@ -117,7 +117,7 @@ impl Sequences {
         Sequences {
             s: seqs,
             aligned: false,
-            alphabet: self.alphabet.clone(),
+            alphabet: *self.alphabet(),
         }
     }
 

Diff for: phylo/src/alphabets/mod.rs

@@ -15,7 +15,7 @@ pub static AMB_CHAR: u8 = b'X';
 pub static GAP: u8 = b'-';
 pub static POSSIBLE_GAPS: &[u8] = b"_*-";
 
-#[derive(Debug, PartialEq, Clone)]
+#[derive(Debug, PartialEq, Clone, Copy)]
 pub struct Alphabet {
     symbols: &'static [u8],
     ambiguous: &'static [u8],
@@ -64,35 +64,21 @@ impl Alphabet {
 }
 
 pub fn detect_alphabet(sequences: &[Record]) -> Alphabet {
-    let dna_alphabet = dna_alphabet();
+    let dna_alphabet = *DNA_ALPHABET;
     for record in sequences.iter() {
         if !dna_alphabet.is_word(record.seq()) {
-            return protein_alphabet();
+            return *PROTEIN_ALPHABET;
         }
     }
     dna_alphabet
 }
 
-pub(crate) fn dna_alphabet() -> Alphabet {
-    Alphabet {
-        symbols: NUCLEOTIDES,
-        ambiguous: AMB_NUCLEOTIDES,
-        index: &NUCLEOTIDE_INDEX,
-        valid_symbols: &VALID_NUCLEOTIDES,
-        char_sets: &NUCLEOTIDE_SETS,
-        name: "DNA",
-    }
+pub fn dna_alphabet() -> Alphabet {
+    *DNA_ALPHABET
 }
 
-pub(crate) fn protein_alphabet() -> Alphabet {
-    Alphabet {
-        symbols: AMINOACIDS,
-        ambiguous: AMB_AMINOACIDS,
-        index: &AMINOACID_INDEX,
-        valid_symbols: &VALID_AMINOACIDS,
-        char_sets: &AMINOACID_SETS,
-        name: "protein",
-    }
+pub fn protein_alphabet() -> Alphabet {
+    *PROTEIN_ALPHABET
 }
 
 lazy_static! {
@@ -120,6 +106,14 @@ lazy_static! {
             .cloned()
             .collect()
     };
+    pub static ref DNA_ALPHABET: Alphabet = Alphabet {
+        symbols: NUCLEOTIDES,
+        ambiguous: AMB_NUCLEOTIDES,
+        index: &NUCLEOTIDE_INDEX,
+        valid_symbols: &VALID_NUCLEOTIDES,
+        char_sets: &NUCLEOTIDE_SETS,
+        name: "DNA",
+    };
 }
 
 fn generic_nucleotide_sets(char: u8) -> FreqVector {
@@ -168,6 +162,14 @@ lazy_static! {
             .cloned()
             .collect()
     };
+    pub static ref PROTEIN_ALPHABET: Alphabet = Alphabet {
+        symbols: AMINOACIDS,
+        ambiguous: AMB_AMINOACIDS,
+        index: &AMINOACID_INDEX,
+        valid_symbols: &VALID_AMINOACIDS,
+        char_sets: &AMINOACID_SETS,
+        name: "protein",
+    };
 }
 
 fn generic_aminoacid_sets(char: u8) -> FreqVector {

Diff for: phylo/src/optimisers/blen_optimiser.rs

@@ -6,12 +6,10 @@ use argmin::solver::brent::BrentOpt;
 use log::{debug, info};
 
 use crate::likelihood::TreeSearchCost;
-use crate::optimisers::PhyloOptimisationResult;
+use crate::optimisers::{PhyloOptimisationResult, SingleValOptResult};
 use crate::tree::NodeIdx;
 use crate::Result;
 
-use super::SingleValOptResult;
-
 pub struct BranchOptimiser<C: TreeSearchCost + Display + Clone> {
     pub(crate) epsilon: f64,
     pub(crate) c: RefCell<C>,
@@ -79,7 +77,7 @@ impl<C: TreeSearchCost + Clone + Display> BranchOptimiser<C> {
 impl<C: TreeSearchCost + Clone + Display> BranchOptimiser<C> {
     pub(crate) fn optimise_branch(&mut self, branch: &NodeIdx) -> Result<SingleValOptResult> {
         let start_blen = self.c.borrow().tree().node(branch).blen;
-        let (start, end) = if start_blen == 0.0 {
+        let (min, max) = if start_blen == 0.0 {
             (0.0, 1.0)
         } else {
             (start_blen * 0.1, start_blen * 10.0)
@@ -88,7 +86,7 @@ impl<C: TreeSearchCost + Clone + Display> BranchOptimiser<C> {
             cost: &mut self.c,
             branch: *branch,
         };
-        let gss = BrentOpt::new(start, end);
+        let gss = BrentOpt::new(min, max);
         let res = Executor::new(optimiser, gss)
             .configure(|_| IterState::new().param(start_blen))
             .run()?;

Diff for: phylo/src/optimisers/blen_optimiser_tests.rs

@@ -6,10 +6,9 @@ use crate::likelihood::TreeSearchCost;
 use crate::optimisers::BranchOptimiser;
 use crate::phylo_info::PhyloInfoBuilder as PIB;
 use crate::pip_model::{PIPCostBuilder as PIPCB, PIPModel};
-use crate::substitution_models::{dna_models::*, SubstModel, SubstitutionCostBuilder as SCB};
-use crate::tree::tree_parser::from_newick;
-
+use crate::substitution_models::{dna_models::*, SubstModel, SubstitutionCostBuilder as SCB, WAG};
 use crate::tree;
+use crate::tree::tree_parser::from_newick;
 
 #[test]
 fn branch_opt_likelihood_increase_pip() {
@@ -97,3 +96,32 @@ fn branch_optimiser_against_phyml() {
     assert_eq!(o.cost.cost(), o.final_cost);
     assert_eq!(c.cost(), o.final_cost);
 }
+
+#[test]
+fn repeated_optimisation_limit() {
+    // This used to create -Inf likelihoods due to too long branch lengths and the probability
+    // turning to 0.0.
+    // This is supposed to run and not crash, no other conditions.
+
+    let fldr = Path::new("./data/");
+    let seq_file = fldr.join("p105.msa.fa");
+    let info = PIB::new(seq_file).build().unwrap();
+
+    let model = PIPModel::<WAG>::new(&[], &[]);
+
+    let mut cost = PIPCB::new(model, info).build().unwrap();
+    let mut prev_cost = f64::NEG_INFINITY;
+    let mut final_cost = TreeSearchCost::cost(&cost);
+    let max_iterations = 100;
+    let epsilon = 1e-5;
+
+    let mut iterations = 0;
+    while final_cost - prev_cost > epsilon && iterations < max_iterations {
+        iterations += 1;
+        prev_cost = final_cost;
+        let branch_o = BranchOptimiser::new(cost.clone()).run().unwrap();
+        assert!(branch_o.final_cost > branch_o.initial_cost);
+        final_cost = branch_o.final_cost;
+        cost = branch_o.cost;
+    }
+}

Diff for: phylo/src/phylo_info/phyloinfo_builder.rs

@@ -5,6 +5,7 @@ use anyhow::bail;
 use log::{info, warn};
 
 use crate::alignment::{AlignmentBuilder, Sequences};
+use crate::alphabets::Alphabet;
 use crate::io::{self, DataError};
 use crate::phylo_info::PhyloInfo;
 use crate::tree::{build_nj_tree, Tree};
@@ -13,6 +14,7 @@ use crate::Result;
 pub struct PhyloInfoBuilder {
     sequence_file: PathBuf,
     tree_file: Option<PathBuf>,
+    alphabet: Option<Alphabet>,
 }
 
 impl PhyloInfoBuilder {
@@ -32,6 +34,7 @@ impl PhyloInfoBuilder {
         PhyloInfoBuilder {
             sequence_file,
             tree_file: None,
+            alphabet: None,
         }
     }
 
@@ -53,6 +56,7 @@ impl PhyloInfoBuilder {
         PhyloInfoBuilder {
             sequence_file,
             tree_file: Some(tree_file),
+            alphabet: None,
         }
     }
 
@@ -92,6 +96,25 @@ impl PhyloInfoBuilder {
         self
     }
 
+    /// Sets the tree file path for the PhyloInfoBuilder struct.
+    /// Returns the PhyloInfoBuilder struct with the tree file path set.
+    ///
+    /// # Arguments
+    /// * `path` - File path to the tree newick file.
+    ///
+    /// # Example
+    /// ```
+    /// use std::path::PathBuf;
+    /// use phylo::alphabets::protein_alphabet;
+    /// use phylo::phylo_info::PhyloInfoBuilder;
+    /// let info = PhyloInfoBuilder::new(PathBuf::from("./data/sequences_DNA_small.fasta")).alphabet(Some(protein_alphabet())).build().unwrap();
+    /// assert_eq!(info.msa.alphabet(), &protein_alphabet());
+    /// ```
+    pub fn alphabet(mut self, alphabet: Option<Alphabet>) -> PhyloInfoBuilder {
+        self.alphabet = alphabet;
+        self
+    }
+
     /// Builds the PhyloInfo struct from the sequence file and the tree file (if provided).
     /// If the provided tree file has more than one tree, only the first tree will be processed.
     /// If no tree file is provided, an NJ tree is built from the sequences.
@@ -119,7 +142,19 @@ impl PhyloInfoBuilder {
             "Reading sequences from file {}",
             self.sequence_file.display()
         );
-        let sequences = Sequences::new(io::read_sequences_from_file(&self.sequence_file)?);
+        let sequences = if self.alphabet.is_none() {
+            info!("No alphabet provided, detecting alphabet from sequences");
+            Sequences::new(io::read_sequences_from_file(&self.sequence_file)?)
+        } else {
+            info!(
+                "Using provided {} alphabet",
+                self.alphabet.as_ref().unwrap()
+            );
+            Sequences::with_alphabet(
+                io::read_sequences_from_file(&self.sequence_file)?,
+                self.alphabet.unwrap(),
+            )
+        };
         info!("{} sequence(s) read successfully", sequences.len());
 
         let tree = match &self.tree_file {

Diff for: phylo/src/phylo_info/tests.rs

@@ -5,6 +5,7 @@ use assert_matches::assert_matches;
 use bio::io::fasta::Record;
 
 use crate::alignment::Sequences;
+use crate::alphabets::{dna_alphabet, protein_alphabet};
 use crate::frequencies;
 use crate::io::{read_sequences_from_file, DataError};
 use crate::phylo_info::{PhyloInfo, PhyloInfoBuilder as PIB};
@@ -425,3 +426,18 @@ fn empirical_frequencies_no_aas() {
         epsilon = 1e-6
     );
 }
+
+#[test]
+fn force_protein_alphabet() {
+    // If data is severely subsampled it might look like DNA even though it's really protein
+    let fldr = PathBuf::from("./data");
+    let info = PIB::new(fldr.join("p226.msa.fa")).build().unwrap();
+    assert_eq!(info.msa.alphabet(), &dna_alphabet());
+
+    let fldr = PathBuf::from("./data");
+    let info = PIB::new(fldr.join("p226.msa.fa"))
+        .alphabet(Some(protein_alphabet()))
+        .build()
+        .unwrap();
+    assert_eq!(info.msa.alphabet(), &protein_alphabet());
+}

Diff for: phylo/src/pip_model/mod.rs

@@ -7,6 +7,7 @@ use std::ops::Mul;
 use std::vec;
 
 use anyhow::bail;
+use lazy_static::lazy_static;
 use log::warn;
 use nalgebra::{DMatrix, DVector};
 
@@ -25,6 +26,10 @@ use crate::Result;
 // (2.0 * PI).ln() / 2.0;
 pub static SHIFT: f64 = 0.9189385332046727;
 
+lazy_static! {
+    pub static ref MINLOGPROB: f64 = (f64::MIN_POSITIVE).ln();
+}
+
 fn log_factorial_shifted(n: usize) -> f64 {
     // An approximation using Stirling's formula, minus constant log(sqrt(2*PI)).
     // Has an absolute error of 3e-4 for n=2, 3e-6 for n=10, 3e-9 for n=100.
@@ -317,7 +322,22 @@ impl<Q: QMatrix> PIPCost<Q> {
         let root_idx = usize::from(&self.info.tree.root);
         let msa_length = self.info.msa.len();
 
-        tmp.pnu[root_idx].map(|x| x.ln()).sum() + tmp.c0_pnu[root_idx]
+        // In certain scenarios (e.g. a completely unrelated sequence, see data/p105.msa.fa)
+        // individual column probabilities become too close to 0.0 (become subnormal)
+        // and the log likelihood becomes -Inf. This is mathematically reasonable, but during branch
+        // length optimisation BrentOpt cannot handle it and proposes NaN branch lengths.
+        // This is a workaround that sets the probability to the smallest posible positive float,
+        // which is equivalent to restricting the log likelihood to f64::MIN.
+        tmp.pnu[root_idx]
+            .map(|x| {
+                if x == 0.0 || x.is_subnormal() {
+                    *MINLOGPROB
+                } else {
+                    x.ln()
+                }
+            })
+            .sum()
+            + tmp.c0_pnu[root_idx]
             - log_factorial_shifted(msa_length)
     }