From db4aa0787233b7ef226a525d29e82217b11b1fbb Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 14 Apr 2026 02:09:48 -0400 Subject: [PATCH 1/9] add custom generation strategies --- src/blop/ax/agent.py | 4 + src/blop/bayesian/kernels.py | 418 ++++++++++++++---------- src/blop/bayesian/models.py | 101 ++++-- src/blop/evaluation/__init__.py | 0 src/blop/evaluation/test_functions.py | 3 + src/blop/tests/ax/test_generation.py | 89 +++++ src/blop/tests/bayesian/test_kernels.py | 137 -------- src/blop/tests/generation/__init__.py | 0 8 files changed, 418 insertions(+), 334 deletions(-) create mode 100644 src/blop/evaluation/__init__.py create mode 100644 src/blop/evaluation/test_functions.py create mode 100644 src/blop/tests/ax/test_generation.py delete mode 100644 src/blop/tests/bayesian/test_kernels.py create mode 100644 src/blop/tests/generation/__init__.py diff --git a/src/blop/ax/agent.py b/src/blop/ax/agent.py index a98f4f04..d414826f 100644 --- a/src/blop/ax/agent.py +++ b/src/blop/ax/agent.py @@ -242,6 +242,7 @@ def __init__( dof_constraints: Sequence[DOFConstraint] | None = None, outcome_constraints: Sequence[OutcomeConstraint] | None = None, checkpoint_path: str | None = None, + generation_strategy: GenerationStrategy = None, **kwargs: Any, ): if any(isinstance(dof.actuator, str) for dof in dofs): @@ -267,6 +268,9 @@ def __init__( self._callbacks: list[CallbackBase] = [OptimizationLogger()] self._callback_router = OptimizationCallbackRouter(self._callbacks) + if generation_strategy is not None: + self.ax_client.set_generation_strategy(generation_strategy) + @classmethod def from_checkpoint( cls, diff --git a/src/blop/bayesian/kernels.py b/src/blop/bayesian/kernels.py index c874ad01..add4abb1 100644 --- a/src/blop/bayesian/kernels.py +++ b/src/blop/bayesian/kernels.py @@ -6,186 +6,248 @@ import torch -class LatentKernel(gpytorch.kernels.Kernel): - @property - def is_stationary(self) -> bool: - return True - num_outputs: int = 1 - batch_inverse_lengthscale: float = 1e6 +import math +import torch +import gpytorch +from gpytorch.kernels import Kernel +from gpytorch.constraints import Interval +from botorch.models.utils.gpytorch_modules import get_matern_kernel_with_gamma_prior + + +class RotatedInputsKernel(Kernel): + """Applies a learned orthogonal rotation (blockwise if desired), then delegates to a base kernel.""" + is_stationary = True + has_lengthscale = False # lengthscale lives in the base kernel def __init__( self, - num_inputs: int = 1, - skew_dims: bool | Iterable[tuple[int, ...]] = True, - priors: bool = True, - scale_output: bool = True, - **kwargs, - ) -> None: - super().__init__() - - self.num_inputs: int = num_inputs - self.scale_output: bool = scale_output - - self.nu: float = kwargs.get("nu", 2.5) - - if type(skew_dims) is bool: - if skew_dims: - self.skew_dims: list[torch.Tensor] = [torch.arange(self.num_inputs)] - else: - self.skew_dims = [torch.arange(num_inputs)] - elif hasattr(skew_dims, "__iter__"): - self.skew_dims = [torch.tensor(np.atleast_1d(skew_group)) for skew_group in skew_dims] + d: int, + batch_shape: torch.Size = torch.Size(), + skew_dims: bool | list[tuple[int, ...]] = True, + ): + super().__init__(batch_shape=batch_shape) + + self.d = d + # Choose groups of dims that can rotate among themselves + if isinstance(skew_dims, bool): + self.groups = [tuple(range(d))] if skew_dims else [] else: - raise ValueError('arg "skew_dims" must be True, False, or an iterable of tuples of ints.') - - # if not all([len(skew_group) >= 2 for skew_group in self.skew_dims]): - # raise ValueError("must have at least two dims per skew group") - skewed_dims = [dim.item() for skew_group in self.skew_dims for dim in skew_group] - if not len(set(skewed_dims)) == len(skewed_dims): - raise ValueError("values in skew_dims must be unique") - if not max(skewed_dims) < self.num_inputs: - raise ValueError("invalid dimension index in skew_dims") - - skew_group_submatrix_indices: list[torch.Tensor] = [] - for dim in range(self.num_outputs): - for skew_group in self.skew_dims: - j, k = skew_group[torch.triu_indices(len(skew_group), len(skew_group), 1)].unsqueeze(1) - i = dim * torch.ones(j.shape).long() - skew_group_submatrix_indices.append(torch.cat((i, j, k), dim=0)) - - self.diag_matrix_indices: list[torch.Tensor] = [ - torch.kron(torch.arange(self.num_outputs), torch.ones(self.num_inputs)).long(), - *2 * [torch.arange(self.num_inputs).repeat(self.num_outputs)], - ] - - self.skew_matrix_indices: tuple[torch.Tensor, ...] = ( - tuple(torch.cat(skew_group_submatrix_indices, dim=1)) - if len(skew_group_submatrix_indices) > 0 - else (torch.tensor([]), torch.tensor([])) - ) - - self.n_skew_entries: int = len(self.skew_matrix_indices[0]) - - lengthscale_constraint = gpytorch.constraints.Positive() - raw_lengthscales_initial = lengthscale_constraint.inverse_transform(torch.tensor(1e-1)) * torch.ones( - self.num_outputs, self.num_inputs, dtype=torch.double - ) - - self.register_parameter(name="raw_lengthscales", parameter=torch.nn.Parameter(raw_lengthscales_initial)) - self.register_constraint(param_name="raw_lengthscales", constraint=lengthscale_constraint) - - if priors: - self.register_prior( - name="lengthscales_prior", - prior=gpytorch.priors.GammaPrior(concentration=3, rate=6), - param_or_closure=lambda m: m.lengthscales, - setting_closure=lambda m, v: m._set_lengthscales(v), - ) - - if self.n_skew_entries > 0: - skew_entries_constraint = gpytorch.constraints.Interval(-2 * np.pi, 2 * np.pi) - skew_entries_initial = torch.zeros((self.num_outputs, self.n_skew_entries), dtype=torch.float64) - self.register_parameter(name="raw_skew_entries", parameter=torch.nn.Parameter(skew_entries_initial)) - self.register_constraint(param_name="raw_skew_entries", constraint=skew_entries_constraint) - - if self.scale_output: - outputscale_constraint = gpytorch.constraints.Positive() - outputscale_prior = gpytorch.priors.GammaPrior(concentration=2, rate=0.15) - - self.register_parameter( - name="raw_outputscale", - parameter=torch.nn.Parameter(torch.ones(1, dtype=torch.double)), - ) - - self.register_constraint("raw_outputscale", constraint=outputscale_constraint) - - self.register_prior( - name="outputscale_prior", - prior=outputscale_prior, - param_or_closure=lambda m: m.outputscale, - setting_closure=lambda m, v: m._set_outputscale(v), - ) - - @property - def lengthscales(self) -> torch.Tensor: - return self.raw_lengthscales_constraint.transform(self.raw_lengthscales) - - @lengthscales.setter - def lengthscales(self, value: torch.Tensor | float | np.ndarray) -> None: - self._set_lengthscales(value) - - @property - def skew_entries(self) -> torch.Tensor: - return self.raw_skew_entries_constraint.transform(self.raw_skew_entries) - - @skew_entries.setter - def skew_entries(self, value: torch.Tensor | float | np.ndarray) -> None: - self._set_skew_entries(value) - - @property - def outputscale(self) -> torch.Tensor: - return self.raw_outputscale_constraint.transform(self.raw_outputscale) - - @outputscale.setter - def outputscale(self, value: torch.Tensor | float | np.ndarray) -> None: - self._set_outputscale(value) - - def _set_lengthscales(self, value: torch.Tensor | float | np.ndarray) -> None: - if not torch.is_tensor(value): - value = torch.as_tensor(value).to(self.raw_lengthscales) - self.initialize(raw_lengthscales=self.raw_lengthscales_constraint.inverse_transform(value)) - - def _set_skew_entries(self, value: torch.Tensor | float | np.ndarray) -> None: - if not torch.is_tensor(value): - value = torch.as_tensor(value).to(self.raw_skew_entries) - self.initialize(raw_skew_entries=self.raw_skew_entries_constraint.inverse_transform(value)) - - def _set_outputscale(self, value: torch.Tensor | float | np.ndarray) -> None: - if not torch.is_tensor(value): - value = torch.as_tensor(value).to(self.raw_outputscale) - self.initialize(raw_outputscale=self.raw_outputscale_constraint.inverse_transform(value)) - - @property - def skew_matrix(self) -> torch.Tensor: - S = torch.zeros((self.num_outputs, self.num_inputs, self.num_inputs), dtype=torch.float64) - if self.n_skew_entries > 0: - # to construct an orthogonal matrix. fun fact: exp(skew(N)) is the generator of SO(N) - S[self.skew_matrix_indices] = self.skew_entries - S += -S.transpose(-1, -2) - return torch.linalg.matrix_exp(S) - - @property - def diag_matrix(self) -> torch.Tensor: - D = torch.zeros((self.num_outputs, self.num_inputs, self.num_inputs), dtype=torch.float64) - D[self.diag_matrix_indices] = self.lengthscales.ravel() ** -1 - return D - - @property - def latent_transform(self) -> torch.Tensor: - return torch.matmul(self.diag_matrix, self.skew_matrix) - - def forward(self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, **params: dict) -> torch.Tensor: - # adapted from the Matern kernel - mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)] - - transform = self.latent_transform.unsqueeze(1) - - trans_x1 = torch.matmul(transform, (x1 - mean).unsqueeze(-1)).squeeze(-1) - trans_x2 = torch.matmul(transform, (x2 - mean).unsqueeze(-1)).squeeze(-1) - - distance = self.covar_dist(trans_x1, trans_x2, diag=diag, **params) - - if self.num_outputs == 1: - distance = distance.squeeze(0) - - outputscale = self.outputscale if self.scale_output else 1.0 - - # special cases of the Matern function - if self.nu == 0.5: - return outputscale * torch.exp(-distance) - if self.nu == 1.5: - return outputscale * (1 + distance) * torch.exp(-distance) - if self.nu == 2.5: - return outputscale * (1 + distance + distance**2 / 3) * torch.exp(-distance) - raise ValueError(f"nu = {self.nu} is not supported") + self.groups = [tuple(g) for g in skew_dims] + + # Base kernel = ScaleKernel(MaternKernel(...)) with your Gamma priors + self.base = get_matern_kernel_with_gamma_prior(ard_num_dims=d, batch_shape=batch_shape) + + # One unconstrained matrix per group; we’ll skew-symmetrize then expm -> orthogonal + self.raw_group_mats = torch.nn.ParameterList() + for i, g in enumerate(self.groups): + k = len(g) + raw = torch.zeros(int(k * (k - 1) / 2), dtype=torch.double) + setattr(self, f"raw_group_entries_{i}", torch.nn.Parameter(raw)) + + # self.register_constraint(f"raw_group_mats.{i}", Interval(-2 * math.pi, 2 * math.pi)) + + def _rotation(self) -> torch.Tensor: + # Start with identity, then fill block rotations + A = torch.zeros((self.d, self.d), dtype=torch.double, device=getattr(self, f"raw_group_entries_0").device if self.groups else None) + A = A.expand(*self.batch_shape, self.d, self.d).clone() if self.batch_shape else A + + for i, g in enumerate(self.groups): + + row, col = torch.triu_indices(len(g), len(g), 1) + A[..., tuple(torch.tensor(g)[row]), tuple(torch.tensor(g)[col])] = getattr(self, f"raw_group_entries_{i}") + + return torch.linalg.matrix_exp(A - A.transpose(-1, -2)) + + def _transform(self, X: torch.Tensor) -> torch.Tensor: + # X is (..., n, d). Right-multiply by R^T to rotate features. + R = self._rotation() + return torch.matmul(X, R.transpose(-1, -2)) + + def forward(self, x1, x2, diag=False, **params): + return self.base(self._transform(x1), self._transform(x2), diag=diag, **params) + + +# class LatentKernel(gpytorch.kernels.Kernel): +# @property +# def is_stationary(self) -> bool: +# return True + +# num_outputs: int = 1 +# batch_inverse_lengthscale: float = 1e6 + +# def __init__( +# self, +# num_inputs: int = 1, +# skew_dims: bool | Iterable[tuple[int, ...]] = True, +# priors: bool = True, +# scale_output: bool = True, +# **kwargs, +# ) -> None: +# super().__init__() + +# self.num_inputs: int = num_inputs +# self.scale_output: bool = scale_output + +# self.nu: float = kwargs.get("nu", 2.5) + +# if type(skew_dims) is bool: +# if skew_dims: +# self.skew_dims: list[torch.Tensor] = [torch.arange(self.num_inputs)] +# else: +# self.skew_dims = [torch.arange(num_inputs)] +# elif hasattr(skew_dims, "__iter__"): +# self.skew_dims = [torch.tensor(np.atleast_1d(skew_group)) for skew_group in skew_dims] +# else: +# raise ValueError('arg "skew_dims" must be True, False, or an iterable of tuples of ints.') + +# # if not all([len(skew_group) >= 2 for skew_group in self.skew_dims]): +# # raise ValueError("must have at least two dims per skew group") +# skewed_dims = [dim.item() for skew_group in self.skew_dims for dim in skew_group] +# if not len(set(skewed_dims)) == len(skewed_dims): +# raise ValueError("values in skew_dims must be unique") +# if not max(skewed_dims) < self.num_inputs: +# raise ValueError("invalid dimension index in skew_dims") + +# skew_group_submatrix_indices: list[torch.Tensor] = [] +# for dim in range(self.num_outputs): +# for skew_group in self.skew_dims: +# j, k = skew_group[torch.triu_indices(len(skew_group), len(skew_group), 1)].unsqueeze(1) +# i = dim * torch.ones(j.shape).long() +# skew_group_submatrix_indices.append(torch.cat((i, j, k), dim=0)) + +# self.diag_matrix_indices: list[torch.Tensor] = [ +# torch.kron(torch.arange(self.num_outputs), torch.ones(self.num_inputs)).long(), +# *2 * [torch.arange(self.num_inputs).repeat(self.num_outputs)], +# ] + +# self.skew_matrix_indices: tuple[torch.Tensor, ...] = ( +# tuple(torch.cat(skew_group_submatrix_indices, dim=1)) +# if len(skew_group_submatrix_indices) > 0 +# else (torch.tensor([]), torch.tensor([])) +# ) + +# self.n_skew_entries: int = len(self.skew_matrix_indices[0]) + +# lengthscale_constraint = gpytorch.constraints.Positive() +# raw_lengthscales_initial = lengthscale_constraint.inverse_transform(torch.tensor(1e-1)) * torch.ones( +# self.num_outputs, self.num_inputs, dtype=torch.double +# ) + +# self.register_parameter(name="raw_lengthscales", parameter=torch.nn.Parameter(raw_lengthscales_initial)) +# self.register_constraint(param_name="raw_lengthscales", constraint=lengthscale_constraint) + +# if priors: +# self.register_prior( +# name="lengthscales_prior", +# prior=gpytorch.priors.GammaPrior(concentration=3, rate=6), +# param_or_closure=lambda m: m.lengthscales, +# setting_closure=lambda m, v: m._set_lengthscales(v), +# ) + +# if self.n_skew_entries > 0: +# skew_entries_constraint = gpytorch.constraints.Interval(-2 * np.pi, 2 * np.pi) +# skew_entries_initial = torch.zeros((self.num_outputs, self.n_skew_entries), dtype=torch.float64) +# self.register_parameter(name="raw_skew_entries", parameter=torch.nn.Parameter(skew_entries_initial)) +# self.register_constraint(param_name="raw_skew_entries", constraint=skew_entries_constraint) + +# if self.scale_output: +# outputscale_constraint = gpytorch.constraints.Positive() +# outputscale_prior = gpytorch.priors.GammaPrior(concentration=2, rate=0.15) + +# self.register_parameter( +# name="raw_outputscale", +# parameter=torch.nn.Parameter(torch.ones(1, dtype=torch.double)), +# ) + +# self.register_constraint("raw_outputscale", constraint=outputscale_constraint) + +# self.register_prior( +# name="outputscale_prior", +# prior=outputscale_prior, +# param_or_closure=lambda m: m.outputscale, +# setting_closure=lambda m, v: m._set_outputscale(v), +# ) + +# @property +# def lengthscales(self) -> torch.Tensor: +# return self.raw_lengthscales_constraint.transform(self.raw_lengthscales) + +# @lengthscales.setter +# def lengthscales(self, value: torch.Tensor | float | np.ndarray) -> None: +# self._set_lengthscales(value) + +# @property +# def skew_entries(self) -> torch.Tensor: +# return self.raw_skew_entries_constraint.transform(self.raw_skew_entries) + +# @skew_entries.setter +# def skew_entries(self, value: torch.Tensor | float | np.ndarray) -> None: +# self._set_skew_entries(value) + +# @property +# def outputscale(self) -> torch.Tensor: +# return self.raw_outputscale_constraint.transform(self.raw_outputscale) + +# @outputscale.setter +# def outputscale(self, value: torch.Tensor | float | np.ndarray) -> None: +# self._set_outputscale(value) + +# def _set_lengthscales(self, value: torch.Tensor | float | np.ndarray) -> None: +# if not torch.is_tensor(value): +# value = torch.as_tensor(value).to(self.raw_lengthscales) +# self.initialize(raw_lengthscales=self.raw_lengthscales_constraint.inverse_transform(value)) + +# def _set_skew_entries(self, value: torch.Tensor | float | np.ndarray) -> None: +# if not torch.is_tensor(value): +# value = torch.as_tensor(value).to(self.raw_skew_entries) +# self.initialize(raw_skew_entries=self.raw_skew_entries_constraint.inverse_transform(value)) + +# def _set_outputscale(self, value: torch.Tensor | float | np.ndarray) -> None: +# if not torch.is_tensor(value): +# value = torch.as_tensor(value).to(self.raw_outputscale) +# self.initialize(raw_outputscale=self.raw_outputscale_constraint.inverse_transform(value)) + +# @property +# def skew_matrix(self) -> torch.Tensor: +# S = torch.zeros((self.num_outputs, self.num_inputs, self.num_inputs), dtype=torch.float64) +# if self.n_skew_entries > 0: +# # to construct an orthogonal matrix. fun fact: exp(skew(N)) is the generator of SO(N) +# S[self.skew_matrix_indices] = self.skew_entries +# S += -S.transpose(-1, -2) +# return torch.linalg.matrix_exp(S) + +# @property +# def diag_matrix(self) -> torch.Tensor: +# D = torch.zeros((self.num_outputs, self.num_inputs, self.num_inputs), dtype=torch.float64) +# D[self.diag_matrix_indices] = self.lengthscales.ravel() ** -1 +# return D + +# @property +# def latent_transform(self) -> torch.Tensor: +# return torch.matmul(self.diag_matrix, self.skew_matrix) + +# def forward(self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, **params: dict) -> torch.Tensor: +# # adapted from the Matern kernel +# mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)] + +# transform = self.latent_transform.unsqueeze(1) + +# trans_x1 = torch.matmul(transform, (x1 - mean).unsqueeze(-1)).squeeze(-1) +# trans_x2 = torch.matmul(transform, (x2 - mean).unsqueeze(-1)).squeeze(-1) + +# distance = self.covar_dist(trans_x1, trans_x2, diag=diag, **params) + +# if self.num_outputs == 1: +# distance = distance.squeeze(0) + +# outputscale = self.outputscale if self.scale_output else 1.0 + +# # special cases of the Matern function +# if self.nu == 0.5: +# return outputscale * torch.exp(-distance) +# if self.nu == 1.5: +# return outputscale * (1 + distance) * torch.exp(-distance) +# if self.nu == 2.5: +# return outputscale * (1 + distance + distance**2 / 3) * torch.exp(-distance) +# raise ValueError(f"nu = {self.nu} is not supported") diff --git a/src/blop/bayesian/models.py b/src/blop/bayesian/models.py index 06f2b803..2a344f8c 100644 --- a/src/blop/bayesian/models.py +++ b/src/blop/bayesian/models.py @@ -1,37 +1,100 @@ +import warnings from typing import Any -import gpytorch -import torch -from botorch.models.gp_regression import SingleTaskGP -from botorch.models.multitask import MultiTaskGP +import botorch # type: ignore[import-untyped] +import gpytorch # type: ignore[import-untyped] +import torch # type: ignore[import-untyped] +from botorch.models.gp_regression import SingleTaskGP # type: ignore[import-untyped] +from botorch.models.multitask import MultiTaskGP # type: ignore[import-untyped] +from botorch.models.transforms.input import InputTransform, Normalize +from botorch.models.transforms.outcome import OutcomeTransform, Standardize +from botorch.utils.types import DEFAULT, _DefaultType +from gpytorch.likelihoods.likelihood import Likelihood +from torch import Tensor from . import kernels +import gpytorch +from gpytorch.priors import NormalPrior +from gpytorch.means import ConstantMean +from botorch.models.gp_regression import SingleTaskGP + class LatentGP(SingleTaskGP): def __init__( self, train_X: torch.Tensor, train_Y: torch.Tensor, + train_Tvar: torch.Tensor = None, + train_Yvar: Tensor | None = None, + likelihood: Likelihood | None = None, + input_transform: InputTransform | None = None, + outcome_transform: OutcomeTransform | _DefaultType | None = DEFAULT, skew_dims: bool | list[tuple[int, ...]] = True, *args: Any, **kwargs: Any, ) -> None: - # Disable outcome transform to avoid shape mismatches with multi-output kernel - super().__init__(train_X, train_Y, *args, outcome_transform=None, **kwargs) - - self.mean_module = gpytorch.means.ConstantMean(constant_prior=gpytorch.priors.NormalPrior(loc=0, scale=1)) - - self.covar_module = kernels.LatentKernel( - num_inputs=train_X.shape[-1], - num_outputs=train_Y.shape[-1], - skew_dims=skew_dims, - priors=True, - scale=True, - **kwargs, - ) - - self.trained: bool = False + + super().__init__(train_X=train_X, + train_Y=train_Y, + input_transform=input_transform, + outcome_transform=outcome_transform, + *args, **kwargs) + + m = train_Y.shape[-1] + aug_batch_shape = train_X.shape[:-2] + (torch.Size([m]) if m > 1 else torch.Size()) + + self.mean_module = ConstantMean(batch_shape=aug_batch_shape, constant_prior=NormalPrior(0.0, 1.0)) + self.covar_module = kernels.RotatedInputsKernel(d=train_X.shape[-1], batch_shape=aug_batch_shape, skew_dims=skew_dims) + + # return SingleTaskGP( + # train_X=train_X, + # train_Y=train_Y, + # mean_module=mean, + # covar_module=covar, + # outcome_transform=None, # keep if you truly need it + # **kwargs, + # ) + + + +# class LatentGP(SingleTaskGP): +# def __init__( +# self, +# train_X: torch.Tensor, +# train_Y: torch.Tensor, +# train_Tvar: torch.Tensor = None, +# train_Yvar: Tensor | None = None, +# likelihood: Likelihood | None = None, +# input_transform: InputTransform | None = None, +# outcome_transform: OutcomeTransform | _DefaultType | None = DEFAULT, +# skew_dims: bool | list[tuple[int, ...]] = True, +# *args: Any, +# **kwargs: Any, +# ) -> None: + +# *batch_shape, n, d = train_X.shape +# input_transform = input_transform or Normalize(d=d) +# # outcome_transform = outcome_transform or Standardize(batch_shape=batch_shape) + +# super().__init__(train_X=train_X, +# train_Y=train_Y, +# input_transform=input_transform, +# outcome_transform=outcome_transform, +# *args, **kwargs) + +# self.mean_module = gpytorch.means.ConstantMean(constant_prior=gpytorch.priors.NormalPrior(loc=0, scale=1)) + +# self.covar_module = kernels.LatentKernel( +# num_inputs=train_X.shape[-1], +# num_outputs=train_Y.shape[-1], +# skew_dims=skew_dims, +# priors=True, +# scale=True, +# **kwargs, +# ) + +# self.trained: bool = False class MultiTaskLatentGP(MultiTaskGP): diff --git a/src/blop/evaluation/__init__.py b/src/blop/evaluation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/blop/evaluation/test_functions.py b/src/blop/evaluation/test_functions.py new file mode 100644 index 00000000..dcdd1664 --- /dev/null +++ b/src/blop/evaluation/test_functions.py @@ -0,0 +1,3 @@ +from ..protocols import EvaluationFunction + + diff --git a/src/blop/tests/ax/test_generation.py b/src/blop/tests/ax/test_generation.py new file mode 100644 index 00000000..55f71ecb --- /dev/null +++ b/src/blop/tests/ax/test_generation.py @@ -0,0 +1,89 @@ +"""Tests for blop.bayesian.kernels module. + +Tests the public API of LatentKernel - a Matérn kernel with learned affine +transformation using SO(N) parameterization for orthogonal rotations. +""" + +from unittest.mock import MagicMock, patch + +import bluesky.plan_stubs as bps +import pytest +from bluesky.run_engine import RunEngine +from bluesky.utils import plan + +from blop.ax.agent import Agent, QueueserverAgent +from blop.ax.dof import DOFConstraint, RangeDOF +from blop.ax.objective import Objective +from blop.ax.optimizer import AxOptimizer +from blop.callbacks.logger import OptimizationLogger +from blop.protocols import AcquisitionPlan, EvaluationFunction + +from ..conftest import MovableSignal, ReadableSignal +from .test_agent import mock_acquisition_plan, mock_evaluation_function +from ..test_plans import _collect_optimize_events + +import numpy as np + + + + +class TestFunctionEvaluation(): + def __init__(self, tiled_client: Container): + self.tiled_client = tiled_client + + def __call__(self, uid: str, suggestions: list[dict]) -> list[dict]: + run = self.tiled_client[uid] + outcomes = [] + x1_data = run["primary/x1"].read() + x2_data = run["primary/x2"].read() + + for suggestion in suggestions: + # Special key to identify a suggestion + suggestion_id = suggestion["_id"] + x1 = suggestion["x1"] + x2 = suggestion["x2"] + + outcomes.append({"test_function_1": 1 - np.exp(-(x1 - 2 * x2 - 1) ** 2 - 1e-3 * (2 * x1 + x2 - 0.5) ** 2), "_id": suggestion_id, + "test_function_2": 1 - np.exp(-1e-3 * (x1 - 2 * x2 - 1) ** 2 - (2 * x1 + x2 - 0.5) ** 2), "_id": suggestion_id}) + + + + # outcomes.append({"test_function": (x1 ** 2 + x2 - 11) ** 2 + (x1 + x2 ** 2 - 7) ** 2, "_id": suggestion_id}) + + return outcomes + +@pytest.fixture(scope="function") +def RE(): + return RunEngine({}) + + + +def test_optimize(RE, mock_acquisition_plan, mock_evaluation_function): + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + readable = ReadableSignal(name="test_readable") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + constraint = DOFConstraint(constraint="x1 + x2 <= 10", x1=dof1, x2=dof2) + objective = Objective(name="test_objective", minimize=False) + agent = Agent( + sensors=[readable], + dofs=[dof1, dof2], + objectives=[objective], + evaluation_function=mock_evaluation_function, + dof_constraints=[constraint], + acquisition_plan=mock_acquisition_plan, + name="test_experiment", + + ) + + + + callback, events = _collect_optimize_events() + RE.subscribe(callback) + try: + RE(agent.optimize(20)) + except RuntimeError: + ... + finally: + RE.unsubscribe(callback) \ No newline at end of file diff --git a/src/blop/tests/bayesian/test_kernels.py b/src/blop/tests/bayesian/test_kernels.py deleted file mode 100644 index 530f3cd0..00000000 --- a/src/blop/tests/bayesian/test_kernels.py +++ /dev/null @@ -1,137 +0,0 @@ -"""Tests for blop.bayesian.kernels module. - -Tests the public API of LatentKernel - a Matérn kernel with learned affine -transformation using SO(N) parameterization for orthogonal rotations. -""" - -import pytest -import torch - -from blop.bayesian.kernels import LatentKernel - - -def test_latent_kernel_init(): - """Can construct LatentKernel with various skew_dims configurations.""" - # Default (all dims rotate together) - kernel1 = LatentKernel(num_inputs=3) - assert kernel1.num_inputs == 3 - assert kernel1.nu == 2.5 # default - - # Custom skew groups - kernel2 = LatentKernel(num_inputs=4, skew_dims=[(0, 1), (2, 3)]) - assert kernel2.num_inputs == 4 - - # No rotation - kernel3 = LatentKernel(num_inputs=2, skew_dims=False) - assert kernel3.num_inputs == 2 - - # With different Matérn smoothness - kernel4 = LatentKernel(num_inputs=2, nu=1.5) - assert kernel4.nu == 1.5 - - -def test_latent_kernel_forward_shapes(): - """forward() returns correct covariance shapes.""" - kernel = LatentKernel(num_inputs=3) - - x1 = torch.randn(10, 3, dtype=torch.float64) - x2 = torch.randn(5, 3, dtype=torch.float64) - - # Square covariance - cov_square = kernel(x1, x1) - assert cov_square.shape == (10, 10) - - # Rectangular covariance - cov_rect = kernel(x1, x2) - assert cov_rect.shape == (10, 5) - - # Diagonal only - cov_diag = kernel(x1, x1, diag=True) - assert cov_diag.shape == (10,) - - -@pytest.mark.parametrize("nu", [0.5, 1.5, 2.5]) -def test_latent_kernel_matern_variants(nu): - """Works with supported Matérn smoothness values (nu=0.5, 1.5, 2.5).""" - kernel = LatentKernel(num_inputs=2, nu=nu) - x = torch.randn(5, 2, dtype=torch.float64) - - # Should compute without error - cov = kernel(x, x).to_dense() - assert cov.shape == (5, 5) - - -def test_latent_kernel_unsupported_nu_raises(): - """Unsupported nu value raises ValueError.""" - kernel = LatentKernel(num_inputs=2, nu=3.5) - x = torch.randn(5, 2, dtype=torch.float64) - - with pytest.raises(ValueError, match="not supported"): - kernel(x, x).to_dense() - - -def test_latent_kernel_invalid_skew_dims(): - """Invalid skew_dims configurations raise appropriate errors.""" - # Duplicate dimensions across groups - with pytest.raises(ValueError, match="unique"): - LatentKernel(num_inputs=3, skew_dims=[(0, 1), (1, 2)]) - - # Dimension index out of range - with pytest.raises(ValueError, match="invalid dimension"): - LatentKernel(num_inputs=3, skew_dims=[(0, 3)]) - - # Invalid type - with pytest.raises((ValueError, TypeError)): - LatentKernel(num_inputs=3, skew_dims="invalid") - - -def test_latent_kernel_forward_deterministic(): - """Forward call is deterministic for same input.""" - kernel = LatentKernel(num_inputs=2, nu=2.5) - x = torch.tensor([[0.0, 0.0], [0.5, 0.5], [1.0, 1.0]], dtype=torch.float64) - - cov1 = kernel(x, x).to_dense() - cov2 = kernel(x, x).to_dense() - - assert torch.allclose(cov1, cov2, atol=1e-10) - - -def test_latent_kernel_golden_values(): - """Validate kernel output against pre-computed expected values. - - FIXME: This kernel is missing the sqrt(2*nu) scaling factor required by the - standard Matérn formula. The correct Matérn 2.5 formula is: - k(r) = σ² * (1 + √5*r + 5r²/3) * exp(-√5*r) - where r = ||x1 - x2|| / lengthscale. The current implementation omits the - √5 factor (and √3 for nu=1.5), causing correlation to decay more slowly - than a true Matérn kernel. This should be fixed when refactoring to use - native BoTorch/GPyTorch kernels. - """ - kernel = LatentKernel(num_inputs=2, skew_dims=False, nu=2.5, scale_output=True) - kernel.lengthscales = torch.tensor([[1.0, 2.0]]) - kernel.outputscale = torch.tensor([1.5]) - - x = torch.tensor( - [ - [0.0, 0.0], - [1.0, 0.0], - [0.0, 2.0], - ], - dtype=torch.float64, - ) - - # Pre-computed expected output for this configuration - expected = torch.tensor( - [ - [1.5000, 1.2876, 1.2876], - [1.2876, 1.5000, 1.1235], - [1.2876, 1.1235, 1.5000], - ], - dtype=torch.float64, - ) - - actual = kernel(x, x).to_dense() - - assert torch.allclose(actual, expected, atol=1e-4), ( - f"Kernel output doesn't match expected values.\nExpected:\n{expected}\nActual:\n{actual}" - ) diff --git a/src/blop/tests/generation/__init__.py b/src/blop/tests/generation/__init__.py new file mode 100644 index 00000000..e69de29b From 284529c7140caff9154d180695c09d5f0e8ab643 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 14 Apr 2026 02:12:56 -0400 Subject: [PATCH 2/9] add generation nodes --- src/blop/bayesian/kernels.py | 21 +++++---------- src/blop/bayesian/models.py | 39 +++++++++++++-------------- src/blop/evaluation/test_functions.py | 3 --- src/blop/tests/ax/test_generation.py | 37 ++++++++++--------------- 4 files changed, 40 insertions(+), 60 deletions(-) diff --git a/src/blop/bayesian/kernels.py b/src/blop/bayesian/kernels.py index add4abb1..17344427 100644 --- a/src/blop/bayesian/kernels.py +++ b/src/blop/bayesian/kernels.py @@ -1,22 +1,12 @@ -from collections.abc import Iterable -import gpytorch -import gpytorch.constraints -import numpy as np import torch - - - -import math -import torch -import gpytorch -from gpytorch.kernels import Kernel -from gpytorch.constraints import Interval from botorch.models.utils.gpytorch_modules import get_matern_kernel_with_gamma_prior +from gpytorch.kernels import Kernel class RotatedInputsKernel(Kernel): """Applies a learned orthogonal rotation (blockwise if desired), then delegates to a base kernel.""" + is_stationary = True has_lengthscale = False # lengthscale lives in the base kernel @@ -49,11 +39,14 @@ def __init__( def _rotation(self) -> torch.Tensor: # Start with identity, then fill block rotations - A = torch.zeros((self.d, self.d), dtype=torch.double, device=getattr(self, f"raw_group_entries_0").device if self.groups else None) + A = torch.zeros( + (self.d, self.d), + dtype=torch.double, + device=self.raw_group_entries_0.device if self.groups else None, + ) A = A.expand(*self.batch_shape, self.d, self.d).clone() if self.batch_shape else A for i, g in enumerate(self.groups): - row, col = torch.triu_indices(len(g), len(g), 1) A[..., tuple(torch.tensor(g)[row]), tuple(torch.tensor(g)[col])] = getattr(self, f"raw_group_entries_{i}") diff --git a/src/blop/bayesian/models.py b/src/blop/bayesian/models.py index 2a344f8c..3e509465 100644 --- a/src/blop/bayesian/models.py +++ b/src/blop/bayesian/models.py @@ -1,24 +1,19 @@ -import warnings from typing import Any -import botorch # type: ignore[import-untyped] import gpytorch # type: ignore[import-untyped] import torch # type: ignore[import-untyped] from botorch.models.gp_regression import SingleTaskGP # type: ignore[import-untyped] from botorch.models.multitask import MultiTaskGP # type: ignore[import-untyped] -from botorch.models.transforms.input import InputTransform, Normalize -from botorch.models.transforms.outcome import OutcomeTransform, Standardize +from botorch.models.transforms.input import InputTransform +from botorch.models.transforms.outcome import OutcomeTransform from botorch.utils.types import DEFAULT, _DefaultType from gpytorch.likelihoods.likelihood import Likelihood +from gpytorch.means import ConstantMean +from gpytorch.priors import NormalPrior from torch import Tensor from . import kernels -import gpytorch -from gpytorch.priors import NormalPrior -from gpytorch.means import ConstantMean -from botorch.models.gp_regression import SingleTaskGP - class LatentGP(SingleTaskGP): def __init__( @@ -34,18 +29,23 @@ def __init__( *args: Any, **kwargs: Any, ) -> None: - - super().__init__(train_X=train_X, - train_Y=train_Y, - input_transform=input_transform, - outcome_transform=outcome_transform, - *args, **kwargs) + + super().__init__( + train_X=train_X, + train_Y=train_Y, + input_transform=input_transform, + outcome_transform=outcome_transform, + *args, + **kwargs, + ) m = train_Y.shape[-1] aug_batch_shape = train_X.shape[:-2] + (torch.Size([m]) if m > 1 else torch.Size()) self.mean_module = ConstantMean(batch_shape=aug_batch_shape, constant_prior=NormalPrior(0.0, 1.0)) - self.covar_module = kernels.RotatedInputsKernel(d=train_X.shape[-1], batch_shape=aug_batch_shape, skew_dims=skew_dims) + self.covar_module = kernels.RotatedInputsKernel( + d=train_X.shape[-1], batch_shape=aug_batch_shape, skew_dims=skew_dims + ) # return SingleTaskGP( # train_X=train_X, @@ -57,7 +57,6 @@ def __init__( # ) - # class LatentGP(SingleTaskGP): # def __init__( # self, @@ -72,14 +71,14 @@ def __init__( # *args: Any, # **kwargs: Any, # ) -> None: - + # *batch_shape, n, d = train_X.shape # input_transform = input_transform or Normalize(d=d) # # outcome_transform = outcome_transform or Standardize(batch_shape=batch_shape) # super().__init__(train_X=train_X, -# train_Y=train_Y, -# input_transform=input_transform, +# train_Y=train_Y, +# input_transform=input_transform, # outcome_transform=outcome_transform, # *args, **kwargs) diff --git a/src/blop/evaluation/test_functions.py b/src/blop/evaluation/test_functions.py index dcdd1664..e69de29b 100644 --- a/src/blop/evaluation/test_functions.py +++ b/src/blop/evaluation/test_functions.py @@ -1,3 +0,0 @@ -from ..protocols import EvaluationFunction - - diff --git a/src/blop/tests/ax/test_generation.py b/src/blop/tests/ax/test_generation.py index 55f71ecb..afe29f6f 100644 --- a/src/blop/tests/ax/test_generation.py +++ b/src/blop/tests/ax/test_generation.py @@ -4,30 +4,20 @@ transformation using SO(N) parameterization for orthogonal rotations. """ -from unittest.mock import MagicMock, patch -import bluesky.plan_stubs as bps +import numpy as np import pytest from bluesky.run_engine import RunEngine -from bluesky.utils import plan -from blop.ax.agent import Agent, QueueserverAgent +from blop.ax.agent import Agent from blop.ax.dof import DOFConstraint, RangeDOF from blop.ax.objective import Objective -from blop.ax.optimizer import AxOptimizer -from blop.callbacks.logger import OptimizationLogger -from blop.protocols import AcquisitionPlan, EvaluationFunction from ..conftest import MovableSignal, ReadableSignal -from .test_agent import mock_acquisition_plan, mock_evaluation_function from ..test_plans import _collect_optimize_events -import numpy as np - - - -class TestFunctionEvaluation(): +class TestFunctionEvaluation: def __init__(self, tiled_client: Container): self.tiled_client = tiled_client @@ -42,22 +32,26 @@ def __call__(self, uid: str, suggestions: list[dict]) -> list[dict]: suggestion_id = suggestion["_id"] x1 = suggestion["x1"] x2 = suggestion["x2"] - - outcomes.append({"test_function_1": 1 - np.exp(-(x1 - 2 * x2 - 1) ** 2 - 1e-3 * (2 * x1 + x2 - 0.5) ** 2), "_id": suggestion_id, - "test_function_2": 1 - np.exp(-1e-3 * (x1 - 2 * x2 - 1) ** 2 - (2 * x1 + x2 - 0.5) ** 2), "_id": suggestion_id}) + outcomes.append( + { + "test_function_1": 1 - np.exp(-((x1 - 2 * x2 - 1) ** 2) - 1e-3 * (2 * x1 + x2 - 0.5) ** 2), + "_id": suggestion_id, + "test_function_2": 1 - np.exp(-1e-3 * (x1 - 2 * x2 - 1) ** 2 - (2 * x1 + x2 - 0.5) ** 2), + "_id": suggestion_id, + } + ) - # outcomes.append({"test_function": (x1 ** 2 + x2 - 11) ** 2 + (x1 + x2 ** 2 - 7) ** 2, "_id": suggestion_id}) - + return outcomes + @pytest.fixture(scope="function") def RE(): return RunEngine({}) - def test_optimize(RE, mock_acquisition_plan, mock_evaluation_function): movable1 = MovableSignal(name="test_movable1") movable2 = MovableSignal(name="test_movable2") @@ -74,11 +68,8 @@ def test_optimize(RE, mock_acquisition_plan, mock_evaluation_function): dof_constraints=[constraint], acquisition_plan=mock_acquisition_plan, name="test_experiment", - ) - - callback, events = _collect_optimize_events() RE.subscribe(callback) try: @@ -86,4 +77,4 @@ def test_optimize(RE, mock_acquisition_plan, mock_evaluation_function): except RuntimeError: ... finally: - RE.unsubscribe(callback) \ No newline at end of file + RE.unsubscribe(callback) From 03395c4d698e4b3dedb88342abe15da077c53b26 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 14 Apr 2026 02:14:29 -0400 Subject: [PATCH 3/9] lol --- src/blop/ax/generation/__init__.py | 0 src/blop/ax/generation/beamline.py | 8 +++++ src/blop/ax/generation/nodes.py | 55 ++++++++++++++++++++++++++++++ 3 files changed, 63 insertions(+) create mode 100644 src/blop/ax/generation/__init__.py create mode 100644 src/blop/ax/generation/beamline.py create mode 100644 src/blop/ax/generation/nodes.py diff --git a/src/blop/ax/generation/__init__.py b/src/blop/ax/generation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/blop/ax/generation/beamline.py b/src/blop/ax/generation/beamline.py new file mode 100644 index 00000000..04036bdf --- /dev/null +++ b/src/blop/ax/generation/beamline.py @@ -0,0 +1,8 @@ + +from ax.generation_strategy.generation_strategy import GenerationStrategy +from .nodes import sobol_node, latent_gp_node + +bl_generation_strategy = GenerationStrategy( + name="Custom Generation Strategy", + nodes=[sobol_node, latent_gp_node] +) diff --git a/src/blop/ax/generation/nodes.py b/src/blop/ax/generation/nodes.py new file mode 100644 index 00000000..6e56e473 --- /dev/null +++ b/src/blop/ax/generation/nodes.py @@ -0,0 +1,55 @@ +from botorch.models.transforms.input import Normalize +from botorch.models.transforms.outcome import Standardize + +from ax.generation_strategy.generation_node import GenerationNode +from ax.generation_strategy.generator_spec import GeneratorSpec +from ax.generation_strategy.transition_criterion import MinTrials +from ax.adapter.registry import Generators +from ax.generators.torch.botorch_modular.surrogate import ModelConfig, SurrogateSpec +from botorch.acquisition.logei import qLogNoisyExpectedImprovement + +from blop.bayesian.models import LatentGP + + +sobol_node = GenerationNode( + name="Sobol", + generator_specs=[ + GeneratorSpec(generator_enum=Generators.SOBOL, model_kwargs={"seed": 0}), + ], + transition_criteria=[ + MinTrials( + threshold=16, + transition_to="LatentGP", + use_all_trials_in_exp=True, + ), + ], + ) + +latent_gp_node = GenerationNode( + name="LatentGP", + generator_specs=[ + GeneratorSpec( + generator_enum=Generators.BOTORCH_MODULAR, + model_kwargs={ + "surrogate_spec": SurrogateSpec( + model_configs=[ + ModelConfig( + botorch_model_class=LatentGP, + input_transform_classes=[Normalize], + # model_options={"skew_dims": True}, + outcome_transform_classes=[Standardize], + ), + ], + ), + "botorch_acqf_class": qLogNoisyExpectedImprovement, + "acquisition_options": {}, + }, + model_gen_kwargs={ + "optimizer_kwargs": { + "num_restarts": 10, + "sequential": True, + }, + }, + ), + ], + ) From d0d6e37e69d379ae536898966e52505f5225b97d Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 7 May 2026 17:36:00 -0400 Subject: [PATCH 4/9] add generation submodule --- src/blop/ax/generation/__init__.py | 9 +++++++++ src/blop/ax/generation/beamline.py | 9 +++------ src/blop/tests/ax/test_generation.py | 8 ++------ 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src/blop/ax/generation/__init__.py b/src/blop/ax/generation/__init__.py index e69de29b..1d92edc3 100644 --- a/src/blop/ax/generation/__init__.py +++ b/src/blop/ax/generation/__init__.py @@ -0,0 +1,9 @@ +from ax.generation_strategy.generation_strategy import GenerationStrategy + +from .nodes import latent_gp_node, sobol_node + + +def get_generation_strategy(name): + + if name == "beamline": + return GenerationStrategy(name="Custom Generation Strategy", nodes=[sobol_node, latent_gp_node]) diff --git a/src/blop/ax/generation/beamline.py b/src/blop/ax/generation/beamline.py index 04036bdf..2ec05023 100644 --- a/src/blop/ax/generation/beamline.py +++ b/src/blop/ax/generation/beamline.py @@ -1,8 +1,5 @@ - from ax.generation_strategy.generation_strategy import GenerationStrategy -from .nodes import sobol_node, latent_gp_node -bl_generation_strategy = GenerationStrategy( - name="Custom Generation Strategy", - nodes=[sobol_node, latent_gp_node] -) +from .nodes import latent_gp_node, sobol_node + +beamline_generation_strategy = GenerationStrategy(name="Custom Generation Strategy", nodes=[sobol_node, latent_gp_node]) diff --git a/src/blop/tests/ax/test_generation.py b/src/blop/tests/ax/test_generation.py index afe29f6f..af0d6578 100644 --- a/src/blop/tests/ax/test_generation.py +++ b/src/blop/tests/ax/test_generation.py @@ -4,7 +4,6 @@ transformation using SO(N) parameterization for orthogonal rotations. """ - import numpy as np import pytest from bluesky.run_engine import RunEngine @@ -18,14 +17,12 @@ class TestFunctionEvaluation: - def __init__(self, tiled_client: Container): + def __init__(self, tiled_client): self.tiled_client = tiled_client def __call__(self, uid: str, suggestions: list[dict]) -> list[dict]: - run = self.tiled_client[uid] + outcomes = [] - x1_data = run["primary/x1"].read() - x2_data = run["primary/x2"].read() for suggestion in suggestions: # Special key to identify a suggestion @@ -36,7 +33,6 @@ def __call__(self, uid: str, suggestions: list[dict]) -> list[dict]: outcomes.append( { "test_function_1": 1 - np.exp(-((x1 - 2 * x2 - 1) ** 2) - 1e-3 * (2 * x1 + x2 - 0.5) ** 2), - "_id": suggestion_id, "test_function_2": 1 - np.exp(-1e-3 * (x1 - 2 * x2 - 1) ** 2 - (2 * x1 + x2 - 0.5) ** 2), "_id": suggestion_id, } From ccb030143a0e77d5cda6eed3b9f6498e5b2fb251 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 7 May 2026 18:58:00 -0400 Subject: [PATCH 5/9] linting fixes --- src/blop/ax/agent.py | 3 +- src/blop/ax/generation/nodes.py | 85 ++++++++++++++++----------------- src/blop/bayesian/kernels.py | 3 +- src/blop/bayesian/models.py | 2 +- 4 files changed, 46 insertions(+), 47 deletions(-) diff --git a/src/blop/ax/agent.py b/src/blop/ax/agent.py index d414826f..85703a43 100644 --- a/src/blop/ax/agent.py +++ b/src/blop/ax/agent.py @@ -6,6 +6,7 @@ from ax import Client from ax.analysis.plotly.surface.contour import ContourPlot from ax.core.types import TParamValue +from ax.generation_strategy.generation_strategy import GenerationStrategy # =============================== # TODO: Remove when Python 3.10 is no longer supported @@ -242,7 +243,7 @@ def __init__( dof_constraints: Sequence[DOFConstraint] | None = None, outcome_constraints: Sequence[OutcomeConstraint] | None = None, checkpoint_path: str | None = None, - generation_strategy: GenerationStrategy = None, + generation_strategy: GenerationStrategy | None = None, **kwargs: Any, ): if any(isinstance(dof.actuator, str) for dof in dofs): diff --git a/src/blop/ax/generation/nodes.py b/src/blop/ax/generation/nodes.py index 6e56e473..a58636a6 100644 --- a/src/blop/ax/generation/nodes.py +++ b/src/blop/ax/generation/nodes.py @@ -1,55 +1,54 @@ -from botorch.models.transforms.input import Normalize -from botorch.models.transforms.outcome import Standardize - +from ax.adapter.registry import Generators from ax.generation_strategy.generation_node import GenerationNode from ax.generation_strategy.generator_spec import GeneratorSpec from ax.generation_strategy.transition_criterion import MinTrials -from ax.adapter.registry import Generators -from ax.generators.torch.botorch_modular.surrogate import ModelConfig, SurrogateSpec +from ax.generators.torch.botorch_modular.surrogate import SurrogateSpec +from ax.generators.torch.botorch_modular.utils import ModelConfig from botorch.acquisition.logei import qLogNoisyExpectedImprovement +from botorch.models.transforms.input import Normalize +from botorch.models.transforms.outcome import Standardize from blop.bayesian.models import LatentGP - sobol_node = GenerationNode( - name="Sobol", - generator_specs=[ - GeneratorSpec(generator_enum=Generators.SOBOL, model_kwargs={"seed": 0}), - ], - transition_criteria=[ - MinTrials( - threshold=16, - transition_to="LatentGP", - use_all_trials_in_exp=True, - ), - ], - ) + name="Sobol", + generator_specs=[ + GeneratorSpec(generator_enum=Generators.SOBOL, model_kwargs={"seed": 0}), + ], + transition_criteria=[ + MinTrials( + threshold=16, + transition_to="LatentGP", + use_all_trials_in_exp=True, + ), + ], +) latent_gp_node = GenerationNode( - name="LatentGP", - generator_specs=[ - GeneratorSpec( - generator_enum=Generators.BOTORCH_MODULAR, - model_kwargs={ - "surrogate_spec": SurrogateSpec( - model_configs=[ - ModelConfig( - botorch_model_class=LatentGP, - input_transform_classes=[Normalize], - # model_options={"skew_dims": True}, - outcome_transform_classes=[Standardize], - ), - ], + name="LatentGP", + generator_specs=[ + GeneratorSpec( + generator_enum=Generators.BOTORCH_MODULAR, + model_kwargs={ + "surrogate_spec": SurrogateSpec( + model_configs=[ + ModelConfig( + botorch_model_class=LatentGP, + input_transform_classes=[Normalize], + # model_options={"skew_dims": True}, + outcome_transform_classes=[Standardize], ), - "botorch_acqf_class": qLogNoisyExpectedImprovement, - "acquisition_options": {}, - }, - model_gen_kwargs={ - "optimizer_kwargs": { - "num_restarts": 10, - "sequential": True, - }, - }, + ], ), - ], - ) + "botorch_acqf_class": qLogNoisyExpectedImprovement, + "acquisition_options": {}, + }, + model_gen_kwargs={ + "optimizer_kwargs": { + "num_restarts": 10, + "sequential": True, + }, + }, + ), + ], +) diff --git a/src/blop/bayesian/kernels.py b/src/blop/bayesian/kernels.py index 17344427..89bf1837 100644 --- a/src/blop/bayesian/kernels.py +++ b/src/blop/bayesian/kernels.py @@ -1,4 +1,3 @@ - import torch from botorch.models.utils.gpytorch_modules import get_matern_kernel_with_gamma_prior from gpytorch.kernels import Kernel @@ -13,7 +12,7 @@ class RotatedInputsKernel(Kernel): def __init__( self, d: int, - batch_shape: torch.Size = torch.Size(), + batch_shape: torch.Size, skew_dims: bool | list[tuple[int, ...]] = True, ): super().__init__(batch_shape=batch_shape) diff --git a/src/blop/bayesian/models.py b/src/blop/bayesian/models.py index 3e509465..23a287ea 100644 --- a/src/blop/bayesian/models.py +++ b/src/blop/bayesian/models.py @@ -35,7 +35,7 @@ def __init__( train_Y=train_Y, input_transform=input_transform, outcome_transform=outcome_transform, - *args, + *args, # noqa **kwargs, ) From 6d6f4a8e5e1591409b7ff4e29057423777ec33f2 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 7 May 2026 19:35:50 -0400 Subject: [PATCH 6/9] fix generation test --- src/blop/bayesian/models.py | 15 ++---- src/blop/evaluation/test_functions.py | 24 ++++++++++ src/blop/tests/ax/test_generation.py | 68 +++++++++------------------ 3 files changed, 52 insertions(+), 55 deletions(-) diff --git a/src/blop/bayesian/models.py b/src/blop/bayesian/models.py index 23a287ea..324e3a8e 100644 --- a/src/blop/bayesian/models.py +++ b/src/blop/bayesian/models.py @@ -1,6 +1,5 @@ from typing import Any -import gpytorch # type: ignore[import-untyped] import torch # type: ignore[import-untyped] from botorch.models.gp_regression import SingleTaskGP # type: ignore[import-untyped] from botorch.models.multitask import MultiTaskGP # type: ignore[import-untyped] @@ -108,18 +107,14 @@ def __init__( ) -> None: super().__init__(train_X, train_Y, task_feature, *args, **kwargs) - self.mean_module = gpytorch.means.ConstantMean(constant_prior=gpytorch.priors.NormalPrior(loc=0, scale=1)) + m = train_Y.shape[-1] + aug_batch_shape = train_X.shape[:-2] + (torch.Size([m]) if m > 1 else torch.Size()) - self.covar_module = kernels.LatentKernel( - num_inputs=self.num_non_task_features, - skew_dims=skew_dims, - priors=True, - scale=True, - **kwargs, + self.mean_module = ConstantMean(batch_shape=aug_batch_shape, constant_prior=NormalPrior(0.0, 1.0)) + self.covar_module = kernels.RotatedInputsKernel( + d=train_X.shape[-1], batch_shape=aug_batch_shape, skew_dims=skew_dims ) - self.trained: bool = False - class LatentConstraintModel(LatentGP): def __init__( diff --git a/src/blop/evaluation/test_functions.py b/src/blop/evaluation/test_functions.py index e69de29b..6dc792e7 100644 --- a/src/blop/evaluation/test_functions.py +++ b/src/blop/evaluation/test_functions.py @@ -0,0 +1,24 @@ +import numpy as np + + +class TestFunctionEvaluation: + def __call__(self, uid: str, suggestions: list[dict]) -> list[dict]: + + outcomes = [] + + for suggestion in suggestions: + suggestion_id = suggestion["_id"] + x1 = suggestion["x1"] + x2 = suggestion["x2"] + + fitness = np.exp(-((x1 - 2 * x2 - 1) ** 2) - 1e-3 * (2 * x1 + x2 - 0.5) ** 2) + fitness += 1e-2 * np.random.standard_normal() + + outcomes.append( + { + "_id": suggestion_id, + "fitness": fitness, + } + ) + + return outcomes diff --git a/src/blop/tests/ax/test_generation.py b/src/blop/tests/ax/test_generation.py index af0d6578..ad9e0520 100644 --- a/src/blop/tests/ax/test_generation.py +++ b/src/blop/tests/ax/test_generation.py @@ -4,66 +4,44 @@ transformation using SO(N) parameterization for orthogonal rotations. """ -import numpy as np import pytest from bluesky.run_engine import RunEngine from blop.ax.agent import Agent -from blop.ax.dof import DOFConstraint, RangeDOF +from blop.ax.dof import RangeDOF +from blop.ax.generation import get_generation_strategy from blop.ax.objective import Objective +from blop.evaluation.test_functions import TestFunctionEvaluation -from ..conftest import MovableSignal, ReadableSignal +from ..conftest import MovableSignal from ..test_plans import _collect_optimize_events -class TestFunctionEvaluation: - def __init__(self, tiled_client): - self.tiled_client = tiled_client - - def __call__(self, uid: str, suggestions: list[dict]) -> list[dict]: - - outcomes = [] - - for suggestion in suggestions: - # Special key to identify a suggestion - suggestion_id = suggestion["_id"] - x1 = suggestion["x1"] - x2 = suggestion["x2"] - - outcomes.append( - { - "test_function_1": 1 - np.exp(-((x1 - 2 * x2 - 1) ** 2) - 1e-3 * (2 * x1 + x2 - 0.5) ** 2), - "test_function_2": 1 - np.exp(-1e-3 * (x1 - 2 * x2 - 1) ** 2 - (2 * x1 + x2 - 0.5) ** 2), - "_id": suggestion_id, - } - ) - - # outcomes.append({"test_function": (x1 ** 2 + x2 - 11) ** 2 + (x1 + x2 ** 2 - 7) ** 2, "_id": suggestion_id}) - - return outcomes - - @pytest.fixture(scope="function") def RE(): return RunEngine({}) -def test_optimize(RE, mock_acquisition_plan, mock_evaluation_function): - movable1 = MovableSignal(name="test_movable1") - movable2 = MovableSignal(name="test_movable2") - readable = ReadableSignal(name="test_readable") - dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") - dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") - constraint = DOFConstraint(constraint="x1 + x2 <= 10", x1=dof1, x2=dof2) - objective = Objective(name="test_objective", minimize=False) +def test_optimize(RE): + + dofs = [ + RangeDOF(actuator=MovableSignal("x1"), bounds=(-2.0, 2.0), parameter_type="float"), + RangeDOF(actuator=MovableSignal("x2"), bounds=(-2.0, 2.0), parameter_type="float"), + ] + + objectives = [ + Objective(name="fitness", minimize=False), + ] + agent = Agent( - sensors=[readable], - dofs=[dof1, dof2], - objectives=[objective], - evaluation_function=mock_evaluation_function, - dof_constraints=[constraint], - acquisition_plan=mock_acquisition_plan, - name="test_experiment", + sensors=[], + dofs=dofs, + objectives=objectives, + evaluation_function=TestFunctionEvaluation(), + name="test", + description="Optimization on a test function", + experiment_type="demo", + generation_strategy=get_generation_strategy("beamline"), ) callback, events = _collect_optimize_events() From 5060217a219e40ab92232e398daed5c1cdf5c79c Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 11 May 2026 15:27:52 -0400 Subject: [PATCH 7/9] add more tests --- src/blop/{ax => }/generation/__init__.py | 5 +++++ src/blop/{ax => }/generation/beamline.py | 0 src/blop/{ax => }/generation/nodes.py | 0 src/blop/tests/ax/test_generation.py | 7 ++++++- 4 files changed, 11 insertions(+), 1 deletion(-) rename src/blop/{ax => }/generation/__init__.py (63%) rename src/blop/{ax => }/generation/beamline.py (100%) rename src/blop/{ax => }/generation/nodes.py (100%) diff --git a/src/blop/ax/generation/__init__.py b/src/blop/generation/__init__.py similarity index 63% rename from src/blop/ax/generation/__init__.py rename to src/blop/generation/__init__.py index 1d92edc3..e6fe05d4 100644 --- a/src/blop/ax/generation/__init__.py +++ b/src/blop/generation/__init__.py @@ -2,8 +2,13 @@ from .nodes import latent_gp_node, sobol_node +all_strategies = ["beamline"] + def get_generation_strategy(name): + if name not in all_strategies: + raise ValueError(f"Invalid strategy '{name}', valid strategies are one of {all_strategies}") + if name == "beamline": return GenerationStrategy(name="Custom Generation Strategy", nodes=[sobol_node, latent_gp_node]) diff --git a/src/blop/ax/generation/beamline.py b/src/blop/generation/beamline.py similarity index 100% rename from src/blop/ax/generation/beamline.py rename to src/blop/generation/beamline.py diff --git a/src/blop/ax/generation/nodes.py b/src/blop/generation/nodes.py similarity index 100% rename from src/blop/ax/generation/nodes.py rename to src/blop/generation/nodes.py diff --git a/src/blop/tests/ax/test_generation.py b/src/blop/tests/ax/test_generation.py index ad9e0520..dabb8a5d 100644 --- a/src/blop/tests/ax/test_generation.py +++ b/src/blop/tests/ax/test_generation.py @@ -9,9 +9,9 @@ from blop.ax.agent import Agent from blop.ax.dof import RangeDOF -from blop.ax.generation import get_generation_strategy from blop.ax.objective import Objective from blop.evaluation.test_functions import TestFunctionEvaluation +from blop.generation import all_strategies, get_generation_strategy from ..conftest import MovableSignal from ..test_plans import _collect_optimize_events @@ -22,6 +22,11 @@ def RE(): return RunEngine({}) +@pytest.mark.parametrize("name", all_strategies) +def test_get_generation_strategy(name): + get_generation_strategy(name) + + def test_optimize(RE): dofs = [ From b773653b2020bc95f1e88b7cad6fbcbc5a360302 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 11 May 2026 17:46:01 -0400 Subject: [PATCH 8/9] even more tests --- src/blop/generation/__init__.py | 11 ++++++----- src/blop/tests/{ax => generation}/test_generation.py | 11 ++++++++++- 2 files changed, 16 insertions(+), 6 deletions(-) rename src/blop/tests/{ax => generation}/test_generation.py (84%) diff --git a/src/blop/generation/__init__.py b/src/blop/generation/__init__.py index e6fe05d4..bfaeeadf 100644 --- a/src/blop/generation/__init__.py +++ b/src/blop/generation/__init__.py @@ -1,14 +1,15 @@ -from ax.generation_strategy.generation_strategy import GenerationStrategy - -from .nodes import latent_gp_node, sobol_node +from .beamline import beamline_generation_strategy all_strategies = ["beamline"] +class InvalidStrategyError(Exception): ... + + def get_generation_strategy(name): if name not in all_strategies: - raise ValueError(f"Invalid strategy '{name}', valid strategies are one of {all_strategies}") + raise InvalidStrategyError(f"Invalid strategy '{name}', valid strategies are one of {all_strategies}") if name == "beamline": - return GenerationStrategy(name="Custom Generation Strategy", nodes=[sobol_node, latent_gp_node]) + return beamline_generation_strategy diff --git a/src/blop/tests/ax/test_generation.py b/src/blop/tests/generation/test_generation.py similarity index 84% rename from src/blop/tests/ax/test_generation.py rename to src/blop/tests/generation/test_generation.py index dabb8a5d..35e38375 100644 --- a/src/blop/tests/ax/test_generation.py +++ b/src/blop/tests/generation/test_generation.py @@ -11,7 +11,7 @@ from blop.ax.dof import RangeDOF from blop.ax.objective import Objective from blop.evaluation.test_functions import TestFunctionEvaluation -from blop.generation import all_strategies, get_generation_strategy +from blop.generation import InvalidStrategyError, all_strategies, get_generation_strategy from ..conftest import MovableSignal from ..test_plans import _collect_optimize_events @@ -27,6 +27,15 @@ def test_get_generation_strategy(name): get_generation_strategy(name) +def test_invalid_strategy(): + try: + get_generation_strategy("invalid_strategy") + except InvalidStrategyError: + pass + except Exception as error: + raise error + + def test_optimize(RE): dofs = [ From 2d2965aab4bef010e78953d0e6ad3e0c4cc8634f Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 11 May 2026 18:06:05 -0400 Subject: [PATCH 9/9] remove old code --- src/blop/bayesian/kernels.py | 185 ----------------------------------- src/blop/bayesian/models.py | 48 --------- 2 files changed, 233 deletions(-) diff --git a/src/blop/bayesian/kernels.py b/src/blop/bayesian/kernels.py index 89bf1837..4c76c1ea 100644 --- a/src/blop/bayesian/kernels.py +++ b/src/blop/bayesian/kernels.py @@ -58,188 +58,3 @@ def _transform(self, X: torch.Tensor) -> torch.Tensor: def forward(self, x1, x2, diag=False, **params): return self.base(self._transform(x1), self._transform(x2), diag=diag, **params) - - -# class LatentKernel(gpytorch.kernels.Kernel): -# @property -# def is_stationary(self) -> bool: -# return True - -# num_outputs: int = 1 -# batch_inverse_lengthscale: float = 1e6 - -# def __init__( -# self, -# num_inputs: int = 1, -# skew_dims: bool | Iterable[tuple[int, ...]] = True, -# priors: bool = True, -# scale_output: bool = True, -# **kwargs, -# ) -> None: -# super().__init__() - -# self.num_inputs: int = num_inputs -# self.scale_output: bool = scale_output - -# self.nu: float = kwargs.get("nu", 2.5) - -# if type(skew_dims) is bool: -# if skew_dims: -# self.skew_dims: list[torch.Tensor] = [torch.arange(self.num_inputs)] -# else: -# self.skew_dims = [torch.arange(num_inputs)] -# elif hasattr(skew_dims, "__iter__"): -# self.skew_dims = [torch.tensor(np.atleast_1d(skew_group)) for skew_group in skew_dims] -# else: -# raise ValueError('arg "skew_dims" must be True, False, or an iterable of tuples of ints.') - -# # if not all([len(skew_group) >= 2 for skew_group in self.skew_dims]): -# # raise ValueError("must have at least two dims per skew group") -# skewed_dims = [dim.item() for skew_group in self.skew_dims for dim in skew_group] -# if not len(set(skewed_dims)) == len(skewed_dims): -# raise ValueError("values in skew_dims must be unique") -# if not max(skewed_dims) < self.num_inputs: -# raise ValueError("invalid dimension index in skew_dims") - -# skew_group_submatrix_indices: list[torch.Tensor] = [] -# for dim in range(self.num_outputs): -# for skew_group in self.skew_dims: -# j, k = skew_group[torch.triu_indices(len(skew_group), len(skew_group), 1)].unsqueeze(1) -# i = dim * torch.ones(j.shape).long() -# skew_group_submatrix_indices.append(torch.cat((i, j, k), dim=0)) - -# self.diag_matrix_indices: list[torch.Tensor] = [ -# torch.kron(torch.arange(self.num_outputs), torch.ones(self.num_inputs)).long(), -# *2 * [torch.arange(self.num_inputs).repeat(self.num_outputs)], -# ] - -# self.skew_matrix_indices: tuple[torch.Tensor, ...] = ( -# tuple(torch.cat(skew_group_submatrix_indices, dim=1)) -# if len(skew_group_submatrix_indices) > 0 -# else (torch.tensor([]), torch.tensor([])) -# ) - -# self.n_skew_entries: int = len(self.skew_matrix_indices[0]) - -# lengthscale_constraint = gpytorch.constraints.Positive() -# raw_lengthscales_initial = lengthscale_constraint.inverse_transform(torch.tensor(1e-1)) * torch.ones( -# self.num_outputs, self.num_inputs, dtype=torch.double -# ) - -# self.register_parameter(name="raw_lengthscales", parameter=torch.nn.Parameter(raw_lengthscales_initial)) -# self.register_constraint(param_name="raw_lengthscales", constraint=lengthscale_constraint) - -# if priors: -# self.register_prior( -# name="lengthscales_prior", -# prior=gpytorch.priors.GammaPrior(concentration=3, rate=6), -# param_or_closure=lambda m: m.lengthscales, -# setting_closure=lambda m, v: m._set_lengthscales(v), -# ) - -# if self.n_skew_entries > 0: -# skew_entries_constraint = gpytorch.constraints.Interval(-2 * np.pi, 2 * np.pi) -# skew_entries_initial = torch.zeros((self.num_outputs, self.n_skew_entries), dtype=torch.float64) -# self.register_parameter(name="raw_skew_entries", parameter=torch.nn.Parameter(skew_entries_initial)) -# self.register_constraint(param_name="raw_skew_entries", constraint=skew_entries_constraint) - -# if self.scale_output: -# outputscale_constraint = gpytorch.constraints.Positive() -# outputscale_prior = gpytorch.priors.GammaPrior(concentration=2, rate=0.15) - -# self.register_parameter( -# name="raw_outputscale", -# parameter=torch.nn.Parameter(torch.ones(1, dtype=torch.double)), -# ) - -# self.register_constraint("raw_outputscale", constraint=outputscale_constraint) - -# self.register_prior( -# name="outputscale_prior", -# prior=outputscale_prior, -# param_or_closure=lambda m: m.outputscale, -# setting_closure=lambda m, v: m._set_outputscale(v), -# ) - -# @property -# def lengthscales(self) -> torch.Tensor: -# return self.raw_lengthscales_constraint.transform(self.raw_lengthscales) - -# @lengthscales.setter -# def lengthscales(self, value: torch.Tensor | float | np.ndarray) -> None: -# self._set_lengthscales(value) - -# @property -# def skew_entries(self) -> torch.Tensor: -# return self.raw_skew_entries_constraint.transform(self.raw_skew_entries) - -# @skew_entries.setter -# def skew_entries(self, value: torch.Tensor | float | np.ndarray) -> None: -# self._set_skew_entries(value) - -# @property -# def outputscale(self) -> torch.Tensor: -# return self.raw_outputscale_constraint.transform(self.raw_outputscale) - -# @outputscale.setter -# def outputscale(self, value: torch.Tensor | float | np.ndarray) -> None: -# self._set_outputscale(value) - -# def _set_lengthscales(self, value: torch.Tensor | float | np.ndarray) -> None: -# if not torch.is_tensor(value): -# value = torch.as_tensor(value).to(self.raw_lengthscales) -# self.initialize(raw_lengthscales=self.raw_lengthscales_constraint.inverse_transform(value)) - -# def _set_skew_entries(self, value: torch.Tensor | float | np.ndarray) -> None: -# if not torch.is_tensor(value): -# value = torch.as_tensor(value).to(self.raw_skew_entries) -# self.initialize(raw_skew_entries=self.raw_skew_entries_constraint.inverse_transform(value)) - -# def _set_outputscale(self, value: torch.Tensor | float | np.ndarray) -> None: -# if not torch.is_tensor(value): -# value = torch.as_tensor(value).to(self.raw_outputscale) -# self.initialize(raw_outputscale=self.raw_outputscale_constraint.inverse_transform(value)) - -# @property -# def skew_matrix(self) -> torch.Tensor: -# S = torch.zeros((self.num_outputs, self.num_inputs, self.num_inputs), dtype=torch.float64) -# if self.n_skew_entries > 0: -# # to construct an orthogonal matrix. fun fact: exp(skew(N)) is the generator of SO(N) -# S[self.skew_matrix_indices] = self.skew_entries -# S += -S.transpose(-1, -2) -# return torch.linalg.matrix_exp(S) - -# @property -# def diag_matrix(self) -> torch.Tensor: -# D = torch.zeros((self.num_outputs, self.num_inputs, self.num_inputs), dtype=torch.float64) -# D[self.diag_matrix_indices] = self.lengthscales.ravel() ** -1 -# return D - -# @property -# def latent_transform(self) -> torch.Tensor: -# return torch.matmul(self.diag_matrix, self.skew_matrix) - -# def forward(self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, **params: dict) -> torch.Tensor: -# # adapted from the Matern kernel -# mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)] - -# transform = self.latent_transform.unsqueeze(1) - -# trans_x1 = torch.matmul(transform, (x1 - mean).unsqueeze(-1)).squeeze(-1) -# trans_x2 = torch.matmul(transform, (x2 - mean).unsqueeze(-1)).squeeze(-1) - -# distance = self.covar_dist(trans_x1, trans_x2, diag=diag, **params) - -# if self.num_outputs == 1: -# distance = distance.squeeze(0) - -# outputscale = self.outputscale if self.scale_output else 1.0 - -# # special cases of the Matern function -# if self.nu == 0.5: -# return outputscale * torch.exp(-distance) -# if self.nu == 1.5: -# return outputscale * (1 + distance) * torch.exp(-distance) -# if self.nu == 2.5: -# return outputscale * (1 + distance + distance**2 / 3) * torch.exp(-distance) -# raise ValueError(f"nu = {self.nu} is not supported") diff --git a/src/blop/bayesian/models.py b/src/blop/bayesian/models.py index 324e3a8e..9a0aa01d 100644 --- a/src/blop/bayesian/models.py +++ b/src/blop/bayesian/models.py @@ -46,54 +46,6 @@ def __init__( d=train_X.shape[-1], batch_shape=aug_batch_shape, skew_dims=skew_dims ) - # return SingleTaskGP( - # train_X=train_X, - # train_Y=train_Y, - # mean_module=mean, - # covar_module=covar, - # outcome_transform=None, # keep if you truly need it - # **kwargs, - # ) - - -# class LatentGP(SingleTaskGP): -# def __init__( -# self, -# train_X: torch.Tensor, -# train_Y: torch.Tensor, -# train_Tvar: torch.Tensor = None, -# train_Yvar: Tensor | None = None, -# likelihood: Likelihood | None = None, -# input_transform: InputTransform | None = None, -# outcome_transform: OutcomeTransform | _DefaultType | None = DEFAULT, -# skew_dims: bool | list[tuple[int, ...]] = True, -# *args: Any, -# **kwargs: Any, -# ) -> None: - -# *batch_shape, n, d = train_X.shape -# input_transform = input_transform or Normalize(d=d) -# # outcome_transform = outcome_transform or Standardize(batch_shape=batch_shape) - -# super().__init__(train_X=train_X, -# train_Y=train_Y, -# input_transform=input_transform, -# outcome_transform=outcome_transform, -# *args, **kwargs) - -# self.mean_module = gpytorch.means.ConstantMean(constant_prior=gpytorch.priors.NormalPrior(loc=0, scale=1)) - -# self.covar_module = kernels.LatentKernel( -# num_inputs=train_X.shape[-1], -# num_outputs=train_Y.shape[-1], -# skew_dims=skew_dims, -# priors=True, -# scale=True, -# **kwargs, -# ) - -# self.trained: bool = False - class MultiTaskLatentGP(MultiTaskGP): def __init__(