Implementation of the power spectrum inference with GPs

[mlefkir]

This is the Python-JAX implementation of a method to infer the power spectral density of irregular time series using Gaussian process regression. The method is described in a forthcoming paper and in the Julia package Pioran.jl, it relies on approximating a bending power-law model in a sum of scalable kernels implemented in tinygp.

Relevant Issue(s)/PR(s)

Provide an overview of the implemented solution or the fix and elaborate on the modifications.

Is there a new dependency introduced by your contribution? If so, please specify.

Any other comments?

[pep8speaks]

Hello @mlefkir! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

• In the file stingray/modeling/gpmodeling.py:

Line 98:80: E501 line too long (85 > 79 characters)
Line 157:80: E501 line too long (95 > 79 characters)
Line 185:80: E501 line too long (92 > 79 characters)
Line 186:80: E501 line too long (93 > 79 characters)
Line 187:80: E501 line too long (83 > 79 characters)
Line 205:80: E501 line too long (86 > 79 characters)
Line 216:80: E501 line too long (97 > 79 characters)
Line 263:80: E501 line too long (83 > 79 characters)
Line 266:80: E501 line too long (92 > 79 characters)
Line 267:80: E501 line too long (99 > 79 characters)
Line 270:80: E501 line too long (89 > 79 characters)
Line 325:80: E501 line too long (81 > 79 characters)
Line 328:80: E501 line too long (86 > 79 characters)
Line 352:80: E501 line too long (94 > 79 characters)
Line 390:80: E501 line too long (91 > 79 characters)
Line 391:80: E501 line too long (88 > 79 characters)
Line 394:80: E501 line too long (94 > 79 characters)
Line 396:80: E501 line too long (87 > 79 characters)
Line 397:80: E501 line too long (86 > 79 characters)
Line 412:80: E501 line too long (87 > 79 characters)
Line 413:80: E501 line too long (86 > 79 characters)
Line 426:80: E501 line too long (94 > 79 characters)
Line 454:80: E501 line too long (92 > 79 characters)
Line 484:80: E501 line too long (81 > 79 characters)
Line 487:80: E501 line too long (86 > 79 characters)
Line 498:80: E501 line too long (95 > 79 characters)
Line 539:80: E501 line too long (92 > 79 characters)
Line 542:80: E501 line too long (91 > 79 characters)
Line 590:80: E501 line too long (86 > 79 characters)
Line 594:80: E501 line too long (80 > 79 characters)
Line 606:80: E501 line too long (94 > 79 characters)
Line 609:80: E501 line too long (96 > 79 characters)
Line 614:80: E501 line too long (89 > 79 characters)
Line 702:80: E501 line too long (89 > 79 characters)
Line 995:80: E501 line too long (92 > 79 characters)
Line 1026:80: E501 line too long (83 > 79 characters)
Line 1044:80: E501 line too long (81 > 79 characters)
Line 1045:80: E501 line too long (89 > 79 characters)
Line 1046:80: E501 line too long (87 > 79 characters)
Line 1194:80: E501 line too long (81 > 79 characters)
Line 1195:80: E501 line too long (89 > 79 characters)
Line 1247:80: E501 line too long (94 > 79 characters)
Line 1317:80: E501 line too long (91 > 79 characters)
Line 1321:80: E501 line too long (92 > 79 characters)

• In the file stingray/modeling/tests/test_gpmodeling.py:

Line 118:80: E501 line too long (81 > 79 characters)
Line 145:80: E501 line too long (90 > 79 characters)
Line 336:80: E501 line too long (94 > 79 characters)
Line 342:80: E501 line too long (80 > 79 characters)
Line 347:80: E501 line too long (83 > 79 characters)
Line 350:80: E501 line too long (91 > 79 characters)
Line 362:80: E501 line too long (98 > 79 characters)
Line 369:80: E501 line too long (91 > 79 characters)
Line 384:80: E501 line too long (91 > 79 characters)
Line 392:80: E501 line too long (90 > 79 characters)
Line 400:80: E501 line too long (91 > 79 characters)
Line 404:80: E501 line too long (83 > 79 characters)
Line 411:80: E501 line too long (89 > 79 characters)
Line 412:80: E501 line too long (96 > 79 characters)

Comment last updated at 2024-10-03 11:36:06 UTC

[matteobachetti]

@mlefkir thanks for your contribution to Stingray, and sorry for my late reply!
I'm not an expert in this kind of methods (@Gaurav17Joshi implemented the current GP model and might have some comments). Would you mind showing the usage of this method through an example?

From the point of view of the requirements for PRs to Stingray:

1. most formatting issues can be solved by running black on the new files. See https://docs.stingray.science/en/stable/contributing.html#coding-style-and-conventions. Don't worry about PEP8 still complaining after black has run ;)
2. A changelog entry is highly appreciated, it helps keeping track of the evolution of Stingray. https://docs.stingray.science/en/stable/contributing.html#updating-and-maintaining-the-changelog

[Gaurav17Joshi]

Yes, @mlefkir, do you have any use case example for this method (Something like a .ipynb notebook). Also @dhuppenkothen, can you also have a look into the usefulness of this method for Stingray, and whether it should go into the same file as the gpmodeling part?

[mlefkir]

@Gaurav17Joshi @matteobachetti I made two examples available here Examples, one uses nested sampling with the jaxns sampler already called in Stingray and the other one uses NUTS with NumPyro.
I will update the changelog at some point.

[matteobachetti]

@mlefkir I'm playing with your PR. Really sorry for the slow progress, but my knowledge of these methods is pretty poor and it takes me a lot of time to just understand how it works, and... my agenda is pretty full 😅 .
Just to guide me in the review: does your model also work with the spectral shapes that @Gaurav17Joshi set up? Is it just working with unevenly sampled data with error bars? What is the difference between your implementation and the existing one?
Also, feel free to advertise your paper when it is ready!

[mlefkir]

@matteobachetti This model is an improvement on available models. Currently, available models have a fixed low and high-frequency slope. For instance, the DRW/Exponential kernel has a Lorentzian power spectrum with a low-frequency slope 0 and a high-frequency slope of -2. This method allows modelling spectral shapes with flexible bends frequencies and slopes which can be between 0 and 4 using a sum of basis functions as shown in the figure below:

This method is designed for Gaussian process regression so it can be used for any Gaussian time series with (or without) error bars. While the algorithm in tinygp is fast there are limitations on the number of points in terms of computational cost so I would use the method only for irregularly sampled data or data with gaps with less than 10,000 points.
This is a simple implementation of the method which should be easier to maintain than the Python library I wrote a few months ago, I intend to archive it at some point.