Separate out Matern kernels (#835)

* update maths * update code * update tests * add NEWS item * precalculate fraction * fix tests * Apply suggestions from code review Co-authored-by: Sam Abbott <[email protected]> --------- Co-authored-by: Sam Abbott <[email protected]>
epiforecasts · Oct 23, 2024 · 3adc6ab · 3adc6ab
1 parent e94756d
commit 3adc6ab
Show file tree

Hide file tree

Showing 4 changed files with 105 additions and 18 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -1,5 +1,9 @@
 # EpiNow2 (development version)
 
+## Model changes
+
+- A bug in the spectral densities of Matern kernels of order other than 3/2 has been fixed and the vignette updated accordingly. By @sbfnk in #835 and reviewed by @seabbs.
+
 ## Package changes
 
 - Users are now informed that `NA` observations (if present implicitly or explicitly) will be treated as missing instead of zero when using the default `obs_opts()`. Options to treat `NA` as zeros or accumulate them are also provided in the message. By @jamesmbaazam in #774 and reviewed by @sbfnk.

diff --git a/inst/stan/functions/gaussian_process.stan b/inst/stan/functions/gaussian_process.stan
@@ -22,22 +22,53 @@ vector diagSPD_EQ(real alpha, real rho, real L, int M) {
 }
 
 /**
-  * Spectral density for Matern kernel
+  * Spectral density for 1/2 Matern (Ornstein-Uhlenbeck) kernel
   *
-  * @param nu Smoothness parameter (1/2, 3/2, or 5/2)
   * @param alpha Scaling parameter
   * @param rho Length scale parameter
   * @param L Length of the interval
   * @param M Number of basis functions
   * @return A vector of spectral densities
   */
-vector diagSPD_Matern(real nu, real alpha, real rho, real L, int M) {
+vector diagSPD_Matern12(real alpha, real rho, real L, int M) {
   vector[M] indices = linspaced_vector(M, 1, M);
-  real factor = 2 * alpha * pow(sqrt(2 * nu) / rho, nu);
-  vector[M] denom = (sqrt(2 * nu) / rho)^2 + pow((pi() / 2 / L) * indices, nu + 0.5);
+  real factor = 2;
+  vector[M] denom = rho * ((1 / rho)^2 + (pi() / 2 / L) * indices);
+  return alpha * sqrt(factor * inv(denom));
+}
+
+/**
+  * Spectral density for 3/2 Matern kernel
+  *
+  * @param alpha Scaling parameter
+  * @param rho Length scale parameter
+  * @param L Length of the interval
+  * @param M Number of basis functions
+  * @return A vector of spectral densities
+  */
+vector diagSPD_Matern32(real alpha, real rho, real L, int M) {
+  vector[M] indices = linspaced_vector(M, 1, M);
+  real factor = 2 * alpha * pow(sqrt(3) / rho, 1.5);
+  vector[M] denom = (sqrt(3) / rho)^2 + pow((pi() / 2 / L) * indices, 2);
   return factor * inv(denom);
 }
 
+/**
+  * Spectral density for 5/2 Matern kernel
+  *
+  * @param alpha Scaling parameter
+  * @param rho Length scale parameter
+  * @param L Length of the interval
+  * @param M Number of basis functions
+  * @return A vector of spectral densities
+  */
+vector diagSPD_Matern52(real alpha, real rho, real L, int M) {
+  vector[M] indices = linspaced_vector(M, 1, M);
+  real factor = 3 * pow(sqrt(5) / rho, 5);
+  vector[M] denom = 2 * (sqrt(5) / rho)^2 + pow((pi() / 2 / L) * indices, 3);
+  return alpha * sqrt(factor * inv(denom));
+}
+
 /**
   * Spectral density for periodic kernel
   *
@@ -143,7 +174,15 @@ vector update_gp(matrix PHI, int M, real L, real alpha,
   } else if (type == 1) {
     diagSPD = diagSPD_Periodic(alpha, rho[1], M);
   } else if (type == 2) {
-    diagSPD = diagSPD_Matern(nu, alpha, rho[1], L, M);
+    if (nu == 0.5) {
+      diagSPD = diagSPD_Matern12(alpha, rho[1], L, M);
+    } else if (nu == 1.5) {
+      diagSPD = diagSPD_Matern32(alpha, rho[1], L, M);
+    } else if (nu == 2.5) {
+      diagSPD = diagSPD_Matern52(alpha, rho[1], L, M);
+    } else {
+      reject("nu must be one of 1/2, 3/2 or 5/2; found nu=", nu);
+    }
   }
   return PHI * (diagSPD .* eta);
 }

diff --git a/tests/testthat/test-stan-guassian-process.R b/tests/testthat/test-stan-guassian-process.R
@@ -24,21 +24,43 @@ test_that("diagSPD_EQ returns correct dimensions and values", {
   expect_equal(result, expected_result, tolerance = 1e-8)
 })
 
-test_that("diagSPD_Matern returns correct dimensions and values", {
-  nu <- 1.5
+test_that("diagSPD_Matern functions return correct dimensions and values", {
   alpha <- 1.0
   rho <- 2.0
   L <- 1.0
   M <- 5
-  result <- diagSPD_Matern(nu, alpha, rho, L, M)
-  expect_equal(length(result), M)
-  expect_true(all(result > 0))  # Expect spectral density to be positive
+  result12 <- diagSPD_Matern12(alpha, rho, L, M)
+  expect_equal(length(result12), M)
+  expect_true(all(result12 > 0))  # Expect spectral density to be positive
+
   # Check specific values for known inputs
   indices <- linspaced_vector(M, 1, M)
-  factor <- 2 * alpha * (sqrt(2 * nu) / rho)^nu
-  denom <- (sqrt(2 * nu) / rho)^2 + (pi / (2 * L) * indices)^(nu + 0.5)
-  expected_result <- factor / denom
-  expect_equal(result, expected_result, tolerance = 1e-8)
+  factor12 <- 2
+  denom12 <- rho * ((1 / rho)^2 + (pi / 2 / L) * indices)
+  expected_result12 <- alpha * sqrt(factor12 / denom12)
+  expect_equal(result12, expected_result12, tolerance = 1e-8)
+
+  result32 <- diagSPD_Matern32(alpha, rho, L, M)
+  expect_equal(length(result32), M)
+  expect_true(all(result32 > 0))  # Expect spectral density to be positive
+
+  # Check specific values for known inputs
+  indices <- linspaced_vector(M, 1, M)
+  factor32 <- 2 * alpha * (sqrt(3) / rho)^(1.5)
+  denom32 <- (sqrt(3) / rho)^2 + (pi / 2 / L * indices)^2
+  expected_result32 <- factor32 / denom32
+  expect_equal(result32, expected_result32, tolerance = 1e-8)
+
+  result52 <- diagSPD_Matern52(alpha, rho, L, M)
+  expect_equal(length(result52), M)
+  expect_true(all(result52 > 0))  # Expect spectral density to be positive
+
+  # Check specific values for known inputs
+  indices <- linspaced_vector(M, 1, M)
+  factor52 <- 3 * (sqrt(5) / rho)^5
+  denom52 <- 2 * (sqrt(5) / rho)^2 + ((pi / 2 / L) * indices)^3
+  expected_result52 <- alpha * sqrt(factor52 / denom52)
+  expect_equal(result52, expected_result52, tolerance = 1e-8)
 })
 
 test_that("diagSPD_Periodic returns correct dimensions and values", {

diff --git a/vignettes/gaussian_process_implementation_details.Rmd b/vignettes/gaussian_process_implementation_details.Rmd
@@ -93,13 +93,35 @@ where $L$ is a positive number termed boundary condition, and $\beta_{j}$ are re
 \beta_j \sim \mathcal{Normal}(0, 1)
 \end{equation}
 
-The function $S_k(x)$ is the spectral density relating to a particular covariance function $k$. In the case of the Matérn 3/2 kernel (the default in `EpiNow2`) this is given by
+The function $S_k(x)$ is the spectral density relating to a particular covariance function $k$.
+In the case of the Matérn kernel of order $\nu$ this is given by
 
 \begin{equation}
-S_k(x) = 4 \alpha^2 \left( \frac{\sqrt{3}}{\rho}\right)^3 \left(\left( \frac{\sqrt{3}}{\rho} \right)^2 + w^2 \right)^{-2}
+S_k(x) = \alpha^2 \frac{2 \sqrt{\pi} \Gamma(\nu + 1/2) (2\nu)^\nu}{\Gamma(\nu) \rho^{2 \nu}} \left( \frac{2 \nu}{\rho^2} + \omega^2 \right)^{-\left( \nu + \frac{1}{2}\right )}
 \end{equation}
 
-and in the case of a squared exponential kernel by
+For $\nu = 3 / 2$ (the default in `EpiNow2`) this simplifies to
+
+\begin{equation}
+    S_k(x) =
+    \alpha^2 \frac{4 \left(\sqrt{3} / \rho \right)^3}{\left(\left(\sqrt{3} / \rho\right)^2 + \omega^2\right)^2} = 
+    \left(\frac{2 \alpha \left(\sqrt{3} / \rho \right)^{\frac{3}{2}}}{\left(\sqrt{3} / \rho\right)^2 + \omega^2} \right)^2
+\end{equation}
+
+For $\nu = 1 / 2$ it is
+
+\begin{equation}
+S_k(x) = \alpha^2 \frac{2}{\rho \left(1 / \rho^2 + \omega^2\right)}
+\end{equation}
+
+and for $\nu = 5 / 2$ it is
+
+\begin{equation}
+S_k(x) =
+    \alpha^2 \frac{3 \left(\sqrt{5} / \rho \right)^5}{2 \left(\left(\sqrt{5} / \rho\right)^2 + \omega^2\right)^3}
+\end{equation}
+
+In the case of a squared exponential the spectral kernel density is given by
 
 \begin{equation}
 S_k(x) = \alpha^2 \sqrt{2\pi} \rho \exp \left( -\frac{1}{2} \rho^2 w^2 \right)