Hello,
I am working with astrophysical models for parameter estimation of Binary Black Holes. The setup requires running a non-homogeneous Poisson Process with the likelihood defined as,
\displaystyle\mathcal{L}=\exp{\left(-\mu(\mathcal{R}, \Lambda)\right)}\prod_{n=1}^{N}\int p(\lambda | d_n)\mathcal{R}p(\lambda | \Lambda)\frac{1}{p(\lambda)}d\lambda\approx \exp{\left(-\mu(\mathcal{R}, \Lambda)\right)}\prod_{n=1}^{N}\frac{\mathcal{R}}{N_i}\sum_{i=1}^{N_i}\frac{p(\lambda_i|\Lambda)}{p(\lambda_i)}
The integral can be approximated to a Monte Carlo Integration, where the \lambda is drawn from the p(\lambda | d_n), and they are pre-computed and stored in files for each n. I am facing trouble in writing all of this in numpyro.
Here \lambda=(m_1,m_2) and \Lambda = (\alpha, k, m_\min, m_\max, \mathcal{R}) and,
\displaystyle p(\lambda|\Lambda)\propto\frac{m_1^{-\alpha-k}m_2^k}{m_1-m_\min}
We have taken k=0. This is the current code I am using for it,
import numpyro
from jax import numpy as jnp
from numpyro import distributions as dist
def model(lambda_n): # , raw_interpolator):
r"""
.. math::
\mathcal{L}(\mathcal{R},\Lambda)\propto\prod_{n=1}^{N}\frac{\mathcal{R}}{N_i}\sum_{i=1}^{N_i}\frac{p(\lambda_i/\Lambda)}{\pi(\lambda_i)}
"""
event, _, _ = lambda_n.shape
alpha_prior = dist.Uniform(-5.0, 5.0)
mmin_prior = dist.Uniform(5.0, 15.0)
mmax_prior = dist.Uniform(30.0, 190.0)
rate_prior = dist.LogUniform(10**-1, 10**6)
alpha = numpyro.sample("alpha", alpha_prior)
mmin = numpyro.sample("mmin", mmin_prior)
mmax = numpyro.sample("mmax", mmax_prior)
rate = numpyro.sample("rate", rate_prior)
ll = 0.0
with numpyro.plate("data", event) as n:
mass_model = Wysocki2019MassModel(
alpha_m=alpha,
k=0,
mmin=mmin,
mmax=mmax,
)
ll_i = 0.0
p = jnp.exp(mass_model.log_prob(lambda_n[n, :, :]))
ll_i = jnp.mean(p, axis=1)
ll_i *= rate
ll += jnp.log(ll_i)
# mean = expval_mc(alpha, mmin, mmax, rate, raw_interpolator)
# ll -= mean
return jnp.exp(ll)
Right now I am avoiding the exponential term but you can advise on it if you can.
Thank you very much!