Remove dead code

Author: ricardoV94

pymc_experimental/gp/pytensor_gp.py

@@ -13,7 +13,7 @@ class GPCovariance(OpFromGraph):
     """OFG representing a GP covariance"""
 
     @staticmethod
-    def square_dist_Xs(X, Xs, ls):
+    def square_dist(X, Xs, ls):
         assert X.ndim == 2, "Complain to Bill about it"
         assert Xs.ndim == 2, "Complain to Bill about it"
 
@@ -24,20 +24,8 @@ def square_dist_Xs(X, Xs, ls):
         Xs2 = pt.sum(pt.square(Xs), axis=-1)
 
         sqd = -2.0 * X @ Xs.mT + (X2[..., :, None] + Xs2[..., None, :])
-        # sqd = -2.0 * pt.dot(X, pt.transpose(Xs)) + (
-        #         pt.reshape(X2, (-1, 1)) + pt.reshape(Xs2, (1, -1))
-        # )
-
         return pt.clip(sqd, 0, pt.inf)
 
-    @staticmethod
-    def square_dist(X, ls):
-        X = X / ls
-        X2 = pt.sum(pt.square(X), axis=-1)
-        sqd = -2.0 * X @ X.mT + (X2[..., :, None] + X2[..., None, :])
-
-        return sqd
-
 
 class ExpQuadCov(GPCovariance):
     """
@@ -46,7 +34,7 @@ class ExpQuadCov(GPCovariance):
 
     @classmethod
     def exp_quad_full(cls, X, Xs, ls):
-        return pt.exp(-0.5 * cls.square_dist_Xs(X, Xs, ls))
+        return pt.exp(-0.5 * cls.square_dist(X, Xs, ls))
 
     @classmethod
     def build_covariance(cls, X, Xs=None, *, ls):
@@ -64,32 +52,10 @@ def build_covariance(cls, X, Xs=None, *, ls):
             return cls(inputs=[X, Xs, ls], outputs=[out])(X, Xs, ls)
 
 
-def ExpQuad(X, X_new=None, *, ls):
+def ExpQuad(X, X_new=None, *, ls=1.0):
     return ExpQuadCov.build_covariance(X, X_new, ls=ls)
 
 
-# class WhiteNoiseCov(GPCovariance):
-#     @classmethod
-#     def white_noise_full(cls, X, sigma):
-#         X_shape = tuple(X.shape)
-#         shape = X_shape[:-1] + (X_shape[-2],)
-#
-#         return _delta(shape, normalize_axis_tuple((-1, -2), X.ndim)) * sigma**2
-#
-#     @classmethod
-#     def build_covariance(cls, X, sigma):
-#         X = pt.as_tensor(X)
-#         sigma = pt.as_tensor(sigma)
-#
-#         ofg = cls(inputs=[X, sigma], outputs=[cls.white_noise_full(X, sigma)])
-#         return ofg(X, sigma)
-
-#
-# def WhiteNoise(X, sigma):
-#     return WhiteNoiseCov.build_covariance(X, sigma)
-#
-
-
 class GP_RV(pm.MvNormal.rv_type):
     name = "gaussian_process"
     signature = "(n),(n,n)->(n)"
@@ -103,7 +69,6 @@ class GP(Continuous):
 
     @classmethod
     def dist(cls, cov, **kwargs):
-        # return Assert(msg="Don't know what a GP_RV is")(False)
         cov = pt.as_tensor(cov)
         mu = pt.zeros(cov.shape[-1])
         return super().dist([mu, cov], **kwargs)
@@ -190,89 +155,3 @@ def _build_conditional(Xnew, f, cov, jitter):
         fgraph.add_output(gp_model_var_star, import_missing=True)
 
     return model_from_fgraph(fgraph, mutate_fgraph=True)
-
-
-# @register_canonicalize
-# @node_rewriter(tracks=[pm.MvNormal.rv_type])
-# def GP_normal_mvnormal_conjugacy(fgraph: FunctionGraph, node):
-#     # TODO: Should this alert users that it can't be applied when the GP is in a deterministic?
-#     gp_rng, gp_size, mu, cov = node.inputs
-#     next_gp_rng, gp_rv = node.outputs
-#
-#     if not isinstance(cov.owner.op, GPCovariance):
-#         return
-#
-#     for client, input_index in fgraph.clients[gp_rv]:
-#         # input_index is 2 because it goes (rng, size, mu, sigma), and we want the mu
-#         # to be the GP we're looking
-#         if isinstance(client.op, pm.Normal.rv_type) and (input_index == 2):
-#             next_normal_rng, normal_rv = client.outputs
-#             normal_rng, normal_size, mu, sigma = client.inputs
-#
-#             if normal_rv.ndim != gp_rv.ndim:
-#                 return
-#
-#             X = cov.owner.inputs[0]
-#
-#             white_noise = WhiteNoiseCov.build_covariance(X, sigma)
-#             white_noise.name = 'WhiteNoiseCov'
-#             cov = cov + white_noise
-#
-#             if not rv_size_is_none(normal_size):
-#                 normal_size = tuple(normal_size)
-#                 new_gp_size = normal_size[:-1]
-#                 core_shape = normal_size[-1]
-#
-#                 cov_shape = (*(None,) * (cov.ndim - 2), core_shape, core_shape)
-#                 cov = pt.specify_shape(cov, cov_shape)
-#
-#             else:
-#                 new_gp_size = None
-#
-#             next_new_gp_rng, new_gp_mvn = pm.MvNormal.dist(cov=cov, rng=gp_rng, size=new_gp_size).owner.outputs
-#             new_gp_mvn.name = 'NewGPMvn'
-#
-#             # Check that the new shape is at least as specific as the shape we are replacing
-#             for new_shape, old_shape in zip(new_gp_mvn.type.shape, normal_rv.type.shape, strict=True):
-#                 if new_shape is None:
-#                     assert old_shape is None
-#
-#             return {
-#                 next_normal_rng: next_new_gp_rng,
-#                 normal_rv: new_gp_mvn,
-#                 next_gp_rng: next_new_gp_rng
-#             }
-#
-#         else:
-#             return None
-#
-# #TODO: Why do I need to register this twice?
-# specialization_ir_rewrites_db.register(
-#     GP_normal_mvnormal_conjugacy.__name__,
-#     GP_normal_mvnormal_conjugacy,
-#     "basic",
-# )
-
-# @node_rewriter(tracks=[pm.MvNormal.rv_type])
-# def GP_normal_marginal_logp(fgraph: FunctionGraph, node):
-#     """
-#     Replace Normal(GP(cov), sigma) -> MvNormal(0, cov + diag(sigma)).
-#     """
-#     rng, size, mu, cov = node.inputs
-#     if cov.owner and cov.owner.op == matrix_inverse:
-#         tau = cov.owner.inputs[0]
-#         return PrecisionMvNormalRV.rv_op(mu, tau, size=size, rng=rng).owner.outputs
-#     return None
-#
-
-# cov_op = GPCovariance()
-# gp_op = GP("vanilla")
-# # SymbolicRandomVariable.register(type(gp_op))
-# prior_from_gp = PriorFromGP()
-#
-# MeasurableVariable.register(type(prior_from_gp))
-#
-#
-# @_get_measurable_outputs.register(type(prior_from_gp))
-# def gp_measurable_outputs(op, node):
-#     return node.outputs

