Skip to content
Snippets Groups Projects
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
dirichlet.py 1.07 KiB
import numpy as np
import pandas as pd
from bilby.core.likelihood import Multinomial
from bilby.core.prior import DirichletPriorDict
from bilby.core.sampler import run_sampler

n_dim = 3
label = "dirichlet_"
priors = DirichletPriorDict(n_dim=n_dim, label=label)

injection_parameters = dict(
    dirichlet_0=1 / 3,
    dirichlet_1=1 / 3,
    dirichlet_2=1 / 3,
)
data = [injection_parameters[label + str(ii)] * 1000 for ii in range(n_dim)]

likelihood = Multinomial(data=data, n_dimensions=n_dim, base=label)

result = run_sampler(
    likelihood=likelihood,
    priors=priors,
    nlive=100,
    label="multinomial",
    injection_parameters=injection_parameters,
)

result.posterior[label + str(n_dim - 1)] = 1 - np.sum(
    [result.posterior[key] for key in priors], axis=0
)
result.plot_corner(parameters=injection_parameters)

samples = priors.sample(10000)
samples[label + str(n_dim - 1)] = 1 - np.sum([samples[key] for key in samples], axis=0)
result.posterior = pd.DataFrame(samples)
result.plot_corner(
    parameters=[key for key in samples], filename="outdir/dirichlet_prior_corner.png"
)