{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eae6df92",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constrained Bayesian fit (Gamma for NC, Beta-Prime for AE) — no regularization, no r-prior\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import optimize\n",
    "from scipy.special import betaln, gammaln\n",
    "import scipy.io as io\n",
    "\n",
    "data_path = \"../data/\"\n",
    "suv   = io.loadmat(data_path + \"suv_percentilesSLOthenUWM.mat\")['lung_SUVperc_COMBINED'][0:58, :, :]\n",
    "flags = io.loadmat(data_path + \"flags_combined.mat\")['flags'][0:58, 3]   # 0=NC, 1=AE\n",
    "\n",
    "# Feature X = max SUV_94 per subject; label y = flags\n",
    "X = np.nanmax(suv[:, :, 94], axis=1).astype(float).ravel()\n",
    "y = np.asarray(flags, int).ravel()\n",
    "\n",
    "# Guard for logs\n",
    "X = np.clip(X, 1e-12, None)\n",
    "p_emp = float(y.mean())\n",
    "\n",
    "# Small helpers\n",
    "\n",
    "def logistic(z):\n",
    "    z = np.clip(z, -60, 60)\n",
    "    return 1.0 / (1.0 + np.exp(-z))\n",
    "\n",
    "def sigmoid(t):\n",
    "    return 1.0 / (1.0 + np.exp(-t))\n",
    "\n",
    "def softplus(t):\n",
    "    t = np.asarray(t, float)\n",
    "    return np.log1p(np.exp(-np.abs(t))) + np.maximum(t, 0.0)\n",
    "\n",
    "# dE(x) pieces for log-odds\n",
    "def dE_ess(x, a, b, s, k, th):\n",
    "    x = np.asarray(x, float)\n",
    "    return (a - k) * np.log(x) - (a + b) * np.log1p(x / s) + x / th\n",
    "\n",
    "def dE_const(a, b, s, k, th):\n",
    "    return -(a * np.log(s)) - betaln(a, b) + k * np.log(th) + gammaln(k)\n",
    "\n",
    "def dE_full(x, a, b, s, k, th):\n",
    "    return dE_ess(x, a, b, s, k, th) + dE_const(a, b, s, k, th)\n",
    "\n",
    "# Monotonicity cap for theta\n",
    "def theta_max(a, b, k, s, eps=1e-12):\n",
    "    A = a - k\n",
    "    if A <= 0:\n",
    "        return np.inf\n",
    "    r = np.sqrt(a + b) - np.sqrt(max(A, eps))\n",
    "    if r <= 1e-12:\n",
    "        return np.inf\n",
    "    return s / (r * r)\n",
    "\n",
    "# φ = [p_raw, b_raw, s_raw, k_raw, d_raw, u_raw]  (unconstrained)\n",
    "def unpack_phi_mono(phi):\n",
    "    p_raw, b_raw, s_raw, k_raw, d_raw, u_raw = phi\n",
    "    p = sigmoid(p_raw)                       # (0,1)\n",
    "    b = softplus(b_raw) + 1e-6               # >0\n",
    "    s = softplus(s_raw) + 1e-6               # >0\n",
    "    k = softplus(k_raw) + 1e-6               # >0\n",
    "    delta = softplus(d_raw) + 1e-6           # >0\n",
    "    a = k + delta                            # enforce a > k\n",
    "    th_cap = theta_max(a, b, k, s)           # theta cap from monotonicity\n",
    "    th = th_cap * sigmoid(u_raw)             # 0 < theta <= th_cap\n",
    "    return p, a, b, s, k, th\n",
    "\n",
    "#  Prior on p \n",
    "TAU = 25.0                                   # shrink toward empirical AE rate\n",
    "alpha = max(TAU * p_emp, 1e-6)\n",
    "beta  = max(TAU * (1.0 - p_emp), 1e-6)\n",
    "\n",
    "prior_r = None\n",
    "#prior_r = (1.01, 1.01)\n",
    "#prior_r = (1.05, 1.05)\n",
    "#prior_r = (3, 3)\n",
    "#prior_r = (1.2, 1.2)\n",
    "# Objective: negative log-posterior (likelihood + Beta prior on p)\n",
    "def neg_post_phi_mono(phi, X, y):\n",
    "    p, a, b, s, k, th = unpack_phi_mono(phi)\n",
    "    eps = 1e-12\n",
    "\n",
    "    L  = (np.log(p) - np.log(1 - p)) + dE_full(X, a, b, s, k, th)\n",
    "    px = logistic(L)\n",
    "    nll = -np.sum(y * np.log(px + eps) + (1 - y) * np.log(1 - px + eps))\n",
    "\n",
    "    # Beta(alpha, beta) prior on p → negative log-prior\n",
    "    npr_p = -((alpha - 1) * np.log(p + eps) + (beta - 1) * np.log(1 - p + eps))\n",
    "\n",
    "    if prior_r is None: return nll + npr_p\n",
    "        \n",
    "    # Prior on r = theta / theta_max (softly avoid boundaries)\n",
    "    thcap = theta_max(a, b, k, s)\n",
    "    if np.isfinite(thcap) and thcap > 0:\n",
    "        r = np.clip(th / thcap, 1e-9, 1 - 1e-9)\n",
    "        # Negative log Beta prior:  -[(α-1)log r + (β-1)log(1-r)]  (const dropped)\n",
    "        npr_r = -((prior_r[0] -1) * np.log(r) + (prior_r[1]-1)* np.log(1.0 - r))\n",
    "    else:\n",
    "        npr_r = 0.0\n",
    "\n",
    "    return nll + npr_p + npr_r\n",
    "\n",
    "# Initialization (stable, simple)\n",
    "def init_phi(X, y):\n",
    "    # Gamma(k, theta) MoM for NC group (only need k0 as a safe size proxy)\n",
    "    X0 = X[y == 0]\n",
    "    m0 = X0.mean() if X0.size else X.mean()\n",
    "    v0 = X0.var()  if X0.size else X.var()\n",
    "    k0 = 2.0 if v0 <= 0 else max((m0**2)/(v0 + 1e-9), 1.5)\n",
    "\n",
    "    # AE median to seed s0\n",
    "    X1 = X[y == 1]\n",
    "    m1 = np.median(X1) if X1.size else np.median(X)\n",
    "\n",
    "    p0 = np.clip(float(y.mean()), 1e-3, 1 - 1e-3)\n",
    "    b0, s0 = 1.5, max(m1, 0.5)\n",
    "\n",
    "    return np.array([\n",
    "        np.log(p0 / (1 - p0)),             # p_raw\n",
    "        np.log(np.expm1(b0) + 1e-9),       # b_raw\n",
    "        np.log(np.expm1(s0) + 1e-9),       # s_raw\n",
    "        np.log(np.expm1(k0) + 1e-9),       # k_raw\n",
    "        np.log(np.expm1(1.0) + 1e-9),      # d_raw  (delta)\n",
    "        -0.2                               # u_raw  (keeps theta a bit below cap initially)\n",
    "    ], float)\n",
    "\n",
    "#  Fit wrapper (one retry)\n",
    "def fit_bayes_mono(X, y, phi_start=None, rng=None):\n",
    "    if rng is None:\n",
    "        rng = np.random.default_rng(0)\n",
    "    if phi_start is None:\n",
    "        phi_start = init_phi(X, y)\n",
    "\n",
    "    obj = lambda phi: neg_post_phi_mono(phi, X, y)\n",
    "\n",
    "    res = optimize.minimize(\n",
    "        obj, phi_start, method=\"L-BFGS-B\",\n",
    "        options={\"maxiter\": 6000, \"ftol\": 1e-9}\n",
    "    )\n",
    "    if not (res.success and np.isfinite(res.fun)):\n",
    "        phi_try = phi_start + rng.normal(0, 0.2, size=phi_start.shape)\n",
    "        res = optimize.minimize(\n",
    "            obj, phi_try, method=\"L-BFGS-B\",\n",
    "            options={\"maxiter\": 6000, \"ftol\": 1e-9}\n",
    "        )\n",
    "    return unpack_phi_mono(res.x), res\n",
    "\n",
    "#  prediction \n",
    "def P_with(theta, x):\n",
    "    p, a, b, s, k, th = theta\n",
    "    L = (np.log(p) - np.log(1 - p)) + dE_full(x, a, b, s, k, th)\n",
    "    return logistic(L)\n",
    "\n",
    "# Run fit + plot\n",
    "theta_hat, res = fit_bayes_mono(X, y)\n",
    "print(\"Optimization success:\", res.success, \" fval:\", float(res.fun))\n",
    "\n",
    "(p,a,b,s,k,th) = theta_hat\n",
    "thcap = theta_max(a, b, k, s)\n",
    "print(\"theta (p,a,b,s,k,theta):\", tuple(float(t) for t in theta_hat), \"ratio(th):\", th/thcap)\n",
    "\n",
    "# x-range (cap right end at 10 for readability)\n",
    "x_lo = max(1e-6, float(X.min()) * 0.8)\n",
    "x_hi = min(10.0, float(X.max()) * 1.2)\n",
    "xg   = np.linspace(x_lo, x_hi, 600)\n",
    "p_curve = P_with(theta_hat, xg)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7.0, 4.6), dpi=140)\n",
    "ax.plot(xg, p_curve, color=\"#000000\", lw=2.2, label=\"P(AE|x) (MAP)\")\n",
    "\n",
    "# overlay data with tiny vertical jitter so points don't overlap\n",
    "rng_plot = np.random.default_rng(999)\n",
    "jit = (rng_plot.random(len(y)) - 0.5) * 0.06\n",
    "ax.scatter(X[y==0], (y + jit)[y==0], s=22, alpha=0.55, color=\"#2ca02c\", edgecolors='none', label='NC')\n",
    "ax.scatter(X[y==1], (y + jit)[y==1], s=26, alpha=0.75, color=\"#ff7f0e\", edgecolors='none', label='AE')\n",
    "\n",
    "ax.set_ylim(-0.05, 1.05)\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('P(AE | x)')\n",
    "\n",
    "if prior_r is None:\n",
    "    ax.set_title('Constrained Bayesian fit (no regularization, prior on p)')\n",
    "else:\n",
    "    ax.set_title(f'Constrained Bayesian fit (no regularization, prior on p and prior r{prior_r})')\n",
    "\n",
    "ax.grid(alpha=0.3)\n",
    "ax.legend(loc='lower right', frameon=False)\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}