{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2d86ef29",
   "metadata": {},
   "source": [
    "As a stepping stone between SATLAS and SATLAS2, an interface has been provided which can mostly be used as a drop-in replacement for code that uses the SATLAS syntax. Note that not all functionalities have been implemented in this fashion. For users that require these functionalities, we recommend migrating to SATLAS2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "31960148",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import time\n",
    "\n",
    "import matplotlib.gridspec as gridspec\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "sys.path.insert(0, '..\\src')\n",
    "\n",
    "import satlas2\n",
    "import satlas as sat\n",
    "\n",
    "\n",
    "def modifiedSqrt(input):\n",
    "    output = np.sqrt(input)\n",
    "    output[input <= 0] = 1e-3\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b35c0c8",
   "metadata": {},
   "source": [
    "# Fitting a single hyperfine spectrum\n",
    "\n",
    "The most common task, and the one this interface is meant for, is fitting a single hyperfine spectrum. A special class in SATLAS2 called *HFSModel* has been created as a replacement for the equivalent SATLAS *HFSModel*. Note that the normal hyperfine spectrum model in SATLAS2 is called *HFS*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "917d2ce6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\megat\\Documents\\git_repositories\\personal\\satlas2\\examples\\..\\src\\satlas2\\models\\hfsModel.py:284: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  shift = phase * n / d\n"
     ]
    }
   ],
   "source": [
    "spin = 3.5\n",
    "J = [0.5, 1.5]\n",
    "A = [9600, 175]\n",
    "B = [0, 315]\n",
    "C = [0, 0]\n",
    "FWHMG = 135\n",
    "FWHML = 101\n",
    "centroid = 480\n",
    "bkg = [100]\n",
    "scale = 90\n",
    "\n",
    "x = np.arange(-17500, -14500, 40)\n",
    "x = np.hstack([x, np.arange(20000, 23000, 40)])\n",
    "\n",
    "rng = np.random.default_rng(0)\n",
    "hfs = satlas2.HFSModel(I=spin,\n",
    "                       J=J,\n",
    "                       ABC=[A[0], A[1], B[0], B[1], C[0], C[1]],\n",
    "                       centroid=centroid,\n",
    "                       fwhm=[FWHMG, FWHML],\n",
    "                       scale=scale,\n",
    "                       background_params=bkg,\n",
    "                       use_racah=True)\n",
    "hfs.set_variation({'Cu': False})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e566f68",
   "metadata": {},
   "source": [
    "The object called *hfs* can be used with the syntax of SATLAS. Generating Poisson-distributed data is done by simply calling the function with frequency values as an argument, and using the result for the NumPy Poisson random number generator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a3ab1656",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = hfs(x)\n",
    "y = rng.poisson(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43246911",
   "metadata": {},
   "source": [
    "In order to demonstrate the difference in performance, the centroid is offset by 100 from the actual value and the fitting is done by both the interface and SATLAS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "87011ba1",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\megat\\.conda\\envs\\satlas2\\Lib\\site-packages\\satlas\\models\\hfsmodel.py:573: RuntimeWarning: invalid value encountered in divide\n",
      "  D = (3*C*(C+1) - 4*I*(I+1)*J*(J+1)) / (2*I*(2*I-1)*J*(2*J-1))\n",
      "C:\\Users\\megat\\.conda\\envs\\satlas2\\Lib\\site-packages\\satlas\\models\\hfsmodel.py:574: RuntimeWarning: invalid value encountered in divide\n",
      "  E = (10*(0.5*C)**3 + 20*(0.5*C)**2 + C*(-3*I*(I+1)*J*(J+1) + I*(I+1) + J*(J+1) + 3) - 5*I*(I+1)*J*(J+1)) / (I*(I-1)*(2*I-1)*J*(J-1)*(2*J-1))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fitting 1 dataset with chisquare (Pearson, satlas2)...\n",
      "define whether you want to see the correlations in display_chisquare_fit(...)\n",
      "[[Fit Statistics]]\n",
      "    # fitting method   = leastsq\n",
      "    # function evals   = 137\n",
      "    # data points      = 150\n",
      "    # variables        = 8\n",
      "    chi-square         = 151.188938\n",
      "    reduced chi-square = 1.06471083\n",
      "    Akaike info crit   = 17.1842512\n",
      "    Bayesian info crit = 41.2693335\n",
      "[[Variables]]\n",
      "    Fit___HFModel__3_5___centroid:  482.548153 +/- 7.56664273 (1.57%) (init = 380)\n",
      "    Fit___HFModel__3_5___Al:        9604.53248 +/- 6.41301473 (0.07%) (init = 9600)\n",
      "    Fit___HFModel__3_5___Au:        176.460908 +/- 2.73509313 (1.55%) (init = 175)\n",
      "    Fit___HFModel__3_5___Bl:        0 (fixed)\n",
      "    Fit___HFModel__3_5___Bu:        348.564601 +/- 19.6945247 (5.65%) (init = 315)\n",
      "    Fit___HFModel__3_5___Cl:        0 (fixed)\n",
      "    Fit___HFModel__3_5___Cu:        0 (fixed)\n",
      "    Fit___HFModel__3_5___FWHMG:     142.382561 +/- 57.6647446 (40.50%) (init = 135)\n",
      "    Fit___HFModel__3_5___FWHML:     100.522928 +/- 63.5247534 (63.19%) (init = 101)\n",
      "    Fit___HFModel__3_5___scale:     89.2398294 +/- 7.15348131 (8.02%) (init = 90)\n",
      "    Fit___HFModel__3_5___Amp3to2:   0.4545455 (fixed)\n",
      "    Fit___HFModel__3_5___Amp3to3:   0.4772727 (fixed)\n",
      "    Fit___HFModel__3_5___Amp3to4:   0.3409091 (fixed)\n",
      "    Fit___HFModel__3_5___Amp4to3:   0.1590909 (fixed)\n",
      "    Fit___HFModel__3_5___Amp4to4:   0.4772727 (fixed)\n",
      "    Fit___HFModel__3_5___Amp4to5:   1 (fixed)\n",
      "    Fit___bkg___p0:                 100.670728 +/- 1.59295185 (1.58%) (init = 100)\n",
      "Fitting 1 dataset with chisquare (Pearson, satlas)...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Chisquare fitting done: 172it [00:00, 200.00it/s]                            \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NDoF: 142, Chisquare: 151.18894, Reduced Chisquare: 1.0647108\n",
      "Akaike Information Criterium: 17.18425, Bayesian Information Criterium: 41.269333\n",
      "Errors scaled with reduced chisquare.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\megat\\.conda\\envs\\satlas2\\Lib\\site-packages\\lmfit\\printfuncs.py:177: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  spercent = f'({abs(par.stderr/par.value):.2%})'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[Variables]]\n",
      "    FWHMG:        142.398641 +/- 57.6603206 (40.49%) (init = 142.3867)\n",
      "    FWHML:        100.507637 +/- 63.5294141 (63.21%) (init = 100.519)\n",
      "    TotalFWHM:    203.616071 +/- 21.3016980 (10.46%) == '0.5346*FWHML+(0.2166*FWHML**2+FWHMG**2)**0.5'\n",
      "    Scale:        89.2388854 +/- 7.15309431 (8.02%) (init = 89.23958)\n",
      "    Saturation:   0 (fixed)\n",
      "    Amp3__2:      0.4546399 (fixed)\n",
      "    Amp3__3:      0.4773649 (fixed)\n",
      "    Amp3__4:      0.3410048 (fixed)\n",
      "    Amp4__3:      0.1591578 (fixed)\n",
      "    Amp4__4:      0.4773975 (fixed)\n",
      "    Amp4__5:      1 (fixed)\n",
      "    Al:           9604.53225 +/- 6.41310259 (0.07%) (init = 9604.532)\n",
      "    Au:           176.461706 +/- 2.73513458 (1.55%) (init = 176.4611)\n",
      "    Bl:           0 (fixed)\n",
      "    Bu:           348.556407 +/- 19.6948460 (5.65%) (init = 348.5624)\n",
      "    Cl:           0 (fixed)\n",
      "    Cu:           0 (fixed)\n",
      "    Centroid:     482.545220 +/- 7.56678472 (1.57%) (init = 482.5474)\n",
      "    Background0:  100.670920 +/- 1.59296491 (1.58%) (init = 100.6708)\n",
      "    N:            0 (fixed)\n",
      "SATLAS2: 0.044 s\n",
      "SATLAS1: 0.875 s\n"
     ]
    }
   ],
   "source": [
    "hfs.params['centroid'].value = centroid - 100\n",
    "# Normal SATLAS implementation\n",
    "hfs1 = sat.HFSModel(spin,\n",
    "                    J, [A[0], A[1], B[0], B[1], C[0], C[1]],\n",
    "                    centroid - 100, [FWHMG, FWHML],\n",
    "                    scale=scale,\n",
    "                    background_params=bkg,\n",
    "                    use_racah=True)\n",
    "hfs1.set_variation({'Cu': False})\n",
    "\n",
    "# Interface fitting\n",
    "print('Fitting 1 dataset with chisquare (Pearson, satlas2)...')\n",
    "start = time.time()\n",
    "satlas2.chisquare_fit(hfs, x, y, modifiedSqrt(y))\n",
    "hfs.display_chisquare_fit(show_correl=False)\n",
    "stop = time.time()\n",
    "dt1 = stop - start\n",
    "\n",
    "# SATLAS fitting\n",
    "print('Fitting 1 dataset with chisquare (Pearson, satlas)...')\n",
    "start = time.time()\n",
    "sat.chisquare_fit(hfs1, x, y, modifiedSqrt(y))\n",
    "hfs1.display_chisquare_fit(show_correl=False, scaled=True)\n",
    "stop = time.time()\n",
    "dt2 = stop - start\n",
    "print('SATLAS2: {:.3} s'.format(dt1))\n",
    "print('SATLAS1: {:.3} s'.format(dt2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90c6cd41",
   "metadata": {},
   "source": [
    "Note that the results are functionally identical: the slight difference is due to a more modern implementation of the least squares fitting routine that is used under the hood by SATLAS2. The speedup by using SATLAS 2 is about a factor 20 for a single spectrum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9f866969",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 119512 characters]\n",
      "text/plain": [
       "<Figure size 1400x900 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_x = x[x<0]\n",
    "right_x = x[x>0]\n",
    "\n",
    "left_y = y[x<0]\n",
    "right_y = y[x>0]\n",
    "\n",
    "fig = plt.figure(constrained_layout=True, figsize=(14, 9))\n",
    "gs = gridspec.GridSpec(nrows=2, ncols=2, figure=fig)\n",
    "ax11 = fig.add_subplot(gs[0, 0])\n",
    "ax11.label_outer()\n",
    "ax12 = fig.add_subplot(gs[0, 1], sharey=ax11)\n",
    "ax12.label_outer()\n",
    "ax21 = fig.add_subplot(gs[1, 0], sharex=ax11)\n",
    "ax21.label_outer()\n",
    "ax22 = fig.add_subplot(gs[1, 1], sharex=ax12, sharey=ax21)\n",
    "ax22.label_outer()\n",
    "\n",
    "ax11.errorbar(left_x, left_y, modifiedSqrt(left_y), fmt='.', label='Artificial data')\n",
    "ax11.plot(left_x, hfs(left_x), '-', label='Fit')\n",
    "ax12.errorbar(right_x, right_y, modifiedSqrt(right_y), fmt='.', label='Artificial data')\n",
    "ax12.plot(right_x, hfs(right_x), '-', label='Fit')\n",
    "\n",
    "ax21.errorbar(left_x, left_y, modifiedSqrt(left_y), fmt='.', label='Artificial data')\n",
    "ax21.plot(left_x, hfs1(left_x), '-', label='SATLAS fit')\n",
    "ax22.errorbar(right_x, right_y, modifiedSqrt(right_y), fmt='.', label='Artificial data')\n",
    "ax22.plot(right_x, hfs1(right_x), '-', label='SATLAS fit')\n",
    "\n",
    "ax11.legend()\n",
    "ax21.legend()\n",
    "\n",
    "ax11.set_ylabel('SATLAS2')\n",
    "ax21.set_ylabel('SATLAS')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab6c15ba",
   "metadata": {},
   "source": [
    "# Overlapping hyperfine spectra\n",
    "\n",
    "The other most common usecase for SATLAS was analysis of spectra with an isomer present, resulting in overlapping spectra. In the SATLAS terminology, this would result in a *SumModel* being used. In SATLAS2, a second *HFS* is simply added to the Source. However, the interface does provide the folllowing functionality:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3c8673d8",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fitting 1 dataset with chisquare (Pearson, satlas2)...\n",
      "[[Fit Statistics]]\n",
      "    # fitting method   = leastsq\n",
      "    # function evals   = 291\n",
      "    # data points      = 175\n",
      "    # variables        = 16\n",
      "    chi-square         = 166.792273\n",
      "    reduced chi-square = 1.04900801\n",
      "    Akaike info crit   = 23.5935582\n",
      "    Bayesian info crit = 74.2301338\n",
      "[[Variables]]\n",
      "    Fit___HFModel__4___centroid:       421.201657 +/- 6.27607829 (1.49%) (init = 300)\n",
      "    Fit___HFModel__4___Al:             5313.73380 +/- 4.60447109 (0.09%) (init = 5300)\n",
      "    Fit___HFModel__4___Au:             105.171460 +/- 1.98883087 (1.89%) (init = 100)\n",
      "    Fit___HFModel__4___Bl:             0 (fixed)\n",
      "    Fit___HFModel__4___Bu:             217.091876 +/- 17.0714679 (7.86%) (init = 230)\n",
      "    Fit___HFModel__4___Cl:             0 (fixed)\n",
      "    Fit___HFModel__4___Cu:             0 (fixed)\n",
      "    Fit___HFModel__4___FWHMG:          29.2819035 +/- 125.489386 (428.56%) (init = 135)\n",
      "    Fit___HFModel__4___FWHML:          161.314301 +/- 33.9739752 (21.06%) (init = 101)\n",
      "    Fit___HFModel__4___scale:          93.5947425 +/- 8.00638611 (8.55%) (init = 90)\n",
      "    Fit___HFModel__4___Amp7_2to5_2:    0.5 (fixed)\n",
      "    Fit___HFModel__4___Amp7_2to7_2:    0.4938272 (fixed)\n",
      "    Fit___HFModel__4___Amp7_2to9_2:    0.3395062 (fixed)\n",
      "    Fit___HFModel__4___Amp9_2to7_2:    0.1728395 (fixed)\n",
      "    Fit___HFModel__4___Amp9_2to9_2:    0.4938272 (fixed)\n",
      "    Fit___HFModel__4___Amp9_2to11_2:   1 (fixed)\n",
      "    Fit___HFModel__7___centroid:      -102.447481 +/- 3.68927947 (3.60%) (init = -200)\n",
      "    Fit___HFModel__7___Al:             3299.70578 +/- 1.65670391 (0.05%) (init = 3300)\n",
      "    Fit___HFModel__7___Au:             60.1867272 +/- 0.66265495 (1.10%) (init = 60)\n",
      "    Fit___HFModel__7___Bl:             0 (fixed)\n",
      "    Fit___HFModel__7___Bu:             278.526941 +/- 10.7269620 (3.85%) (init = 270)\n",
      "    Fit___HFModel__7___Cl:             0 (fixed)\n",
      "    Fit___HFModel__7___Cu:             0 (fixed)\n",
      "    Fit___HFModel__7___FWHMG:          105.982664 +/- 27.4004187 (25.85%) (init = 135)\n",
      "    Fit___HFModel__7___FWHML:          123.594354 +/- 22.6178463 (18.30%) (init = 101)\n",
      "    Fit___HFModel__7___scale:          163.887281 +/- 7.01537309 (4.28%) (init = 160)\n",
      "    Fit___HFModel__7___Amp13_2to11_2:  0.6666667 (fixed)\n",
      "    Fit___HFModel__7___Amp13_2to13_2:  0.5530864 (fixed)\n",
      "    Fit___HFModel__7___Amp13_2to15_2:  0.3358025 (fixed)\n",
      "    Fit___HFModel__7___Amp15_2to13_2:  0.2246914 (fixed)\n",
      "    Fit___HFModel__7___Amp15_2to15_2:  0.5530864 (fixed)\n",
      "    Fit___HFModel__7___Amp15_2to17_2:  1 (fixed)\n",
      "    Fit___bkg___value1:                59.5623399 +/- 1.43149004 (2.40%) (init = 60)\n",
      "    Fit___bkg___value0:                59.6235980 +/- 1.37073255 (2.30%) (init = 60)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Chisquare fitting in progress (166.79425910209244): 379it [00:10, 33.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NDoF: 160, Chisquare: 166.79426, Reduced Chisquare: 1.0424641\n",
      "Akaike Information Criterium: 21.595642, Bayesian Information Criterium: 69.067432\n",
      "Errors scaled with reduced chisquare.\n",
      "[[Variables]]\n",
      "    s0_FWHMG:          28.9780398 +/- 126.283955 (435.79%) (init = 29.03236)\n",
      "    s0_FWHML:          161.420258 +/- 33.7741807 (20.92%) (init = 161.4094)\n",
      "    s0_TotalFWHM:      166.815855 +/- 22.1430745 (13.27%) == '0.5346*s0_FWHML+(0.2166*s0_FWHML**2+s0_FWHMG**2)**0.5'\n",
      "    s0_Scale:          93.5945060 +/- 7.98622420 (8.53%) (init = 93.59192)\n",
      "    s0_Saturation:     0 (fixed)\n",
      "    s0_Amp7_2__5_2:    0.5000937 (fixed)\n",
      "    s0_Amp7_2__7_2:    0.4939217 (fixed)\n",
      "    s0_Amp7_2__9_2:    0.3396039 (fixed)\n",
      "    s0_Amp9_2__7_2:    0.172911 (fixed)\n",
      "    s0_Amp9_2__9_2:    0.4939521 (fixed)\n",
      "    s0_Amp9_2__11_2:   1 (fixed)\n",
      "    s0_Al:             5313.73024 +/- 4.58833543 (0.09%) (init = 5313.729)\n",
      "    s0_Au:             105.169451 +/- 1.98179869 (1.88%) (init = 105.1687)\n",
      "    s0_Bl:             0 (fixed)\n",
      "    s0_Bu:             217.121406 +/- 17.0115404 (7.84%) (init = 217.1222)\n",
      "    s0_Cl:             0 (fixed)\n",
      "    s0_Cu:             0 (fixed)\n",
      "    s0_Centroid:       421.200845 +/- 6.25695848 (1.49%) (init = 421.2003)\n",
      "    s0_Background0:    59.5955667 +/- 1.20655137 (2.02%) (init = 59.59567)\n",
      "    s0_N:              0 (fixed)\n",
      "    s1_FWHMG:          105.967971 +/- 27.3184342 (25.78%) (init = 105.9688)\n",
      "    s1_FWHML:          123.617915 +/- 22.5420792 (18.24%) (init = 123.6177)\n",
      "    s1_TotalFWHM:      186.664575 +/- 11.2649134 (6.03%) == '0.5346*s1_FWHML+(0.2166*s1_FWHML**2+s1_FWHMG**2)**0.5'\n",
      "    s1_Scale:          163.878959 +/- 6.99085568 (4.27%) (init = 163.8787)\n",
      "    s1_Saturation:     0 (fixed)\n",
      "    s1_Amp13_2__11_2:  0.666746 (fixed)\n",
      "    s1_Amp13_2__13_2:  0.5531882 (fixed)\n",
      "    s1_Amp13_2__15_2:  0.3359059 (fixed)\n",
      "    s1_Amp15_2__13_2:  0.2247785 (fixed)\n",
      "    s1_Amp15_2__15_2:  0.55321 (fixed)\n",
      "    s1_Amp15_2__17_2:  1 (fixed)\n",
      "    s1_Al:             3299.70482 +/- 1.65151670 (0.05%) (init = 3299.705)\n",
      "    s1_Au:             60.1860267 +/- 0.66047269 (1.10%) (init = 60.18602)\n",
      "    s1_Bl:             0 (fixed)\n",
      "    s1_Bu:             278.542776 +/- 10.6879643 (3.84%) (init = 278.543)\n",
      "    s1_Cl:             0 (fixed)\n",
      "    s1_Cu:             0 (fixed)\n",
      "    s1_Centroid:      -102.447712 +/- 3.67802373 (3.59%) (init = -102.4477)\n",
      "    s1_Background0:    0 (fixed)\n",
      "    s1_N:              0 (fixed)\n",
      "SATLAS2: 0.133 s\n",
      "SATLAS1: 10.8 s\n"
     ]
    }
   ],
   "source": [
    "J = [0.5, 1.5]\n",
    "FWHMG = 135\n",
    "FWHML = 101\n",
    "\n",
    "spin1 = 4\n",
    "A1 = [5300, 100]\n",
    "B1 = [0, 230]\n",
    "C1 = [0, 0]\n",
    "centroid1 = 400\n",
    "bkg1 = 60\n",
    "scale1 = 90\n",
    "\n",
    "spin2 = 7\n",
    "A2 = [3300, 60]\n",
    "B2 = [0, 270]\n",
    "C2 = [0, 0]\n",
    "centroid2 = -100\n",
    "bkg2 = 60\n",
    "scale2 = 160\n",
    "\n",
    "x = np.arange(-13000, -9000, 40)\n",
    "x = np.hstack([x, np.arange(11000, 14000, 40)])\n",
    "rng = np.random.default_rng(0)\n",
    "\n",
    "# Interface models\n",
    "hfs1 = satlas2.HFSModel(I=spin1,\n",
    "                        J=J,\n",
    "                        ABC=[A1[0], A1[1], B1[0], B1[1], C1[0], C1[1]],\n",
    "                        centroid=centroid1,\n",
    "                        fwhm=[FWHMG, FWHML],\n",
    "                        scale=scale1,\n",
    "                        background_params=[bkg1],\n",
    "                        use_racah=True)\n",
    "hfs1.set_variation({'Cu': False})\n",
    "hfs2 = satlas2.HFSModel(I=spin2,\n",
    "                        J=J,\n",
    "                        ABC=[A2[0], A2[1], B2[0], B2[1], C2[0], C2[1]],\n",
    "                        centroid=centroid2,\n",
    "                        fwhm=[FWHMG, FWHML],\n",
    "                        scale=scale2,\n",
    "                        background_params=[bkg2],\n",
    "                        use_racah=True)\n",
    "hfs2.set_variation({'Cu': False})\n",
    "y = hfs1.f(x) + hfs2.f(x) + satlas2.Polynomial([bkg1]).f(x)\n",
    "y = rng.poisson(y)\n",
    "\n",
    "hfs1.params['centroid'].value = centroid1 - 100\n",
    "hfs2.params['centroid'].value = centroid2 - 100\n",
    "summodel = satlas2.SumModel([hfs1, hfs2], {\n",
    "    'values': [bkg1, bkg2],\n",
    "    'bounds': [0]\n",
    "})\n",
    "\n",
    "# SATLAS implementation\n",
    "hfs3 = sat.HFSModel(spin1,\n",
    "                    J, [A1[0], A1[1], B1[0], B1[1], C1[0], C1[1]],\n",
    "                    centroid - 100, [FWHMG, FWHML],\n",
    "                    scale=scale1,\n",
    "                    background_params=bkg,\n",
    "                    use_racah=True)\n",
    "hfs4 = sat.HFSModel(spin2,\n",
    "                    J, [A2[0], A2[1], B2[0], B2[1], C2[0], C2[1]],\n",
    "                    centroid - 100, [FWHMG, FWHML],\n",
    "                    scale=scale2,\n",
    "                    background_params=[0],\n",
    "                    use_racah=True)\n",
    "hfs3.set_variation({'Cu': False})\n",
    "hfs4.set_variation({'Background0': False, 'Cu': False})\n",
    "summodel2 = hfs3 + hfs4\n",
    "\n",
    "print('Fitting 1 dataset with chisquare (Pearson, satlas2)...')\n",
    "start = time.time()\n",
    "f = satlas2.chisquare_fit(summodel, x, y, modifiedSqrt(y))\n",
    "print(summodel.display_chisquare_fit(show_correl=False))\n",
    "stop = time.time()\n",
    "dt1 = stop - start\n",
    "start = time.time()\n",
    "sat.chisquare_fit(summodel2, x, y, modifiedSqrt(y))\n",
    "summodel2.display_chisquare_fit(show_correl=False, scaled=True)\n",
    "stop = time.time()\n",
    "dt2 = stop - start\n",
    "print('SATLAS2: {:.3} s'.format(dt1))\n",
    "print('SATLAS1: {:.3} s'.format(dt2))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cbe8ba4",
   "metadata": {},
   "source": [
    "The difference in coding implementation is a result of the interface automatically implementing a Step background, where the background is a constant for different regions in *x*-space. Notice here that the speedup due using the SATLAS2 implementation has risen from a factor 20 for a single spectrum to almost a factor 100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f70ea272",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 231928 characters]\n",
      "text/plain": [
       "<Figure size 1400x900 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_x = x[x<0]\n",
    "right_x = x[x>0]\n",
    "\n",
    "left_y = y[x<0]\n",
    "right_y = y[x>0]\n",
    "\n",
    "fig = plt.figure(constrained_layout=True, figsize=(14, 9))\n",
    "gs = gridspec.GridSpec(nrows=2, ncols=2, figure=fig)\n",
    "ax11 = fig.add_subplot(gs[0, 0])\n",
    "ax11.label_outer()\n",
    "ax12 = fig.add_subplot(gs[0, 1], sharey=ax11)\n",
    "ax12.label_outer()\n",
    "ax21 = fig.add_subplot(gs[1, 0], sharex=ax11)\n",
    "ax21.label_outer()\n",
    "ax22 = fig.add_subplot(gs[1, 1], sharex=ax12, sharey=ax21)\n",
    "ax22.label_outer()\n",
    "\n",
    "ax11.errorbar(left_x, left_y, modifiedSqrt(left_y), fmt='.', label='Artificial data')\n",
    "ax11.plot(left_x, hfs1(left_x), '-', label='SATLAS2 fit model 1')\n",
    "ax11.plot(left_x, hfs2(left_x), '-', label='SATLAS2 fit model 2')\n",
    "ax11.plot(left_x, summodel(left_x), '-', label='Sum of models')\n",
    "\n",
    "ax12.errorbar(right_x, right_y, modifiedSqrt(right_y), fmt='.', label='Artificial data')\n",
    "ax12.plot(right_x, hfs1(right_x), '-', label='SATLAS2 fit model 1')\n",
    "ax12.plot(right_x, hfs2(right_x), '-', label='SATLAS2 fit model 2')\n",
    "ax12.plot(right_x, summodel(right_x), '-', label='Sum of models')\n",
    "ax11.legend()\n",
    "\n",
    "ax21.errorbar(left_x, left_y, modifiedSqrt(left_y), fmt='.', label='Artificial data')\n",
    "ax21.plot(left_x, hfs3(left_x), '-', label='SATLAS fit model 1')\n",
    "ax21.plot(left_x, hfs4(left_x), '-', label='SATLAS fit model 2')\n",
    "ax21.plot(left_x, summodel2(left_x), '-', label='Sum of models')\n",
    "\n",
    "ax22.errorbar(right_x, right_y, modifiedSqrt(right_y), fmt='.', label='Artificial data')\n",
    "ax22.plot(right_x, hfs3(right_x), '-', label='SATLAS fit model 1')\n",
    "ax22.plot(right_x, hfs4(right_x), '-', label='SATLAS fit model 2')\n",
    "ax22.plot(right_x, summodel2(right_x), '-', label='Sum of models')\n",
    "ax21.legend()\n",
    "\n",
    "ax11.set_ylabel('SATLAS2')\n",
    "ax21.set_ylabel('SATLAS')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a3bb47e",
   "metadata": {},
   "source": [
    "# Different background for multiplets\n",
    "\n",
    "To demonstrate the convenience of the Step background, the same results are coded with SATLAS, where the use of LinkedModel is required. Note that here, the interface is *not* used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "66bcbe47",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fitting 1 dataset with chisquare (Pearson, satlas2)...\n",
      "[[Fit Statistics]]\n",
      "    # fitting method   = leastsq\n",
      "    # function evals   = 167\n",
      "    # data points      = 175\n",
      "    # variables        = 9\n",
      "    chi-square         = 170.061971\n",
      "    reduced chi-square = 1.02446971\n",
      "    Akaike info crit   = 12.9909634\n",
      "    Bayesian info crit = 41.4740372\n",
      "[[Variables]]\n",
      "    Artificial___HFS___centroid:      428.657492 +/- 7.19104252 (1.68%) (init = 400)\n",
      "    Artificial___HFS___Al:            5308.21780 +/- 5.60047942 (0.11%) (init = 5300)\n",
      "    Artificial___HFS___Au:            104.430048 +/- 2.30813893 (2.21%) (init = 100)\n",
      "    Artificial___HFS___Bl:            0 (fixed)\n",
      "    Artificial___HFS___Bu:            265.552916 +/- 19.1450076 (7.21%) (init = 230)\n",
      "    Artificial___HFS___Cl:            0 (fixed)\n",
      "    Artificial___HFS___Cu:            0 (fixed)\n",
      "    Artificial___HFS___FWHMG:         0.01009164 +/- 457.863719 (4537058.84%) (init = 135)\n",
      "    Artificial___HFS___FWHML:         179.647033 +/- 18.3401608 (10.21%) (init = 101)\n",
      "    Artificial___HFS___scale:         86.1616298 +/- 6.92395679 (8.04%) (init = 90)\n",
      "    Artificial___HFS___Amp7_2to5_2:   0.5 (fixed)\n",
      "    Artificial___HFS___Amp7_2to7_2:   0.4938272 (fixed)\n",
      "    Artificial___HFS___Amp7_2to9_2:   0.3395062 (fixed)\n",
      "    Artificial___HFS___Amp9_2to7_2:   0.1728395 (fixed)\n",
      "    Artificial___HFS___Amp9_2to9_2:   0.4938272 (fixed)\n",
      "    Artificial___HFS___Amp9_2to11_2:  1 (fixed)\n",
      "    Artificial___Step___value1:       62.1558866 +/- 1.12762809 (1.81%) (init = 60)\n",
      "    Artificial___Step___value0:       92.0732461 +/- 1.16029219 (1.26%) (init = 90)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Chisquare fitting done: 361it [00:10, 33.53it/s]                            \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NDoF: 163, Chisquare: 169.0098, Reduced Chisquare: 1.0368699\n",
      "Akaike Information Criterium: 17.904873, Bayesian Information Criterium: 55.882305\n",
      "Errors scaled with reduced chisquare.\n",
      "[[Variables]]\n",
      "    s0_FWHMG:         1.00200757 +/- 5039.81586 (502971.83%) (init = 1.002008)\n",
      "    s0_FWHML:         190.625875 +/- 44.2918367 (23.23%) (init = 190.6259)\n",
      "    s0_TotalFWHM:     190.632114 +/- 31.1034805 (16.32%) == '0.5346*s0_FWHML+(0.2166*s0_FWHML**2+s0_FWHMG**2)**0.5'\n",
      "    s0_Scale:         86.9765425 +/- 10.6152858 (12.20%) (init = 86.97654)\n",
      "    s0_Saturation:    0 (fixed)\n",
      "    s0_Amp7_2__5_2:   0.5000937 (fixed)\n",
      "    s0_Amp7_2__7_2:   0.4939217 (fixed)\n",
      "    s0_Amp7_2__9_2:   0.3396039 (fixed)\n",
      "    s0_Amp9_2__7_2:   0.172911 (fixed)\n",
      "    s0_Amp9_2__9_2:   0.4939521 (fixed)\n",
      "    s0_Amp9_2__11_2:  1 (fixed)\n",
      "    s0_Al:            5308.26995 +/- 5.39979300 (0.10%) (init = 5308.27)\n",
      "    s0_Au:            104.153329 +/- 2.24981898 (2.16%) (init = 104.1533)\n",
      "    s0_Bl:            0 (fixed)\n",
      "    s0_Bu:            267.624392 +/- 18.7374780 (7.00%) (init = 267.6244)\n",
      "    s0_Cl:            0 (fixed)\n",
      "    s0_Cu:            0 (fixed)\n",
      "    s0_Centroid:      428.989898 +/- 7.01443307 (1.64%) (init = 428.9899)\n",
      "    s0_Background0:   91.4973585 +/- 1.37306546 (1.50%) (init = 91.49736)\n",
      "    s0_N:             0 (fixed)\n",
      "    s1_FWHMG:         1.00051876 +/- 2743.17060 (274174.83%) (init = 1.000519)\n",
      "    s1_FWHML:         162.773898 +/- 34.3243957 (21.09%) (init = 162.7739)\n",
      "    s1_TotalFWHM:     162.781002 +/- 31.6370736 (19.44%) == '0.5346*s1_FWHML+(0.2166*s1_FWHML**2+s1_FWHMG**2)**0.5'\n",
      "    s1_Scale:         87.7067266 +/- 11.9631823 (13.64%) (init = 87.70673)\n",
      "    s1_Saturation:    0 (fixed)\n",
      "    s1_Amp7_2__5_2:   0.5000937 (fixed)\n",
      "    s1_Amp7_2__7_2:   0.4939217 (fixed)\n",
      "    s1_Amp7_2__9_2:   0.3396039 (fixed)\n",
      "    s1_Amp9_2__7_2:   0.172911 (fixed)\n",
      "    s1_Amp9_2__9_2:   0.4939521 (fixed)\n",
      "    s1_Amp9_2__11_2:  1 (fixed)\n",
      "    s1_Al:            5308.26995 +/- 5.39979301 (0.10%) == 's0_Al'\n",
      "    s1_Au:            104.153329 +/- 2.24981898 (2.16%) == 's0_Au'\n",
      "    s1_Bl:            0.00000000 +/- 0.00000000 (nan%) == 's0_Bl'\n",
      "    s1_Bu:            267.624392 +/- 18.7374780 (7.00%) == 's0_Bu'\n",
      "    s1_Cl:            0 (fixed)\n",
      "    s1_Cu:            0 (fixed)\n",
      "    s1_Centroid:      428.989898 +/- 7.01443307 (1.64%) == 's0_Centroid'\n",
      "    s1_Background0:   62.7881658 +/- 1.25873883 (2.00%) (init = 62.78817)\n",
      "    s1_N:             0 (fixed)\n",
      "SATLAS2: 0.062 s\n",
      "SATLAS1: 10.8 s\n"
     ]
    }
   ],
   "source": [
    "J = [0.5, 1.5]\n",
    "FWHMG = 135\n",
    "FWHML = 101\n",
    "\n",
    "spin1 = 4\n",
    "A1 = [5300, 100]\n",
    "B1 = [0, 230]\n",
    "C1 = [0, 0]\n",
    "centroid1 = 400\n",
    "bkg1 = 90\n",
    "scale1 = 90\n",
    "\n",
    "x = np.arange(-13000, -9000, 40)\n",
    "x = np.hstack([x, np.arange(11000, 14000, 40)])\n",
    "rng = np.random.default_rng(0)\n",
    "\n",
    "hfs = satlas2.HFS(spin1,\n",
    "                   J=J,\n",
    "                   A=[A1[0], A1[1]],\n",
    "                   B=[B1[0], B1[1]],\n",
    "                   C=[C1[0], C1[1]],\n",
    "                   df=centroid1,\n",
    "                   fwhmg=FWHMG,\n",
    "                   fwhml=FWHML,\n",
    "                   scale=scale1,\n",
    "                   racah=True\n",
    "                  )\n",
    "hfs.params['Cu'].vary = False\n",
    "bkg = satlas2.Step([bkg1, bkg2], [0])\n",
    "\n",
    "y = hfs1.f(x) + bkg.f(x)\n",
    "y = rng.poisson(y)\n",
    "\n",
    "s = satlas2.Source(x, y, yerr=modifiedSqrt, name='Artificial')\n",
    "s.addModel(hfs)\n",
    "s.addModel(bkg)\n",
    "f = satlas2.Fitter()\n",
    "f.addSource(s)\n",
    "\n",
    "hfs2 = sat.HFSModel(spin1,\n",
    "                    J, [A1[0], A1[1], B1[0], B1[1], C1[0], C1[1]],\n",
    "                    centroid - 100, [FWHMG, FWHML],\n",
    "                    scale=scale1,\n",
    "                    background_params=[bkg1],\n",
    "                    use_racah=True)\n",
    "hfs3 = sat.HFSModel(spin1,\n",
    "                    J, [A1[0], A1[1], B1[0], B1[1], C1[0], C1[1]],\n",
    "                    centroid - 100, [FWHMG, FWHML],\n",
    "                    scale=scale1,\n",
    "                    background_params=[bkg1],\n",
    "                    use_racah=True)\n",
    "hfs2.set_variation({'Cu': False})\n",
    "hfs3.set_variation({'Cu': False})\n",
    "linkedmodel = sat.LinkedModel([hfs2, hfs3])\n",
    "linkedmodel.shared = ['Al', 'Au', 'Bl', 'Bu', 'Centroid']\n",
    "linked_x = [x[x<0], x[x>0]]\n",
    "linked_y = [y[x<0], y[x>0]]\n",
    "\n",
    "print('Fitting 1 dataset with chisquare (Pearson, satlas2)...')\n",
    "start = time.time()\n",
    "f.fit()\n",
    "stop = time.time()\n",
    "print(f.reportFit())\n",
    "dt1 = stop - start\n",
    "start = time.time()\n",
    "sat.chisquare_spectroscopic_fit(linkedmodel, linked_x, linked_y, func=modifiedSqrt)\n",
    "stop = time.time()\n",
    "linkedmodel.display_chisquare_fit(show_correl=False, scaled=True)\n",
    "dt2 = stop - start\n",
    "print('SATLAS2: {:.3} s'.format(dt1))\n",
    "print('SATLAS1: {:.3} s'.format(dt2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3983a2f5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 118320 characters]\n",
      "text/plain": [
       "<Figure size 1400x900 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(constrained_layout=True, figsize=(14, 9))\n",
    "gs = gridspec.GridSpec(nrows=2, ncols=2, figure=fig)\n",
    "ax11 = fig.add_subplot(gs[0, 0])\n",
    "ax11.label_outer()\n",
    "ax12 = fig.add_subplot(gs[0, 1], sharey=ax11)\n",
    "ax12.label_outer()\n",
    "ax21 = fig.add_subplot(gs[1, 0], sharex=ax11)\n",
    "ax21.label_outer()\n",
    "ax22 = fig.add_subplot(gs[1, 1], sharex=ax12, sharey=ax21)\n",
    "ax22.label_outer()\n",
    "\n",
    "ax11.errorbar(linked_x[0], linked_y[0], modifiedSqrt(linked_y[0]), fmt='.', label='Artificial data')\n",
    "ax11.plot(linked_x[0], s.evaluate(linked_x[0]), '-', label='Fit')\n",
    "ax12.errorbar(linked_x[1], linked_y[1], modifiedSqrt(linked_y[1]), fmt='.', label='Artificial data')\n",
    "ax12.plot(linked_x[1], s.evaluate(linked_x[1]), '-', label='SATLAS2 fit model 1')\n",
    "ax11.legend()\n",
    "\n",
    "ax21.errorbar(linked_x[0], linked_y[0], modifiedSqrt(linked_y[0]), fmt='.', label='Artificial data')\n",
    "ax21.plot(linked_x[0], linkedmodel.models[0](linked_x[0]), '-', label='Fit')\n",
    "ax22.errorbar(linked_x[1], linked_y[1], modifiedSqrt(linked_y[1]), fmt='.', label='Artificial data')\n",
    "ax22.plot(linked_x[1], linkedmodel.models[1](linked_x[1]), '-', label='Fit')\n",
    "ax21.legend()\n",
    "\n",
    "ax11.set_ylabel('SATLAS2')\n",
    "ax21.set_ylabel('SATLAS')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eb6a959",
   "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
}
