calculate likelihoods also for the batched case

Author: benjamin-lieser

phylo_grad/src/lib.rs

@@ -89,7 +89,7 @@ impl<F: FloatTrait, const DIM: usize> FelsensteinTree<F, DIM> {
     /// If the length of `s` and `sqrt_pi` is 1, it will use a different code path that is optimized for this case and assumes that they are the same for all columns.
     ///
     /// Only the upper diagonal part of `s` is used. The gradients will only be populated in the upper diagonal and the lower diagonal will be filled with zeros.
-    /// 
+    ///
     /// This functions assumes you have already called `bind_leaf_log_p` to bind the log probabilities of the leaves.
     pub fn calculate_gradients(
         &mut self,
@@ -117,9 +117,42 @@ impl<F: FloatTrait, const DIM: usize> FelsensteinTree<F, DIM> {
                 &sqrt_pi[0],
                 tree,
                 d_trans_matrix,
+                false,
             );
         }
-        calculate_column_parallel(&mut self.log_p, s, sqrt_pi, tree)
+        calculate_column_parallel(&mut self.log_p, s, sqrt_pi, tree, false)
+    }
+
+    /// Same as `calculate_gradients`, but only calculates the log likelihoods for each side in the alignment.
+    pub fn calculate_likelihoods(
+        &mut self,
+        s: &[na::SMatrix<F, DIM, DIM>],
+        sqrt_pi: &[na::SVector<F, DIM>],
+    ) -> Vec<F> {
+        let tree = tree::Tree::new(&self.parents, &self.distances, self.num_leaves);
+        // Zero out internal nodes in log_p
+        for log_p in &mut self.log_p {
+            log_p.iter_mut().skip(self.num_leaves).for_each(|p| {
+                *p = na::SVector::<F, DIM>::zeros();
+            });
+        }
+
+        let result = if s.len() == 1 && sqrt_pi.len() == 1 {
+            let mut d_trans_matrix = Vec::new(); // not used in this case
+
+            calculate_column_parallel_single_S(
+                &mut self.log_p,
+                &s[0],
+                &sqrt_pi[0],
+                tree,
+                &mut d_trans_matrix,
+                true,
+            )
+        } else {
+            calculate_column_parallel(&mut self.log_p, s, sqrt_pi, tree, true)
+        };
+
+        return result.log_likelihood;
     }
 
     /// Same as `calculate_gradients`, but it takes also an array of the log_probabilities of the leaves.
@@ -131,56 +164,39 @@ impl<F: FloatTrait, const DIM: usize> FelsensteinTree<F, DIM> {
         log_p: &mut [&mut [na::SVector<F, DIM>]],
     ) -> FelsensteinResult<F, DIM> {
         let tree = tree::Tree::new(&self.parents, &self.distances, self.num_leaves);
-        calculate_column_parallel(
-            log_p,
-            s,
-            sqrt_pi,
-            tree,
-        )
+        calculate_column_parallel(log_p, s, sqrt_pi, tree, false)
     }
 
     /// This function calculates the gradients for a single side in the alignment.
     /// This can be useful if you want to control the parallelization yourself or if you want to calculate the gradients for a single side.
-    /// 
+    ///
     /// log_p is expected to have enough space to hold the log probabilities for all nodes
     pub fn calculate_gradients_single_side(
         &self,
         s: na::SMatrixView<F, DIM, DIM>,
         sqrt_pi: na::SVectorView<F, DIM>,
-        log_p: &mut [na::SVector<F, DIM>]
+        log_p: &mut [na::SVector<F, DIM>],
     ) -> SingleSideResult<F, DIM> {
         let tree = tree::Tree::new(&self.parents, &self.distances, self.num_leaves);
         // zero out internal nodes in log_p
         log_p[self.num_leaves..].iter_mut().for_each(|p| {
             *p = na::SVector::<F, DIM>::zeros();
         });
-        calculate_column(
-            log_p,
-            s.as_view(),
-            sqrt_pi.as_view(),
-            tree,
-            false,
-        )
+        calculate_column(log_p, s.as_view(), sqrt_pi.as_view(), tree, false)
     }
 
     pub fn calculate_likelihood_single_side(
         &self,
         s: na::SMatrixView<F, DIM, DIM>,
         sqrt_pi: na::SVectorView<F, DIM>,
-        log_p: &mut [na::SVector<F, DIM>]
+        log_p: &mut [na::SVector<F, DIM>],
     ) -> F {
         let tree = tree::Tree::new(&self.parents, &self.distances, self.num_leaves);
         // zero out internal nodes in log_p
         log_p[self.num_leaves..].iter_mut().for_each(|p| {
             *p = na::SVector::<F, DIM>::zeros();
         });
-        let result = calculate_column(
-            log_p,
-            s.as_view(),
-            sqrt_pi.as_view(),
-            tree,
-            true,
-        );
+        let result = calculate_column(log_p, s.as_view(), sqrt_pi.as_view(), tree, true);
         result.log_likelihood
     }
 }

