Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add helper functions to calculate required M for hsgp #7058

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
71 changes: 71 additions & 0 deletions pymc/gp/cov.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

from collections import Counter
from functools import reduce
from math import ceil
from operator import add, mul
from typing import Any, Callable, List, Optional, Sequence, Union

Expand Down Expand Up @@ -564,6 +565,10 @@
def power_spectral_density(self, omega: TensorLike) -> TensorVariable:
raise NotImplementedError

@staticmethod
def required_num_eigenvectors(L: float, ls: float) -> int:
raise NotImplementedError

Check warning on line 570 in pymc/gp/cov.py

View check run for this annotation

Codecov / codecov/patch

pymc/gp/cov.py#L570

Added line #L570 was not covered by tests


class ExpQuad(Stationary):
r"""
Expand Down Expand Up @@ -595,6 +600,28 @@
exp = pt.exp(-0.5 * pt.dot(pt.square(omega), pt.square(ls)))
return c * pt.prod(ls) * exp

@staticmethod
def required_num_eigenvectors(L: float, ls: float) -> int:
r"""Number of eigenvectors in Hilbert space that approximate the GP well.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we add the paper (and maybe an equation reference) to understand where to learn more about the factors (for us, the newbies?)


.. math::

m \ge 1.75 \frac{L}{ls}

Parameters
----------
L : float
Approximation bound
ls : float
lengthscale

Returns
-------
int
Number of eigenvectors
"""
return ceil(1.75 * L / ls)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you link the source of these numbers in the docstring?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, you can rebase on main to fix these ❌ tests. I merged Ricardos #7057



class RatQuad(Stationary):
r"""
Expand Down Expand Up @@ -663,6 +690,28 @@
pow = pt.power(5.0 + pt.dot(pt.square(omega), pt.square(ls)), -1 * D52)
return (num / den) * pt.prod(ls) * pow

@staticmethod
def required_num_eigenvectors(L: float, ls: float) -> int:
r"""Number of eigenvectors in Hilbert space that approximate the GP well.
Copy link
Contributor

@bwengals bwengals Dec 10, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably should add some more detail in the docstrings here? Would be kind of mysterious to see this if you're not using HSGPs, or know how to choose L.

Another idea (take it or leave it) would be to have this put into a single function in hsgp_approx.py that does this all in one go, something like

def required_num_eigenvectors(cov_func, L, ls):
    if cov_func.__class__.__name__ == "ExpQuad":
        return ceil(1.75 * L / ls)
    elif ...:
        ...

which maybe is better because the docstring can be in one place and the logic and purpose is isolated in one spot instead of spread out. I wonder if the recommendations may evolve too, since in the paper they didn't try for higher than 1D inputs.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Even more convenient would be to have this function called automatically if m is not passed to HSGP (along with a healthy warning message informing users what is going on).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

and c too (but maybe not in this PR)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I'll do this


.. math::

m \ge 2.65 \frac{L}{ls}

Parameters
----------
L : float
Approximation bound
ls : float
lengthscale

Returns
-------
int
Number of eigenvectors
"""
return ceil(2.65 * L / ls)


class Matern32(Stationary):
r"""
Expand Down Expand Up @@ -702,6 +751,28 @@
pow = pt.power(3.0 + pt.dot(pt.square(omega), pt.square(ls)), -1 * D32)
return (num / den) * pt.prod(ls) * pow

@staticmethod
def required_num_eigenvectors(L: float, ls: float) -> int:
r"""Number of eigenvectors in Hilbert space that approximate the GP well.

.. math::

m \ge 3.42 \frac{L}{ls}

Parameters
----------
L : float
Approximation bound
ls : float
lengthscale

Returns
-------
int
Number of eigenvectors
"""
return ceil(3.42 * L / ls)


class Matern12(Stationary):
r"""
Expand Down
21 changes: 21 additions & 0 deletions tests/gp/test_cov.py
Original file line number Diff line number Diff line change
Expand Up @@ -484,6 +484,13 @@ def test_euclidean_dist(self):
print(result, expected)
npt.assert_allclose(result, expected, atol=1e-5)

def test_required_num_eigenvectors(self):
L = 4
l = 1
m = 7
m_est = pm.gp.cov.ExpQuad.required_num_eigenvectors(L, l)
assert m_est == m


class TestWhiteNoise:
def test_1d(self):
Expand Down Expand Up @@ -572,6 +579,13 @@ def test_psd(self):
)
npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5)

def test_required_num_eigenvectors(self):
L = 4
l = 1
m = 11
m_est = pm.gp.cov.Matern52.required_num_eigenvectors(L, l)
assert m_est == m


class TestMatern32:
def test_1d(self):
Expand All @@ -598,6 +612,13 @@ def test_psd(self):
)
npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5)

def test_required_num_eigenvectors(self):
L = 4
l = 1
m = 14
m_est = pm.gp.cov.Matern32.required_num_eigenvectors(L, l)
assert m_est == m


class TestMatern12:
def test_1d(self):
Expand Down