Sparse Multitask GPR

Exact sparse multitask GPR for incomplete multi-property target matrices.

matgpr.sparse_multitask_gpr

SparseMultitaskObservationData dataclass

Observed entries extracted from a sparse multitask target matrix.

Attributes:

Name	Type	Description
X_observed	ndarray	Feature rows repeated once for each observed task value.
y_observed	ndarray	One-dimensional observed target values.
task_indices	ndarray	Integer task index for each observed value.
sample_indices	ndarray	Original sample-row index for each observed value.
observed_mask	ndarray	Boolean matrix with shape (n_samples, n_tasks) marking finite target observations.
task_names	tuple[str, ...]	Task names in target-column order.
task_observation_counts	ndarray	Number of observed values for each task.

Source code in matgpr/sparse_multitask_gpr.py

@dataclass(frozen=True)
class SparseMultitaskObservationData:
    """Observed entries extracted from a sparse multitask target matrix.

    Attributes
    ----------
    X_observed
        Feature rows repeated once for each observed task value.
    y_observed
        One-dimensional observed target values.
    task_indices
        Integer task index for each observed value.
    sample_indices
        Original sample-row index for each observed value.
    observed_mask
        Boolean matrix with shape ``(n_samples, n_tasks)`` marking finite
        target observations.
    task_names
        Task names in target-column order.
    task_observation_counts
        Number of observed values for each task.
    """

    X_observed: np.ndarray
    y_observed: np.ndarray
    task_indices: np.ndarray
    sample_indices: np.ndarray
    observed_mask: np.ndarray
    task_names: tuple[str, ...]
    task_observation_counts: np.ndarray

SparseMultitaskGPyTorchResult dataclass

Container returned by :func:fit_sparse_multitask_gpytorch_gpr.

Source code in matgpr/sparse_multitask_gpr.py

@dataclass(frozen=True)
class SparseMultitaskGPyTorchResult:
    """Container returned by :func:`fit_sparse_multitask_gpytorch_gpr`."""

    model: ExactSparseMultitaskGPRModel
    likelihood: gpytorch.likelihoods.Likelihood
    observation_data: SparseMultitaskObservationData
    loss_history: list[float]
    target_mean: np.ndarray
    target_std: np.ndarray
    standardize_y: bool
    kernel: str
    ard: bool
    num_tasks: int
    task_names: tuple[str, ...]
    task_covar_rank: int
    noise_mode: str
    standardized_task_noise_variance: np.ndarray
    task_noise_variance: np.ndarray
    task_noise_std: np.ndarray
    standardized_observation_noise_variance: np.ndarray | None
    observation_noise_variance: np.ndarray | None
    device: str
    dtype: torch.dtype

    def predict(
        self,
        X,
        *,
        return_std: bool = True,
        confidence_level: float | None = None,
        include_observation_noise: bool = True,
        prediction_noise_variance=None,
    ) -> MultitaskGPyTorchPrediction:
        """Predict all modeled tasks for new samples."""
        return _predict_sparse_multitask_gpytorch_gpr(
            self.model,
            self.likelihood,
            X,
            device=self.device,
            dtype=self.dtype,
            return_std=return_std,
            confidence_level=confidence_level,
            include_observation_noise=include_observation_noise,
            prediction_noise_variance=prediction_noise_variance,
        )

predict(X, *, return_std=True, confidence_level=None, include_observation_noise=True, prediction_noise_variance=None)

Predict all modeled tasks for new samples.

Source code in matgpr/sparse_multitask_gpr.py

def predict(
    self,
    X,
    *,
    return_std: bool = True,
    confidence_level: float | None = None,
    include_observation_noise: bool = True,
    prediction_noise_variance=None,
) -> MultitaskGPyTorchPrediction:
    """Predict all modeled tasks for new samples."""
    return _predict_sparse_multitask_gpytorch_gpr(
        self.model,
        self.likelihood,
        X,
        device=self.device,
        dtype=self.dtype,
        return_std=return_std,
        confidence_level=confidence_level,
        include_observation_noise=include_observation_noise,
        prediction_noise_variance=prediction_noise_variance,
    )

