Merge pull request #408 from JaxGaussianProcesses/namespace_cleanup

Namespace cleanup

Author: thomaspinder

README.md

@@ -135,10 +135,10 @@ D = gpx.Dataset(X=x, y=y)
 # Construct the prior
 meanf = gpx.mean_functions.Zero()
 kernel = gpx.kernels.RBF()
-prior = gpx.Prior(mean_function=meanf, kernel=kernel)
+prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
 
 # Define a likelihood
-likelihood = gpx.Gaussian(num_datapoints=n)
+likelihood = gpx.likelihoods.Gaussian(num_datapoints = n)
 
 # Construct the posterior
 posterior = prior * likelihood

benchmarks/objectives.py

@@ -22,7 +22,7 @@ def setup(self, n_datapoints: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Gaussian(num_datapoints=self.data.n)
         self.objective = gpx.ConjugateMLL()
         self.posterior = self.prior * self.likelihood
@@ -48,7 +48,7 @@ def setup(self, n_datapoints: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Bernoulli(num_datapoints=self.data.n)
         self.objective = gpx.LogPosteriorDensity()
         self.posterior = self.prior * self.likelihood
@@ -75,7 +75,7 @@ def setup(self, n_datapoints: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Poisson(num_datapoints=self.data.n)
         self.objective = gpx.LogPosteriorDensity()
         self.posterior = self.prior * self.likelihood

benchmarks/predictions.py

@@ -21,7 +21,7 @@ def setup(self, n_test: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Gaussian(num_datapoints=self.data.n)
         self.posterior = self.prior * self.likelihood
         key, subkey = jr.split(key)
@@ -46,7 +46,7 @@ def setup(self, n_test: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Bernoulli(num_datapoints=self.data.n)
         self.posterior = self.prior * self.likelihood
         key, subkey = jr.split(key)
@@ -71,7 +71,7 @@ def setup(self, n_test: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Bernoulli(num_datapoints=self.data.n)
         self.posterior = self.prior * self.likelihood
         key, subkey = jr.split(key)

benchmarks/sparse.py

@@ -19,7 +19,7 @@ def setup(self, n_datapoints: int, n_inducing: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(1)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Gaussian(num_datapoints=self.data.n)
         self.posterior = self.prior * self.likelihood
 

benchmarks/stochastic.py

@@ -20,7 +20,7 @@ def setup(self, n_datapoints: int, n_inducing: int, batch_size: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(1)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Gaussian(num_datapoints=self.data.n)
         self.posterior = self.prior * self.likelihood
 

docs/examples/README.md

@@ -67,7 +67,7 @@ class Prior(AbstractPrior):
         >>>
         >>> meanf = gpx.mean_functions.Zero()
         >>> kernel = gpx.kernels.RBF()
-        >>> prior = gpx.Prior(mean_function=meanf, kernel = kernel)
+        >>> prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
 
     Attributes:
         kernel (Kernel): The kernel function used to parameterise the prior.

docs/examples/barycentres.py

@@ -134,8 +134,13 @@ def fit_gp(x: jax.Array, y: jax.Array) -> tfd.MultivariateNormalFullCovariance:
         y = y.reshape(-1, 1)
     D = gpx.Dataset(X=x, y=y)
 
-    likelihood = gpx.Gaussian(num_datapoints=n)
-    posterior = gpx.Prior(mean_function=gpx.Constant(), kernel=gpx.RBF()) * likelihood
+    likelihood = gpx.likelihoods.Gaussian(num_datapoints=n)
+    posterior = (
+        gpx.gps.Prior(
+            mean_function=gpx.mean_functions.Constant(), kernel=gpx.kernels.RBF()
+        )
+        * likelihood
+    )
 
     opt_posterior, _ = gpx.fit_scipy(
         model=posterior,

docs/examples/bayesian_optimisation.py

@@ -201,9 +201,9 @@ def forrester(x: Float[Array, "N 1"]) -> Float[Array, "N 1"]:
 
 # %%
 def return_optimised_posterior(
-    data: gpx.Dataset, prior: gpx.Module, key: Array
-) -> gpx.Module:
-    likelihood = gpx.Gaussian(
+    data: gpx.Dataset, prior: gpx.base.Module, key: Array
+) -> gpx.base.Module:
+    likelihood = gpx.likelihoods.Gaussian(
         num_datapoints=data.n, obs_stddev=jnp.array(1e-3)
     )  # Our function is noise-free, so we set the observation noise's standard deviation to a very small value
     likelihood = likelihood.replace_trainable(obs_stddev=False)
@@ -230,7 +230,7 @@ def return_optimised_posterior(
 
 mean = gpx.mean_functions.Zero()
 kernel = gpx.kernels.Matern52()
-prior = gpx.Prior(mean_function=mean, kernel=kernel)
+prior = gpx.gps.Prior(mean_function=mean, kernel=kernel)
 opt_posterior = return_optimised_posterior(D, prior, key)
 
 # %% [markdown]
@@ -315,7 +315,7 @@ def optimise_sample(
 
 # %%
 def plot_bayes_opt(
-    posterior: gpx.Module,
+    posterior: gpx.base.Module,
     sample: FunctionalSample,
     dataset: gpx.Dataset,
     queried_x: ScalarFloat,
@@ -401,7 +401,7 @@ def plot_bayes_opt(
     # Generate optimised posterior using previously observed data
     mean = gpx.mean_functions.Zero()
     kernel = gpx.kernels.Matern52()
-    prior = gpx.Prior(mean_function=mean, kernel=kernel)
+    prior = gpx.gps.Prior(mean_function=mean, kernel=kernel)
     opt_posterior = return_optimised_posterior(D, prior, subkey)
 
     # Draw a sample from the posterior, and find the minimiser of it
@@ -543,7 +543,7 @@ def six_hump_camel(x: Float[Array, "N 2"]) -> Float[Array, "N 1"]:
         kernel = gpx.kernels.Matern52(
             active_dims=[0, 1], lengthscale=jnp.array([1.0, 1.0]), variance=2.0
         )
-        prior = gpx.Prior(mean_function=mean, kernel=kernel)
+        prior = gpx.gps.Prior(mean_function=mean, kernel=kernel)
         opt_posterior = return_optimised_posterior(D, prior, subkey)
 
         # Draw a sample from the posterior, and find the minimiser of it
@@ -561,7 +561,8 @@ def six_hump_camel(x: Float[Array, "N 2"]) -> Float[Array, "N 1"]:
         # Evaluate the black-box function at the best point observed so far, and add it to the dataset
         y_star = six_hump_camel(x_star)
         print(
-            f"BO Iteration: {i + 1}, Queried Point: {x_star}, Black-Box Function Value: {y_star}"
+            f"BO Iteration: {i + 1}, Queried Point: {x_star}, Black-Box Function Value:"
+            f" {y_star}"
         )
         D = D + gpx.Dataset(X=x_star, y=y_star)
     bo_experiment_results.append(D)

docs/examples/classification.py

@@ -89,10 +89,10 @@
 # choose a Bernoulli likelihood with a probit link function.
 
 # %%
-kernel = gpx.RBF()
-meanf = gpx.Constant()
-prior = gpx.Prior(mean_function=meanf, kernel=kernel)
-likelihood = gpx.Bernoulli(num_datapoints=D.n)
+kernel = gpx.kernels.RBF()
+meanf = gpx.mean_functions.Constant()
+prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
+likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n)
 
 # %% [markdown]
 # We construct the posterior through the product of our prior and likelihood.
@@ -116,7 +116,7 @@
 # Optax's optimisers.
 
 # %%
-negative_lpd = jax.jit(gpx.LogPosteriorDensity(negative=True))
+negative_lpd = jax.jit(gpx.objectives.LogPosteriorDensity(negative=True))
 
 optimiser = ox.adam(learning_rate=0.01)
 
@@ -345,7 +345,7 @@ def construct_laplace(test_inputs: Float[Array, "N D"]) -> tfd.MultivariateNorma
 num_adapt = 500
 num_samples = 500
 
-lpd = jax.jit(gpx.LogPosteriorDensity(negative=False))
+lpd = jax.jit(gpx.objectives.LogPosteriorDensity(negative=False))
 unconstrained_lpd = jax.jit(lambda tree: lpd(tree.constrain(), D))
 
 adapt = blackjax.window_adaptation(

docs/examples/collapsed_vi.py

@@ -106,10 +106,10 @@
 # this, it is intractable to evaluate.
 
 # %%
-meanf = gpx.Constant()
-kernel = gpx.RBF()
-likelihood = gpx.Gaussian(num_datapoints=D.n)
-prior = gpx.Prior(mean_function=meanf, kernel=kernel)
+meanf = gpx.mean_functions.Constant()
+kernel = gpx.kernels.RBF()
+likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n)
+prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
 posterior = prior * likelihood
 
 # %% [markdown]
@@ -119,15 +119,17 @@
 # inducing points into the constructor as arguments.
 
 # %%
-q = gpx.CollapsedVariationalGaussian(posterior=posterior, inducing_inputs=z)
+q = gpx.variational_families.CollapsedVariationalGaussian(
+    posterior=posterior, inducing_inputs=z
+)
 
 # %% [markdown]
 # We define our variational inference algorithm through `CollapsedVI`. This defines
 # the collapsed variational free energy bound considered in
 # <strong data-cite="titsias2009">Titsias (2009)</strong>.
 
 # %%
-elbo = gpx.CollapsedELBO(negative=True)
+elbo = gpx.objectives.CollapsedELBO(negative=True)
 
 # %% [markdown]
 # For researchers, GPJax has the capacity to print the bibtex citation for objects such
@@ -241,14 +243,14 @@
 # full model.
 
 # %%
-full_rank_model = gpx.Prior(mean_function=gpx.Zero(), kernel=gpx.RBF()) * gpx.Gaussian(
-    num_datapoints=D.n
-)
-negative_mll = jit(gpx.ConjugateMLL(negative=True).step)
+full_rank_model = gpx.gps.Prior(
+    mean_function=gpx.mean_functions.Zero(), kernel=gpx.kernels.RBF()
+) * gpx.likelihoods.Gaussian(num_datapoints=D.n)
+negative_mll = jit(gpx.objectives.ConjugateMLL(negative=True).step)
 # %timeit negative_mll(full_rank_model, D).block_until_ready()
 
 # %%
-negative_elbo = jit(gpx.CollapsedELBO(negative=True).step)
+negative_elbo = jit(gpx.objectives.CollapsedELBO(negative=True).step)
 # %timeit negative_elbo(q, D).block_until_ready()
 
 # %% [markdown]