{
    "cells": [
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "efef8e20-6571-4561-8f0b-6048b57907de",
            "metadata": {},
            "outputs": [],
            "source": [
                "import time\n",
                "import numpy as np\n",
                "np.set_printoptions(suppress=True, precision=6)\n",
                "\n",
                "import matplotlib.pyplot as plt\n",
                "from scipy.special import erf\n",
                "from aare import RisingScurve, FallingScurve\n",
                "from pprint import pprint\n"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "8449f9e9-59ee-4194-b477-d014b8efc93b",
            "metadata": {},
            "source": [
                "## Model curves"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "35772e2c-37c6-4986-a9c0-7dca4a034827",
            "metadata": {},
            "outputs": [],
            "source": [
                "def scurve(x, p): # rising Scurve\n",
                "    p0, p1, p2, p3, p4, p5 = p\n",
                "    return (p0 + p1*x) + 0.5 *(1 + erf((x-p2) / (np.sqrt(2)*p3))) * (p4 + p5*(x-p2))\n",
                "\n",
                "def scurve2(x, p): #falling Scurve\n",
                "    p0, p1, p2, p3, p4, p5 = p\n",
                "    return (p0 + p1*x) + 0.5 *(1 - erf((x-p2) / (np.sqrt(2)*p3))) * (p4 + p5*(x-p2))"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "a318ab93-fe18-4798-9ddb-6ffc2c2e3875",
            "metadata": {},
            "source": [
                "## Generate data (1D)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "06c8fddb-56d9-4f84-8b21-4ffd8b7bd26f",
            "metadata": {},
            "outputs": [],
            "source": [
                "x = np.linspace(0,120, 121)\n",
                "\n",
                "rng = np.random.default_rng(42)\n",
                "0\n",
                "p_true_rising = np.array([100.0, 0.25, 60.0, 6.0, 120.0,  1.0])\n",
                "p_true_falling = np.array([100.0, 0.25, 60.0, 6.0, 120.0, -1.0])\n",
                "\n",
                "y_true_rising = scurve(x, p_true_rising)\n",
                "y_true_falling = scurve2(x, p_true_falling)\n",
                "\n",
                "noise_sigma = 4\n",
                "\n",
                "y_rising = y_true_rising + rng.normal(0, noise_sigma, size=x.shape)\n",
                "# y_err_r = np.full_like(x, noise_sigma)\n",
                "\n",
                "y_falling = y_true_falling + rng.normal(0, noise_sigma, size=x.shape)\n",
                "# y_err_f = np.full_like(x, noise_sigma)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "7e185144-3b72-41f4-80d0-3c3504e48bba",
            "metadata": {},
            "source": [
                "## Plot synthetic data"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "9a0b3292-7c57-4483-aecc-7ef683ff0b62",
            "metadata": {},
            "outputs": [],
            "source": [
                "fig, ax = plt.subplots(1, 2, figsize=(12,4), sharex=True)\n",
                "\n",
                "ax[0].plot(x, y_true_rising, label=\"true\")\n",
                "ax[0].scatter(x, y_rising, s=12, label=\"noisy\")\n",
                "ax[0].set_title(\"Rising S-curve\")\n",
                "ax[0].legend()\n",
                "\n",
                "ax[1].plot(x, y_true_falling, label=\"true\")\n",
                "ax[1].scatter(x, y_falling, s=12, label=\"noisy\")\n",
                "ax[1].set_title(\"Falling S-curve\")\n",
                "ax[1].legend()\n",
                "\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "95c826d1-9223-4fb2-9ada-d9fc1acad2fb",
            "metadata": {},
            "source": [
                "## Fit the Rising S-curve"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "37d2fda4-40b8-4ea2-8f2c-6c3f2384d158",
            "metadata": {},
            "outputs": [],
            "source": [
                "model_r = RisingScurve()\n",
                "print(\"Paramter list of the Scurve:\")\n",
                "# print(model_r.GetParNames())\n",
                "print(model_r.par_names)\n",
                "print(model_r.n_par, '\\n')\n",
                "\n",
                "model_r.FixParameter('p0', 100)  # Fixed p0 = 100 (optimizer does not touch the value)\n",
                "model_r.SetParameter('A', 100) # Start value A = 100\n",
                "\n",
                "print(\"== Tuned fit settings ==\")\n",
                "model_r.compute_errors = True\n",
                "model_r.max_calls = 500\n",
                "model_r.tolerance = 0.01\n",
                "print(f\"max_calls : {model_r.max_calls}\")\n",
                "print(f\"tolerance : {model_r.tolerance}\")\n",
                "print(f\"compute_erros: {model_r.compute_errors}\")\n",
                "print(\"\\n\")\n",
                "print(\"== Results ==\")\n",
                "res_r_munuit_obj = model_r.fit(x, y_rising, np.sqrt(y_rising))\n",
                "\n",
                "print(\"True rising params:         \", p_true_rising)\n",
                "print(\"minuit_grad (obj api ) rising result:\")\n",
                "pprint(res_r_munuit_obj, sort_dicts=False)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "5ebf1328-954a-43ce-953b-c402fd7998f9",
            "metadata": {},
            "source": [
                "## Fit the Falling S-curve"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "5ce22d2d-1adb-47d9-bf36-6cb535fa8b77",
            "metadata": {},
            "outputs": [],
            "source": [
                "model_f = FallingScurve()\n",
                "model_f.SetParameter(2, 20) # set start p2 = 20\n",
                "model_f.SetParameter(4, 90) # set start p4 = 90\n",
                "model_f.SetParLimits(4, 80, 130)\n",
                "model_f.max_calls = 150 # this is limits the number of calls done by the minimizer (increasing the `max_calls` may help if the fit fails) \n",
                "res_f_munuit_obj = model_f.fit(x, y_falling)\n",
                "\n",
                "print(\"True falling params:  \", p_true_falling)\n",
                "print(\"minuit_grad (obj api ) falling result\")\n",
                "pprint(res_f_munuit_obj, sort_dicts=False)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "fe72503a-fc77-4b17-9b17-b98055e2fc87",
            "metadata": {},
            "outputs": [],
            "source": [
                "fig, ax = plt.subplots(1, 2, figsize=(12,4))\n",
                "\n",
                "# rising\n",
                "ax[0].plot(x, y_true_rising, label=\"true\")\n",
                "ax[0].scatter(x, y_rising, s=12, label=\"data\")\n",
                "ax[0].plot(x, model_r(x, res_r_munuit_obj['par']), linewidth=1, color=\"green\", label=\"minuit\")\n",
                "ax[0].set_title(\"Rising S-curve fits\")\n",
                "ax[0].legend()\n",
                "\n",
                "# falling\n",
                "ax[1].plot(x, y_true_falling, label=\"true\")\n",
                "ax[1].scatter(x, y_falling, s=12, label=\"data\")\n",
                "ax[1].plot(x, model_f(x, res_f_munuit_obj['par']), linewidth=1, color=\"green\", label=\"minuit\")\n",
                "ax[1].set_title(\"Falling S-curve fits\")\n",
                "ax[1].legend()\n",
                "\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "2c42ad91-fd59-47ba-b90d-8300f97eb27a",
            "metadata": {},
            "source": [
                "## Quick error check"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "3ba40e9a-0c64-4ef6-ba90-b1116f99b862",
            "metadata": {},
            "outputs": [],
            "source": [
                "print(\"Rising absolute parameter error:\", np.abs(res_r_munuit_obj['par'] - p_true_rising))\n",
                "print(\"Falling absolute parameter error:\", np.abs(res_f_munuit_obj['par'] - p_true_falling))\n"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "c9bbabfd-e24e-43ec-8622-983b3757d7a9",
            "metadata": {},
            "source": [
                "## Benchmark"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "dc0b52b6-9fc3-453d-b6eb-e325a2b8342e",
            "metadata": {},
            "outputs": [],
            "source": [
                "def bench(fn, n_repeats=200):\n",
                "    for _ in range(3):\n",
                "        fn()\n",
                "\n",
                "    t0 = time.perf_counter()\n",
                "    for _ in range(n_repeats):\n",
                "        res = fn()\n",
                "    t1 = time.perf_counter()\n",
                "    return res, (t1 - t0) / n_repeats\n",
                "\n",
                "model_rising = RisingScurve()\n",
                "model_falling = FallingScurve()\n",
                "\n",
                "_, t_rising = bench(lambda: model_rising.fit(x, y_rising))\n",
                "_, t_falling = bench(lambda: model_falling.fit(x, y_falling))\n",
                "\n",
                "print(f\"Rising Minuit2:  {1e3*t_rising:.3f} ms\")\n",
                "print(f\"Falling Minuit2: {1e3*t_falling:.3f} ms\")\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "e694c18a-a203-4132-a472-c93c4d3ce853",
            "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.12.5"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}