phylo_grad/src/run.rs

@@ -178,6 +178,7 @@ pub fn calculate_column_parallel<
     S: &[na::SMatrix<F, DIM, DIM>],
     sqrt_pi: &[na::SVector<F, DIM>],
     tree: Tree<F>,
+    only_likelihood: bool,
 ) -> FelsensteinResult<F, DIM> {
     let col_results = (leaf_log_p, S, sqrt_pi)
         .into_par_iter()
@@ -187,7 +188,7 @@ pub fn calculate_column_parallel<
                 S.as_view(),
                 sqrt_pi.as_view(),
                 tree.clone(),
-                false,
+                only_likelihood,
             ) // The clone is shallow, Tree is cheap to clone
         })
         .collect::<Vec<_>>();
@@ -217,6 +218,7 @@ pub fn calculate_column_parallel_single_S<F: FloatTrait, const DIM: usize>(
     sqrt_pi: &na::SVector<F, DIM>,
     tree: Tree<F>,
     d_trans_matrix: &mut [Vec<na::SMatrix<F, DIM, DIM>>],
+    only_likelihood: bool
 ) -> FelsensteinResult<F, DIM> {
     let L = leaf_log_p.len();
 
@@ -240,12 +242,20 @@ pub fn calculate_column_parallel_single_S<F: FloatTrait, const DIM: usize>(
         .into_par_iter()
         .zip(d_trans_matrix.par_iter_mut())
         .map(|(leaf_log_p, d_trans)| {
-            cacluate_column_single_S(leaf_log_p, &param, &forward_data, tree.clone(), d_trans)
+            cacluate_column_single_S(leaf_log_p, &param, &forward_data, tree.clone(), d_trans, only_likelihood)
         })
         .collect::<Vec<_>>();
 
     let log_likelihood = result.iter().map(|r| r.0).collect::<Vec<_>>();
 
+    if only_likelihood {
+        return FelsensteinResult::<F, DIM> {
+            log_likelihood,
+            grad_s: vec![na::SMatrix::<F, DIM, DIM>::zeros()],
+            grad_sqrt_pi: vec![na::SVector::<F, DIM>::zeros()],
+        };
+    }
+
     let sum_d_log_prior = result.iter().map(|r| r.1).sum::<na::SVector<F, DIM>>();
 
     // We need to skip the root edge, as it does not exist and it will always be the last edge
@@ -289,12 +299,14 @@ fn d_rate_matrix_per_edge<F: FloatTrait, const DIM: usize>(
     sum_d_log_trans
 }
 
+/// In case of only_likelihood=true, d_trans_matrix will not be used
 fn cacluate_column_single_S<F: FloatTrait, const DIM: usize>(
     leaf_log_p: &mut [na::SVector<F, DIM>],
     param: &ParamPrecomp<F, DIM>,
     forward_data: &ForwardData<F, DIM>,
     tree: Tree<F>,
     d_trans_matrix: &mut [na::SMatrix<F, DIM, DIM>],
+    only_likelihood: bool
 ) -> (F, na::SVector<F, DIM>) {
     forward_column(leaf_log_p, tree.parents, forward_data);
     let log_p = leaf_log_p;
@@ -304,6 +316,10 @@ fn cacluate_column_single_S<F: FloatTrait, const DIM: usize>(
 
     let (log_likelihood, grad_log_p_likelihood) =
         final_likelihood(log_p_root.as_view(), log_p_prior.as_view());
+
+    if only_likelihood {
+        return (log_likelihood, na::SVector::<F, DIM>::zeros());
+    }
     let d_log_prior = grad_log_p_likelihood;
     let d_log_p_root = grad_log_p_likelihood;