Generalized Periodic Kernel

[jahall]

This PR is to introduce a generalized form of the periodic kernel. There is no need for us to be stuck with only the periodic form of the squared exponential when we can easily extend to any isotropic stationary kernel e.g. the materns and rational quadratic.

I implemented a similar change in GPFlow a couple of years ago so the changes in this PR very much reflect that. See e.g. periodic.py and stationaries.py from that repo.

NOTE: This is currently in draft mode as I want to get input on whether we should try to extend/adapt the current Periodic kernel (as was done in GPFlow) rather than having a new kernel? And input on the implementation itself...which could almost certainly be improved.

Major / Breaking Changes

• Backward compatible

New features

• Stationary kernel now allows the specification of full_r2 (or full_r) instead of having to override full directly
• Periodic kernel can now accept a base_kernel_class kwarg (defaulting to ExpQuad) and optional base_kernel_kwargs parameters
• Increased type hint coverage for cov funcs

Bugfixes

• ...

Documentation

• ...

Maintenance

• ...

---

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

[welcome]

💖 Thanks for opening this pull request! 💖 The PyMC community really appreciates your time and effort to contribute to the project. Please make sure you have read our Contributing Guidelines and filled in our pull request template to the best of your ability.

[lucianopaz]

Thanks @jahall for the PR! It looks very nice. I added a couple of nitpicky comments above, but I do have one bigger question. Do you want GeneralizedPeriodic to take an instance or a class as base_kernel? I think that the current implementation assumes the former, but I think that the latter might be better. We would have to create the base_kernel instance during initialization, but that doesn't seem like a bad choice to me. That would allow us to keep a consistent __init__ signature across subclasses and also ensure that the base_kernel has consistent input_dim and active_dims

[lucianopaz]

I would try to mimic the initialization signature that the Periodic kernel uses.

Suggested change

	def __init__(self, base_kernel: IsotropicStationary, period):
	def __init__(self, input_dim, base_kernel: IsotropicStationary, period, ls=None, ls_inv=None, active_dims=None):

I'm not sure if we want to keep the ls, ls_inv, etc or just use a **kwargs unpacking though.

By the way, is your idea to pass in an instance of an IsotropicSationary kernel, or the class? In the former case, you would need to make sure that the input_dim, and other parameters match with the base_kernel. In the latter case, you can construct the base_kernel inside of __init__, and make sure that you get consistent active_dims and input_dim.

[jahall]

I like where you're going with this. If we went with this approach could we then simply "enhance" the existing Periodic kernel instead of creating a new one? e.g.

def __init__(self, input_dim, period, ls=None, ls_inv=None, active_dims=None, base_kernel_class: Type[IsotropicKernel] = ExpQuad):
    super().__init__(input_dim, ls, ls_inv, active_dims)
    self.base_kernel_class = base_kernel_class
    self.period = period

@property
def base_kernel(self) -> IsotropicKernel:
    return self.base_kernel_class(self.input_dim, ls=self.ls, ls_inv=self.ls_inv, active_dims=self.active_dims)

...where the instance is instantiated on-demand so we're not storing input_dim, ls, ls_inv in two places? Also means if someone changes any of the attributes on the class, they will be sure to be reflected in the base kernel. But I don't have super strong feelings..

[jahall]

@lucianopaz - re-worked the change such that the original Periodic kernel is now extended in a backward compatible way. And the base kernel is passed as a class.

[lucianopaz]

I really dislike dimshuffle! Could we switch to numpy syntax here?

Suggested change

	f1 = X.dimshuffle(0, "x", 1)
	f2 = Xs.dimshuffle("x", 0, 1)
	f1 = X[..., :, None, :]
	f2 = Xs[..., None, :, :]

you can also use numpy.newaxis if you find that more readable

[ricardoV94]

There's also pytensor.tensor.expand_dims

[jahall]

Yes absolutely, I had just copied and pasted what was currently in Periodic

[lucianopaz]

Should we keep the multiplying 4 or not? I see that GPflow chose to leave it out and have it absorbed by the ls.

My ruff derivation of the argument that goes into r is

$$ y = 2\pi x / period $$

$$ y' = 2\pi x' / period $$

$$ r^2 = (\cos(y) - \cos(y'))^2 + (\sin(y) - \sin(y'))^2 $$

$$ r^2 = 2 - 2\cos(y - y') $$

$$ r^2 = 2 \left[sin\left(\frac{y - y'}{2}\right)^2 + cos\left(\frac{y - y'}{2}\right)^2 - \cos(y - y') \right] $$

$$ r^2 = 4 sin\left(\frac{y - y'}{2}\right)^2 $$

$$ r^2 = 4 sin\left(\frac{\pi(x - x')}{period}\right)^2 $$

From looking at Periodic it seems as though we had dropped it long ago, so it might be irrelevant

[jahall]

I'd prob just go with what is consistent with current implementation.

[dehorsley]

This looks interesting! Sorry for the nit-picking, but I just want to point out those kernels technically aren't isotropic, as the length scale ls can be a vector 😄

[jahall]

@dehorsley This is a great point - let me try and roll this functionality into the Stationary kernel, as it seems to apply to all current stationary kernels...though really it should be limited to isotropic or ARD kernels

[bwengals]

How different is this from this example here? https://www.pymc.io/projects/examples/en/latest/gaussian_processes/GP-MeansAndCovs.html#constructing-periodic-using-warpedinput

[bwengals]

Also, btw, thanks for adding the typing! long overdue

[jahall]

@bwengals In terms of the output values it should be no different that using warped inputs + stationary kernel. It would just be a more efficient implementation, similar to the note here

[ricardoV94]

Nitpick, but python never allows __init__ to return anything, so it must always be None.

I've never seen __init__ return type annotated

[jahall]

My reasoning was as per this (referencing PEP-484) i.e. it allows mypy to ensure __init__ didn't return anything.

[ricardoV94]

According to this comment in that thread you linked, mypy assumes None by default if there are any other typed arguments in __init__: https://stackoverflow.com/a/67693892

[jahall]

Seems it assumes None provided one or more of the args are type-annotated - otherwise its ignored from type-checking. The PEP advises

Note that the return type of init ought to be annotated with -> None. [...]

However, I don't feel too strongly about it, so happy to remove these annotations.

[bwengals]

Gotcha @jahall. Regardless of implementation, I'm all for this functionality being easier to find/use for people. I think it'd be nice to double check that the efficiency claim is true though. Maybe it's cleaner to actually implement Periodic using WarpedInput then? Either way I think the proposed function signature of the improved Periodic should follow that of WarpedInput, Gibbs or ScaledCov.

There are other uses where it'd be great to have the full_r and full_r2 methods available too, like using distances other than Euclidean. Maybe though full_from_distance is a more explicit name? And are both full_r and full_r2 needed?

One thought too if you don't mind, is if you split out the typing improvements into a separate PR I think we could pretty much get that merged in immediately.

[jahall]

@bwengals I've separated the type-hint changes out into this PR which should be ready for review.

I agree that the ideal signature would be inline with WarpedInput, etc i.e. where we pass an already-instantiated base kernel. However, for backward compatibility I assume we'd need to create a separate GeneralizedPeriodic kernel...or maybe WrappedPeriodic?

Also agree that full_r2/full_r should be better named. My current implementation of the period does need to know the difference...but let me revisit that...