adjust cov matrix handling

Author: pluflou

lume_model/base.py

@@ -127,7 +127,6 @@ def recursive_serialize(
             and any(isinstance(ele, torch.nn.Module) for ele in value)
         ):
             # List of transformers
-            print(v[key])
             v[key] = [
                 process_torch_module(
                     value[i], base_key, f"{key}_{i}", file_prefix, save_models, False

lume_model/models/gp_model.py

@@ -11,6 +11,7 @@
 from gpytorch.mlls import ExactMarginalLogLikelihood
 from botorch.models.transforms.input import ReversibleInputTransform
 from botorch.models.transforms.outcome import OutcomeTransform
+from linear_operator.utils.cholesky import psd_safe_cholesky
 
 from lume_model.models.prob_model_base import (
     ProbModelBaseModel,
@@ -30,15 +31,13 @@ class GPModel(ProbModelBaseModel):
         model: A single task GPyTorch model or BoTorch model.
         input_transformers: List of input transformers to apply to the input data. Optional, default is None.
         output_transformers: List of output transformers to apply to the output data. Optional, default is None.
-        jitter: Jitter to add to diagonal of covariance matrix for numerical stability, if matrix is not positive definite. Optional, default is 1e-8.
     """
 
     model: SingleTaskGP | MultiTaskGP  # TODO: any other types?
     input_transformers: list[ReversibleInputTransform | torch.nn.Linear] = None
     output_transformers: list[
         OutcomeTransform | ReversibleInputTransform | torch.nn.Linear
     ] = None
-    jitter: float = 1e-8
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -187,10 +186,8 @@ def _get_predictions(
         distribution = self._get_distribution(posterior)
         # Take mean and covariance of the distribution
         mean, covar = distribution.mean, distribution.covariance_matrix
-        # Transform the output (mean and covariance)
-        if self.output_transformers is not None:
-            mean, covar = self._transform_outputs(mean), self._transform_outputs(covar)
         # Return a dictionary of output variable names to distributions
+        # this untransforms the mean and covariance before returning
         return self._create_output_dict((mean, covar))
 
     def _posterior(self, x):
@@ -257,6 +254,11 @@ def _create_output_dict(
                 _cov[:, :ss, :ss] = cov[:, i * ss : (i + 1) * ss, i * ss : (i + 1) * ss]
 
             _cov = self._check_covariance_matrix(_cov)
+            # Last step is to untransform
+            if self.output_transformers is not None:
+                _mean, _cov = self._transform_outputs(_mean), self._transform_outputs(_cov)
+
+            # TODO: add a check for final covariance matrix to be positive definite?
             output_distributions[name] = MultivariateNormal(_mean, _cov)
 
         return output_distributions
@@ -290,7 +292,7 @@ def _transform_outputs(self, output_tensor: torch.Tensor) -> torch.Tensor:
             if isinstance(transformer, ReversibleInputTransform):
                 output_tensor = transformer.untransform(output_tensor)
             elif isinstance(transformer, OutcomeTransform):
-                output_tensor = transformer.untransform(output_tensor)
+                output_tensor = transformer.untransform(output_tensor)[0]
             else:
                 w, b = transformer.weight, transformer.bias
                 output_tensor = torch.matmul((output_tensor - b), torch.linalg.inv(w.T))
@@ -302,8 +304,9 @@ def _check_covariance_matrix(self, cov: torch.Tensor) -> torch.Tensor:
             torch.linalg.cholesky(cov)
         except torch._C._LinAlgError:
             warnings.warn(
-                f"Covariance matrix is not positive definite. Added jitter of {self.jitter:.1e} to the diagonal."
+                f"Covariance matrix is not positive definite. Attempting to add jitter the diagonal."
             )
-            eps = torch.tensor(self.jitter, **self._tkwargs)
-            cov = cov + torch.eye(cov.shape[-1], **self._tkwargs) * eps
+            l  = psd_safe_cholesky(cov) # determines jitter iteratively
+            cov = l @ l.transpose(-1, -2)
+
         return cov