-
Notifications
You must be signed in to change notification settings - Fork 1
/
20120905b.py
203 lines (190 loc) · 6.45 KB
/
20120905b.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
"""
Try to reproduce fig (5) from a manuscript by Nasrallah.
"""
from StringIO import StringIO
import math
import numpy as np
from scipy import linalg
import Form
import FormOut
import RUtil
from RUtil import mk_call_str
import MatrixUtil
import wfengine
import wfcompens
import multinomstate
def get_form():
"""
@return: the body of a form
"""
return [
Form.ImageFormat(),
]
def get_form_out():
return FormOut.Image('plot')
def rev(x):
return tuple(reversed(list(x)))
def get_plot_array(N_diploid, Nr, theta_values, Ns_values):
"""
@param N_diploid: diploid population size
@param Nr: recombination rate
@param theta_values: mutation rates
@param Ns_values: selection values
@return: arr[i][j] gives time for Ns_values[i] and theta_values[j]
"""
# set up the state space
k = 4
M = multinomstate.get_sorted_states(2*N_diploid, k)
T = multinomstate.get_inverse_map(M)
nstates = M.shape[0]
lmcs = wfengine.get_lmcs(M)
# precompute rate matrices
R_rate = wfcompens.create_recomb(M, T)
M_rate = wfcompens.create_mutation(M, T)
# precompute a recombination probability matrix
R_prob = linalg.expm(Nr * R_rate / float((2*N_diploid)**2))
#
arr = []
for theta in theta_values:
# Compute the expected number of mutation events per generation.
mu = theta / 2
# Precompute the mutation matrix
# and the product of mutation and recombination.
M_prob = linalg.expm(mu * M_rate / float(2*2*N_diploid))
MR_prob = np.dot(M_prob, R_prob)
#
row = []
for Ns in Ns_values:
s = Ns / float(N_diploid)
lps = wfcompens.create_selection(s, M)
S_prob = np.exp(wfengine.create_genic(lmcs, lps, M))
P = np.dot(MR_prob, S_prob)
# compute the stochastic complement
X = linalg.solve(np.eye(nstates - k) - P[k:, k:], P[k:, :k])
H = P[:k, :k] + np.dot(P[:k, k:], X)
# condition on a transition change
#np.fill_diagonal(H, 0)
#H = (H.T / np.sum(H, axis=1)).T
#
#v = MatrixUtil.get_stationary_distribution(H)
#print 'reversibility check (0 if reversible):'
#print H/v - H.T/v
#print
#H_rev = (H*v).T/v
#if not np.allclose(
#[v[i]*H[i,j]-v[j]*H[j,i] for i in range(k) for j in range(k)],
#np.zeros(k)):
#raise ValueError('not reversible')
# break the last state into two states
AB, Ab, aB, ab, abx = 0, 1, 2, 3, 4
J = np.zeros((k+1, k+1))
J[:k, :k] = H
# force ab and abx to be absorbing states
J[ab] = np.zeros(k+1)
J[abx] = np.zeros(k+1)
J[ab, ab] = 1
J[abx, abx] = 1
# connect AB to the new ab state
J[AB, abx] = J[AB, ab]
J[AB, ab] = 0
# Now transition matrix J is in "canonical form"
# because all of the absorbing states are at the end.
B = linalg.solve(np.eye(k-1) - J[:k-1, :k-1], J[:k-1, k-1:])
# compute the absorbing state distribution
row.append(B[0, 1])
arr.append(row)
return arr
def get_plot(
side, Nr, arr, theta_values, Ns_values, xlim, ylim, ylogstr, ylab):
"""
@param side: 'left' or 'right'
"""
colors = ['blue', 'pink', 'green', 'red']
out = StringIO()
print >> out, 'Ns.values <- c', str(tuple(Ns_values))
print >> out, 'ha <- c', str(tuple(arr[0]))
print >> out, 'hb <- c', str(tuple(arr[1]))
print >> out, 'hc <- c', str(tuple(arr[2]))
print >> out, 'hd <- c', str(tuple(arr[3]))
print >> out, mk_call_str('plot', 'Ns.values', 'ha',
type='"l"',
xlab='"Ns"',
ylab=ylab,
log=ylogstr,
main='"Nr=%s"' % Nr,
xlim='c' + str(xlim),
ylim='c' + str(ylim),
col='"%s"' % colors[0],
)
print >> out, mk_call_str(
'lines', 'Ns.values', 'hb', col='"%s"' % colors[1])
print >> out, mk_call_str(
'lines', 'Ns.values', 'hc', col='"%s"' % colors[2])
print >> out, mk_call_str(
'lines', 'Ns.values', 'hd', col='"%s"' % colors[3])
if side == 'left':
print >> out, mk_call_str(
'legend',
'"topleft"',
'c' + str(tuple('%s' % x for x in theta_values)),
title='"theta"',
lty='c' + str(tuple([1]*4)),
lwd='c' + str(tuple([2.5]*4)),
col='c' + str(tuple(colors)),
)
return out.getvalue().rstrip()
def get_response_content(fs):
# define some fixed values
N_diploid = 6
N_hap = 2 * N_diploid
plot_density = 8
# define some mutation rates
theta_values = [0.001, 0.01, 0.1, 1.0]
# define some selection coefficients to plot
Ns_low = 0.0
Ns_high = 3.0
Ns_values = np.linspace(Ns_low, Ns_high, 3*plot_density + 1)
# get the values for each h
Nr_values = (0, 5)
arr_0 = get_plot_array(
N_diploid, Nr_values[0], theta_values, Ns_values)
arr_1 = get_plot_array(
N_diploid, Nr_values[1], theta_values, Ns_values)
ylab='"p"'
# define x and y plot limits
xlim = (Ns_low, Ns_high)
ylim = (np.min((arr_0, arr_1)), np.max((arr_0, arr_1)))
ylogstr = '""'
# http://sphaerula.com/legacy/R/multiplePlotFigure.html
out = StringIO()
print >> out, mk_call_str(
'par',
mfrow='c(1,2)',
oma='c(0,0,2,0)',
)
print >> out, get_plot(
'left', Nr_values[0], arr_0, theta_values, Ns_values,
xlim, ylim, ylogstr, ylab)
print >> out, get_plot(
'right', Nr_values[1], arr_1, theta_values, Ns_values,
xlim, ylim, ylogstr, '""')
print >> out, mk_call_str(
'title',
'"a conditional transition probability, 2N=%s"' % N_hap,
outer='TRUE',
)
script = out.getvalue().rstrip()
# create the R plot image
device_name = Form.g_imageformat_to_r_function[fs.imageformat]
retcode, r_out, r_err, image_data = RUtil.run_plotter_no_table(
script, device_name)
if retcode:
raise RUtil.RError(r_err)
return image_data
def main():
"""
Check some stuff that takes longer to compute.
"""
pass
if __name__ == '__main__':
main()