tk: improved candidate site display

app/orthomam_app.ml

@@ -12,13 +12,8 @@ let () =
       )
     |> Top.eval
     |> Core.(Fn.flip List.take 10)
-    |> List.iter Codepitk.Candidate_site.(fun x ->
-        Option.iter print_string x.alignment_id ;
-        Option.iter (Printf.printf " %f\n") x.score ;
-        let p1, p2 = Codepitk.Candidate_site.profiles x in
-        print_endline @@ [%show: float array] (p1 :> float array) ;
-        print_endline @@ [%show: float array] (p2 :> float array) ;
-        print_endline "====="
+    |> List.iteri (fun i x ->
+        Biotk_croquis.Croquis.Layout.(render `pdf (simple (Codepitk.Candidate_site.draw x)) (`File (Printf.sprintf "candidate_%d.pdf" i)))
       )
   with
   | Failure msg -> prerr_endline msg

lib/tk/candidate_site.ml

@@ -38,11 +38,28 @@ let make ?alignment_id ?score ~condition tree alignment pos =
   in
   { alignment_id ; score ; pos = Some pos ; contents }
 
-module Leaf_data = struct
+module type Alphabet = sig
+  include Alphabet.S
+  val of_char : char -> t option
+end
+
+module AA_or_gap = struct
+  include Alphabet.Make(struct let card = 21 end)
+
+  let of_char = function
+    | '-' -> Some 20
+    | c -> Amino_acid.(of_char c |> Option.map ~f:to_int)
+
+  let to_char = function
+    | 20 -> '-'
+    | i -> Amino_acid.(of_int_exn i |> to_char)
+end
+
+module Leaf_data(A : Alphabet) = struct
   let aa_counts leaf_data =
     Sequence.of_list leaf_data
     |> Sequence.filter_map ~f:(fun ld ->
-        match Amino_acid.of_char ld.state with
+        match A.of_char ld.state with
         | Some aa -> Some aa
         | None -> (
             match ld.state with
@@ -50,33 +67,34 @@ module Leaf_data = struct
             | _ -> invalid_arg "not an AA"
           )
       )
-    |> Amino_acid.counts
-end
-
-let aa_counts s =
-  let data = Tree.leaves s.contents in
-  let data_cond0, data_cond1 = 
-    List.partition_tf data ~f:(fun ld -> 
-        Poly.equal `Ancestral ld.condition
+    |> A.counts
+
+  let aa_counts_by_condition s =
+    let data = Tree.leaves s.contents in
+    let data_cond0, data_cond1 = 
+      List.partition_tf data ~f:(fun ld -> 
+          Poly.equal `Ancestral ld.condition
+        ) in
+    aa_counts data_cond0,
+    aa_counts data_cond1
+
+  let profile_of_counts k =
+    let sum = float (A.Table.fold k ~init:0 ~f:( + )) in
+    A.Table.init (fun aa -> float (A.Table.get k aa) /. sum)
+
+  let _mean_profile s =
+    let k1, k2 = aa_counts_by_condition s in
+    let k = A.Table.(
+        init (fun aa -> get k1 aa + get k2 aa)
       ) in
-  Leaf_data.aa_counts data_cond0,
-  Leaf_data.aa_counts data_cond1
+    profile_of_counts k
 
-let profile_of_counts k =
-  let open Amino_acid in
-  let sum = float (Table.fold k ~init:0 ~f:( + )) in
-  Table.init (fun aa -> float (Table.get k aa) /. sum)
+  let profiles s =
+    let k1, k2 = aa_counts_by_condition s in
+    profile_of_counts k1, profile_of_counts k2
+end
 
-let mean_profile s =
-  let k1, k2 = aa_counts s in
-  let k = Amino_acid.Table.(
-      init (fun aa -> get k1 aa + get k2 aa)
-    ) in
-  profile_of_counts k
 
-let profiles s =
-  let k1, k2 = aa_counts s in
-  profile_of_counts k1, profile_of_counts k2
 
 module Maybe_string = struct
   module Base = struct
