Multi-fidelity/task DKL and GP + code refactoring

docs/source/acquisition.rst

@@ -9,4 +9,4 @@ Acquisition functions
 
 .. autofunction:: gpax.acquisition.UE
 
-.. autofunction:: gpax.acquisition.bUCB
+.. autofunction:: gpax.acquisition.qUCB

docs/source/kernels.rst

@@ -5,4 +5,6 @@ Kernels
 
 .. autofunction:: gpax.kernels.MaternKernel
 
-.. autofunction:: gpax.kernels.PeriodicKernel
+.. autofunction:: gpax.kernels.PeriodicKernel
+
+.. autofunction:: gpax.kernels.NNGPKernel

docs/source/models.rst

@@ -1,34 +1,68 @@
 GPax models
 ===========
 
-Gaussian Processes (Fully Bayesian Implementation)
+Gaussian Processes - Fully Bayesian Implementation
 --------------------------------------------------
-.. autoclass:: gpax.gp.ExactGP
+.. autoclass:: gpax.models.gp.ExactGP
     :members:
     :inherited-members:
     :undoc-members:
     :member-order: bysource
     :show-inheritance:
 
-.. autoclass:: gpax.vgp.vExactGP
+.. autoclass:: gpax.models.vgp.vExactGP
     :members:
     :inherited-members:
     :undoc-members:
     :member-order: bysource
     :show-inheritance:
 
-Deep Kernel Learning (Fully Bayesian Implementation)
+Gaussian Processes - Approximate Bayesian
+------------------------------------------
+.. autoclass:: gpax.models.gp.ExactGP
+    :members:
+    :inherited-members:
+    :undoc-members:
+    :member-order: bysource
+    :show-inheritance:
+
+Deep Kernel Learning - Fully Bayesian Implementation
 ----------------------------------------------------
-.. autoclass:: gpax.dkl.DKL
+.. autoclass:: gpax.models.dkl.DKL
     :members:
     :inherited-members:
     :undoc-members:
     :member-order: bysource
     :show-inheritance:
 
-Deep Kernel Learning (Approximate Bayesian)
+Deep Kernel Learning - Approximate Bayesian
 -------------------------------------------
-.. autoclass:: gpax.vidkl.viDKL
+.. autoclass:: gpax.models.vidkl.viDKL
+    :members:
+    :inherited-members:
+    :undoc-members:
+    :member-order: bysource
+    :show-inheritance:
+
+Infinite-width Bayesian Neural Networks
+----------------------------------------
+.. autoclass:: gpax.models.ibnn.iBNN
+    :members:
+    :inherited-members:
+    :undoc-members:
+    :member-order: bysource
+    :show-inheritance:
+
+Multi-Task Learning
+--------------------
+.. autoclass:: gpax.models.mtgp.MultiTaskGP
+    :members:
+    :inherited-members:
+    :undoc-members:
+    :member-order: bysource
+    :show-inheritance:
+
+.. autoclass:: gpax.models.vi_mtdkl.viMTDKL
     :members:
     :inherited-members:
     :undoc-members:

gpax/__init__.py

@@ -1,12 +1,11 @@
-from . import utils, kernels, acquisition
-from .gp import ExactGP
-from .vgp import vExactGP
-from .bnn import DKL, viDKL, iBNN, vi_iBNN
-from .vigp import viGP
-from .spm import sPM
-from .hypo import sample_next
-
 from .__version__ import version as __version__
+from . import utils
+from . import kernels
+from . import acquisition
+from .hypo import sample_next
+from .models import (DKL, CoregGP, ExactGP, MultiTaskGP, iBNN, vExactGP,
+                     vi_iBNN, viDKL, viGP, viMTDKL)
 
-__all__ = ["utils", "kernels", "acquisition", "ExactGP", "vExactGP", "DKL",
-           "viDKL", "iBNN", "vi_iBNN", "viGP", "sPM", "sample_next", "__version__"]
+__all__ = ["utils", "kernels", "mtkernels", "acquisition", "ExactGP", "vExactGP", "DKL",
+           "viDKL", "iBNN", "vi_iBNN", "MultiTaskGP", "viMTDKL", "viGP", "sPM",
+           "CoregGP", "sample_next", "__version__"]

gpax/acquisition/__init__.py

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

gpax/acquisition.py → gpax/acquisition/acquisition.py

@@ -15,7 +15,7 @@
 import numpy as onp
 import numpyro.distributions as dist
 
-from .gp import ExactGP
+from ..models.gp import ExactGP
 
 
 def EI(rng_key: jnp.ndarray, model: Type[ExactGP],

gpax/hypo.py

@@ -13,8 +13,8 @@
 import numpy as np
 import numpyro
 
-from .gp import ExactGP
-from .spm import sPM
+from .models.gp import ExactGP
+from .models.spm import sPM
 from .utils import get_keys
 
 

gpax/kernels/__init__.py

@@ -0,0 +1,16 @@
+from .kernels import (MaternKernel, NNGPKernel, PeriodicKernel, RBFKernel,
+                      get_kernel, nngp_erf, nngp_relu)
+from .mtkernels import (LCMKernel, MultitaskKernel, MultivariateKernel,
+                        index_kernel)
+
+__all__ = [
+    "RBFKernel",
+    "MaternKernel",
+    "PeriodicKernel",
+    "NNGPKernel",
+    "get_kernel",
+    "index_kernel",
+    "MultitaskKernel",
+    "MultivariateKernel",
+    "LCMKernel"
+]

gpax/kernels.py → gpax/kernels/kernels.py

@@ -28,7 +28,7 @@ def add_jitter(x, jitter=1e-6):
 def square_scaled_distance(X: jnp.ndarray, Z: jnp.ndarray,
                            lengthscale: Union[jnp.ndarray, float] = 1.
                            ) -> jnp.ndarray:
-    """
+    r"""
     Computes a square of scaled distance, :math:`\|\frac{X-Z}{l}\|^2`,
     between X and Z are vectors with :math:`n x num_features` dimensions
     """
@@ -115,7 +115,7 @@ def PeriodicKernel(X: jnp.ndarray, Z: jnp.ndarray,
 
 
 def nngp_erf(x1: jnp.ndarray, x2: jnp.ndarray,
-             var_b: jnp.array, var_w: jnp.array, 
+             var_b: jnp.array, var_w: jnp.array,
              depth: int = 3) -> jnp.array:
     """
     Computes the Neural Network Gaussian Process (NNGP) kernel value for
@@ -187,7 +187,7 @@ def NNGPKernel(activation: str = 'erf', depth: int = 3
 
     Args:
         activation: activation function ('erf' or 'relu')
-        depth: The number of layers in the corresponding infinite-width neural network. 
+        depth: The number of layers in the corresponding infinite-width neural network.
                Controls the level of recursion in the computation.
 
     Returns:
@@ -196,7 +196,7 @@ def NNGPKernel(activation: str = 'erf', depth: int = 3
     nngp_single_pair_ = nngp_relu if activation == 'relu' else nngp_erf
 
     def NNGPKernel_func(X: jnp.ndarray, Z: jnp.ndarray,
-                        params: Dict[str, jnp.ndarray], 
+                        params: Dict[str, jnp.ndarray],
                         noise: jnp.ndarray = 0, **kwargs
                         ) -> jnp.ndarray:
         """