Merge branch 'main' into mc/issue-59

README.md

@@ -24,9 +24,9 @@ gp_model = gpax.ExactGP(1, kernel='RBF')
 # Run Hamiltonian Monte Carlo to obtain posterior samples for the GP model parameters
 gp_model.fit(rng_key, X, y)  # X and y are numpy arrays with dimensions (n, d) and (n,)
 ```
-In the fully Bayesian mode, we get a pair of predictive mean and covariance for each Hamiltonian Monte Carlo sample containing the GP parameters (in this case, the RBF kernel hyperparameters and model noise). Hence, a prediction on new inputs with a trained GP model returns the center of the mass of all the predictive means (```y_pred```) and samples from multivariate normal distributions for all the pairs of predictive means and covariances (```y_sampled```).
+In the fully Bayesian mode, we get a pair of predictive mean and covariance for each Hamiltonian Monte Carlo sample containing the GP parameters (in this case, the RBF kernel hyperparameters and model noise). Hence, a prediction on new inputs with a trained GP model returns the center of the mass of all the predictive means (```posterior_mean```) and samples from multivariate normal distributions for all the pairs of predictive means and covariances (```f_samples```).
 ```python3
-y_pred, y_sampled = gp_model.predict(rng_key_predict, X_test)
+posterior_mean, f_samples = gp_model.predict(rng_key_predict, X_test)
 ```
 
 <img src = "https://user-images.githubusercontent.com/34245227/167945293-8cb5b88a-1f64-4f7d-95ab-26863b90d1e5.jpg" height="60%" width="60%">
@@ -96,7 +96,7 @@ sgp_model = gpax.ExactGP(1, kernel='Matern', mean_fn=piecewise, mean_fn_prior=pi
 # Run MCMC to obtain posterior samples
 sgp_model.fit(rng_key, X, y)
 # Get GP prediction on new/test data
-y_pred, y_sampled = sgp_model.predict(rng_key_predict, X_test)
+posterior_mean, f_samples = sgp_model.predict(rng_key_predict, X_test)
 ```
 
 ![GP_vs_sGP2](https://github.com/ziatdinovmax/gpax/assets/34245227/89de341c-f00c-468c-afe6-c0b1c1140725)
@@ -150,7 +150,7 @@ Note that X has (N, D+1) dimensions where the last column contains task/fidelity
 X_unmeasured2 = np.column_stack((X_full_range, np.ones_like(X_full_range)))
 
 # Make a prediction with the trained model
-y_mean2, y_sampled2 = model.predict(key2, X_unmeasured2, noiseless=True)
+posterior_mean2, f_samples2  = model.predict(key2, X_unmeasured2, noiseless=True)
 ```
 
 ![GP_vs_MTGP](https://github.com/ziatdinovmax/gpax/assets/34245227/5a36d3cd-c904-4345-abc3-b1bea5025cc8)

docs/source/index.rst

@@ -19,9 +19,10 @@ GPax is a small Python package for physics-based Gaussian processes (GPs) built
    :caption: Package Content
 
    models
-   hypo
    acquisition
    kernels
+   priors
+   hypo
    utils
 
 .. toctree::

docs/source/models.rst

@@ -96,3 +96,13 @@ Multi-Task Learning
     :undoc-members:
     :member-order: bysource
     :show-inheritance:
+
+Structured Probabilistic Models
+-------------------------------
+.. autoclass:: gpax.models.spm.sPM
+    :members:
+    :inherited-members:
+    :undoc-members:
+    :member-order: bysource
+    :show-inheritance:
+

docs/source/priors.rst

@@ -0,0 +1,23 @@
+Priors
+======
+
+.. autofunction:: gpax.priors.normal_dist
+
+.. autofunction:: gpax.priors.lognormal_dist
+
+.. autofunction:: gpax.priors.halfnormal_dist
+
+.. autofunction:: gpax.priors.gamma_dist
+
+.. autofunction:: gpax.priors.uniform_dist
+
+.. autofunction:: gpax.priors.auto_normal_priors
+
+.. autofunction:: gpax.priors.auto_lognormal_priors
+
+.. autofunction:: gpax.priors.auto_normal_kernel_priors
+
+.. autofunction:: gpax.priors.auto_lognormal_kernel_priors
+
+
+

docs/source/utils.rst

@@ -1,40 +1,17 @@
 Utilities
 =========
 
-Priors
-------
+Automatic function setters
+--------------------------
 
-.. autofunction:: gpax.utils.normal_dist
+.. autofunction:: gpax.utils.set_fn
 
-.. autofunction:: gpax.utils.lognormal_dist
-
-.. autofunction:: gpax.utils.halfnormal_dist
-
-.. autofunction:: gpax.utils.gamma_dist
-
-.. autofunction:: gpax.utils.uniform_dist
-
-.. autofunction:: gpax.utils.place_normal_prior
-
-.. autofunction:: gpax.utils.place_lognormal_prior
-
-.. autofunction:: gpax.utils.place_halfnormal_prior
-
-.. autofunction:: gpax.utils.place_uniform_prior
-
-.. autofunction:: gpax.utils.place_gamma_prior
+.. autofunction:: gpax.utils.set_kernel_fn
 
 
 Other utilities
 ---------------
 
-.. autoclass:: gpax.models.spm.sPM
-    :members:
-    :inherited-members:
-    :undoc-members:
-    :member-order: bysource
-    :show-inheritance:
-
 .. autofunction:: gpax.utils.dviz
 
 .. autofunction:: gpax.utils.get_keys

gpax/__init__.py

@@ -21,6 +21,7 @@
 )
 
 __all__ = [
+    "priors",
     "utils",
     "kernels",
     "mtkernels",

gpax/models/bnn.py

@@ -32,7 +32,7 @@ def _set_data(self, X: jnp.ndarray, y: Optional[jnp.ndarray] = None
                   ) -> Union[Tuple[jnp.ndarray], jnp.ndarray]:
         X = X if X.ndim > 1 else X[:, None]
         if y is not None:
-            y = y[:, None] if y.ndim < 1 else y
+            y = y[:, None] if y.ndim < 2 else y
             return X, y
         return X
 
@@ -47,7 +47,7 @@ def sample_weights(name: str, in_channels: int, out_channels: int) -> jnp.ndarra
 
 def sample_biases(name: str, channels: int) -> jnp.ndarray:
     """Sampling bias vector"""
-    b = numpyro.sample(name=name, fn=dist.Normal(
+    b = numpyro.sample(name=name, fn=dist.Cauchy(
         loc=jnp.zeros((channels)), scale=jnp.ones((channels))))
     return b
 

gpax/models/spm.py

@@ -16,6 +16,8 @@
 import jax
 import jaxlib
 import jax.numpy as jnp
+import jax.random as jra
+from jax import vmap
 import numpyro
 import numpyro.distributions as dist
 from numpyro.infer import MCMC, NUTS, Predictive, init_to_median
@@ -144,19 +146,44 @@ def sample_from_prior(self, rng_key: jnp.ndarray,
         prior_predictive = Predictive(self.model, num_samples=num_samples)
         samples = prior_predictive(rng_key, X)
         return samples['y']
+
+    def sample_single_posterior_predictive(self, rng_key, X_new, params, n_draws):
+        sigma = params["noise"]
+        loc = self._model(X_new, params)
+        sample = dist.Normal(loc, sigma).sample(rng_key, (n_draws,)).mean(0)
+        return loc, sample
+
+    def _vmap_predict(self, rng_key: jnp.ndarray, X_new: jnp.ndarray,
+                      samples: Optional[Dict[str, jnp.ndarray]] = None, 
+                      n_draws: int = 1, 
+                      ) -> Tuple[jnp.ndarray, jnp.ndarray]:
+        """
+        Helper method to vectorize predictions over posterior samples
+        """
+        if samples is None:
+            samples = self.get_samples(chain_dim=False)
+        num_samples = len(next(iter(samples.values())))
+        vmap_args = (jra.split(rng_key, num_samples), samples)
+
+        predictive = lambda p1, p2: self.sample_single_posterior_predictive(p1, X_new, p2, n_draws)
+        loc, f_samples = vmap(predictive)(*vmap_args)
+
+        return loc, f_samples
 
     def predict(self, rng_key: jnp.ndarray, X_new: jnp.ndarray,
                 samples: Optional[Dict[str, jnp.ndarray]] = None,
+                n: int = 1,
                 filter_nans: bool = False, take_point_predictions_mean: bool = True,
                 device: Type[jaxlib.xla_extension.Device] = None
                 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
         """
-        Make prediction at X_new points using sampled GP hyperparameters
+        Make prediction at X_new points using posterior model parameters
 
         Args:
             rng_key: random number generator key
             X_new: 2D vector with new/'test' data of :math:`n x num_features` dimensionality
             samples: optional posterior samples
+            n: number of samples to draw from normal distribution per single HMC sample
             filter_nans: filter out samples containing NaN values (if any)
             take_point_predictions_mean: take a mean of point predictions (without sampling from the normal distribution)
             device:
@@ -172,10 +199,7 @@ def predict(self, rng_key: jnp.ndarray, X_new: jnp.ndarray,
         if device:
             X_new = jax.device_put(X_new, device)
             samples = jax.device_put(samples, device)
-        predictive = Predictive(
-            self.model, posterior_samples=samples, parallel=True)
-        y_pred = predictive(rng_key, X_new)
-        y_pred, y_sampled = y_pred["mu"], y_pred["y"]
+        y_pred, y_sampled = self._vmap_predict(rng_key, X_new, samples, n)
         if filter_nans:
             y_sampled_ = [y_i for y_i in y_sampled if not jnp.isnan(y_i).any()]
             y_sampled = jnp.array(y_sampled_)

gpax/priors/__init__.py

@@ -0,0 +1 @@
+from .priors import *