diff --git a/lib/CamlPaml/Q.ml b/lib/CamlPaml/Q.ml index 78c95bd..a622491 100644 --- a/lib/CamlPaml/Q.ml +++ b/lib/CamlPaml/Q.ml @@ -28,7 +28,17 @@ let m_of_cm cm = Gsl.Matrix.of_arrays (Array.map (Array.map real_of_complex) (Gs let cm_of_m m = Gsl.Matrix_complex.of_arrays (Array.map (Array.map (fun re -> { Complex.re = re; im = 0. })) (Gsl.Matrix.to_arrays m)) -(* stupid wrappers *) +let v_of_cv cv = Gsl.Vector.of_array (Array.map real_of_complex (Gsl.Vector_complex.to_array cv)) + +let gemm ?c a b = + let m,n = Gsl.Matrix.dims a + let n',p = Gsl.Matrix.dims b + if n <> n' then invalid_arg "CamlPaml.Q.gemm: incompatible dimensions" + let c = match c with + | Some c -> c + | None -> Gsl.Matrix.create m p + Gsl.Blas.gemm ~ta:Gsl.Blas.NoTrans ~tb:Gsl.Blas.NoTrans ~alpha:1. ~a:a ~b:b ~beta:0. ~c:c + c let zgemm ?c a b = let m,n = Gsl.Matrix_complex.dims a @@ -40,6 +50,21 @@ let zgemm ?c a b = Gsl.Blas.Complex.gemm ~ta:Gsl.Blas.NoTrans ~tb:Gsl.Blas.NoTrans ~alpha:Complex.one ~a:a ~b:b ~beta:Complex.zero ~c:c c +let diagm ?c d a = + let m = Gsl.Vector.length d + let (n,p) = Gsl.Matrix.dims a + if m <> n then invalid_arg "CamlPaml.Q.diagm: incompatible dimensions" + let c = match c with + | Some c -> c + | None -> Gsl.Matrix.create m p + + for i = 0 to m-1 do + let d_i = d.{i} + for j = 0 to p-1 do + Bigarray.Array2.unsafe_set c i j (Bigarray.Array2.unsafe_get a i j *. d_i) + + c + let zdiagm ?c d a = let m = Gsl.Vector_complex.length d let (n,p) = Gsl.Matrix_complex.dims a @@ -51,7 +76,7 @@ let zdiagm ?c d a = for i = 0 to m-1 do let d_i = d.{i} for j = 0 to p-1 do - c.{i,j} <- Complex.mul a.{i,j} d_i + Bigarray.Array2.unsafe_set c i j (Complex.mul (Bigarray.Array2.unsafe_get a i j) d_i) c @@ -68,12 +93,21 @@ let zinvm m = | _ -> assert false module Diag = struct - (* diagonalized Q = S*L*S' *) + (* diagonalized Q = S*L*S' + These are real for reversible models, complex for non-reversible models. *) + type eig_r = { + r_s : Gsl.Matrix.matrix; (* S = right eigenvectors (in the columns) *) + r_s' : Gsl.Matrix.matrix; (* S' = left eigenvectors (in the rows) *) + r_l : Gsl.Vector.vector (* diag(L) = eigenvalues *) + } + type eig_nr = { + nr_s : Gsl.Matrix_complex.matrix; (* S = right eigenvectors (in the columns) *) + nr_s' : Gsl.Matrix_complex.matrix; (* S' = left eigenvectors (in the rows) *) + nr_l : Gsl.Vector_complex.vector; (* diag(L) = eigenvalues *) + } type t = { q : Gsl.Matrix.matrix; (* Q *) - s : Gsl.Matrix_complex.matrix; (* S = right eigenvectors (in the columns) *) - s' : Gsl.Matrix_complex.matrix; (* S' = left eigenvectors (in the rows) *) - l : Gsl.Vector_complex.vector; (* diag(L) = eigenvalues *) + eig : [`r of eig_r | `nr of eig_nr]; pi : Gsl.Vector.vector; mutable have_pi : bool; @@ -91,37 +125,49 @@ module Diag = struct let qm = Gsl.Matrix.of_arrays qm let l, s = Gsl.Eigen.nonsymmv ~protect:true (`M qm) let s' = zinvm s - { q = qm; s = s; s' = s'; l = l; - pi = Gsl.Vector.create (fst (Gsl.Matrix.dims qm)); have_pi = false; memoized_to_Pt = None; tol = tol } + + let rev = ref true + let n = Gsl.Vector_complex.length l + for i = 0 to n-1 do if not (check_real ~tol l.{i}) then rev := false + + let eig = + if !rev then + `r { r_s = m_of_cm s; r_s' = m_of_cm s'; r_l = v_of_cv l } + else + `nr { nr_s = s; nr_s' = s'; nr_l = l } + + { q = qm; eig; + pi = Gsl.Vector.create (fst (Gsl.Matrix.dims qm)); have_pi = false; + memoized_to_Pt = None; tol = tol } let to_Q q = Gsl.Matrix.to_arrays q.q - let reversible q = - let rev = ref true - let n = Gsl.Vector_complex.length q.l - for i = 0 to n-1 do if not (check_real ~tol:q.tol q.l.{i}) then rev := false - !rev + let reversible = function + | { eig = `r _ } -> true + | _ -> false - (* TODO figure out what this returns for a non-reversible matrix *) let equilibrium q = if not q.have_pi then + let eig = match q.eig with + | `r eig -> eig + | `nr _ -> failwith "CamlPaml.Q.equilibrium: non-reversible model" + let n = Gsl.Vector.length eig.r_l let min_L = ref infinity let min_Lp = ref (-1) - let n = Gsl.Vector_complex.length q.l for i = 0 to n-1 do - let nm_lz = Complex.norm q.l.{i} - if nm_lz < !min_L then - min_L := nm_lz + let mag_i = abs_float eig.r_l.{i} + if mag_i < !min_L then + min_L := mag_i min_Lp := i assert (!min_Lp >= 0) if (abs_float !min_L) > q.tol then - failwith "CamlPaml.Q.equilibrium: smallest-magnitude eigenvalue is unacceptably large; is this a valid rate matrix?" - let lev = Gsl.Matrix_complex.row q.s' !min_Lp - let mass = ref Complex.zero + failwith (sprintf "CamlPaml.Q.equilibrium: smallest-magnitude eigenvalue %e is unacceptably large; check rate matrix validity or increase tol" !min_L) + let lev = Gsl.Matrix.row eig.r_s' !min_Lp + let mass = ref 0. for i = 0 to n-1 do - mass := Complex.add !mass lev.{i} + mass := !mass +. lev.{i} for i = 0 to n-1 do - q.pi.{i} <- real_of_complex (Complex.div lev.{i} !mass) + q.pi.{i} <- lev.{i} /. !mass q.have_pi <- true Gsl.Vector.to_array q.pi @@ -143,21 +189,34 @@ module Diag = struct if x <= 0. then invalid_arg "CamlPaml.Q.scale: nonpositive scale factor" let xq = Gsl.Matrix.copy q.q Gsl.Matrix.scale xq x - let xl = Gsl.Vector_complex.copy q.l - let cx = { Complex.re = x; im = 0. } - for i = 0 to (Gsl.Vector_complex.length xl) - 1 do - xl.{i} <- Complex.mul xl.{i} cx - { q with q = xq; l = xl; memoized_to_Pt = None } + let xeig = match q.eig with + | `r { r_s; r_s'; r_l } -> + let xl = Gsl.Vector.copy r_l + for i = 0 to (Gsl.Vector.length xl) - 1 do + xl.{i} <- xl.{i} *. x + `r { r_s; r_s'; r_l = xl } + | `nr { nr_s; nr_s'; nr_l } -> + let xl = Gsl.Vector_complex.copy nr_l + let cx = { Complex.re = x; im = 0. } + for i = 0 to (Gsl.Vector_complex.length xl) - 1 do + xl.{i} <- Complex.mul xl.{i} cx + `nr { nr_s; nr_s'; nr_l = xl } + { q with q = xq; eig = xeig; memoized_to_Pt = None } let real_to_Pt q t = if t < 0. then invalid_arg "CamlPaml.Q.to_Pt" - let ct = { Complex.re = t; im = 0. } - let expLt = Gsl.Vector_complex.copy q.l - - for i = 0 to Gsl.Vector_complex.length expLt - 1 do - expLt.{i} <- Complex.exp (Complex.mul ct expLt.{i}) - - let sm = m_of_cm (zgemm q.s (zdiagm expLt q.s')) + let sm = match q.eig with + | `r { r_s; r_s'; r_l } -> + let expLt = Gsl.Vector.copy r_l + for i = 0 to Gsl.Vector.length expLt - 1 do + expLt.{i} <- exp (t *. expLt.{i}) + gemm r_s (diagm expLt r_s') + | `nr { nr_s; nr_s'; nr_l } -> + let ct = { Complex.re = t; im = 0. } + let expLt = Gsl.Vector_complex.copy nr_l + for i = 0 to Gsl.Vector_complex.length expLt - 1 do + expLt.{i} <- Complex.exp (Complex.mul ct expLt.{i}) + m_of_cm (zgemm nr_s (zdiagm expLt nr_s')) let n,_ = Gsl.Matrix.dims sm for i = 0 to n-1 do @@ -193,47 +252,51 @@ module Diag = struct let to_Pt_gc q = q.memoized_to_Pt <- None - let dPt_dt ~q ~t = - let n,_ = Gsl.Matrix.dims q.q - let (( * ),(+),(/)) = Complex.mul, Complex.add, Complex.div - let exp= Complex.exp - let s, l, s' = q.s, q.l, q.s' - let ct = { Complex.re = t; Complex.im = 0. } - Array.init n - fun a -> - Array.init n - fun b -> - let x = ref Complex.zero - for c = 0 to n-1 do - x := !x + (s.{a,c} * l.{c} * (exp (l.{c} * ct)) * s'.{c,b}) - real_of_complex !x + let dPt_dt ~q ~t = match q.eig with + | `r _ -> failwith "not implemented" + | `nr { nr_s; nr_s'; nr_l } -> + let n,_ = Gsl.Matrix.dims q.q + let (( * ),(+),(/)) = Complex.mul, Complex.add, Complex.div + let exp= Complex.exp + let s, l, s' = nr_s, nr_l, nr_s' + let ct = { Complex.re = t; Complex.im = 0. } + Array.init n + fun a -> + Array.init n + fun b -> + let x = ref Complex.zero + for c = 0 to n-1 do + x := !x + (s.{a,c} * l.{c} * (exp (l.{c} * ct)) * s'.{c,b}) + real_of_complex !x (* TODO in richly parameterized models, dQ_dx is likely to be sparse. *) - let dPt_dQ_dx ~q ~t ~dQ_dx = - let n,_ = Gsl.Matrix.dims q.q - let dQt_dx = cm_of_m (Gsl.Matrix.of_arrays dQ_dx) - - if Gsl.Matrix_complex.dims dQt_dx <> Gsl.Matrix.dims q.q then - invalid_arg "CamlPaml.Q.dP_dx: wrong dimension of dQ_dx" - - let ct = { Complex.re = t; Complex.im = 0. } - let exp = Complex.exp - let (( * ),(-),(/)) = Complex.mul, Complex.sub, Complex.div - - (* combos 1,3 or 2 (alone) -- 1,3 matches P&S? *) - Gsl.Matrix_complex.scale dQt_dx ct (* 1 *) - - let f = Gsl.Matrix_complex.create n n - for i = 0 to pred n do - for j = 0 to pred n do - if q.l.{i} = q.l.{j} then - f.{i,j} <- exp (q.l.{i} * ct) (* * ct *) (* 2 *) - else - f.{i,j} <- (exp (q.l.{i} * ct) - exp (q.l.{j} * ct)) / ((q.l.{i} - q.l.{j}) * ct (* 3 *)) - - (* ehh not being too gentle with the allocator/GC here *) - Gsl.Matrix_complex.mul_elements f (zgemm q.s' (zgemm dQt_dx q.s)) - Gsl.Matrix.to_arrays (m_of_cm (zgemm q.s (zgemm f q.s'))) + let dPt_dQ_dx ~q ~t ~dQ_dx = match q.eig with + | `r _ -> failwith "not implemented" + | `nr { nr_s; nr_s'; nr_l } -> + let n,_ = Gsl.Matrix.dims q.q + let dQt_dx = cm_of_m (Gsl.Matrix.of_arrays dQ_dx) + + if Gsl.Matrix_complex.dims dQt_dx <> Gsl.Matrix.dims q.q then + invalid_arg "CamlPaml.Q.dP_dx: wrong dimension of dQ_dx" + + let ct = { Complex.re = t; Complex.im = 0. } + let exp = Complex.exp + let (( * ),(-),(/)) = Complex.mul, Complex.sub, Complex.div + + (* combos 1,3 or 2 (alone) -- 1,3 matches P&S? *) + Gsl.Matrix_complex.scale dQt_dx ct (* 1 *) + + let f = Gsl.Matrix_complex.create n n + for i = 0 to pred n do + for j = 0 to pred n do + if nr_l.{i} = nr_l.{j} then + f.{i,j} <- exp (nr_l.{i} * ct) (* * ct *) (* 2 *) + else + f.{i,j} <- (exp (nr_l.{i} * ct) - exp (nr_l.{j} * ct)) / ((nr_l.{i} - nr_l.{j}) * ct (* 3 *)) + + (* ehh not being too gentle with the allocator/GC here *) + Gsl.Matrix_complex.mul_elements f (zgemm nr_s' (zgemm dQt_dx nr_s)) + Gsl.Matrix.to_arrays (m_of_cm (zgemm nr_s (zgemm f nr_s'))) let logm (m:Gsl.Matrix.matrix) = let n,n' = Gsl.Matrix.dims m