def true_ape(ns):
"""Analytic APE"""
true_ape = []
prior_cov = torch.diag(prior_sds**2)
designs = [
group_assignment_matrix(torch.tensor([n1, N - n1])) for n1 in ns
]
for i in range(len(ns)):
x = designs[i]
posterior_cov = analytic_posterior_cov(prior_cov, x, torch.tensor(1.))
true_ape.append(0.5 * torch.logdet(2 * np.pi * np.e * posterior_cov))
return torch.tensor(true_ape)
def estimated_ape(ns, num_vi_steps):
designs = [
group_assignment_matrix(torch.tensor([n1, N - n1])) for n1 in ns
]
X = torch.stack(designs)
est_ape = vi_ape(
model,
X,
observation_labels="y",
target_labels="w",
vi_parameters={
"guide": guide,
"optim": optim.Adam({"lr": 0.05}),
"loss": TraceEnum_ELBO(
strict_enumeration_warning=False).differentiable_loss,
"num_steps": num_vi_steps
},
is_parameters={"num_samples": 1})
return est_ape
- Donsker-Varadhan
- Barber-Agakov
TODO:
- better guides- allow different levels of amortization
- SVI with BA-style guides
"""
#########################################################################################
# Designs
#########################################################################################
# All design tensors have shape: batch x n x p
# AB test
AB_test_11d_10n_2p = torch.stack(
[group_assignment_matrix(torch.tensor([n, 10 - n])) for n in range(0, 11)])
AB_test_2d_10n_2p = torch.stack(
[group_assignment_matrix(torch.tensor([n, 10 - n])) for n in [0, 5]])
# Design on S^1
item_thetas = torch.linspace(0., np.pi, 10).unsqueeze(-1)
X_circle_10d_1n_2p = torch.stack([item_thetas.cos(), -item_thetas.sin()],
dim=-1)
item_thetas_small = torch.linspace(0., np.pi / 2, 5).unsqueeze(-1)
X_circle_5d_1n_2p = torch.stack(
[item_thetas_small.cos(), -item_thetas_small.sin()], dim=-1)
# Random effects designs
AB_test_reff_6d_10n_12p, AB_sigmoid_design_6d = rf_group_assignments(10)
#########################################################################################
group_normal_guide)
from pyro.contrib.glmm.guides import LinearModelGuide, GuideDV
from pyro.contrib.oed.eig import (vi_ape, naive_rainforth_eig,
donsker_varadhan_eig, barber_agakov_ape)
from pyro.contrib.oed.util import linear_model_ground_truth
from pyro.infer import Trace_ELBO
logger = logging.getLogger(__name__)
#########################################################################################
# Designs
#########################################################################################
# All design tensors have shape: batch x n x p
# AB test
AB_test_2d_10n_2p = torch.stack(
[group_assignment_matrix(torch.tensor([n, 10 - n])) for n in [0, 5]])
# Design on S^1
item_thetas_small = torch.linspace(0., np.pi / 2, 5).unsqueeze(-1)
X_circle_5d_1n_2p = torch.stack(
[item_thetas_small.cos(), -item_thetas_small.sin()], dim=-1)
#########################################################################################
# Models
#########################################################################################
# Linear models
basic_2p_linear_model_sds_10_2pt5, basic_2p_guide = zero_mean_unit_obs_sd_lm(
torch.tensor([10., 2.5]))
_, basic_2p_guide_w1 = zero_mean_unit_obs_sd_lm(torch.tensor([10., 2.5]),
coef_label="w1")
basic_2p_ba_guide = lambda d: LinearModelGuide(d, {"w": 2}) # noqa: E731