inhouse_tdg09: handle gaps

lib/tdg09.ml

@@ -63,9 +63,12 @@ module Inhouse_implementation = struct
     type t = Alignment.t * int
     type species = int
     let get_aa (al, j) sp =
-      match Amino_acid.of_char al.Alignment.sequences.(sp).[j] with
-      | Some aa -> aa
-      | None -> failwith "Unexpected character in alignment"
+      match al.Alignment.sequences.(sp).[j] with
+      | '-' -> None
+      | c ->
+        match Amino_acid.of_char c with
+        | Some aa -> Some aa
+        | None -> failwithf "Unexpected character in alignment: %C" c ()
   end
 
   let reindex_tree alignment (tree : Convergence_tree.t) =

lib/tk/tdg09.ml

@@ -177,7 +177,7 @@ end
 module type Site = sig
   type t
   type species
-  val get_aa : t -> species -> Amino_acid.t
+  val get_aa : t -> species -> Amino_acid.t option
 end
 
 module Make(Branch_info : Branch_info)(Leaf_info : Leaf_info)(Site : Site with type species = Leaf_info.species) = struct
@@ -201,7 +201,7 @@ module Make(Branch_info : Branch_info)(Leaf_info : Leaf_info)(Site : Site with t
         fun b -> (f (Branch_info.length b) :> mat)
       in
       let leaf_state = aa_of_leaf_info site in
-      CTMC.pruning tree ~transition_matrix ~leaf_state ~root_frequencies:pi
+      CTMC.pruning_with_missing_values tree ~transition_matrix ~leaf_state ~root_frequencies:pi
 
     let clip f param =
       if Float.(param.(0) > 3.) then Float.infinity
@@ -246,7 +246,7 @@ module Make(Branch_info : Branch_info)(Leaf_info : Leaf_info)(Site : Site with t
         fun b -> (f (Branch_info.length b) :> mat)
       in
       let leaf_state = aa_of_leaf_info site in
-      CTMC.pruning tree ~transition_matrix ~leaf_state ~root_frequencies:(stationary_distribution :> Vector.t)
+      CTMC.pruning_with_missing_values tree ~transition_matrix ~leaf_state ~root_frequencies:(stationary_distribution :> Vector.t)
 
     let counts xs =
       Amino_acid.Table.init (fun aa -> List.count xs ~f:(Amino_acid.equal aa))
@@ -309,7 +309,7 @@ module Make(Branch_info : Branch_info)(Leaf_info : Leaf_info)(Site : Site with t
     let inner_maximum_log_likelihood ?debug ?mode ~exchangeability_matrix tree site =
       let counts =
         Tree.leaves tree
-        |> List.map ~f:(aa_of_leaf_info site)
+        |> List.filter_map ~f:(aa_of_leaf_info site)
         |> counts
       in
       let schema = param_schema ?mode counts in
@@ -370,7 +370,7 @@ module Make(Branch_info : Branch_info)(Leaf_info : Leaf_info)(Site : Site with t
       in
       let root_frequencies = (param.stationary_distribution0 :> Vector.t) in
       let leaf_state = aa_of_leaf_info site in
-      CTMC.pruning tree ~transition_matrix ~leaf_state ~root_frequencies
+      CTMC.pruning_with_missing_values tree ~transition_matrix ~leaf_state ~root_frequencies
 
     let tuple_map (x, y) ~f = (f x, f y)
 
@@ -381,7 +381,7 @@ module Make(Branch_info : Branch_info)(Leaf_info : Leaf_info)(Site : Site with t
           | `Ancestral -> true
           | `Convergent -> false
         )
-      |> tuple_map ~f:(List.map ~f:(aa_of_leaf_info site))
+      |> tuple_map ~f:(List.filter_map ~f:(aa_of_leaf_info site))
       |> tuple_map ~f:Model2.counts
 
     let initial_param schema tree site =
@@ -406,7 +406,7 @@ module Make(Branch_info : Branch_info)(Leaf_info : Leaf_info)(Site : Site with t
     let inner_maximum_log_likelihood ?debug ?mode ?model2_opt ~exchangeability_matrix tree site =
       let schema =
         Tree.leaves tree
-        |> List.map ~f:(aa_of_leaf_info site)
+        |> List.filter_map ~f:(aa_of_leaf_info site)
         |> Model2.counts
         |> Model2.param_schema ?mode
       in
@@ -464,7 +464,7 @@ module Pack = struct
   module Site = struct
     type t = Amino_acid.t array
     type species = int
-    let get_aa = Array.get
+    let get_aa x i = Some (Array.get x i)
     let of_simulation s =
       Tree.leaves s
       |> Array.of_list

lib/tk/tdg09.mli

@@ -116,7 +116,7 @@ end
 module type Site = sig
   type t
   type species
-  val get_aa : t -> species -> Amino_acid.t
+  val get_aa : t -> species -> Amino_acid.t option
 end
 
 module Make(Branch_info : Branch_info)(Leaf_info : Leaf_info)(Site : Site with type species = Leaf_info.species) :