def _extract_samples(quantities): x = quantities.x samples = quantities.all_samples return ( B.concat(-x[::-1][:-1], x), B.concat(samples[::-1, :][:-1, :], samples, axis=0), )
def project(self, x, y): """Project data. Args: x (matrix): Locations of data. y (matrix): Observations of data. Returns: tuple: Tuple containing the locations of the projection, the projection, weights associated with the projection, and a regularisation term. """ n = B.shape(x)[0] available = ~B.isnan(B.to_numpy(y)) # Optimise the case where all data is available. if B.all(available): return self._project_pattern(x, y, (True,) * self.p) # Extract patterns. patterns = list(set(map(tuple, list(available)))) if len(patterns) > 30: warnings.warn( f"Detected {len(patterns)} patterns, which is more " f"than 30 and can be slow.", category=UserWarning, ) # Per pattern, find data points that belong to it. patterns_inds = [[] for _ in range(len(patterns))] for i in range(n): patterns_inds[patterns.index(tuple(available[i]))].append(i) # Per pattern, perform the projection. proj_xs = [] proj_ys = [] proj_ws = [] total_reg = 0 for pattern, pattern_inds in zip(patterns, patterns_inds): proj_x, proj_y, proj_w, reg = self._project_pattern( B.take(x, pattern_inds), B.take(y, pattern_inds), pattern ) proj_xs.append(proj_x) proj_ys.append(proj_y) proj_ws.append(proj_w) total_reg = total_reg + reg return ( B.concat(*proj_xs, axis=0), B.concat(*proj_ys, axis=0), B.concat(*proj_ws, axis=0), total_reg, )
def __call__(self, prev, x): prev_h, _ = prev # Gate logic: z = self.f_z(B.concat(prev_h, x, axis=1)) r = self.f_r(B.concat(prev_h, x, axis=1)) h = (1 - z) * prev_h + z * self.f_h(B.concat(r * prev_h, x, axis=1)) y = h return h, y
def test_concat(dense1, dense2, diag1, diag2): with AssertDenseWarning("concatenating <dense>, <dense>, <diagonal>..."): res = B.concat(dense1, dense2, diag1, diag2, axis=1) dense_res = B.concat(B.dense(dense1), B.dense(dense2), B.dense(diag1), B.dense(diag2), axis=1) approx(res, dense_res) assert isinstance(res, Dense)
def test_combine(): x1 = B.linspace(0, 2, 10) x2 = B.linspace(2, 4, 10) m = Measure() p1 = GP(EQ(), measure=m) p2 = GP(Matern12(), measure=m) y1 = p1(x1).sample() y2 = p2(x2).sample() # Check the one-argument case. assert_equal_normals(combine(p1(x1, 1)), p1(x1, 1)) fdd_combined, y_combined = combine((p1(x1, 1), B.squeeze(y1))) assert_equal_normals(fdd_combined, p1(x1, 1)) approx(y_combined, y1) # Check the two-argument case. fdd_combined = combine(p1(x1, 1), p2(x2, 2)) assert_equal_normals( fdd_combined, Normal(B.block_diag(p1(x1, 1).var, p2(x2, 2).var)), ) fdd_combined, y_combined = combine((p1(x1, 1), B.squeeze(y1)), (p2(x2, 2), y2)) assert_equal_normals( fdd_combined, Normal(B.block_diag(p1(x1, 1).var, p2(x2, 2).var)), ) approx(y_combined, B.concat(y1, y2, axis=0))
def test_ess(): # Construct a prior and a likelihood. prior = Normal(np.array([[0.6, 0.3], [0.3, 0.6]])) lik = Normal( np.array([[0.2], [0.3]]), np.array([[1, 0.2], [0.2, 1]]), ) # Perform sampling. sampler = ESS(lik.logpdf, prior.sample) num_samples = 30_000 samples = B.concat(*sampler.sample(num=num_samples), axis=1) samples_mean = B.mean(samples, axis=1)[:, None] samples_cov = ( B.matmul(samples - samples_mean, samples - samples_mean, tr_b=True) / num_samples) # Compute posterior statistics. prec_prior = B.inv(prior.var) prec_lik = B.inv(lik.var) cov = B.inv(prec_prior + prec_lik) mean = cov @ (prec_prior @ prior.mean + prec_lik @ lik.mean) approx(samples_cov, cov, atol=5e-2) approx(samples_mean, mean, atol=5e-2)
def test_momean(x): prior = Measure() p1 = GP(lambda x: 2 * x, 1 * EQ(), measure=prior) p2 = GP(1, 2 * EQ().stretch(2), measure=prior) m = MultiOutputMean(prior, p1, p2) ms = prior.means # Check representation. assert str(m) == "MultiOutputMean(<lambda>, 1)" # Check computation. approx(m(x), B.concat(ms[p1](x), ms[p2](x), axis=0)) approx(m(p1(x)), ms[p1](x)) approx(m(p2(x)), ms[p2](x)) approx(m(MultiInput(p2(x), p1(x))), B.concat(ms[p2](x), ms[p1](x), axis=0))
def test_logpdf_missing_data(): # Setup model. m = 3 noise = 1e-2 latent_noises = 2e-2 * B.ones(m) kernels = [0.5 * EQ().stretch(0.75) for _ in range(m)] x = B.linspace(0, 10, 20) # Concatenate two orthogonal matrices, to make the missing data # approximation exact. u1 = B.svd(B.randn(m, m))[0] u2 = B.svd(B.randn(m, m))[0] u = Dense(B.concat(u1, u2, axis=0) / B.sqrt(2)) s_sqrt = Diagonal(B.rand(m)) # Construct a reference model. oilmm_pp = ILMMPP(kernels, u @ s_sqrt, noise, latent_noises) # Sample to generate test data. y = oilmm_pp.sample(x, latent=False) # Throw away data, but retain orthogonality. y[5:10, 3:] = np.nan y[10:, :3] = np.nan # Construct OILMM to test. oilmm = OILMM(kernels, u, s_sqrt, noise, latent_noises) # Check that evidence is still exact. approx(oilmm_pp.logpdf(x, y), oilmm.logpdf(x, y), atol=1e-7)
def _attempt_diagonal(rows): # Check whether the result is diagonal. # Check that the blocks form a square. if not all([len(row) == len(rows) for row in rows]): return None # Collect the diagonal blocks. diagonal_blocks = [] for r in range(len(rows)): for c in range(len(rows[0])): block_shape = B.shape(rows[r][c]) if r == c: # Keep track of all the diagonal blocks. diagonal_blocks.append(rows[r][c]) # All blocks on the diagonal must be diagonal or zero. if not isinstance(rows[r][c], (Diagonal, Zero)): return None # All blocks on the diagonal must be square. if not block_shape[0] == block_shape[1]: return None else: # All blocks not on the diagonal must be zero. if not isinstance(rows[r][c], Zero): return None return Diagonal(B.concat(*[B.diag(x) for x in diagonal_blocks]))
def __call__(self, prev, x): prev_h, prev_y = prev # Gate logic: h = self.f_h(B.concat(prev_h, x, axis=1)) y = h return h, y
def concat(*elements: AbstractMatrix, axis=0): if structured(*elements): elements_str = ", ".join(map(str, elements[:3])) if len(elements) > 3: elements_str += "..." warn_upmodule( f"Concatenating {elements_str}: converting to dense.", category=ToDenseWarning, ) return Dense(B.concat(*(B.dense(el) for el in elements), axis=axis))
def pack(*objs: B.Numeric): """Pack objects. Args: *objs (tensor): Objects to pack. Returns: tensor: Vector representation of the objects. """ return B.concat(*[B.flatten(obj) for obj in objs], axis=0)
def compute_I_uz(model, t): """Compute the :math:`I_{uz,t_i}` matrix for :math:`t_i` in `t`. Args: model (:class:`.gprv.GPRV`): Model. t (vector): Time points :math:`t_i` of data. Returns: tensor: Value of :math:`I_{uz,t_i}`, with shape `(len(t), len(model.t_u), len(model.ms))`. """ # Compute sorting permutation. perm = np.argsort(model.ms) inverse_perm = invert_perm(perm) # Sort to allow for simple concatenation. ms = model.ms[perm] # Part of the factor is absorbed in the integrals. factor = ( model.alpha_t * model.gamma_t * B.exp(-model.gamma * model.t_u[None, :, None]) ) # Compute I(l, u, 0). k = ms[ms == 0] I_uz0 = ( _integral_lu0( model, t[:, None, None] - model.t_u[None, :, None], t[:, None, None] ) + k[None, None, :] ) # This is all zeros, but performs broadcasting. # Compute I(l, u, k) for 0 < k < M + 1. k = ms[(0 < ms) * (ms <= model.m_max)] I_uz_leM = _integral_luk_leq_M( model, t[:, None, None] - model.t_u[None, :, None], t[:, None, None], k[None, None, :], ) # Compute I(l, u, k) for k > M. k = ms[ms > model.m_max] I_uz_gtM = _integral_luk_g_M( model, t[:, None, None] - model.t_u[None, :, None], t[:, None, None], k[None, None, :], ) # Construct result. result = B.concat(I_uz0, I_uz_leM, I_uz_gtM, axis=2) # Undo sorting and return. return factor * B.take(result, inverse_perm, axis=2)
def test_mom(): x = B.linspace(0, 1, 10) prior = Measure() p1 = GP(lambda x: 2 * x, 1 * EQ(), measure=prior) p2 = GP(1, 2 * EQ().stretch(2), measure=prior) m = MultiOutputMean(prior, p1, p2) ms = prior.means # Check dimensionality. assert dimensionality(m) == 2 # Check representation. assert str(m) == "MultiOutputMean(<lambda>, 1)" # Check computation. approx(m(x), B.concat(ms[p1](x), ms[p2](x), axis=0)) approx(m(p1(x)), ms[p1](x)) approx(m(p2(x)), ms[p2](x)) approx(m((p2(x), p1(x))), B.concat(ms[p2](x), ms[p1](x), axis=0))
def sample(self, x, latent=False): """Sample from the model. Args: x (matrix): Locations to sample at. latent (bool, optional): Sample noiseless processes. Defaults to `False`. Returns: matrix: Sample. """ fdds = [f(x, *(() if latent else (n,))) for f, n in zip(self.fs, self.noises)] return B.concat(*[fdd.sample() for fdd in fdds], axis=1)
def inverse_transform(x): chol = B.cholesky(B.reg(x)) return B.concat(B.log(B.diag(chol)), B.tril_to_vec(chol, offset=-1))
def test_mokernel(x1, x2, x3): m = Measure() p1 = GP(1 * EQ(), measure=m) p2 = GP(2 * EQ().stretch(2), measure=m) k = MultiOutputKernel(m, p1, p2) ks = m.kernels # Check representation. assert str(k) == "MultiOutputKernel(EQ(), 2 * (EQ() > 2))" # Input versus input: approx( k(x1, x2), B.concat2d( [ks[p1, p1](x1, x2), ks[p1, p2](x1, x2)], [ks[p2, p1](x1, x2), ks[p2, p2](x1, x2)], ), ) approx( k.elwise(x1, x3), B.concat(ks[p1, p1].elwise(x1, x3), ks[p2, p2].elwise(x1, x3), axis=0), ) # Input versus `FDD`: approx(k(p1(x1), x2), B.concat(ks[p1, p1](x1, x2), ks[p1, p2](x1, x2), axis=1)) approx(k(p2(x1), x2), B.concat(ks[p2, p1](x1, x2), ks[p2, p2](x1, x2), axis=1)) approx(k(x1, p1(x2)), B.concat(ks[p1, p1](x1, x2), ks[p2, p1](x1, x2), axis=0)) approx(k(x1, p2(x2)), B.concat(ks[p1, p2](x1, x2), ks[p2, p2](x1, x2), axis=0)) with pytest.raises(ValueError): k.elwise(x1, p2(x3)) with pytest.raises(ValueError): k.elwise(p1(x1), x3) # `FDD` versus `FDD`: approx(k(p1(x1), p1(x2)), ks[p1](x1, x2)) approx(k(p1(x1), p2(x2)), ks[p1, p2](x1, x2)) approx(k.elwise(p1(x1), p1(x3)), ks[p1].elwise(x1, x3)) approx(k.elwise(p1(x1), p2(x3)), ks[p1, p2].elwise(x1, x3)) # `MultiInput` versus input: approx( k(MultiInput(p2(x1), p1(x2)), x1), B.concat2d( [ks[p2, p1](x1, x1), ks[p2, p2](x1, x1)], [ks[p1, p1](x2, x1), ks[p1, p2](x2, x1)], ), ) approx( k(x1, MultiInput(p2(x1), p1(x2))), B.concat2d( [ks[p1, p2](x1, x1), ks[p1, p1](x1, x2)], [ks[p2, p2](x1, x1), ks[p2, p1](x1, x2)], ), ) with pytest.raises(ValueError): k.elwise(MultiInput(p2(x1), p1(x3)), p2(x1)) with pytest.raises(ValueError): k.elwise(p2(x1), MultiInput(p2(x1), p1(x3))) # `MultiInput` versus `FDD`: approx( k(MultiInput(p2(x1), p1(x2)), p2(x1)), B.concat(ks[p2, p2](x1, x1), ks[p1, p2](x2, x1), axis=0), ) approx( k(p2(x1), MultiInput(p2(x1), p1(x2))), B.concat(ks[p2, p2](x1, x1), ks[p2, p1](x1, x2), axis=1), ) with pytest.raises(ValueError): k.elwise(MultiInput(p2(x1), p1(x3)), p2(x1)) with pytest.raises(ValueError): k.elwise(p2(x1), MultiInput(p2(x1), p1(x3))) # `MultiInput` versus `MultiInput`: approx( k(MultiInput(p2(x1), p1(x2)), MultiInput(p2(x1))), B.concat(ks[p2, p2](x1, x1), ks[p1, p2](x2, x1), axis=0), ) with pytest.raises(ValueError): k.elwise(MultiInput(p2(x1), p1(x3)), MultiInput(p2(x1))) approx( k.elwise(MultiInput(p2(x1), p1(x3)), MultiInput(p2(x1), p1(x3))), B.concat(ks[p2, p2].elwise(x1, x1), ks[p1, p1].elwise(x3, x3), axis=0), )
def compute_I_hz(model, t): """Compute the :math:`I_{hz,t_i}` matrix for :math:`t_i` in `t`. Args: model (:class:`.gprv.GPRV`): Model. t (vector): Time points of data. Returns: tensor: Value of :math:`I_{hz,t_i}`, with shape `(len(t), len(model.ms), len(model.ms))`. """ # Compute sorting permutation. perm = np.argsort(model.ms) inverse_perm = invert_perm(perm) # Sort to allow for simple concatenation. m_max = model.m_max ms = model.ms[perm] # Construct I_hz for m,n <= M. ns = ms[ms <= m_max] I_0_cos_1 = _I_hx_0_cos( model, -ns[None, :, None] + ns[None, None, :], t[:, None, None] ) I_0_cos_2 = _I_hx_0_cos( model, ns[None, :, None] + ns[None, None, :], t[:, None, None] ) I_hz_mnleM = 0.5 * (I_0_cos_1 + I_0_cos_2) # Construct I_hz for m,n > M. ns = ms[ms > m_max] - m_max I_0_cos_1 = _I_hx_0_cos( model, -ns[None, :, None] + ns[None, None, :], t[:, None, None] ) I_0_cos_2 = _I_hx_0_cos( model, ns[None, :, None] + ns[None, None, :], t[:, None, None] ) I_hz_mngtM = 0.5 * (I_0_cos_1 - I_0_cos_2) # Construct I_hz for 0 < m <= M and n > M. ns = ms[(0 < ms) * (ms <= m_max)] ns2 = ms[ms > m_max] # Do not subtract M! I_0_sin_1 = _I_hx_0_sin( model, ns[None, :, None] + ns2[None, None, :], t[:, None, None] ) I_0_sin_2 = _I_hx_0_sin( model, -ns[None, :, None] + ns2[None, None, :], t[:, None, None] ) I_hz_mleM_ngtM = 0.5 * (I_0_sin_1 + I_0_sin_2) # Construct I_hz for m = 0 and n > M. ns = ms[ms == 0] ns2 = ms[ms > m_max] # Do not subtract M! I_hz_0_gtM = _I_hx_0_sin( model, ns[None, :, None] + ns2[None, None, :], t[:, None, None] ) # Concatenate to form I_hz for m <= M and n > M. I_hz_mleM_ngtM = B.concat(I_hz_0_gtM, I_hz_mleM_ngtM, axis=1) # Compute the other half by transposing. I_hz_mgtM_nleM = B.transpose(I_hz_mleM_ngtM, perm=(0, 2, 1)) # Construct result. result = B.concat( B.concat(I_hz_mnleM, I_hz_mleM_ngtM, axis=2), B.concat(I_hz_mgtM_nleM, I_hz_mngtM, axis=2), axis=1, ) # Undo sorting. result = B.take(result, inverse_perm, axis=1) result = B.take(result, inverse_perm, axis=2) return result
models = model() # Perform sampling. if args.train: ks, us, fs = sample(model, t, noise_f) wd.save((ks, us, fs), "samples.pickle") else: ks, us, fs = wd.load("samples.pickle") # Plot. plt.figure(figsize=(15, 4)) for i, (k, u, f) in enumerate(zip(ks, us, fs)): plt.subplot(3, 5, 1 + i) plt.plot( B.concat(-t[::-1][:-1], t), B.concat(u[:-1] * 0, u), lw=1, ) if hasattr(model, "t_u"): plt.scatter(model.t_u, model.t_u * 0, s=5, marker="o", c="black") # plt.xlabel("Time (s)") if i == 0: plt.ylabel("$h$") plt.xlim(-6, 6) plt.ylim(-0.5, 1) tweak(legend=False) plt.subplot(3, 5, 6 + i) plt.plot( B.concat(-t[::-1][:-1], t),
starting from the origin. k (vector): Kernel. n_zero (int, optional): Zero padding. Defaults to `2_000`. db (bool, optional): Convert to decibel. Defaults to `False`. Returns: vector: PSD, correctly scaled. """ # Convert to NumPy for compatibility with frameworks. t, k = B.to_numpy(t, k) if t[0] != 0: raise ValueError("Time points must start at zero.") # Perform zero padding. k = B.concat(k, B.zeros(n_zero)) # Symmetrise and Fourier transform. k_symmetric = B.concat(k, k[1:-1][::-1]) psd = np.fft.fft(k_symmetric) freqs = np.fft.fftfreq(len(psd)) / (t[1] - t[0]) # Should be real and positive, but the numerics may not be in our favour. psd = np.abs(np.real(psd)) # Now scale appropriately: the total power should equal `k[0]`. total_power = np.trapz(y=psd, x=freqs) psd /= total_power / k[0] # Convert to dB. if db:
def diag(a: Diagonal, b: Diagonal): return Diagonal(B.concat(a.diag, b.diag))
def test_mok(): x1 = B.linspace(0, 1, 10) x2 = B.linspace(1, 2, 5) x3 = B.linspace(1, 2, 10) m = Measure() p1 = GP(EQ(), measure=m) p2 = GP(2 * EQ().stretch(2), measure=m) k = MultiOutputKernel(m, p1, p2) ks = m.kernels # Check dimensionality. assert dimensionality(k) == 2 # Check representation. assert str(k) == "MultiOutputKernel(EQ(), 2 * (EQ() > 2))" # Input versus input: approx( k(x1, x2), B.concat2d( [ks[p1, p1](x1, x2), ks[p1, p2](x1, x2)], [ks[p2, p1](x1, x2), ks[p2, p2](x1, x2)], ), ) approx( k.elwise(x1, x3), B.concat(ks[p1, p1].elwise(x1, x3), ks[p2, p2].elwise(x1, x3), axis=0), ) # Input versus `FDD`: approx(k(p1(x1), x2), B.concat(ks[p1, p1](x1, x2), ks[p1, p2](x1, x2), axis=1)) approx(k(p2(x1), x2), B.concat(ks[p2, p1](x1, x2), ks[p2, p2](x1, x2), axis=1)) approx(k(x1, p1(x2)), B.concat(ks[p1, p1](x1, x2), ks[p2, p1](x1, x2), axis=0)) approx(k(x1, p2(x2)), B.concat(ks[p1, p2](x1, x2), ks[p2, p2](x1, x2), axis=0)) with pytest.raises(ValueError): k.elwise(x1, p2(x3)) with pytest.raises(ValueError): k.elwise(p1(x1), x3) # `FDD` versus `FDD`: approx(k(p1(x1), p1(x2)), ks[p1](x1, x2)) approx(k(p1(x1), p2(x2)), ks[p1, p2](x1, x2)) approx(k.elwise(p1(x1), p1(x3)), ks[p1].elwise(x1, x3)) approx(k.elwise(p1(x1), p2(x3)), ks[p1, p2].elwise(x1, x3)) # Multiple inputs versus input: approx( k((p2(x1), p1(x2)), x1), B.concat2d( [ks[p2, p1](x1, x1), ks[p2, p2](x1, x1)], [ks[p1, p1](x2, x1), ks[p1, p2](x2, x1)], ), ) approx( k(x1, (p2(x1), p1(x2))), B.concat2d( [ks[p1, p2](x1, x1), ks[p1, p1](x1, x2)], [ks[p2, p2](x1, x1), ks[p2, p1](x1, x2)], ), ) with pytest.raises(ValueError): k.elwise((p2(x1), p1(x3)), p2(x1)) with pytest.raises(ValueError): k.elwise(p2(x1), (p2(x1), p1(x3))) # Multiple inputs versus `FDD`: approx( k((p2(x1), p1(x2)), p2(x1)), B.concat(ks[p2, p2](x1, x1), ks[p1, p2](x2, x1), axis=0), ) approx( k(p2(x1), (p2(x1), p1(x2))), B.concat(ks[p2, p2](x1, x1), ks[p2, p1](x1, x2), axis=1), ) with pytest.raises(ValueError): k.elwise((p2(x1), p1(x3)), p2(x1)) with pytest.raises(ValueError): k.elwise(p2(x1), (p2(x1), p1(x3))) # Multiple inputs versus multiple inputs: approx( k((p2(x1), p1(x2)), (p2(x1))), B.concat(ks[p2, p2](x1, x1), ks[p1, p2](x2, x1), axis=0), ) with pytest.raises(ValueError): k.elwise((p2(x1), p1(x3)), (p2(x1),)) approx( k.elwise((p2(x1), p1(x3)), (p2(x1), p1(x3))), B.concat(ks[p2, p2].elwise(x1, x1), ks[p1, p1].elwise(x3, x3), axis=0), )
def elwise(k, x: tuple, y: tuple): if len(x) != len(y): raise ValueError('"elwise" must be called with similarly sized tuples.') return B.concat(*[elwise(k, xi, yi) for xi, yi in zip(x, y)], axis=-2)
def _pad_zero_row(a): zeros = B.zeros(B.dtype(a), 1, B.shape(a)[1]) return B.concat(a, zeros, axis=0)
def _pad_zero_col(a): zeros = B.zeros(B.dtype(a), B.shape(a)[0], 1) return B.concat(a, zeros, axis=1)