{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sampling Galactic binaries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we will walk through the process of building a likelihood callable to feed into\n",
    "an MCMC sampler. We will see\n",
    "\n",
    "1. How to generate a mock data with [FastGB](https://gitlab.in2p3.fr/LISA/fastgb),\n",
    "2. How to generate a frequency domain waveform in the format of {class}`~.types.ProjectedWaveform`.\n",
    "\n",
    "Some snippets of code are adapted from the [LISA Data Generation and Analysis Workshop tutorial](https://indico.in2p3.fr/event/33255/contributions/142249/attachments/87343/132314/go).\n",
    "\n",
    "## Dependencies\n",
    "- [numpy](https://numpy.org/): the fundamental package for scientific computing with Python.\n",
    "- [fastgb](https://gitlab.in2p3.fr/LISA/fastgb): a fast waveform generator for Galactic binaries in the LISA band.\n",
    "- [LISA Orbits](https://gitlab.in2p3.fr/lisa-simulation/orbits): a Python package to generate LISA orbits."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's get started!\n",
    "\n",
    "A GB source is described by 8 parameters. We can use a [`NamedTuple`](https://docs.python.org/3/library/typing.html#typing.NamedTuple) to encapsulate them. This is not necessary, but it makes the code more readable and maintainable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from typing import NamedTuple\n",
    "\n",
    "import lovelyplots\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import numpy.typing as npt\n",
    "\n",
    "plt.style.use([\"paper\", \"use_mathtext\", \"colors5\"])\n",
    "del lovelyplots\n",
    "\n",
    "Param = float | np.floating | npt.NDArray[np.floating]\n",
    "\n",
    "\n",
    "class GBSource(NamedTuple):\n",
    "    \"\"\"Encapsulator for GB source parameters.\"\"\"\n",
    "\n",
    "    amp: Param\n",
    "    \"\"\"Amplitude\"\"\"\n",
    "\n",
    "    f0: Param\n",
    "    \"\"\"Frequency\"\"\"\n",
    "\n",
    "    fdot: Param\n",
    "    \"\"\"Frequency derivative\"\"\"\n",
    "\n",
    "    phi0: Param\n",
    "    \"\"\"Initial phase\"\"\"\n",
    "\n",
    "    inc: Param\n",
    "    \"\"\"Inclination angle\"\"\"\n",
    "\n",
    "    psi: Param\n",
    "    \"\"\"Polarisation angle\"\"\"\n",
    "\n",
    "    lam: Param\n",
    "    \"\"\"Ecliptic longitude\"\"\"\n",
    "\n",
    "    beta: Param\n",
    "    \"\"\"Ecliptic latitude\"\"\"\n",
    "\n",
    "    def asdict(self):\n",
    "        \"\"\"Convert to dictionary.\"\"\"\n",
    "        return self._asdict()\n",
    "\n",
    "    @classmethod\n",
    "    def keys(cls):\n",
    "        \"\"\"Return the keys.\"\"\"\n",
    "        return cls._fields"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To build the waveform, we can wrap the fast GB generator with a class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastgb.fastgb import FastGB\n",
    "\n",
    "import typed_lisa_toolkit as tlt\n",
    "\n",
    "\n",
    "class GBGen:\n",
    "    \"\"\"Class to generate GB waveforms.\"\"\"\n",
    "\n",
    "    def __init__(self, fastgb: FastGB):\n",
    "        self.fastgb = fastgb\n",
    "\n",
    "    def __call__(self, params: GBSource):\n",
    "        \"\"\"Generate a GB waveform.\"\"\"\n",
    "        params_dict = params.asdict()\n",
    "        # FastGB accepts the parameters in a fixed order\n",
    "        fastgb_param_order = [\"f0\", \"fdot\", \"amp\", \"beta\", \"lam\", \"psi\", \"inc\", \"phi0\"]\n",
    "        params_array = np.column_stack([params_dict[key] for key in fastgb_param_order])\n",
    "        X, Y, Z, kmin = self.fastgb.get_fd_tdixyz(params_array, tdi2=False)\n",
    "        A, E, T = tlt.shop.xyz2aet(X=X, Y=Y, Z=Z)\n",
    "        # Get the frequencies\n",
    "        df = 1 / self.fastgb.T\n",
    "        f_min, f_max = df * kmin, df * (kmin + self.fastgb.N)\n",
    "        freqs = np.linspace(f_min, f_max, self.fastgb.N, axis=-1, endpoint=False)\n",
    "\n",
    "        for i in range(freqs.shape[0]):\n",
    "            if i != 0 and not np.array_equal(freqs[0], freqs[i]):\n",
    "                msg = \"Frequencies are not the same across all sources.\"\n",
    "                raise ValueError(msg)\n",
    "\n",
    "        A, E, T, freqs = (\n",
    "            A[:, None, None, None, :],\n",
    "            E[:, None, None, None, :],\n",
    "            T[:, None, None, None, :],\n",
    "            tlt.linspace_from_array(freqs[0]),\n",
    "        )\n",
    "\n",
    "        # Encapsulate the result\n",
    "        result = tlt.projected_waveform(\n",
    "            {\n",
    "                \"A\": tlt.frequency_series(freqs, A),\n",
    "                \"E\": tlt.frequency_series(freqs, E),\n",
    "                \"T\": tlt.frequency_series(freqs, T),\n",
    "            },\n",
    "        )\n",
    "        return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us consider one year of observation with LISA, with cadence of 5 seconds. Note that the observation time in data analysis must be an integer multiple of the cadence, so not exactly one year strictly speaking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lisaconstants  # this is a dependency of lisaorbits\n",
    "import lisaorbits\n",
    "\n",
    "T_obs, cadence = lisaconstants.ASTRONOMICAL_YEAR, 5\n",
    "T_obs = np.rint(T_obs / cadence) * cadence\n",
    "\n",
    "orbits = lisaorbits.EqualArmlengthOrbits()\n",
    "gb_gen = GBGen(FastGB(orbits=orbits, T=T_obs, delta_t=cadence, N=256))\n",
    "\n",
    "# Prepare a GB source to inject\n",
    "inj_source = GBSource(\n",
    "    amp=1.8326657820908912e-23,\n",
    "    f0=0.006009098866294652,\n",
    "    fdot=2.755088767708786e-16,\n",
    "    phi0=2.504337144699295,\n",
    "    inc=0.8694902600959342,\n",
    "    psi=5.053288418125767,\n",
    "    lam=5.17647019318475,\n",
    "    beta=0.6865491022703035,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To generate the mock data, we need to insert the generated waveform into the frequency grid corresponding to the observation time and cadence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_frequencies = tlt.axis(np.fft.rfftfreq(int(np.rint(T_obs / cadence)), d=cadence))\n",
    "\n",
    "inj_wav = gb_gen(inj_source)\n",
    "inj_data = tlt.fsdata(inj_wav)\n",
    "inj_data = inj_data.get_embedded((data_frequencies,))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the waveform:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = tlt.plot(inj_wav, plot_mode=\"semilogy\", set_legend=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And also in the full frequency range:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = tlt.plot(inj_data, set_legend=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To build the likelihood callable, we can use {class}`~.types.FDWhittleLikelihood`. The last piece we need is the noise power spectral density. We use the 7 parameter noise model from the Gee-Moo global-fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def parametric_noise_7p(pars, frqs, option=\"A\", arm_len=2.5e9, *, log_amp=False):\n",
    "    \"\"\"Return noise psd for 7 parameter noise model.\"\"\"\n",
    "    psd_acc, psd_oms, amp, freq1, freq2, alpha, fknee = pars\n",
    "    if log_amp:\n",
    "        psd_acc = 10**psd_acc\n",
    "        psd_oms = 10**psd_oms\n",
    "        amp = 10**amp\n",
    "\n",
    "    omega = 2.0 * np.pi * frqs\n",
    "\n",
    "    psd_pm = (\n",
    "        psd_acc\n",
    "        * (1.0 + (0.4e-3 / frqs) ** 2)\n",
    "        * (1.0 + (frqs / 8.0e-3) ** 4)\n",
    "        / (lisaconstants.SPEED_OF_LIGHT * omega) ** 2\n",
    "    )\n",
    "\n",
    "    psd_op = (\n",
    "        psd_oms\n",
    "        * (1.0 + (2.0e-3 / frqs) ** 4)\n",
    "        * (omega / lisaconstants.SPEED_OF_LIGHT) ** 2\n",
    "    )\n",
    "\n",
    "    x = omega * (arm_len / lisaconstants.SPEED_OF_LIGHT)\n",
    "\n",
    "    if option in [\"A\", \"E\"]:\n",
    "        psd_instr = (\n",
    "            8.0\n",
    "            * np.sin(x) ** 2\n",
    "            * (\n",
    "                2.0 * psd_pm * (3.0 + 2.0 * np.cos(x) + np.cos(2 * x))\n",
    "                + psd_op * (2.0 + np.cos(x))\n",
    "            )\n",
    "        )\n",
    "\n",
    "    elif option == \"X\":\n",
    "        psd_instr = (\n",
    "            16.0 * np.sin(x) ** 2 * (2.0 * (1.0 + np.cos(x) ** 2) * psd_pm + psd_op)\n",
    "        )\n",
    "\n",
    "    elif option == \"T\":\n",
    "        psd_instr = (\n",
    "            16.0 * psd_op * (1.0 - np.cos(x)) * np.sin(x) ** 2\n",
    "            + 128.0 * psd_pm * np.sin(x) ** 2 * np.sin(0.5 * x) ** 4\n",
    "        )\n",
    "\n",
    "    else:\n",
    "        msg = \"Only channels A, E, T, X are available.\"\n",
    "        raise NotImplementedError(msg)\n",
    "\n",
    "    if option != \"T\":\n",
    "        psd_gal = (\n",
    "            amp\n",
    "            * frqs ** (-7.0 / 3.0)\n",
    "            * np.exp(-((frqs / freq1) ** alpha))\n",
    "            * 0.5\n",
    "            * (1.0 + np.tanh(-(frqs - fknee) / freq2))\n",
    "        )\n",
    "        psd_gal = 4.0 * (x * np.sin(x)) ** 2 * psd_gal\n",
    "\n",
    "        if option in [\"A\", \"E\"]:\n",
    "            psd_gal = 1.5 * psd_gal\n",
    "    else:\n",
    "        psd_gal = 0.0\n",
    "\n",
    "    return psd_instr + psd_gal\n",
    "\n",
    "\n",
    "class P7Noise:\n",
    "    \"\"\"7 parameter noise model.\"\"\"\n",
    "\n",
    "    def __init__(self, pars, *, log_amp=False):\n",
    "        self.pars = pars\n",
    "        self.log_amp = log_amp\n",
    "\n",
    "    def psd(self, frequencies, option: str = \"A\"):\n",
    "        \"\"\"Return the parametric noise PSD.\"\"\"\n",
    "        if frequencies[0] == 0:\n",
    "            frequencies[0] = 1e-10 * frequencies[1]  # avoid division by zero\n",
    "        return parametric_noise_7p(\n",
    "            self.pars,\n",
    "            frequencies,\n",
    "            option=option,\n",
    "            log_amp=self.log_amp,\n",
    "        )\n",
    "\n",
    "\n",
    "# Set the noise parameters\n",
    "noise_params = [\n",
    "    5.55620285e-30,\n",
    "    5.44814245e-23,\n",
    "    1.44305853e-44,\n",
    "    3.13562154e-03,\n",
    "    1.88869275e-03,\n",
    "    1.78907291e00,\n",
    "    2.46721224e-03,\n",
    "]\n",
    "\n",
    "p7noise = P7Noise(noise_params, log_amp=False)\n",
    "\n",
    "from typed_lisa_toolkit import make_sdm, noise_model\n",
    "\n",
    "channel_names = (\"A\", \"E\", \"T\")\n",
    "ipsd_diag = np.stack(\n",
    "    [np.squeeze(1 / p7noise.psd(data_frequencies.ax, option=c)) for c in channel_names],\n",
    "    axis=-1,\n",
    ")\n",
    "sdm = make_sdm(\n",
    "    ipsd_diag,\n",
    "    frequencies=data_frequencies,\n",
    "    channel_names=channel_names,\n",
    "    is_diagonal=True,\n",
    ")\n",
    "noisemodel = noise_model(sdm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we have the likelihood callable, which can be fed into an MCMC sampler. Let's check the result with the injected parameters (it should be close to 0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typed_lisa_toolkit import whittle\n",
    "from typed_lisa_toolkit.types import FDNoiseModel, FSData\n",
    "\n",
    "\n",
    "def src_wrap(x: np.ndarray):\n",
    "    \"\"\"Return a GB source.\"\"\"\n",
    "    # If x is of shape (n, ) reshape it to (1,n)\n",
    "    if len(x.shape) == 1:\n",
    "        x = x.reshape((1, -1))\n",
    "    params = np.split(x, axis=1, indices_or_sections=x.shape[1])\n",
    "    src = GBSource(*params)\n",
    "    return src\n",
    "\n",
    "\n",
    "class LogLikelihood:\n",
    "    \"\"\"Log likelihood computation class.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        data: FSData,\n",
    "        noisemodel: FDNoiseModel,\n",
    "        waveform_generator: GBGen,\n",
    "    ):\n",
    "        self.likelihood = whittle(data, noisemodel)\n",
    "        self.waveform_generator = waveform_generator\n",
    "\n",
    "    def __call__(self, x: np.ndarray):\n",
    "        \"\"\"Return the log-likelihood.\"\"\"\n",
    "        source = src_wrap(x)\n",
    "        return self.likelihood.get_log_likelihood(self.waveform_generator(source))\n",
    "\n",
    "\n",
    "log_likelihood = LogLikelihood(inj_data, noisemodel.reset(), gb_gen)\n",
    "\n",
    "log_likelihood(\n",
    "    np.array(list(inj_source)),\n",
    ").sum()  # Should be zero modulo numerical errors for the injected source"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}