{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3: Perform fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{warning}\n",
    "\n",
    "The {doc}`pycompwa </index>` is no longer maintained. Use the [ComPWA](https://compwa-org.rtfd.io) packages [QRules](https://qrules.rtfd.io), [AmpForm](https://ampform.rtfd.io), and [TensorWaves](https://tensorwaves.rtfd.io) instead!\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As explained in the {doc}`previous step <step2>`, an {class}`~.Intensity` object behaves just like a mathematical function that takes a {class}`~.DataSet` as an argument and returns a list of intensities (real numbers). We now want to optimize the parameters of the intensity model so that it matches  the distribution of our data sample. This is what we call 'performing a fit'.\n",
    "\n",
    ":::{note}\n",
    "\n",
    "We first load the relevant data from the previous steps.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycompwa.ui as pwa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particle_list = pwa.read_particles(\"model.xml\")\n",
    "kinematics = pwa.create_helicity_kinematics(\"model.xml\", particle_list)\n",
    "kinematics.create_all_subsystems()\n",
    "data_sample = pwa.read_root_data(input_file=\"generated_data.root\")\n",
    "phsp_sample = pwa.read_root_data(input_file=\"generated_phsp.root\")\n",
    "intensity_builder = pwa.IntensityBuilderXML(\n",
    "    \"model.xml\", particle_list, kinematics, phsp_sample\n",
    ")\n",
    "intensity = intensity_builder.create_intensity()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.1 Define estimator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To perform a fit, you need to define an **estimator**. An estimator is a measure for the discrepancy between the intensity model and the data distribution to which you fit. To create an estimator we therefore generally need:\n",
    "\n",
    "- an {class}`~.Intensity` instance\n",
    "- a {class}`~.DataSet` to which we fit that intensity\n",
    "\n",
    "The data sample of momentum tuples can be converted to a {class}`~.DataSet` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_set = kinematics.convert(data_sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create an estimator. In this example, we use an **unbinned log likelihood** estimator. This estimator is used most commonly in high-energy physics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    estimator,\n",
    "    initial_parameters,\n",
    ") = pwa.create_unbinned_log_likelihood_function_tree_estimator(\n",
    "    intensity, data_set\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that you not only receive a estimator object, but also a list of fit parameters. You use this list of fit parameters to initialize the optimizer later on.\n",
    "\n",
    "The parameters contains the following info:\n",
    "\n",
    "- initial value\n",
    "- whether they are free or fixed\n",
    "- optional upper and lower limits\n",
    "- errors that give certain optimizers hints on the step size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2 Modify parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fit parameters in this list have been initialized with the values stated in the XML file or with default values if unspecified, but they can easily be modified in Python as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\nthis parameter is initially not fixed:\")\n",
    "print(initial_parameters[8])\n",
    "\n",
    "initial_parameters[7].is_fixed = True\n",
    "initial_parameters[8].is_fixed = True\n",
    "print(\"\\nand now it is fixed:\")\n",
    "print(initial_parameters[8])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make the fit process a bit more interesting, we modify one of the parameters to a different initial value than the true value:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_value = initial_parameters[12].value\n",
    "initial_parameters[12]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_parameters[12].value = 2.0\n",
    "initial_parameters[12]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.3 Optimize the intensity model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's time to perform the fit! This is quite simple: just create an optimizer instance of your choice, here [Minuit2](https://root.cern.ch/doc/master/Minuit2Page.html), and call its {meth}`~.optimize` method to start the fitting process.\n",
    "\n",
    "The computation time depends on the complexity of the model, number of data events, the size of the phase space sample, and the number of free parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"data size\", len(data_sample.events))\n",
    "print(\"phsp size\", len(phsp_sample.events))\n",
    "print(\n",
    "    \"free parameters:\",\n",
    "    len([par for par in initial_parameters if not par.is_fixed]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "With 14 free parameters, this fit takes about 15 seconds on an Intel(R) Core(TM) i7-6820HQ 2.70GHz CPU.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = pwa.MinuitIF()\n",
    "fit_result = optimizer.optimize(estimator, initial_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_result.fit_duration_in_seconds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check whether the fit parameters are close to the true values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "value = fit_result.final_parameters[12].value\n",
    "error = fit_result.final_parameters[12].error\n",
    "if abs(value - original_value) < 3 * error[0]:\n",
    "    print(\"This value is close to\", original_value, \"again:\")\n",
    "    print(value, \"+-\", error[0])  # upper and lower error are equal in hesse\n",
    "else:\n",
    "    raise Exception(\"Value does not lie within 3 sigma!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{warning}\n",
    "\n",
    "The intensity instance still needs to be notified about this optimal set of parameters. They can be applied with {meth}`~.updateParametersFrom`.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intensity.updateParametersFrom(fit_result.final_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.4 Exporting and importing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optimizing an intensity model can take a long time, so it is important that you store the fit result in the end. We again write to an XML file (but note that this has a different format than that of the amplitude model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_result.write(\"fit_result.xml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imported_fit_result = pwa.load(\"fit_result.xml\")\n",
    "\n",
    "exported_value = fit_result.final_parameters[12].value\n",
    "imported_value = imported_fit_result.final_parameters[12].value\n",
    "\n",
    "if exported_value != imported_value:\n",
    "    raise Exception\n",
    "else:\n",
    "    print(imported_fit_result.final_parameters[12])"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Raw Cell Format",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
