TrajectorySampler amenable to any Kernel/IndVar #606

SebastianPopescu · 2022-08-16T15:21:01Z

Rewrite of Decoupled Trajectory Sampling such that it also works for SeparateIndependent kernels and SeparteIndependentInducingVariables. The goal is to eventually use this for heteroskedastic likelihoods.

hstojic

seems ok in terms of the code, I haven't checked the maths, left some minor comments

more general notes:

there are lots of if-else statements, I wonder if it would be a bit cleaner with kernel type dependent methods - perhaps @uri-granta has some suggestions here?
you are missing unit tests for all of this...

P.S I don't know if it matters but gpflow is at the moment changing how heteroskedastic modeling is done

tests/requirements.txt

trieste/models/gpflux/sampler.py

hstojic · 2022-09-16T09:00:43Z

trieste/models/gpflux/sampler.py


        q_mu = self._layer.q_mu  # [M, P]
        q_sqrt = self._layer.q_sqrt  # [P, M, M]
+
+        #NOTE -- I don't understand why in the original code this approach was used and not the gpflow.covariances.Kuu dispatcher


leave this as a comment for @sebastianober here in github rather than as a comment in the code :)

hstojic · 2022-09-16T09:07:01Z

trieste/models/gpflux/sampler.py

+        # Build the RFF objects
+        if isinstance(self._kernel, list):
+
+            self._rff = defaultdict()


you'll need to provide a type here most likely (see below), though not sure why you need defaultdict here, it seems like you can simply do dict()...

self._rff: DefaultDict[int, RFF] = defaultdict()

SebastianPopescu · 2022-09-16T09:58:48Z

thanks for the comments @hstojic , I will take care of adding unit tests now

I don't think that the change in heteroskedastic modelling in GPflow would influence anything here since this is just TrajectorySampling.

tests/unit/models/gpflux/test_sampler.py

trieste/models/gpflux/architectures.py

hstojic · 2022-09-16T23:52:40Z

trieste/models/gpflux/architectures.py

+    lengthscales = [2.0] * input_dim
+    return SquaredExponential(lengthscales=lengthscales, variance=variance)
+
+def build_constant_input_dim_flexible_deep_gp(X: np.ndarray, num_layers: int, config: Config, separate_case: bool, output_dim: bool) -> DeepGP:


this seems like it should go to gpflux rather than being here in trieste, if its generally useful and not just for trieste needs

I will probably move this architecture to tests..since that's the only place it will be ever used

hstojic · 2022-09-16T23:53:20Z

trieste/models/gpflux/architectures.py

+
+
+@dataclass
+class Config:


Config seems like a copy-paste, do we need to redefine it here?

no..it should just be an import

hstojic · 2022-09-16T23:54:52Z

trieste/models/gpflux/builders.py

+    Build a :class:`~gpflux.models.DeepGP` model with sensible initial parameters. We found the
+    default configuration used here to work well in most situation, but it should not be taken as a
+    universally good solution.


whats the key difference with the ohter builder?

commented a bit below

hstojic · 2022-09-16T23:57:34Z

tests/unit/models/gpflux/test_generalized_sampler.py

+"""
+In this module, we test the *behaviour* of Trieste models against reference GPflux models (thus
+implicitly assuming the latter are correct).
+*NOTE:* Where GPflux models are used as the underlying model in an Trieste model, we should
+*not* test that the underlying model is used in any particular way. To do so would break
+encapsulation. For example, we should *not* test that methods on the GPflux models are called
+(except in the rare case that such behaviour is an explicitly documented behaviour of the
+Trieste model).
+"""


is this file an accident or there should be another test file for sampler?

SebastianPopescu · 2022-09-20T11:02:26Z

trieste/models/gpflux/builders.py

+        whiten=True,  # whiten = False not supported yet in GPflux for this model
+    )
+
+    model = build_constant_input_dim_flexible_deep_gp(query_points, num_layers, config, True, dim_output)


