tk/Candidate_site: equivalent but more strictly typed reformulation

lib/tk/candidate_site.ml

@@ -40,23 +40,17 @@ let make ?alignment_id ?score ~condition tree alignment pos =
 
 module Leaf_data = struct
   let aa_counts leaf_data =
-    let leaves =
-      leaf_data
-      |> List.filter_map ~f:(fun ld ->
-          match Amino_acid.of_char ld.state with
-          | Some aa -> Some (Amino_acid.to_int aa)
-          | None -> (
-              match ld.state with
-              | '-' -> None
-              | _ -> invalid_arg "not an AA"
-            )
-        )
-    in
-    let r = Array.create ~len:Amino_acid.card 0 in
-    List.iter leaves ~f:(fun aa ->
-        r.(aa) <- r.(aa) + 1
-      ) ;
-    r
+    Sequence.of_list leaf_data
+    |> Sequence.filter_map ~f:(fun ld ->
+        match Amino_acid.of_char ld.state with
+        | Some aa -> Some aa
+        | None -> (
+            match ld.state with
+            | '-' -> None
+            | _ -> invalid_arg "not an AA"
+          )
+      )
+    |> Amino_acid.counts
 end
 
 let aa_counts s =
@@ -69,12 +63,15 @@ let aa_counts s =
   Leaf_data.aa_counts data_cond1
 
 let profile_of_counts k =
-  let sum = float (Array.fold k ~init:0 ~f:( + )) in
-  Array.map k ~f:(fun k -> float k /. sum)
+  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 mean_profile s =
   let k1, k2 = aa_counts s in
-  let k = Array.map2_exn k1 k2 ~f:( + ) in
+  let k = Amino_acid.Table.(
+      init (fun aa -> get k1 aa + get k2 aa)
+    ) in
   profile_of_counts k
 
 let profiles s =
@@ -251,10 +248,10 @@ let draw s =
           | `Convergent -> Gg.Color.red 
         in
         Picture.blend [
-          Picture.scale ~center:`bbox_center ~sx:0.95 ~sy:0.95 (draw_profile ~col profile) ;
+          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 (Array.length profile) /. ratio)
+            ~xmin:0. ~xmax:(float Amino_acid.card /. ratio)
             ~ymin:0. ~ymax:1. () ;
         ]
       )

lib/tk/candidate_site.mli

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