Inhouse LMM partial implementation

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
+
+let merge_correlations time_from_ancestor ((dist_l, l) : correlations)
+    ((dist_r, r) : correlations) =
+  let dist_lr =
+    String.Set.fold l ~init:[] ~f:(fun acc e ->
+        String.Set.fold r ~init:acc ~f:(fun acc f ->
+            (e, f, time_from_ancestor) :: acc))
+  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 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
+  in
+  fst (tree 0. t)
+
+let design_matrix ~m ~n ~aa_at_site (al : Alignment.t) =
+  L.Mat.init_rows m n (fun i j ->
+      if Char.(al.sequences.(i).[j] = aa_at_site.(j)) then 1. else 0.)
+
+let solve ~y ~_X_ ~_T_ =
+  let _Xtilde_ = L.Mat.mul _X_ _T_ in
+  let ytilde = L.gemv _T_ y in
+  L.gemm
+    (L.getri (L.gemm ~transa:`T _Xtilde_ _Xtilde_))
+    (L.gemv ~trans:`T _Xtilde_ ytilde)
+
+let lrt_on_one_site ~alignment:al ~phenotypes:_ ~_C_ ~site =
+  let m = Alignment.nrows al in
+  let aa_at_site =
+    Alignment.residues al ~column:site |> Char.Set.to_list |> Array.of_list
+  in
+  let n = Array.length aa_at_site in
+  let _X_ = design_matrix ~m ~n ~aa_at_site al in
+  ()
+
+let phenotypes_of_tree t =
+  Convergence_tree.leaves t
+  |> List.map ~f:(fun (_, condition) ->
+         match condition with `Ancestral -> 0. | `Convergent -> 1.)
+  |> Array.of_list |> L.Vec.of_array
+
+let lrt ~alignment ~tree =
+  let phenotypes = phenotypes_of_tree tree in
+  Array.init (Alignment.ncols alignment) ~f:(fun site ->
+      lrt_on_one_site ~alignment ~phenotypes ~site)
+  |> assert false

lib/toolbox/inhouse_lmm.mli

+open Phylogenetics
+(**
+    Model: y = Xb + u + e
+    u ~ N(0, λΣ)
+
+    H0: b = 0
+*)
+
+val lrt :
+  alignment:Alignment.t -> tree:Convergence_tree.u -> Result_table.t