@@ -134,10 +152,11 @@ let draw_profile ~col xs =
   List.init n ~f:(fun i ->
       let x = float i /. ratio in
       let xtext = (x +. x +. 1. /. (ratio +. 1.)) /. 2. in
-      let aa = Amino_acid.(to_char (of_int_exn i)) |> Char.to_string in
+      let aa = AA_or_gap.(to_char (of_int_exn i)) |> Char.to_string in
       Picture.blend [
         Picture.text ~x:xtext ~y:0. ~valign:`base ~halign:`middle ~size:0.2 aa ;
         Picture.rect ~fill:col ~xmin:x ~xmax:(x +. 1. /. (ratio +. 1.)) ~ymin:0. ~ymax:xs.(i) () ;
+        Picture.rect ~fill:Gg.Color.white ~xmin:x ~xmax:(x +. 1. /. (ratio +. 1.)) ~ymin:0. ~ymax:1. () ;
       ]
     )
   |> Picture.blend
@@ -156,15 +175,17 @@ let draw_species_names xs =
 
 let draw s =
   let open Biotk_croquis.Croquis in
-  let f = function
+  let f =
+    let module LD = Leaf_data(AA_or_gap) in
+    function
     | [] -> assert false
     | ld :: _ as data ->
       ld.condition,
-      Leaf_data.aa_counts data |> profile_of_counts,
+      LD.(aa_counts data |> profile_of_counts),
       List.map data ~f:(fun ld -> ld.species)
   in
   let tree = collapsed_tree ~f s.contents in
-  let profiles = Tree.leaves tree in
+  let leaf_profiles = Tree.leaves tree in
   let species_names =
     tree
     |> Tree.leaves
@@ -173,7 +194,7 @@ let draw s =
   in
   let tree_picture = Tree_croquis.make tree in
   let profile_picture =
-    List.map profiles ~f:(fun (cond, profile, _) ->
+    List.map leaf_profiles ~f:(fun (cond, profile, _) ->
         let col = match cond with
           | `Ancestral -> Gg.Color.blue
           | `Convergent -> Gg.Color.red 
@@ -182,10 +203,31 @@ let draw s =
           Picture.scale ~center:`bbox_center ~sx:0.95 ~sy:0.95 (draw_profile ~col (profile :> float array)) ;
           Picture.rect
             ~draw:Gg.Color.black
-            ~xmin:0. ~xmax:(float Amino_acid.card /. ratio)
+            ~xmin:0. ~xmax:(float AA_or_gap.card /. ratio)
             ~ymin:0. ~ymax:1. () ;
         ]
       )
     |> Picture.vstack ~align:`left
   in
-  Picture.hstack ~align:`centered [ tree_picture ; profile_picture ; species_names ]
+  let estimated_profiles =
+    let module LD = Leaf_data(Amino_acid) in
+    let ap, cp = LD.profiles s in
+    Picture.vstack [
+      Picture.rect ~fill:Gg.Color.white ~xmin:0. ~xmax:1. ~ymin:0. ~ymax:1. () ;
+      draw_profile ~col:Gg.Color.blue (ap :> float array) ;
+      draw_profile ~col:Gg.Color.red  (cp :> float array) ;
+      Picture.rect ~fill:Gg.Color.white ~xmin:0. ~xmax:1. ~ymin:0. ~ymax:1. () ;
+    ]
+  in
+  let info_txt =
+    sprintf "%s%s%s"
+      (Option.value ~default:"" s.alignment_id)
+      (Option.value_map ~default:"" ~f:(sprintf " (%d)") s.pos)
+      (Option.value_map ~default:"" ~f:(sprintf " s = %f") s.score)
+  in
+  let info = Picture.text ~x:0. ~y:0. ~size:0.75 info_txt in
+  Picture.vstack ~align:`centered [
+    info ;
+    estimated_profiles ;
+    Picture.hstack ~align:`centered [ tree_picture ; profile_picture ; species_names ] ;
+  ]

lib/tk/candidate_site.mli

@@ -26,12 +26,6 @@ val make :
   int ->
   t
 
-val aa_counts : t -> int Amino_acid.table * int Amino_acid.table
-
-val mean_profile : t -> float Amino_acid.table
-
-val profiles : t -> float Amino_acid.table * float Amino_acid.table
-
 val group_by_gene : t list -> (string option * (int * t) list) list
 
 val collapsed_tree :