Add changepoint kernel #357

jakeyeung · 2023-08-10T07:12:55Z

Type of changes

Checklist

I've formatted the new code by running poetry run pre-commit run --all-files --show-diff-on-failure before committing.
I've added tests for new code.
I've added docstrings for the new code.

Description

Add change point kernel for two RBFs.

Issue Number: N/A
Related to Q&A in Discussion #337

github-actions

Thank you for opening your first PR into GPJax!

If you have not heard from us in a while, please feel free to ping @gpjax/developers or anyone who has commented on the PR. Most of our reviewers are volunteers and sometimes things fall through the cracks.

You can also join us on Slack for real-time discussion.

For details on testing, writing docs, and our review process, please see the developer guide

We strive to be a welcoming and open project. Please follow our Code of Conduct.

thomaspinder

Nice work @jakeyeung ! I've left some comments that we should resolve before merging. It will also need some unit tests before we're able to merge.

thomaspinder · 2023-08-12T04:27:43Z

gpjax/kernels/stationary/changepoint.py

+#!/usr/bin/env python
+"""
+AUTHOR: Jake Yeung ([email protected])
+CREATED ON:  2023-07-28
+LAST CHANGE: see git log
+LICENSE:     Apache License
+"""
+


Could you change this to the standard GPJax boilerplate please?

Should be done now.

gpjax/kernels/stationary/changepoint.py

kernels. Update new imports, check they are sorted with isort

jakeyeung

Replaced ChangePointRBF to SwitchKernel, which accepts a list of two kernels.

jakeyeung · 2023-08-14T14:37:04Z

gpjax/kernels/stationary/changepoint.py

+#!/usr/bin/env python
+"""
+AUTHOR: Jake Yeung ([email protected])
+CREATED ON:  2023-07-28
+LAST CHANGE: see git log
+LICENSE:     Apache License
+"""
+


Should be done now.

gpjax/kernels/stationary/changepoint.py

thomaspinder · 2023-08-16T03:48:28Z

Nice work @jakeyeung - this PR is looking good. If we can add some unit tests to this, then I think we'll be good to merge.

jakeyeung · 2023-08-25T09:20:57Z

Nice work @jakeyeung - this PR is looking good. If we can add some unit tests to this, then I think we'll be good to merge.

Hi @thomaspinder, glad to know it's progressing. Any ideas of some unit tests that would be good, both ones standard for all kernels as well as ones specific to the changepoint?

kernels. Update new imports, check they are sorted with isort

…changepoint_kernel

jakeyeung · 2023-10-10T12:17:32Z

Hi @thomaspinder, sorry for the delay. I have now added a test_changepoint_kernel() function under test_stationary.py. Hope everything is alright now to merge.

daniel-dodd · 2023-10-10T12:29:35Z

Nice work @jakeyeung! Just to check: have you rebased from the main branch - the change log seems huge here? Would just make it easier to see what you have added. :)

jakeyeung · 2023-10-10T20:14:30Z

Nice work @jakeyeung! Just to check: have you rebased from the main branch - the change log seems huge here? Would just make it easier to see what you have added. :)

Hi Daniel, yeah I think what happened was that git fetch upstream and git rebase upstream/main updated a bunch of files, and then after poetry run pre-commit run --all-files changed a lot of files that weren't added by me.

I reverted one step back before this change to a bunch of files not added by me due to poetry run pre-commit run --all-files, and it should now just contain changes to three files.

daniel-dodd · 2023-10-10T22:33:37Z

gpjax/kernels/stationary/changepoint.py

+    tswitch: ScalarFloat = param_field(
+        jnp.array(1.0), bijector=tfb.Identity(), trainable=False
+    )


Is the changepoint location itself not learnt here (trainable=False)?

Also what happens if we have multiple changepoints (does this kernel work recursively by passing a ChangePoint kernel to kernels?)

