Add HSGP Latent GP approximation #6458
Conversation
Codecov Report
Additional details and impacted files@@ Coverage Diff @@
## main #6458 +/- ##
==========================================
+ Coverage 92.02% 92.04% +0.01%
==========================================
Files 92 93 +1
Lines 15563 15719 +156
==========================================
+ Hits 14322 14468 +146
- Misses 1241 1251 +10
|
Co-authored-by: Michael Osthege <[email protected]>
|
Have a type hint question, what's the best type for something like |
It's true that we don't type hint "tensor-like" in most places. We should probably add a convenient type alias somewhere.. Even if you decide to not type hint the tensor-like, type hints for other kwargs in the signature should be added, because IIRC mypy automatically skips functions that don't have any type hints in their signature |
|
Ah thanks @michaelosthege that makes sense, and explains why I couldn't find that type somewhere else in the codebase. Going with your suggestion |
|
@bwengals where do you see this PR w.r.t. to the finish 🏁 line? |
|
Pretty close! I'd thought things were wrapping up, but taking a few iterations to settle on the exact API. Other than your suggestions (thank you btw), I think I'd like to improve the tests for |
Re API I'd be interested in drawing posterior predictive samples at high resolution Maybe you already have one, but a test for this use case would be great |
Co-authored-by: Michael Osthege <[email protected]>
Co-authored-by: Michael Osthege <[email protected]>
|
thanks a ton for your help on that @michaelosthege. I refactored a bit and tried to take your suggestions. I think handling either one of |
pymc/gp/cov.py
Outdated
| elif np.asarray(value).squeeze().shape == (): | ||
| return np.squeeze(value) | ||
| elif isinstance(value, numbers.Real): | ||
| return value |
There was a problem hiding this comment.
Is the third branch not test-covered because np.asarray(value).squeeze().shape == () applies to numbers already?
There was a problem hiding this comment.
yep, looks like it. I'll remove it
pymc/gp/hsgp.py
Outdated
| elif self._parameterization == "centered": | ||
| return self.mean_func(Xnew) + phi[:, i:] @ beta | ||
|
|
||
| def conditional(self, name: str, Xnew: TensorVariable, *args, **kwargs): |
There was a problem hiding this comment.
| def conditional(self, name: str, Xnew: TensorVariable, *args, **kwargs): | |
| def conditional(self, name: str, Xnew: TensorVariable, **kwargs): |
The args are not used within the function!
pymc/gp/hsgp.py
Outdated
| Optional arguments such as `dims`. | ||
| """ | ||
| fnew = self._build_conditional(Xnew) | ||
| return pm.Deterministic(name, fnew, dims=kwargs.get("dims")) |
There was a problem hiding this comment.
| return pm.Deterministic(name, fnew, dims=kwargs.get("dims")) | |
| return pm.Deterministic(name, fnew, **kwargs) |
This way other kwargs will be forwarded too. If you don't want that, I suggeste having only dims in the signature and not a **kwargs
pymc/gp/hsgp.py
Outdated
| def prior(self, name: str, X: TensorVariable, *args, **kwargs): | ||
| R""" | ||
| Returns the (approximate) GP prior distribution evaluated over the input locations `X`. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| name: string | ||
| Name of the random variable | ||
| X: array-like | ||
| Function input values. | ||
| dims: None | ||
| Dimension name for the GP random variable. | ||
| """ |
There was a problem hiding this comment.
See similar comments below:
- args/kwargs getting dropped silently
- Parameters section formatting
- Docstring doesn't match signature
There was a problem hiding this comment.
Right, I added the args/kwargs issues you're pointing to make mypy pass. The code originally just passed dims like the docstring says. It's also the only input used. Without doing this mypy complains that the HSGP function signatures don't match the signatures of the base class. Is there a third option? I could have HSGP not be a subclass the base gp class, which seems weird to me because it is a subclass.
Is it actually a bad thing if args/kwargs get dropped silently? At least with kwargs isnt that somewhat expected? The docstring says it just takes dims, so what else should a user expect?
There was a problem hiding this comment.
I'd expect it to take any of the kwargs of a typical PyMC distribution, for example dims, initval...
Can you make the signature of the base class more specific? (Specifying just dims, not *args, **kwargs..)
Otherwise you can go with a # type: ignore comment, but be aware that this violates the Liskov subtitution principle. Best check https://mypy.readthedocs.io/en/stable/common_issues.html#incompatible-overrides
There was a problem hiding this comment.
Great, added the ignores. Maybe in a later PR I can refactor the GP module to not violate the liskov substitution principle.
There was a problem hiding this comment.
args/kwargs arent getting dropped silently now, and the docstrings match the signature. I also added a Returns text to the docstring for prior_linearized.
I held off on adding Returns to the docstrings for the rest of the methods because the rest of the GP module methods don't have a returns sections either. I think it's good to have of course, but out of scope to add it for everything here. I also would like to defer refactoring how Base is used so it doesn't violate the Liskov sub. principle you pointed out because that also touches the entire gp submodule. It would be good to fix, but it's been that way for several years now without issue because users don't use Base. But I agree, structurally it could be better.
pymc/gp/cov.py
Outdated
| r""" | ||
| Base class for all kernels/covariance functions. | ||
| def _verify_scalar(value): | ||
| if ( |
There was a problem hiding this comment.
Why not call at/np.squeeze and capture the errors that both emit when the inputs are not allowed to be squeezed?
pymc/gp/cov.py
Outdated
| isinstance(value, pytensor.compile.SharedVariable) | ||
| and value.get_value().squeeze().shape == () | ||
| ): | ||
| return at.squeeze(value) |
There was a problem hiding this comment.
This is slightly incorrect. If I pass pt.shared(np.ones(1)) it will fail and not be captured by your error.
If I pass pt.shared((1), shape=(1,)) then it will work.
I suggest just calling squeeze directly. Also inputs could be constants (e.g., from pm.ConstantData) and would not meet either branch, no?
There was a problem hiding this comment.
This functions is only used by exponentiated kernels now (which aren't usable by HSGP). I tried to use _verify_scalar for other HSGP stuff earlier in this PR (that I cant recall) but it got refactored out.
Thanks for catching this potential regression. In practice though, I think its extremely unlikely anyone would do anything more sophisticated than
cov_func = eta**2 * pm.gp.cov.ExpQuad(1, ls)**2 # or cubed or something
_verify_scalar is only handling the power, **2. I'm also OK with not checking/supporting type checking here. Basically usage wise there's really no chance that exponent is anything other than "2" or maybe "3". I've not see this really in the wild in models, more that its just kinda neat that you can exponentiate kernels (because products of kernels are kernels) so why not support it.
How about I roll things back to how they are in master for this, which is,
def __pow__(self, other):
if (
isinstance(other, pytensor.compile.SharedVariable)
and other.get_value().squeeze().shape == ()
):
other = at.squeeze(other)
return Exponentiated(self, other)
elif isinstance(other, Number):
return Exponentiated(self, other)
elif np.asarray(other).squeeze().shape == ():
other = np.squeeze(other)
return Exponentiated(self, other)
raise ValueError("A covariance function can only be exponentiated by a scalar value")The errors aren't captured, but users will see what happens when squeeze is attempted.
There was a problem hiding this comment.
No strong preference, but why is this not enough?
def __pow__(self, other):
other = as_tensor_variable(other).squeeze()
if not other.ndim == 0:
raise ValueError(...)
return Exponentiated(self, other)There was a problem hiding this comment.
alright that's why they pay you the big bucks, changed it to this
Since pymc-devs#6458, Covariance is now the base class for kernels/covariance functions with input_dim and active_dims, which does not include WhiteNoise and Constant kernels.
This PR add's the basics required for the HSGP GP approximation. It replaces #6036, much thanks to @ferrine for rebasing and cleaning that PR up! This PR also modifies the covariance function classes to allow an additional method
power_spectral_densitywhich is needed for HSGP (or random fourier features). Any covariance function that defines apower_spectral_densitymethod will work. Like with the other GP implementations, one can also add covariances,eta1**2 * pm.gp.cov.ExpQuad(2, ls=ls1) + eta2**2 * pm.gp.cov.Matern52(2, ls=ls2). The implementation also works for any number of input dimensions.Not part of this PR, but plan on adding in later:
PeriodiccovariancemandLorc.Checklist
Major / Breaking Changes
New features
Bugfixes
Documentation
Maintenance