diff --git a/autoemulate/core/compare.py b/autoemulate/core/compare.py index d8e47b587..f5fe2dd3f 100644 --- a/autoemulate/core/compare.py +++ b/autoemulate/core/compare.py @@ -539,8 +539,6 @@ def fit_from_reinitialized( neural networks. """ - from autoemulate.emulators import get_emulator_class - transformed_emulator_params = ( transformed_emulator_params or self.transformed_emulator_params ) diff --git a/autoemulate/emulators/gaussian_process/exact.py b/autoemulate/emulators/gaussian_process/exact.py index 7b7f30b5a..d68181363 100644 --- a/autoemulate/emulators/gaussian_process/exact.py +++ b/autoemulate/emulators/gaussian_process/exact.py @@ -38,7 +38,13 @@ rbf_times_linear, rq_kernel, ) -from .mean import constant_mean, linear_mean, poly_mean, zero_mean +from .mean import ( + constant_mean, + linear_mean, + partially_learnable_mean, + poly_mean, + zero_mean, +) class GaussianProcess(GaussianProcessEmulator, gpytorch.models.ExactGP): @@ -284,6 +290,7 @@ def get_tune_params(): zero_mean, linear_mean, poly_mean, + partially_learnable_mean, ], "covar_module_fn": [ rbf, diff --git a/autoemulate/emulators/gaussian_process/mean.py b/autoemulate/emulators/gaussian_process/mean.py index 986f2d85b..07eed7bb9 100644 --- a/autoemulate/emulators/gaussian_process/mean.py +++ b/autoemulate/emulators/gaussian_process/mean.py @@ -1,6 +1,9 @@ +from collections.abc import Callable + import torch from gpytorch.means import ConstantMean, LinearMean, ZeroMean +from .partially_learnable import PartiallyLearnableMean from .poly_mean import PolyMean @@ -93,3 +96,43 @@ def poly_mean(n_features: int, n_outputs: torch.Size | None) -> PolyMean: if n_outputs is not None else PolyMean(degree=2, input_size=n_features) ) + + +def partially_learnable_mean( + n_features: int, + n_outputs: torch.Size | None, + mean_func: Callable = torch.sin, + known_dim: int = 0, +) -> PartiallyLearnableMean: + """ + PartiallyLearnableMean module with known function for one dimension. + + Parameters + ---------- + n_features: int + Number of input features. + n_outputs: torch.Size | None + Batch shape of the mean. If None, the mean is not initialized with a batch + shape. + mean_func: callable + Function to apply to the known dimension. Defaults to torch.sin. + known_dim: int + Dimension index for the known function. Defaults to 0. + + Returns + ------- + PartiallyLearnableMean + The initialized PartiallyLearnableMean module. + """ + return ( + PartiallyLearnableMean( + mean_func=mean_func, + known_dim=known_dim, + input_size=n_features, + batch_shape=n_outputs, + ) + if n_outputs is not None + else PartiallyLearnableMean( + mean_func=mean_func, known_dim=known_dim, input_size=n_features + ) + ) diff --git a/autoemulate/emulators/gaussian_process/partially_learnable.py b/autoemulate/emulators/gaussian_process/partially_learnable.py new file mode 100644 index 000000000..320e82c8e --- /dev/null +++ b/autoemulate/emulators/gaussian_process/partially_learnable.py @@ -0,0 +1,72 @@ +from collections.abc import Callable + +import gpytorch +import torch +from gpytorch.means import LinearMean + +from autoemulate.core.types import TensorLike + + +class PartiallyLearnableMean(gpytorch.means.Mean): + """ + A mixed mean module that combines a known function with learnable components. + + Parameters + ---------- + mean_func : Callable + A function that takes a tensor and returns a tensor of the same shape. + This function will be applied to the specified dimension. + known_dim : int + The dimension index to which the custom mean function will be applied. + Default is 0 (first dimension). + input_size : int + The total number of input features. + batch_shape : torch.Size | None + Optional batch dimension for multi-task GPs. + """ + + def __init__( + self, + mean_func: Callable, + known_dim: int = 0, + input_size: int = 1, + batch_shape: torch.Size | None = None, + ): + super().__init__() + + # Store the custom function and dimension + self.mean_func = mean_func + self.known_dim = known_dim + self.input_size = input_size + + if batch_shape is None: + batch_shape = torch.Size() + + # Create indices for learnable dimensions (all except known_dim) + self.learnable_dims = [i for i in range(input_size) if i != known_dim] + + # Only create linear mean if there are learnable dimensions + self.linear_mean = LinearMean( + input_size=len(self.learnable_dims), batch_shape=batch_shape + ) + + def forward(self, x: TensorLike) -> TensorLike: + """Forward pass through the partially learnable mean module.""" + # Apply custom mean function to the known dimension + known_part = self.mean_func(x[..., self.known_dim]) + + learnable_data = x[..., self.learnable_dims] + learnable_part = self.linear_mean( + learnable_data + ) # this part could be replaced with other function / NN + return known_part + learnable_part + + def __repr__(self) -> str: + """Return string representation of the PartiallyLearnableMean module.""" + func_name = getattr(self.mean_func, "__name__", "mean_func") + return ( + f"PartiallyLearnableMean(" + f"mean_func={func_name}, " + f"known_dim={self.known_dim}, " + f"input_size={self.input_size})" + ) diff --git a/autoemulate/experimental/exploratory/universal_kriging.ipynb b/autoemulate/experimental/exploratory/universal_kriging.ipynb new file mode 100644 index 000000000..982b3e14e --- /dev/null +++ b/autoemulate/experimental/exploratory/universal_kriging.ipynb @@ -0,0 +1,237 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "# Demo for Universal Kriging implementation\n", + "\n", + "This terminology is borrowed from this paper \n", + "https://arxiv.org/pdf/2408.02331\n", + "\n", + "1. Simple Kriging: Known mean function, no noise\n", + "2. Ordinary Kriging: Unknown constant mean, no noise\n", + "3. Universal Kriging: Unknown mean as linear combination of known functions\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from autoemulate.simulations.projectile import Projectile\n", + "\n", + "projectile = Projectile(log_level=\"error\")\n", + "n_samples = 50\n", + "x = projectile.sample_inputs(n_samples).float()\n", + "y = projectile.forward_batch(x).float()\n", + "x.shape, y.shape" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from autoemulate import AutoEmulate \n", + "from autoemulate.emulators import GaussianProcess\n" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "# define custom mean \n", + "for example here we try to incorporates some knowledge of projectile motion physics into the GP. With drag, there is no simple closed form solution and simulaiton is solved numerically. for no drag : $ R = v_0^2 sin(2\\theta)/g$ \n", + "\n", + "cutom_mean is returning the `mean_module.PartiallyLearnableMean` class which has `projectile_mean` mean_func\n", + "\n", + "and we replace `mean_module.partially_learnable_mean`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "import autoemulate.emulators.gaussian_process.mean as mean_module\n", + "\n", + "def projectile_mean(x):\n", + " return x**2/9.8\n", + "\n", + "def custom_mean(n_features, n_outputs):\n", + " return mean_module.PartiallyLearnableMean(\n", + " mean_func=projectile_mean,\n", + " known_dim=0,\n", + " input_size=n_features,\n", + " batch_shape=n_outputs\n", + " )\n", + "\n", + "mean_module.partially_learnable_mean = custom_mean\n" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "This means that here , \n", + "\n", + "```python \n", + " return {\n", + " \"mean_module_fn\": [\n", + " constant_mean,\n", + " zero_mean,\n", + " linear_mean,\n", + " poly_mean,\n", + " partially_learnable_mean,\n", + " ],\n", + "```\n", + "\n", + "it itterates over these and our updated `partially_learnable_mean` and choose the best, as here the result is not the best for `partially_learnable_mean` I also check this for a case without tuning " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "ae = AutoEmulate(x, y, models=[GaussianProcess], log_level=\"error\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "ae.summarise()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "ae.plot(0)\n" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "# No tuning \n", + "here ` mean_module.partially_learnable_mean` is forced without tuning " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "ae_2 = AutoEmulate(\n", + " x, y, \n", + " models=[GaussianProcess],\n", + " model_tuning=False,\n", + " model_params={\"mean_module_fn\": mean_module.partially_learnable_mean},\n", + " log_level=\"error\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "ae_2.summarise()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "ae_2.plot(0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/transforms/test_transforms_utils.py b/tests/transforms/test_transforms_utils.py index e841f2cc4..88518daf4 100644 --- a/tests/transforms/test_transforms_utils.py +++ b/tests/transforms/test_transforms_utils.py @@ -1,11 +1,11 @@ import pytest +import torch +from autoemulate.transforms.utils import make_positive_definite @pytest.mark.parametrize("scale", [1e-6, 1.0, 1e6]) def test_make_positive_definite(scale): """Test that make_positive_definite can handle various cases of cov matrices.""" - import torch - from autoemulate.transforms.utils import make_positive_definite # Try large scale for dim in [5, 20, 100]: