注意
转到末尾 下载完整的示例代码。
示例：生态数据CJS捕获-再捕获模型
本示例移植自 [8]。
我们将展示如何在生态学中用于分析动物捕获-再捕获数据的 Cormack-Jolly-Seber (CJS) [4, 5, 6] 模型的几种变体。有关这些模型的讨论，请参阅参考文献 [1]。
我们使用了两个数据集
参考文献 [2] 中的欧洲河乌 (Cinclus cinclus) 数据（这是挪威的国鸟）。
参考文献 [3] 中的草地田鼠数据。
与 [7] 中的 Stan 实现进行比较。
参考文献
Kery, M., & Schaub, M. (2011). 使用 WinBUGS 进行贝叶斯种群分析：一个分层视角。Academic Press.
Lebreton, J.D., Burnham, K.P., Clobert, J., & Anderson, D.R. (1992). 使用标记动物建模生存和检验生物学假设：带有案例研究的统一方法。Ecological monographs, 62(1), 67-118.
Nichols, Pollock, Hines (1984) 在小型哺乳动物种群研究中使用稳健的捕获-再捕获设计：宾夕法尼亚草地田鼠的野外实例。Acta Theriologica 29:357-365.
Cormack, R.M., 1964. 从标记动物的目击估算生存率。Biometrika 51, 429-438.
Jolly, G.M., 1965. 从捕获-再捕获数据中获取明确估计（包括死亡和迁入-随机模型）。Biometrika 52, 225-247.
Seber, G.A.F., 1965. 关于多次再捕获普查的注意事项。Biometrika 52, 249-259.
https://github.com/stan-dev/example-models/tree/master/BPA/Ch.07
import argparse
import os
from jax import random
import jax.numpy as jnp
from jax.scipy.special import expit, logit
import numpyro
from numpyro import handlers
from numpyro.contrib.control_flow import scan
import numpyro.distributions as dist
from numpyro.examples.datasets import DIPPER_VOLE, load_dataset
from numpyro.infer import HMC, MCMC, NUTS
from numpyro.infer.reparam import LocScaleReparam
我们的第一个也是最简单的 CJS 模型变体只有两个连续（标量）潜在随机变量：i) 生存概率 phi；ii) 再捕获概率 rho。这些被视为固定效应，没有时间和个体/群体变异。
def model_1(capture_history, sex):
N, T = capture_history.shape
phi = numpyro.sample("phi", dist.Uniform(0.0, 1.0)) # survival probability
rho = numpyro.sample("rho", dist.Uniform(0.0, 1.0)) # recapture probability
def transition_fn(carry, y):
first_capture_mask, z = carry
with numpyro.plate("animals", N, dim=-1):
with handlers.mask(mask=first_capture_mask):
mu_z_t = first_capture_mask * phi * z + (1 - first_capture_mask)
# NumPyro exactly sums out the discrete states z_t.
z = numpyro.sample(
"z",
dist.Bernoulli(dist.util.clamp_probs(mu_z_t)),
infer={"enumerate": "parallel"},
)
mu_y_t = rho * z
numpyro.sample(
"y", dist.Bernoulli(dist.util.clamp_probs(mu_y_t)), obs=y
)
first_capture_mask = first_capture_mask | y.astype(bool)
return (first_capture_mask, z), None
z = jnp.ones(N, dtype=jnp.int32)
# we use this mask to eliminate extraneous log probabilities
# that arise for a given individual before its first capture.
first_capture_mask = capture_history[:, 0].astype(bool)
# NB swapaxes: we move time dimension of `capture_history` to the front to scan over it
scan(
transition_fn,
(first_capture_mask, z),
jnp.swapaxes(capture_history[:, 1:], 0, 1),
)
在我们的第二个模型变体中，对于捕获数据的 T 个时间段中的 T-1 个，存在一个随时间变化的生存概率 phi_t；每个 phi_t 都被视为固定效应。
def model_2(capture_history, sex):
N, T = capture_history.shape
rho = numpyro.sample("rho", dist.Uniform(0.0, 1.0)) # recapture probability
def transition_fn(carry, y):
first_capture_mask, z = carry
# note that phi_t needs to be outside the plate, since
# phi_t is shared across all N individuals
phi_t = numpyro.sample("phi", dist.Uniform(0.0, 1.0))
with numpyro.plate("animals", N, dim=-1):
with handlers.mask(mask=first_capture_mask):
mu_z_t = first_capture_mask * phi_t * z + (1 - first_capture_mask)
# NumPyro exactly sums out the discrete states z_t.
z = numpyro.sample(
"z",
dist.Bernoulli(dist.util.clamp_probs(mu_z_t)),
infer={"enumerate": "parallel"},
)
mu_y_t = rho * z
numpyro.sample(
"y", dist.Bernoulli(dist.util.clamp_probs(mu_y_t)), obs=y
)
first_capture_mask = first_capture_mask | y.astype(bool)
return (first_capture_mask, z), None
z = jnp.ones(N, dtype=jnp.int32)
# we use this mask to eliminate extraneous log probabilities
# that arise for a given individual before its first capture.
first_capture_mask = capture_history[:, 0].astype(bool)
# NB swapaxes: we move time dimension of `capture_history` to the front to scan over it
scan(
transition_fn,
(first_capture_mask, z),
jnp.swapaxes(capture_history[:, 1:], 0, 1),
)
在我们的第三个模型变体中，对于捕获数据的 T 个时间段中的 T-1 个，存在一个生存概率 phi_t（与模型 2 类似），但这里每个 phi_t 都被视为随机效应。
def model_3(capture_history, sex):
N, T = capture_history.shape
phi_mean = numpyro.sample(
"phi_mean", dist.Uniform(0.0, 1.0)
) # mean survival probability
phi_logit_mean = logit(phi_mean)
# controls temporal variability of survival probability
phi_sigma = numpyro.sample("phi_sigma", dist.Uniform(0.0, 10.0))
rho = numpyro.sample("rho", dist.Uniform(0.0, 1.0)) # recapture probability
def transition_fn(carry, y):
first_capture_mask, z = carry
with handlers.reparam(config={"phi_logit": LocScaleReparam(0)}):
phi_logit_t = numpyro.sample(
"phi_logit", dist.Normal(phi_logit_mean, phi_sigma)
)
phi_t = expit(phi_logit_t)
with numpyro.plate("animals", N, dim=-1):
with handlers.mask(mask=first_capture_mask):
mu_z_t = first_capture_mask * phi_t * z + (1 - first_capture_mask)
# NumPyro exactly sums out the discrete states z_t.
z = numpyro.sample(
"z",
dist.Bernoulli(dist.util.clamp_probs(mu_z_t)),
infer={"enumerate": "parallel"},
)
mu_y_t = rho * z
numpyro.sample(
"y", dist.Bernoulli(dist.util.clamp_probs(mu_y_t)), obs=y
)
first_capture_mask = first_capture_mask | y.astype(bool)
return (first_capture_mask, z), None
z = jnp.ones(N, dtype=jnp.int32)
# we use this mask to eliminate extraneous log probabilities
# that arise for a given individual before its first capture.
first_capture_mask = capture_history[:, 0].astype(bool)
# NB swapaxes: we move time dimension of `capture_history` to the front to scan over it
scan(
transition_fn,
(first_capture_mask, z),
jnp.swapaxes(capture_history[:, 1:], 0, 1),
)
在我们的第四个模型变体中，我们包含了按性别（雄性、雌性）分组的固定效应。
def model_4(capture_history, sex):
N, T = capture_history.shape
# survival probabilities for males/females
phi_male = numpyro.sample("phi_male", dist.Uniform(0.0, 1.0))
phi_female = numpyro.sample("phi_female", dist.Uniform(0.0, 1.0))
# we construct a N-dimensional vector that contains the appropriate
# phi for each individual given its sex (female = 0, male = 1)
phi = sex * phi_male + (1.0 - sex) * phi_female
rho = numpyro.sample("rho", dist.Uniform(0.0, 1.0)) # recapture probability
def transition_fn(carry, y):
first_capture_mask, z = carry
with numpyro.plate("animals", N, dim=-1):
with handlers.mask(mask=first_capture_mask):
mu_z_t = first_capture_mask * phi * z + (1 - first_capture_mask)
# NumPyro exactly sums out the discrete states z_t.
z = numpyro.sample(
"z",
dist.Bernoulli(dist.util.clamp_probs(mu_z_t)),
infer={"enumerate": "parallel"},
)
mu_y_t = rho * z
numpyro.sample(
"y", dist.Bernoulli(dist.util.clamp_probs(mu_y_t)), obs=y
)
first_capture_mask = first_capture_mask | y.astype(bool)
return (first_capture_mask, z), None
z = jnp.ones(N, dtype=jnp.int32)
# we use this mask to eliminate extraneous log probabilities
# that arise for a given individual before its first capture.
first_capture_mask = capture_history[:, 0].astype(bool)
# NB swapaxes: we move time dimension of `capture_history` to the front to scan over it
scan(
transition_fn,
(first_capture_mask, z),
jnp.swapaxes(capture_history[:, 1:], 0, 1),
)
在我们的最终模型变体中，我们对生存概率 phi 同时包含固定的分组效应和固定的时间效应：logit(phi_t) = beta_group + gamma_t 我们需要注意模型不要过参数化；为此，我们实际上让一个标量 beta 有效地编码了雄性和雌性生存概率的差异。
def model_5(capture_history, sex):
N, T = capture_history.shape
# phi_beta controls the survival probability differential
# for males versus females (in logit space)
phi_beta = numpyro.sample("phi_beta", dist.Normal(0.0, 10.0))
phi_beta = sex * phi_beta
rho = numpyro.sample("rho", dist.Uniform(0.0, 1.0)) # recapture probability
def transition_fn(carry, y):
first_capture_mask, z = carry
phi_gamma_t = numpyro.sample("phi_gamma", dist.Normal(0.0, 10.0))
phi_t = expit(phi_beta + phi_gamma_t)
with numpyro.plate("animals", N, dim=-1):
with handlers.mask(mask=first_capture_mask):
mu_z_t = first_capture_mask * phi_t * z + (1 - first_capture_mask)
# NumPyro exactly sums out the discrete states z_t.
z = numpyro.sample(
"z",
dist.Bernoulli(dist.util.clamp_probs(mu_z_t)),
infer={"enumerate": "parallel"},
)
mu_y_t = rho * z
numpyro.sample(
"y", dist.Bernoulli(dist.util.clamp_probs(mu_y_t)), obs=y
)
first_capture_mask = first_capture_mask | y.astype(bool)
return (first_capture_mask, z), None
z = jnp.ones(N, dtype=jnp.int32)
# we use this mask to eliminate extraneous log probabilities
# that arise for a given individual before its first capture.
first_capture_mask = capture_history[:, 0].astype(bool)
# NB swapaxes: we move time dimension of `capture_history` to the front to scan over it
scan(
transition_fn,
(first_capture_mask, z),
jnp.swapaxes(capture_history[:, 1:], 0, 1),
)
执行推断
models = {
name[len("model_") :]: model
for name, model in globals().items()
if name.startswith("model_")
}
def run_inference(model, capture_history, sex, rng_key, args):
if args.algo == "NUTS":
kernel = NUTS(model)
elif args.algo == "HMC":
kernel = HMC(model)
mcmc = MCMC(
kernel,
num_warmup=args.num_warmup,
num_samples=args.num_samples,
num_chains=args.num_chains,
progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
)
mcmc.run(rng_key, capture_history, sex)
mcmc.print_summary()
return mcmc.get_samples()
def main(args):
# load data
if args.dataset == "dipper":
capture_history, sex = load_dataset(DIPPER_VOLE, split="dipper", shuffle=False)[
1
]()
elif args.dataset == "vole":
if args.model in ["4", "5"]:
raise ValueError(
"Cannot run model_{} on meadow voles data, since we lack sex "
"information for these animals.".format(args.model)
)
(capture_history,) = load_dataset(DIPPER_VOLE, split="vole", shuffle=False)[1]()
sex = None
else:
raise ValueError("Available datasets are 'dipper' and 'vole'.")
N, T = capture_history.shape
print(
"Loaded {} capture history for {} individuals collected over {} time periods.".format(
args.dataset, N, T
)
)
model = models[args.model]
rng_key = random.PRNGKey(args.rng_seed)
run_inference(model, capture_history, sex, rng_key, args)
if __name__ == "__main__":
assert numpyro.__version__.startswith("0.18.0")
parser = argparse.ArgumentParser(
description="CJS capture-recapture model for ecological data"
)
parser.add_argument(
"-m",
"--model",
default="1",
type=str,
help="one of: {}".format(", ".join(sorted(models.keys()))),
)
parser.add_argument("-d", "--dataset", default="dipper", type=str)
parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
parser.add_argument("--num-chains", nargs="?", default=1, type=int)
parser.add_argument(
"--rng_seed", default=0, type=int, help="random number generator seed"
)
parser.add_argument(
"--algo", default="NUTS", type=str, help='whether to run "NUTS" or "HMC"'
)
args = parser.parse_args()
main(args)