nbs/04_cls.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp cls"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# neos.cls\n",
    "\n",
    "> Module containing the differentiable calculation of CLs values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import config\n",
    "from jax.experimental import stax\n",
    "import pyhf\n",
    "\n",
    "# avoid those precision errors!\n",
    "config.update(\"jax_enable_x64\", True)\n",
    "\n",
    "pyhf.set_backend(pyhf.tensor.jax_backend())\n",
    "\n",
    "from neos.fit import get_solvers\n",
    "from neos.transforms import *\n",
    "from neos.models import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "def cls_maker(nn_model_maker, solver_kwargs):\n",
    "    '''\n",
    "    Wraps the construction of the `cls_jax` method.\n",
    "    \n",
    "    Args:\n",
    "            nn_model_maker: Function that returns a Model object using the nn parameters.\n",
    "\n",
    "    Returns:\n",
    "            cls_jax: A callable function that takes the parameters of the observable as argument, \n",
    "            and returns an expected CLs value from testing the background-only model against the\n",
    "            nominal signal hypothesis (or whatever the value of 'test_mu' is)\n",
    "    '''\n",
    "    @jax.jit\n",
    "    def cls_jax(nn_params, test_mu):\n",
    "        g_fitter, c_fitter = get_solvers(nn_model_maker, **solver_kwargs)\n",
    "\n",
    "        m, bonlypars = nn_model_maker(nn_params)\n",
    "        exp_data = m.expected_data(bonlypars)\n",
    "        #print(f'exp_data: {exp_data}')\n",
    "        bounds = m.config.suggested_bounds()\n",
    "\n",
    "        # map these\n",
    "        initval = jnp.asarray([test_mu, 1.0])\n",
    "        transforms = solver_kwargs.get(\"pdf_transform\", False)\n",
    "        if transforms:\n",
    "            initval = to_inf_vec(initval, bounds)\n",
    "\n",
    "        # the constrained fit\n",
    "\n",
    "        #print('fitting constrained with init val %s setup %s', initval,[test_mu, nn_params])\n",
    "\n",
    "        numerator = (\n",
    "            to_bounded_vec(c_fitter(initval, [test_mu, nn_params]), bounds)\n",
    "            if transforms\n",
    "            else c_fitter(initval, [test_mu, nn_params])\n",
    "        )\n",
    "        \n",
    "        # don't have to fit these -- we know them for expected limits!\n",
    "        denominator = bonlypars  # to_bounded_vec(g_fitter(initval, [test_mu, nn_params]), bounds) if transforms else g_fitter(initval, [test_mu, nn_params])\n",
    "\n",
    "        # print(f\"constrained fit: {numerator}\")\n",
    "        # print(f\"global fit: {denominator}\")\n",
    "\n",
    "        # compute test statistic (lambda(µ))\n",
    "        profile_likelihood = -2 * (\n",
    "            m.logpdf(numerator, exp_data)[0] - m.logpdf(denominator, exp_data)[0]\n",
    "        )\n",
    "\n",
    "        # in exclusion fit zero out test stat if best fit µ^ is larger than test µ\n",
    "        muhat = denominator[0]\n",
    "        sqrtqmu = jnp.sqrt(\n",
    "            jnp.where(muhat < test_mu, profile_likelihood, 0.0)\n",
    "        )\n",
    "        # print(f\"sqrt(q(mu)): {sqrtqmu}\")\n",
    "        # compute CLs\n",
    "        nullval = sqrtqmu\n",
    "        altval = 0\n",
    "        CLsb = 1 - pyhf.tensorlib.normal_cdf(nullval)\n",
    "        CLb = 1 - pyhf.tensorlib.normal_cdf(altval)\n",
    "        CLs = CLsb / CLb\n",
    "        return CLs\n",
    "\n",
    "    return cls_jax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "def pyhf_cls_maker(nn_model_maker, solver_kwargs):\n",
    "#     @jax.jit\n",
    "    def cls_jax(nn_params, test_mu):\n",
    "        g_fitter, c_fitter = get_solvers(nn_model_maker, **solver_kwargs)\n",
    "\n",
    "        m, bonlypars = nn_model_maker(nn_params)\n",
    "        exp_data = m.expected_data(bonlypars)\n",
    "        #print(f'exp_data: {exp_data}')\n",
    "#         bounds = m.config.suggested_bounds()[0]\n",
    "        \n",
    "        names = m.config.par_order\n",
    "        bounds = m.config.suggested_bounds()\n",
    "\n",
    "        # map these\n",
    "        initval = jnp.asarray([test_mu, 1.0])\n",
    "        transforms = solver_kwargs.get(\"pdf_transform\", False)\n",
    "        if transforms:\n",
    "            initval = to_inf_vec(initval, bounds)\n",
    "\n",
    "        # the constrained fit\n",
    "\n",
    "        #print('fitting constrained with init val %s setup %s', initval,[test_mu, nn_params])\n",
    "\n",
    "        numerator = (\n",
    "            to_bounded_vec(c_fitter(initval, [test_mu, nn_params]), bounds)\n",
    "            if transforms\n",
    "            else c_fitter(initval, [test_mu, nn_params])\n",
    "        )\n",
    "\n",
    "        denominator = bonlypars  # to_bounded_vec(g_fitter(initval, [test_mu, nn_params]), bounds) if transforms else g_fitter(initval, [test_mu, nn_params])\n",
    "\n",
    "        # print(f\"constrained fit: {numerator}\")\n",
    "        # print(f\"global fit: {denominator}\")\n",
    "\n",
    "        # compute test statistic (lambda(µ))\n",
    "        profile_likelihood = -2 * (\n",
    "            m.logpdf(numerator, exp_data)[0] - m.logpdf(denominator, exp_data)[0]\n",
    "        )\n",
    "\n",
    "        # in exclusion fit zero out test stat if best fit µ^ is larger than test µ\n",
    "        muhat = denominator[0]\n",
    "        sqrtqmu = jnp.sqrt(\n",
    "            jnp.where(muhat < test_mu, profile_likelihood, 0.0)\n",
    "        )\n",
    "        # print(f\"sqrt(q(mu)): {sqrtqmu}\")\n",
    "        # compute CLs\n",
    "        nullval = sqrtqmu\n",
    "        altval = 0\n",
    "        CLsb = 1 - pyhf.tensorlib.normal_cdf(nullval)\n",
    "        CLb = 1 - pyhf.tensorlib.normal_cdf(altval)\n",
    "        CLs = CLsb / CLb\n",
    "        return CLs\n",
    "\n",
    "    return cls_jax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}