Add function view_site() to detection pipeline

Work in progress, is broken right now

lib/detection_pipeline.ml

@@ -39,6 +39,11 @@ module type S = sig
     ?eps:float ->
     query ->
     text file
+  val view_site : 
+    query -> 
+    convergent_species:string list -> 
+    site_pos:int -> 
+    pdf file
 end
 
 module Make(Q : Query) = struct
@@ -160,4 +165,24 @@ module Make(Q : Query) = struct
   let ds_tree d   = (dn_ds_dnds_trees d).ds_tsv
   let dnds_tree d = (dn_ds_dnds_trees d).dnds_tsv
 
+  let%pworkflow view_site query ~convergent_species ~site_pos =
+    let tree_path = [%path tree ~branch_length_unit:`Amino_acid query] in
+    let alignment_path = [%path amino_acid_alignment query] in
+    let module CS = Reviewphiltrans_toolbox.Candidate_site in
+    let convergent_species = [%param convergent_species] in
+    let site_pos = [%param site_pos] in
+    let condition s =
+      if List.mem s convergent_species then `Convergent
+      else `Ancestral
+    in
+    let tree = Phylogenetics.Newick.from_file tree_path in
+    let alignment = 
+      Biotk.Fasta.from_file alignment_path 
+      |> Base.Result.ok_or_failwith
+      |> snd
+    in
+    CS.make ~condition tree alignment site_pos
+    |> CS.draw
+    |> Biotk_croquis.Croquis.Layout.simple
+    |> Base.Fn.flip Biotk_croquis.Croquis.Layout.render_pdf [%dest]
 end

lib/detection_pipeline.mli

@@ -39,6 +39,11 @@ module type S = sig
     ?eps:float ->
     query ->
     text file
+  val view_site : 
+    query -> 
+    convergent_species:string list -> 
+    site_pos:int -> 
+    pdf file
 end
 
 module Make(Q : Query) : S with type query := Q.t

lib/orthomam.ml

@@ -503,9 +503,13 @@ let%workflow ranking_of_results ~alignment_ids ~convergent_species (alignments :
             ~node:(fun _ -> ())
             ~branch:(fun bi -> Option.value_exn bi.length)
             ~leaf:(fun ni ->
-                let name = Option.value_exn ni.name in
-                let aa = List.Assoc.find_exn leaf_values name ~equal:String.equal in
-                name, aa, String.Set.mem convergent_species name
+                let species = Option.value_exn ni.name in
+                let state = List.Assoc.find_exn leaf_values species ~equal:String.equal in
+                let condition = 
+                  if String.Set.mem convergent_species species then `Convergent
+                  else `Ancestral
+                in
+                { Reviewphiltrans_toolbox.Candidate_site.species ; state ; condition }
               )
         )
     in
@@ -539,7 +543,11 @@ let%pworkflow draw_site q pos =
   let open Reviewphiltrans_toolbox in
   let open Biotk_croquis in
   let tree = Phylogenetics.Newick.from_file tree_fn in
-  let condition n = List.mem convergent_species n ~equal:String.equal in
+  let condition n = 
+    if List.mem convergent_species n ~equal:String.equal 
+    then `Convergent
+    else `Ancestral
+  in
   match Biotk.Fasta.from_file fasta_fn with
   | Ok (_, items) ->
     let picture =

lib/toolbox/candidate_site.ml

 open Core_kernel
 open Phylogenetics
 
-type tree = (unit, string * char * bool, float) Tree.t
+type leaf_data = {
+  species : string ;
+  state : char ;
+  condition : [`Ancestral | `Convergent] ;
+}
+
+type branch_length = float
+
+type tree = (unit, leaf_data, branch_length) Tree.t
 
 type t = {
   alignment_id : string option ;
@@ -22,9 +30,9 @@ let make ?alignment_id ?score ~condition tree alignment pos =
           ~node:(fun _ -> ())
           ~branch:(fun bi -> Option.value_exn bi.length)
           ~leaf:(fun ni ->
-              let name = Option.value_exn ni.name in
-              let aa = List.Assoc.find_exn leaf_values name ~equal:String.equal in
-              name, aa, condition name
+              let species = Option.value_exn ni.name in
+              let state = List.Assoc.find_exn leaf_values species ~equal:String.equal in
+              { species ; state ; condition = condition species }
             )
       )
   in
@@ -34,11 +42,11 @@ module Leaf_data = struct
   let aa_counts leaf_data =
     let leaves =
       leaf_data
-      |> List.filter_map ~f:(fun (_, aa, _) ->
-          match Amino_acid.of_char aa with
+      |> List.filter_map ~f:(fun ld ->
+          match Amino_acid.of_char ld.state with
           | Some aa -> Some (Amino_acid.to_int aa)
           | None -> (
-              match aa with
+              match ld.state with
               | '-' -> None
               | _ -> invalid_arg "not an AA"
             )
@@ -53,7 +61,10 @@ end
 
 let aa_counts s =
   let data = Tree.leaves s.contents in
-  let data_cond0, data_cond1 = List.partition_tf data ~f:trd3 in
+  let data_cond0, data_cond1 = 
+    List.partition_tf data ~f:(fun ld -> 
+        Poly.equal `Ancestral ld.condition
+      ) in
   Leaf_data.aa_counts data_cond0,
   Leaf_data.aa_counts data_cond1
 
