Add means to EEG experiment

wesselb · Jan 26, 2021 · cd4469d · cd4469d
1 parent 33615d2
commit cd4469d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/experiments/eeg.py b/experiments/eeg.py
@@ -5,7 +5,7 @@
 import wbml.out as out
 from stheno.jax import GP, Matern32, Delta
 from varz.jax import Vars
-from varz.spec import parametrised, Positive
+from varz.spec import parametrised, Positive, Unbounded
 from wbml.data.eeg import load
 from wbml.experiment import WorkingDirectory
 from wbml.plot import tweak
@@ -38,14 +38,15 @@
 def model(
     vs,
     z,
+    means: Unbounded(shape=(m,)),
     variances: Positive = 1 * B.ones(m),
     scales: Positive = np.array([0.5, 0.1, 0.05]),
     noises: Positive = 1e-2 * B.ones(m),
 ):
     logpdf = 0
     for i in range(m):
         kernel = variances[i] * Matern32().stretch(scales[i]) + noises[i] * Delta()
-        logpdf += GP(kernel)(x).logpdf(z[:, i])
+        logpdf += GP(means[i], kernel)(x).logpdf(z[:, i])
     return logpdf