Metadata-Version: 2.3
Name: gpytorch-mogp
Version: 0.1.2
Summary: A package which extends GPyTorch with correlated multi-output GPs
Project-URL: Documentation, https://github.com/dnv-opensource/gpytorch-mogp#readme
Project-URL: Issues, https://github.com/dnv-opensource/gpytorch-mogp/issues
Project-URL: Source, https://github.com/dnv-opensource/gpytorch-mogp
Author-email: Magnus Kristiansen <magnus.kristiansen@dnv.com>
License: MIT License
        
        Copyright (c) 2024 [DNV](https://www.dnv.com) [open source](https://github.com/dnv-opensource)
        
        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.
Classifier: Development Status :: 3 - Alpha
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.11
Requires-Dist: einops>=0.8.0
Requires-Dist: gpytorch~=1.11
Provides-Extra: all
Requires-Dist: jupyter>=1.0.0; extra == 'all'
Requires-Dist: matplotlib>=3.8.4; extra == 'all'
Requires-Dist: ruff>=0.4.2; extra == 'all'
Provides-Extra: dev
Requires-Dist: ruff>=0.4.2; extra == 'dev'
Provides-Extra: examples
Requires-Dist: jupyter>=1.0.0; extra == 'examples'
Requires-Dist: matplotlib>=3.8.4; extra == 'examples'
Description-Content-Type: text/markdown

# gpytorch-mogp

[![PyPI - Version](https://img.shields.io/pypi/v/gpytorch-mogp.svg)](https://pypi.org/project/gpytorch-mogp)
[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/gpytorch-mogp.svg)](https://pypi.org/project/gpytorch-mogp)
[![license](https://img.shields.io/pypi/l/gpytorch-mogp.svg)](https://github.com/dnv-opensource/gpytorch-mogp/blob/main/LICENSE)

`gpytorch-mogp` is a Python package that extends [GPyTorch](https://gpytorch.ai/) to support Multiple-Output Gaussian
processes where the correlation between the outputs is known.

> **Note:** The package is currently in an early stage of development. Expect bugs and breaking changes in future
> versions. Please create an issue if you encounter any problems or have any suggestions.

-----

**Table of Contents**

- [Installation](#installation)
- [Usage](#usage)
- [Known Issues](#known-issues)
- [Glossary](#glossary)
- [TODO](#todo)
- [Contributing](#contributing)
- [License](#license)

## Installation

```shell
pip install gpytorch-mogp
```

> **Note:** If you want to use GPU acceleration, you should manually install
> [PyTorch](https://pytorch.org/get-started/locally/) with CUDA support before running the above command.

If you want to run the [examples](examples), you should install the package from source instead. First, clone the
repository:

```shell
git clone https://github.com/dnv-opensource/gpytorch-mogp.git
```

Then install the package with the `examples` dependencies:

```shell
pip install .[examples]
```

To run the [comparison notebook](examples/comparison_with_rapid_models.ipynb), you also need to install
[`rapid-models`](https://github.com/RaPiD-models/rapid_models_dev2023) from source.

> **Note:** The version of `rapid-modles` linked to above is private at the time of writing, so you may not be able to
> install it.

## Usage

Usage is similar to the official [Multitask GP Regression](https://docs.gpytorch.ai/en/latest/examples/03_Multitask_Exact_GPs/Multitask_GP_Regression.html)
example.

The package provides a custom `MultiOutputKernel` module that is used similarly to `MultitaskKernel`, and a custom
`FixedNoiseMultiOutputGaussianLikelihood` that is used similarly to `MultitaskGaussianLikelihood`. The
`MultiOutputKernel` wraps one or more base kernels, producing a joint covariance matrix for the outputs. The
`FixedNoiseMultiOutputGaussianLikelihood` allows for fixed noise to be added to the joint covariance matrix.

```python
import gpytorch
from gpytorch_mogp import MultiOutputKernel, FixedNoiseMultiOutputGaussianLikelihood

# Define a multi-output GP model
class MultiOutputGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super().__init__(train_x, train_y, likelihood)
        # Reuse the `MultitaskMean` module from `gpytorch`
        self.mean_module = gpytorch.means.MultitaskMean(gpytorch.means.ConstantMean(), num_tasks=2)
        # Use the custom `MultiOutputKernel` module
        self.covar_module = MultiOutputKernel(gpytorch.kernels.MaternKernel(), num_outputs=2)

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        # Reuse the `MultitaskMultivariateNormal` distribution from `gpytorch`
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)

# Training data
train_x = ... # (n,) or (n, num_inputs)
train_y = ... # (n, num_outputs)
train_noise = ...  # (n, num_outputs, num_outputs)

# Initialize the model
likelihood = FixedNoiseMultiOutputGaussianLikelihood(noise=train_noise, num_tasks=2)
model = MultiOutputGPModel(train_x, train_y, likelihood)

# Training
model.train()
...

# Testing
model.eval()
test_x = ... # (m,) or (m, num_inputs)
test_noise = ...  # (m, num_outputs, num_outputs)
f_preds = model(test_x)
y_preds = likelihood(model(test_x), noise=test_noise)
```
> **Note:** `MultiOutputKernel` currently uses `num_outputs`, while `MultitaskMean` and
> `FixedNoiseMultiOutputGaussianLikelihood` use `num_tasks`. In this example, they mean the same thing. However, 
> multi-output and multi-task can have different meanings in other contexts. The reason for this inconsistency is that
> the `num_tasks` argument comes from `gpytorch`, while the `num_outputs` argument comes from `gpytorch-mogp`
> (`FixedNoiseMultiOutputGaussianLikelihood` inherits it from `gpytorch`). This inconsistency may be addressed in future
> versions of the package.

> **Note:** The `MultiOutputKernel` produces an [_interleaved_ block diagonal covariance matrix](#glossary) by default,
> as that is the convention used in `gpytorch`. If you want to produce a [_non-interleaved_ block diagonal covariance
> matrix](#glossary), you can pass `interleaved=False` to the `MultiOutputKernel` constructor.
>
> `MultitaskMultivariateNormal` expects an interleaved block diagonal covariance matrix by default. If you want to use a
> non-interleaved block diagonal covariance matrix, you can pass `interleaved=False` to the
> `MultitaskMultivariateNormal` constructor.
>
> `FixedNoiseMultiOutputGaussianLikelihood` expects the same noise structure regardless of interleaving, but will
> internally apply interleaving to the noise before adding it to the covariance matrix. Interleaving is applied by
> default. If you want to avoid this, you can pass `interleaved=False` to the `FixedNoiseMultiOutputGaussianLikelihood`
> constructor. The likelihood should always use the same interleaving setting as the kernel.
> 
> **WARNING:** The `interleaved=False` option is not working as expected at the moment. Avoid using it for now.

See also the [example notebooks/scripts](examples) for more usage examples.

See the [comparison notebook](examples/comparison_with_rapid_models.ipynb) for a comparison between the `gpytorch-mogp`
package and the `rapid-models` package (demonstrating that the two packages produce the same results for one example).

## Known Issues
- Constructing a GP with `interleaved=False` in `MultiOutputKernel`, `MultitaskMultivariateNormal`, and
    `FixedNoiseMultiOutputGaussianLikelihood` does not work as expected. Avoid using `interleaved=False` for now.
- The output of `MultiOutputKernel` is dense at the moment, as indexing into a block [linear operator]
    (https://github.com/cornellius-gp/linear_operator) is not working as expected.
- Important docstrings are missing.
- The code style is not consistent with the `gpytorch` code style (note that `gpytorch` itself does not have a good and
    consistent code style).
- There are currently no tests for the package, except for some doctests in the docstrings. The primary goal is to
    replicate the behavior of `rapid-models`, which is currently checked by comparing the results from _one_ example.
- Most issues are probably unknown at this point. Please create an issue if you encounter any problems ...

## Glossary
- **Multi-output GP model**: A Gaussian process model that produces multiple outputs for a given input. I.e. a model
    that predicts a vector-valued output for a given input.

- **Multi-output kernel**: A kernel that wraps one or more base kernels, producing a joint covariance matrix for the
    outputs. The joint covariance matrix is typically block diagonal, with each block corresponding to the output of a
    single base kernel. 

- **Block matrix**: A matrix that is partitioned into blocks. E.g. a block matrix with $2$ blocks per side of size $3$ 
    is structured like:
```math
\begin{bmatrix}
A & B \\
C & D
\end{bmatrix} = \begin{bmatrix}
a_{11} & a_{12} & a_{13} & b_{11} & b_{12} & b_{13} \\
a_{21} & a_{22} & a_{23} & b_{21} & b_{22} & b_{23} \\
a_{31} & a_{32} & a_{33} & b_{31} & b_{32} & b_{33} \\
c_{11} & c_{12} & c_{13} & d_{11} & d_{12} & d_{13} \\
c_{21} & c_{22} & c_{23} & d_{21} & d_{22} & d_{23} \\
c_{31} & c_{32} & c_{33} & d_{31} & d_{32} & d_{33}
\end{bmatrix}
```

- **Interleaved block matrix**: A block matrix where the blocks are interleaved. E.g. a block matrix with $2$ blocks per
    side of size $3$ is structured like:
```math
\begin{bmatrix}
a_{11} & b_{11} & a_{12} & b_{12} & a_{13} & b_{13} \\
c_{11} & d_{11} & c_{12} & d_{12} & c_{13} & d_{13} \\
a_{21} & b_{21} & a_{22} & b_{22} & a_{23} & b_{23} \\
c_{21} & d_{21} & c_{22} & d_{22} & c_{23} & d_{23} \\
a_{31} & b_{31} & a_{32} & b_{32} & a_{33} & b_{33} \\
c_{31} & d_{31} & c_{32} & d_{32} & c_{33} & d_{33}
\end{bmatrix}
```

- **(Non-interleaved) block diagonal covariance matrix**: A joint covariance matrix that is block diagonal, with each
    block corresponding to the output of a single base kernel. E.g. for a multi-output kernel with two base kernels,
    $`K_{\alpha}`$ and $`K_{\beta}`$, the joint covariance matrix is given by:
```math
K(\mathbf{X}, \mathbf{X}^*) = \begin{bmatrix}
K_{\alpha}(\mathbf{X}, \mathbf{X}^*) & \mathbf{0} \\
\mathbf{0} & K_{\beta}(\mathbf{X}, \mathbf{X}^*)
\end{bmatrix} = \begin{bmatrix}
K_{\alpha}(\mathbf{x}_1, \mathbf{x}_1^*) & \cdots & K_{\alpha}(\mathbf{x}_1, \mathbf{x}_n^*)\\
\vdots & \ddots & \vdots & & \mathbf{0} & \\
K_{\alpha}(\mathbf{x}_N, \mathbf{x}_1^*) & \cdots & K_{\alpha}(\mathbf{x}_N, \mathbf{x}_n^*) & & &\\
& & & K_{\beta}(\mathbf{x}_1, \mathbf{x}_1^*) & \cdots & K_{\beta}(\mathbf{x}_1, \mathbf{x}_n^*)\\
& \mathbf{0} & & \vdots & \ddots & \vdots\\
& & & K_{\beta}(\mathbf{x}_N, \mathbf{x}_1^*) & \cdots & K_{\beta}(\mathbf{x}_N, \mathbf{x}_n^*)
\end{bmatrix}
```

- **Interleaved block diagonal covariance matrix**: Similar to a block diagonal covariance matrix, but with the blocks
    interleaved:

```math
K(\mathbf{X}, \mathbf{X}^*) = \begin{bmatrix}
K_{\alpha}(\mathbf{x}_1, \mathbf{x}_1^*) & \mathbf{0} & \cdots & K_{\alpha}(\mathbf{x}_1, \mathbf{x}_n^*) & \mathbf{0} \\
\mathbf{0} & K_{\beta}(\mathbf{x}_1, \mathbf{x}_1^*) & \cdots & \mathbf{0} & K_{\beta}(\mathbf{x}_1, \mathbf{x}_n^*) \\
\vdots & \vdots & \ddots & \vdots & \vdots \\
K_{\alpha}(\mathbf{x}_N, \mathbf{x}_1^*) & \mathbf{0} & \cdots & K_{\alpha}(\mathbf{x}_N, \mathbf{x}_n^*) & \mathbf{0} \\
\mathbf{0} & K_{\beta}(\mathbf{x}_N, \mathbf{x}_1^*) & \cdots & \mathbf{0} & K_{\beta}(\mathbf{x}_N, \mathbf{x}_n^*)
\end{bmatrix}
```

## Contributing

Contributions are welcome! Please create an issue or a pull request if you have any suggestions or improvements.

To get started contributing, clone the repository:

```shell
git clone https://github.com/dnv-opensource/gpytorch-mogp.git
```

Then install the package in editable mode with `all` dependencies:
```shell
pip install -e .[all]
```

## License

`gpytorch-mogp` is distributed under the terms of the [MIT](https://spdx.org/licenses/MIT.html) license.
