gpjax 0.13.0__py3-none-any.whl → 0.13.2__py3-none-any.whl

gpjax/__init__.py

@@ -40,7 +40,7 @@ __license__ = "MIT"
 __description__ = "Gaussian processes in JAX and Flax"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.13.0"
+__version__ = "0.13.2"
 
 __all__ = [
     "gps",

gpjax/kernels/computations/basis_functions.py

@@ -6,9 +6,7 @@ from jaxtyping import Float
 import gpjax
 from gpjax.kernels.computations.base import AbstractKernelComputation
 from gpjax.linalg import (
-    Dense,
     Diagonal,
-    psd,
 )
 from gpjax.typing import Array
 
@@ -27,9 +25,9 @@ class BasisFunctionComputation(AbstractKernelComputation):
         z2 = self.compute_features(kernel, y)
         return self.scaling(kernel) * jnp.matmul(z1, z2.T)
 
-    def _gram(self, kernel: K, inputs: Float[Array, "N D"]) -> Dense:
+    def _gram(self, kernel: K, inputs: Float[Array, "N D"]) -> Float[Array, "N N"]:
         z1 = self.compute_features(kernel, inputs)
-        return psd(Dense(self.scaling(kernel) * jnp.matmul(z1, z1.T)))
+        return self.scaling(kernel) * jnp.matmul(z1, z1.T)
 
     def diagonal(self, kernel: K, inputs: Float[Array, "N D"]) -> Diagonal:
         r"""For a given kernel, compute the elementwise diagonal of the

gpjax/kernels/computations/eigen.py

@@ -15,7 +15,6 @@
 
 
 import beartype.typing as tp
-import jax.numpy as jnp
 from jaxtyping import (
     Float,
     Num,
@@ -39,17 +38,4 @@ class EigenKernelComputation(AbstractKernelComputation):
     def _cross_covariance(
         self, kernel: Kernel, x: Num[Array, "N D"], y: Num[Array, "M D"]
     ) -> Float[Array, "N M"]:
-        # Transform the eigenvalues of the graph Laplacian according to the
-        # RBF kernel's SPDE form.
-        S = jnp.power(
-            kernel.eigenvalues
-            + 2
-            * kernel.smoothness.value
-            / kernel.lengthscale.value
-            / kernel.lengthscale.value,
-            -kernel.smoothness.value,
-        )
-        S = jnp.multiply(S, kernel.num_vertex / jnp.sum(S))
-        # Scale the transform eigenvalues by the kernel variance
-        S = jnp.multiply(S, kernel.variance.value)
-        return kernel(x, y, S=S)
+        return kernel(x, y)

gpjax/kernels/non_euclidean/graph.py

@@ -25,7 +25,10 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     EigenKernelComputation,
 )
-from gpjax.kernels.non_euclidean.utils import jax_gather_nd
+from gpjax.kernels.non_euclidean.utils import (
+    calculate_heat_semigroup,
+    jax_gather_nd,
+)
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.parameters import (
     Parameter,
@@ -98,14 +101,12 @@ class GraphKernel(StationaryKernel):
 
         super().__init__(active_dims, lengthscale, variance, n_dims, compute_engine)
 
-    def __call__(  # TODO not consistent with general kernel interface
+    def __call__(
         self,
         x: Int[Array, "N 1"],
-        y: Int[Array, "N 1"],
-        *,
-        S,
-        **kwargs,
+        y: Int[Array, "M 1"],
     ):
+        S = calculate_heat_semigroup(self)
         Kxx = (jax_gather_nd(self.eigenvectors, x) * S.squeeze()) @ jnp.transpose(
             jax_gather_nd(self.eigenvectors, y)
         )  # shape (n,n)

gpjax/kernels/non_euclidean/utils.py

@@ -13,6 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
+import beartype.typing as tp
+import jax.numpy as jnp
 from jaxtyping import (
     Float,
     Int,
@@ -20,6 +24,9 @@ from jaxtyping import (
 
 from gpjax.typing import Array
 
+if tp.TYPE_CHECKING:
+    from gpjax.kernels.non_euclidean.graph import GraphKernel
+
 
 def jax_gather_nd(
     params: Float[Array, " N *rest"], indices: Int[Array, " M 1"]
@@ -41,3 +48,26 @@ def jax_gather_nd(
     """
     tuple_indices = tuple(indices[..., i] for i in range(indices.shape[-1]))
     return params[tuple_indices]
+
+
+def calculate_heat_semigroup(kernel: GraphKernel) -> Float[Array, "N M"]:
+    r"""Returns the rescaled heat semigroup, S
+
+    Args:
+        kernel: instance of the graph kernel
+
+    Returns:
+        S
+    """
+    S = jnp.power(
+        kernel.eigenvalues
+        + 2
+        * kernel.smoothness.value
+        / kernel.lengthscale.value
+        / kernel.lengthscale.value,
+        -kernel.smoothness.value,
+    )
+    S = jnp.multiply(S, kernel.num_vertex / jnp.sum(S))
+    # Scale the transform eigenvalues by the kernel variance
+    S = jnp.multiply(S, kernel.variance.value)
+    return S

gpjax/variational_families.py

@@ -19,7 +19,10 @@ import beartype.typing as tp
 from flax import nnx
 import jax.numpy as jnp
 import jax.scipy as jsp
-from jaxtyping import Float
+from jaxtyping import (
+    Float,
+    Int,
+)
 
 from gpjax.dataset import Dataset
 from gpjax.distributions import GaussianDistribution
@@ -108,6 +111,7 @@ class AbstractVariationalGaussian(AbstractVariationalFamily[L]):
         self,
         posterior: AbstractPosterior[P, L],
         inducing_inputs: tp.Union[
+            Int[Array, "N D"],
             Float[Array, "N D"],
             Real,
         ],
@@ -140,7 +144,7 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
     def __init__(
         self,
         posterior: AbstractPosterior[P, L],
-        inducing_inputs: Float[Array, "N D"],
+        inducing_inputs: tp.Union[Int[Array, "N D"], Float[Array, "N D"]],
         variational_mean: tp.Union[Float[Array, "N 1"], None] = None,
         variational_root_covariance: tp.Union[Float[Array, "N N"], None] = None,
         jitter: ScalarFloat = 1e-6,
@@ -156,6 +160,12 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         self.variational_mean = Real(variational_mean)
         self.variational_root_covariance = LowerTriangular(variational_root_covariance)
 
+    def _fmt_Kzt_Ktt(self, Kzt, Ktt):
+        return Kzt, Ktt
+
+    def _fmt_inducing_inputs(self):
+        return self.inducing_inputs.value
+
     def prior_kl(self) -> ScalarFloat:
         r"""Compute the prior KL divergence.
 
@@ -178,7 +188,7 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         # Unpack variational parameters
         variational_mean = self.variational_mean.value
         variational_sqrt = self.variational_root_covariance.value
-        inducing_inputs = self.inducing_inputs.value
+        inducing_inputs = self._fmt_inducing_inputs()
 
         # Unpack mean function and kernel
         mean_function = self.posterior.prior.mean_function
@@ -202,7 +212,9 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
 
         return q_inducing.kl_divergence(p_inducing)
 
-    def predict(self, test_inputs: Float[Array, "N D"]) -> GaussianDistribution:
+    def predict(
+        self, test_inputs: tp.Union[Int[Array, "N D"], Float[Array, "N D"]]
+    ) -> GaussianDistribution:
         r"""Compute the predictive distribution of the GP at the test inputs t.
 
         This is the integral $q(f(t)) = \int p(f(t)\mid u) q(u) \mathrm{d}u$, which
@@ -222,7 +234,7 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         # Unpack variational parameters
         variational_mean = self.variational_mean.value
         variational_sqrt = self.variational_root_covariance.value
-        inducing_inputs = self.inducing_inputs.value
+        inducing_inputs = self._fmt_inducing_inputs()
 
         # Unpack mean function and kernel
         mean_function = self.posterior.prior.mean_function
@@ -241,6 +253,8 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         Kzt = kernel.cross_covariance(inducing_inputs, test_points)
         test_mean = mean_function(test_points)
 
+        Kzt, Ktt = self._fmt_Kzt_Ktt(Kzt, Ktt)
+
         # Lz⁻¹ Kzt
         Lz_inv_Kzt = solve(Lz, Kzt)
 
@@ -259,8 +273,10 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
             - jnp.matmul(Lz_inv_Kzt.T, Lz_inv_Kzt)
             + jnp.matmul(Ktz_Kzz_inv_sqrt, Ktz_Kzz_inv_sqrt.T)
         )
+
         if hasattr(covariance, "to_dense"):
             covariance = covariance.to_dense()
+
         covariance = add_jitter(covariance, self.jitter)
         covariance = Dense(covariance)
 
@@ -269,6 +285,53 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         )
 
 
+class GraphVariationalGaussian(VariationalGaussian[L]):
+    r"""A variational Gaussian defined over graph-structured inducing inputs.
+
+    This subclass adapts the :class:`VariationalGaussian` family to the
+    case where the inducing inputs are discrete graph node indices rather
+    than continuous spatial coordinates.
+
+    The main differences are:
+      * Inducing inputs are integer node IDs.
+      * Kernel matrices are ensured to be dense and 2D.
+    """
+
+    def __init__(
+        self,
+        posterior: AbstractPosterior[P, L],
+        inducing_inputs: Int[Array, "N D"],
+        variational_mean: tp.Union[Float[Array, "N 1"], None] = None,
+        variational_root_covariance: tp.Union[Float[Array, "N N"], None] = None,
+        jitter: ScalarFloat = 1e-6,
+    ):
+        super().__init__(
+            posterior,
+            inducing_inputs,
+            variational_mean,
+            variational_root_covariance,
+            jitter,
+        )
+        self.inducing_inputs = self.inducing_inputs.value.astype(jnp.int64)
+
+    def _fmt_Kzt_Ktt(self, Kzt, Ktt):
+        Ktt = Ktt.to_dense() if hasattr(Ktt, "to_dense") else Ktt
+        Kzt = Kzt.to_dense() if hasattr(Kzt, "to_dense") else Kzt
+        Ktt = jnp.atleast_2d(Ktt)
+        Kzt = (
+            jnp.transpose(jnp.atleast_2d(Kzt)) if Kzt.ndim < 2 else jnp.atleast_2d(Kzt)
+        )
+        return Kzt, Ktt
+
+    def _fmt_inducing_inputs(self):
+        return self.inducing_inputs
+
+    @property
+    def num_inducing(self) -> int:
+        """The number of inducing inputs."""
+        return self.inducing_inputs.shape[0]
+
+
 class WhitenedVariationalGaussian(VariationalGaussian[L]):
     r"""The whitened variational Gaussian family of probability distributions.
 
@@ -811,6 +874,7 @@ __all__ = [
     "AbstractVariationalFamily",
     "AbstractVariationalGaussian",
     "VariationalGaussian",
+    "GraphVariationalGaussian",
     "WhitenedVariationalGaussian",
     "NaturalVariationalGaussian",
     "ExpectationVariationalGaussian",

{gpjax-0.13.0.dist-info → gpjax-0.13.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.13.0
+Version: 0.13.2
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues

{gpjax-0.13.0.dist-info → gpjax-0.13.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-gpjax/__init__.py,sha256=dB99f4sP0k5SfxABCyVMYkLVf-UVCXY-iR6153RRpv8,1641
+gpjax/__init__.py,sha256=jpuNxARmW3gOmXFAhJvrXeMk_rvYDn49E-8wy3ATZKg,1641
 gpjax/citation.py,sha256=pwFS8h1J-LE5ieRS0zDyuwhmQHNxkFHYE7iSMlVNmQc,3928
 gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
 gpjax/distributions.py,sha256=iKmeQ_NN2CIjRiuOeJlwEGASzGROi4ZCerVi1uY7zRM,7758
@@ -12,21 +12,21 @@ gpjax/objectives.py,sha256=GvKbDIPqYjsc9FpiTccmZwRdHr2lCykgfxI9BX9I_GA,15362
 gpjax/parameters.py,sha256=hnyIKr6uIzd7Kb3KZC9WowR88ruQwUvdcto3cx2ZDv4,6756
 gpjax/scan.py,sha256=jStQvwkE9MGttB89frxam1kaeXdWih7cVxkGywyaeHQ,5365
 gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
-gpjax/variational_families.py,sha256=TJGGkwkE805X4PQb-C32FxvD9B_OsFLWf6I-ZZvOUWk,29628
+gpjax/variational_families.py,sha256=x4VnUh9GiW77ijnDEwrYzH0aWFOW6NLPQXEp--I9R-g,31566
 gpjax/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
 gpjax/kernels/base.py,sha256=4vUV6_wnbV64EUaOecVWE9Vz4G5Dg6xUMNsQSis8A1o,11561
 gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
 gpjax/kernels/approximations/rff.py,sha256=WJ2Bw4RRW5bLXGKfStxph65H-ozDelvt_4mxXQm-g7w,4074
 gpjax/kernels/computations/__init__.py,sha256=uTVkqvnZVesFLDN92h0ZR0jfR69Eo2WyjOlmSYmCPJ8,1379
 gpjax/kernels/computations/base.py,sha256=L6K0roxZbrYeJKxEw-yaTiK9Mtcv0YtZfWI2Xnau7i8,3616
-gpjax/kernels/computations/basis_functions.py,sha256=_SFv4Tiwne40bxr1uVYpEjjZgjIQHKseLmss2Zgl1L4,2484
+gpjax/kernels/computations/basis_functions.py,sha256=Y6OS4UEvo-Mdn7yRMWPde5rOvewDHIleaUIVBOpCTd4,2466
 gpjax/kernels/computations/constant_diagonal.py,sha256=JkQhLj7cK48IhOER4ivkALNhD1oQleKe-Rr9BtUJ6es,1984
 gpjax/kernels/computations/dense.py,sha256=vnW6XKQe4_gzpXRWTctxhgMA9-9TebdtiXzAqh_-j6g,1392
 gpjax/kernels/computations/diagonal.py,sha256=k1KqW0DwWRIBvbb7jzcKktXRfhXbcos3ncWrFplJ4W0,1768
-gpjax/kernels/computations/eigen.py,sha256=NTHm-cn-RepYuXFrvXo2ih7Gtu1YR_pAg4Jb7IhE_o8,1930
+gpjax/kernels/computations/eigen.py,sha256=LuwYVPK0AuDGNSccPmAS8wyDJ2ngkRbc6kNZzqiJEOg,1380
 gpjax/kernels/non_euclidean/__init__.py,sha256=RT7puRPqCTpyxZ16q596EuOQEQi1LK1v3J9_fWz1NlY,790
-gpjax/kernels/non_euclidean/graph.py,sha256=xTrx6ro8ubRXgM7Wgg6NmOyyEjEcGhzydY7KXueknCc,4120
-gpjax/kernels/non_euclidean/utils.py,sha256=z42aw8ga0zuREzHawemR9okttgrAUPmq-aN5HMt4SuY,1578
+gpjax/kernels/non_euclidean/graph.py,sha256=MukL4ZghUywkau4tVWmHsrAjzTcl9kHN6onEAvxSvjc,4109
+gpjax/kernels/non_euclidean/utils.py,sha256=vxjRX209tOL3jKsw69hQfVtBqkASDZ__4tJOGI4oLjo,2342
 gpjax/kernels/nonstationary/__init__.py,sha256=YpWQfOy_cqOKc5ezn37vqoK3Z6jznYiJz28BD_8F7AY,930
 gpjax/kernels/nonstationary/arccosine.py,sha256=cqb8sqaNwW3fEbrA7MY9OF2KJFTkxHhqwmQtABE3G8w,5408
 gpjax/kernels/nonstationary/linear.py,sha256=UIMoCq2hg6dQKr4J5UGiiPqotBleQuYfy00Ia1NaMOo,2571
@@ -46,7 +46,7 @@ gpjax/linalg/__init__.py,sha256=F8mxk_9Zc2nFd7Q-unjJ50_6rXEKzZj572WsU_jUKqI,547
 gpjax/linalg/operations.py,sha256=xvhOy5P4FmUCPWjIVNdg1yDXaoFQ48anFUfR-Tnfr6k,6480
 gpjax/linalg/operators.py,sha256=arxRGwcoAy_RqUYqBpZ3XG6OXbjShUl7m8sTpg85npE,11608
 gpjax/linalg/utils.py,sha256=fKV8G_iKZVhNkNvN20D_dQEi93-8xosGbXBP-v7UEyo,2020
-gpjax-0.13.0.dist-info/METADATA,sha256=LZllpjzOW4QDzpUX-haSNrLMXzkiREGpaYveG6jeufE,10400
-gpjax-0.13.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-gpjax-0.13.0.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
-gpjax-0.13.0.dist-info/RECORD,,
+gpjax-0.13.2.dist-info/METADATA,sha256=_dkTRApmtHexaDFVnH8FhpOf_YkIyci2Gihg9elVeII,10400
+gpjax-0.13.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+gpjax-0.13.2.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
+gpjax-0.13.2.dist-info/RECORD,,