gpjax 0.11.0__py3-none-any.whl → 0.11.2__py3-none-any.whl

gpjax/__init__.py

@@ -32,14 +32,15 @@ from gpjax.citation import cite
 from gpjax.dataset import Dataset
 from gpjax.fit import (
     fit,
+    fit_lbfgs,
     fit_scipy,
 )
 
 __license__ = "MIT"
-__description__ = "Didactic Gaussian processes in JAX"
+__description__ = "Gaussian processes in JAX and Flax"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.11.0"
+__version__ = "0.11.2"
 
 __all__ = [
     "base",
@@ -56,5 +57,6 @@ __all__ = [
     "fit",
     "Module",
     "param_field",
+    "fit_lbfgs",
     "fit_scipy",
 ]

gpjax/citation.py

@@ -8,7 +8,12 @@ from beartype.typing import (
     Dict,
     Union,
 )
-from jaxlib.xla_extension import PjitFunction
+
+try:
+    # safely removable once jax>=0.6.0
+    from jaxlib.xla_extension import PjitFunction
+except ModuleNotFoundError:
+    from jaxlib._jax import PjitFunction
 
 from gpjax.kernels import (
     RFF,
@@ -45,7 +50,7 @@ class AbstractCitation:
 
 
 class NullCitation(AbstractCitation):
-    def __str__(self) -> str:
+    def as_str(self) -> str:
         return (
             "No citation available. If you think this is an error, please open a pull"
             " request."

gpjax/fit.py

@@ -127,7 +127,6 @@ def fit(  # noqa: PLR0913
         _check_verbose(verbose)
 
     # Model state filtering
-
     graphdef, params, *static_state = nnx.split(model, Parameter, ...)
 
     # Parameters bijection to unconstrained space
@@ -253,6 +252,110 @@ def fit_scipy(  # noqa: PLR0913
     return model, history
 
 
+def fit_lbfgs(
+    *,
+    model: Model,
+    objective: Objective,
+    train_data: Dataset,
+    params_bijection: tp.Union[dict[Parameter, Transform], None] = DEFAULT_BIJECTION,
+    max_iters: int = 100,
+    safe: bool = True,
+    max_linesearch_steps: int = 32,
+    gtol: float = 1e-5,
+) -> tuple[Model, jax.Array]:
+    r"""Train a Module model with respect to a supplied Objective function.
+
+    Uses Optax's LBFGS implementation and a jax.lax.while loop.
+
+     Args:
+         model: the model Module to be optimised.
+         objective: The objective function that we are optimising with
+             respect to.
+         train_data (Dataset): The training data to be used for the optimisation.
+         max_iters (int): The maximum number of optimisation steps to run. Defaults
+             to 500.
+         safe (bool): Whether to check the types of the inputs.
+         max_linesearch_steps (int): The maximum number of linesearch steps to use
+            for finding the stepsize.
+        gtol (float): Terminate the optimisation if the L2 norm of the gradient is
+            below this threshold.
+
+     Returns:
+         A tuple comprising the optimised model and final loss.
+    """
+    if safe:
+        # Check inputs
+        _check_model(model)
+        _check_train_data(train_data)
+        _check_num_iters(max_iters)
+
+    # Model state filtering
+    graphdef, params, *static_state = nnx.split(model, Parameter, ...)
+
+    # Parameters bijection to unconstrained space
+    if params_bijection is not None:
+        params = transform(params, params_bijection, inverse=True)
+
+    # Loss definition
+    def loss(params: nnx.State) -> ScalarFloat:
+        params = transform(params, params_bijection)
+        model = nnx.merge(graphdef, params, *static_state)
+        return objective(model, train_data)
+
+    # Initialise optimiser
+    optim = ox.lbfgs(
+        linesearch=ox.scale_by_zoom_linesearch(
+            max_linesearch_steps=max_linesearch_steps,
+            initial_guess_strategy="one",
+        )
+    )
+    opt_state = optim.init(params)
+    loss_value_and_grad = ox.value_and_grad_from_state(loss)
+
+    # Optimisation step.
+    def step(carry):
+        params, opt_state = carry
+
+        # Using optax's value_and_grad_from_state is more efficient given LBFGS uses a linesearch
+        # See https://optax.readthedocs.io/en/latest/api/utilities.html#optax.value_and_grad_from_state
+        loss_val, loss_gradient = loss_value_and_grad(params, state=opt_state)
+        updates, opt_state = optim.update(
+            loss_gradient,
+            opt_state,
+            params,
+            value=loss_val,
+            grad=loss_gradient,
+            value_fn=loss,
+        )
+        params = ox.apply_updates(params, updates)
+
+        return params, opt_state
+
+    def continue_fn(carry):
+        _, opt_state = carry
+        n = ox.tree_utils.tree_get(opt_state, "count")
+        g = ox.tree_utils.tree_get(opt_state, "grad")
+        g_l2_norm = ox.tree_utils.tree_l2_norm(g)
+        return (n == 0) | ((n < max_iters) & (g_l2_norm >= gtol))
+
+    # Optimisation loop
+    params, opt_state = jax.lax.while_loop(
+        continue_fn,
+        step,
+        (params, opt_state),
+    )
+    final_loss = ox.tree_utils.tree_get(opt_state, "value")
+
+    # Parameters bijection to constrained space
+    if params_bijection is not None:
+        params = transform(params, params_bijection)
+
+    # Reconstruct model
+    model = nnx.merge(graphdef, params, *static_state)
+
+    return model, final_loss
+
+
 def get_batch(train_data: Dataset, batch_size: int, key: KeyArray) -> Dataset:
     """Batch the data into mini-batches. Sampling is done with replacement.
 

gpjax/kernels/nonstationary/arccosine.py

@@ -23,7 +23,10 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import (
+    NonNegativeReal,
+    PositiveReal,
+)
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -91,9 +94,9 @@ class ArcCosine(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarArray], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarArray], self.variance)
 
         if isinstance(bias_variance, nnx.Variable):
             self.bias_variance = bias_variance

gpjax/kernels/nonstationary/linear.py

@@ -23,7 +23,7 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import NonNegativeReal
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -64,9 +64,9 @@ class Linear(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarArray], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarArray], self.variance)
 
     def __call__(
         self,

gpjax/kernels/nonstationary/polynomial.py

@@ -23,7 +23,10 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import (
+    NonNegativeReal,
+    PositiveReal,
+)
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -76,9 +79,9 @@ class Polynomial(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarArray], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarArray], self.variance)
 
         self.name = f"Polynomial (degree {self.degree})"
 

gpjax/kernels/stationary/base.py

@@ -25,7 +25,10 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import (
+    NonNegativeReal,
+    PositiveReal,
+)
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -85,11 +88,11 @@ class StationaryKernel(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
 
             # static typing
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarFloat], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarFloat], self.variance)
 
     @property
     def spectral_density(self) -> npd.Normal | npd.StudentT:

gpjax/likelihoods.py

@@ -28,7 +28,7 @@ from gpjax.integrators import (
     GHQuadratureIntegrator,
 )
 from gpjax.parameters import (
-    PositiveReal,
+    NonNegativeReal,
     Static,
 )
 from gpjax.typing import (
@@ -134,7 +134,7 @@ class Gaussian(AbstractLikelihood):
         self,
         num_datapoints: int,
         obs_stddev: tp.Union[
-            ScalarFloat, Float[Array, "#N"], PositiveReal, Static
+            ScalarFloat, Float[Array, "#N"], NonNegativeReal, Static
         ] = 1.0,
         integrator: AbstractIntegrator = AnalyticalGaussianIntegrator(),
     ):
@@ -148,8 +148,8 @@ class Gaussian(AbstractLikelihood):
                 likelihoods. Must be an instance of `AbstractIntegrator`. For the Gaussian likelihood, this defaults to
                 the `AnalyticalGaussianIntegrator`, as the expected log likelihood can be computed analytically.
         """
-        if not isinstance(obs_stddev, (PositiveReal, Static)):
-            obs_stddev = PositiveReal(jnp.asarray(obs_stddev))
+        if not isinstance(obs_stddev, (NonNegativeReal, Static)):
+            obs_stddev = NonNegativeReal(jnp.asarray(obs_stddev))
         self.obs_stddev = obs_stddev
 
         super().__init__(num_datapoints, integrator)

gpjax/mean_functions.py

@@ -207,5 +207,5 @@ SumMeanFunction = ft.partial(
     CombinationMeanFunction, operator=ft.partial(jnp.sum, axis=0)
 )
 ProductMeanFunction = ft.partial(
-    CombinationMeanFunction, operator=ft.partial(jnp.sum, axis=0)
+    CombinationMeanFunction, operator=ft.partial(jnp.prod, axis=0)
 )

gpjax/parameters.py

@@ -82,6 +82,14 @@ class Parameter(nnx.Variable[T]):
         self._tag = tag
 
 
+class NonNegativeReal(Parameter[T]):
+    """Parameter that is non-negative."""
+
+    def __init__(self, value: T, tag: ParameterTag = "non_negative", **kwargs):
+        super().__init__(value=value, tag=tag, **kwargs)
+        _safe_assert(_check_is_non_negative, self.value)
+
+
 class PositiveReal(Parameter[T]):
     """Parameter that is strictly positive."""
 
@@ -143,6 +151,7 @@ class LowerTriangular(Parameter[T]):
 
 DEFAULT_BIJECTION = {
     "positive": npt.SoftplusTransform(),
+    "non_negative": npt.SoftplusTransform(),
     "real": npt.IdentityTransform(),
     "sigmoid": npt.SigmoidTransform(),
     "lower_triangular": FillTriangularTransform(),
@@ -164,6 +173,13 @@ def _check_is_arraylike(value: T) -> None:
         )
 
 
+@checkify.checkify
+def _check_is_non_negative(value):
+    checkify.check(
+        jnp.all(value >= 0), "value needs to be non-negative, got {value}", value=value
+    )
+
+
 @checkify.checkify
 def _check_is_positive(value):
     checkify.check(

{gpjax-0.11.0.dist-info → gpjax-0.11.2.dist-info}/METADATA

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

{gpjax-0.11.0.dist-info → gpjax-0.11.2.dist-info}/RECORD

@@ -1,16 +1,16 @@
-gpjax/__init__.py,sha256=wXJtQa_3W7wZEw_t1Dk0uHUzNQQDv8QzsVbnwXCMXcQ,1654
-gpjax/citation.py,sha256=f2Hzj5MLyCE7l0hHAzsEQoTORZH5hgV_eis4uoBiWvE,3811
+gpjax/__init__.py,sha256=ylCFMtXwcMS2zxm4pI3KsnRdnX6bdh26TSdTfUh9l9o,1686
+gpjax/citation.py,sha256=pwFS8h1J-LE5ieRS0zDyuwhmQHNxkFHYE7iSMlVNmQc,3928
 gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
 gpjax/distributions.py,sha256=8LWmfmRVHOX29Uy8PkKFi2UhcCiunuu-4TMI_5-krHc,9299
-gpjax/fit.py,sha256=STwpeqSuu2pgT6uZU7xd7koPZbAjPDzhcZ8nHfozR7Q,11538
+gpjax/fit.py,sha256=R4TIPvBNHYSg9vBVp6is_QYENldRLIU_FklGE85C-aA,15046
 gpjax/gps.py,sha256=97lYGrsmsufQxKEd8qz5wPNvui6FKXTF_Ps-sMFIjnY,31246
 gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
-gpjax/likelihoods.py,sha256=VcCibgihaskmvNJT4kuPa7ehgjlnR9LgMz_2KJJvHY0,9296
+gpjax/likelihoods.py,sha256=99oTZoWld1M7vxgGM0pNY5Hnt2Ajd2lQNqawzrLmwtk,9308
 gpjax/lower_cholesky.py,sha256=3pnHaBrlGckFsrfYJ9Lsbd0pGmO7NIXdyY4aGm48MpY,1952
-gpjax/mean_functions.py,sha256=gIPz7exEhish3yeJQxZp5Q_jlf2-gCE-KVAnL2Rumkc,6489
+gpjax/mean_functions.py,sha256=-sVYO1_LWE8f34rllUOuaT5sgGGAdxo99v5kRo2d4oM,6490
 gpjax/numpyro_extras.py,sha256=-vWJ7SpZVNhSdCjjrlxIkovMFrM1IzpsMJK3B4LioGE,3411
 gpjax/objectives.py,sha256=I_ZqnwTNYIAUAZ9KQNenIl0ish1jDOXb7KaNmjz3Su4,15340
-gpjax/parameters.py,sha256=Vj1xzrziSLxfBSqyc-BacyKBwkbE9Sjq4b1HV5HZiOg,6507
+gpjax/parameters.py,sha256=H-DiXmotdBZCbf-GOjRaJoS_isk3GgFrpKHTq5GpnoA,6998
 gpjax/scan.py,sha256=jStQvwkE9MGttB89frxam1kaeXdWih7cVxkGywyaeHQ,5365
 gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
 gpjax/variational_families.py,sha256=Y9J1H91tXPm_hMy3ri_PgjAxqc_3r-BqKV83HRvB_m4,28295
@@ -29,11 +29,11 @@ gpjax/kernels/non_euclidean/__init__.py,sha256=RT7puRPqCTpyxZ16q596EuOQEQi1LK1v3
 gpjax/kernels/non_euclidean/graph.py,sha256=K4WIdX-dx1SsWuNHZnNjHFw8ElKZxGcReUiA3w4aCOI,4204
 gpjax/kernels/non_euclidean/utils.py,sha256=z42aw8ga0zuREzHawemR9okttgrAUPmq-aN5HMt4SuY,1578
 gpjax/kernels/nonstationary/__init__.py,sha256=YpWQfOy_cqOKc5ezn37vqoK3Z6jznYiJz28BD_8F7AY,930
-gpjax/kernels/nonstationary/arccosine.py,sha256=UCTVJEhTZFQjARGFsYMImLnTDyTyxobIL5f2LiAHkPI,5822
-gpjax/kernels/nonstationary/linear.py,sha256=UKDHFCQzKWDMYo76qcb5-ujjnP2_iL-1tcN017xjK48,2562
-gpjax/kernels/nonstationary/polynomial.py,sha256=7SDMfEcBCqnRn9xyj4iGcYLNvYJZiveN3uLZ_h12p10,3257
+gpjax/kernels/nonstationary/arccosine.py,sha256=2WV6aM0Z3-xXZnoPw-77n2CW62n-AZuJy-7AQ9xrMco,5858
+gpjax/kernels/nonstationary/linear.py,sha256=UIMoCq2hg6dQKr4J5UGiiPqotBleQuYfy00Ia1NaMOo,2571
+gpjax/kernels/nonstationary/polynomial.py,sha256=arP8DK0jnBOaayDWcFvHF0pdu9FVhwzXdqjnHUAL2VI,3293
 gpjax/kernels/stationary/__init__.py,sha256=j4BMTaQlIx2kNAT1Dkf4iO2rm-f7_oSVWNrk1bN0tqE,1406
-gpjax/kernels/stationary/base.py,sha256=FlsXMsXyZ5cI80jbsIo8Jv-H6gsV3C7v6plIhyCl-GI,7042
+gpjax/kernels/stationary/base.py,sha256=25qDqpZP4gNtzbyzDCW-6u7rJfMqkg0dW88XUmTTupU,7078
 gpjax/kernels/stationary/matern12.py,sha256=DGjqw6VveYsyy0TrufyJJvCei7p9slnm2f0TgRGG7_U,1773
 gpjax/kernels/stationary/matern32.py,sha256=laLsJWJozJzpYHBzlkPUq0rWxz1eWEwGC36P2nPJuaQ,1966
 gpjax/kernels/stationary/matern52.py,sha256=VSByD2sb7k-DzRFjaz31P3Rtc4bPPhHvMshrxZNFnns,2019
@@ -43,7 +43,7 @@ gpjax/kernels/stationary/rational_quadratic.py,sha256=dYONp3i4rnKj3ET8UyxAKXv6UO
 gpjax/kernels/stationary/rbf.py,sha256=euHUs6FdfRICQcabAWE4MX-7GEDr2TxgZWdFQiXr9Bw,1690
 gpjax/kernels/stationary/utils.py,sha256=6BI9EBcCzeeKx-XH-MfW1ORmtU__tPX5zyvfLhpkBsU,2180
 gpjax/kernels/stationary/white.py,sha256=TkdXXZCCjDs7JwR_gj5uvn2s1wyfRbe1vyHhUMJ8jjI,2212
-gpjax-0.11.0.dist-info/METADATA,sha256=eSWVc5y9WNrUmKpaOVq1CcHjrKjMwlmSvwovN9h9aCk,8558
-gpjax-0.11.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-gpjax-0.11.0.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
-gpjax-0.11.0.dist-info/RECORD,,
+gpjax-0.11.2.dist-info/METADATA,sha256=lTQVlrUbkxI7fU9Gdnac_eoNRyjCHEoEuEvvWbKmDqM,8558
+gpjax-0.11.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+gpjax-0.11.2.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
+gpjax-0.11.2.dist-info/RECORD,,