From bf7ba08d3da0cbc8070acc53b9344ceaf735fa7c Mon Sep 17 00:00:00 2001
From: Ben Mares <services-git-throwaway1@tensorial.com>
Date: Wed, 21 Aug 2024 18:26:38 +0200
Subject: [PATCH] Minor refactoring to make periodic HSGP look similar

---
 pymc/gp/hsgp_approx.py | 41 +++++++++++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 12 deletions(-)

diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py
index 760262276a..38726e438b 100644
--- a/pymc/gp/hsgp_approx.py
+++ b/pymc/gp/hsgp_approx.py
@@ -443,17 +443,22 @@ def prior(
         """
         phi, sqrt_psd = self.prior_linearized(X)
 
+        size = self._m_star - int(self._drop_first)
         if self._parametrization == "noncentered":
             self._beta = pm.Normal(
                 f"{name}_hsgp_coeffs_",
-                size=self._m_star - int(self._drop_first),
+                size=size,
                 dims=hsgp_coeffs_dims,
             )
             self._sqrt_psd = sqrt_psd
             f = self.mean_func(X) + phi @ (self._beta * self._sqrt_psd)
 
         elif self._parametrization == "centered":
-            self._beta = pm.Normal(f"{name}_hsgp_coeffs_", sigma=sqrt_psd, dims=hsgp_coeffs_dims)
+            self._beta = pm.Normal(
+                f"{name}_hsgp_coeffs_",
+                sigma=sqrt_psd,
+                dims=hsgp_coeffs_dims,
+            )
             f = self.mean_func(X) + phi @ self._beta
 
         self.f = pm.Deterministic(name, f, dims=gp_dims)
@@ -678,7 +683,12 @@ def prior_linearized(self, X: TensorLike):
         psd = self.scale * self.cov_func.power_spectral_density_approx(J)
         return (phi_cos, phi_sin), psd
 
-    def prior(self, name: str, X: TensorLike, dims: str | None = None):  # type: ignore
+    def prior(
+        self,
+        name: str,
+        X: TensorLike,
+        dims: str | None = None,
+    ):  # type: ignore
         R"""
         Returns the (approximate) GP prior distribution evaluated over the input locations `X`.
         For usage examples, refer to `pm.gp.Latent`.
@@ -695,16 +705,23 @@ def prior(self, name: str, X: TensorLike, dims: str | None = None):  # type: ign
         (phi_cos, phi_sin), psd = self.prior_linearized(X)
 
         m = self._m
-        self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(m * 2 - 1))
-        # The first eigenfunction for the sine component is zero
-        # and so does not contribute to the approximation.
-        f = (
-            self.mean_func(X)
-            + phi_cos @ (psd * self._beta[:m])  # type: ignore
-            + phi_sin[..., 1:] @ (psd[1:] * self._beta[m:])  # type: ignore
-        )
+        gp_dims = dims
+        size = 2 * m - 1
+        parametrization = "noncentered"
+        if parametrization == "noncentered":
+            self._beta = pm.Normal(
+                f"{name}_hsgp_coeffs_",
+                size=size,
+            )
+            # The first eigenfunction for the sine component is zero
+            # and so does not contribute to the approximation.
+            f = (
+                self.mean_func(X)
+                + phi_cos @ (self._beta[:m] * psd)  # type: ignore
+                + phi_sin[..., 1:] @ (self._beta[m:] * psd[1:])  # type: ignore
+            )
 
-        self.f = pm.Deterministic(name, f, dims=dims)
+        self.f = pm.Deterministic(name, f, dims=gp_dims)
         return self.f
 
     def _build_conditional(self, Xnew):