gpjax 0.10.0__tar.gz → 0.10.2__tar.gz

gpjax-0.10.2/.cursorrules

@@ -0,0 +1,37 @@
+You are an AI assistant specialized in Python development and machine learning. Your approach emphasizes:
+
+Clear project structure with separate directories for source code, tests, docs, and config.
+
+Modular design with distinct files for models, services, controllers, and utilities.
+
+Configuration management using environment variables.
+
+Robust error handling and logging, including context capture.
+
+Comprehensive testing with pytest.
+
+Detailed documentation using docstrings and README files.
+
+Code style consistency using Ruff.
+
+CI/CD implementation with GitHub Actions or GitLab CI.
+
+AI-friendly coding practices:
+
+You provide code snippets and explanations tailored to these principles, optimizing for clarity and AI-assisted development.
+
+Follow the following rules:
+
+For any python file, be sure to ALWAYS add typing annotations to each function or class. Be sure to include return types when necessary. Add descriptive docstrings to all python functions and classes as well. Please use pep257 convention. Update existing docstrings if need be.
+
+Make sure you keep any comments that exist in a file.
+
+When writing tests, make sure that you ONLY use pytest or pytest plugins, do NOT use the unittest module. All tests should have typing annotations as well. All tests should be in ./tests. Be sure to create all necessary files and folders. If you are creating files inside of ./tests or ./src/goob_ai, be sure to make a init.py file if one does not exist.
+
+All tests should be fully annotated and should contain docstrings. Be sure to import the following if TYPE_CHECKING:
+
+from _pytest.capture import CaptureFixture
+from _pytest.fixtures import FixtureRequest
+from _pytest.logging import LogCaptureFixture
+from _pytest.monkeypatch import MonkeyPatch
+from pytest_mock.plugin import MockerFixture

{gpjax-0.10.0 → gpjax-0.10.2}/PKG-INFO

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

{gpjax-0.10.0 → gpjax-0.10.2}/docs/sharp_bits.md

@@ -175,3 +175,60 @@ mini-batch optimisation of the parameters of your sparse Gaussian process model.
 model will scale linearly in the batch size and quadratically in the number of inducing
 points. We demonstrate its use in
 [our sparse stochastic variational inference notebook](_examples/uncollapsed_vi.md).
