From 5119ee9626a3d4bcd7eb23c3b7af7fc8ff247a09 Mon Sep 17 00:00:00 2001 From: marjanfamili Date: Wed, 20 Aug 2025 15:28:40 +0100 Subject: [PATCH 1/8] partially_learnable.py was added which has Universal Kriging with partial prior knowledge. this class combines combines known function on one dimension with learnable linear function on remaining dimensions. It supports custum functions for the known component. It should be compatible with multi-task GPs through batch_shape parameter --- .../gaussian_process/partially_learnable.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 autoemulate/emulators/gaussian_process/partially_learnable.py diff --git a/autoemulate/emulators/gaussian_process/partially_learnable.py b/autoemulate/emulators/gaussian_process/partially_learnable.py new file mode 100644 index 000000000..694bd7276 --- /dev/null +++ b/autoemulate/emulators/gaussian_process/partially_learnable.py @@ -0,0 +1,76 @@ +import torch +import gpytorch +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 for one dimension + with a learnable linear function for the remaining dimensions. + + This implements Universal Kriging where part of the mean function is known + and part is learned from data. + + 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) + 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})" + ) \ No newline at end of file From 5b749b91d962338124339645c7478514ebccebac Mon Sep 17 00:00:00 2001 From: marjanfamili Date: Wed, 20 Aug 2025 15:37:32 +0100 Subject: [PATCH 2/8] added the partial mean function to the mean file --- .../emulators/gaussian_process/mean.py | 43 ++++++++++++++++++- .../gaussian_process/partially_learnable.py | 1 + 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/autoemulate/emulators/gaussian_process/mean.py b/autoemulate/emulators/gaussian_process/mean.py index 986f2d85b..7bc995d1b 100644 --- a/autoemulate/emulators/gaussian_process/mean.py +++ b/autoemulate/emulators/gaussian_process/mean.py @@ -2,7 +2,7 @@ from gpytorch.means import ConstantMean, LinearMean, ZeroMean from .poly_mean import PolyMean - +from .partially_learnable import PartiallyLearnableMean def constant_mean(n_features: int | None, n_outputs: torch.Size | None) -> ConstantMean: """ @@ -93,3 +93,44 @@ 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 + ) + ) \ No newline at end of file diff --git a/autoemulate/emulators/gaussian_process/partially_learnable.py b/autoemulate/emulators/gaussian_process/partially_learnable.py index 694bd7276..ba57cb2f6 100644 --- a/autoemulate/emulators/gaussian_process/partially_learnable.py +++ b/autoemulate/emulators/gaussian_process/partially_learnable.py @@ -1,5 +1,6 @@ import torch import gpytorch +from typing import Callable from gpytorch.means import LinearMean from autoemulate.core.types import TensorLike From 5900702ddb128e71270d946d96ba193491c00427 Mon Sep 17 00:00:00 2001 From: marjanfamili Date: Wed, 20 Aug 2025 16:12:30 +0100 Subject: [PATCH 3/8] added the mean function to the GP workflow --- autoemulate/emulators/gaussian_process/exact.py | 3 ++- autoemulate/emulators/gaussian_process/mean.py | 4 ++-- autoemulate/emulators/gaussian_process/partially_learnable.py | 1 - 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/autoemulate/emulators/gaussian_process/exact.py b/autoemulate/emulators/gaussian_process/exact.py index 7b7f30b5a..5cab47eee 100644 --- a/autoemulate/emulators/gaussian_process/exact.py +++ b/autoemulate/emulators/gaussian_process/exact.py @@ -38,7 +38,7 @@ rbf_times_linear, rq_kernel, ) -from .mean import constant_mean, linear_mean, poly_mean, zero_mean +from .mean import constant_mean, linear_mean, poly_mean, zero_mean, partially_learnable_mean class GaussianProcess(GaussianProcessEmulator, gpytorch.models.ExactGP): @@ -284,6 +284,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 7bc995d1b..3c3bc5d67 100644 --- a/autoemulate/emulators/gaussian_process/mean.py +++ b/autoemulate/emulators/gaussian_process/mean.py @@ -1,6 +1,6 @@ import torch from gpytorch.means import ConstantMean, LinearMean, ZeroMean - +from typing import Callable from .poly_mean import PolyMean from .partially_learnable import PartiallyLearnableMean @@ -97,7 +97,7 @@ def poly_mean(n_features: int, n_outputs: torch.Size | None) -> PolyMean: def partially_learnable_mean( n_features: int, n_outputs: torch.Size | None, - mean_func: callable = torch.sin, + mean_func: Callable = torch.sin, known_dim: int = 0 ) -> PartiallyLearnableMean: """ diff --git a/autoemulate/emulators/gaussian_process/partially_learnable.py b/autoemulate/emulators/gaussian_process/partially_learnable.py index ba57cb2f6..694bd7276 100644 --- a/autoemulate/emulators/gaussian_process/partially_learnable.py +++ b/autoemulate/emulators/gaussian_process/partially_learnable.py @@ -1,6 +1,5 @@ import torch import gpytorch -from typing import Callable from gpytorch.means import LinearMean from autoemulate.core.types import TensorLike From 2680ae3c92d3ac1cc42bfba89500b75b6208230b Mon Sep 17 00:00:00 2001 From: marjanfamili Date: Wed, 20 Aug 2025 22:44:03 +0100 Subject: [PATCH 4/8] added example for reviewer --- .../gaussian_process/partially_learnable.py | 3 +- .../emulation/03_universal_kriging.ipynb | 394 ++++++++++++++++++ 2 files changed, 396 insertions(+), 1 deletion(-) create mode 100644 docs/tutorials/emulation/03_universal_kriging.ipynb diff --git a/autoemulate/emulators/gaussian_process/partially_learnable.py b/autoemulate/emulators/gaussian_process/partially_learnable.py index 694bd7276..561b2d5d4 100644 --- a/autoemulate/emulators/gaussian_process/partially_learnable.py +++ b/autoemulate/emulators/gaussian_process/partially_learnable.py @@ -1,5 +1,6 @@ import torch import gpytorch +from typing import Callable from gpytorch.means import LinearMean from autoemulate.core.types import TensorLike @@ -28,7 +29,7 @@ class PartiallyLearnableMean(gpytorch.means.Mean): def __init__( self, - mean_func: callable, + mean_func: Callable, known_dim: int = 0, input_size: int = 1, batch_shape: torch.Size | None = None, diff --git a/docs/tutorials/emulation/03_universal_kriging.ipynb b/docs/tutorials/emulation/03_universal_kriging.ipynb new file mode 100644 index 000000000..b8ecf59a3 --- /dev/null +++ b/docs/tutorials/emulation/03_universal_kriging.ipynb @@ -0,0 +1,394 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9bfa9cf4", + "metadata": {}, + "source": [] + }, + { + "cell_type": "markdown", + "id": "cfdb4327", + "metadata": {}, + "source": [ + "# Demo for Universal Kriging implementation" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "1316c711", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(torch.Size([50, 2]), torch.Size([50, 1]))" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "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": "8010b1d7", + "metadata": {}, + "source": [] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "70844b6b", + "metadata": {}, + "outputs": [], + "source": [ + "from autoemulate import AutoEmulate \n", + "from autoemulate.emulators import GaussianProcess\n" + ] + }, + { + "cell_type": "markdown", + "id": "2b19dcc5", + "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": 12, + "id": "0dea1216", + "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": "a73c073d", + "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": 13, + "id": "275accfa", + "metadata": {}, + "outputs": [], + "source": [ + "ae = AutoEmulate(x, y, models=[GaussianProcess], log_level=\"error\")" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "caa2a2af", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
model_namex_transformsy_transformsparamsrmse_testr2_testr2_test_stdr2_trainr2_train_std
0GaussianProcess[StandardizeTransform()][StandardizeTransform()]{'mean_module_fn': <function partially_learnab...383.5930180.9800650.1467730.999840.000533
\n", + "
" + ], + "text/plain": [ + " model_name x_transforms y_transforms \\\n", + "0 GaussianProcess [StandardizeTransform()] [StandardizeTransform()] \n", + "\n", + " params rmse_test r2_test \\\n", + "0 {'mean_module_fn': " + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ae.plot(0)\n" + ] + }, + { + "cell_type": "markdown", + "id": "4f3a1108", + "metadata": {}, + "source": [ + "# No tuning \n", + "here ` mean_module.partially_learnable_mean` is forced without tuning " + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "6fc28e25", + "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": 17, + "id": "57329c99", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
model_namex_transformsy_transformsparamsrmse_testr2_testr2_test_stdr2_trainr2_train_std
0GaussianProcess[StandardizeTransform()][StandardizeTransform()]{'mean_module_fn': <function custom_mean at 0x...273.0160520.9464130.0854910.9999060.000065
\n", + "
" + ], + "text/plain": [ + " model_name x_transforms y_transforms \\\n", + "0 GaussianProcess [StandardizeTransform()] [StandardizeTransform()] \n", + "\n", + " params rmse_test r2_test \\\n", + "0 {'mean_module_fn': " + ] + }, + "execution_count": 18, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ae_2.plot(0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7cdbd380", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bcccb3d3", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a873b068", + "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 +} From e5a15c6264248ca1ee70170e9c2bb4aa44ece8e7 Mon Sep 17 00:00:00 2001 From: marjanfamili Date: Thu, 21 Aug 2025 11:21:04 +0100 Subject: [PATCH 5/8] more documentation --- .../emulators/gaussian_process/exact.py | 8 +- .../emulators/gaussian_process/mean.py | 24 +- .../gaussian_process/partially_learnable.py | 37 +- .../experimental/universal_kriging.ipynb | 402 ++++++++++++++++++ .../emulation/03_universal_kriging.ipynb | 394 ----------------- 5 files changed, 438 insertions(+), 427 deletions(-) create mode 100644 autoemulate/experimental/universal_kriging.ipynb delete mode 100644 docs/tutorials/emulation/03_universal_kriging.ipynb diff --git a/autoemulate/emulators/gaussian_process/exact.py b/autoemulate/emulators/gaussian_process/exact.py index 5cab47eee..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, partially_learnable_mean +from .mean import ( + constant_mean, + linear_mean, + partially_learnable_mean, + poly_mean, + zero_mean, +) class GaussianProcess(GaussianProcessEmulator, gpytorch.models.ExactGP): diff --git a/autoemulate/emulators/gaussian_process/mean.py b/autoemulate/emulators/gaussian_process/mean.py index 3c3bc5d67..07eed7bb9 100644 --- a/autoemulate/emulators/gaussian_process/mean.py +++ b/autoemulate/emulators/gaussian_process/mean.py @@ -1,8 +1,11 @@ +from collections.abc import Callable + import torch from gpytorch.means import ConstantMean, LinearMean, ZeroMean -from typing import Callable -from .poly_mean import PolyMean + from .partially_learnable import PartiallyLearnableMean +from .poly_mean import PolyMean + def constant_mean(n_features: int | None, n_outputs: torch.Size | None) -> ConstantMean: """ @@ -94,11 +97,12 @@ def poly_mean(n_features: int, n_outputs: torch.Size | None) -> PolyMean: else PolyMean(degree=2, input_size=n_features) ) + def partially_learnable_mean( - n_features: int, - n_outputs: torch.Size | None, + n_features: int, + n_outputs: torch.Size | None, mean_func: Callable = torch.sin, - known_dim: int = 0 + known_dim: int = 0, ) -> PartiallyLearnableMean: """ PartiallyLearnableMean module with known function for one dimension. @@ -114,7 +118,7 @@ def partially_learnable_mean( 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 @@ -125,12 +129,10 @@ def partially_learnable_mean( mean_func=mean_func, known_dim=known_dim, input_size=n_features, - batch_shape=n_outputs + batch_shape=n_outputs, ) if n_outputs is not None else PartiallyLearnableMean( - mean_func=mean_func, - known_dim=known_dim, - input_size=n_features + mean_func=mean_func, known_dim=known_dim, input_size=n_features ) - ) \ No newline at end of file + ) diff --git a/autoemulate/emulators/gaussian_process/partially_learnable.py b/autoemulate/emulators/gaussian_process/partially_learnable.py index 561b2d5d4..a4acf320a 100644 --- a/autoemulate/emulators/gaussian_process/partially_learnable.py +++ b/autoemulate/emulators/gaussian_process/partially_learnable.py @@ -1,6 +1,6 @@ -import torch +from collections.abc import Callable import gpytorch -from typing import Callable +import torch from gpytorch.means import LinearMean from autoemulate.core.types import TensorLike @@ -9,13 +9,13 @@ class PartiallyLearnableMean(gpytorch.means.Mean): """ A mixed mean module that combines a known function for one dimension with a learnable linear function for the remaining dimensions. - + This implements Universal Kriging where part of the mean function is known and part is learned from data. - + Parameters ---------- - mean_func : callable + 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 @@ -26,7 +26,7 @@ class PartiallyLearnableMean(gpytorch.means.Mean): batch_shape : torch.Size | None Optional batch dimension for multi-task GPs. """ - + def __init__( self, mean_func: Callable, @@ -35,43 +35,38 @@ def __init__( 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 + input_size=len(self.learnable_dims), batch_shape=batch_shape ) - def forward(self, x: TensorLike) -> TensorLike: - """ - Forward pass through the partially learnable mean module. - """ + """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) + 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') + 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})" - ) \ No newline at end of file + ) diff --git a/autoemulate/experimental/universal_kriging.ipynb b/autoemulate/experimental/universal_kriging.ipynb new file mode 100644 index 000000000..3f1b7cf41 --- /dev/null +++ b/autoemulate/experimental/universal_kriging.ipynb @@ -0,0 +1,402 @@ +{ + "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": 1, + "id": "2", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(torch.Size([50, 2]), torch.Size([50, 1]))" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "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": 2, + "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": 3, + "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": 4, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "ae = AutoEmulate(x, y, models=[GaussianProcess], log_level=\"error\")" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "9", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
model_namex_transformsy_transformsparamsrmse_testr2_testr2_test_stdr2_trainr2_train_std
0GaussianProcess[StandardizeTransform()][StandardizeTransform()]{'mean_module_fn': <function zero_mean at 0x33...422.483490.6665551.9029980.9998910.000173
\n", + "
" + ], + "text/plain": [ + " model_name x_transforms y_transforms \\\n", + "0 GaussianProcess [StandardizeTransform()] [StandardizeTransform()] \n", + "\n", + " params rmse_test r2_test \\\n", + "0 {'mean_module_fn': " + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "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": 7, + "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": 8, + "id": "13", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
model_namex_transformsy_transformsparamsrmse_testr2_testr2_test_stdr2_trainr2_train_std
0GaussianProcess[StandardizeTransform()][StandardizeTransform()]{'mean_module_fn': <function custom_mean at 0x...362.603912-87.338058283.282410.9998910.000077
\n", + "
" + ], + "text/plain": [ + " model_name x_transforms y_transforms \\\n", + "0 GaussianProcess [StandardizeTransform()] [StandardizeTransform()] \n", + "\n", + " params rmse_test r2_test \\\n", + "0 {'mean_module_fn': " + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "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/docs/tutorials/emulation/03_universal_kriging.ipynb b/docs/tutorials/emulation/03_universal_kriging.ipynb deleted file mode 100644 index b8ecf59a3..000000000 --- a/docs/tutorials/emulation/03_universal_kriging.ipynb +++ /dev/null @@ -1,394 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "9bfa9cf4", - "metadata": {}, - "source": [] - }, - { - "cell_type": "markdown", - "id": "cfdb4327", - "metadata": {}, - "source": [ - "# Demo for Universal Kriging implementation" - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "id": "1316c711", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "(torch.Size([50, 2]), torch.Size([50, 1]))" - ] - }, - "execution_count": 10, - "metadata": {}, - "output_type": "execute_result" - } - ], - "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": "8010b1d7", - "metadata": {}, - "source": [] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "70844b6b", - "metadata": {}, - "outputs": [], - "source": [ - "from autoemulate import AutoEmulate \n", - "from autoemulate.emulators import GaussianProcess\n" - ] - }, - { - "cell_type": "markdown", - "id": "2b19dcc5", - "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": 12, - "id": "0dea1216", - "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": "a73c073d", - "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": 13, - "id": "275accfa", - "metadata": {}, - "outputs": [], - "source": [ - "ae = AutoEmulate(x, y, models=[GaussianProcess], log_level=\"error\")" - ] - }, - { - "cell_type": "code", - "execution_count": 14, - "id": "caa2a2af", - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
model_namex_transformsy_transformsparamsrmse_testr2_testr2_test_stdr2_trainr2_train_std
0GaussianProcess[StandardizeTransform()][StandardizeTransform()]{'mean_module_fn': <function partially_learnab...383.5930180.9800650.1467730.999840.000533
\n", - "
" - ], - "text/plain": [ - " model_name x_transforms y_transforms \\\n", - "0 GaussianProcess [StandardizeTransform()] [StandardizeTransform()] \n", - "\n", - " params rmse_test r2_test \\\n", - "0 {'mean_module_fn': " - ] - }, - "execution_count": 15, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "ae.plot(0)\n" - ] - }, - { - "cell_type": "markdown", - "id": "4f3a1108", - "metadata": {}, - "source": [ - "# No tuning \n", - "here ` mean_module.partially_learnable_mean` is forced without tuning " - ] - }, - { - "cell_type": "code", - "execution_count": 16, - "id": "6fc28e25", - "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": 17, - "id": "57329c99", - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
model_namex_transformsy_transformsparamsrmse_testr2_testr2_test_stdr2_trainr2_train_std
0GaussianProcess[StandardizeTransform()][StandardizeTransform()]{'mean_module_fn': <function custom_mean at 0x...273.0160520.9464130.0854910.9999060.000065
\n", - "
" - ], - "text/plain": [ - " model_name x_transforms y_transforms \\\n", - "0 GaussianProcess [StandardizeTransform()] [StandardizeTransform()] \n", - "\n", - " params rmse_test r2_test \\\n", - "0 {'mean_module_fn': " - ] - }, - "execution_count": 18, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "ae_2.plot(0)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7cdbd380", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "bcccb3d3", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "a873b068", - "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 -} From 07223fc4203d2e84703beb6088de8ab44be22af7 Mon Sep 17 00:00:00 2001 From: Edwin Brown Date: Thu, 21 Aug 2025 11:47:22 +0100 Subject: [PATCH 6/8] Move import outside of test function --- tests/transforms/test_transforms_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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]: From 1473e52a4e2357a6e7fe1fb4cc6757fefb7a6e4a Mon Sep 17 00:00:00 2001 From: marjanfamili Date: Thu, 21 Aug 2025 12:13:17 +0100 Subject: [PATCH 7/8] moved notebook to exploratory --- autoemulate/core/compare.py | 2 - .../gaussian_process/partially_learnable.py | 15 +- .../exploratory/universal_kriging.ipynb | 237 +++++++++++ .../experimental/universal_kriging.ipynb | 402 ------------------ 4 files changed, 246 insertions(+), 410 deletions(-) create mode 100644 autoemulate/experimental/exploratory/universal_kriging.ipynb delete mode 100644 autoemulate/experimental/universal_kriging.ipynb 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/partially_learnable.py b/autoemulate/emulators/gaussian_process/partially_learnable.py index a4acf320a..bd34269a3 100644 --- a/autoemulate/emulators/gaussian_process/partially_learnable.py +++ b/autoemulate/emulators/gaussian_process/partially_learnable.py @@ -1,17 +1,18 @@ 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 for one dimension - with a learnable linear function for the remaining dimensions. - - This implements Universal Kriging where part of the mean function is known - and part is learned from data. + A mixed mean module that combines a known function with learnable components. + Specifically, it applies a known function to one dimension and uses a learnable + linear function for the remaining dimensions. This implements Universal Kriging + where part of the mean function is known and part is learned from data. Parameters ---------- @@ -58,7 +59,9 @@ def forward(self, x: TensorLike) -> TensorLike: 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 + 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: 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/autoemulate/experimental/universal_kriging.ipynb b/autoemulate/experimental/universal_kriging.ipynb deleted file mode 100644 index 3f1b7cf41..000000000 --- a/autoemulate/experimental/universal_kriging.ipynb +++ /dev/null @@ -1,402 +0,0 @@ -{ - "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": 1, - "id": "2", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "(torch.Size([50, 2]), torch.Size([50, 1]))" - ] - }, - "execution_count": 1, - "metadata": {}, - "output_type": "execute_result" - } - ], - "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": 2, - "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": 3, - "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": 4, - "id": "8", - "metadata": {}, - "outputs": [], - "source": [ - "ae = AutoEmulate(x, y, models=[GaussianProcess], log_level=\"error\")" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "9", - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
model_namex_transformsy_transformsparamsrmse_testr2_testr2_test_stdr2_trainr2_train_std
0GaussianProcess[StandardizeTransform()][StandardizeTransform()]{'mean_module_fn': <function zero_mean at 0x33...422.483490.6665551.9029980.9998910.000173
\n", - "
" - ], - "text/plain": [ - " model_name x_transforms y_transforms \\\n", - "0 GaussianProcess [StandardizeTransform()] [StandardizeTransform()] \n", - "\n", - " params rmse_test r2_test \\\n", - "0 {'mean_module_fn': " - ] - }, - "execution_count": 6, - "metadata": {}, - "output_type": "execute_result" - } - ], - "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": 7, - "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": 8, - "id": "13", - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
model_namex_transformsy_transformsparamsrmse_testr2_testr2_test_stdr2_trainr2_train_std
0GaussianProcess[StandardizeTransform()][StandardizeTransform()]{'mean_module_fn': <function custom_mean at 0x...362.603912-87.338058283.282410.9998910.000077
\n", - "
" - ], - "text/plain": [ - " model_name x_transforms y_transforms \\\n", - "0 GaussianProcess [StandardizeTransform()] [StandardizeTransform()] \n", - "\n", - " params rmse_test r2_test \\\n", - "0 {'mean_module_fn': " - ] - }, - "execution_count": 9, - "metadata": {}, - "output_type": "execute_result" - } - ], - "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 -} From 7085be1a25fdfbf7dd68a626c66cef4adb54b52a Mon Sep 17 00:00:00 2001 From: marjanfamili Date: Thu, 21 Aug 2025 12:30:38 +0100 Subject: [PATCH 8/8] changed docstring --- autoemulate/emulators/gaussian_process/partially_learnable.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/autoemulate/emulators/gaussian_process/partially_learnable.py b/autoemulate/emulators/gaussian_process/partially_learnable.py index bd34269a3..320e82c8e 100644 --- a/autoemulate/emulators/gaussian_process/partially_learnable.py +++ b/autoemulate/emulators/gaussian_process/partially_learnable.py @@ -10,9 +10,6 @@ class PartiallyLearnableMean(gpytorch.means.Mean): """ A mixed mean module that combines a known function with learnable components. - Specifically, it applies a known function to one dimension and uses a learnable - linear function for the remaining dimensions. This implements Universal Kriging - where part of the mean function is known and part is learned from data. Parameters ----------