SparseTaskConstantMean

Bases: Module

Learn one constant mean offset per task for sparse observations.

Source code in matgpr/sparse_multitask_gpr.py

class SparseTaskConstantMean(torch.nn.Module):
    """Learn one constant mean offset per task for sparse observations."""

    def __init__(self, num_tasks: int):
        super().__init__()
        self.register_parameter(
            name="task_means",
            param=torch.nn.Parameter(torch.zeros(int(num_tasks))),
        )

    def forward(self, x: torch.Tensor, task_indices: torch.Tensor) -> torch.Tensor:
        del x
        return self.task_means[task_indices.long()]

ExactSparseMultitaskGPRModel

Bases: ExactGP

Exact sparse multitask GPR using observed (x, task) pairs.

Unlike :class:matgpr.ExactMultitaskGPRModel, this model does not require every task to be observed for every sample. Each finite target entry is represented as one scalar observation with covariance

k((x, i), (x', j)) = k_x(x, x') k_task(i, j).

Source code in matgpr/sparse_multitask_gpr.py

class ExactSparseMultitaskGPRModel(gpytorch.models.ExactGP):
    """Exact sparse multitask GPR using observed ``(x, task)`` pairs.

    Unlike :class:`matgpr.ExactMultitaskGPRModel`, this model does not require
    every task to be observed for every sample. Each finite target entry is
    represented as one scalar observation with covariance

    ``k((x, i), (x', j)) = k_x(x, x') k_task(i, j)``.
    """

    def __init__(
        self,
        train_x: torch.Tensor,
        train_task_indices: torch.Tensor,
        train_y: torch.Tensor,
        likelihood: gpytorch.likelihoods.Likelihood,
        *,
        num_tasks: int,
        task_covar_rank: int = 1,
        kernel: str = "matern",
        ard_num_dims: int | None = None,
        mean_module: torch.nn.Module | None = None,
    ):
        super().__init__((train_x, train_task_indices), train_y, likelihood)
        self.num_tasks = int(num_tasks)
        self.task_covar_rank = int(task_covar_rank)
        self.mean_module = mean_module or SparseTaskConstantMean(self.num_tasks)
        self.data_covar_module = gpytorch.kernels.ScaleKernel(
            _make_gpytorch_base_kernel(kernel, ard_num_dims=ard_num_dims)
        )
        self.task_covar_module = gpytorch.kernels.IndexKernel(
            num_tasks=self.num_tasks,
            rank=self.task_covar_rank,
        )
        # Buffers so per-task target standardization survives ``state_dict()``.
        self.register_buffer(
            "target_mean",
            torch.zeros(self.num_tasks, dtype=train_x.dtype, device=train_x.device),
        )
        self.register_buffer(
            "target_std",
            torch.ones(self.num_tasks, dtype=train_x.dtype, device=train_x.device),
        )

    def forward(
        self,
        x: torch.Tensor,
        task_indices: torch.Tensor,
    ) -> gpytorch.distributions.MultivariateNormal:
        task_indices = task_indices.long().reshape(-1)
        mean_x = _call_sparse_mean_module(self.mean_module, x, task_indices)
        covar_x = self.data_covar_module(x).mul(self.task_covar_module(task_indices))
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

prepare_sparse_multitask_observations(X, y, *, task_names=None, min_observations_per_task=2)

Convert a target matrix with missing entries into observed task values.

Missing targets are represented by NaN values in y. Feature values must already be finite because every observed task for a row shares the same feature vector.

Source code in matgpr/sparse_multitask_gpr.py

def prepare_sparse_multitask_observations(
    X,
    y,
    *,
    task_names: Sequence[str] | None = None,
    min_observations_per_task: int = 2,
) -> SparseMultitaskObservationData:
    """Convert a target matrix with missing entries into observed task values.

    Missing targets are represented by ``NaN`` values in ``y``. Feature values
    must already be finite because every observed task for a row shares the
    same feature vector.
    """
    X_array = _to_2d_numpy(X, "X")
    y_array = _to_2d_numpy(y, "y", allow_nan=True)
    if X_array.shape[0] != y_array.shape[0]:
        raise ValueError("X and y must contain the same number of samples")
    if y_array.shape[1] < 2:
        raise ValueError("Sparse multitask GPR requires at least two target columns")
    if np.isinf(y_array).any():
        raise ValueError("y must not contain infinite values")

    task_names_resolved = _resolve_sparse_task_names(task_names, y, y_array.shape[1])
    min_observations_per_task = _validate_min_observations_per_task(min_observations_per_task)
    observed_mask = np.isfinite(y_array)
    task_counts = observed_mask.sum(axis=0).astype(int)
    low_count_tasks = [
        name
        for name, count in zip(task_names_resolved, task_counts, strict=True)
        if count < min_observations_per_task
    ]
    if low_count_tasks:
        raise ValueError(
            "Each task must have at least "
            f"{min_observations_per_task} observed values; low-count task(s): "
            f"{low_count_tasks}"
        )

    sample_indices, task_indices = np.where(observed_mask)
    if sample_indices.size < 2:
        raise ValueError("At least two observed target values are required")

    return SparseMultitaskObservationData(
        X_observed=X_array[sample_indices].copy(),
        y_observed=y_array[sample_indices, task_indices].astype(float, copy=True),
        task_indices=task_indices.astype(int, copy=True),
        sample_indices=sample_indices.astype(int, copy=True),
        observed_mask=observed_mask.copy(),
        task_names=task_names_resolved,
        task_observation_counts=task_counts,
    )

fit_sparse_multitask_gpytorch_gpr(X_train, y_train, *, task_names=None, task_covar_rank=1, kernel='matern', ard=True, mean_module=None, lr=0.01, training_iter=1000, initial_noise=0.1, initial_task_noises=None, known_noise_variance=None, noise_mode='shared', noise_lower_bound=0.0001, standardize_y=True, min_observations_per_task=2, device='cpu', dtype=torch.float64, verbose=False, log_every=100)

Fit exact sparse multitask GPR from a target matrix with NaN gaps.

Set noise_mode="shared" for one learned observation-noise variance across all observed task values, or noise_mode="task" for one learned noise variance per target task. Use noise_mode="known" with known_noise_variance when each observed target entry has a known measurement-noise variance in original target units.

Source code in matgpr/sparse_multitask_gpr.py

def fit_sparse_multitask_gpytorch_gpr(
    X_train,
    y_train,
    *,
    task_names: Sequence[str] | None = None,
    task_covar_rank: int = 1,
    kernel: str = "matern",
    ard: bool = True,
    mean_module: torch.nn.Module | None = None,
    lr: float = 0.01,
    training_iter: int = 1000,
    initial_noise: float | None = 0.1,
    initial_task_noises: Sequence[float] | Mapping[str, float] | None = None,
    known_noise_variance=None,
    noise_mode: str = "shared",
    noise_lower_bound: float = 1e-4,
    standardize_y: bool = True,
    min_observations_per_task: int = 2,
    device: str = "cpu",
    dtype: torch.dtype = torch.float64,
    verbose: bool = False,
    log_every: int = 100,
) -> SparseMultitaskGPyTorchResult:
    """Fit exact sparse multitask GPR from a target matrix with ``NaN`` gaps.

    Set ``noise_mode="shared"`` for one learned observation-noise variance
    across all observed task values, or ``noise_mode="task"`` for one learned
    noise variance per target task. Use ``noise_mode="known"`` with
    ``known_noise_variance`` when each observed target entry has a known
    measurement-noise variance in original target units.
    """
    _validate_training_options(lr=lr, training_iter=training_iter, log_every=log_every)
    observation_data = prepare_sparse_multitask_observations(
        X_train,
        y_train,
        task_names=task_names,
        min_observations_per_task=min_observations_per_task,
    )
    train_x = _to_tensor(observation_data.X_observed, device=device, dtype=dtype)
    train_task_indices = torch.as_tensor(
        observation_data.task_indices,
        dtype=torch.long,
        device=device,
    )
    train_y = _to_tensor(observation_data.y_observed, device=device, dtype=dtype)
    num_tasks = len(observation_data.task_names)
    task_covar_rank = validate_task_covar_rank(task_covar_rank, num_tasks)

    target_mean, target_std = _sparse_target_statistics(
        train_y,
        train_task_indices,
        num_tasks=num_tasks,
        standardize_y=standardize_y,
    )
    train_y_model = (train_y - target_mean[train_task_indices]) / target_std[train_task_indices]
    observation_noise_variance, standardized_observation_noise_variance = (
        _resolve_sparse_known_noise_variance(
            known_noise_variance,
            observation_data=observation_data,
            target_std=target_std.detach().cpu().numpy(),
        )
    )

    noise_mode = _validate_sparse_noise_mode(noise_mode)
    likelihood = _make_sparse_likelihood(
        noise_mode=noise_mode,
        num_tasks=num_tasks,
        task_names=observation_data.task_names,
        initial_noise=initial_noise,
        initial_task_noises=initial_task_noises,
        standardized_known_noise_variance=standardized_observation_noise_variance,
        noise_lower_bound=noise_lower_bound,
        device=device,
        dtype=dtype,
    )
    if mean_module is not None:
        mean_module = mean_module.to(device=device, dtype=dtype)

    model = ExactSparseMultitaskGPRModel(
        train_x,
        train_task_indices,
        train_y_model,
        likelihood,
        num_tasks=num_tasks,
        task_covar_rank=task_covar_rank,
        kernel=kernel,
        ard_num_dims=train_x.shape[1] if ard else None,
        mean_module=mean_module,
    ).to(device=device, dtype=dtype)
    model.target_mean = target_mean.detach()
    model.target_std = target_std.detach()
    model.standardize_y = standardize_y
    model.task_names = observation_data.task_names
    model.noise_mode = noise_mode
    model.standardized_observation_noise_variance = (
        None
        if standardized_observation_noise_variance is None
        else standardized_observation_noise_variance.copy()
    )
    model.observation_noise_variance = (
        None if observation_noise_variance is None else observation_noise_variance.copy()
    )

    model.train()
    likelihood.train()

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    marginal_log_likelihood = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)
    loss_history: list[float] = []

    for iteration in range(training_iter):
        optimizer.zero_grad()
        output = model(train_x, train_task_indices)
        loss = -_sparse_marginal_log_likelihood(
            marginal_log_likelihood,
            output,
            train_y_model,
            train_task_indices,
            likelihood=likelihood,
        )
        loss.backward()
        optimizer.step()
        loss_history.append(float(loss.detach().cpu().item()))

        if verbose and _should_log_iteration(iteration, training_iter, log_every):
            _print_sparse_multitask_training_status(
                iteration, training_iter, loss, likelihood, model
            )

    model.training_loss_history = loss_history
    standardized_noise_variance = _standardized_sparse_task_noise_variance(
        likelihood,
        noise_mode=noise_mode,
        num_tasks=num_tasks,
        task_indices=observation_data.task_indices,
        standardized_observation_noise_variance=standardized_observation_noise_variance,
    )
    task_noise_variance = standardized_noise_variance * target_std.detach().cpu().numpy() ** 2
    task_noise_std = np.sqrt(task_noise_variance)
    model.standardized_task_noise_variance = standardized_noise_variance.copy()
    model.task_noise_variance = task_noise_variance.copy()
    model.task_noise_std = task_noise_std.copy()

    return SparseMultitaskGPyTorchResult(
        model=model,
        likelihood=likelihood,
        observation_data=observation_data,
        loss_history=loss_history,
        target_mean=target_mean.detach().cpu().numpy(),
        target_std=target_std.detach().cpu().numpy(),
        standardize_y=standardize_y,
        kernel=kernel,
        ard=ard,
        num_tasks=num_tasks,
        task_names=observation_data.task_names,
        task_covar_rank=task_covar_rank,
        noise_mode=noise_mode,
        standardized_task_noise_variance=standardized_noise_variance,
        task_noise_variance=task_noise_variance,
        task_noise_std=task_noise_std,
        standardized_observation_noise_variance=standardized_observation_noise_variance,
        observation_noise_variance=observation_noise_variance,
        device=device,
        dtype=dtype,
    )