+
+## JIT compilation
+
+There are a subset of operations in GPJax that are not JIT compatible by default. This
+is because we have assertions in place to check the properties of the parameters. For
+example, we check that the lengthscale parameter that a user provides is positive. This
+makes for a better user experience as we can provide more informative error messages;
+however, JIT compiling functions wherein these assertions are made will break the code.
+As an example, consider the following code:
+
+```python
+import jax
+import jax.numpy as jnp
+import gpjax as gpx
+
+x = jnp.linspace(0, 1, 10)[:, None]
+
+def compute_gram(lengthscale):
+    k = gpx.kernels.RBF(active_dims=[0], lengthscale=lengthscale, variance=jnp.array(1.0))
+    return k.gram(x)
+
+compute_gram(1.0)
+```
+
+so far so good. However, if we try to JIT compile this function, we will get an error:
+
+```python
+jit_compute_gram = jax.jit(compute_gram)
+try:
+    jit_compute_gram(1.0)
+except Exception as e:
+    print(e)
+```
+
+This error is due to the fact that the `RBF` kernel contains an assertion that checks
+that the lengthscale is positive. It does not matter that the assertion is satisfied;
+the very presence of the assertion will break JIT compilation.
+
+To resolve this, we can use the `checkify` decorator to remove the assertion. This will
+allow the function to be JIT compiled.
+
+```python
+from jax.experimental import checkify
+
+jit_compute_gram = jax.jit(checkify.checkify(compute_gram))
+error, value = jit_compute_gram(1.0)
+```
+By virtue of the `checkify.checkify`, a tuple is returned where the first element is the
+output of the assertion, and the second element is the value of the function. 
+
+This design is not perfect, and in an ideal world we would not enforce the user to wrap
+their code in `checkify.checkify`. We are actively looking into cleaner ways to provide
+guardrails in a less intrusive manner. However, for now, should you try to JIT compile
+a component of GPJax wherein there is an assertion, you will need to wrap the function
+in `checkify.checkify` as shown above.
+
+For more on `checkify`, please see the [JAX Checkify Doc](https://docs.jax.dev/en/latest/debugging/checkify_guide.html).

gpjax-0.10.2/examples/oak_example.py

@@ -0,0 +1,216 @@
+# -*- coding: utf-8 -*-
+# ---
+# jupyter:
+#   jupytext:
+#     cell_metadata_filter: -all
+#     custom_cell_magics: kql
+#     text_representation:
+#       extension: .py
+#       format_name: percent
+#       format_version: '1.3'
+#       jupytext_version: 1.16.7
+#   kernelspec:
+#     display_name: docs
+#     language: python
+#     name: python3
+# ---
+
+# %% [markdown]
+# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# %%
+"""Example of using the OrthogonalAdditiveKernel."""
+
+# %%
+import jax
+from jax import config
+import jax.numpy as jnp
+from matplotlib.colors import ListedColormap
+import matplotlib.pyplot as plt
+import optax
+
+import gpjax as gpx
+from gpjax.dataset import Dataset
+from gpjax.kernels import (
+    RBF,
+    OrthogonalAdditiveKernel,
+)
+from gpjax.typing import KeyArray
+
+config.update("jax_enable_x64", True)  # Enable Float64 precision
+
+
+# %%
+def f(x: jnp.ndarray) -> jnp.ndarray:
+    """Additive function with mixed dependencies:
+    f(x) = sin(π*x₁) + 2*cos(2π*x₂) + 0.5*sin(3π*x₁*x₂)
+
+    Args:
+        x: Input points array with shape (..., 2)
+
+    Returns:
+        Function values at the input points
+    """
+    return (
+        jnp.sin(jnp.pi * x[..., 0])
+        + 2.0 * jnp.cos(2.0 * jnp.pi * x[..., 1])
+        + 0.5 * jnp.sin(3.0 * jnp.pi * x[..., 0] * x[..., 1])
+    )
+
+
+# %%
+def generate_data(
+    key: KeyArray, n_train: int = 100, noise_std: float = 0.1
+) -> tuple[Dataset, jnp.ndarray, jnp.ndarray]:
+    """Generate synthetic training data.
+
+    Args:
+        key: JAX PRNG key for random number generation
+        n_train: Number of training points to generate
+        noise_std: Standard deviation of Gaussian observation noise
+
+    Returns:
+        Tuple of (training_data, X_test, meshgrid_for_plotting)
+    """
+    key1, key2, key3 = jax.random.split(key, 3)
+
+    # Generate training data
+    X_train = jax.random.uniform(key1, (n_train, 2))
+    y_train = f(X_train) + noise_std * jax.random.normal(key2, (n_train,))
+
+    training_data = Dataset(X=X_train, y=y_train[:, None])
+
+    # Generate test points for prediction
+    n_test = 20
+    x_range = jnp.linspace(0.0, 1.0, n_test)
+    X1, X2 = jnp.meshgrid(x_range, x_range)
+    X_test = jnp.vstack([X1.flatten(), X2.flatten()]).T
+
+    return training_data, X_test, (X1, X2)
+
+
+# %%
+def main():
+    # Set random seed for reproducibility
+    key = jax.random.PRNGKey(42)
+
+    # Generate synthetic training data
+    training_data, X_test, (X1, X2) = generate_data(key, n_train=100, noise_std=0.1)
+
+    # Create base kernel (RBF)
+    base_kernel = RBF(lengthscale=0.2)
+
+    # Create OAK kernel with second-order interactions
+    oak_kernel = OrthogonalAdditiveKernel(
+        base_kernel=base_kernel,
+        dim=2,
+        quad_deg=20,
+        second_order=True,
+    )
+
+    # Create a GP prior model
+    prior = gpx.gps.Prior(
+        mean_function=gpx.mean_functions.Zero(),
+        kernel=oak_kernel,
+    )
+
+    # Create a likelihood
+    likelihood = gpx.likelihoods.Gaussian(num_datapoints=training_data.n)
+
+    # Create the posterior
+    posterior = prior * likelihood
+
+    # Create parameter optimizer
+    optimizer = optax.adam(learning_rate=0.01)
+
+    # Define objective function for training
+    def objective(model, data):
+        return -model.mll(model.params, data)
+
+    # Optimize hyperparameters
+    opt_posterior, history = gpx.fit(
+        model=posterior,
+        objective=objective,
+        train_data=training_data,
+        optim=optimizer,
+        num_iters=300,
+        key=key,
+        verbose=True,
+    )
+
+    # Plot training curve
+    plt.figure(figsize=(10, 4))
+    plt.subplot(1, 2, 1)
+    plt.plot(history)
+    plt.title("Negative Log Marginal Likelihood")
+    plt.xlabel("Iteration")
+    plt.ylabel("NLML")
+
+    # Get posterior predictions
+    latent_dist = opt_posterior.predict(params=opt_posterior.params, x=X_test)
+    predictive_dist = opt_posterior.likelihood.condition(
+        latent_dist, opt_posterior.params
+    )
+    mu = predictive_dist.mean().reshape(X1.shape)
+    std = predictive_dist.stddev().reshape(X1.shape)
+
+    # Plot predictions
+    plt.subplot(1, 2, 2)
+    plt.contourf(X1, X2, mu, 50, cmap="viridis")
+    plt.colorbar(label="Predicted Mean")
+    plt.scatter(
+        training_data.X[:, 0],
+        training_data.X[:, 1],
+        c=training_data.y,
+        cmap=ListedColormap(["red", "blue"]),
+        alpha=0.6,
+        s=20,
+        edgecolors="k",
+    )
+    plt.title("OAK GP Predictions")
+    plt.xlabel("$x_1$")
+    plt.ylabel("$x_2$")
+
+    plt.tight_layout()
+    plt.savefig("oak_example.png", dpi=300)
+    plt.show()
+
+    # Print learned kernel parameters
+    print("\nLearned Parameters:")
+    print(f"Offset coefficient: {opt_posterior.params.kernel.offset.value}")
+    print(f"First-order coefficients: {opt_posterior.params.kernel.coeffs_1.value}")
+
+    # Analyze the importance of each dimension
+    importance_1st_order = opt_posterior.params.kernel.coeffs_1.value
+    total_importance = jnp.sum(importance_1st_order)
+    relative_importance = importance_1st_order / total_importance
+
+    print("\nRelative Importance of Input Dimensions:")
+    for i, imp in enumerate(relative_importance):
+        print(f"Dimension {i + 1}: {imp:.4f}")
+
+    if opt_posterior.params.kernel.coeffs_2 is not None:
+        # Analyze second-order interactions
+        coeffs_2 = opt_posterior.params.kernel.coeffs_2
+        print("\nSecond-order Interaction Coefficient:")
+        print(f"{coeffs_2[0, 1]:.4f}")
+
+
+# %%
+if __name__ == "__main__":
+    main()
+
+# %%

{gpjax-0.10.0 → gpjax-0.10.2}/gpjax/__init__.py

@@ -39,7 +39,7 @@ __license__ = "MIT"
 __description__ = "Didactic Gaussian processes in JAX"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.10.0"
+__version__ = "0.10.2"
 
 __all__ = [
     "base",

{gpjax-0.10.0 → gpjax-0.10.2}/gpjax/kernels/base.py

@@ -32,6 +32,7 @@ from gpjax.kernels.computations import (
 from gpjax.parameters import (
     Parameter,
     Real,
+    Static,
 )
 from gpjax.typing import (
     Array,
@@ -220,7 +221,9 @@ class Constant(AbstractKernel):
     def __init__(
         self,
         active_dims: tp.Union[list[int], slice, None] = None,
-        constant: tp.Union[ScalarFloat, Parameter[ScalarFloat]] = jnp.array(0.0),
+        constant: tp.Union[
+            ScalarFloat, Parameter[ScalarFloat], Static[ScalarFloat]
+        ] = jnp.array(0.0),
         compute_engine: AbstractKernelComputation = DenseKernelComputation(),
     ):
         if isinstance(constant, Parameter):

{gpjax-0.10.0 → gpjax-0.10.2}/gpjax/kernels/nonstationary/polynomial.py

@@ -46,7 +46,7 @@ class Polynomial(AbstractKernel):
         self,
         active_dims: tp.Union[list[int], slice, None] = None,
         degree: int = 2,
-        shift: tp.Union[ScalarFloat, nnx.Variable[ScalarArray]] = 0.0,
+        shift: tp.Union[ScalarFloat, nnx.Variable[ScalarArray]] = 1.0,
         variance: tp.Union[ScalarFloat, nnx.Variable[ScalarArray]] = 1.0,
         n_dims: tp.Union[int, None] = None,
         compute_engine: AbstractKernelComputation = DenseKernelComputation(),

{gpjax-0.10.0 → gpjax-0.10.2}/gpjax/mean_functions.py

@@ -28,6 +28,7 @@ from jaxtyping import (
 from gpjax.parameters import (
     Parameter,
     Real,
+    Static
 )
 from gpjax.typing import (
     Array,
@@ -130,9 +131,9 @@ class Constant(AbstractMeanFunction):
     """
 
     def __init__(
-        self, constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter] = 0.0
+        self, constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter, Static] = 0.0
     ):
-        if isinstance(constant, Parameter):
+        if isinstance(constant, Parameter) or isinstance(constant, Static):
             self.constant = constant
         else:
             self.constant = Real(jnp.array(constant))
@@ -158,7 +159,7 @@ class Zero(Constant):
     """
 
     def __init__(self):
-        super().__init__(constant=jnp.array(0.0))
+        super().__init__(constant=Static(jnp.array(0.0)))
 
 
 class CombinationMeanFunction(AbstractMeanFunction):

{gpjax-0.10.0 → gpjax-0.10.2}/gpjax/parameters.py

@@ -1,6 +1,7 @@
 import typing as tp
 
 from flax import nnx
+from jax.experimental import checkify
 import jax.numpy as jnp
 import jax.tree_util as jtu
 from jax.typing import ArrayLike
@@ -84,8 +85,7 @@ class PositiveReal(Parameter[T]):
 
     def __init__(self, value: T, tag: ParameterTag = "positive", **kwargs):
         super().__init__(value=value, tag=tag, **kwargs)
-
-        _check_is_positive(self.value)
+        _safe_assert(_check_is_positive, self.value)
 
 
 class Real(Parameter[T]):
@@ -101,7 +101,17 @@ class SigmoidBounded(Parameter[T]):
     def __init__(self, value: T, tag: ParameterTag = "sigmoid", **kwargs):
         super().__init__(value=value, tag=tag, **kwargs)
 
-        _check_in_bounds(self.value, 0.0, 1.0)
+        # Only perform validation in non-JIT contexts
+        if (
+            not isinstance(value, jnp.ndarray)
+            or not getattr(value, "aval", None) is None
+        ):
+            _safe_assert(
+                _check_in_bounds,
+                self.value,
+                low=jnp.array(0.0),
+                high=jnp.array(1.0),
+            )
 
 
 class Static(nnx.Variable[T]):
@@ -120,8 +130,13 @@ class LowerTriangular(Parameter[T]):
     def __init__(self, value: T, tag: ParameterTag = "lower_triangular", **kwargs):
         super().__init__(value=value, tag=tag, **kwargs)
 
-        _check_is_square(self.value)
-        _check_is_lower_triangular(self.value)
+        # Only perform validation in non-JIT contexts
+        if (
+            not isinstance(value, jnp.ndarray)
+            or not getattr(value, "aval", None) is None
+        ):
+            _safe_assert(_check_is_square, self.value)
+            _safe_assert(_check_is_lower_triangular, self.value)
 
 
 DEFAULT_BIJECTION = {
@@ -132,36 +147,83 @@ DEFAULT_BIJECTION = {
 }
 
 
-def _check_is_arraylike(value: T):
+def _check_is_arraylike(value: T) -> None:
+    """Check if a value is array-like.
+
+    Args:
+        value: The value to check.
+
+    Raises:
+        TypeError: If the value is not array-like.
+    """
     if not isinstance(value, (ArrayLike, list)):
         raise TypeError(
             f"Expected parameter value to be an array-like type. Got {value}."
         )
 
 
-def _check_is_positive(value: T):
-    if jnp.any(value < 0):
-        raise ValueError(
-            f"Expected parameter value to be strictly positive. Got {value}."
-        )
+@checkify.checkify
+def _check_is_positive(value):
+    checkify.check(
+        jnp.all(value > 0), "value needs to be positive, got {value}", value=value
+    )
 
 
-def _check_is_square(value: T):
-    if value.shape[0] != value.shape[1]:
-        raise ValueError(
-            f"Expected parameter value to be a square matrix. Got {value}."
-        )
+@checkify.checkify
+def _check_is_square(value: T) -> None:
+    """Check if a value is a square matrix.
 
+    Args:
+        value: The value to check.
 
-def _check_is_lower_triangular(value: T):
-    if not jnp.all(jnp.tril(value) == value):
-        raise ValueError(
-            f"Expected parameter value to be a lower triangular matrix. Got {value}."
-        )
+    Raises:
+        ValueError: If the value is not a square matrix.
+    """
+    checkify.check(
+        value.shape[0] == value.shape[1],
+        "value needs to be a square matrix, got {value}",
+        value=value,
+    )
 
 
-def _check_in_bounds(value: T, low: float, high: float):
-    if jnp.any((value < low) | (value > high)):
-        raise ValueError(
-            f"Expected parameter value to be bounded between {low} and {high}. Got {value}."
-        )
+@checkify.checkify
+def _check_is_lower_triangular(value: T) -> None:
+    """Check if a value is a lower triangular matrix.
+
+    Args:
+        value: The value to check.
+
+    Raises:
+        ValueError: If the value is not a lower triangular matrix.
+    """
+    checkify.check(
+        jnp.all(jnp.tril(value) == value),
+        "value needs to be a lower triangular matrix, got {value}",
+        value=value,
+    )
+
+
+@checkify.checkify
+def _check_in_bounds(value: T, low: T, high: T) -> None:
+    """Check if a value is bounded between low and high.
+
+    Args:
+        value: The value to check.
+        low: The lower bound.
+        high: The upper bound.
+
+    Raises:
+        ValueError: If any element of value is outside the bounds.
+    """
+    checkify.check(
+        jnp.all((value >= low) & (value <= high)),
+        "value needs to be bounded between {low} and {high}, got {value}",
+        value=value,
+        low=low,
+        high=high,
+    )
+
+
+def _safe_assert(fn: tp.Callable[[tp.Any], None], value: T, **kwargs) -> None:
+    error, _ = fn(value, **kwargs)
+    checkify.check_error(error)

{gpjax-0.10.0 → gpjax-0.10.2}/tests/test_mean_functions.py

@@ -5,17 +5,20 @@ config.update("jax_enable_x64", True)
 
 
 import jax.numpy as jnp
+import jax.random as jr
 from jaxtyping import (
     Array,
     Float,
 )
 import pytest
 
+import gpjax as gpx
 from gpjax.mean_functions import (
     AbstractMeanFunction,
     Constant,
     Zero,
 )
+from gpjax.parameters import Static
 
 
 def test_abstract() -> None:
@@ -49,38 +52,28 @@ def test_constant(constant: Float[Array, " Q"]) -> None:
     ).all()
 
 
-# TODO: rewrite this test after work on fit
-# def test_zero_mean_remains_zero() -> None:
-#     key = jr.PRNGKey(123)
-
-#     x = jr.uniform(key=key, minval=0, maxval=1, shape=(20, 1))
-#     y = jnp.full((20, 1), 50, dtype=jnp.float64)  # Dataset with non-zero mean
-#     D = gpx.Dataset(X=x, y=y)
-
-#     kernel = gpx.kernels.Constant(constant=jnp.array(0.0))
-#     kernel = kernel.replace_trainable(
-#         constant=False
-#     )  # Prevent kernel from modelling non-zero mean
-#     meanf = Zero()
-#     prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
-#     likelihood = gpx.likelihoods.Gaussian(
-#         num_datapoints=D.n, obs_stddev=jnp.array(1e-3)
-#     )
-#     likelihood = likelihood.replace_trainable(obs_stddev=False)
-#     posterior = prior * likelihood
-
-#     negative_mll = gpx.objectives.ConjugateMLL(negative=True)
-#     opt_posterior, _ = gpx.fit(
-#         model=posterior,
-#         objective=negative_mll,
-#         train_data=D,
-#         optim=ox.adam(learning_rate=0.5),
-#         num_iters=1000,
-#         safe=True,
-#         key=key,
-#     )
-
-#     assert opt_posterior.prior.mean_function.constant == 0.0
+def test_zero_mean_remains_zero() -> None:
+    key = jr.PRNGKey(123)
+
+    x = jr.uniform(key=key, minval=0, maxval=1, shape=(20, 1))
+    y = jnp.full((20, 1), 50, dtype=jnp.float64)  # Dataset with non-zero mean
+    D = gpx.Dataset(X=x, y=y)
+
+    constant = Static(jnp.array(0.0))
+    kernel = gpx.kernels.Constant(constant=constant)
+    meanf = Zero()
+    prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
+    likelihood = gpx.likelihoods.Gaussian(
+        num_datapoints=D.n, obs_stddev=jnp.array(1e-3)
+    )
+    posterior = prior * likelihood
+
+    opt_posterior, _ = gpx.fit_scipy(
+        model=posterior,
+        objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
+        train_data=D,
+    )
+    assert opt_posterior.prior.mean_function.constant.value == 0.0
 
 
 def test_initialising_zero_mean_with_constant_raises_error():