{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting for variable sources\n",
    "\n",
    "_Weights used:_\n",
    "* Spatial\n",
    "* Temporal\n",
    "* Spectral\n",
    "\n",
    "This simulated example source is pulsed and has a sweeping PA. This notebook determines the PD and EVPA of the sweep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[93m>>> PyXSPEC is not installed, you will no be able to use it.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import leakagelib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll load the data as usual, and cut to within 280 arcsec of the center"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      ">>> Reading (in memory) /opt/homebrew/anaconda3/lib/python3.12/site-packages/ixpeobssim/caldb/ixpe/xrt/bcf/vign/ixpe_d1_obssim20240101_vign_v013.fits...\n"
     ]
    }
   ],
   "source": [
    "datas = leakagelib.IXPEData.load_all_detectors_with_path(\"data\", \"pulse\")\n",
    "for data in datas:\n",
    "    data.iterative_centroid_center()\n",
    "    data.retain(data.evt_energies > 2)\n",
    "    data.retain(data.evt_energies < 8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We plan to implement phase weights, but LeakageLib records times, not phases. You should calculate them yourself (e.g. with PINT, or load phases from a separate file), then replace the `evt_times` field with these phases. For this mock data, the frequency is exactly 10 Hz, so the phases are just $10 t$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WARNING: only run this box once.\n",
    "\n",
    "for data in datas:\n",
    "    # It's convenient to overwrite the \"times\" list with a list of phases. This source sweeps with frequency of 10 Hz, so multiplying by 10 gives the phase\n",
    "    data.evt_times *= 10\n",
    "    data.evt_times = np.fmod(data.evt_times, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create the point source and background sources as usual, setting their spectra. For this fit, we'll fix the background to be unpolarized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "6462 events were cut for being outside the region of interest.\n",
      ">>> Reading (in memory) /opt/homebrew/anaconda3/lib/python3.12/site-packages/ixpeobssim/caldb/ixpe/gpd/cpf/arf/ixpe_d1_obssim20240101_v013.arf...\n",
      ">>> Using cached xEffectiveArea object at /opt/homebrew/anaconda3/lib/python3.12/site-packages/ixpeobssim/caldb/ixpe/gpd/cpf/arf/ixpe_d1_obssim20240101_v013.arf...\n"
     ]
    }
   ],
   "source": [
    "settings = leakagelib.FitSettings(datas)\n",
    "settings.apply_circular_roi(280)\n",
    "\n",
    "settings.add_point_source(\"src\")\n",
    "settings.fix_flux(\"src\", 1)\n",
    "\n",
    "settings.add_background(\"bkg\")\n",
    "settings.fix_qu(\"bkg\", (0, 0))\n",
    "settings.set_initial_flux(\"bkg\", 1)\n",
    "\n",
    "settings.set_spectrum(\"bkg\", lambda e: e**-2.5)\n",
    "settings.set_spectrum(\"src\", lambda e: e**-1.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To implement phase weights, we need to tell `LeakageLib` the source's light curve. This source was simulated with a sine squared light curve. Note that normalization doesn't matter.\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**Warning** \n",
    "\n",
    "If you cut based on phase e.g. for an on-off fit, you will need to supply the `duty_cycle` argument in `set_lightcurve`.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "settings.set_lightcurve(\"src\", lambda ph: np.sin(ph * np.pi)**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can perform the fit and see what we get."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Data set pulse DU 1 had no exposure map loaded. Please load an exposure map if you are fitting to events in the vignetted portion.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FITTED PARAMETERS:\n",
      "Source\tParam\n",
      "src:\tq\n",
      "src:\tu\n",
      "bkg:\tf\n",
      "\n",
      "FIXED PARAMETERS:\n",
      "Source\tParam\tValue\n",
      "src:\tf\t1\n",
      "bkg:\tq\t0\n",
      "bkg:\tu\t0\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "FitResult:\n",
       "\tq (src) = -0.0033 +/- 0.0721\n",
       "\tu (src) = -0.0027 +/- 0.0721\n",
       "\tf (bkg) = 2.1789 +/- 0.0417\n",
       "\n",
       "Polarization:\n",
       "\tPD (src): 0.0043 +/- 0.0721\n",
       "\tPA (src): -70.5491 deg +/- 481.0701\n",
       "Likelihood 19526.389495555064, dof 15927\n",
       "Optimization terminated successfully."
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fitter = leakagelib.Fitter(settings)\n",
    "print(fitter)\n",
    "const_result = fitter.fit()\n",
    "const_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The best-fit polarization degree was quite low! This is because the source's true PA sweeps, and we modeled it as constant. We need to make a sweeping PA model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting a custom polarization model\n",
    "\n",
    "We'll test a model where the EVPA rotates by 2 pi every pulse, and PD is constant. This model will have two parameters: a constant PD and a EVPA at phase zero. Let's create those parameters. The `add_param` function allows us to do this, and give initial fit values and bounds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "settings.add_param(\"sweep-PD\", 0.1, [0, 1])\n",
    "settings.add_param(\"sweep-PA\", 0, [-np.pi, np.pi])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to create a function which takes in the phase and gives the expected Q and U polarization, and tell the fitter to use this function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_fn(ph, fit_data, param_array):\n",
    "    pd = fit_data.param_to_value(param_array, \"sweep-PD\")\n",
    "    pa = fit_data.param_to_value(param_array, \"sweep-PA\")\n",
    "    q = pd * np.cos(2 * (ph * 2*np.pi + pa))\n",
    "    u = pd * np.sin(2 * (ph * 2*np.pi + pa))\n",
    "    return q, u\n",
    "settings.set_model_fn(\"src\", model_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `set_model_fn` automatically tells the fitter to no longer use the src q and src u params. You can see that by printing the fitter object, which displays the free parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Data set pulse DU 1 had no exposure map loaded. Please load an exposure map if you are fitting to events in the vignetted portion.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FITTED PARAMETERS:\n",
      "Source\tParam\n",
      "bkg:\tf\n",
      "None:\tsweep-PD\n",
      "None:\tsweep-PA\n",
      "\n",
      "FIXED PARAMETERS:\n",
      "Source\tParam\tValue\n",
      "src:\tq\t0\n",
      "src:\tu\t0\n",
      "src:\tf\t1\n",
      "bkg:\tq\t0\n",
      "bkg:\tu\t0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "fitter = leakagelib.Fitter(settings)\n",
    "print(fitter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running the fit will only fit for f, sweep-PD, and sweep-PA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FitResult:\n",
       "\tf (bkg) = 2.1787 +/- 0.0417\n",
       "\tsweep-PD = 0.4722 +/- 0.0711\n",
       "\tsweep-PA = 0.0244 +/- 0.0763\n",
       "\n",
       "Polarization:\n",
       "Likelihood 19548.00863653022, dof 15927\n",
       "Optimization terminated successfully."
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sweep_result = fitter.fit()\n",
    "sweep_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The PD of the sweeping fit was higher. The likelihood was also higher, indicating the sweep model provided a better fit. The true source polarization had PD = 0.5 and PA = 0, which the fit agrees with."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}