Repository: pymc-devs/pytensor

Issue:
Add helper for `special.kn` that just calls `special.kv` under the hood
### After

I wish for 
```python
pytensor.tensor.math.kn(0, 1.0)
```
to compute the modified Bessel function of the second kind, of order zero, at a generic point on its domain (e.g. `1.0` here).

### Context for the issue:

I am unable to use Scipy's implementation `scipy.special.kn` ([link](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.kn.html)) of the modified Bessel function of the second kind of integer order `n` (`n=0` is what I need) in PyMC since it does not handle autodiff.

I see there are many Bessel functions implemented already in pytensor, but the modified Bessel function of the second kind, of order zero, is not. I am kindly requesting an implementation of this for my project modeling the spread of chemical plumes ([link](https://simonensemble.github.io/posts/2024-05-28-chemical-plume/)) to search for the source of a chemical plume using Bayesian inference in PyMC.

Discussion:
@SimonEnsemble: the `pytensor.tensor.math` docs [here](https://www.pymc.io/projects/docs/en/stable/_modules/pytensor/tensor/math.html) were missing the integer order one. but I see the real order one work with zero, so I'm in business, I think.

```
# Define the expression
expr = pytensor.tensor.kv(0, 0.1)

# Compile the expression into a function
f = pytensor.function([], expr)

# Call the function to get the numerical value
result = f()
result
```

matches
```
scipy.special.kn(0, 0.1)
```

@ricardoV94: @SimonEnsemble we should add a helper for `kn` that just calls `kv` so future users don't have to do the same work you had to do here

Solution patch:
diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py
index 2ba2ab158..714f597b3 100644
--- a/pytensor/tensor/math.py
+++ b/pytensor/tensor/math.py
@@ -2453,6 +2453,11 @@ def kv(v, x):
     return kve(v, x) * exp(-x)
 
 
+def kn(n, x):
+    """Modified Bessel function of the second kind of integer order v."""
+    return kv(n, x)
+
+
 @scalar_elemwise
 def sigmoid(x):
     """Logistic sigmoid function (1 / (1 + exp(-x)), also known as expit or inverse logit"""
@@ -4337,6 +4342,7 @@ __all__ = [
     "i1",
     "iv",
     "ive",
+    "kn",
     "kv",
     "kve",
     "sigmoid",


Test patch:
diff --git a/tests/tensor/test_math_scipy.py b/tests/tensor/test_math_scipy.py
index 8f7095020..e15293e97 100644
--- a/tests/tensor/test_math_scipy.py
+++ b/tests/tensor/test_math_scipy.py
@@ -18,7 +18,7 @@ from pytensor import function, grad
 from pytensor import tensor as pt
 from pytensor.compile.mode import get_default_mode
 from pytensor.configdefaults import config
-from pytensor.tensor import gammaincc, inplace, kv, kve, vector
+from pytensor.tensor import gammaincc, inplace, kn, kv, kve, vector
 from tests import unittest_tools as utt
 from tests.tensor.utils import (
     _good_broadcast_unary_chi2sf,
@@ -1220,3 +1220,17 @@ def test_kv():
         out.eval({v: test_v, x: test_x}),
         scipy.special.kv(test_v[:, None], test_x[None, :]),
     )
+
+
+def test_kn():
+    n = vector("n")
+    x = vector("x")
+
+    out = kn(n[:, None], x[None, :])
+    test_n = np.array([-3, 4, 0, 5], dtype=n.type.dtype)
+    test_x = np.linspace(0, 512, 10, dtype=x.type.dtype)
+
+    np.testing.assert_allclose(
+        out.eval({n: test_n, x: test_x}),
+        scipy.special.kn(test_n[:, None], test_x[None, :]),
+    )