simulation_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a59919c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import seaborn as sns\n",
    "from torch import distributions\n",
    "\n",
    "from torch import nn, optim\n",
    "import math\n",
    "from tqdm import tqdm\n",
    "\n",
    "from pyro import distributions as dist\n",
    "from pyro.distributions import transforms as T\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e16156a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "mmin = 4.59\n",
    "mmax = 86.22"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "245bace9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pm1(m1, alpha=2.63, mum=33.07, sigmam=5.69, lpeak=0.1):\n",
    "    m1p = m1.clone()\n",
    "    m1p = (alpha - 1) / mmin * (m1p / mmin) ** -alpha\n",
    "    m1p[m1 < mmin] = 1e-9    \n",
    "    m1g = m1.clone()\n",
    "    m1g = 1 / sigmam * 1 / math.sqrt(math.pi * 2) * torch.exp(-0.5 * (m1g - mum) ** 2 / sigmam ** 2) \n",
    "    return (1 - lpeak) * m1p + lpeak * m1g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d248d56c",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = torch.linspace(0, 100, 1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aced5415",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(m, pm1(m).log())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2adb9165",
   "metadata": {},
   "outputs": [],
   "source": [
    "def trapezoid(fx, deltax):\n",
    "    return (fx[:-1] + fx[1:]) / 2 * deltax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3725d1fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.linspace(mmin, 1024, 32768)\n",
    "dx = x[1] - x[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b06ed74",
   "metadata": {},
   "outputs": [],
   "source": [
    "trapezoid(pm1(x), dx).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02cb5c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pq(q, beta=1.26):\n",
    "    qc = q.clone()\n",
    "    qc = qc ** beta\n",
    "    qc[q <= 0] = 1e-9\n",
    "    qc[q > 1] = 1e-9\n",
    "    return qc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5388766c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def px(x, q, mux0=0.19, gamma=1.08, log10sx0=-1.06, alpha=-0.46, beta=-0.83):\n",
    "    mux = mux0 + alpha*(q - 0.5)\n",
    "    log10sx = log10sx0 + beta * (q - 0.5)\n",
    "    sx = 10 ** log10sx\n",
    "    return torch.exp(-0.5 * (x - mux) ** 2 / sx ** 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0053233e",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.linspace(-3, 3, 1024)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94b712df",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(x, px(x, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fbbb7cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.interpolate import interp1d\n",
    "from astropy.cosmology import Planck18 as cosmo\n",
    "\n",
    "z_axis = np.linspace(0,10,100000)\n",
    "dVdz = cosmo.differential_comoving_volume(z_axis).value/1e9*4*np.pi # Astropy dVcdz is per stradian\n",
    "dVdz_interp = interp1d(z_axis,dVdz)\n",
    "\n",
    "def z_distribution_unnormalized(z):\n",
    "    z = z.clone().numpy()\n",
    "    return torch.from_numpy(dVdz_interp(z)*(1+z)**1.7)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56544cb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def rejection_sampling(fn, rng, num_samples=1024, batch_size=32768, max_log_prob=0, n_dims=1):\n",
    "    z = torch.Tensor()\n",
    "    i = 0\n",
    "    while len(z) < num_samples:\n",
    "        eps = torch.rand((num_samples, n_dims))\n",
    "        min, max = rng\n",
    "        z_ = eps * (max - min) + min\n",
    "        prob = torch.rand(num_samples)\n",
    "\n",
    "        logprob = fn(z_).squeeze()\n",
    "        assert not torch.any(logprob > max_log_prob), (logprob.max(), max_log_prob)\n",
    "        prob_ = torch.exp(logprob - max_log_prob)\n",
    "        accept = prob_ > prob\n",
    "        z = torch.cat([z, z_[accept, :]])\n",
    "        i += 1\n",
    "    return z[:num_samples]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43e0f502",
   "metadata": {},
   "outputs": [],
   "source": [
    "m1 = rejection_sampling(lambda x: pm1(x).log(), (mmin, mmax))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1680d9fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "q = rejection_sampling(lambda x: pq(x).log(), (0, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2b7eafd",
   "metadata": {},
   "outputs": [],
   "source": [
    "Xeff = []\n",
    "for q_ in q:\n",
    "    Xeff.append(rejection_sampling(lambda x: px(x, q_).log(), (-1, 1), max_log_prob=0, num_samples=1))\n",
    "    \n",
    "Xeff = torch.cat(Xeff)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa60478a",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.kdeplot(Xeff.squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a39924c",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = rejection_sampling(lambda x: z_distribution_unnormalized(x).log(), (0, 10), max_log_prob=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e74c432a",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.cat([m1, q, Xeff, z], dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef5df107",
   "metadata": {},
   "outputs": [],
   "source": [
    "import corner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2596dee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = corner.corner(data.numpy(), labels=[r\"$m_1\\, [M_{\\odot}]$\", r\"$q$\", r\"$\\mathcal{X}_{eff}$\", r\"$z$\"], label_kwargs=dict(fontsize=22))\n",
    "\n",
    "for ax in fig.get_axes():\n",
    "    ax.tick_params(axis='both', labelsize=16)\n",
    "plt.tight_layout()\n",
    "\n",
    "# fig.subplots_adjust(right=1.5,top=1.5)\n",
    "plt.savefig('corner.pdf', bbox_inches='tight')\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a194e207",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(data, 'data.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c33f525f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}