{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Your first distributional model\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/StrudelDoodleS/superglm/blob/master/docs/tutorials/distributional-model.ipynb)\n",
    "[Download the notebook](https://raw.githubusercontent.com/StrudelDoodleS/superglm/master/docs/tutorials/distributional-model.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q superglm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15d7f321",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "logging.getLogger(\"matplotlib.font_manager\").setLevel(logging.ERROR)\n",
    "candidates = (Path(\"../_static/superglm.mplstyle\"), Path(\"docs/_static/superglm.mplstyle\"))\n",
    "for candidate in candidates:\n",
    "    if candidate.exists():\n",
    "        plt.style.use(str(candidate))\n",
    "        break\n",
    "else:\n",
    "    # No docs tree at all means Colab or a bare download: keep matplotlib's defaults.\n",
    "    if any(candidate.parent.is_dir() for candidate in candidates):\n",
    "        raise FileNotFoundError(\"superglm.mplstyle: run this page from its own directory\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "`SuperLSS` fits several parameters of a response distribution together. A\n",
    "Gaussian model, for example, can let both the mean and the standard deviation\n",
    "vary with the input columns. Each parameter gets its own predictor.\n",
    "\n",
    "This walkthrough creates sample data, fits two models and compares their\n",
    "predictions on held-out rows. It needs NumPy, pandas and SuperGLM, with no data\n",
    "download.\n",
    "\n",
    "## Create some data\n",
    "\n",
    "Here the response mean depends on age and region. Its standard deviation also\n",
    "increases with age. These are simulated continuous measurements, so a Gaussian\n",
    "response is appropriate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from superglm import GaussianLS, SuperLSS, cat, s\n",
    "\n",
    "rng = np.random.default_rng(42)\n",
    "n = 800\n",
    "frame = pd.DataFrame(\n",
    "    {\n",
    "        \"age\": rng.uniform(18, 80, n),\n",
    "        \"region\": rng.choice([\"North\", \"South\"], n),\n",
    "    }\n",
    ")\n",
    "age = frame[\"age\"].to_numpy()\n",
    "south = (frame[\"region\"] == \"South\").to_numpy()\n",
    "mean = 10 + 0.08 * (age - 45) + 0.002 * (age - 45) ** 2 + 1.5 * south\n",
    "sd = 0.8 + 0.018 * (age - 18)\n",
    "y = rng.normal(mean, sd)\n",
    "\n",
    "X_train, X_test = frame.iloc[:600], frame.iloc[600:]\n",
    "y_train, y_test = y[:600], y[600:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "The split happens before fitting. Both models below learn their spline bases\n",
    "and smoothing parameters from the training rows.\n",
    "\n",
    "## Declare the predictors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "family = GaussianLS()\n",
    "\n",
    "model = SuperLSS(\n",
    "    family,\n",
    "    family.location(s(\"age\", kind=\"cr\", k=8), cat(\"region\")),\n",
    "    family.scale(s(\"age\", kind=\"cr\", k=6)),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "The family comes first. Its two helper calls describe the predictors:\n",
    "\n",
    "- `location(...)` models the Gaussian conditional mean. Age has a smooth\n",
    "  effect and region has a categorical effect.\n",
    "- `scale(...)` models the Gaussian standard deviation. It has its own smooth\n",
    "  age effect and no region effect.\n",
    "\n",
    "These declarations are configuration. `model` owns the fit. The age smooths\n",
    "have separate coefficients because they belong to different predictors.\n",
    "\n",
    "`kind=\"cr\"` selects a cubic regression spline. `k` sets its basis size;\n",
    "smoothing determines how much of that flexibility the fit uses. Use a bare\n",
    "string, such as `\"age\"`, when you want a numeric linear term instead.\n",
    "Categories are always explicit with `cat(...)`, including categories stored\n",
    "as numbers.\n",
    "\n",
    "Every parameter must have a declaration. An empty call such as\n",
    "`family.scale()` estimates an intercept-only predictor. It does not fix the\n",
    "parameter to a numeric value, and omitting the call is an error. Each predictor\n",
    "includes an intercept unless you pass `intercept=False` to its helper.\n",
    "\n",
    "Use helpers on the same family instance that you pass to `SuperLSS`. You may\n",
    "put the declarations in any order; names identify their parameters.\n",
    "\n",
    "## Fit the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit_reml(X_train, y_train, outer=\"efs+newton\")\n",
    "print(model.summary())\n",
    "print(model.diagnose())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "`fit_reml` estimates the coefficients and smoothing parameters jointly.\n",
    "`outer=\"efs+newton\"` adds Newton refinement to the smoothing updates. On this\n",
    "example, the default EFS updates stop after rejecting a proposal; the\n",
    "refinement reaches a stationary fit with the same convergence tolerances.\n",
    "`diagnose()` reports the stopping evidence. Use `fit` when you want to hold\n",
    "smoothing parameters fixed. Both fitting methods update the model and return it.\n",
    "\n",
    "The summary separates terms by predictor. An age effect in `location` changes\n",
    "the conditional mean. An age effect in `scale` changes the spread. The Gaussian\n",
    "scale link is `log(scale - scale_floor)`, so its coefficients act on that\n",
    "linked quantity.\n",
    "\n",
    "## Predict the mean, spread and a quantile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = model.predict_parameters(X_test)\n",
    "predicted_mean = model.predict(X_test)\n",
    "upper_quantile = model.predict_quantile(X_test, 0.95)\n",
    "\n",
    "predictions = parameters.assign(\n",
    "    observed=y_test,\n",
    "    predicted_mean=predicted_mean,\n",
    "    q95=upper_quantile,\n",
    ")\n",
    "print(predictions.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "For this family, `parameters` has `location` and `scale` columns. They contain\n",
    "the mean and standard deviation on the response scale. `predict` returns the\n",
    "same mean as the `location` column. `predict_link` is available when you need\n",
    "the linear predictors before applying the inverse links.\n",
    "\n",
    "The 95th percentile describes the upper part of each row's predictive\n",
    "distribution. It is not a confidence bound on the estimated mean.\n",
    "\n",
    "## Compare against constant spread\n",
    "\n",
    "Keep the same mean terms and estimate one standard deviation for all rows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "constant_scale = SuperLSS(\n",
    "    family,\n",
    "    family.location(s(\"age\", kind=\"cr\", k=8), cat(\"region\")),\n",
    "    family.scale(),\n",
    ").fit_reml(X_train, y_train, outer=\"efs+newton\")\n",
    "\n",
    "held_out_loss = pd.Series(\n",
    "    {\n",
    "        \"varying_scale\": model.scores(X_test, y_test, which=(\"log\",))[\"log\"].mean(),\n",
    "        \"constant_scale\": constant_scale.scores(X_test, y_test, which=(\"log\",))[\"log\"].mean(),\n",
    "    },\n",
    "    name=\"mean_negative_log_likelihood\",\n",
    ")\n",
    "print(held_out_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "A lower mean negative log-likelihood is better on these held-out rows. It\n",
    "scores the predicted distribution, so spread matters as well as mean. For\n",
    "your own data, make the split respect time or group boundaries when those\n",
    "matter.\n",
    "\n",
    "Reusing `family` does not share fitted coefficients. Constructing and fitting\n",
    "`constant_scale` leaves the first model's fit intact. The models each copy\n",
    "their configuration at construction.\n",
    "\n",
    "## Choose another response family\n",
    "\n",
    "Choose the family for the response you have. `GammaLS` models strictly\n",
    "positive values. `TweedieLSS` admits both zeros and positive values.\n",
    "`NegativeBinomialLS` models overdispersed counts. Their predictor names differ\n",
    "because their parameters differ.\n",
    "\n",
    "For example, a Tweedie declaration uses three helpers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "from superglm import TweedieLSS\n",
    "\n",
    "tweedie = TweedieLSS()\n",
    "tweedie_model = SuperLSS(\n",
    "    tweedie,\n",
    "    tweedie.mu(s(\"age\", kind=\"cr\", k=8), cat(\"region\")),\n",
    "    tweedie.phi(s(\"age\", kind=\"cr\", k=6)),\n",
    "    tweedie.p(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "This declares the mean, dispersion and power predictors. Their names in\n",
    "results and offsets are `mean`, `dispersion` and `power`. The empty `p()` call\n",
    "estimates one power value for all rows. This code only constructs the model;\n",
    "fit it to a response for which the Tweedie law is appropriate.\n",
    "\n",
    "See [family predictor names](../how-to/fit-a-distributional-model.md#family-predictor-names)\n",
    "for all nine families, including what their scale and shape parameters mean.\n",
    "For fitting options and return values, use the\n",
    "[SuperLSS API reference](../api/distributional.md). The\n",
    "[distributional model guide](../how-to/fit-a-distributional-model.md) covers weights,\n",
    "offsets, interactions and discrete fitting; the\n",
    "[checking guide](../how-to/check-a-distributional-fit.md) covers calibration and\n",
    "predictive diagnostics."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "default_lexer": "ipython3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}