{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "Hbol",
   "metadata": {
    "marimo": {
     "config": {
      "hide_code": true
     },
     "md_prefix": "r"
    }
   },
   "source": [
    "# Fitting Many Datasets with `MetalogFitMethod.Feasible`\n",
    "\n",
    "Ordinary least squares can return metalog coefficients whose density goes negative\n",
    "somewhere, especially with many terms. `MetalogFitMethod.Feasible` returns the best\n",
    "least-squares fit *among valid metalogs* (the coefficient vector a* of Baucells,\n",
    "Chrisman, Keelin and Xu, 2025), so every fit is a proper distribution.\n",
    "\n",
    "This notebook fits a batch of datasets with `Feasible`:\n",
    "\n",
    "1. **Fixed number of terms**: one call to `fit_grid_datasets`, which vmaps `fit`\n",
    "   over the datasets.\n",
    "2. **Different numbers of terms**: one `fit_grid_datasets` call per term count, then\n",
    "   pick the best term count for each dataset by KS distance.\n",
    "\n",
    "`fit_grid` cannot be used here: it varies the number of terms inside `vmap`, while the\n",
    "feasible solver needs it fixed at trace time, so it raises a `ValueError` for\n",
    "`Feasible`. `fit_grid_datasets` compiles once per term count instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "MJUe",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-17T19:53:05.400015Z",
     "iopub.status.busy": "2026-10-17T19:53:05.399803Z",
     "iopub.status.idle": "2026-10-17T19:53:06.154943Z",
     "shell.execute_reply": "2026-10-17T19:53:06.154558Z"
    }
   },
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "from scipy.stats import gamma, lognorm, lomax, weibull_min\n",
    "\n",
    "from metalog_jax.base import (\n",
    "    MetalogBoundedness,\n",
    "    MetalogFitMethod,\n",
    "    MetalogInputData,\n",
    "    MetalogParameters,\n",
    ")\n",
    "from metalog_jax.grid_search import extract_metalog, fit_grid, fit_grid_datasets\n",
    "from metalog_jax.utils import DEFAULT_Y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "vblA",
   "metadata": {
    "marimo": {
     "config": {
      "hide_code": true
     },
     "md_prefix": "r"
    }
   },
   "source": [
    "---\n",
    "## Data\n",
    "\n",
    "Five positive datasets of 200 samples each, including a heavy tail (Lomax) and a\n",
    "two-peaked mixture. Each is summarized by its empirical quantiles at `DEFAULT_Y`\n",
    "and stacked into `(n_datasets, n_quantiles)` arrays. All fits use a lower bound of 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bkHC",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-17T19:53:06.156381Z",
     "iopub.status.busy": "2026-10-17T19:53:06.156287Z",
     "iopub.status.idle": "2026-10-17T19:53:06.314717Z",
     "shell.execute_reply": "2026-10-17T19:53:06.314335Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batched x shape: (5, 105)\n",
      "Batched y shape: (5, 105)\n"
     ]
    }
   ],
   "source": [
    "_rng = np.random.default_rng(5)\n",
    "_bimodal = np.where(\n",
    "    _rng.random(200) < 0.6, _rng.normal(3.0, 0.5, 200), _rng.normal(7.0, 1.0, 200)\n",
    ")\n",
    "_samples = {\n",
    "    \"Lognormal\": lognorm(s=0.9).rvs(size=200, random_state=1),\n",
    "    \"Weibull\": weibull_min(c=1.5, scale=2).rvs(size=200, random_state=2),\n",
    "    \"Gamma\": gamma(a=2).rvs(size=200, random_state=3),\n",
    "    \"Lomax\": lomax(c=2.5).rvs(size=200, random_state=4),\n",
    "    \"Bimodal\": _bimodal,\n",
    "}\n",
    "dataset_names = list(_samples)\n",
    "\n",
    "_data = [\n",
    "    MetalogInputData.from_values(jnp.asarray(v), DEFAULT_Y, False)\n",
    "    for v in _samples.values()\n",
    "]\n",
    "batched_x = jnp.stack([d.x for d in _data])\n",
    "batched_y = jnp.stack([d.y for d in _data])\n",
    "\n",
    "print(f\"Batched x shape: {batched_x.shape}\")\n",
    "print(f\"Batched y shape: {batched_y.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "PKri",
   "metadata": {
    "marimo": {
     "config": {
      "hide_code": true
     },
     "md_prefix": "r"
    }
   },
   "source": [
    "---\n",
    "## Case 1: Fixed Number of Terms\n",
    "\n",
    "First, fit all datasets with OLS. `fit_grid` returns every fit, and its `feasible`\n",
    "mask says which are valid distributions. The mask comes from `check_feasibility`,\n",
    "which is exact: it finds every local minimum of the density and tests both tails,\n",
    "so it also catches negative density outside the probabilities that `fit` checks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "Xref",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-17T19:53:06.318720Z",
     "iopub.status.busy": "2026-10-17T19:53:06.318657Z",
     "iopub.status.idle": "2026-10-17T19:57:21.363272Z",
     "shell.execute_reply": "2026-10-17T19:57:21.362791Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OLS fits that are valid distributions: 2 of 5\n"
     ]
    }
   ],
   "source": [
    "FIXED_TERMS = 10\n",
    "\n",
    "params_ols = MetalogParameters(\n",
    "    boundedness=MetalogBoundedness.STRICTLY_LOWER_BOUND,\n",
    "    lower_bound=0.0,\n",
    "    upper_bound=0.0,\n",
    "    method=MetalogFitMethod.OLS,\n",
    "    num_terms=FIXED_TERMS,\n",
    ")\n",
    "\n",
    "# fit_grid returns every OLS fit, and marks which are valid distributions\n",
    "result_ols = fit_grid(batched_x, batched_y, params_ols)\n",
    "ols_feasible = result_ols.feasible\n",
    "print(\n",
    "    f\"OLS fits that are valid distributions: {int(ols_feasible.sum())} of {len(ols_feasible)}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "SFPL",
   "metadata": {
    "marimo": {
     "config": {
      "hide_code": true
     },
     "md_prefix": "r"
    }
   },
   "source": [
    "Now fit the same batch with `MetalogFitMethod.Feasible`. `fit_grid_datasets` vmaps\n",
    "`fit` over the datasets, so this is one compiled call for all of them. Where the OLS\n",
    "fit is already feasible, `Feasible` returns the same coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "BYtC",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-17T19:57:21.364616Z",
     "iopub.status.busy": "2026-10-17T19:57:21.364533Z",
     "iopub.status.idle": "2026-10-17T19:57:27.178457Z",
     "shell.execute_reply": "2026-10-17T19:57:27.178118Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Coefficients shape: (5, 10)\n",
      "KS Distances shape: (5,)\n",
      "\n",
      "Dataset     OLS valid   OLS KS  Feasible valid  Feasible KS\n",
      "Lognormal       False   0.0286            True       0.0286\n",
      "Weibull          True   0.0190            True       0.0190\n",
      "Gamma            True   0.0286            True       0.0286\n",
      "Lomax           False   0.0381            True       0.0381\n",
      "Bimodal         False   0.0571            True       0.0667\n"
     ]
    }
   ],
   "source": [
    "params_feasible = MetalogParameters(\n",
    "    boundedness=MetalogBoundedness.STRICTLY_LOWER_BOUND,\n",
    "    lower_bound=0.0,\n",
    "    upper_bound=0.0,\n",
    "    method=MetalogFitMethod.Feasible,\n",
    "    num_terms=FIXED_TERMS,\n",
    ")\n",
    "\n",
    "result_feasible = fit_grid_datasets(batched_x, batched_y, params_feasible)\n",
    "feasible_ok = result_feasible.feasible\n",
    "\n",
    "print(f\"Coefficients shape: {result_feasible.metalog.a.shape}\")\n",
    "print(f\"KS Distances shape: {result_feasible.ks_dist.shape}\")\n",
    "print()\n",
    "print(\n",
    "    f\"{'Dataset':<10} {'OLS valid':>10} {'OLS KS':>8} {'Feasible valid':>15} {'Feasible KS':>12}\"\n",
    ")\n",
    "for _i, _name in enumerate(dataset_names):\n",
    "    print(\n",
    "        f\"{_name:<10} {str(bool(ols_feasible[_i])):>10} {float(result_ols.ks_dist[_i]):>8.4f}\"\n",
    "        f\" {str(bool(feasible_ok[_i])):>15} {float(result_feasible.ks_dist[_i]):>12.4f}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "RGSE",
   "metadata": {
    "marimo": {
     "config": {
      "hide_code": true
     },
     "md_prefix": "r"
    }
   },
   "source": [
    "An infeasible OLS fit can have a *lower* KS distance than the feasible fit, because\n",
    "the KS distance only compares quantiles at the data points. It is still not a valid\n",
    "distribution: its density is negative somewhere, so its quantile function is not\n",
    "monotone.\n",
    "\n",
    "---\n",
    "## Case 2: Different Numbers of Terms\n",
    "\n",
    "The feasible solver needs the number of terms fixed when it compiles, so loop over\n",
    "the term counts in Python and batch the datasets inside each call. Each term count\n",
    "compiles once. Stacking the KS distances gives a `(n_datasets, n_terms)` grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "Kclp",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-17T19:57:27.179714Z",
     "iopub.status.busy": "2026-10-17T19:57:27.179656Z",
     "iopub.status.idle": "2026-10-17T19:57:36.894016Z",
     "shell.execute_reply": "2026-10-17T19:57:36.893475Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KS grid shape (n_datasets, n_terms): (5, 5)\n",
      "\n",
      "Dataset          4 terms       6 terms       8 terms      10 terms      12 terms\n",
      "Lognormal         0.0381        0.0286        0.0286        0.0286        0.0190\n",
      "Weibull           0.0476        0.0286        0.0286        0.0190        0.0190\n",
      "Gamma             0.0571        0.0381        0.0381        0.0286        0.0190\n",
      "Lomax             0.0476        0.0286        0.0286        0.0381        0.0381\n",
      "Bimodal           0.1333        0.0857        0.0667        0.0667        0.0571\n"
     ]
    }
   ],
   "source": [
    "num_terms_list = [4, 6, 8, 10, 12]\n",
    "\n",
    "# One compiled call per term count; each batches all datasets\n",
    "results_by_terms = {\n",
    "    _k: fit_grid_datasets(\n",
    "        batched_x, batched_y, params_feasible.replace(num_terms=_k)\n",
    "    )\n",
    "    for _k in num_terms_list\n",
    "}\n",
    "ks_grid = jnp.stack([results_by_terms[_k].ks_dist for _k in num_terms_list], axis=1)\n",
    "\n",
    "print(f\"KS grid shape (n_datasets, n_terms): {ks_grid.shape}\")\n",
    "print()\n",
    "print(f\"{'Dataset':<10}\" + \"\".join(f\"{_k:>8} terms\" for _k in num_terms_list))\n",
    "for _i, _name in enumerate(dataset_names):\n",
    "    print(f\"{_name:<10}\" + \"\".join(f\"{float(_v):>14.4f}\" for _v in ks_grid[_i]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "emfo",
   "metadata": {
    "marimo": {
     "config": {
      "hide_code": true
     },
     "md_prefix": "r"
    }
   },
   "source": [
    "Pick the term count with the lowest KS distance for each dataset (ties go to the\n",
    "fewest terms), then use `extract_metalog` to pull out a ready-to-use `Metalog`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "Hstk",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-17T19:57:36.895258Z",
     "iopub.status.busy": "2026-10-17T19:57:36.895199Z",
     "iopub.status.idle": "2026-10-17T19:57:53.534874Z",
     "shell.execute_reply": "2026-10-17T19:57:53.533828Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset     Terms       KS   Median      P90  Valid\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Lognormal      12   0.0190    1.110    2.944   True\n",
      "Weibull        12   0.0190    1.451    3.156   True\n",
      "Gamma          12   0.0190    1.717    4.349   True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Lomax           6   0.0286    0.368    1.785   True\n",
      "Bimodal        12   0.0571    3.435    7.393   True\n"
     ]
    }
   ],
   "source": [
    "best_term_idx = jnp.argmin(ks_grid, axis=1)\n",
    "best_metalogs = {}\n",
    "\n",
    "print(\n",
    "    f\"{'Dataset':<10} {'Terms':>6} {'KS':>8} {'Median':>8} {'P90':>8} {'Valid':>6}\"\n",
    ")\n",
    "for _i, _name in enumerate(dataset_names):\n",
    "    _k = num_terms_list[int(best_term_idx[_i])]\n",
    "    _m = extract_metalog(results_by_terms[_k], _i)\n",
    "    best_metalogs[_name] = _m\n",
    "    _median, _p90 = _m.ppf(jnp.array([0.5, 0.9]))\n",
    "    _valid = bool(results_by_terms[_k].feasible[_i])\n",
    "    print(\n",
    "        f\"{_name:<10} {_k:>6} {float(ks_grid[_i, int(best_term_idx[_i])]):>8.4f}\"\n",
    "        f\" {float(_median):>8.3f} {float(_p90):>8.3f} {str(_valid):>6}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f14adf66",
   "metadata": {},
   "source": [
    "`extract_metalog` returns an ordinary `Metalog`, so each best fit prints its\n",
    "fitted equation. The fits are bounded below by 0, so Q(y) = 0 + exp(M(y))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fd7cb840",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Lognormal\n",
      "Metalog: 12 terms, Feasible, lower bound 0\n",
      "  Q(y) = 0 + exp(M(y))\n",
      "  M(y) = 0.104676\n",
      "         + 10.289 ln(y/(1-y))\n",
      "         + 11.8326 (y - 0.5) ln(y/(1-y))\n",
      "         - 39.3163 (y - 0.5)\n",
      "         - 46.3022 (y - 0.5)^2\n",
      "         - 74.8519 (y - 0.5)^2 ln(y/(1-y))\n",
      "         + 251.515 (y - 0.5)^3\n",
      "         - 65.0183 (y - 0.5)^3 ln(y/(1-y))\n",
      "         + 174.852 (y - 0.5)^4\n",
      "         + 135.934 (y - 0.5)^4 ln(y/(1-y))\n",
      "         - 322.459 (y - 0.5)^5\n",
      "         + 73.0509 (y - 0.5)^5 ln(y/(1-y))\n",
      "--- Weibull\n",
      "Metalog: 12 terms, Feasible, lower bound 0\n",
      "  Q(y) = 0 + exp(M(y))\n",
      "  M(y) = 0.372291\n",
      "         - 1.09448 ln(y/(1-y))\n",
      "         - 1.66134 (y - 0.5) ln(y/(1-y))\n",
      "         + 5.98058 (y - 0.5)\n",
      "         + 5.56858 (y - 0.5)^2\n",
      "         + 3.69382 (y - 0.5)^2 ln(y/(1-y))\n",
      "         - 5.0922 (y - 0.5)^3\n",
      "         + 10.8938 (y - 0.5)^3 ln(y/(1-y))\n",
      "         - 31.7413 (y - 0.5)^4\n",
      "         + 3.06623 (y - 0.5)^4 ln(y/(1-y))\n",
      "         - 10.6074 (y - 0.5)^5\n",
      "         - 17.6535 (y - 0.5)^5 ln(y/(1-y))\n",
      "--- Gamma\n",
      "Metalog: 12 terms, Feasible, lower bound 0\n",
      "  Q(y) = 0 + exp(M(y))\n",
      "  M(y) = 0.540348\n",
      "         - 11.5595 ln(y/(1-y))\n",
      "         + 0.307369 (y - 0.5) ln(y/(1-y))\n",
      "         + 48.0526 (y - 0.5)\n",
      "         - 4.34273 (y - 0.5)^2\n",
      "         + 69.5203 (y - 0.5)^2 ln(y/(1-y))\n",
      "         - 212.118 (y - 0.5)^3\n",
      "         + 3.92031 (y - 0.5)^3 ln(y/(1-y))\n",
      "         + 11.9898 (y - 0.5)^4\n",
      "         - 90.9757 (y - 0.5)^4 ln(y/(1-y))\n",
      "         + 145.198 (y - 0.5)^5\n",
      "         - 24.9057 (y - 0.5)^5 ln(y/(1-y))\n",
      "--- Lomax\n",
      "Metalog: 6 terms, Feasible, lower bound 0\n",
      "  Q(y) = 0 + exp(M(y))\n",
      "  M(y) = -1.00038\n",
      "         - 0.30998 ln(y/(1-y))\n",
      "         - 0.517191 (y - 0.5) ln(y/(1-y))\n",
      "         + 4.14349 (y - 0.5)\n",
      "         + 1.36477 (y - 0.5)^2\n",
      "         + 2.38814 (y - 0.5)^2 ln(y/(1-y))\n",
      "--- Bimodal\n",
      "Metalog: 12 terms, Feasible, lower bound 0\n",
      "  Q(y) = 0 + exp(M(y))\n",
      "  M(y) = 1.23413\n",
      "         + 11.131 ln(y/(1-y))\n",
      "         + 15.6573 (y - 0.5) ln(y/(1-y))\n",
      "         - 42.9342 (y - 0.5)\n",
      "         - 55.1576 (y - 0.5)^2\n",
      "         - 77.7504 (y - 0.5)^2 ln(y/(1-y))\n",
      "         + 255.59 (y - 0.5)^3\n",
      "         - 91.0345 (y - 0.5)^3 ln(y/(1-y))\n",
      "         + 214.656 (y - 0.5)^4\n",
      "         + 133.956 (y - 0.5)^4 ln(y/(1-y))\n",
      "         - 323.198 (y - 0.5)^5\n",
      "         + 115.723 (y - 0.5)^5 ln(y/(1-y))\n"
     ]
    }
   ],
   "source": [
    "for _name, _m in best_metalogs.items():\n",
    "    print(f\"--- {_name}\")\n",
    "    _m.print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "nWHF",
   "metadata": {
    "marimo": {
     "config": {
      "hide_code": true
     },
     "md_prefix": "r"
    }
   },
   "source": [
    "---\n",
    "## Summary\n",
    "\n",
    "- `MetalogFitMethod.Feasible` guarantees every fit is a valid distribution, and\n",
    "  returns the OLS coefficients unchanged when OLS is already valid.\n",
    "- Use `fit_grid_datasets` to fit a batch of datasets with `Feasible` in one call.\n",
    "- To compare term counts, loop over them in Python with one `fit_grid_datasets` call\n",
    "  each. `fit_grid` raises a `ValueError` for `Feasible`.\n",
    "- Every grid result marks its valid fits in `feasible`, from `check_feasibility`,\n",
    "  an exact validity check that you can also run on any coefficient vector."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "metalog_jax",
   "language": "python",
   "name": "metalog_jax"
  },
  "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.13.7"
  },
  "marimo": {
   "header": "\"\"\"Batch feasible fitting demonstration (Baucells et al. 2025).\"\"\"\n\n",
   "marimo_version": "0.25.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}