Add static principal axes kernel metric rotation

[nathyglobal]

Summary

Adds support for a static principal-axes rotation of the GP kernel metric.

The rotation can be constructed from either:

• a covariance matrix, or
• samples/chains, optionally with weights or log weights.

User-supplied rotation inputs are assumed to be in physical parameter coordinates by default and are converted to BOBE's unit-hypercube coordinates before GP construction. Inputs already expressed in unit coordinates can instead be supplied with rotation_space="unit".

The rotation may be applied to the full parameter space or to selected dimensions. The transform is applied inside the kernel before ARD lengthscale scaling, so the original GP and classifier input coordinates are unchanged.

No rotation centre is required because the RBF and Matérn kernels are stationary and depend only on pairwise coordinate differences, for which a common translation cancels.

Changes

• Added PrincipalAxesTransform for covariance- and sample-based kernel metric rotations.
• Added physical-to-unit scaling of rotation covariance matrices and samples.
• Added optional kernel input transforms to the RBF and Matérn kernels.
• Added rotation support to GP and GPwithClassifier.
• Added support for full or reduced covariance/sample inputs when rotating selected dimensions.
• Added support for saving, loading and copying rotated GP states.
• Added tests covering transforms, coordinate scaling, kernels, GP integration, classifier-GP integration and serialization.

Testing

• 65/65 focused kernel-rotation tests pass.
• Previous full test suite: 106/107 tests passed.
• The remaining failure was the existing test_bobe_wipstd_2d optimisation failure and was unrelated to these changes.

Validation

Draft PR while additional validation is performed against the existing likelihood benchmarks from the earlier kernel-rotation implementation.

[Ameek94]

Wouldn't it be simpler to ask for a rotation matrix and centre instead of asking for the full samples and weights? This would simplify the input handling in the GP, transforms as well as the tests.

[nathyglobal]

Wouldn't it be simpler to ask for a rotation matrix and centre instead of asking for the full samples and weights? This would simplify the input handling in the GP, transforms as well as the tests.

It would probably simplify the internal handling a bit, but I think it would make the interface more complicated for the user. In practice they are much more likely to already have a covariance matrix or posterior samples/weights, so asking for a rotation matrix would just move the coordinate rescaling and egendecomposition into user code. We also don't actually need a rotation centre here, since for a stationary kernel any common translation cancels in the pairwise differences.

[Ameek94]

We also don't actually need a rotation centre here, since for a stationary kernel any common translation cancels in the pairwise differences.

That makes sense, I was still thinking in terms of our other rotation approach.

I would say say we don't ask for samples and weights, only the covmat since most cosmo samplers already provide it, otherwise np.cov handles it. It would make the transforms and related tests a lot smaller.

[Ameek94]

There is a bug in how we handle kernel rotation + batch acquisition. Our current method discards the kernel rotation if n_batch>1.

if n_batch > 1:
    # Create dummy GP without classifier functionality.
    # For now, we do not use batching for EI/LogEI.
    dummy_gp = GP(
        train_x=gp.train_x,
        train_y=gp.train_y * gp.y_std + gp.y_mean,
        noise=gp.noise,
        kernel=gp.kernel_name,
        lengthscales=gp.lengthscales,
        kernel_variance=gp.kernel_variance,
    )

The fix would be to reconstruct the dummy GP from the state_dict instead

    dummy_gp = GP.from_state_dict(gp.state_dict())

which retains the rotation. This would explain why kernel-rotation did not seem to work as well as the parameter-space rotation in our earlier tests. With this fix I see much better results (added to our draft).

[nathyglobal]

Thanks for catching that! But I believe this was me forgetting to port over that part of the code, here's how it was written in my testing version:

if n_batch > 1:
            # Create dummy GP without classifier functionality, for now we do not use batching for EI/LogEI
            state = gp.state_dict()
            dummy_gp = GP.from_state_dict(state)
                        
            dummy_gp.update(x_next, dummy_gp.predict_mean_single(x_next))
            for i in range(1,n_batch):
                x_next, acq_val_next = self.get_next_point(dummy_gp, acq_kwargs=acq_kwargs,
                                                        maxiter=maxiter,
                                                        n_restarts=n_restarts,
                                                        verbose=verbose,
                                                        early_stop_patience=early_stop_patience,
                                                        rng=rng)
                x_batch.append(x_next)
                acq_vals.append(acq_val_next)

                mu = dummy_gp.predict_mean_single(x_next)
                dummy_gp.update(x_next, mu)

        return np.array(x_batch), np.array(acq_vals)

So if there was better results, I'm not quite certain why.