train_sparse_multitask_gpytorch_gpr(X_train, y_train, *, task_names=None, task_covar_rank=1, kernel='matern', ard=True, mean_module=None, lr=0.01, training_iter=1000, initial_noise=0.1, initial_task_noises=None, known_noise_variance=None, noise_mode='shared', noise_lower_bound=0.0001, standardize_y=True, min_observations_per_task=2, device='cpu', dtype=torch.float64, verbose=False, log_every=100, return_result=False)

Train sparse multitask GPR and return (model, likelihood) by default.

Source code in matgpr/sparse_multitask_gpr.py

def train_sparse_multitask_gpytorch_gpr(
    X_train,
    y_train,
    *,
    task_names: Sequence[str] | None = None,
    task_covar_rank: int = 1,
    kernel: str = "matern",
    ard: bool = True,
    mean_module: torch.nn.Module | None = None,
    lr: float = 0.01,
    training_iter: int = 1000,
    initial_noise: float | None = 0.1,
    initial_task_noises: Sequence[float] | Mapping[str, float] | None = None,
    known_noise_variance=None,
    noise_mode: str = "shared",
    noise_lower_bound: float = 1e-4,
    standardize_y: bool = True,
    min_observations_per_task: int = 2,
    device: str = "cpu",
    dtype: torch.dtype = torch.float64,
    verbose: bool = False,
    log_every: int = 100,
    return_result: bool = False,
):
    """Train sparse multitask GPR and return ``(model, likelihood)`` by default."""
    result = fit_sparse_multitask_gpytorch_gpr(
        X_train,
        y_train,
        task_names=task_names,
        task_covar_rank=task_covar_rank,
        kernel=kernel,
        ard=ard,
        mean_module=mean_module,
        lr=lr,
        training_iter=training_iter,
        initial_noise=initial_noise,
        initial_task_noises=initial_task_noises,
        known_noise_variance=known_noise_variance,
        noise_mode=noise_mode,
        noise_lower_bound=noise_lower_bound,
        standardize_y=standardize_y,
        min_observations_per_task=min_observations_per_task,
        device=device,
        dtype=dtype,
        verbose=verbose,
        log_every=log_every,
    )
    if return_result:
        return result
    return result.model, result.likelihood

predict_sparse_multitask_gpytorch_gpr(model, likelihood, X, *, device='cpu', dtype=torch.float64, return_std=True, confidence_level=None, include_observation_noise=True, prediction_noise_variance=None, return_prediction=False)

Predict every modeled task for each new feature row.

Source code in matgpr/sparse_multitask_gpr.py

def predict_sparse_multitask_gpytorch_gpr(
    model: ExactSparseMultitaskGPRModel,
    likelihood: gpytorch.likelihoods.Likelihood,
    X,
    *,
    device: str = "cpu",
    dtype: torch.dtype = torch.float64,
    return_std: bool = True,
    confidence_level: float | None = None,
    include_observation_noise: bool = True,
    prediction_noise_variance=None,
    return_prediction: bool = False,
):
    """Predict every modeled task for each new feature row."""
    prediction = _predict_sparse_multitask_gpytorch_gpr(
        model,
        likelihood,
        X,
        device=device,
        dtype=dtype,
        return_std=return_std,
        confidence_level=confidence_level,
        include_observation_noise=include_observation_noise,
        prediction_noise_variance=prediction_noise_variance,
    )
    if return_prediction:
        return prediction
    if return_std:
        return prediction.mean, prediction.std
    return prediction.mean