<jupyter_start><jupyter_text>Normal Normal<jupyter_code>class NormalNormal:
    @bm.random_variable
    def mu(self):
        return dist.Normal(0, 1)

    @bm.random_variable
    def x(self, i):
        return dist.Normal(self.mu(), 1)


model = NormalNormal()
queries = [model.mu()]
obs = {
    model.x(0): torch.tensor(10.0),
}


def true_log_prob(mu):
    tau_0 = 1.0**-2
    tau_1 = 1.0**-2
    return dist.Normal(
        (tau_1 / (tau_0 + tau_1)) * 10.0,
        (tau_0 + tau_1)**-1,
    ).log_prob(mu)
vi = MeanFieldVariationalInference()
vi_dicts = vi.infer(queries, obs, num_iter=500, lr=1e-3, num_elbo_mc_samples=1000)
mu = torch.linspace(0, 10, steps=100).unsqueeze(1)
p_normal = torch.exp(vi_dicts(model.mu()).log_prob(mu))
sns.relplot(data=pd.DataFrame({
    "mu": mu.flatten().detach().numpy(),
    "p_normal": p_normal.flatten().detach().numpy(),
    "truth":  torch.exp(true_log_prob(mu)).flatten().detach().numpy(),
}).melt(id_vars=["mu"]), x="mu", y="value", hue="variable", kind="line")<jupyter_output><empty_output><jupyter_text>Fit the base distribution parameters too this time<jupyter_code>model = NormalNormal()
queries = [model.mu()]
obs = {
    model.x(0): torch.tensor(10.0),
}
vi = MeanFieldVariationalInference()
# SEE https://github.com/stefanwebb/flowtorch/issues/45, loc/scale must not be scalars
vi_dicts = vi.infer(queries, obs, base_dist=dist.Normal, base_args={
    'loc': nn.Parameter(torch.tensor([0.])),
    'scale': torch.tensor([1.]),
}, num_iter=100, lr=1e-3)
mu = torch.linspace(0, 10, steps=100).unsqueeze(1)
p_normal = torch.exp(vi_dicts(model.mu()).log_prob(mu))
sns.relplot(data=pd.DataFrame({
    "mu": mu.flatten().detach().numpy(),
    "p_normal": p_normal.flatten().detach().numpy(),
    "truth":  torch.exp(true_log_prob(mu)).flatten().detach().numpy(),
}).melt(id_vars=["mu"]), x="mu", y="value", hue="variable", kind="line")<jupyter_output>Training iterations:   0%|          | 0/100 [00:00<?, ?it/s]
Training iterations:  13%|█▎        | 13/100 [00:00<00:00, 123.33it/s]
Training iterations:  27%|██▋       | 27/100 [00:00<00:00, 132.02it/s]
Training iterations:  41%|████      | 41/100 [00:00<00:00, 135.45it/s]
Training iterations:  55%|█████▌    | 55/100 [00:00<00:00, 130.95it/s]
Training iterations:  69%|██████▉   | 69/100 [00:00<00:00, 83.54it/s] 
Training iterations:  80%|████████  | 80/100 [00:00<00:00, 81.45it/s]
Training iterations:  91%|█████████ | 91/100 [00:00<00:00, 87.09it/s]
Training iterations: 100%|██████████| 100/100 [00:01<00:00, 97.51it/s]<jupyter_text>Using a heavy-tailed base distribution<jupyter_code>vi_dicts_student = vi.infer(
    queries, obs, 
    base_dist=dist.StudentT,
    base_args={
        #'df': torch.tensor(1000000.0),
        'df': nn.Parameter(torch.tensor([10.0])),
        'loc': torch.tensor([0.0]),
        #'loc': nn.Parameter(torch.tensor(0.0)),
        #'scale': torch.tensor(1.0),
        #'scale': nn.Parameter(torch.tensor(1.0)),
    },
    num_iter=100, lr=1e-3)

