Some improvement and additions to MO kernels

[Crown421]

I have been using KernelFunctions in my research, but have been building some more or less hacky solution on top of it to facilitate Multi-Output GPs.

Here I would like to move some of that work into this package to streamline my own code and get it ready for publication.

The following is a start, with more things to come soon. I have added tests as I thought appropriate, but ran into a segfault caused by the AD test function. I hope that this was a local problem.

[st--]

Thanks for your interest in KernelFunctions:) could you provide a bit more context on what hadn't been working for you / what you're trying to achieve / work around with your changes?

[Crown421]

Thanks for your interest in KernelFunctions:) could you provide a bit more context on what hadn't been working for you / what you're trying to achieve / work around with your changes?

Sorry, of course. For this set of changes, my changes aim to make sure that the kernels are type stable, as well as propose a specific implementation for the kernelmatrix function that is substantially faster than the general method.
Regarding the first point, I had noticed that the independent kernel in its current form always return a Float64, regardless of the input type, even though the scalar kernels return a type consistent with the input.

In additional commits, I would like to introduce a new type of MO kernel, the Group-Invariant-Matrix (GIM) kernel (Reisert, 2007).

[willtebbutt]

Thanks for uploading the timing data -- pretty clear win with the new versions. This kind of optimisation is exactly the kind of PR we're keen to have, so please rename kernelmatrix2 to kernelmatrix when you get a minute so that we can push this PR forward.

[st--]

How about similar definitions for kernelmatrix! using kron!?

This is probably a source of inconsistency that is likely to pop up again in the future: would it perhaps be better to generally define kernelmatrix in terms of kernelmatrix!, then people only have to implement kernelmatrix! - unless of course you want to return lazy versions (e.g. Kronecker.jl) in the not-in-place method?

[willtebbutt]

That would be a nice way to do it, but unfortunately kernelmatrix! tends not to play nicely with reverse-mode AD due to a lack of support for in-place operations. Consequently, we need the kernelmatrix implementations to be free of in-place operations.

How about similar definitions for kernelmatrix! using kron!?

This makes sense to me though.

[Crown421]

Worked in most of the changes, will look at the lazy Kronecker product now.

[Crown421]

It seems that Kronecker.jl is already an optional dependency of KernelFunctions. I have to think about how to implement the matrix kernels in a way that works with the optional dependency.

[willtebbutt]

It seems that Kronecker.jl is already an optional dependency of KernelFunctions. I have to think about how to implement the matrix kernels in a way that works with the optional dependency.

Please do feel free not to worry about it for now if it's not high on your priority list -- we can always revisit it later.

[st--]

@Crown421 you'll have to accept reviewdog's formatting suggestions (or run JuliaFormatter locally):)

[Crown421]

Ok, did some reworking using traits, now there is less code duplication, and getting the lazy Kronecker product has to be made explicit.

[willtebbutt]

Many thanks for reducing the size of this PR, it's now much more easily digestible.

I've left a couple of minor style-related things, but I think they're fairly inconsequential. The main thing I'd like to get sorted is the argument names for _kronkernelmatrix, and the function's name itself.

[willtebbutt]

I agree with @st-- regarding argument names in particular -- Kfeatures and Koutputs are quite a bit more informative than Ktmp and B. The function name also makes more sense to me.

[Crown421]

Ok, I have removed the in place tests, and changed a few more small things.

I agree with @st-- regarding argument names in particular -- Kfeatures and Koutputs are quite a bit more informative than Ktmp and B. The function name also makes more sense to me.

Re function names, the helper function should not really be visible to a user (which I tried to indicate with the leading _, as is done in Distances.jl). I am not sure how to improve the function name further.
Regarding argument names, I looked at a few papers, and it is not clear to me why B could be called Koutputs. The paper cited in the docs uses B and calls it the "coregionalization matrix", so I would personally prefer to keep it consistent with the cited paper.

[willtebbutt]

Regarding argument names, I looked at a few papers, and it is not clear to me why B could be called Koutputs. The paper cited in the docs uses B and calls it the "coregionalization matrix", so I would personally prefer to keep it consistent with the cited paper.

The rationale is that B represents the "between output" kernel matrix, while Ktmp represents the "between feature" kernel matrix (features as opposed to inputs). I can see your point though, so we'll stick with the current argument names.

Could you please remove the temporary_script.jl file and bump the patch version. I'll approve once that's done, and merge + tag after that once CI passes.

[willtebbutt]

Actually, one other quick thing.

I am not sure how to improve the function name further.

I think sticking the word "helper" on the end was a good move. Could you please rename to _kernelmatrix_kron_helper though? We go with underscores wherever possible, as mandated by the style guide we follow (somehow kernelmatrix slipped through, and now we're stuck with it).

[Crown421]

Alright, done all of those things. I am sorry that this was such a long process, I learned quite a bit though. Will do the pull requests in a bit.
@willtebbutt

[willtebbutt]

No problem at all. Thanks for sticking with it and working through everything. The first PR always takes as while -- your follow up PRs will almost certainly progress much faster now that we all know what's going on better.

[willtebbutt]

@Crown421 looks like there's some usings missing in the tests or something like that.

[Crown421]

@Crown421 looks like there's some usings missing in the tests or something like that.

I think I fixed it?

[willtebbutt]

LGTM. Will merge + tag.