Exercise-2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BAT.jl: Background + Signal Analysis "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "In this exercise, we will get familiar with fitting a model to datasets of events, measured from your experiment, which are supplied as histograms. \n",
    "\n",
    "You will have access to two different datasets, which you will use separately to construct knowledge of the background of your experiment and later on to look for a signal using a Bayesian approach. The first dataset contains a large population of events and this is the one you will use to perform a Bayesian fit of the model for the background of your experiment. \n",
    "\n",
    "After you have obtained a model of the background, you can use the second dataset, which has a low event population, to test whether a signal is present by computing the evidence for both the background model only and the background model with a superimposed signal. \n",
    "\n",
    "Given the evidence for both hypotheses, you will decide whether a signal is present using the  [Bayes factor](https://en.wikipedia.org/wiki/Bayes_factor).\n",
    "\n",
    "***\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: This tutorial is written in the style of a cloze text. Parts that have to be filled out by you are marked `#=FILLME=#`.\n",
    "\n",
    "The [the tutorial section](https://bat.github.io/BAT.jl/stable/tutorial/) of the BAT.jl documentation should prove helpful.\n",
    "\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start we need to import some packages that we will need for our analysis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using BAT\n",
    "using Random\n",
    "using Distributions\n",
    "using IntervalSets\n",
    "using ValueShapes\n",
    "using Plots\n",
    "using StatsBase\n",
    "using LinearAlgebra\n",
    "using Statistics\n",
    "using ArraysOfArrays\n",
    "\n",
    "gr(palette = :inferno, size=(750,500))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Generate Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us generate `N_background=1500` samples for the background only data. And `N=440` samples for background+signal data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins=0:0.5:30\n",
    "N_background = 1500\n",
    "λ = 10\n",
    "\n",
    "data_background = rand(Truncated(Exponential(λ), 0, 30), N_background)\n",
    "hist_background = append!(Histogram(bins), data_background);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_signal = 40\n",
    "N_s_background = 400\n",
    "μ_true = 15.0\n",
    "σ_true = 1.0\n",
    "\n",
    "data_signal =  vcat(\n",
    "    rand(Truncated(Exponential(λ), 0, 30), N_s_background),\n",
    "    rand(Normal( μ_true, σ_true), N_signal)\n",
    ")\n",
    "\n",
    "hist_signal = append!(Histogram(bins), data_signal);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(normalize(hist_background, mode=:density), alpha=0.9, color=1, fillalpha = 0.05, linewidth = 0.0, label=\"\")\n",
    "plot!(normalize(hist_background, mode=:density), seriestype = :steps, alpha=0.9, color=1,  label=\"Background Data\")\n",
    "plot!(normalize(hist_signal, mode=:density), alpha=0.9, color=3, label=\"Signal Data\")\n",
    "plot!(xlabel = \"x\", ylabel=\"counts\", title=\"Histogram of the Data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Definition of Model Functions & Likelihood"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given the data histograms, we have to declare the model functions we want to fit. You can notice that you can use both a scalar ($\\lambda$) or a vector of variables ($A$) to define the parameters of your model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function background_model(p::NamedTuple{(:A, :λ)}, x::Real)\n",
    "    p.A[1] + p.A[2]*exp(-p.λ*x)\n",
    "end\n",
    "\n",
    "function background_signal_model(p::NamedTuple{(:A, :λ, :D, :x_0, :σ)}, x::Real)\n",
    "    p.A[1] + p.A[2]*exp(-p.λ*x) + (p.D/(sqrt(2*pi*p.σ^2)))*exp(-((x - p.x_0)^2)/p.σ^2)  \n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we need to define a likelihood function for our problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function make_hist_likelihood(h::Histogram, f::Function)\n",
    "    params -> begin\n",
    "        # Histogram counts for each bin as an array:\n",
    "        counts = h.weights\n",
    "\n",
    "        # Histogram binning, has length (length(counts) + 1):\n",
    "        binning = h.edges[1]\n",
    "\n",
    "        # sum log-likelihood value over bins:\n",
    "        ll_value::Float64 = 0.0\n",
    "        \n",
    "        for i in eachindex(counts)\n",
    "            #=FILLME=#\n",
    "            \n",
    "            # Hint: Hava a look at the tutorial in the BAT.jl documentation.\n",
    "            \n",
    "            ll_value += #=FILLME=#\n",
    "        end\n",
    "\n",
    "        return ll_value\n",
    "    end \n",
    "end "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "algorithm = MetropolisHastings()\n",
    "nsamples = 2*10^5\n",
    "nchains = 5\n",
    "ENV[\"JULIA_INFO\"] = \"BAT\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Background Analysis: "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll perform a Bayesian fit of the background histogram, using BAT, to infer the model parameters from the data histogram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "likelihood_bm = make_hist_likelihood(hist_background, background_model);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to choose a sensible prior for the fit. Try an interval from 0 to 1000 for the entries of a `A` (note that `A` is a vector-shaped parameter) and an interval from 0 to 20 for `λ`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior_bm = NamedTupleDist(\n",
    "    A = [#=FILLME=#, #=FILLME=#],\n",
    "    λ = #=FILLME=#\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parshapes_bm = valshape(prior_bm)\n",
    "posterior_bm = PosteriorDensity(likelihood_bm, prior_bm);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_bm, stats_bm = bat_sample(posterior_bm, (nsamples, nchains), algorithm)\n",
    "stats_bm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mode_parms_bm = parshapes_bm(stats_bm.mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(samples_bm, params=[1,2,3], param_labels=[\"A_1\",\"A_2\",\"\\$\\\\lambda\\$\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(normalize(hist_background, mode=:density), alpha=0.9, color=1, fillalpha = 0.05, linewidth = 0.0, label=\"\")\n",
    "plot!(normalize(hist_background, mode=:density), seriestype = :steps, alpha=0.9, color=1,  label=\"Background Data\")\n",
    "\n",
    "plot!(xlabel = \"x\", ylabel=\"counts\", title=\"Background Analysis\")\n",
    "\n",
    "plot!(\n",
    "    0:0.1:30, x -> background_model(mode_parms_bm, x),\n",
    "    label = \"Background Model\", lw=1.5, color=1,  linestyle = :dot\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Histogram as a prior: "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We would like now to use posterior distribution for $\\lambda$ as a prior for signal+background fit. Let us construct histogram for $\\lambda$ by using MCMC samples "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = flatview(samples_bm.params)[3,:]\n",
    "λ_hist = fit(Histogram, samples, FrequencyWeights(samples_bm.weight), nbins = 150)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Background + Signal Analysis: "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us use the second dataset, which has much fewer samples, to infer parameters of the signal model. Let us define a prior for the fit parameters including histogram for $\\lambda$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "likelihood_signal = make_hist_likelihood(hist_signal, background_signal_model);\n",
    "\n",
    "prior_signal = NamedTupleDist(\n",
    "    A = [0.0 .. 1000.0, 0 .. 10000.0],\n",
    "    λ = BAT.HistogramAsUvDistribution(λ_hist),\n",
    "    D = 0.0..1000, \n",
    "    x_0 = 0 .. 30, \n",
    "    σ = 0.5 .. 10.0\n",
    ");\n",
    "\n",
    "parshapes_signal = valshape(prior_signal)\n",
    "posterior_signal = PosteriorDensity(likelihood_signal, prior_signal);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_signal, stats_signal = bat_sample(posterior_signal, (nsamples, nchains), algorithm)\n",
    "stats_signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(samples_signal, params=[1,2,3], param_labels=[\"A_1\",\"A_2\",\"\\$\\\\lambda\\$\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(samples_signal, params=[4,5,6], param_labels=[\"D\",\"x_0\",\"\\$\\\\sigma\\$\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(samples_signal, 3)\n",
    "plot!(prior_signal, 3)\n",
    "plot!(ylabel=\"P(\\\\lambda)\", xlabel=\"\\\\lambda\", title=\"Knowledge Update\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mode_parms_signal = parshapes_signal(stats_signal.mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(normalize(hist_signal, mode=:density), fillalpha=0.65, color=3, label=\"Signal Data\")\n",
    "plot!(\n",
    "    0:0.1:30, x -> background_signal_model(mode_parms_signal, x),\n",
    "    label = \"Signal Model\", lw=2, color=6, linestyle = :dash\n",
    ")\n",
    "plot!(xlabel = \"x\", ylabel=\"counts\", title=\"Signal Analysis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Bayes factor for the signal model can be computed in the following way"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evidence_signal = bat_integrate(samples_signal).integral"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Signal vs. No signal: "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now analyze the same dataset under the assumption that we have no signal events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "likelihood_nosignal = make_hist_likelihood(hist_signal, background_model);\n",
    "\n",
    "prior_nosignal = NamedTupleDist(\n",
    "    A = [#=FILLME=#, #=FILLME=#],\n",
    "    λ = #=FILLME=#\n",
    ")\n",
    "\n",
    "parshapes_nosignal = valshape(prior_nosignal)\n",
    "posterior_nosignal = PosteriorDensity(likelihood_nosignal, prior_nosignal);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_nosignal, stats_nosignal = bat_sample(posterior_nosignal, (nsamples, nchains), algorithm)\n",
    "stats_nosignal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mode_parms_nosignal = parshapes_nosignal(stats_nosignal.mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(normalize(hist_signal, mode=:density), fillalpha=0.65, color=3, label=\"Signal Data\")\n",
    "plot!(\n",
    "    0:0.1:30, x -> background_model(mode_parms_nosignal, x),\n",
    "    label = \"Background Model\", lw=2, color=1,  linestyle = :dot\n",
    ")\n",
    "plot!(\n",
    "    0:0.1:30, x -> background_signal_model(mode_parms_signal, x),\n",
    "    label = \"Signal Model\", lw=2, color=6, linestyle = :dash\n",
    ")\n",
    "plot!(xlabel = \"x\", ylabel=\"counts\", title=\"Signal Analysis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Bayes factor for this case is equal to"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evidence_nosignal = bat_integrate(samples_nosignal).integral"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let us compare Bayes factors for these two assumptions: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evidence_signal / evidence_nosignal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    Tasks: \n",
    "    \n",
    "    1) Change the number of signal/background events and see for which values we no longer detect the signal. Explore how the binning of the data effects analysis. \n",
    "    \n",
    "    2) Assume that we do not have the dataset with the background only measurements, and we use a flat prior for the parameter lambda. Explore how this effects MCMC sampling. "
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "hide_input": false,
  "kernelspec": {
   "display_name": "Julia 1.3.0-rc4",
   "language": "julia",
   "name": "julia-1.3"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.3.0"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}