Example2_Run_simulations_in_parallel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This is a demonstration of how to run the physics simulation of a jet energy loss in a Relativistic Heavy Ion Collision. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The parameters we are interested in the physics model are as below. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " - Relevent to Parton Energy loss rates\n",
    "     - alpha_S\n",
    "     - mD_factor\n",
    "     - exponent_inel\n",
    "     - exponent_el"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import run_simple_energy_loss\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulation=run_simple_energy_loss.run_simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### simulation function\n",
    " - It will take a numpy array with dimension M x N as input.\n",
    "    - M number of design points\n",
    "    - N number of model parameters of the model\n",
    " - It will return a M x P numpy array.\n",
    "     - M number of design points\n",
    "     - P number of Momentem bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior_ranges= {'alpha_s':(0.05,0.5), 'exponent_inel':(-3,3), 'exponent_el':(-3,3), 'scale_inel':(0.01, 0.5), 'scale_el':(0.01, 0.5)}\n",
    "default_values= {'alpha_s': 0.2, 'exponent_inel': -1, 'exponent_el': 1, 'scale_inel': 0.2, 'scale_el': 0.2}\n",
    "name_to_position =  {'alpha_s':0, 'exponent_inel':1, 'exponent_el':2, 'scale_inel':3, 'scale_el':4}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_design(param_name, n_design):\n",
    "    \"\"\"Returns a design matrix by only changing one model parameter and keeping all other \n",
    "    parameters fixed at the defaul value\"\"\"\n",
    "    design_matrix = []\n",
    "    for i in default_values:\n",
    "        if i == param_name:\n",
    "            temp = np.linspace(*prior_ranges[i],n_design)\n",
    "        else:\n",
    "            temp = np.full((n_design,), default_values[i])\n",
    "        design_matrix.append(temp)\n",
    "    design = np.array(design_matrix).T\n",
    "    #print(f'Shape of the design is {design.shape}')\n",
    "    return design"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_design = 10\n",
    "p_mom=(1, 10, 10)\n",
    "pt_list = np.linspace(*p_mom)\n",
    "show_pt = [3, 5, 7]\n",
    "\n",
    "# Plot how R_AA changes only one parameter change while keeping all other fixed at the default value.\n",
    "def sim_1D(param):\n",
    "    print(f'Working on {param} design set \\n')\n",
    "    i=name_to_position[param]\n",
    "    design = make_design(param, n_design)\n",
    "    observables=simulation(design,*p_mom)\n",
    "    return observables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of cores in the computer 8\n",
      "Working on exponent_inel design set \n",
      "Working on alpha_s design set \n",
      "Working on exponent_el design set \n",
      "Working on scale_inel design set \n",
      "Working on scale_el design set \n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Working on 1/10 designWorking on 1/10 designWorking on 1/10 designWorking on 1/10 designWorking on 1/10 design\n",
      "\n",
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/dananjayaliyanage/miniconda3/envs/parton_loss/lib/python3.6/site-packages/scipy/integrate/_ode.py:1180: UserWarning: dopri5: step size becomes too small\n",
      "  self.messages.get(istate, unexpected_istate_msg)))\n",
      "/Users/dananjayaliyanage/miniconda3/envs/parton_loss/lib/python3.6/site-packages/scipy/integrate/_ode.py:1180: UserWarning: dopri5: step size becomes too small\n",
      "  self.messages.get(istate, unexpected_istate_msg)))\n",
      "/Users/dananjayaliyanage/miniconda3/envs/parton_loss/lib/python3.6/site-packages/scipy/integrate/_ode.py:1180: UserWarning: dopri5: step size becomes too small\n",
      "  self.messages.get(istate, unexpected_istate_msg)))\n",
      "/Users/dananjayaliyanage/miniconda3/envs/parton_loss/lib/python3.6/site-packages/scipy/integrate/_ode.py:1180: UserWarning: dopri5: step size becomes too small\n",
      "  self.messages.get(istate, unexpected_istate_msg)))\n",
      "/Users/dananjayaliyanage/miniconda3/envs/parton_loss/lib/python3.6/site-packages/scipy/integrate/_ode.py:1180: UserWarning: dopri5: step size becomes too small\n",
      "  self.messages.get(istate, unexpected_istate_msg)))\n",
      "/Users/dananjayaliyanage/miniconda3/envs/parton_loss/lib/python3.6/site-packages/scipy/integrate/_ode.py:1180: UserWarning: dopri5: larger nsteps is needed\n",
      "  self.messages.get(istate, unexpected_istate_msg)))\n",
      "/Users/dananjayaliyanage/miniconda3/envs/parton_loss/lib/python3.6/site-packages/scipy/integrate/_ode.py:1180: UserWarning: dopri5: larger nsteps is needed\n",
      "  self.messages.get(istate, unexpected_istate_msg)))\n",
      "/Users/dananjayaliyanage/miniconda3/envs/parton_loss/lib/python3.6/site-packages/scipy/integrate/_ode.py:1180: UserWarning: dopri5: larger nsteps is needed\n",
      "  self.messages.get(istate, unexpected_istate_msg)))\n",
      "/Users/dananjayaliyanage/miniconda3/envs/parton_loss/lib/python3.6/site-packages/scipy/integrate/_ode.py:1180: UserWarning: dopri5: larger nsteps is needed\n",
      "  self.messages.get(istate, unexpected_istate_msg)))\n",
      "/Users/dananjayaliyanage/miniconda3/envs/parton_loss/lib/python3.6/site-packages/scipy/integrate/_ode.py:1180: UserWarning: dopri5: larger nsteps is needed\n",
      "  self.messages.get(istate, unexpected_istate_msg)))\n",
      "/Users/dananjayaliyanage/git/parton_energy_loss_qcd_bjorken/solver_rk.py:78: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  res_quad2b=scipy.integrate.quad(lambda u, p=p: p*integrand_middle(p,u), 0, delta, limit=npts, epsabs=epsabs, epsrel=epsrel)\n",
      "/Users/dananjayaliyanage/git/parton_energy_loss_qcd_bjorken/solver_rk.py:78: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  res_quad2b=scipy.integrate.quad(lambda u, p=p: p*integrand_middle(p,u), 0, delta, limit=npts, epsabs=epsabs, epsrel=epsrel)\n",
      "/Users/dananjayaliyanage/git/parton_energy_loss_qcd_bjorken/solver_rk.py:78: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  res_quad2b=scipy.integrate.quad(lambda u, p=p: p*integrand_middle(p,u), 0, delta, limit=npts, epsabs=epsabs, epsrel=epsrel)\n",
      "/Users/dananjayaliyanage/git/parton_energy_loss_qcd_bjorken/solver_rk.py:78: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  res_quad2b=scipy.integrate.quad(lambda u, p=p: p*integrand_middle(p,u), 0, delta, limit=npts, epsabs=epsabs, epsrel=epsrel)\n",
      "/Users/dananjayaliyanage/git/parton_energy_loss_qcd_bjorken/solver_rk.py:78: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  res_quad2b=scipy.integrate.quad(lambda u, p=p: p*integrand_middle(p,u), 0, delta, limit=npts, epsabs=epsabs, epsrel=epsrel)\n",
      "/Users/dananjayaliyanage/git/parton_energy_loss_qcd_bjorken/solver_rk.py:73: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  res_quad2a=scipy.integrate.quad(vec_integrand, pmin, p, limit=npts, epsabs=epsabs, epsrel=epsrel)\n",
      "/Users/dananjayaliyanage/git/parton_energy_loss_qcd_bjorken/solver_rk.py:73: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  res_quad2a=scipy.integrate.quad(vec_integrand, pmin, p, limit=npts, epsabs=epsabs, epsrel=epsrel)\n",
      "/Users/dananjayaliyanage/git/parton_energy_loss_qcd_bjorken/solver_rk.py:73: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  res_quad2a=scipy.integrate.quad(vec_integrand, pmin, p, limit=npts, epsabs=epsabs, epsrel=epsrel)\n",
      "/Users/dananjayaliyanage/git/parton_energy_loss_qcd_bjorken/solver_rk.py:73: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  res_quad2a=scipy.integrate.quad(vec_integrand, pmin, p, limit=npts, epsabs=epsabs, epsrel=epsrel)\n",
      "/Users/dananjayaliyanage/git/parton_energy_loss_qcd_bjorken/solver_rk.py:73: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  res_quad2a=scipy.integrate.quad(vec_integrand, pmin, p, limit=npts, epsabs=epsabs, epsrel=epsrel)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working on 2/10 design\n",
      "Working on 2/10 design\n",
      "Working on 2/10 design\n",
      "Working on 2/10 design\n",
      "Working on 2/10 design\n",
      "Working on 3/10 design\n",
      "Working on 3/10 design\n",
      "Working on 3/10 design\n",
      "Working on 3/10 design\n",
      "Working on 3/10 design\n",
      "Working on 4/10 design\n",
      "Working on 4/10 design\n",
      "Working on 4/10 design\n",
      "Working on 4/10 design\n",
      "Working on 4/10 design\n",
      "Working on 5/10 design\n",
      "Working on 5/10 design\n",
      "Working on 5/10 design\n",
      "Working on 5/10 design\n",
      "Working on 5/10 design\n",
      "Working on 6/10 design\n",
      "Working on 6/10 design\n",
      "Working on 6/10 design\n",
      "Working on 6/10 design\n",
      "Working on 6/10 design\n",
      "Working on 7/10 design\n",
      "Working on 7/10 design\n",
      "Working on 7/10 design\n",
      "Working on 7/10 design\n",
      "Working on 7/10 design\n",
      "Working on 8/10 design\n",
      "Working on 8/10 design\n",
      "Working on 8/10 design\n",
      "Working on 8/10 design\n",
      "Working on 8/10 design\n",
      "Working on 9/10 design\n",
      "Working on 9/10 design\n",
      "Working on 9/10 design\n",
      "Working on 9/10 design\n"
     ]
    }
   ],
   "source": [
    "import multiprocessing\n",
    "from multiprocessing import Pool\n",
    "n_cpu = multiprocessing.cpu_count()\n",
    "print(f'Number of cores in the computer {n_cpu}')\n",
    "with Pool(n_cpu-1) as pool:\n",
    "    st = time.time()\n",
    "    obs_matrix=pool.map(sim_1D, prior_ranges.keys())\n",
    "    et = time.time()\n",
    "print(f'Total run time for the simulations {(et-st)/60:.2f} minutes')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "observations = np.array(obs_matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for param in prior_ranges.keys():\n",
    "    i=name_to_position[param]\n",
    "    design = make_design(param, n_design)\n",
    "    observables= observations[i,:,:]\n",
    "    fig, ax = plt.subplots()\n",
    "    for ii,pT in enumerate(pt_list):\n",
    "        if pT in show_pt:\n",
    "            ax.plot(design[:,i].flatten(),observables[:,ii], label=r'$p_T =$' +f'{pT:.2f} GeV')\n",
    "\n",
    "    ax.set_xlabel(f'{param}')\n",
    "    ax.set_ylabel(r'$R_{AA}$')\n",
    "    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.2),shadow=True, ncol=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}