Support for UQ with conformal prediction

[sgreenbury]

Closes #848.
Contributes towards #589.
(merge after #917)

This PR:

• Updates the base emulator API to enable the specification of validation data (to be used here as calibration data) (Let user pass train/val/test data split #905)
• Adds a base class wrapper for emulators to provide conformal prediction (analogous to Ensemble)
• Adds a ConformalMLP class (analogous to EnsembleMLP). This supports both constant width and quantile regression derived UQ.
• The quantile regression version uses an MLP with a quantile training loss to fit the upper and lower quantiles and then uses a constant per-target correction based on performance on the calibration training data
• Adds ConformalMLP to registry and re-exports (not adding as default emulator)
• Updates API to support passing n_samples - this is to enable this to be controlled at lower levels of the API from the high-level AutoEmulate.
• Adds output_to_tensors method on ConversionMixin to avoid repeated mapping of outputs to tensors
• Fixes indentation of docstrings in MLP subclasses

Question:

• Currently the two methods for deriving the intervals are part of the same base class with a parameter passed as a keyword argument - this is similar to the way the kernel fn is passed for GPs. Would it be preferable to have two classes? Or perhaps a general emulator subclass factory extending the one currently that only supports GPs.

Next steps:

• Tests for passing validation data
• Check UQ performance of ConformalMLP (e.g. on projectile and other simulators)
  • For projectile "split" constant-width intervals does not provide good UQ while "quantile" is an improvement but less good than e.g. a GP.

[radka-j]

I don't think this PR closes #905 since we still don't let the user specify the test data set (and use all x as train data).

[radka-j]

This looks great! All my comments are really just about naming/docstrings.

[sgreenbury]

I don't think this PR closes #905 since we still don't let the user specify the test data set (and use all x as train data).

Ah I see - #905 is distinct from updates to the Emulator API required in #589. I'll update the top-level comment as contributes towards #589. Are the API changes in #905 specifically towards the AutoEmulate class?

[sgreenbury]

From discussion with @radka-j:

• Add reference from paper above in notes docstring (we can mention that you can do the quantile regression in different ways uses regression in paper but we're using MLP here)
• Extending to methods like the scaling can be in conformal base - add comment to explain and open an issue (different methodologies such as distance to training data and even MLP, to be explored)
• Change to constant
• Update to specific subclasses
• Update API to support distribution type specification and docs explaining assumption

[radka-j]

I was thinking about our discussion yesterday about evaluating conformal intervals using MSLL. I only had a brief look but I don't think we can treat the conformal intervals as a distribution - we don't know if it is normal but also don't have a reason to assume it is uniform. We only have the 2 quantiles we compute (and the mean).

[sgreenbury]

I was thinking about our discussion yesterday about evaluating conformal intervals using MSLL. I only had a brief look but I don't think we can treat the conformal intervals as a distribution - we don't know if it is normal but also don't have a reason to assume it is uniform. We only have the 2 quantiles we compute (and the mean).

Yes good point - adding a distribution s imposing an additional assumption so that we can represent as a ProbabilisticEmulator. One option could be to provide API to support specifying / overriding the distributional assumption and some docs to describe the design choice.

[radka-j]

Yes good point - adding a distribution s imposing an additional assumption so that we can represent as a ProbabilisticEmulator. One option could be to provide API to support specifying / overriding the distributional assumption and some docs to describe the design choice.

I think that sounds really good.

This discussion made me think about whether in the case of Ensembles we also shouldn't assume normality. One option would be to introduce something like an EmpiricalDistribution that wraps around the ensemble predictions and can compute all the different empirical quantiles. Sampling in this case could be sampling from the ensemble with replacement. I think assuming normality is common so we can leave it as is but this might be more accurate - maybe we open an issue to discuss?

[radka-j]

Related to some of the discussion here, I also think we might want to update our plots to show 95% intervals (rather than 2 sigma). Can create an issue for that if you agree.

[sgreenbury]

Related to some of the discussion here, I also think we might want to update our plots to show 95% intervals (rather than 2 sigma). Can create an issue for that if you agree.

Yes good point given we have other distributions than normals potentially - let's update this.

[EdwinB12]

Can I check where we are with this PR?

[radka-j]

@EdwinB12 waiting for me to do a second review

[radka-j]

At the meeting on Wednesday a suggestion was made to plot calibration of the conformal predictor on a test data set to check it is behaving as expected - I think this is standard. In that instance a user would probably want to loop over different $\alpha$ values, construct those intervals and then evaluate how much data is contained within each. Is our current implementation suited to this?

[radka-j]

After a discussion with @sgreenbury we've agreed on the following design:

# Q: possibly namedtuple instead for clarity?
IntervalLike = tuple[TensorLike, TensorLike]  # lower/upper interval

#  NOTE: call predict_mean in predict
class IntervalEmulator(Emulator):   

    # Q: Possibly abstractmethod?
    # NOTE: use alpha to determine lower/upper quantile pair (e.g., 0.05 and 0.95 quantiles for alpha=0.1)
    def predict_interval(self, x: TensorLike, alpha=0.05) -> IntervalLike:
         raise NotImplementedException("Should be implemented in subclass")

class Conformal(InervalEmulator):
    emulator: Emulator

    # Q: custom fit? could refit here?
    def _fit(self, x: TensorLike, ..., refit: bool = True): 
        ...

    def predict_interval(self, x: TensorLike, alpha=0.05) -> IntervalLike:
        ...

class ConformalConstant(Conformal):
    ... 

class ConformalizedQR(Conformal):
    alpha: float | list[float] # self.alphas = [alpha] if isinstance(alpha, float)

     ...

    # NOTE: alpha has to be in self.alphas
    def predict_interval(self, x: TensorLike, alpha=0.05) -> IntervalLike:
        ...