From 9f18c439b1029f9680889a1815f90aff8d80849d Mon Sep 17 00:00:00 2001
From: Max Kochurov <maxim.v.kochurov@gmail.com>
Date: Sat, 9 Dec 2023 16:52:32 +0000
Subject: [PATCH] add helper functions to calculate required M for hsgp

---
 pymc/gp/cov.py       | 71 ++++++++++++++++++++++++++++++++++++++++++++
 tests/gp/test_cov.py | 21 +++++++++++++
 2 files changed, 92 insertions(+)

diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py
index 268b678505..9b7e282ace 100644
--- a/pymc/gp/cov.py
+++ b/pymc/gp/cov.py
@@ -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
 
@@ -564,6 +565,10 @@ def full_from_distance(self, dist: TensorLike, squared: bool = False) -> TensorV
     def power_spectral_density(self, omega: TensorLike) -> TensorVariable:
         raise NotImplementedError
 
+    @staticmethod
+    def required_num_eigenvectors(L: float, ls: float) -> int:
+        raise NotImplementedError
+
 
 class ExpQuad(Stationary):
     r"""
@@ -595,6 +600,28 @@ def power_spectral_density(self, omega: TensorLike) -> TensorVariable:
         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.
+
+        .. 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)
+
 
 class RatQuad(Stationary):
     r"""
@@ -663,6 +690,28 @@ def power_spectral_density(self, omega: TensorLike) -> TensorVariable:
         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.
+
+        .. 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"""
@@ -702,6 +751,28 @@ def power_spectral_density(self, omega: TensorLike) -> TensorVariable:
         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"""
diff --git a/tests/gp/test_cov.py b/tests/gp/test_cov.py
index e9d6744624..df4d9acefe 100644
--- a/tests/gp/test_cov.py
+++ b/tests/gp/test_cov.py
@@ -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):
@@ -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):
@@ -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):