Fix mean

[marjanfamili]

Universal Kriging implementation

This terminology is borrowed from this paper
https://arxiv.org/pdf/2408.02331

1. Simple Kriging: Known mean function, no noise
2. Ordinary Kriging: Unknown constant mean, no noise
3. Universal Kriging: Unknown mean as linear combination of known functions

A notebook is added to exploratory to help run it

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[sgreenbury]

I think this is looking great, thanks @marjanfamili! Adding a few comments:

• This should work well for composing in a GP using as the mean_module_fn along with a covar_module_fn setting the covariance along the known axis as 0. Perhaps it would be worth us adding an example with a GP including this aspect here (i.e. only considers the covariance along the non-fixed axes). It would be interesting to see the affect on the performance in the projectile case too.
• Perhaps it's worth adding a test for the PartiallyLearnableMean for the expected behaviour when passing data through its .forward()
• I wonder if it would be worth generalizing this to take a list of callables along specified axes? This would probably make more sense in a new PR as can extend the implementation here.

[sgreenbury]

From discussion with @marjanfamili, we'll return to this next week following looking more at:

• Combining with a kernel that goes to 0 along the direction of the known axis
• Update the tutorial here with the kernel
• Confirming approach regarding how the learned weights can be used linearly with distance from axis