-
-
Notifications
You must be signed in to change notification settings - Fork 5k
/
_complexstuff.pxd
162 lines (138 loc) · 4.65 KB
/
_complexstuff.pxd
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
# -*-cython-*-
#
# Common functions required when doing complex arithmetic with Cython.
#
import cython
cimport numpy as np
from libc.math cimport (
isnan, isinf, isfinite, log1p, fabs, exp, log, sin, cos, sqrt, pow
)
cdef extern from "npy_2_complexcompat.h":
void NPY_CSETREAL(np.npy_cdouble *c, double real) nogil
void NPY_CSETIMAG(np.npy_cdouble *c, double imag) nogil
cdef extern from "_complexstuff.h":
double npy_cabs(np.npy_cdouble z) nogil
double npy_carg(np.npy_cdouble z) nogil
np.npy_cdouble npy_clog(np.npy_cdouble z) nogil
np.npy_cdouble npy_cexp(np.npy_cdouble z) nogil
np.npy_cdouble npy_csin(np.npy_cdouble z) nogil
np.npy_cdouble npy_ccos(np.npy_cdouble z) nogil
np.npy_cdouble npy_csqrt(np.npy_cdouble z) nogil
np.npy_cdouble npy_cpow(np.npy_cdouble x, np.npy_cdouble y) nogil
ctypedef double complex double_complex
ctypedef fused number_t:
double
double_complex
ctypedef union _complex_pun:
np.npy_cdouble npy
double_complex c99
cdef inline np.npy_cdouble npy_cdouble_from_double_complex(
double_complex x) noexcept nogil:
cdef _complex_pun z
z.c99 = x
return z.npy
cdef inline double_complex double_complex_from_npy_cdouble(
np.npy_cdouble x) noexcept nogil:
cdef _complex_pun z
z.npy = x
return z.c99
cdef inline bint zisnan(number_t x) noexcept nogil:
if number_t is double_complex:
return isnan(x.real) or isnan(x.imag)
else:
return isnan(x)
cdef inline bint zisfinite(number_t x) noexcept nogil:
if number_t is double_complex:
return isfinite(x.real) and isfinite(x.imag)
else:
return isfinite(x)
cdef inline bint zisinf(number_t x) noexcept nogil:
if number_t is double_complex:
return not zisnan(x) and not zisfinite(x)
else:
return isinf(x)
cdef inline double zreal(number_t x) noexcept nogil:
if number_t is double_complex:
return x.real
else:
return x
cdef inline double zabs(number_t x) noexcept nogil:
if number_t is double_complex:
return npy_cabs(npy_cdouble_from_double_complex(x))
else:
return fabs(x)
cdef inline double zarg(double complex x) noexcept nogil:
return npy_carg(npy_cdouble_from_double_complex(x))
cdef inline number_t zlog(number_t x) noexcept nogil:
cdef np.npy_cdouble r
if number_t is double_complex:
r = npy_clog(npy_cdouble_from_double_complex(x))
return double_complex_from_npy_cdouble(r)
else:
return log(x)
cdef inline number_t zexp(number_t x) noexcept nogil:
cdef np.npy_cdouble r
if number_t is double_complex:
r = npy_cexp(npy_cdouble_from_double_complex(x))
return double_complex_from_npy_cdouble(r)
else:
return exp(x)
cdef inline number_t zsin(number_t x) noexcept nogil:
cdef np.npy_cdouble r
if number_t is double_complex:
r = npy_csin(npy_cdouble_from_double_complex(x))
return double_complex_from_npy_cdouble(r)
else:
return sin(x)
cdef inline number_t zcos(number_t x) noexcept nogil:
cdef np.npy_cdouble r
if number_t is double_complex:
r = npy_ccos(npy_cdouble_from_double_complex(x))
return double_complex_from_npy_cdouble(r)
else:
return cos(x)
cdef inline number_t zsqrt(number_t x) noexcept nogil:
cdef np.npy_cdouble r
if number_t is double_complex:
r = npy_csqrt(npy_cdouble_from_double_complex(x))
return double_complex_from_npy_cdouble(r)
else:
return sqrt(x)
cdef inline number_t zpow(number_t x, double y) noexcept nogil:
cdef np.npy_cdouble r, z
# FIXME
if number_t is double_complex:
NPY_CSETREAL(&z, y)
NPY_CSETIMAG(&z, 0.0)
r = npy_cpow(npy_cdouble_from_double_complex(x), z)
return double_complex_from_npy_cdouble(r)
else:
return pow(x, y)
cdef inline double_complex zpack(double zr, double zi) noexcept nogil:
cdef np.npy_cdouble z
NPY_CSETREAL(&z, zr)
NPY_CSETIMAG(&z, zi)
return double_complex_from_npy_cdouble(z)
@cython.cdivision(True)
cdef inline double complex zlog1(double complex z) noexcept nogil:
"""
Compute log, paying special attention to accuracy around 1. We
implement this ourselves because some systems (most notably the
Travis CI machines) are weak in this regime.
"""
cdef:
int n
double complex coeff = -1
double complex res = 0
double tol = 2.220446092504131e-16
if zabs(z - 1) > 0.1:
return zlog(z)
z = z - 1
if z == 0:
return 0
for n in range(1, 17):
coeff *= -z
res += coeff/n
if zabs(res/coeff) < tol:
break
return res