{
  "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fits with shared parameters\n",
    "\n",
    "We demonstrate how to simultaneously fit two datasets with different models that shares a common parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from iminuit import Minuit\n",
    "from iminuit.cost import UnbinnedNLL\n",
    "from iminuit.util import describe\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "from numba_stats import norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate two data sets which are fitted simultaneously\n",
    "rng = np.random.default_rng(1)\n",
    "\n",
    "width = 2.0\n",
    "data1 = rng.normal(0, width, size=1000)\n",
    "data2 = rng.normal(5, width, size=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use two pdfs with different names for non-shared parameters,\n",
    "# so that they are fitted independently\n",
    "\n",
    "def pdf1(x, μ_1, σ):\n",
    "    return norm.pdf(x, μ_1, σ)\n",
    "\n",
    "def pdf2(x, μ_2, σ):\n",
    "    return norm.pdf(x, μ_2, σ)\n",
    "\n",
    "# combine two log-likelihood functions by adding them\n",
    "lh = UnbinnedNLL(data1, pdf1) + UnbinnedNLL(data2, pdf2)\n",
    "\n",
    "print(f\"{describe(lh)=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `σ` parameter is shared between the data sets, while the means of the two normal distributions are independently fitted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(cost, xe, minuit, ax, **style):\n",
    "    signature = describe(cost)\n",
    "    data = cost.data\n",
    "    \n",
    "    values = minuit.values[signature]\n",
    "    errors = minuit.errors[signature]\n",
    "\n",
    "    cx = (xe[1:] + xe[:-1]) / 2\n",
    "\n",
    "    ym = np.diff(norm.cdf(xe, *values)) * np.sum(w)\n",
    "    t = []\n",
    "    for n, v, e in zip(signature, values, errors):\n",
    "        t.append(f\"${n} = {v:.3f} ± {e:.3f}$\")\n",
    "    ax.plot(cx, ym, label=\"\\n\".join(t), **style)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = Minuit(lh, μ_1=1, μ_2=2, σ=1)\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(14, 5))\n",
    "\n",
    "hists = [np.histogram(lhi.data, bins=50) for lhi in lh]\n",
    "\n",
    "# draw data and model with initial parameters\n",
    "for lhi, (w, xe), axi in zip(lh, hists, ax):\n",
    "    cx = (xe[1:] + xe[:-1]) / 2\n",
    "    axi.errorbar(cx, w, np.sqrt(w), fmt=\"ok\", capsize=0, zorder=0)\n",
    "    plot(lhi, xe, m, axi, ls=\"--\")\n",
    "\n",
    "m.migrad()\n",
    "\n",
    "# draw model with fitted parameters\n",
    "for lhi, (w, xe), axi in zip(lh, hists, ax):\n",
    "    plot(lhi, xe, m, axi)\n",
    "    axi.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dashed line shows the initial model before the fit, the solid line shows the model after the fit. Note that the σ parameter is shared."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.14 ('venv': venv)",
   "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.10.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "bdbf20ff2e92a3ae3002db8b02bd1dd1b287e934c884beb29a73dced9dbd0fa3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}