Trajectory samplers now work with mean functions

[henrymoss]

@sebastianober noticed that I had been stupid and forgotten to take into account mean functions in my decoupled sampling implementation.

This is now fixed and I have added extra tests to make sure this cant happen again!

[henrymoss]

@uri-granta at the moment the high level FeatureDecompositionTrajectorySampler is defined for any probabilistic model. However, these samplers only really work for very specific models (i.e. FeatureDecompositionInternalDataModel and FeatureDecompositionInducingPointModel).

Is there a way to specify this requirement for the high level FeatureDecompositionTrajectorySampler rather than at the lower level RandomFourierFeatureTrajectorySampler and DecoupledTrajectorySampler. At the moment I had to do a slightly weird lazy init for the mean function in the init of FeatureDecompositionTrajectorySampler as general probabilistic models do not have mean functions.

[uri-granta]

Looks good! (though I wouldn't trust me to catch any other omissions where the mean function needs to be accounted for but isn't)

[uri-granta]

Is there a way to specify this requirement for the high level FeatureDecompositionTrajectorySampler rather than at the lower level RandomFourierFeatureTrajectorySampler and DecoupledTrajectorySampler.

Not sure I fully understand. If FeatureDecompositionTrajectorySampler only supports the same model types as DecoupledTrajectorySampler then why can't you just define it in the same way:

class FeatureDecompositionTrajectorySampler(
    TrajectorySampler[
        Union[
            FeatureDecompositionInducingPointModel,
            FeatureDecompositionInternalDataModel,
        ]],
    ABC,
):