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.
This commit is contained in:
Fabian Joswig 2024-01-07 17:22:05 +01:00 committed by GitHub
parent 1360942a7b
commit 5122e260ea
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 37 additions and 1 deletions

View file

@ -487,5 +487,6 @@ from . import linalg
from . import mpm
from . import roots
from . import integrate
from . import special
from .version import __version__

23
pyerrors/special.py Normal file
View file

@ -0,0 +1,23 @@
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)))

12
tests/special_test.py Normal file
View file

@ -0,0 +1,12 @@
import numpy as np
import scipy
import pyerrors as pe
import pytest
from autograd import jacobian
from numdifftools import Jacobian as num_jacobian
def test_kn():
for n in np.arange(0, 10):
for val in np.linspace(0.1, 7.3, 10):
assert np.isclose(num_jacobian(lambda x: scipy.special.kn(n, x))(val), jacobian(lambda x: pe.special.kn(n, x))(val), rtol=1e-10, atol=1e-10)