Test/inference

packages/treetime/src/commands/timetree/inference/forward_pass.rs

@@ -18,8 +18,8 @@ pub fn propagate_distributions_forward(graph: &GraphAncestral) -> Result<(), Rep
 fn propagate_distributions_forward_single_node(
   node: &mut GraphNodeForward<NodeAncestral, EdgeAncestral, ()>,
 ) -> Result<(), Report> {
-  set_likely_time(node);
   refine_distribution_from_parent(node)?;
+  set_likely_time(node);
   Ok(())
 }
 

packages/treetime/src/commands/timetree/inference/runner.rs

@@ -208,6 +208,7 @@ mod tests {
     //   --dates data/flu/h3n2/20/metadata.tsv \
     //   --branch-length-mode input \
     //   --sequence-len 1400 \
+    //   --keep-polytomies \
     //   --outdir tmp/python_treetime_baseline
     //
     // Output: tmp/python_treetime_baseline/dates.tsv
@@ -216,6 +217,7 @@ mod tests {
     // for branch length distributions (same as this test), enabling direct comparison.
     let expected = btreemap! {
       o!("NODE_0000017") => 1996.974064,
+      o!("NODE_0000018") => 1997.116240,
       o!("NODE_0000012") => 1998.499705,
       o!("A/Canterbury/58/2000|CY009150|09/05/2000|New_Zealand||H3N2/8-1416") => 2000.681725,
       o!("NODE_0000011") => 1998.763998,

packages/treetime/src/distribution/distribution_convolution.rs

@@ -26,7 +26,7 @@ pub fn distribution_convolution(a: &Distribution, b: &Distribution) -> Result<Di
       convolution_range_range(a, b) //
     },
     (Distribution::Point(a), Distribution::Function(b)) | (Distribution::Function(b), Distribution::Point(a)) => {
-      Ok(convolution_point_function(a, b)?) //
+      Ok(Distribution::Function(convolution_point_function(a, b)?))
     },
     (Distribution::Range(a), Distribution::Function(b)) | (Distribution::Function(b), Distribution::Range(a)) => {
       Ok(convolution_range_function(a, b)) //
@@ -73,20 +73,32 @@ fn convolution_point_range(p: &DistributionPoint<f64>, r: &DistributionRange<f64
 fn convolution_point_function(
   p: &DistributionPoint<f64>,
   f: &DistributionFunction<f64>,
-) -> Result<Distribution, Report> {
+) -> Result<DistributionFunction<f64>, Report> {
   let t = f.t().map(|t| t + p.t());
   let y = f.y().map(|y| y * p.amplitude());
-  Distribution::function(t, y)
+  DistributionFunction::new(t, y)
 }
 
 fn convolution_range_function(r: &DistributionRange<f64>, f: &DistributionFunction<f64>) -> Distribution {
-  let t_out = f.t().clone();
-  let mut y_out = Array1::zeros(f.y().len());
+  // split in a convolution with
+  // - a point distribution (taking care of the shift + amplitude)
+  // - an interval centered on zero and of a fixed width (taking care of the smoothing)
+
+  let shift = (r.start() + r.end()) / 2.0;
+  let amplitude = r.amplitude();
+  let width = r.end() - r.start();
+
+  let point_distr = DistributionPoint::new(shift, amplitude);
+  let shifted_function = convolution_point_function(&point_distr, f).unwrap();
 
-  for (i, &ti) in f.t().iter().enumerate() {
-    let mask = f.t().mapv(|x| (x >= ti - r.end()) && (x <= ti - r.start()));
-    let filtered_y = f.y() * &mask.mapv(|x| if x { 1.0 } else { 0.0 });
-    y_out[i] = r.amplitude() * filtered_y.sum();
+  // Convolution with a range centered on zero and of given width
+  let t_out = shifted_function.t().clone();
+  let mut y_out = Array1::zeros(shifted_function.y().len());
+
+  for (i, &ti) in shifted_function.t().iter().enumerate() {
+    let mask = shifted_function.t().mapv(|x| (x - ti).abs() <= width / 2.0);
+    let filtered_y = shifted_function.y() * &mask.mapv(|x| if x { 1.0 } else { 0.0 });
+    y_out[i] = filtered_y.sum();
   }
 
   Distribution::function(t_out, y_out).unwrap()
@@ -280,8 +292,8 @@ mod tests {
     let f = Distribution::function(x, y).unwrap();
     let actual = distribution_convolution(&r, &f).unwrap();
 
-    let x = array![0.0, 1.0, 2.0, 3.0, 4.0, 5.0];
-    let y = array![0.0, 0.0, 2.0, 2.0, 6.0, 6.0];
+    let x = array![2.0, 3.0, 4.0, 5.0, 6.0, 7.0];
+    let y = array![2.0, 2.0, 6.0, 6.0, 6.0, 2.0];
     let expected = Distribution::function(x, y).unwrap();
 
     assert_eq!(expected, actual);