Add GP Wrapped Periodic Kernel #6742

jahall · 2023-05-29T14:24:19Z

This PR is a result of the conversation from the generalized periodic PR.

New features

A full_from_distance(dist, squared=False) method available on all Stationary kernels
A WrappedPeriodic kernel following the pattern of WarpedInput, ScaledCov, and other kernels which transforms which accept a base kernel
- The implementation is more efficient than the warped input method outlined here as per screenshot below

Performance

Example outputs

📚 Documentation preview 📚: https://pymc--6742.org.readthedocs.build/en/6742/

jahall · 2023-05-29T14:35:04Z

@bwengals I think something like this would be good. I guess one outstanding question I have is whether we

Keep the multiplication by 4 - inline with the warping approach and original derivation
Drop the multiplication by 4 - inline with the current Periodic kernel (and the Periodic kernel in GPFlow)

TBH I don't quite get the intuition behind why the same length scale leads to more rapid variations in the periodic versions of these functions - see e.g. these for period=1.0 and ls=0.25

bwengals · 2023-06-02T03:25:20Z

TBH I don't quite get the intuition behind why the same length scale leads to more rapid variations in the periodic versions of these functions - see e.g. these for period=1.0 and ls=0.25

Me neither. I'm hesitant to change the constant for the existing Periodic, because it will make any ones models that use it subtly wrong for not much gain. However, if it's better then maybe now is the perfect time to put it into the WrappedPeriodic with the docstring making it very clear what's going on.

I did some timing tests on your PR too, and I get about equal timings for WrappedPeriodic and WarpedInput, with Periodic being a little faster than both. Either way, I'm excited about adding this class since it makes it much easier.

With your refactor of Stationary, I think it'd be pretty straightforward to add a distance_func argument to Stationary and subclasses where the preexisting euclidean_distance is the default. What do you think? Tagging @lucianopaz

jahall · 2023-06-02T09:12:01Z

I'm hesitant to change the constant for the existing Periodic, because it will make any ones models that use it subtly wrong for not much gain. However, if it's better then maybe now is the perfect time to put it into the WrappedPeriodic with the docstring making it very clear what's going on.

@bwengals I am in favour of not altering the existing Periodic, but keeping the multiplication by 4 in the WrappedPeriodic (along with clarity in the docstring) I did some digging and looks like:

sklearn has the x4
gpytorch has the x4
gpflow drops the x4

I did some timing tests on your PR too, and I get about equal timings for WrappedPeriodic and WarpedInput, with Periodic being a little faster than both. Either way, I'm excited about adding this class since it makes it much easier.

Hmm, let me look into that. I mean it would make sense to me for it to be a bit more efficient as we're avoiding doubling the input space...but maybe in most cases the gain is negligible.

With your refactor of Stationary, I think it'd be pretty straightforward to add a distance_func argument to Stationary and subclasses where the preexisting euclidean_distance is the default. What do you think?

Sounds great - happy to look into it. In a follow-on PR?

bwengals · 2023-06-02T18:48:47Z

pymc/gp/cov.py

+    def __init__(
+        self,
+        input_dim: int,
+        cov_func: Stationary,
+        period,
+        active_dims: Optional[Sequence[int]] = None,
+    ):  


Can WrappedPeriodic take input_dim and active_dims from cov_func? That way these don't need to be repeated.

That makes sense. My only concern would be it is then the only Covariance subclass which doesn't take those params on init.

bwengals · 2023-06-02T18:55:57Z

pymc/gp/cov.py

@@ -812,6 +824,52 @@ def full(self, X, Xs=None):
    def diag(self, X):
        X, _ = self._slice(X, None)
        return self.cov_func(self.w(X, self.args), diag=True)
+
+
+class WrappedPeriodic(Covariance):


I think you had GeneralizedPeriodic originally as the name, why the switch? I think GeneralizedPeriodic makes it a bit clearer what it's doing.

I felt it captured better what it was doing i.e. you use it to wrap up an existing kernel to make it periodic. I think a good name might be a verb (like Add or Prod) since it acts on an existing kernel...but I don't know what that verb would be :) Periodify... But I don't mind moving back to GeneralizedPeriodic.

Makes sense. I guess Wrapped is more describes what the code does, and Generalized describes what the kernel is. Either way makes sense.

bwengals · 2023-06-02T19:16:00Z

pymc/gp/cov.py

+    Wrap a stationary covariance function to make it periodic.
+
+    This is done by warping the input with the function
+
+    .. math::
+        \mathbf{u}(x) = \left(
+            \mathrm{sin} \left( \frac{2\pi x}{T} \right),
+            \mathrm{cos} \left( \frac{2\pi x}{T} \right)
+        \right)
+


