{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fd526b28",
   "metadata": {},
   "source": [
    "# Making a model\n",
    "\n",
    "As illustrated in the architecture page, SATLAS2 is built on the concept of a Fitter object, to which Source objects are assigned. These Source objects contain both experimental data, but are also assigned one or more Models, which should be fitted to the data.\n",
    "\n",
    "The base aspect of these Models is that two things should be implemented: a *params* attribute which is a dictionary containing all the Parameters that the Model needs, and a method *f(x)*, where the response for the given *x* values is returned. It is recommended that the first line of the method is *x = self.transform(x)* in order to correctly handle any transformation functions that are required.\n",
    "\n",
    "In this tutorial, a new Model will be created to model a sine wave with exponential damping called *ExpSine*, just as an example. In the preamble, we just import all the libraries that we will need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "334dee65",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.insert(0, '..\\src')\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import satlas2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afcccf3c",
   "metadata": {},
   "source": [
    "The new class will be a subclass of the SATLAS2 Model class, and it will contain SATLAS2 Parameters in its *params* dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d74ea3ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ExpSine(satlas2.Model):\n",
    "    def __init__(self, A, lamda, omega, name='ExpSine', prefunc=None):\n",
    "        super().__init__(name, prefunc=prefunc)\n",
    "        self.params={\n",
    "            'amplitude': satlas2.Parameter(value=A,min=0,max=np.inf,vary=True),\n",
    "            'lambda': satlas2.Parameter(value=lamda, min=0, max=np.inf, vary=True),\n",
    "            'omega': satlas2.Parameter(value=omega, min=0, max=np.inf, vary=True)\n",
    "        }\n",
    "    \n",
    "    def f(self, x):\n",
    "        x = self.transform(x)\n",
    "        a = self.params['amplitude'].value\n",
    "        l = self.params['lambda'].value\n",
    "        o = self.params['omega'].value\n",
    "        return a*np.exp(-l*x)*np.sin(o*x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edd0df21",
   "metadata": {},
   "source": [
    "This defines a sine wave of angular frequency ω with an exponential decay with decay constant λ. Note that, in the parameter name in *__init__*, *lamda* is used instead of *lambda* since the last one is a keyword in Python.\n",
    "\n",
    "With this new Model, we can use this to both simulate what a dataset would look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a0725112",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'amplitude': 7+/-0 (inf max, 0 min, vary=True, correl={}),\n",
       " 'lambda': 1.5+/-0 (inf max, 0 min, vary=True, correl={}),\n",
       " 'omega': 4+/-0 (inf max, 0 min, vary=True, correl={})}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 27884 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "amplitude = 7\n",
    "lamda = 1.5\n",
    "omega = 4\n",
    "model = ExpSine(amplitude, lamda, omega, name='MyModel')\n",
    "\n",
    "x = np.linspace(0, 4, 100)\n",
    "y = model.f(x)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_axes([0.1, 0.1, 0.8, 0.8])\n",
    "ax.plot(x, y)\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "ax.set_title('My ExpSine Model')\n",
    "model.params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cb75319",
   "metadata": {},
   "source": [
    "And this model can easily be used with Source and Fitter to fit to data. First, we generate some data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2f10c717",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28532 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_x = np.linspace(0, 4, 20)\n",
    "data_y = model.f(data_x) + np.random.randn(data_x.shape[0])*0.5\n",
    "yerr = np.ones(data_y.shape)*0.5\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_axes([0.1, 0.1, 0.8, 0.8])\n",
    "ax.errorbar(data_x, data_y, yerr=yerr, fmt='o', label='Data')\n",
    "ax.plot(x, y, label='Initial guess')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "_ = ax.legend(loc=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d673f8a",
   "metadata": {},
   "source": [
    "We assign this data to a Source, add the ExpSine model to this Source, and pass it to a Fitter to fit this. Since this requires a normal chisquare fit, no extra arguments are required for the fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5a93d8e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[Fit Statistics]]\n",
      "    # fitting method   = leastsq\n",
      "    # function evals   = 53\n",
      "    # data points      = 20\n",
      "    # variables        = 3\n",
      "    chi-square         = 28.7327890\n",
      "    reduced chi-square = 1.69016406\n",
      "    Akaike info crit   = 13.2461335\n",
      "    Bayesian info crit = 16.2333303\n",
      "[[Variables]]\n",
      "    Datafile1___MyModel___amplitude:  5.39734439 +/- 1.20120658 (22.26%) (init = 7)\n",
      "    Datafile1___MyModel___lambda:     1.36883125 +/- 0.40630320 (29.68%) (init = 1.5)\n",
      "    Datafile1___MyModel___omega:      4.44140916 +/- 0.25391170 (5.72%) (init = 4)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41024 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "datasource = satlas2.Source(data_x, data_y, yerr=yerr, name='Datafile1')\n",
    "datasource.addModel(model)\n",
    "f = satlas2.Fitter()\n",
    "f.addSource(datasource)\n",
    "f.fit()\n",
    "print(f.reportFit())\n",
    "fig = plt.figure()\n",
    "ax = fig.add_axes([0.1, 0.1, 0.8, 0.8])\n",
    "ax.errorbar(data_x, data_y, yerr=yerr, fmt='o', label='Data')\n",
    "ax.plot(x, y, label='Initial guess')\n",
    "ax.plot(x, model.f(x), label='Fit')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "_ = ax.legend(loc=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "660ea3dd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
