Commit 1e38b720 authored by Louis Duchemin's avatar Louis Duchemin
Browse files

Tdg09 : fixes matrix exponential computation

parent 12017b70
......@@ -28,3 +28,9 @@
(libraries codepi)
(preprocess
(pps ppx_jane)))
(executable
(name matrix_exponential_bench)
(public_name matrix_exp)
(modules matrix_exponential_bench)
(libraries codepi phylogenetics core_bench))
\ No newline at end of file
open Phylogenetics
open Codepitk.Tdg09
open Core_bench
let () =
let stationary_distribution = Amino_acid.random_profile 0.5 in
let exchangeability_matrix = Rate_matrix.Amino_acid.make (fun aa_i aa_j->
let i, j = Amino_acid.(to_int aa_i, to_int aa_j) in
if i <= j then float_of_int i *. 0.01 +. float_of_int j *. 0.001
else float_of_int j *. 0.01 +. float_of_int i *. 0.001
) in
let p = {
stationary_distribution;
Evolution_model.exchangeability_matrix;
scale=1.;
} in
let t = 100. in
let transition_matrix_diag = Evolution_model.transition_probability_matrix p in
let transition_matrix_exp =
let m = Evolution_model.rate_matrix p in
fun t -> Amino_acid.Matrix.(expm (scal_mul t m))
in
Core.Command.run (Bench.make_command [
Bench.Test.create ~name:"diag" (fun () -> transition_matrix_diag t) ;
Bench.Test.create ~name:"exp" (fun () -> transition_matrix_exp t) ;
])
\ No newline at end of file
......@@ -20,17 +20,46 @@ module Evolution_model = struct
p.exchangeability_matrix.Amino_acid.%{i, j} *.
p.stationary_distribution.Amino_acid.%(j)
)
let transition_probability_matrix p =
let module V = Amino_acid.Vector in
let module M = Amino_acid.Matrix in
let m = rate_matrix p in
let (d_vec, transform_matrix) = Amino_acid.Matrix.diagonalize m in
let rev_trans_matrix = Amino_acid.Matrix.inverse transform_matrix in
let sqrt_pi = V.map p.stationary_distribution ~f:Float.sqrt in
let diag_pi = M.diagm sqrt_pi in
let diag_pi_inv = V.map sqrt_pi ~f:(fun v -> 1. /. v) |> M.diagm in
let m' = M.(dot diag_pi @@ dot m diag_pi_inv) in
let (d_vec, step_transform_matrix) = M.diagonalize m' in
let transform_matrix = M.dot diag_pi_inv step_transform_matrix in
let rev_transform_matrix = M.dot (M.transpose step_transform_matrix) diag_pi in
fun t ->
let exp_matrix = Amino_acid.Vector.(exp (scal_mul t d_vec))
|> Amino_acid.Matrix.diagm in
Amino_acid.Matrix.(dot (dot transform_matrix exp_matrix) rev_trans_matrix)
(* Amino_acid.Matrix.(expm (scal_mul t m)) *)
Amino_acid.Matrix.(dot transform_matrix @@ dot exp_matrix rev_transform_matrix)
let test_diagonal_matrix_exponential () =
let stationary_distribution = Amino_acid.random_profile 0.5 in
let exchangeability_matrix = Rate_matrix.Amino_acid.make (fun aa_i aa_j->
let i, j = Amino_acid.(to_int aa_i, to_int aa_j) in
if i <= j then float_of_int i *. 0.01 +. float_of_int j *. 0.001
else float_of_int j *. 0.01 +. float_of_int i *. 0.001
) in
let p = {
stationary_distribution;
exchangeability_matrix;
scale=1.;
} in
let t = 100. in
let diag_exp_matrix = transition_probability_matrix p t in
let m = rate_matrix p in
let exp_matrix = Amino_acid.Matrix.(expm (scal_mul t m)) in
Amino_acid.Matrix.robust_equal ~tol:1e-10 diag_exp_matrix exp_matrix
let%test "Matrix exponential through diagonalisation matches naive implementation" =
test_diagonal_matrix_exponential ()
end
let choose_aa p =
Amino_acid.Table.of_vector p
|> Amino_acid.Table.choose
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment