Coverage for /usr/lib/python3/dist-packages/mpmath/functions/theta.py: 2%
895 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1from .functions import defun, defun_wrapped
3@defun
4def _jacobi_theta2(ctx, z, q):
5 extra1 = 10
6 extra2 = 20
7 # the loops below break when the fixed precision quantities
8 # a and b go to zero;
9 # right shifting small negative numbers by wp one obtains -1, not zero,
10 # so the condition a**2 + b**2 > MIN is used to break the loops.
11 MIN = 2
12 if z == ctx.zero:
13 if (not ctx._im(q)):
14 wp = ctx.prec + extra1
15 x = ctx.to_fixed(ctx._re(q), wp)
16 x2 = (x*x) >> wp
17 a = b = x2
18 s = x2
19 while abs(a) > MIN:
20 b = (b*x2) >> wp
21 a = (a*b) >> wp
22 s += a
23 s = (1 << (wp+1)) + (s << 1)
24 s = ctx.ldexp(s, -wp)
25 else:
26 wp = ctx.prec + extra1
27 xre = ctx.to_fixed(ctx._re(q), wp)
28 xim = ctx.to_fixed(ctx._im(q), wp)
29 x2re = (xre*xre - xim*xim) >> wp
30 x2im = (xre*xim) >> (wp-1)
31 are = bre = x2re
32 aim = bim = x2im
33 sre = (1<<wp) + are
34 sim = aim
35 while are**2 + aim**2 > MIN:
36 bre, bim = (bre * x2re - bim * x2im) >> wp, \
37 (bre * x2im + bim * x2re) >> wp
38 are, aim = (are * bre - aim * bim) >> wp, \
39 (are * bim + aim * bre) >> wp
40 sre += are
41 sim += aim
42 sre = (sre << 1)
43 sim = (sim << 1)
44 sre = ctx.ldexp(sre, -wp)
45 sim = ctx.ldexp(sim, -wp)
46 s = ctx.mpc(sre, sim)
47 else:
48 if (not ctx._im(q)) and (not ctx._im(z)):
49 wp = ctx.prec + extra1
50 x = ctx.to_fixed(ctx._re(q), wp)
51 x2 = (x*x) >> wp
52 a = b = x2
53 c1, s1 = ctx.cos_sin(ctx._re(z), prec=wp)
54 cn = c1 = ctx.to_fixed(c1, wp)
55 sn = s1 = ctx.to_fixed(s1, wp)
56 c2 = (c1*c1 - s1*s1) >> wp
57 s2 = (c1 * s1) >> (wp - 1)
58 cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
59 s = c1 + ((a * cn) >> wp)
60 while abs(a) > MIN:
61 b = (b*x2) >> wp
62 a = (a*b) >> wp
63 cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
64 s += (a * cn) >> wp
65 s = (s << 1)
66 s = ctx.ldexp(s, -wp)
67 s *= ctx.nthroot(q, 4)
68 return s
69 # case z real, q complex
70 elif not ctx._im(z):
71 wp = ctx.prec + extra2
72 xre = ctx.to_fixed(ctx._re(q), wp)
73 xim = ctx.to_fixed(ctx._im(q), wp)
74 x2re = (xre*xre - xim*xim) >> wp
75 x2im = (xre*xim) >> (wp - 1)
76 are = bre = x2re
77 aim = bim = x2im
78 c1, s1 = ctx.cos_sin(ctx._re(z), prec=wp)
79 cn = c1 = ctx.to_fixed(c1, wp)
80 sn = s1 = ctx.to_fixed(s1, wp)
81 c2 = (c1*c1 - s1*s1) >> wp
82 s2 = (c1 * s1) >> (wp - 1)
83 cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
84 sre = c1 + ((are * cn) >> wp)
85 sim = ((aim * cn) >> wp)
86 while are**2 + aim**2 > MIN:
87 bre, bim = (bre * x2re - bim * x2im) >> wp, \
88 (bre * x2im + bim * x2re) >> wp
89 are, aim = (are * bre - aim * bim) >> wp, \
90 (are * bim + aim * bre) >> wp
91 cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
92 sre += ((are * cn) >> wp)
93 sim += ((aim * cn) >> wp)
94 sre = (sre << 1)
95 sim = (sim << 1)
96 sre = ctx.ldexp(sre, -wp)
97 sim = ctx.ldexp(sim, -wp)
98 s = ctx.mpc(sre, sim)
99 #case z complex, q real
100 elif not ctx._im(q):
101 wp = ctx.prec + extra2
102 x = ctx.to_fixed(ctx._re(q), wp)
103 x2 = (x*x) >> wp
104 a = b = x2
105 prec0 = ctx.prec
106 ctx.prec = wp
107 c1, s1 = ctx.cos_sin(z)
108 ctx.prec = prec0
109 cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
110 cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
111 snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
112 snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
113 #c2 = (c1*c1 - s1*s1) >> wp
114 c2re = (c1re*c1re - c1im*c1im - s1re*s1re + s1im*s1im) >> wp
115 c2im = (c1re*c1im - s1re*s1im) >> (wp - 1)
116 #s2 = (c1 * s1) >> (wp - 1)
117 s2re = (c1re*s1re - c1im*s1im) >> (wp - 1)
118 s2im = (c1re*s1im + c1im*s1re) >> (wp - 1)
119 #cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
120 t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
121 t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
122 t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
123 t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
124 cnre = t1
125 cnim = t2
126 snre = t3
127 snim = t4
128 sre = c1re + ((a * cnre) >> wp)
129 sim = c1im + ((a * cnim) >> wp)
130 while abs(a) > MIN:
131 b = (b*x2) >> wp
132 a = (a*b) >> wp
133 t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
134 t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
135 t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
136 t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
137 cnre = t1
138 cnim = t2
139 snre = t3
140 snim = t4
141 sre += ((a * cnre) >> wp)
142 sim += ((a * cnim) >> wp)
143 sre = (sre << 1)
144 sim = (sim << 1)
145 sre = ctx.ldexp(sre, -wp)
146 sim = ctx.ldexp(sim, -wp)
147 s = ctx.mpc(sre, sim)
148 # case z and q complex
149 else:
150 wp = ctx.prec + extra2
151 xre = ctx.to_fixed(ctx._re(q), wp)
152 xim = ctx.to_fixed(ctx._im(q), wp)
153 x2re = (xre*xre - xim*xim) >> wp
154 x2im = (xre*xim) >> (wp - 1)
155 are = bre = x2re
156 aim = bim = x2im
157 prec0 = ctx.prec
158 ctx.prec = wp
159 # cos(z), sin(z) with z complex
160 c1, s1 = ctx.cos_sin(z)
161 ctx.prec = prec0
162 cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
163 cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
164 snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
165 snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
166 c2re = (c1re*c1re - c1im*c1im - s1re*s1re + s1im*s1im) >> wp
167 c2im = (c1re*c1im - s1re*s1im) >> (wp - 1)
168 s2re = (c1re*s1re - c1im*s1im) >> (wp - 1)
169 s2im = (c1re*s1im + c1im*s1re) >> (wp - 1)
170 t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
171 t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
172 t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
173 t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
174 cnre = t1
175 cnim = t2
176 snre = t3
177 snim = t4
178 n = 1
179 termre = c1re
180 termim = c1im
181 sre = c1re + ((are * cnre - aim * cnim) >> wp)
182 sim = c1im + ((are * cnim + aim * cnre) >> wp)
183 n = 3
184 termre = ((are * cnre - aim * cnim) >> wp)
185 termim = ((are * cnim + aim * cnre) >> wp)
186 sre = c1re + ((are * cnre - aim * cnim) >> wp)
187 sim = c1im + ((are * cnim + aim * cnre) >> wp)
188 n = 5
189 while are**2 + aim**2 > MIN:
190 bre, bim = (bre * x2re - bim * x2im) >> wp, \
191 (bre * x2im + bim * x2re) >> wp
192 are, aim = (are * bre - aim * bim) >> wp, \
193 (are * bim + aim * bre) >> wp
194 #cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
195 t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
196 t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
197 t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
198 t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
199 cnre = t1
200 cnim = t2
201 snre = t3
202 snim = t4
203 termre = ((are * cnre - aim * cnim) >> wp)
204 termim = ((aim * cnre + are * cnim) >> wp)
205 sre += ((are * cnre - aim * cnim) >> wp)
206 sim += ((aim * cnre + are * cnim) >> wp)
207 n += 2
208 sre = (sre << 1)
209 sim = (sim << 1)
210 sre = ctx.ldexp(sre, -wp)
211 sim = ctx.ldexp(sim, -wp)
212 s = ctx.mpc(sre, sim)
213 s *= ctx.nthroot(q, 4)
214 return s
216@defun
217def _djacobi_theta2(ctx, z, q, nd):
218 MIN = 2
219 extra1 = 10
220 extra2 = 20
221 if (not ctx._im(q)) and (not ctx._im(z)):
222 wp = ctx.prec + extra1
223 x = ctx.to_fixed(ctx._re(q), wp)
224 x2 = (x*x) >> wp
225 a = b = x2
226 c1, s1 = ctx.cos_sin(ctx._re(z), prec=wp)
227 cn = c1 = ctx.to_fixed(c1, wp)
228 sn = s1 = ctx.to_fixed(s1, wp)
229 c2 = (c1*c1 - s1*s1) >> wp
230 s2 = (c1 * s1) >> (wp - 1)
231 cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
232 if (nd&1):
233 s = s1 + ((a * sn * 3**nd) >> wp)
234 else:
235 s = c1 + ((a * cn * 3**nd) >> wp)
236 n = 2
237 while abs(a) > MIN:
238 b = (b*x2) >> wp
239 a = (a*b) >> wp
240 cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
241 if nd&1:
242 s += (a * sn * (2*n+1)**nd) >> wp
243 else:
244 s += (a * cn * (2*n+1)**nd) >> wp
245 n += 1
246 s = -(s << 1)
247 s = ctx.ldexp(s, -wp)
248 # case z real, q complex
249 elif not ctx._im(z):
250 wp = ctx.prec + extra2
251 xre = ctx.to_fixed(ctx._re(q), wp)
252 xim = ctx.to_fixed(ctx._im(q), wp)
253 x2re = (xre*xre - xim*xim) >> wp
254 x2im = (xre*xim) >> (wp - 1)
255 are = bre = x2re
256 aim = bim = x2im
257 c1, s1 = ctx.cos_sin(ctx._re(z), prec=wp)
258 cn = c1 = ctx.to_fixed(c1, wp)
259 sn = s1 = ctx.to_fixed(s1, wp)
260 c2 = (c1*c1 - s1*s1) >> wp
261 s2 = (c1 * s1) >> (wp - 1)
262 cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
263 if (nd&1):
264 sre = s1 + ((are * sn * 3**nd) >> wp)
265 sim = ((aim * sn * 3**nd) >> wp)
266 else:
267 sre = c1 + ((are * cn * 3**nd) >> wp)
268 sim = ((aim * cn * 3**nd) >> wp)
269 n = 5
270 while are**2 + aim**2 > MIN:
271 bre, bim = (bre * x2re - bim * x2im) >> wp, \
272 (bre * x2im + bim * x2re) >> wp
273 are, aim = (are * bre - aim * bim) >> wp, \
274 (are * bim + aim * bre) >> wp
275 cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
277 if (nd&1):
278 sre += ((are * sn * n**nd) >> wp)
279 sim += ((aim * sn * n**nd) >> wp)
280 else:
281 sre += ((are * cn * n**nd) >> wp)
282 sim += ((aim * cn * n**nd) >> wp)
283 n += 2
284 sre = -(sre << 1)
285 sim = -(sim << 1)
286 sre = ctx.ldexp(sre, -wp)
287 sim = ctx.ldexp(sim, -wp)
288 s = ctx.mpc(sre, sim)
289 #case z complex, q real
290 elif not ctx._im(q):
291 wp = ctx.prec + extra2
292 x = ctx.to_fixed(ctx._re(q), wp)
293 x2 = (x*x) >> wp
294 a = b = x2
295 prec0 = ctx.prec
296 ctx.prec = wp
297 c1, s1 = ctx.cos_sin(z)
298 ctx.prec = prec0
299 cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
300 cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
301 snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
302 snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
303 #c2 = (c1*c1 - s1*s1) >> wp
304 c2re = (c1re*c1re - c1im*c1im - s1re*s1re + s1im*s1im) >> wp
305 c2im = (c1re*c1im - s1re*s1im) >> (wp - 1)
306 #s2 = (c1 * s1) >> (wp - 1)
307 s2re = (c1re*s1re - c1im*s1im) >> (wp - 1)
308 s2im = (c1re*s1im + c1im*s1re) >> (wp - 1)
309 #cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
310 t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
311 t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
312 t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
313 t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
314 cnre = t1
315 cnim = t2
316 snre = t3
317 snim = t4
318 if (nd&1):
319 sre = s1re + ((a * snre * 3**nd) >> wp)
320 sim = s1im + ((a * snim * 3**nd) >> wp)
321 else:
322 sre = c1re + ((a * cnre * 3**nd) >> wp)
323 sim = c1im + ((a * cnim * 3**nd) >> wp)
324 n = 5
325 while abs(a) > MIN:
326 b = (b*x2) >> wp
327 a = (a*b) >> wp
328 t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
329 t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
330 t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
331 t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
332 cnre = t1
333 cnim = t2
334 snre = t3
335 snim = t4
336 if (nd&1):
337 sre += ((a * snre * n**nd) >> wp)
338 sim += ((a * snim * n**nd) >> wp)
339 else:
340 sre += ((a * cnre * n**nd) >> wp)
341 sim += ((a * cnim * n**nd) >> wp)
342 n += 2
343 sre = -(sre << 1)
344 sim = -(sim << 1)
345 sre = ctx.ldexp(sre, -wp)
346 sim = ctx.ldexp(sim, -wp)
347 s = ctx.mpc(sre, sim)
348 # case z and q complex
349 else:
350 wp = ctx.prec + extra2
351 xre = ctx.to_fixed(ctx._re(q), wp)
352 xim = ctx.to_fixed(ctx._im(q), wp)
353 x2re = (xre*xre - xim*xim) >> wp
354 x2im = (xre*xim) >> (wp - 1)
355 are = bre = x2re
356 aim = bim = x2im
357 prec0 = ctx.prec
358 ctx.prec = wp
359 # cos(2*z), sin(2*z) with z complex
360 c1, s1 = ctx.cos_sin(z)
361 ctx.prec = prec0
362 cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
363 cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
364 snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
365 snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
366 c2re = (c1re*c1re - c1im*c1im - s1re*s1re + s1im*s1im) >> wp
367 c2im = (c1re*c1im - s1re*s1im) >> (wp - 1)
368 s2re = (c1re*s1re - c1im*s1im) >> (wp - 1)
369 s2im = (c1re*s1im + c1im*s1re) >> (wp - 1)
370 t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
371 t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
372 t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
373 t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
374 cnre = t1
375 cnim = t2
376 snre = t3
377 snim = t4
378 if (nd&1):
379 sre = s1re + (((are * snre - aim * snim) * 3**nd) >> wp)
380 sim = s1im + (((are * snim + aim * snre)* 3**nd) >> wp)
381 else:
382 sre = c1re + (((are * cnre - aim * cnim) * 3**nd) >> wp)
383 sim = c1im + (((are * cnim + aim * cnre)* 3**nd) >> wp)
384 n = 5
385 while are**2 + aim**2 > MIN:
386 bre, bim = (bre * x2re - bim * x2im) >> wp, \
387 (bre * x2im + bim * x2re) >> wp
388 are, aim = (are * bre - aim * bim) >> wp, \
389 (are * bim + aim * bre) >> wp
390 #cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
391 t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
392 t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
393 t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
394 t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
395 cnre = t1
396 cnim = t2
397 snre = t3
398 snim = t4
399 if (nd&1):
400 sre += (((are * snre - aim * snim) * n**nd) >> wp)
401 sim += (((aim * snre + are * snim) * n**nd) >> wp)
402 else:
403 sre += (((are * cnre - aim * cnim) * n**nd) >> wp)
404 sim += (((aim * cnre + are * cnim) * n**nd) >> wp)
405 n += 2
406 sre = -(sre << 1)
407 sim = -(sim << 1)
408 sre = ctx.ldexp(sre, -wp)
409 sim = ctx.ldexp(sim, -wp)
410 s = ctx.mpc(sre, sim)
411 s *= ctx.nthroot(q, 4)
412 if (nd&1):
413 return (-1)**(nd//2) * s
414 else:
415 return (-1)**(1 + nd//2) * s
417@defun
418def _jacobi_theta3(ctx, z, q):
419 extra1 = 10
420 extra2 = 20
421 MIN = 2
422 if z == ctx.zero:
423 if not ctx._im(q):
424 wp = ctx.prec + extra1
425 x = ctx.to_fixed(ctx._re(q), wp)
426 s = x
427 a = b = x
428 x2 = (x*x) >> wp
429 while abs(a) > MIN:
430 b = (b*x2) >> wp
431 a = (a*b) >> wp
432 s += a
433 s = (1 << wp) + (s << 1)
434 s = ctx.ldexp(s, -wp)
435 return s
436 else:
437 wp = ctx.prec + extra1
438 xre = ctx.to_fixed(ctx._re(q), wp)
439 xim = ctx.to_fixed(ctx._im(q), wp)
440 x2re = (xre*xre - xim*xim) >> wp
441 x2im = (xre*xim) >> (wp - 1)
442 sre = are = bre = xre
443 sim = aim = bim = xim
444 while are**2 + aim**2 > MIN:
445 bre, bim = (bre * x2re - bim * x2im) >> wp, \
446 (bre * x2im + bim * x2re) >> wp
447 are, aim = (are * bre - aim * bim) >> wp, \
448 (are * bim + aim * bre) >> wp
449 sre += are
450 sim += aim
451 sre = (1 << wp) + (sre << 1)
452 sim = (sim << 1)
453 sre = ctx.ldexp(sre, -wp)
454 sim = ctx.ldexp(sim, -wp)
455 s = ctx.mpc(sre, sim)
456 return s
457 else:
458 if (not ctx._im(q)) and (not ctx._im(z)):
459 s = 0
460 wp = ctx.prec + extra1
461 x = ctx.to_fixed(ctx._re(q), wp)
462 a = b = x
463 x2 = (x*x) >> wp
464 c1, s1 = ctx.cos_sin(ctx._re(z)*2, prec=wp)
465 c1 = ctx.to_fixed(c1, wp)
466 s1 = ctx.to_fixed(s1, wp)
467 cn = c1
468 sn = s1
469 s += (a * cn) >> wp
470 while abs(a) > MIN:
471 b = (b*x2) >> wp
472 a = (a*b) >> wp
473 cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
474 s += (a * cn) >> wp
475 s = (1 << wp) + (s << 1)
476 s = ctx.ldexp(s, -wp)
477 return s
478 # case z real, q complex
479 elif not ctx._im(z):
480 wp = ctx.prec + extra2
481 xre = ctx.to_fixed(ctx._re(q), wp)
482 xim = ctx.to_fixed(ctx._im(q), wp)
483 x2re = (xre*xre - xim*xim) >> wp
484 x2im = (xre*xim) >> (wp - 1)
485 are = bre = xre
486 aim = bim = xim
487 c1, s1 = ctx.cos_sin(ctx._re(z)*2, prec=wp)
488 c1 = ctx.to_fixed(c1, wp)
489 s1 = ctx.to_fixed(s1, wp)
490 cn = c1
491 sn = s1
492 sre = (are * cn) >> wp
493 sim = (aim * cn) >> wp
494 while are**2 + aim**2 > MIN:
495 bre, bim = (bre * x2re - bim * x2im) >> wp, \
496 (bre * x2im + bim * x2re) >> wp
497 are, aim = (are * bre - aim * bim) >> wp, \
498 (are * bim + aim * bre) >> wp
499 cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
500 sre += (are * cn) >> wp
501 sim += (aim * cn) >> wp
502 sre = (1 << wp) + (sre << 1)
503 sim = (sim << 1)
504 sre = ctx.ldexp(sre, -wp)
505 sim = ctx.ldexp(sim, -wp)
506 s = ctx.mpc(sre, sim)
507 return s
508 #case z complex, q real
509 elif not ctx._im(q):
510 wp = ctx.prec + extra2
511 x = ctx.to_fixed(ctx._re(q), wp)
512 a = b = x
513 x2 = (x*x) >> wp
514 prec0 = ctx.prec
515 ctx.prec = wp
516 c1, s1 = ctx.cos_sin(2*z)
517 ctx.prec = prec0
518 cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
519 cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
520 snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
521 snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
522 sre = (a * cnre) >> wp
523 sim = (a * cnim) >> wp
524 while abs(a) > MIN:
525 b = (b*x2) >> wp
526 a = (a*b) >> wp
527 t1 = (cnre*c1re - cnim*c1im - snre*s1re + snim*s1im) >> wp
528 t2 = (cnre*c1im + cnim*c1re - snre*s1im - snim*s1re) >> wp
529 t3 = (snre*c1re - snim*c1im + cnre*s1re - cnim*s1im) >> wp
530 t4 = (snre*c1im + snim*c1re + cnre*s1im + cnim*s1re) >> wp
531 cnre = t1
532 cnim = t2
533 snre = t3
534 snim = t4
535 sre += (a * cnre) >> wp
536 sim += (a * cnim) >> wp
537 sre = (1 << wp) + (sre << 1)
538 sim = (sim << 1)
539 sre = ctx.ldexp(sre, -wp)
540 sim = ctx.ldexp(sim, -wp)
541 s = ctx.mpc(sre, sim)
542 return s
543 # case z and q complex
544 else:
545 wp = ctx.prec + extra2
546 xre = ctx.to_fixed(ctx._re(q), wp)
547 xim = ctx.to_fixed(ctx._im(q), wp)
548 x2re = (xre*xre - xim*xim) >> wp
549 x2im = (xre*xim) >> (wp - 1)
550 are = bre = xre
551 aim = bim = xim
552 prec0 = ctx.prec
553 ctx.prec = wp
554 # cos(2*z), sin(2*z) with z complex
555 c1, s1 = ctx.cos_sin(2*z)
556 ctx.prec = prec0
557 cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
558 cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
559 snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
560 snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
561 sre = (are * cnre - aim * cnim) >> wp
562 sim = (aim * cnre + are * cnim) >> wp
563 while are**2 + aim**2 > MIN:
564 bre, bim = (bre * x2re - bim * x2im) >> wp, \
565 (bre * x2im + bim * x2re) >> wp
566 are, aim = (are * bre - aim * bim) >> wp, \
567 (are * bim + aim * bre) >> wp
568 t1 = (cnre*c1re - cnim*c1im - snre*s1re + snim*s1im) >> wp
569 t2 = (cnre*c1im + cnim*c1re - snre*s1im - snim*s1re) >> wp
570 t3 = (snre*c1re - snim*c1im + cnre*s1re - cnim*s1im) >> wp
571 t4 = (snre*c1im + snim*c1re + cnre*s1im + cnim*s1re) >> wp
572 cnre = t1
573 cnim = t2
574 snre = t3
575 snim = t4
576 sre += (are * cnre - aim * cnim) >> wp
577 sim += (aim * cnre + are * cnim) >> wp
578 sre = (1 << wp) + (sre << 1)
579 sim = (sim << 1)
580 sre = ctx.ldexp(sre, -wp)
581 sim = ctx.ldexp(sim, -wp)
582 s = ctx.mpc(sre, sim)
583 return s
585@defun
586def _djacobi_theta3(ctx, z, q, nd):
587 """nd=1,2,3 order of the derivative with respect to z"""
588 MIN = 2
589 extra1 = 10
590 extra2 = 20
591 if (not ctx._im(q)) and (not ctx._im(z)):
592 s = 0
593 wp = ctx.prec + extra1
594 x = ctx.to_fixed(ctx._re(q), wp)
595 a = b = x
596 x2 = (x*x) >> wp
597 c1, s1 = ctx.cos_sin(ctx._re(z)*2, prec=wp)
598 c1 = ctx.to_fixed(c1, wp)
599 s1 = ctx.to_fixed(s1, wp)
600 cn = c1
601 sn = s1
602 if (nd&1):
603 s += (a * sn) >> wp
604 else:
605 s += (a * cn) >> wp
606 n = 2
607 while abs(a) > MIN:
608 b = (b*x2) >> wp
609 a = (a*b) >> wp
610 cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
611 if nd&1:
612 s += (a * sn * n**nd) >> wp
613 else:
614 s += (a * cn * n**nd) >> wp
615 n += 1
616 s = -(s << (nd+1))
617 s = ctx.ldexp(s, -wp)
618 # case z real, q complex
619 elif not ctx._im(z):
620 wp = ctx.prec + extra2
621 xre = ctx.to_fixed(ctx._re(q), wp)
622 xim = ctx.to_fixed(ctx._im(q), wp)
623 x2re = (xre*xre - xim*xim) >> wp
624 x2im = (xre*xim) >> (wp - 1)
625 are = bre = xre
626 aim = bim = xim
627 c1, s1 = ctx.cos_sin(ctx._re(z)*2, prec=wp)
628 c1 = ctx.to_fixed(c1, wp)
629 s1 = ctx.to_fixed(s1, wp)
630 cn = c1
631 sn = s1
632 if (nd&1):
633 sre = (are * sn) >> wp
634 sim = (aim * sn) >> wp
635 else:
636 sre = (are * cn) >> wp
637 sim = (aim * cn) >> wp
638 n = 2
639 while are**2 + aim**2 > MIN:
640 bre, bim = (bre * x2re - bim * x2im) >> wp, \
641 (bre * x2im + bim * x2re) >> wp
642 are, aim = (are * bre - aim * bim) >> wp, \
643 (are * bim + aim * bre) >> wp
644 cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
645 if nd&1:
646 sre += (are * sn * n**nd) >> wp
647 sim += (aim * sn * n**nd) >> wp
648 else:
649 sre += (are * cn * n**nd) >> wp
650 sim += (aim * cn * n**nd) >> wp
651 n += 1
652 sre = -(sre << (nd+1))
653 sim = -(sim << (nd+1))
654 sre = ctx.ldexp(sre, -wp)
655 sim = ctx.ldexp(sim, -wp)
656 s = ctx.mpc(sre, sim)
657 #case z complex, q real
658 elif not ctx._im(q):
659 wp = ctx.prec + extra2
660 x = ctx.to_fixed(ctx._re(q), wp)
661 a = b = x
662 x2 = (x*x) >> wp
663 prec0 = ctx.prec
664 ctx.prec = wp
665 c1, s1 = ctx.cos_sin(2*z)
666 ctx.prec = prec0
667 cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
668 cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
669 snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
670 snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
671 if (nd&1):
672 sre = (a * snre) >> wp
673 sim = (a * snim) >> wp
674 else:
675 sre = (a * cnre) >> wp
676 sim = (a * cnim) >> wp
677 n = 2
678 while abs(a) > MIN:
679 b = (b*x2) >> wp
680 a = (a*b) >> wp
681 t1 = (cnre*c1re - cnim*c1im - snre*s1re + snim*s1im) >> wp
682 t2 = (cnre*c1im + cnim*c1re - snre*s1im - snim*s1re) >> wp
683 t3 = (snre*c1re - snim*c1im + cnre*s1re - cnim*s1im) >> wp
684 t4 = (snre*c1im + snim*c1re + cnre*s1im + cnim*s1re) >> wp
685 cnre = t1
686 cnim = t2
687 snre = t3
688 snim = t4
689 if (nd&1):
690 sre += (a * snre * n**nd) >> wp
691 sim += (a * snim * n**nd) >> wp
692 else:
693 sre += (a * cnre * n**nd) >> wp
694 sim += (a * cnim * n**nd) >> wp
695 n += 1
696 sre = -(sre << (nd+1))
697 sim = -(sim << (nd+1))
698 sre = ctx.ldexp(sre, -wp)
699 sim = ctx.ldexp(sim, -wp)
700 s = ctx.mpc(sre, sim)
701 # case z and q complex
702 else:
703 wp = ctx.prec + extra2
704 xre = ctx.to_fixed(ctx._re(q), wp)
705 xim = ctx.to_fixed(ctx._im(q), wp)
706 x2re = (xre*xre - xim*xim) >> wp
707 x2im = (xre*xim) >> (wp - 1)
708 are = bre = xre
709 aim = bim = xim
710 prec0 = ctx.prec
711 ctx.prec = wp
712 # cos(2*z), sin(2*z) with z complex
713 c1, s1 = ctx.cos_sin(2*z)
714 ctx.prec = prec0
715 cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
716 cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
717 snre = s1re = ctx.to_fixed(ctx._re(s1), wp)
718 snim = s1im = ctx.to_fixed(ctx._im(s1), wp)
719 if (nd&1):
720 sre = (are * snre - aim * snim) >> wp
721 sim = (aim * snre + are * snim) >> wp
722 else:
723 sre = (are * cnre - aim * cnim) >> wp
724 sim = (aim * cnre + are * cnim) >> wp
725 n = 2
726 while are**2 + aim**2 > MIN:
727 bre, bim = (bre * x2re - bim * x2im) >> wp, \
728 (bre * x2im + bim * x2re) >> wp
729 are, aim = (are * bre - aim * bim) >> wp, \
730 (are * bim + aim * bre) >> wp
731 t1 = (cnre*c1re - cnim*c1im - snre*s1re + snim*s1im) >> wp
732 t2 = (cnre*c1im + cnim*c1re - snre*s1im - snim*s1re) >> wp
733 t3 = (snre*c1re - snim*c1im + cnre*s1re - cnim*s1im) >> wp
734 t4 = (snre*c1im + snim*c1re + cnre*s1im + cnim*s1re) >> wp
735 cnre = t1
736 cnim = t2
737 snre = t3
738 snim = t4
739 if(nd&1):
740 sre += ((are * snre - aim * snim) * n**nd) >> wp
741 sim += ((aim * snre + are * snim) * n**nd) >> wp
742 else:
743 sre += ((are * cnre - aim * cnim) * n**nd) >> wp
744 sim += ((aim * cnre + are * cnim) * n**nd) >> wp
745 n += 1
746 sre = -(sre << (nd+1))
747 sim = -(sim << (nd+1))
748 sre = ctx.ldexp(sre, -wp)
749 sim = ctx.ldexp(sim, -wp)
750 s = ctx.mpc(sre, sim)
751 if (nd&1):
752 return (-1)**(nd//2) * s
753 else:
754 return (-1)**(1 + nd//2) * s
756@defun
757def _jacobi_theta2a(ctx, z, q):
758 """
759 case ctx._im(z) != 0
760 theta(2, z, q) =
761 q**1/4 * Sum(q**(n*n + n) * exp(j*(2*n + 1)*z), n=-inf, inf)
762 max term for minimum (2*n+1)*log(q).real - 2* ctx._im(z)
763 n0 = int(ctx._im(z)/log(q).real - 1/2)
764 theta(2, z, q) =
765 q**1/4 * Sum(q**(n*n + n) * exp(j*(2*n + 1)*z), n=n0, inf) +
766 q**1/4 * Sum(q**(n*n + n) * exp(j*(2*n + 1)*z), n, n0-1, -inf)
767 """
768 n = n0 = int(ctx._im(z)/ctx._re(ctx.log(q)) - 1/2)
769 e2 = ctx.expj(2*z)
770 e = e0 = ctx.expj((2*n+1)*z)
771 a = q**(n*n + n)
772 # leading term
773 term = a * e
774 s = term
775 eps1 = ctx.eps*abs(term)
776 while 1:
777 n += 1
778 e = e * e2
779 term = q**(n*n + n) * e
780 if abs(term) < eps1:
781 break
782 s += term
783 e = e0
784 e2 = ctx.expj(-2*z)
785 n = n0
786 while 1:
787 n -= 1
788 e = e * e2
789 term = q**(n*n + n) * e
790 if abs(term) < eps1:
791 break
792 s += term
793 s = s * ctx.nthroot(q, 4)
794 return s
796@defun
797def _jacobi_theta3a(ctx, z, q):
798 """
799 case ctx._im(z) != 0
800 theta3(z, q) = Sum(q**(n*n) * exp(j*2*n*z), n, -inf, inf)
801 max term for n*abs(log(q).real) + ctx._im(z) ~= 0
802 n0 = int(- ctx._im(z)/abs(log(q).real))
803 """
804 n = n0 = int(-ctx._im(z)/abs(ctx._re(ctx.log(q))))
805 e2 = ctx.expj(2*z)
806 e = e0 = ctx.expj(2*n*z)
807 s = term = q**(n*n) * e
808 eps1 = ctx.eps*abs(term)
809 while 1:
810 n += 1
811 e = e * e2
812 term = q**(n*n) * e
813 if abs(term) < eps1:
814 break
815 s += term
816 e = e0
817 e2 = ctx.expj(-2*z)
818 n = n0
819 while 1:
820 n -= 1
821 e = e * e2
822 term = q**(n*n) * e
823 if abs(term) < eps1:
824 break
825 s += term
826 return s
828@defun
829def _djacobi_theta2a(ctx, z, q, nd):
830 """
831 case ctx._im(z) != 0
832 dtheta(2, z, q, nd) =
833 j* q**1/4 * Sum(q**(n*n + n) * (2*n+1)*exp(j*(2*n + 1)*z), n=-inf, inf)
834 max term for (2*n0+1)*log(q).real - 2* ctx._im(z) ~= 0
835 n0 = int(ctx._im(z)/log(q).real - 1/2)
836 """
837 n = n0 = int(ctx._im(z)/ctx._re(ctx.log(q)) - 1/2)
838 e2 = ctx.expj(2*z)
839 e = e0 = ctx.expj((2*n + 1)*z)
840 a = q**(n*n + n)
841 # leading term
842 term = (2*n+1)**nd * a * e
843 s = term
844 eps1 = ctx.eps*abs(term)
845 while 1:
846 n += 1
847 e = e * e2
848 term = (2*n+1)**nd * q**(n*n + n) * e
849 if abs(term) < eps1:
850 break
851 s += term
852 e = e0
853 e2 = ctx.expj(-2*z)
854 n = n0
855 while 1:
856 n -= 1
857 e = e * e2
858 term = (2*n+1)**nd * q**(n*n + n) * e
859 if abs(term) < eps1:
860 break
861 s += term
862 return ctx.j**nd * s * ctx.nthroot(q, 4)
864@defun
865def _djacobi_theta3a(ctx, z, q, nd):
866 """
867 case ctx._im(z) != 0
868 djtheta3(z, q, nd) = (2*j)**nd *
869 Sum(q**(n*n) * n**nd * exp(j*2*n*z), n, -inf, inf)
870 max term for minimum n*abs(log(q).real) + ctx._im(z)
871 """
872 n = n0 = int(-ctx._im(z)/abs(ctx._re(ctx.log(q))))
873 e2 = ctx.expj(2*z)
874 e = e0 = ctx.expj(2*n*z)
875 a = q**(n*n) * e
876 s = term = n**nd * a
877 if n != 0:
878 eps1 = ctx.eps*abs(term)
879 else:
880 eps1 = ctx.eps*abs(a)
881 while 1:
882 n += 1
883 e = e * e2
884 a = q**(n*n) * e
885 term = n**nd * a
886 if n != 0:
887 aterm = abs(term)
888 else:
889 aterm = abs(a)
890 if aterm < eps1:
891 break
892 s += term
893 e = e0
894 e2 = ctx.expj(-2*z)
895 n = n0
896 while 1:
897 n -= 1
898 e = e * e2
899 a = q**(n*n) * e
900 term = n**nd * a
901 if n != 0:
902 aterm = abs(term)
903 else:
904 aterm = abs(a)
905 if aterm < eps1:
906 break
907 s += term
908 return (2*ctx.j)**nd * s
910@defun
911def jtheta(ctx, n, z, q, derivative=0):
912 if derivative:
913 return ctx._djtheta(n, z, q, derivative)
915 z = ctx.convert(z)
916 q = ctx.convert(q)
918 # Implementation note
919 # If ctx._im(z) is close to zero, _jacobi_theta2 and _jacobi_theta3
920 # are used,
921 # which compute the series starting from n=0 using fixed precision
922 # numbers;
923 # otherwise _jacobi_theta2a and _jacobi_theta3a are used, which compute
924 # the series starting from n=n0, which is the largest term.
926 # TODO: write _jacobi_theta2a and _jacobi_theta3a using fixed-point
928 if abs(q) > ctx.THETA_Q_LIM:
929 raise ValueError('abs(q) > THETA_Q_LIM = %f' % ctx.THETA_Q_LIM)
931 extra = 10
932 if z:
933 M = ctx.mag(z)
934 if M > 5 or (n == 1 and M < -5):
935 extra += 2*abs(M)
936 cz = 0.5
937 extra2 = 50
938 prec0 = ctx.prec
939 try:
940 ctx.prec += extra
941 if n == 1:
942 if ctx._im(z):
943 if abs(ctx._im(z)) < cz * abs(ctx._re(ctx.log(q))):
944 ctx.dps += extra2
945 res = ctx._jacobi_theta2(z - ctx.pi/2, q)
946 else:
947 ctx.dps += 10
948 res = ctx._jacobi_theta2a(z - ctx.pi/2, q)
949 else:
950 res = ctx._jacobi_theta2(z - ctx.pi/2, q)
951 elif n == 2:
952 if ctx._im(z):
953 if abs(ctx._im(z)) < cz * abs(ctx._re(ctx.log(q))):
954 ctx.dps += extra2
955 res = ctx._jacobi_theta2(z, q)
956 else:
957 ctx.dps += 10
958 res = ctx._jacobi_theta2a(z, q)
959 else:
960 res = ctx._jacobi_theta2(z, q)
961 elif n == 3:
962 if ctx._im(z):
963 if abs(ctx._im(z)) < cz * abs(ctx._re(ctx.log(q))):
964 ctx.dps += extra2
965 res = ctx._jacobi_theta3(z, q)
966 else:
967 ctx.dps += 10
968 res = ctx._jacobi_theta3a(z, q)
969 else:
970 res = ctx._jacobi_theta3(z, q)
971 elif n == 4:
972 if ctx._im(z):
973 if abs(ctx._im(z)) < cz * abs(ctx._re(ctx.log(q))):
974 ctx.dps += extra2
975 res = ctx._jacobi_theta3(z, -q)
976 else:
977 ctx.dps += 10
978 res = ctx._jacobi_theta3a(z, -q)
979 else:
980 res = ctx._jacobi_theta3(z, -q)
981 else:
982 raise ValueError
983 finally:
984 ctx.prec = prec0
985 return res
987@defun
988def _djtheta(ctx, n, z, q, derivative=1):
989 z = ctx.convert(z)
990 q = ctx.convert(q)
991 nd = int(derivative)
993 if abs(q) > ctx.THETA_Q_LIM:
994 raise ValueError('abs(q) > THETA_Q_LIM = %f' % ctx.THETA_Q_LIM)
995 extra = 10 + ctx.prec * nd // 10
996 if z:
997 M = ctx.mag(z)
998 if M > 5 or (n != 1 and M < -5):
999 extra += 2*abs(M)
1000 cz = 0.5
1001 extra2 = 50
1002 prec0 = ctx.prec
1003 try:
1004 ctx.prec += extra
1005 if n == 1:
1006 if ctx._im(z):
1007 if abs(ctx._im(z)) < cz * abs(ctx._re(ctx.log(q))):
1008 ctx.dps += extra2
1009 res = ctx._djacobi_theta2(z - ctx.pi/2, q, nd)
1010 else:
1011 ctx.dps += 10
1012 res = ctx._djacobi_theta2a(z - ctx.pi/2, q, nd)
1013 else:
1014 res = ctx._djacobi_theta2(z - ctx.pi/2, q, nd)
1015 elif n == 2:
1016 if ctx._im(z):
1017 if abs(ctx._im(z)) < cz * abs(ctx._re(ctx.log(q))):
1018 ctx.dps += extra2
1019 res = ctx._djacobi_theta2(z, q, nd)
1020 else:
1021 ctx.dps += 10
1022 res = ctx._djacobi_theta2a(z, q, nd)
1023 else:
1024 res = ctx._djacobi_theta2(z, q, nd)
1025 elif n == 3:
1026 if ctx._im(z):
1027 if abs(ctx._im(z)) < cz * abs(ctx._re(ctx.log(q))):
1028 ctx.dps += extra2
1029 res = ctx._djacobi_theta3(z, q, nd)
1030 else:
1031 ctx.dps += 10
1032 res = ctx._djacobi_theta3a(z, q, nd)
1033 else:
1034 res = ctx._djacobi_theta3(z, q, nd)
1035 elif n == 4:
1036 if ctx._im(z):
1037 if abs(ctx._im(z)) < cz * abs(ctx._re(ctx.log(q))):
1038 ctx.dps += extra2
1039 res = ctx._djacobi_theta3(z, -q, nd)
1040 else:
1041 ctx.dps += 10
1042 res = ctx._djacobi_theta3a(z, -q, nd)
1043 else:
1044 res = ctx._djacobi_theta3(z, -q, nd)
1045 else:
1046 raise ValueError
1047 finally:
1048 ctx.prec = prec0
1049 return +res