pyerrors/pyerrors/special.py
Fabian Joswig 5122e260ea
Special function module (#221)
* [fix] Fixed intendation for flake8

* [feat] special module with Bessel functions added.

* [feat] init adapted.

* [tests] First test for special module added.

* [tests] Check kn special function derivative explicitly vs numerical
derivative of scipy function.

* [feat] Added remaining autograd scipy special functions to the special
module.

* [feat] Imports corrected, docstring added.
2024-01-07 17:22:05 +01:00

23 lines
1 KiB
Python

import scipy
import numpy as np
from autograd.extend import primitive, defvjp
from autograd.scipy.special import j0, y0, j1, y1, jn, yn, i0, i1, iv, ive, beta, betainc, betaln
from autograd.scipy.special import polygamma, psi, digamma, gamma, gammaln, gammainc, gammaincc, gammasgn, rgamma, multigammaln
from autograd.scipy.special import erf, erfc, erfinv, erfcinv, logit, expit, logsumexp
__all__ = ["beta", "betainc", "betaln",
"polygamma", "psi", "digamma", "gamma", "gammaln", "gammainc", "gammaincc", "gammasgn", "rgamma", "multigammaln",
"kn", "j0", "y0", "j1", "y1", "jn", "yn", "i0", "i1", "iv", "ive",
"erf", "erfc", "erfinv", "erfcinv", "logit", "expit", "logsumexp"]
@primitive
def kn(n, x):
"""Modified Bessel function of the second kind of integer order n"""
if int(n) != n:
raise TypeError("The order 'n' needs to be an integer.")
return scipy.special.kn(n, x)
defvjp(kn, None, lambda ans, n, x: lambda g: - g * 0.5 * (kn(np.abs(n - 1), x) + kn(n + 1, x)))