tests/test_gp.py

@@ -7,17 +7,7 @@
 from pymc_experimental.gp.pytensor_gp import GP, ExpQuad, conditional_gp
 
 
-def test_exp_quad():
-    x = pt.arange(3)[:, None]
-    ls = pt.ones(())
-    cov = ExpQuad(x, ls=ls).eval()
-    expected_distance = np.array([[0.0, 1.0, 4.0], [1.0, 0.0, 1.0], [4.0, 1.0, 0.0]])
-
-    np.testing.assert_allclose(cov, np.exp(-0.5 * expected_distance))
-
-
-# @pytest.fixture(scope="session")
-def latent_model():
+def build_latent_model():
     with pm.Model() as m:
         X = pm.Data("X", np.arange(3)[:, None])
         y = np.full(3, np.pi)
@@ -31,7 +21,7 @@ def latent_model():
     return m
 
 
-def latent_model_old_API():
+def build_latent_model_old_API():
     with pm.Model() as m:
         X = pm.Data("X", np.arange(3)[:, None])
         y = np.full(3, np.pi)
@@ -46,9 +36,18 @@ def latent_model_old_API():
     return m, gp_class
 
 
+def test_exp_quad():
+    x = pt.arange(3)[:, None]
+    ls = pt.ones(())
+    cov = ExpQuad(x, ls=ls).eval()
+    expected_distance = np.array([[0.0, 1.0, 4.0], [1.0, 0.0, 1.0], [4.0, 1.0, 0.0]])
+
+    np.testing.assert_allclose(cov, np.exp(-0.5 * expected_distance))
+
+
 def test_latent_model_prior():
-    m = latent_model()
-    ref_m, _ = latent_model_old_API()
+    m = build_latent_model()
+    ref_m, _ = build_latent_model_old_API()
 
     prior = pm.draw(m["gp"], draws=1000)
     prior_ref = pm.draw(ref_m["gp"], draws=1000)
@@ -67,10 +66,10 @@ def test_latent_model_prior():
 
 
 def test_latent_model_logp():
-    m = latent_model()
+    m = build_latent_model()
     ip = m.initial_point()
 
-    ref_m, _ = latent_model_old_API()
+    ref_m, _ = build_latent_model_old_API()
 
     np.testing.assert_allclose(
         m.compile_logp()(ip),
@@ -89,7 +88,7 @@ def test_latent_model_conditional(inline):
 
     new_x = np.array([3, 4])[:, None]
 
-    m = latent_model()
+    m = build_latent_model()
     with m:
         pm.Deterministic("gp_exp", m["gp"].exp())
 
@@ -100,7 +99,7 @@ def test_latent_model_conditional(inline):
             progressbar=False,
         ).posterior_predictive
 
-    ref_m, ref_gp_class = latent_model_old_API()
+    ref_m, ref_gp_class = build_latent_model_old_API()
     with ref_m:
         gp_star = ref_gp_class.conditional("gp_star", Xnew=new_x)
         pred_ref = pm.sample_posterior_predictive(