The idea started with implementing a basic change-point kernel similar to what GPflow has (See Change-point section in https://gpflow.github.io/GPflow/develop/notebooks/getting_started/kernels.html).

When I first implemented it, I was thinking of a fixed parameter, but I don't see any reason why it has to be that way.

At the moment this piece of code only works with two kernels, so one changepoint, but one can also imagine extending this to a list of N kernels with (N-1) changepoint locations.

daniel-dodd · 2023-10-10T22:36:48Z

gpjax/kernels/stationary/changepoint.py

+
+@dataclass
+class ChangePoint(AbstractKernel):
+    r"""A change point kernel


Is there a citation for your implementation of the kernel? We could add this to our citation.py. Also probably good to give a paper in the docstring. :)

The idea started with the GPflow documentation, but I guess if you're thinking about a paper it would be Saatci et al 2010 : (https://icml.cc/Conferences/2010/papers/170.pdf).

I have now added it the Saatci et al 2010 reference to citation.py and the doc strings.

daniel-dodd · 2023-10-10T22:38:32Z

gpjax/kernels/stationary/changepoint.py

+    def __call__(
+        self,
+        x: Float[Array, " D"],
+        y: Float[Array, " D"],
+    ) -> ScalarFloat:


Maybe I'd be clear the intended dimensions is one here? (correct me if I'm wrong). i.e., Float[Array, " 1"] instead of Float[Array, " D"] .

They should be 1-dimensional, I've now fixed now.

daniel-dodd · 2023-10-10T22:44:44Z

gpjax/kernels/stationary/changepoint.py

+
+        def get_function_index(x, y, tswitch=1):


I would probably drop the default argument here.

Suggested change

def get_function_index(x, y, tswitch=1):

def get_function_index(x, y, tswitch):

daniel-dodd · 2023-10-10T22:52:21Z

gpjax/kernels/stationary/changepoint.py

+            r"""
+            Specify four possible indices given x, y, and tswitch.
+
+            Args:
+                    x: Left hand argument of kernel function's call
+                    y: Right hand argument of kernel function's call
+                    tswitch: point at which to change to a different kernel
+            """
+            # Four possible indexes: 0, 1, 2, 3
+            indx = 3  # if indx doesn't get set to 0, 1, or 2, then default 3
+            # lessthan means that at tswitch, you are already switched
+            cond1 = jnp.less(x, tswitch)
+            cond2 = jnp.less(y, tswitch)
+            indx = jnp.where(jnp.logical_and(cond1, cond2), 0, indx)
+            indx = jnp.where(jnp.logical_and(jnp.invert(cond1), cond2), 1, indx)
+            indx = jnp.where(jnp.logical_and(cond1, jnp.invert(cond2)), 2, indx)
+            return indx.squeeze().astype("uint8")


I would comment somewhere here (or use the docstring) the purpose - i.e., the index is going to select which kernel function to compute.

But I am a bit confused as to why we need 4 indexes? Why don't we shorten to three? As the flst could be shortened to three given k_zero is repeated? So e.g., we could have flst = [self.kernels[0].__call__, k_zero, self.kernels[1].__call__]. Just curious to know. :)

I have now added in words the purpose of this function. I have also renamed flst to kernel_options to be more explicit.

The four indices can be reduced to three. I found the code easier to read to have four indices corresponding to four quadrants in a 2D covariance matrix. But returning just three indices might make more sense if this function accepts more than two kernels.

e.g. visualizing the gram matrix of the ChangePoint kernel you can see the four quadrants. You can imagine with more than two kernels it would be a bunch of blocks along the diagonal of the matrix, with 0 everywhere else.

daniel-dodd · 2023-10-10T22:55:19Z

gpjax/kernels/stationary/changepoint.py

+        def k_zero(x, y):
+            r"""Return 0 covariance"""
+            out = jnp.float64(0)
+            return out.squeeze()


This is nice. But the only concern I have is with jnp.float64 potentially clashing with JAX's default float32. Could this not work with jnp.array(0.0) (or even 0.0)? Or does this break the jax.lax.switch or something? :)

If problems persist. You could always use a Constant(0.0) kernel here.

I don't know the details of jax enough to know how the different types play with other types and conditionals. What would be a good way to test which type is best to put in here?

daniel-dodd · 2023-10-10T23:02:36Z

gpjax/kernels/stationary/changepoint.py

+        flst = [self.kernels[0].__call__, k_zero, k_zero, self.kernels[1].__call__]
+        K = jax.lax.switch(indx, flst, x, y)


Maybe on naming. I don't know what flst means. Maybe something like kernel_options.

Also, why do we need self.kernels[0].__call__ and self.kernels[1].__call__. Does this not work?

kernel_options = self.kernels + [k_zero].

Would this help with extending this kernel to more than two kernels? Or is there some issue with the computing this indexing?

Yeah this would work. I just need to rewrite the get_function_index() so it returns an indx that selects a kernel, otherwise set indx = len(kernels) + 1 to select k_zero.

In order to return the right index, this new get_function_index() would check whether both x and y are greater than or equal to a vector of locations (returning a boolean vector of same size as locations), then return the greatest index for which this is True.

I can try out an initial implementation, but if you have code suggestions would be happy to hear a clean way to code this.

daniel-dodd · 2023-10-10T23:03:51Z

gpjax/kernels/stationary/changepoint.py

+    def __post_init__(self):
+        # Add kernels to a list, flattening out instances of this class therein, as in GPFlow kernels.
+        kernels_list: List[AbstractKernel] = []
+
+        for kernel in self.kernels:
+            if not isinstance(kernel, AbstractKernel):
+                raise TypeError("can only combine Kernel instances")  # pragma: no cover
+
+            if isinstance(kernel, self.__class__):
+                kernels_list.extend(kernel.kernels)
+            else:
+                kernels_list.append(kernel)
+
+        self.kernels = kernels_list


This is great. But we should probably through an error if we specify more than two kernels (with the current implementation)?

daniel-dodd · 2023-10-10T23:04:43Z

tests/test_kernels/test_stationary.py

+@pytest.mark.parametrize("tswitch_param", [0.1, 0.25, 0.5, 0.75, 0.9])
+def test_changepoint_kernel(kernel: AbstractKernel, tswitch_param: float) -> None:
+    n_kerns = 2


This looks good. The only thing is, I would like to see (a) an error is thrown if n_kerns is wrong and (b) this is tested.

Sorry to ask basic question, how does write a test to make sure that a proper error is thrown?

daniel-dodd · 2023-10-10T23:09:26Z

Thanks @jakeyeung. Overall this looks really good. Quite excited to give this a go on some data soon. Mainly asking some clarification questions in the above review. Just the only main points is, I would suggest it is made clear that this is for 1D data. This needs to be in the docstrings and in the typing if possible. I would like to see a check for the number of kernels. If this kernel is just for two kernels (and not more), then maybe instead of kernels as a list, it might be more sensible to have e.g., before_kernel, after_kernel, before and after the changepoint. Maybe I'd also change naming of tswitch to location of changepoint (or locations if this supports multiple kernels as an input).

implementation.

jakeyeung added 2 commits August 10, 2023 08:31

Init commit: changepoint kernel for two RBFs with independent parameters

6bc8889

Rename changepoint_kernel to just changepoint

a94ac98

github-actions bot reviewed Aug 10, 2023

View reviewed changes

thomaspinder requested changes Aug 12, 2023

View reviewed changes

jakeyeung added 6 commits August 14, 2023 14:54

Replace boilerplate with standard GPJax boilerplate for 2023

4e4fab4

Rename function GetFunctionIndex to get_function_index

6f3c01b

Add SwitchPoint kernel, where you can initialize with a list of two

2fc53e4

kernels. Update new imports, check they are sorted with isort

Remove commented out import

34598df

Document what is kernels and what is tswitch in SwitchKernel

21521ad

Remove ChangePointRBF, since SwitchKernel covers the case of two RBFs

46366ad

jakeyeung commented Aug 14, 2023

View reviewed changes

Rename SwitchKernel to ChangePoint to align with literature

d779b93

jakeyeung added 11 commits October 10, 2023 11:58

Init commit: changepoint kernel for two RBFs with independent parameters

b71e5f2

Rename changepoint_kernel to just changepoint

0d25c5f

Replace boilerplate with standard GPJax boilerplate for 2023

f424d6c

Rename function GetFunctionIndex to get_function_index

6dede3c

Add SwitchPoint kernel, where you can initialize with a list of two

25bc0b7

kernels. Update new imports, check they are sorted with isort

Remove commented out import

b961221

Document what is kernels and what is tswitch in SwitchKernel

dd4feea

Remove ChangePointRBF, since SwitchKernel covers the case of two RBFs

f1ee522

Rename SwitchKernel to ChangePoint to align with literature

db264be

Add unit test for changepoint kernel: same as combination kernel.

efeb977

Merge branch 'changepoint_kernel' of github.com:jakeyeung/GPJax into …

19a58da

…changepoint_kernel

jakeyeung force-pushed the changepoint_kernel branch from 533347f to 19a58da Compare October 10, 2023 20:16

daniel-dodd reviewed Oct 10, 2023

View reviewed changes

jakeyeung added 6 commits October 11, 2023 14:20

Add citation for ChangePoint kernel.

443269d

Change inputs x and y to be 1-dimensional as intended.

07f25c3

Remove default argument

309d51d

Document the purpose of get_function_index()

46c52f6

Rename variables flst to kernel_options

127bd6c

Raise error if exactly two kernels are not supplied in current

ddcd673

implementation.

github-actions bot added the stale label Sep 2, 2024

github-actions bot closed this Sep 9, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add changepoint kernel #357

Add changepoint kernel #357

jakeyeung commented Aug 10, 2023 •

edited

Loading

github-actions bot left a comment

thomaspinder left a comment

thomaspinder Aug 12, 2023

jakeyeung Aug 14, 2023

jakeyeung left a comment

jakeyeung Aug 14, 2023

thomaspinder commented Aug 16, 2023

jakeyeung commented Aug 25, 2023

jakeyeung commented Oct 10, 2023

daniel-dodd commented Oct 10, 2023

jakeyeung commented Oct 10, 2023 •

edited

Loading

daniel-dodd Oct 10, 2023 •

edited

Loading

jakeyeung Oct 11, 2023 •

edited

Loading

daniel-dodd Oct 10, 2023

jakeyeung Oct 11, 2023 •

edited

Loading

daniel-dodd Oct 10, 2023

jakeyeung Oct 11, 2023

daniel-dodd Oct 10, 2023

daniel-dodd Oct 10, 2023

jakeyeung Oct 11, 2023 •

edited

Loading

daniel-dodd Oct 10, 2023

daniel-dodd Oct 10, 2023

jakeyeung Oct 11, 2023

daniel-dodd Oct 10, 2023

jakeyeung Oct 11, 2023 •

edited

Loading

daniel-dodd Oct 10, 2023

daniel-dodd Oct 10, 2023 •

edited

Loading

jakeyeung Oct 11, 2023

daniel-dodd commented Oct 10, 2023 •

edited

Loading


	def get_function_index(x, y, tswitch=1):

	def get_function_index(x, y, tswitch):

		flst = [self.kernels[0].__call__, k_zero, k_zero, self.kernels[1].__call__]
		K = jax.lax.switch(indx, flst, x, y)

Add changepoint kernel #357

Add changepoint kernel #357

Conversation

jakeyeung commented Aug 10, 2023 • edited Loading

Type of changes

Checklist

Description

github-actions bot left a comment

Choose a reason for hiding this comment

thomaspinder left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jakeyeung left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

thomaspinder commented Aug 16, 2023

jakeyeung commented Aug 25, 2023

jakeyeung commented Oct 10, 2023

daniel-dodd commented Oct 10, 2023

jakeyeung commented Oct 10, 2023 • edited Loading

daniel-dodd Oct 10, 2023 • edited Loading

Choose a reason for hiding this comment

jakeyeung Oct 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jakeyeung Oct 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jakeyeung Oct 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jakeyeung Oct 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

daniel-dodd Oct 10, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

daniel-dodd commented Oct 10, 2023 • edited Loading

jakeyeung commented Aug 10, 2023 •

edited

Loading

jakeyeung commented Oct 10, 2023 •

edited

Loading

daniel-dodd Oct 10, 2023 •

edited

Loading

jakeyeung Oct 11, 2023 •

edited

Loading

jakeyeung Oct 11, 2023 •

edited

Loading

jakeyeung Oct 11, 2023 •

edited

Loading

jakeyeung Oct 11, 2023 •

edited

Loading

daniel-dodd Oct 10, 2023 •

edited

Loading

daniel-dodd commented Oct 10, 2023 •

edited

Loading