Add power_spectral_density to RatQuad covariance kernel for HSGP support

[fonnesbeck]

Description

PSD method and test added. A simple HSGP fits successfully with the resulting RatQuad kernel.

Related Issue

• Closes #
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.49%. Comparing base (f69e159) to head (46343e2).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #7973   +/-   ##
=======================================
  Coverage   91.49%   91.49%           
=======================================
  Files         116      116           
  Lines       18962    18971    +9     
=======================================
+ Hits        17349    17358    +9     
  Misses       1613     1613           

Files with missing lines	Coverage Δ
pymc/gp/cov.py	98.68% <100.00%> (+0.02%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[fonnesbeck]

BTW, one limitation here is that pytensor.kv does not appear to support gradient calculation so NUTS can't be fully supported. There is an alternative parameterization that may not need Bessel functions, unless said support is coming.

[jessegrabowski]

BTW, one limitation here is that pytensor.kv does not appear to support gradient calculation so NUTS can't be fully supported. There is an alternative parameterization that may not need Bessel functions, unless said support is coming.

Does jax or tfp have the gradients we could use as the basis for an implementation?

[fonnesbeck]

Apparently not.

Looks like a numerical approximation (via finite differences) is as close as you can get.

[jessegrabowski]

They have an expression for the 2nd input: https://github.com/tensorflow/probability/blob/732e7d16235a38b639a3b2cd00e76a4add45118c/tensorflow_probability/python/math/bessel.py#L1057

[ricardoV94]

In the past I've implemented a couple gradients using stan math / wolfram as reference. Never had a case where tfp was ahead of the curve for these special functions

[ricardoV94]

Gradient wrt second input is the trivial one in these functions, we also have it

[fonnesbeck]

There has been some useful discussion on Jax: jax-ml/jax#12402

[fonnesbeck]

FWIW, it does work without a fully differentiable function -- it just uses a slice sampler for alpha.

ratquad_hsgp.html

[bwengals]

Will this work for a 2 or higher dimensional X?

[fonnesbeck]

It does