gpjax 0.9.3__py3-none-any.whl → 0.9.5__py3-none-any.whl

gpjax/__init__.py

@@ -19,7 +19,6 @@ from beartype.roar import BeartypeDecorHintPep585DeprecationWarning
 filterwarnings("ignore", category=BeartypeDecorHintPep585DeprecationWarning)
 
 from gpjax import (
-    decision_making,
     gps,
     integrators,
     kernels,
@@ -40,11 +39,10 @@ __license__ = "MIT"
 __description__ = "Didactic Gaussian processes in JAX"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.9.3"
+__version__ = "0.9.5"
 
 __all__ = [
     "base",
-    "decision_making",
     "gps",
     "integrators",
     "kernels",

gpjax/citation.py

@@ -10,11 +10,6 @@ from beartype.typing import (
 )
 from jaxlib.xla_extension import PjitFunction
 
-from gpjax.decision_making.test_functions import (
-    Forrester,
-    LogarithmicGoldsteinPrice,
-)
-from gpjax.decision_making.utility_functions import ThompsonSampling
 from gpjax.kernels import (
     RFF,
     ArcCosine,
@@ -149,41 +144,3 @@ def _(tree) -> PaperCitation:
         booktitle="Advances in neural information processing systems",
         citation_type="article",
     )
-
-
-####################
-# Decision making citations
-####################
-@cite.register(ThompsonSampling)
-def _(tree) -> PaperCitation:
-    return PaperCitation(
-        citation_key="wilson2020efficiently",
-        title="Efficiently sampling functions from Gaussian process posteriors",
-        authors="Wilson, James and Borovitskiy, Viacheslav and Terenin, Alexander and Mostowsky, Peter and Deisenroth, Marc",
-        year="2020",
-        booktitle="International Conference on Machine Learning",
-        citation_type="article",
-    )
-
-
-@cite.register(Forrester)
-def _(tree) -> BookCitation:
-    return BookCitation(
-        citation_key="forrester2008engineering",
-        authors="Forrester, Alexander and Sobester, Andras and Keane, Andy",
-        title="Engineering design via surrogate modelling: a practical guide",
-        year="2008",
-        publisher="John Wiley & Sons",
-    )
-
-
-@cite.register(LogarithmicGoldsteinPrice)
-def _(tree) -> PaperCitation:
-    return PaperCitation(
-        citation_key="picheny2013benchmark",
-        authors="Picheny, Victor and Wagner, Tobias and Ginsbourger, David",
-        title="A benchmark of kriging-based infill criteria for noisy optimization",
-        year="2013",
-        booktitle="Structural and multidisciplinary optimization",
-        citation_type="article",
-    )

gpjax/distributions.py

@@ -162,7 +162,9 @@ class GaussianDistribution(tfd.Distribution):
 
         return vmap(affine_transformation)(Z)
 
-    def sample(self, seed: KeyArray, sample_shape: Tuple[int, ...]):  # pylint: disable=useless-super-delegation
+    def sample(
+        self, seed: KeyArray, sample_shape: Tuple[int, ...]
+    ):  # pylint: disable=useless-super-delegation
         r"""See `Distribution.sample`."""
         return self._sample_n(
             seed, sample_shape[0]

gpjax/gps.py

@@ -652,7 +652,8 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
         """
         super().__init__(prior=prior, likelihood=likelihood, jitter=jitter)
 
-        latent = latent or jr.normal(key, shape=(self.likelihood.num_datapoints, 1))
+        if latent is None:
+            latent = jr.normal(key, shape=(self.likelihood.num_datapoints, 1))
 
         # TODO: static or intermediate?
         self.latent = latent if isinstance(latent, Parameter) else Real(latent)

gpjax/variational_families.py

@@ -149,12 +149,14 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
     ):
         super().__init__(posterior, inducing_inputs, jitter)
 
-        self.variational_mean = Real(
-            variational_mean or jnp.zeros((self.num_inducing, 1))
-        )
-        self.variational_root_covariance = LowerTriangular(
-            variational_root_covariance or jnp.eye(self.num_inducing)
-        )
+        if variational_mean is None:
+            variational_mean = jnp.zeros((self.num_inducing, 1))
+
+        if variational_root_covariance is None:
+            variational_root_covariance = jnp.eye(self.num_inducing)
+
+        self.variational_mean = Real(variational_mean)
+        self.variational_root_covariance = LowerTriangular(variational_root_covariance)
 
     def prior_kl(self) -> ScalarFloat:
         r"""Compute the prior KL divergence.
@@ -378,12 +380,14 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
     ):
         super().__init__(posterior, inducing_inputs, jitter)
 
-        self.natural_vector = Static(
-            natural_vector or jnp.zeros((self.num_inducing, 1))
-        )
-        self.natural_matrix = Static(
-            natural_matrix or -0.5 * jnp.eye(self.num_inducing)
-        )
+        if natural_vector is None:
+            natural_vector = jnp.zeros((self.num_inducing, 1))
+
+        if natural_matrix is None:
+            natural_matrix = -0.5 * jnp.eye(self.num_inducing)
+
+        self.natural_vector = Static(natural_vector)
+        self.natural_matrix = Static(natural_matrix)
 
     def prior_kl(self) -> ScalarFloat:
         r"""Compute the KL-divergence between our current variational approximation
@@ -540,13 +544,14 @@ class ExpectationVariationalGaussian(AbstractVariationalGaussian[L]):
     ):
         super().__init__(posterior, inducing_inputs, jitter)
 
-        # must come after super().__init__
-        self.expectation_vector = Static(
-            expectation_vector or jnp.zeros((self.num_inducing, 1))
-        )
-        self.expectation_matrix = Static(
-            expectation_matrix or jnp.eye(self.num_inducing)
-        )
+        if expectation_vector is None:
+            expectation_vector = jnp.zeros((self.num_inducing, 1))
+
+        if expectation_matrix is None:
+            expectation_matrix = jnp.eye(self.num_inducing)
+
+        self.expectation_vector = Static(expectation_vector)
+        self.expectation_matrix = Static(expectation_matrix)
 
     def prior_kl(self) -> ScalarFloat:
         r"""Evaluate the prior KL-divergence.

{gpjax-0.9.3.dist-info → gpjax-0.9.5.dist-info}/METADATA

@@ -1,13 +1,13 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: gpjax
-Version: 0.9.3
+Version: 0.9.5
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
 Project-URL: Source, https://github.com/JaxGaussianProcesses/GPJax
 Author-email: Thomas Pinder <tompinder@live.co.uk>
-License-Expression: Apache-2.0
-License-File: LICENSE
+License: MIT
+License-File: LICENSE.txt
 Keywords: gaussian-processes jax machine-learning bayesian
 Classifier: Development Status :: 4 - Beta
 Classifier: Programming Language :: Python
@@ -19,10 +19,9 @@ Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: <3.13,>=3.10
 Requires-Dist: beartype>0.16.1
 Requires-Dist: cola-ml==0.0.5
-Requires-Dist: flax>=0.8.5
+Requires-Dist: flax<0.10.0
 Requires-Dist: jax<0.4.28
 Requires-Dist: jaxlib<0.4.28
-Requires-Dist: jaxopt==0.8.2
 Requires-Dist: jaxtyping>0.2.10
 Requires-Dist: numpy<2.0.0
 Requires-Dist: optax>0.2.1
@@ -103,23 +102,23 @@ helped to shape GPJax into the package it is today.
 
 ## Notebook examples
 
-> - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/examples/regression/)
-> - [**Classification**](https://docs.jaxgaussianprocesses.com/examples/classification/)
-> - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/collapsed_vi/)
-> - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/uncollapsed_vi/)
-> - [**Laplace Approximation**](https://docs.jaxgaussianprocesses.com/examples/classification/#laplace-approximation)
-> - [**Inference on Non-Euclidean Spaces**](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
-> - [**Inference on Graphs**](https://docs.jaxgaussianprocesses.com/examples/graph_kernels/)
-> - [**Pathwise Sampling**](https://docs.jaxgaussianprocesses.com/examples/spatial/)
-> - [**Learning Gaussian Process Barycentres**](https://docs.jaxgaussianprocesses.com/examples/barycentres/)
-> - [**Deep Kernel Regression**](https://docs.jaxgaussianprocesses.com/examples/deep_kernels/)
-> - [**Poisson Regression**](https://docs.jaxgaussianprocesses.com/examples/poisson/)
-> - [**Bayesian Optimisation**](https://docs.jaxgaussianprocesses.com/examples/bayesian_optimisation/)
+> - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/_examples/regression/)
+> - [**Classification**](https://docs.jaxgaussianprocesses.com/_examples/classification/)
+> - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/_examples/collapsed_vi/)
+> - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/_examples/uncollapsed_vi/)
+> - [**Laplace Approximation**](https://docs.jaxgaussianprocesses.com/_examples/classification/#laplace-approximation)
+> - [**Inference on Non-Euclidean Spaces**](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/#custom-kernel)
+> - [**Inference on Graphs**](https://docs.jaxgaussianprocesses.com/_examples/graph_kernels/)
+> - [**Pathwise Sampling**](https://docs.jaxgaussianprocesses.com/_examples/spatial/)
+> - [**Learning Gaussian Process Barycentres**](https://docs.jaxgaussianprocesses.com/_examples/barycentres/)
+> - [**Deep Kernel Regression**](https://docs.jaxgaussianprocesses.com/_examples/deep_kernels/)
+> - [**Poisson Regression**](https://docs.jaxgaussianprocesses.com/_examples/poisson/)
+> - [**Bayesian Optimisation**](https://docs.jaxgaussianprocesses.com/_examples/bayesian_optimisation/)
 
 ## Guides for customisation
 >
-> - [**Custom kernels**](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
-> - [**UCI regression**](https://docs.jaxgaussianprocesses.com/examples/yacht/)
+> - [**Custom kernels**](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/#custom-kernel)
+> - [**UCI regression**](https://docs.jaxgaussianprocesses.com/_examples/yacht/)
 
 ## Conversion between `.ipynb` and `.py`
 Above examples are stored in [examples](docs/examples) directory in the double
@@ -180,7 +179,7 @@ optimiser = ox.adam(learning_rate=1e-2)
 # Obtain Type 2 MLEs of the hyperparameters
 opt_posterior, history = gpx.fit(
     model=posterior,
-    objective=gpx.objectives.conjugate_mll,
+    objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
     train_data=D,
     optim=optimiser,
     num_iters=500,

{gpjax-0.9.3.dist-info → gpjax-0.9.5.dist-info}/RECORD

@@ -1,9 +1,9 @@
-gpjax/__init__.py,sha256=UNfvnpRhJEZfz9qYjzWYNSSh3xg1FqNFst_A6xl_nfE,1697
-gpjax/citation.py,sha256=R4Pmvjt0ndA0avEDSvIbxDxKapkRRYXWX7RRWBvZCRQ,5306
+gpjax/__init__.py,sha256=T-2EbsNxg5VcdTeSH_G-mWwNcMTJVqbdI55gl9HMvG8,1653
+gpjax/citation.py,sha256=f2Hzj5MLyCE7l0hHAzsEQoTORZH5hgV_eis4uoBiWvE,3811
 gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
-gpjax/distributions.py,sha256=zxkSEZIlTg0PHvvgj0BQuIFEg-ugx6_NkEwSsbqWUM0,9325
+gpjax/distributions.py,sha256=X48FJr3reop9maherdMVt7-XZOm2f26T8AJt_IKM_oE,9339
 gpjax/fit.py,sha256=OHv8jUHxa1ndpqMERSDRtYtUDzubk9rMPVIhfCiIH5Q,11551
-gpjax/gps.py,sha256=NO18geRfcjo4mA3PGkuGont_Mj_yRqfvWzJqYmoKwiY,31225
+gpjax/gps.py,sha256=97lYGrsmsufQxKEd8qz5wPNvui6FKXTF_Ps-sMFIjnY,31246
 gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
 gpjax/likelihoods.py,sha256=DOyV1L0ompkpeImMTiOOiWLJfqSqvDX_acOumuFqPEc,9234
 gpjax/lower_cholesky.py,sha256=3pnHaBrlGckFsrfYJ9Lsbd0pGmO7NIXdyY4aGm48MpY,1952
@@ -12,21 +12,7 @@ gpjax/objectives.py,sha256=XwkPyL_iovTNKpKGVNt0Lt2_OMTJitSPhuyCtUrJpbc,15383
 gpjax/parameters.py,sha256=Z4Wy3gEzPZG23-dtqC437_ZWnd_sPe9LcLCKn21ZBvA,4886
 gpjax/scan.py,sha256=jStQvwkE9MGttB89frxam1kaeXdWih7cVxkGywyaeHQ,5365
 gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
-gpjax/variational_families.py,sha256=JO78dywHNH9__hjJkrP2ASb1L3C9aEBOW7fd0run-e4,27918
-gpjax/decision_making/__init__.py,sha256=SDuPQl80lJ7nhfRsiB_7c22wCMiQO5ehSNohxUGnB7w,2170
-gpjax/decision_making/decision_maker.py,sha256=S4pOXrWcEHy0NDA0gfWzhk7pG0NJfaPpMXvq03yTy0g,13915
-gpjax/decision_making/posterior_handler.py,sha256=UgXf1Gu7GMh2YDSmiSWJIzmWlFW06KTS44HYz3mazZQ,5905
-gpjax/decision_making/search_space.py,sha256=bXwtMOhHZ2klnABpXm5Raxe7b0NTRDjo_cN3ecbk53Y,3545
-gpjax/decision_making/utility_maximizer.py,sha256=VT2amwSJbB64IL_MiWNl9ZgjcqO757qK6NW2gUBKsqs,5965
-gpjax/decision_making/utils.py,sha256=5j1GO5kcmG2laZR39NjhqgEjRekAWWzrnREv_5Zct_Y,2367
-gpjax/decision_making/test_functions/__init__.py,sha256=GDCY9_kaAnxDWwzo1FkdxnDx-80MErAHchbGybT9xYs,1109
-gpjax/decision_making/test_functions/continuous_functions.py,sha256=oL5ZQkvmbC3u9rEvSYI2DRAN3r7Ynf7wRZQlUWjKjt0,5612
-gpjax/decision_making/test_functions/non_conjugate_functions.py,sha256=cfo3xQOzB5ajMjjl0YFfNlJClkAcY7ZbT23UyBYEofQ,2955
-gpjax/decision_making/utility_functions/__init__.py,sha256=xXI-4JKWAfTJ7XZ1vRDpqtb91MNzSPD0lP6xo0tOc7o,1445
-gpjax/decision_making/utility_functions/base.py,sha256=FOqrsRDmtHiCVl6IHr12-AEYBLStzMT5EBs-F92e1Og,3882
-gpjax/decision_making/utility_functions/expected_improvement.py,sha256=H6hjC-lj1oiHf2BomeQqroORQ7vtcOngiDAWxRwkNbg,4481
-gpjax/decision_making/utility_functions/probability_of_improvement.py,sha256=O_rHH1yR34JJlpAueSDJ_yo95fPI2aAGkwphS8snBYk,5220
-gpjax/decision_making/utility_functions/thompson_sampling.py,sha256=S-Yyn-9jsKkaXTvKFBP4sG_eCCKApGbHao5RR5tqXAo,4353
+gpjax/variational_families.py,sha256=s1rk7PtNTjQPabmVu-jBsuJBoqsxAAXwKFZJOEswkNQ,28161
 gpjax/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
 gpjax/kernels/base.py,sha256=abkj3zidsBs7YSkYEfjeJ5jTs1YyDCPoBM2ZzqaqrgI,11561
 gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
@@ -56,7 +42,7 @@ gpjax/kernels/stationary/rational_quadratic.py,sha256=dYONp3i4rnKj3ET8UyxAKXv6UO
 gpjax/kernels/stationary/rbf.py,sha256=G13gg5phO7ite7D9QgoCy7gB2_y0FM6GZhgFW4RL6Xw,1734
 gpjax/kernels/stationary/utils.py,sha256=Xa9EEnxgFqEi08ZSFAZYYHhJ85_3Ac-ZUyUk18B63M4,2225
 gpjax/kernels/stationary/white.py,sha256=TkdXXZCCjDs7JwR_gj5uvn2s1wyfRbe1vyHhUMJ8jjI,2212
-gpjax-0.9.3.dist-info/METADATA,sha256=A-tRR4wxz_YizCra4tZdOfRD1aUGJ_5sKgzL7Ax81B0,9976
-gpjax-0.9.3.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-gpjax-0.9.3.dist-info/licenses/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
-gpjax-0.9.3.dist-info/RECORD,,
+gpjax-0.9.5.dist-info/METADATA,sha256=T-OvGAyBe1N_QW6F9RbV-sx8wBJSAYQjpildBdhotS0,9967
+gpjax-0.9.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+gpjax-0.9.5.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
+gpjax-0.9.5.dist-info/RECORD,,

{gpjax-0.9.3.dist-info → gpjax-0.9.5.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.25.0
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any

gpjax-0.9.5.dist-info/licenses/LICENSE.txt

@@ -0,0 +1,19 @@
+(C) Copyright 2019 Hewlett Packard Enterprise Development LP
+
+Permission is hereby granted, free of charge, to any person obtaining a
+copy of this software and associated documentation files (the "Software"),
+to deal in the Software without restriction, including without limitation
+the rights to use, copy, modify, merge, publish, distribute, sublicense,
+and/or sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included
+in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
+OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+OTHER DEALINGS IN THE SOFTWARE.

gpjax/decision_making/__init__.py

@@ -1,63 +0,0 @@
-# Copyright 2023 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.
-# ==============================================================================
-from gpjax.decision_making.decision_maker import (
-    AbstractDecisionMaker,
-    UtilityDrivenDecisionMaker,
-)
-from gpjax.decision_making.posterior_handler import PosteriorHandler
-from gpjax.decision_making.search_space import (
-    AbstractSearchSpace,
-    ContinuousSearchSpace,
-)
-from gpjax.decision_making.test_functions import (
-    AbstractContinuousTestFunction,
-    Forrester,
-    LogarithmicGoldsteinPrice,
-    Quadratic,
-)
-from gpjax.decision_making.utility_functions import (
-    AbstractSinglePointUtilityFunctionBuilder,
-    AbstractUtilityFunctionBuilder,
-    SinglePointUtilityFunction,
-    ThompsonSampling,
-    UtilityFunction,
-)
-from gpjax.decision_making.utility_maximizer import (
-    AbstractSinglePointUtilityMaximizer,
-    AbstractUtilityMaximizer,
-    ContinuousSinglePointUtilityMaximizer,
-)
-from gpjax.decision_making.utils import build_function_evaluator
-
-__all__ = [
-    "AbstractUtilityFunctionBuilder",
-    "AbstractUtilityMaximizer",
-    "AbstractDecisionMaker",
-    "AbstractSearchSpace",
-    "AbstractSinglePointUtilityFunctionBuilder",
-    "AbstractSinglePointUtilityMaximizer",
-    "UtilityFunction",
-    "build_function_evaluator",
-    "ContinuousSinglePointUtilityMaximizer",
-    "ContinuousSearchSpace",
-    "UtilityDrivenDecisionMaker",
-    "AbstractContinuousTestFunction",
-    "Forrester",
-    "LogarithmicGoldsteinPrice",
-    "PosteriorHandler",
-    "Quadratic",
-    "SinglePointUtilityFunction",
-    "ThompsonSampling",
-]