{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "38ce6ee6-9e86-4108-96cb-dbfb63e297cf",
   "metadata": {},
   "source": [
    "## Creating a new workflow for oscillating data\n",
    "\n",
    "In this tutorial we will consider a toy model to demonstrate how to extend the `quickBayes` package to solve new problems. At present all of the models must have the form:\n",
    "\n",
    "$$\n",
    "y(x) = g(x) + \\sum_{j=1}{N} f_j(x),\n",
    "$$\n",
    "\n",
    "where $g(x)$ is a function that is present in all of the models (e.g. background), $f_j$ is the $j^\\mathrm{th}$ instance of the repeated function (e.g. sin) and $N$ is the maximum number of repeated functions. \n",
    "\n",
    "The first step is to read the test data and inspect it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "228acbb8-4d33-4e0a-964e-4039a3f9459e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "x, y, e = np.loadtxt('tutorial3.npy')\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set(xlabel='Time ($\\mu s$)', ylabel='', title='sin data');\n",
    "ax.errorbar(x, y, e, fmt='kx', label='data');\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a8320fd-2389-420a-8be0-f2dcc218092f",
   "metadata": {},
   "source": [
    "The `quickBayes` package does not currently have a `sin` fitting function as standard. Hence, we will create a new fitting function. Every fitting function must inherit the `BaseFitFunction` as this defines some useful methods that allow this class to be much simpler. \n",
    "\n",
    "The `__init__` has a `super` that take the number of arguments (3), the prefix, the starting guess as a list, the lower bounds as a list and the upper bounds as a list. \n",
    "\n",
    "Next are the properties, these are used to create the text in the reporting. Each fitting parameter will need its own property. The inclusion of `self._prefix` allows for the full string to have some context, such as how many repeated functions are present. \n",
    "\n",
    "The `__call__` method is a simple evaluation of the fit function.\n",
    "\n",
    "The `read_from_report` loads the results in from a report dictionary. This is done by repeating the `_read_report` for each parameter. The order of these must match those used in the `report` method. the `report` method adds the name and value for the fit parameter to the results dictionary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f09db3b-17e8-4a32-80be-6a31a9e555b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from quickBayes.functions.base import BaseFitFunction\n",
    "from numpy import ndarray\n",
    "import numpy as np\n",
    "\n",
    "class Sin(BaseFitFunction):\n",
    "    def __init__(self, prefix=''):\n",
    "        \"\"\"\n",
    "        :param prefix: prefix for function parameters in report\n",
    "        \"\"\"\n",
    "        super().__init__(3, prefix, [1.75, 2.25, 0.22], [0., 0, 0], [5., 5., 2.])\n",
    "\n",
    "    @property\n",
    "    def amp(self):\n",
    "        return str(f'{self._prefix} Amplitude')\n",
    "    \n",
    "    @property\n",
    "    def freq(self):\n",
    "        return str(f'{self._prefix} Frequency')\n",
    "        \n",
    "    @property\n",
    "    def phase(self):\n",
    "        return str(f'{self._prefix} phase')\n",
    "        \n",
    "    def __call__(self, x: ndarray, amp, omega, phi):\n",
    "        return amp*np.sin(omega*x + phi)\n",
    "\n",
    "    def read_from_report(self, report_dict,\n",
    "                         index=0):\n",
    "        \"\"\"\n",
    "        Read the parameters from the results dict\n",
    "        :param report_dict: the dict of results\n",
    "        :param index: the index to get results from\n",
    "        :return the parameters\n",
    "        \"\"\"\n",
    "        return [self._read_report(report_dict, self.amp, index),\n",
    "                self._read_report(report_dict, self.freq, index),\n",
    "                self._read_report(report_dict, self.phase, index)]\n",
    "\n",
    "    def report(self, report_dict, amp, omega, phi):\n",
    "        \"\"\"\n",
    "        reporting method\n",
    "        :param report_dict: dict of parameters\n",
    "        :param c: constant\n",
    "        :return dict of parameters, including BG\n",
    "        \"\"\"\n",
    "        report_dict = self._add_to_report(self.amp,\n",
    "                                          amp, report_dict)\n",
    "        report_dict = self._add_to_report(self.freq,\n",
    "                                          omega, report_dict)\n",
    "        report_dict = self._add_to_report(self.phase,\n",
    "                                          phi, report_dict)\n",
    "        return report_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7ed7d47-6a1f-4f07-aed7-db577d119d91",
   "metadata": {},
   "source": [
    "The second thing that is needed to analyse the above data is a workflow. This will inherit the `ModelSelectionWorkflow`, which contains the common aspects of a model selection calculation. The `preprocess_data` method prepares the data for analysis, in this example it will just be a simple crop. The `super` sets the `x, y, e` values of the data we want to investigate.\n",
    "\n",
    "The `_update_function` method defines what the workflow should do to add a repeated fitting function. Hence, it creates a `Sin` object and then adds it to the input function, which is then returned. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d10f6aa-e68f-4efa-af3f-bedb36ed12c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from quickBayes.utils.crop_data import crop\n",
    "from quickBayes.workflow.model_selection.template import ModelSelectionWorkflow\n",
    "\n",
    "class SinSelection(ModelSelectionWorkflow):\n",
    "\n",
    "    def preprocess_data(self, x_data,\n",
    "                        y_data, e_data,\n",
    "                        start_x, end_x):\n",
    "        \"\"\"\n",
    "        The preprocessing needed for the data.\n",
    "        This crops and stores the data.\n",
    "        :param x_data: the x data to fit to\n",
    "        :param y_data: the y data to fit to\n",
    "        :param e_data: the errors for the y data\n",
    "        :param start_x: the start x value\n",
    "        :param end_x: the end x value\n",
    "        \"\"\"\n",
    "        sx, sy, se = crop(x_data, y_data, e_data,\n",
    "                          start_x, end_x)\n",
    "        super().preprocess_data(sx, sy, se)\n",
    "\n",
    "    @staticmethod\n",
    "    def _update_function(in_func: BaseFitFunction) -> BaseFitFunction:\n",
    "        \"\"\"\n",
    "        This method adds a exponential decay to the fitting\n",
    "        function.\n",
    "        :param func: the fitting function that needs modifying\n",
    "        :return the modified fitting function\n",
    "        \"\"\"\n",
    "\n",
    "        function = Sin()\n",
    "        in_func.add_function(function)\n",
    "        return in_func"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "563b0534-d688-4af9-a0fe-640efa3e45bf",
   "metadata": {},
   "source": [
    "As with previous examples, we start by defining the parameters for the problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75bd297c-3efd-49c2-a13c-811b657bacb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = {}\n",
    "results_errors = {}\n",
    "start_x = 0.5\n",
    "end_x = 6.3\n",
    "max_features = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5ecd018-8595-4d02-8cd0-13fe9cad4349",
   "metadata": {},
   "source": [
    "The next step is to setup the workflow, this will look almost identical to the other examples. This makes it easy to swap and change between different workflows. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "680232f2-1721-4aeb-b534-733b5d77af1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = SinSelection(results, results_errors)\n",
    "workflow.preprocess_data(x, y, e, start_x, end_x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80ba98c2-1bc9-4f2a-8da6-50d81250498b",
   "metadata": {},
   "source": [
    "In this case the unique functions are equal to zero. Therefore, we can just use an empty `CompositeFunction`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "415e9381-4544-4a85-8564-37694f4fe4a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from quickBayes.functions.composite import CompositeFunction\n",
    "\n",
    "func = CompositeFunction()\n",
    "workflow.set_scipy_engine(func.get_guess(), *func.get_bounds())\n",
    "func = workflow.execute(max_features, func, func.get_guess())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db4bc6c8-faa7-403f-b41e-63c7bc206199",
   "metadata": {},
   "source": [
    "Finally we get the loglikelihoods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df0f0c52-84b6-41db-ac5b-f817fbffd424",
   "metadata": {},
   "outputs": [],
   "source": [
    "results, results_errors = workflow.get_parameters_and_errors\n",
    "\n",
    "for key in results.keys():\n",
    "    if 'log' in key:\n",
    "        print(key, results[key][0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22f354c2-9dd9-4556-b8ee-ebf1b2e7bf05",
   "metadata": {},
   "source": [
    "To get the Bayesian P-value we first calculate the normalisation, by adding up the probabilities (note that the $\\log$ are base 10). Then we can report the P-values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27327f2c-7935-4db9-a09a-f5dde4c4b0fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 0\n",
    "for key in results.keys():\n",
    "    if 'log' in key:\n",
    "        N += 10**results[key][0]\n",
    "\n",
    "for key in results.keys():\n",
    "    if 'log' in key:\n",
    "        print(f'P-value for {key.split(\":\")}', (10**results[key][0])/N)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac5510ad-20af-4754-8970-a4c3b87e2f15",
   "metadata": {},
   "source": [
    "It is clear that two sin waves are more likely for this example. However, one sin wave has a reasonable probability. We will plot the data along with the results for one and two sin waves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "820cc0c7-883c-4fe1-8f1e-ad601fb4d976",
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_engine = workflow.fit_engine\n",
    "x1, y1, e1, _, _ = fit_engine.get_fit_values(0)\n",
    "\n",
    "fit_engine = workflow.fit_engine\n",
    "x2, y2, e2, _, _ = fit_engine.get_fit_values(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef7368f-a684-451f-9116-e2dd183f55fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.errorbar(x, y, e, fmt='ok', label='data')\n",
    "ax.set(xlabel='time ($\\mu s)$', ylabel='Asymmetry', title='input data');\n",
    "ax.errorbar(x1, y1, e1, fmt='b--', label='1 sin wave')\n",
    "ax.errorbar(x2, y2, e2, fmt='r--', label='2 sin wave')\n",
    "ax.set(xlabel='Time ($\\mu s$)', ylabel='', title='sin data');\n",
    "\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3a6c427-14e7-4494-89d7-2696bb599eea",
   "metadata": {},
   "source": [
    "The plot shows that it is not unreasonable to have one sin wave, it does under and over estimate the values at the two extremes of the data. However, the two sin waves always has a good quality fit as expected. \n",
    "\n",
    "The simulated data did use two sin waves. Hence, `quickBayes` was able to identify the correct model to use for this data. "
   ]
  }
 ],
 "metadata": {
  "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}