#)
p_student = torch.exp(vi_dicts_student(model.mu()).log_prob(mu))
sns.relplot(data=pd.DataFrame({
    "mu": mu.flatten().detach().numpy(),
    "p_normal": p_normal.flatten().detach().numpy(),
    "p_student": p_student.flatten().detach().numpy(),
    "truth":  torch.exp(true_log_prob(mu)).flatten().detach().numpy(),
}).melt(id_vars=["mu"]), x="mu", y="value", hue="variable", kind="line")<jupyter_output><empty_output><jupyter_text>Scale Mixture Representation for StudentT[Ref](https://proceedings.neurips.cc/paper/2009/file/13fe9d84310e77f13a6d184dbf1232f3-Paper.pdf), implementation from https://www.johndcook.com/t_normal_mixture.pdf<jupyter_code>df = 4.0


@bm.random_variable
def W_inv(i):
    return dist.Gamma(concentration=0.5*df, rate=0.5*df)


@bm.random_variable
def y(i):
    return dist.Normal(0, 1.0 / torch.sqrt(W_inv(i)))


vi = MeanFieldVariationalInference()
vi_student_df = vi.infer(
    [y(0)], {},
    base_dist=dist.StudentT,
    base_args={
        'df': nn.Parameter(torch.tensor([df])),
        #'loc': torch.tensor(0.),
        #'scale': torch.tensor(1.),
    },
    num_iter=100, lr=1e-3
)
vi_student_oracle = vi.infer(
    [y(0)], {},
    base_dist=dist.StudentT,
    base_args={
        'df': torch.tensor([df]),
        #'loc': torch.tensor(0.),
        #'scale': torch.tensor(1.),
    },
    num_iter=100, lr=1e-3
)
vi_normal = vi.infer(
    [y(0)], {},
    num_iter=100, lr=1e-3)


def true_log_prob(x):
    return dist.StudentT(df=df).log_prob(x)
x = torch.linspace(-5, 5, steps=100).unsqueeze(1)
p_true = torch.exp(true_log_prob(x))
sns.relplot(data=pd.DataFrame({
    "y": x.flatten().detach().numpy(),
    "vi_student_df": torch.exp(vi_student_df(y(0)).log_prob(x)).flatten().detach().numpy(),
    "vi_student_oracle": torch.exp(vi_student_oracle(y(0)).log_prob(x)).flatten().detach().numpy(),
    "vi_normal": torch.exp(vi_normal(y(0)).log_prob(x)).flatten().detach().numpy(),
    "truth": p_true.flatten().detach().numpy(),
}).melt(id_vars=["y"]), x="y", y="value", hue="variable", kind="line")<jupyter_output><empty_output><jupyter_text>Look on semilog to see tail approximation error better<jupyter_code>g = sns.relplot(data=pd.DataFrame({
    "y": x.flatten().detach().numpy(),
    "vi_student_df": torch.exp(vi_student_df(y(0)).log_prob(x)).flatten().detach().numpy(),
    "vi_student_oracle": torch.exp(vi_student_oracle(y(0)).log_prob(x)).flatten().detach().numpy(),
    "vi_normal": torch.exp(vi_normal(y(0)).log_prob(x)).flatten().detach().numpy(),
    "truth": p_true.flatten().detach().numpy(),
}).melt(id_vars=["y"]), x="y", y="value", hue="variable", kind="line")
g.set(yscale='log')<jupyter_output><empty_output><jupyter_text>Monte carlo samples from ancestral (ground truth)<jupyter_code>samples = bm.SingleSiteAncestralMetropolisHastings().infer([y(0)], {}, num_samples=100, num_chains=1)
sns.distplot(samples[y(0)].flatten().detach().numpy())<jupyter_output>Samples collected:   0%|          | 0/100 [00:00<?, ?it/s]
Samples collected:  57%|█████▋    | 57/100 [00:00<00:00, 563.04it/s]
Samples collected: 100%|██████████| 100/100 [00:00<00:00, 572.71it/s]
{'fraction_non_alphanumeric': 0.2174983585029547, 'fraction_numerical': 0.061720288903479976, 'mean_word_length': 4.731890874882408, 'pattern_counts': {'":': 20, '<': 35, '<?xml version=': 0, '>': 23, 'https://': 3, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '568549', 'license': 'MIT', 'n_tokens_mistral': 2755, 'n_tokens_neox': 2470, 'n_words': 339, 'path': 'src/beanmachine/ppl/experimental/MeanFieldVI.ipynb', 'repo_name': 'michaeltingley/beanmachine-1', 'size': '6092'}