Anisotropic Modeling

[alecmdunton]

We need to add a feature that allows for anisotropic modeling. This will involve changing the optimization chassis and creating functionality that takes individual distance tensors, weights them, and produces the distance tensor for the full dataset.

[alecmdunton]

For everyone's awareness, anisotropy is currently supported in the PyTorch implementation of the library. One must simply created an embedding which is a diagonal linear mapping to deform the data anisotropically (each diagonal element corresponds to the length scale parameter along each feature dimension).

[bwpriest]

PR #107 created a package, MuyGPyS.gp.distortion that holds the IsotropicDistortion logic. We will need to create an AnisotropicDistortion class therein, although it will be more involved because it has parameters (which can be optimization targets). This class will need get_opt_fn and get_optim_params methods that will have to hook into those of RBF/Matern. IsotropicDistortion will need its own version of those functions, although they will effectively be no-ops because there are no parameters.

[alecmdunton]

I am thinking that we will want to represent the self.length_scales attribute of an AnisotropicDistortion as a TensorHyperparameter.

[bwpriest]

We'll need to think about how to encourage users to not learn a separate length scale in a stationary kernel when we're using anisotropy. Should we pull the length scale fully into the distortion model? If so, I don't know if it makes sense to use TensorHyperparameter.

[alecmdunton]

from typing import Callable
def _F2_func(distortion: Callable) -> Callable:
    def _F2(distortion):
        return np.sum(np.exp2(distortion), axis=-1)
    return _F2

def _l2_func(distortion: Callable) -> Callable:
    def _l2(distortion):
        return np.sqrt(np.sum(np.exp2(distortion), axis=-1))
    return _l2

def _scale(length_scales: np.ndarray) -> Callable:
    def scaling_fn(diffs: np.ndarray) -> np.ndarray: 
        return diffs / length_scales
    return scaling_fn

def _dist(metric: str, distortion: Callable) -> Callable:
    if metric == "l2":
        return _l2_func(distortion)
    elif metric == "F2":
        return _F2_func(distortion)
    else:
        raise ValueError(f"Metric {metric} is not supported!")

dist_fn = _dist(metric,_scale(length_scales))

[alecmdunton]

I am looking at building the self._dist_fn in AnisotropicDistortion like this. It will involve creating some new functions, and I was trying to think of the best place to store them, if not just within the class itself.

I also prefer your idea of pulling the length_scales into the distortion model. We don't want to use a TensorHyperparameter anyway because they don't currently support optimization.

I am also realizing that your code in MuyGPyS.gp.distortion.embed may already do what my code above does.

[alecmdunton]

Do we want BenchmarkGP to support anisotropic modeling? Looking at it now it should actually be really straightforward.

[bwpriest]

Do we want BenchmarkGP to support anisotropic modeling? Looking at it now it should actually be really straightforward.

Ultimately yes, but I think we'll need to rewrite how the BenchmarkGP creates distance tensors. We can leave that for a later push unless you have momentum and want to do it now.

[alecmdunton]

I was thinking that broadcasting would take care of the rescaling done within BenchmarkGP, e.g., test / self.length_scale. That would require me to rework the naming convention I currently have in AnisotropicDistortion, so that self.length_scale would be an mm.ndarray in the AnisotropicDistortion class but would be a Hyperparameter in the IsotropicDistortion class.

[bwpriest]

We should not make software decisions to accommodate BenchmarkGP, because it is not a user-facing class. It is fine to do suboptimal stuff therein, although I would prefer that we modify it to use IsotropicDistortion and AnisotropicDistortion in the same way that MuyGPS does to make it more maintainable. Again, we can leave that for later if you want to.

[alecmdunton]

How extensive do we want to make the test coverage for anisotropy? I have updated tests/gp.py, and am now looking at some of the examples in tests/predict.py and tests/optimize.py. At a minimum we have to stick to problems with lower dimensionality, e.g., Heaton. I am also wondering if it is a necessary addition to anything in tests/backend.

[bwpriest]

I'd ignore tests/predict.py. We'll probably drop that one or refactor it in the relatively near future. It will be hard/impossible to add anything meaningful to tests/optimize.py if the loss functions are not sensitive to length scales. We've only ever tested the recovery of sigma_sq and nu therein, and it is kinda broken right now anyway. We definitely need tests in tests/backend though. We probably want to copy all the tests that use isotropic kernels using anisotropic kernels, although that will be a lot of keystrokes. It might be worth figuring out if there is a way to automate those tests - i.e. give a DistortionModel and a NoiseModel to a function that then creates and runs all of the relevant tests. That will likely be a little nontrivial, and would probably be best left to a subsequent PR.

[alecmdunton]

Thanks for the response - I had similar feelings about each of those test scripts but wanted to get your opinion before I issue a draft PR.

[bwpriest]

Addressed with PR #127