It might be nice to add something like, "the GeneralizedPeriodic kernel constructs periodic kernels from any Stationary kernel"

Also, I think it'd be nice to add a note that describes and gives the code that makes this function equivalent to Periodic, but mention in that case using that Periodic might be a bit faster.

Also, the function $u(x)$ is defined, but without context I'd have to know where to look this up. Could you point to a reference or maybe add a bit more detail here (or both)?

Have addressed these in latest commit.

Thank you! Super nice

bwengals · 2023-06-02T19:28:23Z

I played around with the lengthscale * 4 issue, and it looks to me like the way the current Periodic and what GPflow does makes more sense than in the original derivation. I would guess this is the reason for the change, right? Then the lengthscale on Periodic and ExpQuad and the other stationary kernels have the same interpretation and scaling. Is this what you were showing in the plot above? There is a note in the Periodic docstring that points this out.

Here's the code I used to take a look at this:

ls = 0.5
period = 5

cov1 = pm.gp.cov.ExpQuad(1, ls=ls)
K1 = cov1(t)
s1 = pm.draw(pm.MvNormal.dist(mu=np.zeros(len(t)), cov=K1), 10)

cov3 = pm.gp.cov.Periodic(1, ls=ls, period=period)
K3 = cov3(t)
s3 = pm.draw(pm.MvNormal.dist(mu=np.zeros(len(t)), cov=K3), 10)

plt.plot(t, s1.T, color="b");
plt.plot(t, s3.T, color="k");
plt.xlim([0, 5]); # one period

Then for the timing tests I did here's the code I used:

import pymc as pm
import pytensor.tensor as pt
import numpy as np
import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings("ignore", category=UserWarning)

### new
cov_exp = pm.gp.cov.ExpQuad(1, ls=ls/4)
cov1 = pm.gp.cov.WrappedPeriodic(1, cov_func=cov_exp, period=period)
cov1(t).eval();  # eval once so pytensor compilation doesnt count in timing

### using WarpedInput
def mapping(x, T):
    c = 2.0 * np.pi * (1.0 / T)
    u = pt.concatenate((pt.sin(c * x), pt.cos(c * x)), 1)
    return u
cov_exp2 = pm.gp.cov.ExpQuad(2, ls=ls)
cov2 = pm.gp.cov.WarpedInput(1, cov_func=cov_exp, warp_func=mapping, args=(period, ))
cov2(t).eval();

### Existing periodic
cov3 = pm.gp.cov.Periodic(1, ls=ls, period=period)
cov3(t).eval();

Then using @timeit magic:

But that's just my machine, and didn't try using jax or numba.

Sounds great - happy to look into it. In a follow-on PR?

Yup totally OK of course

bwengals · 2023-06-02T19:28:58Z

Also, could you add a test for the new kernel?

jahall · 2023-06-05T19:46:15Z

I played around with the lengthscale * 4 issue, and it looks to me like the way the current Periodic and what GPflow does makes more sense than in the original derivation. I would guess this is the reason for the change, right? Then the lengthscale on Periodic and ExpQuad and the other stationary kernels have the same interpretation and scaling. Is this what you were showing in the plot above?

Certainly the current periodic definition / gpflow produces variations that more closely follow the non-periodic version...but even then the variations still seem a little more rapid than the non-periodic version...but that's just from eye-balling. Either way, I think I'll drop the constant as then the definition is at least consistent within pymc, and more inline with non-periodic as you say.

jahall · 2023-06-05T20:02:47Z

Then for timing tests here's the code I used.

Ok, I copy / pasted your code, set t = np.linspace(0, 5, 500)[:, None] and still see warped input taking more than double the time

...increasing to 1000 data points I get cov1=87ms, cov2=426ms and cov3=73ms. The above with pytensor=2.11.3, running on Windows with no jax / numba.

codecov · 2023-06-14T00:02:28Z

Codecov Report

Merging #6742 (5fe3fba) into main (4a65148) will decrease coverage by 11.91%.
The diff coverage is 100.00%.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #6742       +/-   ##
===========================================
- Coverage   92.05%   80.15%   -11.91%     
===========================================
  Files          95       95               
  Lines       16280    16298       +18     
===========================================
- Hits        14986    13063     -1923     
- Misses       1294     3235     +1941

Impacted Files	Coverage Δ
pymc/gp/cov.py	`97.99% <100.00%> (+0.08%)`	⬆️

... and 32 files with indirect coverage changes

bwengals · 2023-06-14T00:19:56Z

Hey @jahall so sorry for the delay on my end, got caught up in other stuff. This looks awesome -- will approve when all the checks are green.

