"""Contains the implementations for surrogate modelling methods using GPyTorch
Implements the supported surrogate modelling methods using the GPyTorch library.
"""
import numpy as np
import pandas as pd
import torch
from matplotlib import pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from ..base import SurrogateBase
from ..estimators import get_single_task_exact_gp
[docs]
class GPyTorchSurrogateModel(SurrogateBase):
"""The GPyTorch surrogate modelling method class."""
[docs]
def execute(self) -> None:
"""Execute the simulation calibration procedure."""
surrogate_kwargs = self.get_calibration_func_kwargs()
n_samples = self.specification.n_samples
X = self.specification.X
if X is None:
X = self.sample_parameters(n_samples)
Y = self.specification.Y
if Y is None:
Y = self.calibration_func_wrapper(
X,
self,
self.specification.observed_data,
self.names,
self.data_types,
surrogate_kwargs,
)
self.Y_shape = Y.shape
if self.specification.flatten_Y and len(self.Y_shape) > 1:
X = self.extend_X(X, self.Y_shape[1])
Y = Y.flatten()
X_scaler = MinMaxScaler()
X_scaler.fit(X)
self.X_scaler = X_scaler
if torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
X = torch.tensor(X_scaler.transform(X), dtype=torch.double)
Y = torch.tensor(Y, dtype=torch.double)
model = get_single_task_exact_gp(
lr=self.specification.lr,
max_epochs=self.specification.n_iterations,
device=device,
)
model.fit(X, Y)
self.emulator = model
self.device = device
self.X = X
self.Y = Y
[docs]
def analyze(self) -> None:
"""Analyze the results of the simulation calibration procedure."""
task, time_now, experiment_name, outdir = self.prepare_analyze()
model = self.emulator
X = self.X
Y = self.Y
n_samples = self.specification.n_samples
X_sample = self.parameters.sample(n_samples, rule="sobol").T
if self.specification.flatten_Y and len(self.Y_shape) > 1:
X_sample = self.extend_X(X_sample, self.Y_shape[1])
X_sample = torch.tensor(self.X_scaler.transform(X_sample), dtype=torch.double)
Y_sample = model.predict(X_sample, return_std=False)
names = self.names.copy()
if X.shape[1] > len(names):
names.append("_dummy_index")
df = pd.DataFrame(X, columns=names)
if self.specification.use_shap and outdir is not None:
outfile = self.join(
outdir,
f"{time_now}-{task}-{experiment_name}-param-importances.png",
)
self.calculate_shap_importances(
df, self.emulator, names, self.specification.test_size, outfile
)
output_labels = self.specification.output_labels
if output_labels is None:
output_labels = ["output"]
output_label = output_labels[0]
if len(self.Y_shape) == 1:
df[f"simulated_{output_label}"] = Y
fig, axes = plt.subplots(
nrows=len(self.names), figsize=self.specification.figsize
)
if not isinstance(axes, np.ndarray):
axes = [axes]
for i, parameter_name in enumerate(self.names):
df.plot.scatter(
parameter_name,
f"simulated_{output_label}",
ax=axes[i],
title=f"simulated_{output_label} against {parameter_name}",
)
self.present_fig(fig, outdir, time_now, task, experiment_name, "plot-slice")
fig, axes = plt.subplots(nrows=2, figsize=self.specification.figsize)
df = pd.DataFrame(
{
"index": np.arange(0, Y.shape[0], 1),
"simulated": Y,
"emulated": Y_sample,
}
)
df.plot.scatter(
"index", "simulated", ax=axes[0], title=f"Simulated {output_label}"
)
df.plot.scatter(
"index", "emulated", ax=axes[1], title=f"Emulated {output_label}"
)
self.present_fig(
fig, outdir, time_now, task, experiment_name, f"emulated-{output_label}"
)
else:
if self.specification.flatten_Y:
df[f"simulated_{output_label}"] = Y
fig, axes = plt.subplots(
nrows=len(self.names), figsize=self.specification.figsize
)
if not isinstance(axes, np.ndarray):
axes = [axes]
for i, parameter_name in enumerate(self.names):
df.plot.scatter(
parameter_name,
f"simulated_{output_label}",
ax=axes[i],
title=f"simulated_{output_label} against {parameter_name}",
)
self.present_fig(
fig, outdir, time_now, task, experiment_name, "plot-slice"
)
reshaped_Y_sample = Y_sample.reshape(self.Y_shape)
Y = self.Y.reshape(self.Y_shape).detach().cpu().numpy()
index = np.arange(0, reshaped_Y_sample.shape[1], 1)
fig, axes = plt.subplots(nrows=2, figsize=self.specification.figsize)
for i in range(Y.shape[0]):
axes[0].plot(index, Y[i])
axes[0].set_title(f"Simulated {output_label}")
for i in range(reshaped_Y_sample.shape[0]):
axes[1].plot(index, reshaped_Y_sample[i])
axes[1].set_title(f"Emulated {output_label}")
self.present_fig(
fig,
outdir,
time_now,
task,
experiment_name,
f"emulated-{output_label}",
)