fix(kernels): correct q=2 coefficient in PiecewisePolynomialKernel (#2756)

The q=2 branch of _get_cov used (j + 4*j + 3) instead of (j**2 + 4*j + 3),
yielding 5j+3 in place of j^2+4j+3 and producing materially wrong covariance
values for the default smoothness setting. Fix the coefficient to match
Rasmussen & Williams Eq. 4.21 (and the kernel's own docstring and test
reference). Add a q=2 numeric test with closely-spaced inputs so the
polynomial term is actually exercised (existing tests only used q=0 with
out-of-support distances, masking the bug).

Co-authored-by: Geoff Pleiss <gpleiss@gmail.com>

Author: devteamaegis

gpytorch/kernels/piecewise_polynomial_kernel.py

@@ -18,7 +18,7 @@ def _get_cov(r: Tensor, j: int, q: int) -> Tensor:
     if q == 1:
         return (j + 1) * r + 1
     if q == 2:
-        return 1 + (j + 2) * r + ((j + 4 * j + 3) / 3.0) * (r**2)
+        return 1 + (j + 2) * r + ((j**2 + 4 * j + 3) / 3.0) * (r**2)
     if q == 3:
         return (
             1

test/kernels/test_piecewise_polynomial_kernel.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import math
 import unittest
 
 import torch
@@ -64,6 +65,28 @@ def test_fmax(r, j, q):
         actual = torch.cat([actual[i].diagonal(dim1=-1, dim2=-2).unsqueeze(0) for i in range(actual.size(0))])
         self.assertLess(torch.norm(res - actual), 1e-5)
 
+    def test_computes_piecewise_polynomial_kernel_q2(self):
+        # Points must be close enough that r < 1 (i.e. inside the kernel's
+        # compact support), otherwise (1 - r)_+ zeros out the polynomial term
+        # and the q=2 coefficient is never exercised.
+        a = torch.tensor([[0.0, 0.0], [0.1, 0.2], [0.3, 0.1]], dtype=torch.float)
+        b = torch.tensor([[0.05, 0.05], [0.2, 0.1], [0.15, 0.25]], dtype=torch.float)
+        D = a.shape[-1]
+        kernel = PiecewisePolynomialKernel(q=2)
+        kernel.eval()
+
+        # Reconstruct the reference value using the kernel's own distance (so the
+        # only thing under test is the q=2 covariance polynomial coefficient).
+        r = kernel.covar_dist(a.div(kernel.lengthscale), b.div(kernel.lengthscale))
+        self.assertTrue((r < 1).any())  # ensure the polynomial term is actually exercised
+        j = math.floor(D / 2.0) + kernel.q + 1
+        fmax = torch.clamp(1 - r, min=0.0).pow(j + kernel.q)
+        # Closed form from Rasmussen & Williams Eq. 4.21
+        cov = 1 + (j + 2) * r + ((j**2 + 4 * j + 3) / 3.0) * r**2
+        actual = fmax * cov
+        res = kernel(a, b).to_dense()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
     def test_piecewise_polynomial_kernel_batch(self):
         a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
         b = torch.tensor([[0, 2, 1], [-1, 2, 0]], dtype=torch.float).view(2, 3, 1)