Source code for aepsych.factory.pairwise
# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
from typing import List, Tuple, Union
import gpytorch
from aepsych.config import Config
from aepsych.factory.default import (
_get_default_cov_function,
_get_default_mean_function,
default_mean_covar_factory,
)
from aepsych.kernels.pairwisekernel import PairwiseKernel
[docs]def pairwise_mean_covar_factory(
config: Config,
) -> Tuple[gpytorch.means.ConstantMean, gpytorch.kernels.ScaleKernel]:
"""Creates a mean and covariance function for pairwise GPs.
Args:
config (Config): Config object containing bounds.
Returns:
Tuple[gpytorch.means.ConstantMean, gpytorch.kernels.ScaleKernel]: A tuple containing
the mean function (ConstantMean) and the covariance function (ScaleKernel)."""
lb = config.gettensor("common", "lb")
ub = config.gettensor("common", "ub")
assert lb.shape[0] == ub.shape[0], "bounds shape mismatch!"
assert lb.shape[0] >= 2, "PairwiseKernel requires at least 2 dimensions!"
config_dim = lb.shape[0]
shared_dims: Union[List[int], None] = config.getlist(
"pairwise_mean_covar_factory", "shared_dims", fallback=None
)
if shared_dims is not None:
shared_dims = [int(d) for d in shared_dims]
assert len(shared_dims) < config_dim, "length of shared_dims must be < dim!"
for dim in shared_dims:
assert dim < len(shared_dims)
else:
shared_dims = []
base_mean_covar_factory = config.getobj(
"pairwise_mean_covar_factory",
"base_mean_covar_factory",
fallback=default_mean_covar_factory,
)
if base_mean_covar_factory is not default_mean_covar_factory:
raise NotImplementedError(
"Only default_mean_covar_factory is supported for the base factor of pairwise_mean_covar_factory right now!"
)
if len(shared_dims) > 0:
active_dims = [i for i in range(config_dim) if i not in shared_dims]
assert (
len(active_dims) % 2 == 0
), "dimensionality of non-shared dims must be even!"
mean = _get_default_mean_function(config)
cov1 = _get_default_cov_function(
config, len(active_dims) // 2, stimuli_per_trial=1
)
cov2 = _get_default_cov_function(
config, len(shared_dims), active_dims=shared_dims, stimuli_per_trial=1
)
covar = PairwiseKernel(cov1, active_dims=active_dims) * cov2
else:
assert config_dim % 2 == 0, "dimensionality must be even!"
mean = _get_default_mean_function(config)
cov = _get_default_cov_function(config, config_dim // 2, stimuli_per_trial=1)
covar = PairwiseKernel(cov)
return mean, covar