{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Multiband parametric fit: RAINBOW model\n",
    "\n",
    "[`RainbowFit`](../api/#light_curve.RainbowFit) fits a physical **blackbody + bolometric-light-curve** model to multi-band\n",
    "photometric time series. It is designed for extragalactic transients such as supernovae,\n",
    "kilonovae, or tidal disruption events, where the SED is well approximated by a Planck\n",
    "function whose temperature and luminosity evolve with time.\n",
    "\n",
    "The method is described in [Russeil et al. 2024](https://ui.adsabs.harvard.edu/abs/2024A%26A...683A.251R/abstract).\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "[`RainbowFit`](../api/#light_curve.RainbowFit) requires [`iminuit`](https://iminuit.readthedocs.io/) ≥ 2.21.\n",
    "Install it together with all optional dependencies or separately.\n",
    "We would also need `matplotlib` for plotting, but of course it is not needed for the fitting itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install \"light-curve[full]\" matplotlib\n",
    "# OR\n",
    "# %pip install light-curve iminuit matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from light_curve import RainbowFit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": "## Generating synthetic supernova data\n\nWe simulate a Bazin-shaped bolometric light curve with a sigmoid temperature evolution,\nobserved in four LSST-like bands (*g r i z*). Observations are flux densities — not\nmagnitudes — because `RainbowFit` works in linear flux space.\n\nThe sigmoid temperature is parametrized by the **mid temperature** `T` = $(T_\\text{min}+T_\\text{max})/2$\nand a dimensionless **relative amplitude** `T_amplitude` = $(T_\\text{max}-T_\\text{min})/(T_\\text{max}+T_\\text{min})$,\nso $T_\\text{max} = T(1 + \\texttt{T\\_amplitude})$ and $T_\\text{min} = T(1 - \\texttt{T\\_amplitude})$. Here the\nsource cools between 15 000 K and 5 000 K, i.e. `T` = 10 000 K and `T_amplitude` = 0.5.\n\n### Physical parameters\n\n| Parameter | Symbol | Value |\n|-----------|--------|-------|\n| Peak time | $t_0$ | 60 000 (MJD) |\n| Amplitude | $A$ | 1.0 |\n| Rise time | $t_\\text{rise}$ | 5 d |\n| Fall time | $t_\\text{fall}$ | 30 d |\n| Mid temperature | $T = (T_\\text{min}+T_\\text{max})/2$ | 10 000 K |\n| Relative amplitude | $T_\\text{amplitude} = \\frac{T_\\text{max}-T_\\text{min}}{T_\\text{max}+T_\\text{min}}$ | 0.5 (i.e. $T_\\text{min}$ = 5 000 K, $T_\\text{max}$ = 15 000 K) |\n| Colour timescale | $t_\\text{color}$ | 10 d |"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": "rng = np.random.default_rng(0)\n\n# Effective wavelengths in Angstrom for griz\nband_wave_aa = {\"g\": 4770.0, \"r\": 6231.0, \"i\": 7625.0, \"z\": 9134.0}\n\n# True parameters\nreference_time = 60_000.0\namplitude      = 1.0\nrise_time      = 5.0\nfall_time      = 30.0\nTmax           = 15_000.0  # hot (early) temperature\nTmin           = 5_000.0   # cool (late) floor\nt_color        = 10.0\n\n# Sigmoid temperature fit parameters: mid temperature and relative amplitude\nT_mid          = 0.5 * (Tmin + Tmax)                # = 10 000 K\nT_amplitude    = (Tmax - Tmin) / (Tmax + Tmin)      # = 0.5\n\n# Per-band baseline flux (constant offset)\nbaselines = {b: 0.3 * amplitude + rng.exponential(scale=0.3 * amplitude) for b in band_wave_aa}\n\n# Build the feature object first so we can use .model() to generate data\nfeature = RainbowFit.from_angstrom(\n    band_wave_aa,\n    with_baseline=True,\n    bolometric=\"bazin\",\n    temperature=\"sigmoid\",\n)\n\n# True parameter vector, in feature.names order:\n# common + bolometric + temperature + baselines + chi2_placeholder\ntrue_params = [\n    reference_time, amplitude, rise_time, fall_time,  # Bazin bolometric\n    T_mid, T_amplitude, t_color,                       # sigmoid temperature: T (mid), T_amplitude, t_color\n    *baselines.values(),                               # per-band baselines\n    1.0,                                               # reduced chi^2 placeholder\n]\n\n# Random observation times covering ±3 rise/fall windows\nn_obs = 1000\nt    = np.sort(rng.uniform(reference_time - 3 * rise_time, reference_time + 3 * fall_time, n_obs))\nband = rng.choice(list(band_wave_aa), size=n_obs)\n\n# True model flux\nflux = feature.model(t, band, *true_params)\n\n# Poisson-like noise: S/N=10 at minimum flux\nflux_err = np.sqrt(flux * np.min(flux)) / 10.0\nflux = flux + rng.normal(0.0, flux_err)\n\nprint(f\"Generated {n_obs} observations in bands {list(band_wave_aa)}\")"
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Fitting the light curve\n",
    "\n",
    "Call `feature(t, flux, sigma=flux_err, band=band)`. The result is a NumPy array of\n",
    "fit parameters followed by the **reduced χ²** of the fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = feature(t, flux, sigma=flux_err, band=band)\n",
    "\n",
    "# Parameter names are stored in feature.names\n",
    "print(\"Parameter names:\", feature.names)\n",
    "print()\n",
    "for name, val, true in zip(feature.names, result, true_params):\n",
    "    print(f\"  {name:25s}  fitted={val:10.3f}  true={true:10.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Visualising the fit\n",
    "\n",
    "Plot the data and the best-fit model curve for each band."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = {\"g\": \"#2ca02c\", \"r\": \"#d62728\", \"i\": \"#9467bd\", \"z\": \"#8c564b\"}\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(9, 4))\n",
    "\n",
    "t_grid = np.linspace(t.min(), t.max(), 400)\n",
    "\n",
    "for b in band_wave_aa:\n",
    "    mask = band == b\n",
    "    ax.errorbar(\n",
    "        t[mask], flux[mask], yerr=flux_err[mask],\n",
    "        fmt=\".\", color=colors[b], alpha=0.4, label=f\"{b} data\",\n",
    "    )\n",
    "    band_grid = np.full(len(t_grid), b)\n",
    "    ax.plot(t_grid, feature.model(t_grid, band_grid, *result), color=colors[b], lw=2)\n",
    "\n",
    "ax.set_xlabel(\"Time (MJD)\")\n",
    "ax.set_ylabel(\"Flux density\")\n",
    "ax.set_title(\"RainbowFit — Bazin bolometric, sigmoid temperature\")\n",
    "ax.legend(ncol=2, fontsize=8)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Parameter uncertainties\n",
    "\n",
    "`fit_and_get_errors` returns a `(params, errors)` tuple. The errors are estimated by\n",
    "`iminuit` from the Hessian of the chi-squared surface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "params, errors = feature.fit_and_get_errors(t, flux, sigma=flux_err, band=band)\n",
    "\n",
    "print(f\"{'Parameter':25s}  {'Value':>10s}  {'Error':>10s}  {'True':>10s}\")\n",
    "print(\"-\" * 62)\n",
    "for name, val, err, true in zip(feature.names, params, errors, true_params):\n",
    "    print(f\"{name:25s}  {val:10.3f}  {err:10.3f}  {true:10.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Peak time\n",
    "\n",
    "The bolometric peak time can differ from `reference_time` for asymmetric light curves.\n",
    "`feature.peak_time(params)` computes it analytically from the bolometric term."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_peak = feature.peak_time(params)\n",
    "print(f\"Fitted peak time : {t_peak:.2f} MJD\")\n",
    "print(f\"True   peak time : {feature.peak_time(np.array(true_params)):.2f} MJD\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": "## Choosing different models\n\n`RainbowFit` accepts pluggable **bolometric**, **temperature**, and **spectral** terms.\n\n### Bolometric terms\n\n| Name | Parameters | Description |\n|------|-----------|-------------|\n| `'bazin'` (default) | `reference_time, amplitude, rise_time, fall_time` | Bazin 2011 rise-and-fall |\n| `'sigmoid'` | `reference_time, amplitude, rise_time` | Logistic rise only |\n| `'linexp'` | `reference_time, amplitude, rise_time` | Linear rise, exponential decay |\n| `'doublexp'` | `reference_time, amplitude, time1, time2, p` | Exponential rise, decay with two exponentials |\n\n### Temperature terms\n\nThe sigmoid terms use the mid temperature `T` = $(T_\\text{min}+T_\\text{max})/2$ and a relative\namplitude `T_amplitude` = $(T_\\text{max}-T_\\text{min})/(T_\\text{max}+T_\\text{min})$\n(`T_amplitude` = 0 is a constant temperature; a weak prior anchors it there unless the data demand a\ntemperature change). Strictly positive temperature corresponds to the independent box `T > 0`,\n`-1 < T_amplitude < 1`.\n\n| Name | Parameters | Description |\n|------|-----------|-------------|\n| `'sigmoid'` (default) | `T, T_amplitude, t_color` | Temperature drops sigmoidally between $T(1+\\texttt{T\\_amplitude})$ and $T(1-\\texttt{T\\_amplitude})$ |\n| `'constant'` | `T` | Fixed temperature |\n| `'delayed_sigmoid'` | `T, T_amplitude, t_color, t_delay` | Sigmoid delayed relative to the bolometric peak |\n\n### Spectral terms\n\n`'planck'` is a pure blackbody. The other terms describe SEDs that *deviate* from a blackbody\n(UV blanketing or power-law spectra); each reduces to Planck at a **null** parameter value and\ncarries a weak Gaussian prior anchoring it there, so a true blackbody is recovered unbiased.\n\n| Name | Parameters | Description |\n|------|-----------|-------------|\n| `'planck'` (default) | *(none)* | Standard blackbody $B_\\nu(T)$ — no extra parameters |\n| `'blanketed'` | `lambda_scale` | UV-extincted blackbody $B_\\nu(T)\\,e^{-\\tau}$, $\\tau = I\\,e^{-\\lambda/\\lambda_s}$; the extinction reach $\\lambda_s$ is anchored to the characteristic temperature `T` (shared with the temperature term) so the blanketing depth stays fixed as the source cools |\n| `'modified_bb'` | `beta` | Planck tilted by a power law, $B_\\nu(T)\\,(\\lambda/\\lambda_\\text{ref})^{\\beta}$; $\\beta = 0 \\Rightarrow$ Planck |\n| `'logparabola'` | `sp_a, sp_b` | Planck × $e^{a L + b L^2}$, $L = \\ln(\\lambda/\\lambda_\\text{ref})$; $a = b = 0 \\Rightarrow$ Planck |\n| `'genwien'` | `spec_k` | Generalized Wien $B_\\nu \\propto \\nu^3 e^{-x^{\\,\\texttt{spec\\_k}}}$, $x = h\\nu/k_B T$; $\\texttt{spec\\_k} = 1 \\Rightarrow$ Wien tail |\n\nBoth backends accept an `optimizer` argument: `\"iminuit\"` (default, robust Migrad) or\n`\"least_squares\"` (scipy Trust Region Reflective, sharing the same analytic Jacobian — usually\nfaster, with an automatic fall-back to iminuit when needed)."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constant temperature — fewer parameters, useful for bluer transients\n",
    "feature_const_T = RainbowFit.from_angstrom(\n",
    "    band_wave_aa,\n",
    "    with_baseline=False,\n",
    "    bolometric=\"bazin\",\n",
    "    temperature=\"constant\",\n",
    ")\n",
    "result_const = feature_const_T(t, flux, sigma=flux_err, band=band)\n",
    "print(\"Parameters (constant T model):\")\n",
    "for name, val in zip(feature_const_T.names, result_const):\n",
    "    print(f\"  {name:25s}  {val:.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Blanketed blackbody — adds UV extinction parameters: lambda_scale\n",
    "feature_blanketed = RainbowFit.from_angstrom(\n",
    "    band_wave_aa,\n",
    "    with_baseline=False,\n",
    "    bolometric=\"bazin\",\n",
    "    temperature=\"sigmoid\",\n",
    "    spectral=\"blanketed\",\n",
    ")\n",
    "result_blanketed = feature_blanketed(t, flux, sigma=flux_err, band=band)\n",
    "print(\"Parameters (blanketed spectral model):\")\n",
    "for name, val in zip(feature_blanketed.names, result_blanketed):\n",
    "    print(f\"  {name:25s}  {val:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "### Non-blackbody SED terms and the fast optimizer\n",
    "\n",
    "Beyond `'blanketed'`, the spectral terms `'modified_bb'`, `'logparabola'`, and `'genwien'`\n",
    "describe SEDs that depart from a pure blackbody (a power-law tilt, a log-parabola, or a\n",
    "generalized-Wien blue cutoff). Each reduces to Planck at a *null* parameter value\n",
    "(`beta = 0`, `sp_a = sp_b = 0`, `spec_k = 1`), and a weak prior anchors it there — so fitting\n",
    "them to genuinely blackbody data just recovers that null rather than inventing a deviation.\n",
    "\n",
    "The fit itself can run on scipy's `least_squares` backend via `optimizer=\"least_squares\"`,\n",
    "which shares the analytic Jacobian of the default `iminuit` path and is usually faster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Each deviation term reduces to a pure blackbody at a \"null\" parameter value, with a weak\n",
    "# prior anchoring it there. Fitting them to the (blackbody) data above recovers that null —\n",
    "# i.e. no spurious deviation is introduced.\n",
    "shape_params = {\"modified_bb\": [\"beta\"], \"logparabola\": [\"sp_a\", \"sp_b\"], \"genwien\": [\"spec_k\"]}\n",
    "for spectral, shape in shape_params.items():\n",
    "    feat = RainbowFit.from_angstrom(\n",
    "        band_wave_aa, with_baseline=True,\n",
    "        bolometric=\"bazin\", temperature=\"sigmoid\", spectral=spectral,\n",
    "    )\n",
    "    res = feat(t, flux, sigma=flux_err, band=band)\n",
    "    shape_vals = \", \".join(f\"{p}={res[feat.names.index(p)]:+.3f}\" for p in shape)\n",
    "    print(f\"{spectral:12s}  reduced chi2 = {res[-1]:.2f}   {shape_vals}\")\n",
    "\n",
    "# The same model fitted with the scipy least_squares backend (shares the analytic Jacobian).\n",
    "feature_ls = RainbowFit.from_angstrom(\n",
    "    band_wave_aa, with_baseline=True,\n",
    "    bolometric=\"bazin\", temperature=\"sigmoid\", optimizer=\"least_squares\",\n",
    ")\n",
    "result_ls = feature_ls(t, flux, sigma=flux_err, band=band)\n",
    "print(f\"\\n{'least_squares':12s}  reduced chi2 = {result_ls[-1]:.2f}  (matches the default iminuit fit)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "## See also\n",
    "\n",
    "- [RainbowFit API reference](../api/#light_curve.RainbowFit)\n",
    "- [Feature table](../../) — all extractors\n",
    "- [Parametric fits tutorial](../../fitting/) — single-band transient fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}