tk/Convergence_tree: cleaner interface

lib/inhouse_lmm.ml

@@ -6,7 +6,10 @@ let%pworkflow[@version 2] test alignment tree =
     Alignment.from_fasta [%path alignment]
     |> Rresult.R.get_ok
   in
-  let tree = Convergence_tree.from_file [%path tree] in
+  let tree =
+    Convergence_tree.from_file [%path tree]
+    |> Result.get_ok
+  in
   Inhouse_lmm.test ~alignment ~tree
   |> Inhouse_lmm.result_table_of_test
   |> Result_table.to_file ~output:[%dest]

lib/toolbox/convergence_tree.ml

 open Core_kernel
 open Phylogenetics
 
-type t = Newick.tree
+type condition = [`Ancestral | `Convergent]
+
+type branch_info = {
+  condition : condition ;
+  length : float ;
+}
+
+type t = (unit, string, branch_info) Tree.t
 
 module Tags = struct
   let condition_label = "Condition"
@@ -31,45 +38,48 @@ end
 let condition_of_branch_info (bi : Newick.branch_info) =
   Tags.condition bi.tags
 
-  type u =
-  ( Newick.node_info,
-    Newick.node_info,
-    Phylogenetics_convergence.Simulator.Branch_info.t )
-  Tree.t
+let of_newick_tree t =
+  let open Phylogenetics.Newick in
+  try
+    let node _ = () in
+    let leaf (l : Newick.node_info) = match l.name with
+      | Some n -> n
+      | None -> failwith "missing leaf name"
+    in
+    let branch b =
+      let length = match b.length with
+        | None -> failwith "missing branch length"
+        | Some bl -> bl
+      in
+      let condition =
+        match
+          List.Assoc.find ~equal:String.equal b.tags "Condition"
+        with
+        | Some s -> (
+            match s with
+            | "0" -> `Ancestral
+            | "1" -> `Convergent
+            | _ -> failwithf "Invalid condition: %s" s () )
+        | None -> failwith "Missing Condition tag"
+      in
+      { length; condition }
+    in
+    Tree.map t ~node ~leaf ~branch
+    |> Result.return
+  with Failure msg -> Result.fail (`Msg msg)
 
 let from_file fn =
   let module BI = Phylogenetics_convergence.Simulator.Branch_info in
   Newick.from_file fn
-  |> Newick.with_inner_tree ~f:(fun t ->
-         Tree.map t ~node:Fn.id ~leaf:Fn.id
-           ~branch:
-             Phylogenetics.Newick.(
-               fun b ->
-                 let length = Option.value_exn b.length in
-                 let condition =
-                   match
-                     List.Assoc.find ~equal:String.equal b.tags "Condition"
-                   with
-                   | Some s -> (
-                       match s with
-                       | "0" -> `Ancestral
-                       | "1" -> `Convergent
-                       | _ -> failwithf "Invalid condition: %s" s () )
-                   | None -> failwith "Missing Condition tag"
-                 in
-                 { BI.length; condition }))
+  |> Newick.with_inner_tree ~f:of_newick_tree
 
 let leaves tree =
   let rec node condition t acc =
     match t with
     | Tree.Node n -> List1.fold_right n.branches ~init:acc ~f:branch
-    | Leaf { Newick.name = Some n } -> (n, condition) :: acc
-    | Leaf { name = None } ->
-        failwith "leaves_condition: missing leaf name in nhx"
+    | Leaf species -> (species, condition) :: acc
   and branch (Tree.Branch b) acc =
-    node
-      (Phylogenetics_convergence.Simulator.Branch_info.condition b.data)
-      b.tip acc
+    node b.data.condition b.tip acc
   in
   node `Ancestral tree []
 
@@ -91,9 +101,9 @@ let rec transfer_condition_to_branches t =
             (transfer_condition_to_branches b.tip))
       |> Tree.node (fst n.data)
 