@@ -87,7 +98,7 @@ let group_by_gene xs =
 let tree_add_cond_at_leaves (t : _ Tree.t) =
   let open Tree in
   let rec tree = function
-    | Leaf (id, aa, c) -> Some c, (Leaf (id, aa, c))
+    | Leaf l -> Some l.condition, (Leaf l)
     | Node n ->
       let List1.Cons (h, t) = List1.map n.branches ~f:branch in
       let cond = List.fold t ~init:(fst h) ~f:(fun acc (maybe_cond, _) ->
@@ -95,7 +106,7 @@ let tree_add_cond_at_leaves (t : _ Tree.t) =
           | None, _
           | _, None -> None
           | Some c1, Some c2 ->
-            if Bool.(c1 = c2) then acc else None
+            if Poly.equal c1 c2 then acc else None
         )
       in
       cond, Tree.node cond (List1.cons (snd h) (List.map t ~f:snd))
@@ -214,10 +225,10 @@ let draw s =
   let open Biotk_croquis.Croquis in
   let f = function
     | [] -> assert false
-    | (_, _, cond) :: _ as data ->
-      cond,
+    | ld :: _ as data ->
+      ld.condition,
       Leaf_data.aa_counts data |> profile_of_counts,
-      List.map data ~f:fst3
+      List.map data ~f:(fun ld -> ld.species)
   in
   let tree = collapsed_tree ~f s.contents in
   let profiles = Tree.leaves tree in
@@ -235,7 +246,10 @@ let draw s =
   in
   let profile_picture =
     List.map profiles ~f:(fun (cond, profile, _) ->
-        let col = if cond then Gg.Color.blue else Gg.Color.red in
+        let col = match cond with
+          | `Ancestral -> Gg.Color.blue
+          | `Convergent -> Gg.Color.red 
+        in
         Picture.blend [
           Picture.scale ~center:`bbox_center ~sx:0.95 ~sy:0.95 (draw_profile ~col profile) ;
           Picture.rect

lib/toolbox/candidate_site.mli

 open Phylogenetics
 
-type tree = (unit, string * char * bool, float) Tree.t
+type leaf_data = {
+  species : string ;
+  state : char ;
+  condition : [`Ancestral | `Convergent] ;
+}
+
+type branch_length = float
+
+type tree = (unit, leaf_data, branch_length) Tree.t
 
 type t = {
   alignment_id : string option ;
@@ -12,7 +20,7 @@ type t = {
 val make :
   ?alignment_id:string ->
   ?score:float ->
-  condition:(string -> bool) ->
+  condition:(string -> [`Ancestral | `Convergent]) ->
   Newick_ast.t ->
   Biotk.Fasta.item list ->
   int ->
@@ -25,7 +33,7 @@ val profiles : t -> float array * float array
 val group_by_gene : t list -> (string option * (int * t) list) list
 
 val collapsed_tree :
-  f:((string * char * bool) list -> 'a) ->
+  f:(leaf_data list -> 'a) ->
   tree ->
   (unit, 'a, float) Tree.t