@hstojic this is the only difference..if you can think of a more efficient way of doing this, that would be grand

you said you need build_constant_input_dim_flexible_deep_gp only for the tests, meaning that the builder build_vanilla_flexible_deep_gp is needed only for a test, right? you can then move the whole thing there?

otherwise there are indeed more efficient ways to do this, e.g. having a private function _build_deep_gp which takes a function as an input, either build_constant_input_dim_flexible_deep_gp or build_constant_input_dim_deep_gp and then you have two public functions build_vanilla_flexible_deep_gp and build_vanilla_deep_gp that call the private one with these two different functions

…astian.p/trajectory_Sampler

henrymoss · 2023-01-09T09:30:00Z

trieste/models/gpflux/sampler.py

 class DeepGaussianProcessDecoupledLayer(ABC):
    """
    Layer that samples an approximate decoupled trajectory for a GPflux
-    :class:`~gpflux.layers.GPLayer` using Matheron's rule (:cite:`wilson2020efficiently`). Note
-    that the only multi-output kernel that is supported is a
-    :class:`~gpflow.kernels.SharedIndependent` kernel.
+    :class:`~gpflux.layers.GPLayer` using Matheron's rule (:cite:`wilson2020efficiently`).
+    Supports multi-output kernels of :class:`~gpflow.kernels.SharedIndependent`


why is this a layer rather than just a sampler?

henrymoss · 2023-01-09T09:30:22Z

trieste/models/gpflux/sampler.py

-        ]  # [N, B, L + M, 1]
-
+
+        # TODO -- probably have to re-write unflatten to accomodate for this case as well


henrymoss · 2023-01-09T09:30:28Z

trieste/models/gpflux/sampler.py

+            feature_evaluations.append(
+                unflatten(flattened_feature_evaluations[counter, :, :])[..., None]
+            )  # [N, B, L + M, 1] if Shared or [N, B, L + M, P] if separate
+            # TODO -- check that this is acutally true


henrymoss · 2023-01-09T09:30:33Z

trieste/models/gpflux/sampler.py

+            # TODO -- check that this is acutally true
+        feature_evaluations = tf.concat(feature_evaluations, axis=-1)
+
+        # TODO -- should probably introduce a tf.debugging.assert_equal just to be sure


henrymoss · 2023-01-09T09:31:09Z

trieste/models/gpflux/sampler.py

+        """
        Kmm = self._kernel.K(inducing_points, inducing_points)  # [M, M]
        Kmm += tf.eye(tf.shape(inducing_points)[0], dtype=Kmm.dtype) * DEFAULTS.JITTER