-let reset_transitions (tree : t) =
+let reset_transitions (tree : Newick.tree) =
   let rec aux mother_condition tree =
-    match (tree : t) with
+    match (tree : Newick.tree) with
     | Leaf _ as l -> l
     | Node n ->
         let branches =
@@ -134,7 +144,7 @@ let length_on_each_condition branches =
       | _ -> ()) ;
   A.to_alist acc
 
-let remove_nodes_with_single_child (tree : t) =
+let remove_nodes_with_single_child (tree : Newick.tree) =
   Tree.simplify_node_with_single_child tree
     ~merge_branch_data:(fun branches ->
       let condition_stats = length_on_each_condition branches in

lib/toolbox/convergence_tree.mli

 open Core_kernel
 open Phylogenetics
 
-type t = Newick.tree
+type condition = [`Ancestral | `Convergent]
 
-type u =
-  ( Newick.node_info,
-    Newick.node_info,
-    Phylogenetics_convergence.Simulator.Branch_info.t )
-  Tree.t
+type branch_info = {
+  condition : condition ;
+  length : float ;
+}
 
-val from_file : string -> u
+type t = (unit, string, branch_info) Tree.t
 
-val leaves : u -> (string * [ `Ancestral | `Convergent ]) list
+val of_newick_tree : Newick.tree -> (t, [> `Msg of string]) result
+
+val from_file : string -> (t, [> `Msg of string]) result
+
+val leaves : t -> (string * condition) list
 
 val infer_binary_condition_on_branches :
-  ?gain_relative_cost:float -> t -> convergent_leaves:String.Set.t -> t
+  ?gain_relative_cost:float ->
+  Newick.tree ->
+  convergent_leaves:String.Set.t ->
+  Newick.tree
 
-val remove_nodes_with_single_child : t -> t
+val remove_nodes_with_single_child : Newick.tree -> Newick.tree

lib/toolbox/inhouse_lmm.ml

 open Core_kernel
 open Phylogenetics
 module L = Lacaml.D
-module BI = Phylogenetics_convergence.Simulator.Branch_info
 
 type correlations = (string * string * float) list * String.Set.t
 
@@ -14,26 +13,25 @@ let merge_correlations time_from_ancestor ((dist_l, l) : correlations)
   in
   (List.concat [ dist_l; dist_r; dist_lr ], String.Set.union l r)
 
-let correlations (t : Convergence_tree.u) : (string * string * float) list
+let correlations (t : Convergence_tree.t) : (string * string * float) list
     =
   let rec tree time_from_ancestor = function
     | Tree.Leaf l ->
-        let l = Option.value_exn l.Newick.name in
         ([ (l, l, time_from_ancestor) ], String.Set.singleton l)
     | Node n ->
         List1.map n.branches ~f:(branch time_from_ancestor)
         |> List1.reduce ~f:(merge_correlations time_from_ancestor)
   and branch time_from_ancestor (Branch b) =
-    tree (time_from_ancestor +. b.data.BI.length) b.tip
+    tree (time_from_ancestor +. b.data.Convergence_tree.length) b.tip
   in
   fst (tree 0. t)
 
-let index_correlations (t : Convergence_tree.u) cors =
+let index_correlations (t : Convergence_tree.t) cors =
   let index = String.Table.create () in
   let leaves = Tree.leaves t in
   List.iteri leaves ~f:(fun idx l ->
       String.Table.set index
-        ~key:(Option.value_exn l.Newick.name)
+        ~key:l
         ~data:idx) ;
   List.map cors ~f:(fun (l1, l2, c) ->
       let f = String.Table.find index in

lib/toolbox/inhouse_lmm.mli

@@ -7,6 +7,6 @@ open Phylogenetics
 *)
 
 val test :
-  alignment:Alignment.t -> tree:Convergence_tree.u -> float option array
+  alignment:Alignment.t -> tree:Convergence_tree.t -> float option array
 
 val result_table_of_test : float option array -> Result_table.t