Add CirculantNormal distribution.

[tillahoffmann]

This PR adds a CirculantNormal distribution which is a multivariate normal distribution where the covariance matrix has circular boundary conditions. It uses the Fourier transform introduced in #1762 to evaluate the log probability, scaling with n * log(n) instead of the usual n ** 3. The PR

• Implements the CirculantNormal distribution.
• Adds a PackRealFastFourierCoefficientsTransform transform. This is required to transform a real vector of size n to a complex vector of size n // 2 + 1 which can be passed to the rfft method. See Add complex constraint and real Fourier transform. #1762 (comment) for a brief initial discussion.
• Adds a positive_definite_circulant_vector constraint. This is really just required for the test/test_distributions.py::test_distribution_constraints tests to pass.
• Adds an example notebook, illustrating how the CirculantNormal can be used to accelerate Gaussian process inference.
• Fix the Jacobian of the Fourier transform (which turns out not to be the identity but depends on the norm parameter).

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[juanitorduz]

It would be nice to put the legends outside the plots as in https://stackoverflow.com/questions/4700614/how-to-put-the-legend-outside-the-plot :)

Reply via ReviewNB

[tillahoffmann]

I think the failing tests are due to issues with the data download and not related to the changes in this PR. There are a few further failing tests due to numerics (maybe associated with a new jax release?).

[fehiepsi]

Could you add reference for this distribution or the circulant covariance matrix?

I don't understand why we call it circulant covariance matrix. The covariance matrix in MVN is symmetric. But the circulant matrix is not symmetric.

[tillahoffmann]

I've added a bit more background and a reference. I unfortunately couldn't find a non-paywalled link.

[tillahoffmann]

There are still a few failing numerical tests. Is it possible the gof tests are a bit strict? E.g., increasing the sample size to 20k for all distributions results in the following tests failing locally.

FAILED test/test_gof.py::test_gof[GaussianRandomWalk-None-params32] - assert 0.0 > 2e-05
FAILED test/test_gof.py::test_gof[GaussianRandomWalk-None-params33] - assert 0.0 > 2e-05
FAILED test/test_gof.py::test_gof[GaussianStateSpace-None-params35] - assert 0.0 > 2e-05
FAILED test/test_gof.py::test_gof[GaussianCopulaBeta-None-params37] - assert 1.623809683759525e-05 > 2e-05
FAILED test/test_gof.py::test_gof[MultivariateStudentT-_multivariate_t_to_scipy-params97] - assert 1.0300558915599112e-05 > 2e-05
FAILED test/test_gof.py::test_gof[VonMises-<lambda>-params161] - assert 0.0 > 2e-05

[fehiepsi]

I think I understand what is a circulant matrix (from the wikipedia page) but I don't understand why this is a MVN distribution (it seems to be a transformed normal distribution based on the implementation) and why this matrix is called covariance matrix, and what "positive-definite" means here. Could you clarify?

It would be helpful to add formula in the docstrings like https://num.pyro.ai/en/stable/distributions.html#numpyro.distributions.continuous.Levy

[fehiepsi]

The implementation looks great to me! Same for the tutorial. Apart from adding more description to the definition of the distribution, could you also add a thumbnail for it https://github.com/pyro-ppl/numpyro/tree/master/docs/source/_static/img/tutorials

[fehiepsi]

Is it possible the gof tests are a bit strict?

I'm not sure. It seems not robust to multivariate distributions. You can skip it because we have tests for the bijections.

[tillahoffmann]

The implementation looks great to me! Same for the tutorial. Apart from adding more description to the definition of the distribution, could you also add a thumbnail for it https://github.com/pyro-ppl/numpyro/tree/master/docs/source/_static/img/tutorials

Ah, I forgot the thumbnail. Will add that in a future PR.