Funny thing about the timings, must just be the systems we're on? Either way the new kernel is faster so that's nice.

I think I do see what you mean about the factor of four, the variation might be a bit faster but I have to squint... hard to say.

bwengals

lgtm

jahall · 2023-06-14T06:27:45Z

@bwengals Merging the type changes first will be a bit easier for me I reckon..

ferrine · 2023-06-22T09:43:12Z

Does this kernel play well with HSGP? Just curious

jahall · 2023-06-22T13:05:17Z

Does this kernel play well with HSGP? Just curious

@ferrine Sadly there is not a power spectral decomposition for the SE-based periodic kernel (a requirement for a kernel to be compatible with the HSGP approximation)...however, after a bit of digging around in the Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming paper referenced in the HSGP docs, it seems that Appendix B has an approximate method for doing a decomposition for the periodic kernel...worth investigating!

jahall · 2023-07-19T15:25:39Z

@ricardoV94 I think this is good to go too...just not sure about the random code-cov issue...

ricardoV94 · 2023-07-19T15:51:38Z

@ricardoV94 I think this is good to go too...just not sure about the random code-cov issue...

That happens too often

ricardoV94 · 2023-07-19T15:53:07Z

Thanks @jahall!

Joseph Hall added 2 commits May 29, 2023 15:02

Introducing the WrappedPeriodic kernel

fd1ac33

Move the 4 outside of sum

cc4eea9

jahall marked this pull request as ready for review June 1, 2023 10:12

bwengals self-requested a review June 2, 2023 01:58

bwengals reviewed Jun 2, 2023

View reviewed changes

Joseph Hall added 3 commits June 5, 2023 21:04

Inherit input_dim and active_dims from cov_func

59d59cd

Improved docstring

d3b0586

Improved docstring and added tests

9eee89a

jahall requested a review from bwengals June 5, 2023 20:43

Fix test class name

34d3ee8

bwengals approved these changes Jun 14, 2023

View reviewed changes

Joseph Hall added 4 commits July 19, 2023 14:10

Merge latest master and type updates

3c74942

Fix mypy, linting and test cov

9f36319

Add typing to WrappedPeriodic

539061c

Bring WrappedPeriodic signature in line with other Covariance funcs

5fe3fba

ricardoV94 added enhancements GP Gaussian Process labels Jul 19, 2023

ricardoV94 changed the title ~~Wrapped Periodic Kernel~~ Add GP Wrapped Periodic Kernel Jul 19, 2023

ricardoV94 merged commit 82c6318 into pymc-devs:main Jul 19, 2023

Uh oh!

Add GP Wrapped Periodic Kernel #6742

Add GP Wrapped Periodic Kernel #6742

Uh oh!

Conversation

jahall commented May 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

New features

Performance

Example outputs

Uh oh!

jahall commented May 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bwengals commented Jun 2, 2023

Uh oh!

jahall commented Jun 2, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bwengals Jun 2, 2023

Choose a reason for hiding this comment

Uh oh!

jahall Jun 5, 2023

Choose a reason for hiding this comment

Uh oh!

bwengals Jun 2, 2023

Choose a reason for hiding this comment

Uh oh!

jahall Jun 5, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bwengals Jun 14, 2023

Choose a reason for hiding this comment

Uh oh!

bwengals Jun 2, 2023

Choose a reason for hiding this comment

Uh oh!

jahall Jun 5, 2023

Choose a reason for hiding this comment

Uh oh!

bwengals Jun 14, 2023

Choose a reason for hiding this comment

Uh oh!

bwengals commented Jun 2, 2023

Uh oh!

bwengals commented Jun 2, 2023

Uh oh!

jahall commented Jun 5, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jahall commented Jun 5, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Jun 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

bwengals commented Jun 14, 2023

Uh oh!

bwengals left a comment

Choose a reason for hiding this comment

Uh oh!

jahall commented Jun 14, 2023

Uh oh!

ferrine commented Jun 22, 2023

Uh oh!

jahall commented Jun 22, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jahall commented Jul 19, 2023

Uh oh!

ricardoV94 commented Jul 19, 2023

Uh oh!

ricardoV94 commented Jul 19, 2023

Uh oh!

Uh oh!

jahall commented May 29, 2023 •

edited

Loading

jahall commented May 29, 2023 •

edited

Loading

jahall commented Jun 2, 2023 •

edited

Loading

jahall Jun 5, 2023 •

edited

Loading

jahall commented Jun 5, 2023 •

edited

Loading

jahall commented Jun 5, 2023 •

edited

Loading

codecov bot commented Jun 14, 2023 •

edited

Loading

jahall commented Jun 22, 2023 •

edited

Loading