Skip to content

Commit 3b54a5b

Browse files
committed
Added eigenvalue-based tests for PSD in stationary kernels
1 parent ebd836e commit 3b54a5b

File tree

1 file changed

+81
-32
lines changed

1 file changed

+81
-32
lines changed

tests/gp/test_cov.py

Lines changed: 81 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -459,16 +459,33 @@ def test_inv_lengthscale(self):
459459
Kd = cov(X, diag=True).eval()
460460
npt.assert_allclose(np.diag(K), Kd, atol=1e-5)
461461

462-
def test_psd(self):
463-
# compare to simple 1d formula
464-
X = np.linspace(0, 1, 10)[:, None]
465-
omega = np.linspace(0, 2, 50)
466-
ell = 2.0
467-
true_1d_psd = np.sqrt(2 * np.pi * np.square(ell)) * np.exp(-0.5 * np.square(ell * omega))
468-
test_1d_psd = (
469-
pm.gp.cov.ExpQuad(1, ls=ell).power_spectral_density(omega[:, None]).flatten().eval()
470-
)
471-
npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5)
462+
def test_psd_eigenvalues(self):
463+
"""Test PSD implementation using Rayleigh quotients."""
464+
ls = 0.1
465+
N = 1000
466+
L = 10.0
467+
dx = L / N
468+
X = np.linspace(0, L, N)[:, None]
469+
470+
with pm.Model():
471+
cov = pm.gp.cov.ExpQuad(1, ls=ls)
472+
473+
K = cov(X).eval()
474+
475+
freqs = np.fft.fftfreq(N, d=dx)
476+
omegas = 2 * np.pi * freqs
477+
478+
j = np.arange(N)
479+
modes = np.exp(2j * np.pi * np.outer(np.arange(N), j) / N)
480+
numerator = np.diag(modes @ K @ modes.conj().T).real
481+
rayleigh_quotient = numerator / N
482+
483+
psd = cov.power_spectral_density(omegas[:, None]).eval()
484+
psd_scaled = psd.flatten() / dx
485+
486+
# Trim boundaries where numerical error concentrates
487+
trim = N // 10
488+
npt.assert_allclose(psd_scaled[trim:-trim], rayleigh_quotient[trim:-trim], atol=1e-2)
472489

473490
def test_euclidean_dist(self):
474491
X = np.arange(0, 3)[:, None]
@@ -590,17 +607,33 @@ def test_1d(self):
590607
Kd = cov(X, diag=True).eval()
591608
npt.assert_allclose(np.diag(K), Kd, atol=1e-5)
592609

593-
def test_psd(self):
594-
# compare to simple 1d formula
595-
X = np.linspace(0, 1, 10)[:, None]
596-
omega = np.linspace(0, 2, 50)
597-
ell = 2.0
598-
lamda = np.sqrt(5) / ell
599-
true_1d_psd = (16.0 / 3.0) * np.power(lamda, 5) * np.power(lamda**2 + omega**2, -3)
600-
test_1d_psd = (
601-
pm.gp.cov.Matern52(1, ls=ell).power_spectral_density(omega[:, None]).flatten().eval()
602-
)
603-
npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5)
610+
def test_psd_eigenvalues(self):
611+
"""Test PSD implementation using Rayleigh quotients."""
612+
ls = 0.1
613+
N = 1000
614+
L = 10.0
615+
dx = L / N
616+
X = np.linspace(0, L, N)[:, None]
617+
618+
with pm.Model():
619+
cov = pm.gp.cov.Matern52(1, ls=ls)
620+
621+
K = cov(X).eval()
622+
623+
freqs = np.fft.fftfreq(N, d=dx)
624+
omegas = 2 * np.pi * freqs
625+
626+
j = np.arange(N)
627+
modes = np.exp(2j * np.pi * np.outer(np.arange(N), j) / N)
628+
numerator = np.diag(modes @ K @ modes.conj().T).real
629+
rayleigh_quotient = numerator / N
630+
631+
psd = cov.power_spectral_density(omegas[:, None]).eval()
632+
psd_scaled = psd.flatten() / dx
633+
634+
# Trim boundaries where numerical error concentrates
635+
trim = N // 10
636+
npt.assert_allclose(psd_scaled[trim:-trim], rayleigh_quotient[trim:-trim], atol=1e-2)
604637

605638

606639
class TestMatern32:
@@ -616,17 +649,33 @@ def test_1d(self):
616649
Kd = cov(X, diag=True).eval()
617650
npt.assert_allclose(np.diag(K), Kd, atol=1e-5)
618651

619-
def test_psd(self):
620-
# compare to simple 1d formula
621-
X = np.linspace(0, 1, 10)[:, None]
622-
omega = np.linspace(0, 2, 50)
623-
ell = 2.0
624-
lamda = np.sqrt(3) / ell
625-
true_1d_psd = 4 * np.power(lamda, 3) * np.power(lamda**2 + omega**2, -2)
626-
test_1d_psd = (
627-
pm.gp.cov.Matern32(1, ls=ell).power_spectral_density(omega[:, None]).flatten().eval()
628-
)
629-
npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5)
652+
def test_psd_eigenvalues(self):
653+
"""Test PSD implementation using Rayleigh quotients."""
654+
ls = 0.1
655+
N = 1000
656+
L = 10.0
657+
dx = L / N
658+
X = np.linspace(0, L, N)[:, None]
659+
660+
with pm.Model():
661+
cov = pm.gp.cov.Matern32(1, ls=ls)
662+
663+
K = cov(X).eval()
664+
665+
freqs = np.fft.fftfreq(N, d=dx)
666+
omegas = 2 * np.pi * freqs
667+
668+
j = np.arange(N)
669+
modes = np.exp(2j * np.pi * np.outer(np.arange(N), j) / N)
670+
numerator = np.diag(modes @ K @ modes.conj().T).real
671+
rayleigh_quotient = numerator / N
672+
673+
psd = cov.power_spectral_density(omegas[:, None]).eval()
674+
psd_scaled = psd.flatten() / dx
675+
676+
# Trim boundaries where numerical error concentrates
677+
trim = N // 10
678+
npt.assert_allclose(psd_scaled[trim:-trim], rayleigh_quotient[trim:-trim], atol=1e-2)
630679

631680

632681
class TestMatern12:

0 commit comments

Comments
 (0)