1"""
2This is an ML estimate of HKY85 molecular evolutionary parameters.
3"""
4
5import time
6
7import numpy as np
8import algopy
9from scipy import optimize, linalg
10
11
12# AGCT subsititution counts between human and chimp mitochondrial coding dna.
13g_data = np.array([
14        [2954, 141, 17, 16],
15        [165, 1110, 5, 2],
16        [18, 4, 3163, 374],
17        [15, 2, 310, 2411],
18        ],dtype=float)
19
20def transform_params(Y):
21    X = algopy.exp(Y)
22    tsrate, tvrate = X[0], X[1]
23    v_unnormalized = algopy.zeros(4, dtype=X)
24    v_unnormalized[0] = X[2]
25    v_unnormalized[1] = X[3]
26    v_unnormalized[2] = X[4]
27    v_unnormalized[3] = 1.0
28    v = v_unnormalized / algopy.sum(v_unnormalized)
29    return tsrate, tvrate, v
30
31def eval_f_orig(Y):
32    """ function as sent in the email by alex """
33    a, b, v = transform_params(Y)
34    Q = np.array([
35        [0, a, b, b],
36        [a, 0, b, b],
37        [b, b, 0, a],
38        [b, b, a, 0],
39        ])
40    Q = np.dot(Q, np.diag(v))
41    Q -= np.diag(np.sum(Q, axis=1))
42    S = np.log(np.dot(np.diag(v), linalg.expm(Q)))
43    return -np.sum(S * g_data)
44
45def eval_f(Y):
46    """ some reformulations to make eval_f_orig
47        compatible with algopy
48
49        missing: support for scipy.linalg.expm
50
51        i.e., this function can't be differentiated with algopy
52
53    """
54
55    a, b, v = transform_params(Y)
56
57    Q = algopy.zeros((4,4), dtype=Y)
58    Q[0,0] = 0;    Q[0,1] = a;    Q[0,2] = b;    Q[0,3] = b;
59    Q[1,0] = a;    Q[1,1] = 0;    Q[1,2] = b;    Q[1,3] = b;
60    Q[2,0] = b;    Q[2,1] = b;    Q[2,2] = 0;    Q[2,3] = a;
61    Q[3,0] = b;    Q[3,1] = b;    Q[3,2] = a;    Q[3,3] = 0;
62
63    Q = Q * v
64    Q -= algopy.diag(algopy.sum(Q, axis=1))
65    P = algopy.expm(Q)
66    S = algopy.log(algopy.dot(algopy.diag(v), P))
67    return -algopy.sum(S * g_data)
68
69def eval_f_eigh(Y):
70    """ some reformulations to make eval_f_orig
71        compatible with algopy
72
73        replaced scipy.linalg.expm by a symmetric eigenvalue decomposition
74
75        this function **can** be differentiated with algopy
76
77    """
78    a, b, v = transform_params(Y)
79
80    Q = algopy.zeros((4,4), dtype=Y)
81    Q[0,0] = 0;    Q[0,1] = a;    Q[0,2] = b;    Q[0,3] = b;
82    Q[1,0] = a;    Q[1,1] = 0;    Q[1,2] = b;    Q[1,3] = b;
83    Q[2,0] = b;    Q[2,1] = b;    Q[2,2] = 0;    Q[2,3] = a;
84    Q[3,0] = b;    Q[3,1] = b;    Q[3,2] = a;    Q[3,3] = 0;
85
86    Q = algopy.dot(Q, algopy.diag(v))
87    Q -= algopy.diag(algopy.sum(Q, axis=1))
88    va = algopy.diag(algopy.sqrt(v))
89    vb = algopy.diag(1./algopy.sqrt(v))
90    W, U = algopy.eigh(algopy.dot(algopy.dot(va, Q), vb))
91    M = algopy.dot(U, algopy.dot(algopy.diag(algopy.exp(W)), U.T))
92    P = algopy.dot(vb, algopy.dot(M, va))
93    S = algopy.log(algopy.dot(algopy.diag(v), P))
94    return -algopy.sum(S * g_data)
95
96def eval_grad_f_eigh(Y):
97    """
98    compute the gradient of f in the forward mode of AD
99    """
100    Y = algopy.UTPM.init_jacobian(Y)
101    retval = eval_f_eigh(Y)
102    return algopy.UTPM.extract_jacobian(retval)
103
104def eval_hess_f_eigh(Y):
105    """
106    compute the hessian of f in the forward mode of AD
107    """
108    Y = algopy.UTPM.init_hessian(Y)
109    retval = eval_f_eigh(Y)
110    hessian = algopy.UTPM.extract_hessian(5, retval)
111    return hessian
112
113def eval_grad_f(Y):
114    """
115    compute the gradient of f in the forward mode of AD
116    """
117    Y = algopy.UTPM.init_jacobian(Y)
118    retval = eval_f(Y)
119    return algopy.UTPM.extract_jacobian(retval)
120
121def eval_hess_f(Y):
122    """
123    compute the hessian of f in the forward mode of AD
124    """
125    Y = algopy.UTPM.init_hessian(Y)
126    retval = eval_f(Y)
127    hessian = algopy.UTPM.extract_hessian(5, retval)
128    return hessian
129
130
131def main():
132    Y = np.zeros(5)
133
134
135    print('--------------------------------')
136    print(' simple check (functions)       ')
137    print('--------------------------------')
138    print('eval_f_orig(Y)\n', eval_f_orig(Y))
139    print('eval_f(Y)\n', eval_f(Y))
140    print('eval_f_eigh(Y)\n', eval_f_eigh(Y))
141    print('eval_f_eigh(Y) - eval_f(Y)\n', eval_f_eigh(Y) - eval_f(Y))
142    print('--------------------------------')
143    print()
144
145
146    print('--------------------------------')
147    print(' simple check (gradients)       ')
148    print('--------------------------------')
149    print('eval_grad_f(Y)\n', eval_grad_f(Y))
150    print('eval_grad_f_eigh(Y)\n', eval_grad_f_eigh(Y))
151    print('eval_grad_f_eigh(Y) - eval_grad_f(Y)\n', eval_grad_f_eigh(Y) - eval_grad_f(Y))
152    print('--------------------------------')
153    print()
154
155
156    print('--------------------------------')
157    print(' simple check (hessians)        ')
158    print('--------------------------------')
159    print('eval_hess_f(Y)\n', eval_hess_f(Y))
160    print('eval_hess_f_eigh(Y)\n', eval_hess_f_eigh(Y))
161    print('eval_hess_f_eigh(Y) - eval_hess_f(Y)\n', eval_hess_f_eigh(Y) - eval_hess_f(Y))
162    print('--------------------------------')
163    print()
164
165    tm = time.time()
166    results = optimize.fmin(
167            eval_f, Y,
168            maxiter=10000, maxfun=10000, full_output=True)
169    tsrate, tvrate, v = transform_params(results[0])
170    print('--------------------------------')
171    print('time:', time.time() - tm)
172    print('results output from fmin:', results)
173    print('estimated transition rate parameter:', tsrate)
174    print('estimated transversion rate parameter:', tvrate)
175    print('estimated stationary distribution:', v)
176    print('--------------------------------')
177    print()
178
179    tm = time.time()
180    results = optimize.fmin_ncg(
181        #eval_f_eigh,             # obj. function
182        eval_f,
183        Y,                  # initial value
184        #eval_grad_f_eigh,   # gradient of obj. function
185        eval_grad_f,
186        fhess_p=None,
187        #fhess=eval_hess_f_eigh, # hessian of obj. function
188        fhess=eval_hess_f,
189        args=(),
190        #avextol=1e-04,
191        avextol=1e-05,
192        epsilon=1.4901161193847656e-08,
193        #maxiter=10,
194        maxiter=100,
195        full_output=True,
196        disp=1,
197        retall=0,
198        callback=None)
199
200    tsrate, tvrate, v = transform_params(results[0])
201    #hess = eval_hess_f_eigh(results[0])
202    hess = eval_hess_f(results[0])
203    print('--------------------------------')
204    print('time:', time.time() - tm)
205    print('results output from fmin:', results)
206    #print 'objective function value:', eval_f_eigh(results[0])
207    print('objective function value:', eval_f(results[0]))
208    #print 'gradient:', eval_grad_f_eigh(results[0])
209    print('gradient:', eval_grad_f(results[0]))
210    print('hessian:', hess)
211    print('hess - hess.T:', hess - hess.T)
212    print('eigvalsh(hess):', linalg.eigvalsh(hess))
213    print('inverse of hessian:', linalg.inv(hess))
214    print('determinant of hessian:', linalg.det(hess))
215    print('estimated transition rate parameter:', tsrate)
216    print('estimated transversion rate parameter:', tvrate)
217    print('estimated stationary distribution:', v)
218    print('--------------------------------')
219
220
221if __name__ == '__main__':
222    main()
223