forked from bluescarni/stark_weierstrass
-
Notifications
You must be signed in to change notification settings - Fork 0
/
weierstrass_elliptic.py
316 lines (315 loc) · 10.1 KB
/
weierstrass_elliptic.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
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
# -*- coding: iso-8859-1 -*-
# Copyright (C) 2013 by Francesco Biscani
# bluescarni@gmail.com
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the
# Free Software Foundation, Inc.,
# 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
class weierstrass_elliptic(object):
# TODO:
# - for testing, use results from Wolfram alpha with all possible branches (Delta positive and negative, g3 positive and negative)
# and both real and complex arguments.
def __cubic_roots(self,a,c,d):
from mpmath import mpf, polyroots
assert(all([isinstance(_,mpf) for _ in [a,c,d]]))
Delta = -4 * a * c*c*c - 27 * a*a * d*d
self.__Delta = Delta
# NOTE: replace with exact calculation of cubic roots.
proots, err = polyroots([a,0,c,d],error=True,maxsteps=100)
if Delta < 0:
# NOTE: here and below we ignore any residual imaginary part that we know must come from numerical artefacts.
# Sort the roots following the convention. proots[1] is the first complex root, proots[0] is the real one.
# The complex root with negative imaginary part is e3.
if proots[2].imag <= 0:
e1,e2,e3 = proots[1],proots[0].real,proots[2]
else:
e1,e2,e3 = proots[2],proots[0].real,proots[1]
else:
# The convention in this case is to sort in descending order.
e1,e2,e3 = sorted([_.real for _ in proots],reverse = True)
return e1,e2,e3
def __compute_periods(self):
# A+S 18.9.
from mpmath import sqrt, ellipk, mpc, pi, mpf
Delta = self.Delta
e1, e2, e3 = self.__roots
if Delta > 0:
m = (e2 - e3) / (e1 - e3)
Km = ellipk(m)
Kpm = ellipk(1 - m)
om = Km / sqrt(e1 - e3)
omp = mpc(0,1) * om * Kpm / Km
elif Delta < 0:
# NOTE: the expression in the sqrt has to be real and positive, as e1 and e3 are
# complex conjugate and e2 is real.
H2 = (sqrt((e2 - e3) * (e2 - e1))).real
assert(H2 > 0)
m = mpf(1) / mpf(2) - 3 * e2 / (4 * H2)
Km = ellipk(m)
Kpm = ellipk(1 - m)
om2 = Km / sqrt(H2)
om2p = mpc(0,1) * Kpm * om2 / Km
om = (om2 - om2p) / 2
omp = (om2 + om2p) / 2
else:
g2, g3 = self.__invariants
if g2 == 0 and g3 == 0:
om = mpf('+inf')
omp = mpc(0,'+inf')
else:
# NOTE: here there is no need for the dichotomy on the sign of g3 because
# we are already working in a regime in which g3 >= 0 by definition.
c = e1 / 2
om = 1 / sqrt(12 * c) * pi()
omp = mpc(0,'+inf')
return 2 * om, 2 * omp
def __compute_user_periods(self):
from mpmath import mpc
Delta = self.Delta
# NOTE: here there is no need to handle Delta == 0 separately,
# as it falls under the case of om purely real and omp purely imaginary.
if Delta >= 0:
T1, T3 = self.__periods
if self.__ng3:
return T3.imag, T1.real * mpc(0,1)
else:
return T1, T3
else:
T1, T3 = (sum(self.__periods)).real, self.__periods[1]
if self.__ng3:
return 2 * T3.imag, mpc(-T3.imag,T3.real)
else:
return T1, T3
def __compute_user_invariants(self):
if self.__ng3:
return self.__invariants[0],-self.__invariants[1]
else:
return self.__invariants
def __init__(self,g2,g3):
from mpmath import mpf
g2 = mpf(g2)
# Handle negative g3.
if g3 < 0:
g3 = -mpf(g3)
self.__ng3 = True
else:
g3 = mpf(g3)
self.__ng3 = False
assert(g3 >= 0)
# Store for future use.
self.__invariants = (g2,g3)
self.__user_invariants = self.__compute_user_invariants()
self.__roots = self.__cubic_roots(mpf(4),-g2,-g3)
self.__periods = self.__compute_periods()
self.__user_periods = self.__compute_user_periods()
@property
def invariants(self):
from copy import deepcopy
return deepcopy(self.__user_invariants)
@property
def Delta(self):
from copy import deepcopy
return deepcopy(self.__Delta)
@property
def periods(self):
from copy import deepcopy
return deepcopy(self.__user_periods)
@property
def roots(self):
from copy import deepcopy
if self.__ng3:
from functools import cmp_to_key
# g3 < 0 means that all roots change sign.
retval = [-x for x in self.__roots]
# Sort by imaginary part, then real.
return sorted(retval,key = cmp_to_key(lambda z1, z2: z1.imag - z2.imag if z1.imag != z2.imag else z1.real - z2.real),reverse=True)
else:
return deepcopy(self.__roots)
def __repr__(self):
retval = 'Invariants:\t' + str(self.invariants) + '\n'
# NOTE: the Delta does not change for differences in sign of g_3.
retval += 'Delta:\t\t' + str(self.Delta) + '\n'
retval += 'Periods:\t' + str(self.periods) + '\n'
retval += 'Roots:\t\t' + str(self.roots) + '\n'
return retval
def P(self,z):
# A+S 18.9.
from mpmath import sqrt, mpc, sin, ellipfun, mpf
Delta = self.Delta
e1, e2, e3 = self.__roots
if self.__ng3:
z = mpc(0,1) * z
if Delta > 0:
zs = sqrt(e1 - e3) * z
m = (e2 - e3) / (e1 - e3)
retval = e3 + (e1 - e3) / ellipfun('sn',u=zs,m=m)**2
elif Delta < 0:
H2 = (sqrt((e2 - e3) * (e2 - e1))).real
assert(H2 > 0)
m = mpf(1) / mpf(2) - 3 * e2 / (4 * H2)
zp = 2 * z * sqrt(H2)
retval = e2 + H2 * (1 + ellipfun('cn',u=zp,m=m)) / (1 - ellipfun('cn',u=zp,m=m))
else:
g2, g3 = self.__invariants
if g2 == 0 and g3 == 0:
retval = 1 / (z**2)
else:
c = e1 / 2
retval = -c + 3 * c / (sin(sqrt(3 * c) * z))**2
if self.__ng3:
return -retval
else:
return retval
def Pprime(self,z):
# A+S 18.9.
from mpmath import ellipfun, sqrt, cos, sin, mpc, mpf
Delta = self.Delta
e1, e2, e3 = self.__roots
if self.__ng3:
z = mpc(0,1) * z
if Delta > 0:
zs = sqrt(e1 - e3) * z
m = (e2 - e3) / (e1 - e3)
retval = -2 * sqrt((e1 - e3)**3) * ellipfun('cn',u=zs,m=m) * ellipfun('dn',u=zs,m=m) / (ellipfun('sn',u=zs,m=m)**3)
elif Delta < 0:
H2 = (sqrt((e2 - e3) * (e2 - e1))).real
assert(H2 > 0)
m = mpf(1) / mpf(2) - 3 * e2 / (4 * H2)
zp = 2 * z * sqrt(H2)
retval = -4 * sqrt(H2**3) * ellipfun('sn',u=zp,m=m) * ellipfun('dn',u=zp,m=m) / ((1 - ellipfun('cn',u=zp,m=m))**2)
else:
g2, g3 = self.__invariants
if g2 == 0 and g3 == 0:
retval = -2 / (z**3)
else:
c = e1 / 2
A = sqrt(3 * c)
retval = -6 * c * A * cos(A * z) / (sin(A * z))**3
if self.__ng3:
return mpc(0,-1) * retval
else:
return retval
def zeta(self,z):
# A+S 18.10.
from mpmath import pi, jtheta, exp, mpc, cot, sqrt
Delta = self.Delta
e1, _, _ = self.__roots
om = self.__periods[0] / 2
omp = self.__periods[1] / 2
if self.__ng3:
z = mpc(0,1) * z
if Delta > 0:
tau = omp / om
# NOTE: here q has to be real.
q = (exp(mpc(0,1) * pi() * tau)).real
eta = -(pi()**2 * jtheta(n=1,z=0,q=q,derivative=3)) / (12 * om * jtheta(n=1,z=0,q=q,derivative=1))
v = (pi() * z) / (2 * om)
retval = (eta * z) / om + (pi() * jtheta(n=1,z=v,q=q,derivative=1)) / (2 * om * jtheta(n=1,z=v,q=q))
elif Delta < 0:
om2 = om + omp
om2p = omp - om
tau2 = om2p / (2 * om2)
# NOTE: here q will be pure imaginary.
q = mpc(0,(mpc(0,1) * exp(mpc(0,1) * pi() * tau2)).imag)
eta2 = -(pi()**2 * jtheta(n=1,z=0,q=q,derivative=3)) / (12 * om2 * jtheta(n=1,z=0,q=q,derivative=1))
v = (pi() * z) / (2 * om2)
retval = (eta2 * z) / om2 + (pi() * jtheta(n=1,z=v,q=q,derivative=1)) / (2 * om2 * jtheta(n=1,z=v,q=q))
else:
g2, g3 = self.__invariants
if g2 == 0 and g3 == 0:
retval = 1 / z
else:
c = e1 / 2
A = sqrt(3 * c)
retval = c*z + A * cot(A*z)
if self.__ng3:
return mpc(0,1) * retval
else:
return retval
def sigma(self,z):
# A+S 18.10.
from mpmath import pi, jtheta, exp, mpc, sqrt, sin
Delta = self.Delta
e1, _, _ = self.__roots
om = self.__periods[0] / 2
omp = self.__periods[1] / 2
if self.__ng3:
z = mpc(0,1) * z
if Delta > 0:
tau = omp / om
q = (exp(mpc(0,1) * pi() * tau)).real
eta = -(pi()**2 * jtheta(n=1,z=0,q=q,derivative=3)) / (12 * om * jtheta(n=1,z=0,q=q,derivative=1))
v = (pi() * z) / (2 * om)
retval = (2 * om) / pi() * exp((eta * z**2)/(2 * om)) * jtheta(n=1,z=v,q=q)/jtheta(n=1,z=0,q=q,derivative=1)
elif Delta < 0:
om2 = om + omp
om2p = omp - om
tau2 = om2p / (2 * om2)
q = mpc(0,(mpc(0,1) * exp(mpc(0,1) * pi() * tau2)).imag)
eta2 = -(pi()**2 * jtheta(n=1,z=0,q=q,derivative=3)) / (12 * om2 * jtheta(n=1,z=0,q=q,derivative=1))
v = (pi() * z) / (2 * om2)
retval = (2 * om2) / pi() * exp((eta2 * z**2)/(2 * om2)) * jtheta(n=1,z=v,q=q)/jtheta(n=1,z=0,q=q,derivative=1)
else:
g2, g3 = self.__invariants
if g2 == 0 and g3 == 0:
retval = z
else:
c = e1 / 2
A = sqrt(3 * c)
retval = (1 / A) * sin(A*z) * exp((c*z**2) / 2)
if self.__ng3:
return mpc(0,-1) * retval
else:
return retval
def Pinv(self,P):
from mpmath import ellipf, sqrt, asin, acos, mpc, mpf
Delta = self.Delta
e1, e2, e3 = self.__roots
if self.__ng3:
P = -P
if Delta > 0:
m = (e2 - e3) / (e1 - e3)
retval = (1 / sqrt(e1 - e3)) * ellipf(asin(sqrt((e1 - e3)/(P - e3))),m=m)
elif Delta < 0:
H2 = (sqrt((e2 - e3) * (e2 - e1))).real
assert(H2 > 0)
m = mpf(1) / mpf(2) - 3 * e2 / (4 * H2)
retval = 1 / (2 * sqrt(H2)) * ellipf(acos((e2-P+H2)/(e2-P-H2)),m=m)
else:
g2, g3 = self.__invariants
if g2 == 0 and g3 == 0:
retval = 1 / sqrt(P)
else:
c = e1 / 2
retval = (1 / sqrt(3 * c)) * asin(sqrt((3 * c)/(P + c)))
if self.__ng3:
retval /= mpc(0,1)
alpha, beta, _, _ = self.reduce_to_fpp(retval)
T1, T2 = self.periods
return T1 * alpha + T2 * beta
def reduce_to_fpp(self,z):
# TODO: need to check what happens here when periods are infinite.
from mpmath import floor, matrix, lu_solve
T1, T2 = self.periods
R1, R2 = T1.real, T2.real
I1, I2 = T1.imag, T2.imag
A = matrix([[R1,R2],[I1,I2]])
b = matrix([z.real,z.imag])
x = lu_solve(A,b)
N = int(floor(x[0]))
M = int(floor(x[1]))
alpha = x[0] - N
beta = x[1] - M
assert(alpha >= 0 and beta >= 0)
return alpha,beta,N,M