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
|
# An implementation of the digamma function for complex arguments.
#
# Author: Josh Wilson
#
# Distributed under the same license as Scipy.
#
# Sources:
# [1] "The Digital Library of Mathematical Functions", dlmf.nist.gov
#
# [2] mpmath (version 0.19), http://mpmath.org
#
import cython
from libc.math cimport ceil, fabs, M_PI
from _complexstuff cimport number_t, nan, zlog, zabs, zdiv
from _trig cimport sinpi, cospi
cimport sf_error
cdef extern from "cephes.h":
double zeta(double x, double q) nogil
double psi(double x) nogil
# Use the asymptotic series for z away from the negative real axis
# with abs(z) > smallabsz.
DEF smallabsz = 16
# Use the reflection principle for z with z.real < 0 that are within
# smallimag of the negative real axis.
DEF smallimag = 6
# Relative tolerance for series
DEF tol = 2.220446092504131e-16
# All of the following were computed with mpmath
# Location of the positive root
DEF posroot = 1.4616321449683623
# Value of the positive root
DEF posrootval = -9.2412655217294275e-17
# Location of the negative root
DEF negroot = -0.504083008264455409
# Value of the negative root
DEF negrootval = 7.2897639029768949e-17
cdef inline double digamma(double z) nogil:
"""
Wrap Cephes' psi to take advantage of the series expansions
around the two smallest zeros.
"""
if zabs(z - posroot) < 0.5:
return zeta_series(z, posroot, posrootval)
elif zabs(z - negroot) < 0.3:
return zeta_series(z, negroot, negrootval)
else:
return psi(z)
@cython.cdivision(True)
cdef inline double complex cdigamma(double complex z) nogil:
"""
Compute the digamma function for complex arguments. The strategy
is:
- Around the two zeros closest to the origin (posroot and negroot)
use a Taylor series with precomputed zero order coefficient.
- If close to the origin, use a recurrence relation to step away
from the origin.
- If close to the negative real axis, use the reflection formula
to move to the right halfplane.
- If |z| is large (> 16), use the asymptotic series.
- If |z| is small, use a recurrence relation to make |z| large
enough to use the asymptotic series.
"""
cdef:
int n
double absz = zabs(z)
double complex res = 0
double complex init
if z.real <= 0 and ceil(z.real) == z:
# Poles
sf_error.error("digamma", sf_error.SINGULAR, NULL)
return nan + 1j*nan
elif zabs(z - negroot) < 0.3:
# First negative root
return zeta_series(z, negroot, negrootval)
if z.real < 0 and fabs(z.imag) < smallabsz:
# Reflection formula for digamma. See
#
# http://dlmf.nist.gov/5.5#E4
#
res -= M_PI*cospi(z)/sinpi(z)
z = 1 - z
absz = zabs(z)
if absz < 0.5:
# Use one step of the recurrence relation to step away from
# the pole.
res -= 1/z
z += 1
absz = zabs(z)
if zabs(z - posroot) < 0.5:
res += zeta_series(z, posroot, posrootval)
elif absz > smallabsz:
res += asymptotic_series(z)
elif z.real >= 0:
n = <int>(smallabsz - absz) + 1
init = asymptotic_series(z + n)
res += backward_recurrence(z + n, init, n)
else:
# z.real < 0, absz < smallabsz, and z.imag > smallimag
n = <int>(smallabsz - absz) - 1
init = asymptotic_series(z - n)
res += forward_recurrence(z - n, init, n)
return res
@cython.cdivision(True)
cdef inline double complex forward_recurrence(double complex z,
double complex psiz,
int n) nogil:
"""
Compute digamma(z + n) using digamma(z) using the recurrence
relation
digamma(z + 1) = digamma(z) + 1/z.
See http://dlmf.nist.gov/5.5#E2
"""
cdef:
int k
double complex res = psiz
for k in range(n):
res += 1/(z + k)
return res
@cython.cdivision(True)
cdef inline double complex backward_recurrence(double complex z,
double complex psiz,
int n) nogil:
"""
Compute digamma(z - n) using digamma(z) and a recurrence
relation.
"""
cdef:
int k
double complex res = psiz
for k in range(1, n + 1):
res -= 1/(z - k)
return res
@cython.cdivision(True)
cdef inline double complex asymptotic_series(double complex z) nogil:
"""
Evaluate digamma using an asymptotic series. See
http://dlmf.nist.gov/5.11#E2
"""
cdef:
int k = 1
# The Bernoulli numbers B_2k for 1 <= k <= 16.
double *bernoulli2k = [
0.166666666666666667, -0.0333333333333333333,
0.0238095238095238095, -0.0333333333333333333,
0.0757575757575757576, -0.253113553113553114,
1.16666666666666667, -7.09215686274509804,
54.9711779448621554, -529.124242424242424,
6192.12318840579710, -86580.2531135531136,
1425517.16666666667, -27298231.0678160920,
601580873.900642368, -15116315767.0921569]
double complex rzz = zdiv(zdiv(1, z), z)
double complex zfac = 1
double complex term
double complex res
res = zlog(z) - zdiv(1.0, 2*z)
for k in range(1, 17):
zfac *= rzz
term = -bernoulli2k[k-1]*zfac/(2*k)
res += term
if zabs(term) < tol*zabs(res):
break
return res
cdef inline number_t zeta_series(number_t z, double root, double rootval) nogil:
"""
The coefficients of the Taylor series for digamma at any point can
be expressed in terms of the Hurwitz zeta function. If we
precompute the floating point number closest to a zero and the 0th
order Taylor coefficient at that point then we can compute higher
order coefficients without loss of accuracy using zeta (the zeros
are simple) and maintain high-order accuracy around the zeros.
"""
cdef:
int n
number_t res = rootval
number_t coeff = -1
number_t term
z = z - root
for n in range(1, 100):
coeff *= -z
term = coeff*zeta(n + 1, root)
res += term
if zabs(term) < tol*zabs(res):
break
return res
|