+        """


henrymoss · 2023-01-09T09:31:58Z

trieste/models/gpflux/sampler.py

+                tf.debugging.assert_shapes(
+                    [
+                        (u_sample, ["B", "M", "P"]),
+                    ]
+                )
+


this would be nicer with check shapes!

henrymoss · 2023-01-09T09:32:44Z

trieste/models/gpflux/sampler.py

+                    self._rff[counter].b.assign(self._rff[counter]._sample_bias(tf.shape(self._rff[
+                    counter].b), dtype=self._rff[counter]._dtype))
+                    self._rff[counter].W.assign(self._rff[counter]._sample_weights(tf.shape(
+                    self._rff[counter].W), dtype=self._rff[counter]._dtype))
+                else:
+                    self._rff[counter].b.assign(self._rff[counter]._bias_init(tf.shape(self._rff[
+                    counter].b), dtype=self._rff[counter]._dtype))
+                    self._rff[counter].W.assign(self._rff[counter]._weights_init(tf.shape(self._rff[


henrymoss · 2023-01-09T09:33:02Z

trieste/models/gpflux/sampler.py

+            # fourier_feature_eval = []
+            # for counter, ker in enumerate(self._kernel):
+            #    fourier_feature_eval.append(self._rff[counter].__call__(x))  # [N, L]
+
+            # fourier_feature_eval = tf.stack(fourier_feature_eval, axis = 0)  # [P, N, L]


this PR needs some tidying!

hstojic · 2023-04-13T19:19:16Z

@sebastianober has made a different set of changes in gpflux to accomodate separate independent kernels, and made corresponding changes here in GPflux package, so I'll close this PR - but we cannot handle sep ind inducing variables, so we might need to return to this at some point

TrajectorySampler amenable to any Kernel/IndVar

7a70f1d

SebastianPopescu requested review from hstojic and sebastianober August 16, 2022 15:21

clean up of initial codebase

0cb7be6

hstojic reviewed Sep 16, 2022

View reviewed changes

latest iteration on it

0827c14

hstojic reviewed Sep 16, 2022

View reviewed changes

tests/unit/models/gpflux/test_sampler.py Outdated Show resolved Hide resolved

hstojic reviewed Sep 16, 2022

View reviewed changes

trieste/models/gpflux/architectures.py Show resolved Hide resolved

hstojic reviewed Sep 16, 2022

View reviewed changes

SebastianPopescu commented Sep 20, 2022

View reviewed changes

SebastianPopescu and others added 14 commits October 4, 2022 14:10

tests are passing

88533c3

Merge branch 'develop' into sebastian.p/trajectory_Sampler

5fb4e10

solved merge conflicts

f029b51

formatting

9fab7e4

constriants

b9cab8d

fix tensorflow constraint

ad1fbc8

fix tensorflow constraint

f37e343

contrs and requir

1f6c7c7

Merge branch 'develop' of github.com:secondmind-labs/trieste into seb…

6acffef

…astian.p/trajectory_Sampler

resolved conflicss

15865e7

format

d729ecd

removed gpflux folder

a538561

some additional cleaning

3292acf

merged builder into one; started cleaning testing code

b293195

hstojic requested review from uri-granta and henrymoss November 25, 2022 15:53

henrymoss reviewed Jan 9, 2023

View reviewed changes

uri-granta removed their request for review March 16, 2023 09:37

hstojic closed this Apr 13, 2023

hstojic mentioned this pull request Apr 13, 2023

allowing SeparateIndependent inducing points in GPflux samplers #720

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

TrajectorySampler amenable to any Kernel/IndVar #606

TrajectorySampler amenable to any Kernel/IndVar #606

SebastianPopescu commented Aug 16, 2022

hstojic left a comment

hstojic Sep 16, 2022

hstojic Sep 16, 2022

SebastianPopescu Sep 16, 2022

SebastianPopescu commented Sep 16, 2022

hstojic Sep 16, 2022

SebastianPopescu Sep 20, 2022

hstojic Sep 20, 2022

hstojic Sep 16, 2022

SebastianPopescu Sep 20, 2022

hstojic Sep 16, 2022

SebastianPopescu Sep 20, 2022

hstojic Sep 16, 2022

SebastianPopescu Sep 20, 2022

hstojic Sep 20, 2022 •

edited

Loading

hstojic Sep 20, 2022

henrymoss Jan 9, 2023

henrymoss Jan 9, 2023

henrymoss Jan 9, 2023

henrymoss Jan 9, 2023

henrymoss Jan 9, 2023

henrymoss Jan 9, 2023

henrymoss Jan 9, 2023

henrymoss Jan 9, 2023

hstojic commented Apr 13, 2023

		] # [N, B, L + M, 1]


		# TODO -- probably have to re-write unflatten to accomodate for this case as well

TrajectorySampler amenable to any Kernel/IndVar #606

TrajectorySampler amenable to any Kernel/IndVar #606

Conversation

SebastianPopescu commented Aug 16, 2022

hstojic left a 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

SebastianPopescu commented Sep 16, 2022

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

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

hstojic Sep 20, 2022 • 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

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

hstojic commented Apr 13, 2023

hstojic Sep 20, 2022 •

edited

Loading