Kronecker 结构协方差#
PyMC 包含了具有Kronecker结构协方差的模型的实现。这种模式结构使得高斯过程模型能够处理更大规模的数据集。当
输入数据的维度是两个或更多（\(\mathbf{x} \in \mathbb{R}^{d}\,, d \ge 2\)）
该过程在每个维度或维度集上的影响是可分离的
核函数可以写成维度上的乘积，不包含交叉项：
与上述协方差函数对应的协方差矩阵可以用克罗内克积来表示
这些实现支持以下Kronecker积的性质以加速计算，\((\mathbf{K}_1 \otimes \mathbf{K}_2)^{-1} = \mathbf{K}_{1}^{-1} \otimes \mathbf{K}_{2}^{-1}\)，即和的逆等于逆的和。如果\(K_1\)是\(n \times n\)，而\(K_2\)是\(m \times m\)，那么\(\mathbf{K}_1 \otimes \mathbf{K}_2\)是\(mn \times mn\)。对于\(m\)和\(n\)即使大小适中，这种逆运算也变得难以高效完成。对两个矩阵进行逆运算，一个是\(n \times n\)，另一个是\(m \times m\)，则要容易得多。
这种结构在时空数据中很常见。鉴于协方差矩阵中存在Kronecker结构，此实现是精确的——不是对完整高斯过程的近似。PyMC包含两个实现，它们遵循与gp.Marginal和gp.Latent相同的模式。对于数据似然为高斯分布的Kronecker结构协方差，请使用gp.MarginalKron。对于数据似然为非高斯分布的Kronecker结构协方差，请使用gp.LatentKron。
我们的实现遵循Saatchi的论文。gp.MarginalKron遵循“算法16”使用特征分解，而gp.LatentKron遵循“算法14”，并使用Cholesky分解。
使用 MarginalKron 解决二维空间问题#
以下是 gp.MarginalKron 使用的一个典型示例。与 gp.Marginal 类似，该模型假设潜在的 GP 是未观察到的，但观察到的是 GP 与正态分布噪声的总和。
对于模拟数据集，我们从具有二维输入的高斯过程中抽取一个样本，其协方差具有Kronecker结构。然后我们使用gp.MarginalKron来恢复用于模拟数据的未知高斯过程超参数\(\theta\)。
示例#
我们将模拟一个二维数据集，并将其显示为散点图，点的颜色根据大小变化。这两个维度分别标记为x1和x2。这可能是一个空间数据集，例如。由于点位于二维网格上，协方差将具有Kronecker结构。
import arviz as az
import matplotlib as mpl
import numpy as np
import pymc as pm
az.style.use("arviz-whitegrid")
plt = mpl.pyplot
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
seed = sum(map(ord, "gpkron"))
rng = np.random.default_rng(seed)
# One dimensional column vectors of inputs
n1, n2 = (50, 30)
x1 = np.linspace(0, 5, n1)
x2 = np.linspace(0, 3, n2)
# make cartesian grid out of each dimension x1 and x2
X = pm.math.cartesian(x1[:, None], x2[:, None])
l1_true = 0.8
l2_true = 1.0
eta_true = 1.0
# Although we could, we don't exploit kronecker structure to draw the sample
cov = (
eta_true**2
* pm.gp.cov.Matern52(2, l1_true, active_dims=[0])
* pm.gp.cov.Cosine(2, ls=l2_true, active_dims=[1])
)
K = cov(X).eval()
f_true = rng.multivariate_normal(np.zeros(X.shape[0]), K, 1).flatten()
sigma_true = 0.5
y = f_true + sigma_true * rng.standard_normal(X.shape[0])
沿x2维度的长度尺度比沿x1方向的长度尺度更长（l1_true < l2_true）。
fig = plt.figure(figsize=(12, 6))
cmap = "terrain"
norm = mpl.colors.Normalize(vmin=-3, vmax=3)
plt.scatter(X[:, 0], X[:, 1], s=35, c=y, marker="o", norm=norm, cmap=cmap)
plt.colorbar()
plt.xlabel("x1"), plt.ylabel("x2")
plt.title("Simulated dataset");
这个数据集中有1500个数据点。如果不使用Kronecker分解，找到MAP估计的速度会慢得多。
由于两个协方差是乘积，我们只需要一个尺度参数 eta 来建模乘积协方差函数。
# this implementation takes a list of inputs for each dimension as input
Xs = [x1[:, None], x2[:, None]]
with pm.Model() as model:
# Set priors on the hyperparameters of the covariance
ls1 = pm.Gamma("ls1", alpha=2, beta=2)
ls2 = pm.Gamma("ls2", alpha=2, beta=2)
eta = pm.HalfNormal("eta", sigma=2)
# Specify the covariance functions for each Xi
# Since the covariance is a product, only scale one of them by eta.
# Scaling both overparameterizes the covariance function.
cov_x1 = pm.gp.cov.Matern52(1, ls=ls1) # cov_x1 must accept X1 without error
cov_x2 = eta**2 * pm.gp.cov.Cosine(1, ls=ls2) # cov_x2 must accept X2 without error
# Specify the GP. The default mean function is `Zero`.
gp = pm.gp.MarginalKron(cov_funcs=[cov_x1, cov_x2])
# Set the prior on the variance for the Gaussian noise
sigma = pm.HalfNormal("sigma", sigma=2)
# Place a GP prior over the function f.
y_ = gp.marginal_likelihood("y", Xs=Xs, y=y, sigma=sigma)
with model:
mp = pm.find_MAP(method="BFGS")
mp
{'ls1_log__': array(0.13716063),
'ls2_log__': array(-0.0004206),
'eta_log__': array(0.34822276),
'sigma_log__': array(-0.65839497),
'ls1': array(1.14701237),
'ls2': array(0.99957949),
'eta': array(1.41654776),
'sigma': array(0.51768156)}
接下来我们使用地图点 mp 在外部区域的原始网格之外进行外推。 我们也可以进行插值。 在新输入中，预测所需的区域没有网格限制。 需要注意的是，在当前的实现中，这些点具有网格结构并不会带来任何效率提升。 包含外推的图表如下所示。 原始数据如前所示用圆圈标记，但外推的后验均值用方块标记。
x1new = np.linspace(5.1, 7.1, 20)
x2new = np.linspace(-0.5, 3.5, 40)
Xnew = pm.math.cartesian(x1new[:, None], x2new[:, None])
with model:
mu, var = gp.predict(Xnew, point=mp, diag=True)
fig = plt.figure(figsize=(12, 6))
cmap = "terrain"
norm = mpl.colors.Normalize(vmin=-3, vmax=3)
m = plt.scatter(X[:, 0], X[:, 1], s=30, c=y, marker="o", norm=norm, cmap=cmap)
plt.colorbar(m)
plt.scatter(Xnew[:, 0], Xnew[:, 1], s=30, c=mu, marker="s", norm=norm, cmap=cmap)
plt.ylabel("x2"), plt.xlabel("x1")
plt.title("observed data 'y' (circles) with predicted mean (squares)");
LatentKron#
与gp.Latent实现类似，gp.LatentKron实现指定了一个Kronecker结构的GP，无论上下文如何。它可以与任何似然函数一起使用，或者可以用于建模方差或其他一些未观察到的过程。语法与gp.Latent完全相同。
模型#
为了与边际似然进行比较，我们使用了与之前相同的示例，其中噪声是正态分布的，但高斯过程本身并未被边缘化。相反，它是直接使用NUTS进行采样的。需要注意的是，gp.LatentKron不需要像gp.MarginalKron那样的高斯似然；相反，任何似然都是可接受的。
在这里，我们需要为我们的先验提供更多的信息，至少是为高斯过程超参数的先验。这是使用高斯过程时的一般规则：尽可能使用信息丰富的先验，因为采样长度尺度和幅度是一项具有挑战性的工作，所以你希望使采样器的工作尽可能简单。
在这里，幸运的是，我们有很多关于振幅和长度尺度的信息——是我们创造了它们 ;) 所以我们可以固定它们，但我们将展示如何在您自己的模型中编码这些先验知识，例如，使用截断正态分布：
with pm.Model() as model:
# Set priors on the hyperparameters of the covariance
ls1 = pm.TruncatedNormal("ls1", lower=0.5, upper=1.5, mu=1, sigma=0.5)
ls2 = pm.TruncatedNormal("ls2", lower=0.5, upper=1.5, mu=1, sigma=0.5)
eta = pm.HalfNormal("eta", sigma=0.5)
# Specify the covariance functions for each Xi
cov_x1 = pm.gp.cov.Matern52(1, ls=ls1)
cov_x2 = eta**2 * pm.gp.cov.Cosine(1, ls=ls2)
# Specify the GP. The default mean function is `Zero`
gp = pm.gp.LatentKron(cov_funcs=[cov_x1, cov_x2])
# Place a GP prior over the function f
f = gp.prior("f", Xs=Xs)
# Set the prior on the variance for the Gaussian noise
sigma = pm.HalfNormal("sigma", sigma=0.5)
y_ = pm.Normal("y_", mu=f, sigma=sigma, observed=y)
pm.model_to_graphviz(model)
with model:
idata = pm.sample(nuts_sampler="numpyro", target_accept=0.9, tune=1500, draws=1500)
idata.sample_stats.diverging.sum().data
array(0)
后验收敛#
未知长度尺度参数、协方差缩放 eta 和白噪声 sigma 的后验分布如下所示。垂直线是用于生成原始数据集的真实值：
var_names = ["ls1", "ls2", "eta", "sigma"]
az.plot_posterior(
idata,
var_names=var_names,
ref_val=[l1_true, l2_true, eta_true, sigma_true],
grid=(2, 2),
figsize=(12, 6),
);
我们可以看到在这些情况下采样是多么具有挑战性。在这里，一切都很顺利，因为我们仔细选择了先验——特别是在这个模拟案例中，参数没有实际的解释。
轨迹图看起来是什么样的？
az.plot_trace(idata, var_names=var_names);
一切顺利，所以让我们继续进行样本外预测！
样本外预测#
x1new = np.linspace(5.1, 7.1, 20)[:, None]
x2new = np.linspace(-0.5, 3.5, 40)[:, None]
Xnew = pm.math.cartesian(x1new, x2new)
x1new.shape, x2new.shape, Xnew.shape
((20, 1), (40, 1), (800, 2))
with model:
fnew = gp.conditional("fnew", Xnew, jitter=1e-6)
with model:
ppc = pm.sample_posterior_predictive(idata, var_names=["fnew"], compile_kwargs={"mode": "JAX"})
/Users/alex_andorra/mambaforge/envs/pymc-examples/lib/python3.12/site-packages/pytensor/link/jax/linker.py:28: UserWarning: The RandomType SharedVariables [RandomGeneratorSharedVariable(<Generator(PCG64) at 0x17B535D20>)] will not be used in the compiled JAX graph. Instead a copy will be used.
warnings.warn(
Sampling: [fnew]
2024-05-27 15:30:56.157723: E external/xla/xla/service/slow_operation_alarm.cc:65] Constant folding an instruction is taking > 1s:
%reduce.3 = f64[1500,800]{1,0} reduce(f64[1,1500,800]{2,1,0} %broadcast.59, f64[] %constant.76), dimensions={0}, to_apply=%region_2.198, metadata={op_name="jit(jax_funcified_fgraph)/jit(main)/reduce_sum[axes=(0,)]" source_file="/var/folders/m_/brf3tky55f3gf6dy8w7c6s0w0000gn/T/tmpnlsq8yzu" source_line=77}
This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.
If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
2024-05-27 15:30:56.305845: E external/xla/xla/service/slow_operation_alarm.cc:133] The operation took 1.153176s
Constant folding an instruction is taking > 1s:
%reduce.3 = f64[1500,800]{1,0} reduce(f64[1,1500,800]{2,1,0} %broadcast.59, f64[] %constant.76), dimensions={0}, to_apply=%region_2.198, metadata={op_name="jit(jax_funcified_fgraph)/jit(main)/reduce_sum[axes=(0,)]" source_file="/var/folders/m_/brf3tky55f3gf6dy8w7c6s0w0000gn/T/tmpnlsq8yzu" source_line=77}
This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.
If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
下面我们展示了原始数据集为彩色圆圈，条件样本的均值为彩色方块。结果与gp.MarginalKron实现的结果非常接近。
fig = plt.figure(figsize=(14, 7))
m = plt.scatter(X[:, 0], X[:, 1], s=30, c=y, marker="o", norm=norm, cmap=cmap)
plt.colorbar(m)
plt.scatter(
Xnew[:, 0],
Xnew[:, 1],
s=30,
c=np.mean(ppc.posterior_predictive["fnew"].sel(chain=0), axis=0),
marker="s",
norm=norm,
cmap=cmap,
)
plt.ylabel("x2"), plt.xlabel("x1")
plt.title("observed data 'y' (circles) with mean of conditional, or predicted, samples (squares)");
接下来，我们绘制了用圆形标记表示的原始数据集，以及用方形标记表示的条件分布中的四个样本。正如我们所见，预测分布中的变化水平导致了fnew值中截然不同的模式。然而，这些样本显示了正确的相关结构——我们在y轴上看到了明显的正弦模式，在x轴上看到了近端相关结构。观察到的数据中的模式与条件分布无缝融合。
fig, axs = plt.subplots(2, 2, figsize=(24, 16))
axs = axs.ravel()
for i, ax in enumerate(axs):
ax.axis("off")
ax.scatter(X[:, 0], X[:, 1], s=20, c=y, marker="o", norm=norm, cmap=cmap)
ax.scatter(
Xnew[:, 0],
Xnew[:, 1],
s=20,
c=ppc.posterior_predictive["fnew"].sel(chain=0)[i],
marker="s",
norm=norm,
cmap=cmap,
)
ax.set_title(f"Sample {i+1}", fontsize=24)
水印#
%load_ext watermark
%watermark -n -u -v -iv -w -p pytensor,xarray
Last updated: Mon May 27 2024
Python implementation: CPython
Python version : 3.12.2
IPython version : 8.22.2
pytensor: 2.20.0
xarray : 2024.3.0
numpy : 1.26.4
pymc : 5.15.0+14.gfd11cf012
arviz : 0.17.1
matplotlib: 3.8.3
Watermark: 2.4.3
许可证声明#
本示例库中的所有笔记本均在MIT许可证下提供，该许可证允许修改和重新分发，前提是保留版权和许可证声明。
引用 PyMC 示例#
要引用此笔记本，请使用Zenodo为pymc-examples仓库提供的DOI。
重要
许多笔记本是从其他来源改编的：博客、书籍……在这种情况下，您应该引用原始来源。
同时记得引用代码中使用的相关库。
这是一个BibTeX的引用模板：
@incollection{citekey,
author = "<notebook authors, see above>",
title = "<notebook title>",
editor = "PyMC Team",
booktitle = "PyMC examples",
doi = "10.5281/zenodo.5654871"
}
渲染后可能看起来像：