test: dump intermediate graph states

Author: ivan-aksamentov

packages/treetime-convolution/src/grid_fn.rs

@@ -25,74 +25,40 @@ impl<T: InterpElem> PartialEq for GridFn<T> {
 
 impl<T: InterpElem> Eq for GridFn<T> {}
 
-impl<T: InterpElem> Serialize for GridFn<T> {
+impl<T> Serialize for GridFn<T>
+where
+  T: InterpElem + Serialize,
+{
   fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
   where
     S: Serializer,
   {
     use serde::ser::SerializeStruct;
     let mut state = serializer.serialize_struct("GridFn", 2)?;
-    state.serialize_field("x", self.x())?;
-    state.serialize_field("y", self.y())?;
+    state.serialize_field("x", &self.x().to_vec())?;
+    state.serialize_field("y", &self.y().to_vec())?;
     state.end()
   }
 }
 
-impl<'de> Deserialize<'de> for GridFn<f64> {
+impl<'de, T> Deserialize<'de> for GridFn<T>
+where
+  T: InterpElem + Deserialize<'de>,
+{
   fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
   where
     D: Deserializer<'de>,
   {
-    use serde::de::{self, MapAccess, Visitor};
-    use std::fmt;
-
     #[derive(Deserialize)]
-    #[serde(field_identifier, rename_all = "lowercase")]
-    enum Field {
-      X,
-      Y,
-    }
-
-    struct GridFnVisitor;
-
-    impl<'de> Visitor<'de> for GridFnVisitor {
-      type Value = GridFn<f64>;
-
-      fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
-        formatter.write_str("struct GridFn")
-      }
-
-      fn visit_map<V>(self, mut map: V) -> Result<GridFn<f64>, V::Error>
-      where
-        V: MapAccess<'de>,
-      {
-        let mut x: Option<Array1<f64>> = None;
-        let mut y: Option<Array1<f64>> = None;
-        while let Some(key) = map.next_key()? {
-          match key {
-            Field::X => {
-              if x.is_some() {
-                return Err(de::Error::duplicate_field("x"));
-              }
-              x = Some(map.next_value()?);
-            },
-            Field::Y => {
-              if y.is_some() {
-                return Err(de::Error::duplicate_field("y"));
-              }
-              y = Some(map.next_value()?);
-            },
-          }
-        }
-        let x = x.ok_or_else(|| de::Error::missing_field("x"))?;
-        let y = y.ok_or_else(|| de::Error::missing_field("y"))?;
-
-        GridFn::new(x, y).map_err(|e| de::Error::custom(format!("Failed to create GridFn: {e}")))
-      }
+    struct GridFnHelper<T> {
+      x: Vec<T>,
+      y: Vec<T>,
     }
 
-    const FIELDS: &[&str] = &["x", "y"];
-    deserializer.deserialize_struct("GridFn", FIELDS, GridFnVisitor)
+    let helper = GridFnHelper::<T>::deserialize(deserializer)?;
+    let x_array = Array1::from_vec(helper.x);
+    let y_array = Array1::from_vec(helper.y);
+    GridFn::new(x_array, y_array).map_err(serde::de::Error::custom)
   }
 }
 

packages/treetime-convolution/src/lib.rs

@@ -3,10 +3,9 @@ pub mod grid_fn;
 pub mod testing;
 
 use num_traits::{Num, NumCast};
-use serde::{Deserialize, Serialize};
 use std::fmt::Debug;
 
-pub trait InterpElem: Num + NumCast + Debug + Send + PartialOrd + Copy + Serialize + for<'de> Deserialize<'de> {}
+pub trait InterpElem: Num + NumCast + Debug + Send + PartialOrd + Copy {}
 
 impl InterpElem for f64 {}
 

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

@@ -159,13 +159,21 @@ mod tests {
   use ordered_float::OrderedFloat;
   use pretty_assertions::assert_eq;
   use std::collections::BTreeMap;
+  use std::path::Path;
+  use treetime_io::json::{JsonPretty, json_write_file};
 
   #[test]
   fn test_timetree_flu_h3n2_poisson() -> Result<(), Report> {
     const CLOCK_RATE_MU: f64 = 0.0028;
     const SEQUENCE_LENGTH_L: usize = 1400;
     const GRID_SIZE: usize = 200;
 
+    fn dump_graph(graph: &GraphAncestral, filename: &str) -> Result<(), Report> {
+      let OUTPUT_DIR = Path::new("../../tmp/test_timetree_flu_h3n2_poisson");
+      json_write_file(OUTPUT_DIR.join(filename), graph, JsonPretty(true))?;
+      Ok(())
+    }
+
     // Rerooted H3N2 flu tree obtained from:
     //
     // cargo run --bin=treetime -- clock \
@@ -261,10 +269,13 @@ mod tests {
     load_date_constraints(&input_dates, &graph)?;
 
     create_poisson_branch_distributions(&graph, CLOCK_RATE_MU, SEQUENCE_LENGTH_L, GRID_SIZE)?;
+    dump_graph(&graph, "001_after_create_poisson_branch_distributions.json")?;
 
     propagate_distributions_backward(&graph)?;
+    dump_graph(&graph, "002_after_propagate_distributions_backward.json")?;
 
     propagate_distributions_forward(&graph)?;
+    dump_graph(&graph, "003_after_propagate_distributions_forward.json")?;
 
     let actual: BTreeMap<String, f64> = graph
       .get_nodes()

packages/treetime/src/distribution/distribution.rs

@@ -17,8 +17,6 @@ pub enum Distribution {
   Empty,
   Point(DistributionPoint<f64>),
   Range(DistributionRange<f64>),
-
-  #[serde(skip)]
   Function(DistributionFunction<f64>),
 }
 

packages/treetime/src/distribution/distribution_function.rs

@@ -1,9 +1,10 @@
 use eyre::Report;
 use ndarray::Array1;
 use ndarray_stats::QuantileExt;
+use serde::{Deserialize, Serialize};
 use treetime_convolution::{GridFn, InterpElem};
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct DistributionFunction<T: InterpElem> {
   grid_fn: GridFn<T>,
 }