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