qra/Bayesian_Zahra.py

#BetaPrime vs Gamma, hard-mono fit WITH constant + weak regularization

import numpy as np
import matplotlib.pyplot as plt
from scipy import optimize
from scipy.special import betaln, gammaln
import scipy.io as io

# ---------- Data ----------
data_path = "../data/"

# === Load data ===
data_path = "../data/"
suv = io.loadmat(data_path + "suv_percentilesSLOthenUWM.mat")['lung_SUVperc_COMBINED'][0:58, :, :]
flags = io.loadmat(data_path + "flags_combined.mat")['flags'][0:58, 3]

X = np.nanmax(suv[:, :, 94], axis=1).reshape(-1)
X_NC = X[flags == 0]
X_AE = X[flags == 1]


y = np.array([1]*len(X_AE) + [0]*len(X_NC), int) # 1=AE, 0=NC
X = np.concatenate([X_AE, X_NC], axis=0)
# sanity checks now that X,y actually exist
n = len(y); n1 = int(y.sum()); p_emp = n1 / n
rng = np.random.default_rng(12345)

# ---------- Helpers ----------
def logistic(z):
    z = np.clip(z, -60, 60)
    return 1.0/(1.0+np.exp(-z))

def sigmoid(t):
    return 1.0/(1.0+np.exp(-t))

def softplus(t):
    t = np.asarray(t, float)
    return np.log1p(np.exp(-np.abs(t))) + np.maximum(t, 0.0)

# --- Eq: logit P(AE|x) = log(p/(1-p)) + dE(x) ---
# where dE(x) = dE_ess(x) + C(params)

def dE_ess(x, a, b, s, k, th):
    # Essential (x-dependent) terms:
    # (a - k) * log(x) - (a + b) * log(1 + x/s) + x/th
    x = np.asarray(x, float)
    return (a - k) * np.log(x) - (a + b) * np.log1p(x/s) + x/th

def dE_const(a, b, s, k, th):
    # Constant (parameter-only) terms:
    # C = -a*log(s) - log B(a,b) + k*log(th) + log Γ(k)
    return -(a*np.log(s)) - betaln(a, b) + k*np.log(th) + gammaln(k)

def dE_full(x, a, b, s, k, th):
    # Total evidence term: dE(x) = dE_ess(x) + C
    return dE_ess(x, a, b, s, k, th) + dE_const(a, b, s, k, th)

# ---------- Global monotonicity cap for theta ----------
def theta_max(a, b, k, s, eps=1e-12):
    A = a - k
    if A <= 0:
        return np.inf
    r = np.sqrt(a + b) - np.sqrt(max(A, eps))
    if r <= 1e-12:
        return np.inf
    return s/(r*r)

# phi = [p_raw, b_raw, s_raw, k_raw, d_raw, u_raw]
def unpack_phi_mono(phi):
    p_raw, b_raw, s_raw, k_raw, d_raw, u_raw = phi
    
    # Map raw parameters into valid constrained space:
    # p = sigmoid(p_raw) ∈ (0,1) → AE prior (class prior)
    p = sigmoid(p_raw)                     
    
    # b = softplus(b_raw) > 0 → Beta–Prime shape parameter
    b = softplus(b_raw) + 1e-6             
    
    # s = softplus(s_raw) > 0 → Beta–Prime scale parameter
    s = softplus(s_raw) + 1e-6             
    
    # k = softplus(k_raw) > 0 → Gamma shape parameter
    k = softplus(k_raw) + 1e-6             
    
    # delta = softplus(d_raw) > 0; a = k + delta > k
    delta = softplus(d_raw) + 1e-6         
    a = k + delta                          
    
    # θ (theta) is constrained: 0 < θ ≤ θ_max(a,b,k,s)
    th_cap = theta_max(a, b, k, s)         
    th = th_cap * sigmoid(u_raw)            # map u_raw ∈ R into (0, th_cap]
    
    return p, a, b, s, k, th

# ---------- Priors ----------
# Beta prior on p centered at empirical rate
TAU = 25.0   # reduce to ~5 if you want it weaker
alpha = max(TAU * float(p_emp), 1e-6)
beta  = max(TAU * (1.0 - float(p_emp)), 1e-6)

# Weak log-normal shrinkage on positive parameters
def nlog_lognormal(x, mu, sigma, eps=1e-12):
    # -log LogNormal(x | mu, sigma) up to additive const
    x = np.maximum(x, eps)
    lx = np.log(x)
    return 0.5 * ((lx - mu)/sigma)**2 + lx

# ---------- Objective ----------
def neg_post_phi_mono_WITH_CONST_REG(phi, X, y):
    p, a, b, s, k, th = unpack_phi_mono(phi)
    eps = 1e-12

    # Likelihood with constant included
    z  = (np.log(p) - np.log(1-p)) + dE_full(X, a, b, s, k, th)
    px = logistic(z)
    nll = -np.sum(y*np.log(px + eps) + (1-y)*np.log(1 - px + eps))

    # Prior on p ~ Beta(alpha, beta)
    npr_p = -((alpha-1)*np.log(p + eps) + (beta-1)*np.log(1 - p + eps))

    # --- Regularization (weak priors) ---
    # AE median m1: use AE median if present; otherwise overall median.
    if (y == 1).any():
        m1 = np.median(X[y == 1])
    else:
        m1 = np.median(X)

    reg = 0.0  # total penalty starts at zero

    # 1) Gamma shape k (>0): very weak prior centered at 2 (σ=1.2).
    reg += nlog_lognormal(k, mu=np.log(2.0), sigma=1.2)

    # 2) Beta-Prime shape b (>0): same weak prior.
    reg += nlog_lognormal(b, mu=np.log(2.0), sigma=1.2)

    # 3) Beta-Prime scale s (>0): center near AE median (tighter σ=0.5).
    reg += nlog_lognormal(s, mu=np.log(max(m1, 1e-6)), sigma=0.5)

    # 4) Left-tail gap delta = a - k (>0): center around ~1.5 (σ=0.5)
    delta = a - k
    reg += nlog_lognormal(delta, mu=np.log(1.5), sigma=0.5)

    # Keep theta away from the boundary: Beta(3,3) on r = th/th_cap
    thcap = theta_max(a, b, k, s)
    if np.isfinite(thcap) and thcap > 0:
        r = np.clip(th/thcap, 1e-9, 1-1e-9)
        npr_r = -((3-1)*np.log(r) + (3-1)*np.log(1 - r))
    else:
        npr_r = 0.0

    return nll + npr_p + reg + npr_r

# ---------- Initialization ----------
def init_phi(X, y):
    # Method-of-moments init for Gamma(k, theta) using NC data (y==0)
    # ref: https://en.wikipedia.org/wiki/Gamma_distribution#Estimation_of_parameters
    X0 = X[y==0]
    m0 = X0.mean() if X0.size else X.mean()  # sample mean
    v0 = X0.var()  if X0.size else X.var()   # sample variance

    if v0 <= 0:
        # If variance is degenerate, pick a safe, sane starting point
        k0, th0 = 2.0, max(m0/2, 0.1)
    else:
        # MoM: k = m^2 / v, theta = v / m, with small eps and lower bounds
        k0  = max((m0**2)/(v0 + 1e-9), 1.5)
        th0 = max(v0/(m0 + 1e-9), 0.3)

    X1 = X[y==1]
    m1 = np.median(X1) if X1.size else np.median(X)


# Empirical AE rate as a starting prior for p. Clip away from 0/1 so logit is finite.
# p0 = n_AE / n, but truncated to [1e-3, 1-1e-3] to avoid infinities in log(p/(1-p)).
    p0 = np.clip(float(y.mean()), 1e-3, 1 - 1e-3)
    
    
    
# Simple, stable seeds for AE Beta–Prime:
#   b0 = 1.5   → mild shape; not too spiky, not too flat.
#   s0 = max(m1, 0.5) → anchor scale near the AE median, but don’t go tiny.
    b0, s0 = 1.5, max(m1, 0.5)


# Pack raw parameters φ for the optimizer.
# We optimize in an unconstrained space and map with:
#   p      = sigmoid(p_raw)
#   b,s,k  = softplus(raw) + 1e-6
#   a      = k + softplus(delta_raw) + 1e-6
#   theta  = theta_max * sigmoid(u_raw)
# To “invert” softplus for the initial guess we use log(expm1(v)) which is the exact inverse
# of softplus when you define softplus(t) = log(1 + exp(t)). The +1e-9 is just numerical padding.
    raw = np.array([
        np.log(p0/(1-p0)),                # p_raw
        np.log(np.expm1(b0) + 1e-9),      # b_raw
        np.log(np.expm1(s0) + 1e-9),      # s_raw
        np.log(np.expm1(k0) + 1e-9),      # k_raw
        np.log(np.expm1(1.0) + 1e-9),     # delta_raw
        -0.2                              # u_raw (keeps theta a bit below cap initially)
    ], float) 
    return raw

# ---------- Fitting ----------
def fit_hard_mono_WITH_CONST_REG(X, y, phi_start=None, maxtries=6, jitter=0.3, rng=None):
    if rng is None:
        rng = np.random.default_rng(12345)
    if phi_start is None:
        phi_start = init_phi(X, y)
    phi = phi_start.copy()
    last_err = None
    for _ in range(maxtries):
        res = optimize.minimize(
            neg_post_phi_mono_WITH_CONST_REG, phi, args=(X, y),
            method="L-BFGS-B",
            options=dict(maxiter=12000, ftol=1e-10)
        )
        if res.success and np.isfinite(res.fun):
            return unpack_phi_mono(res.x), res
        last_err = res
        phi = phi + rng.normal(0, jitter, size=phi.shape)
    raise RuntimeError(f"Fit failed. Last status: {getattr(last_err, 'message', 'n/a')}")

# ---------- Convenience ----------
def P_with(theta, x):
    p, a, b, s, k, th = theta
    L = (np.log(p) - np.log(1-p)) + dE_full(x, a, b, s, k, th)
    return logistic(L)

def diag_report(theta, X):
    p, a, b, s, k, th = theta
    thcap = theta_max(a, b, k, s)
    C = dE_const(a, b, s, k, th)
    logit_p = np.log(p) - np.log(1 - p)
    A = a - k
    den = np.sqrt(a + b) - np.sqrt(max(A, 1e-12))
    xs = np.inf if den <= 1e-12 else s*np.sqrt(max(A,1e-12))/den
    print({
        "p": p, "a": a, "b": b, "s": s, "k": k, "theta": th,
        "theta_max": thcap, "theta/theta_max": (th/thcap if np.isfinite(thcap) else np.nan),
        "logit(p)": logit_p, "C": C, "x* (bottleneck)": xs
    })

def plot_s_shape(theta, X, y, rng=None, ax=None, label='P(AE | x)'):
    if rng is None:
        rng = np.random.default_rng(0)
    if ax is None:
        fig, ax = plt.subplots(figsize=(7, 4.5))

    x_lo = max(1e-6, float(X.min())*0.8)
    x_hi = float(X.max())*1.2
    xg = np.linspace(x_lo, x_hi, 600)
    pg = P_with(theta, xg)

    ax.plot(xg, pg, lw=2, label=label)
    jit = (rng.random(len(X)) - 0.5) * 0.06
    y_jit = y + jit
    ax.scatter(X[y==0], y_jit[y==0], s=22, alpha=0.35, label='NC (y=0)', edgecolors='none')
    ax.scatter(X[y==1], y_jit[y==1], s=28, alpha=0.60, label='AE (y=1)', edgecolors='none')

    ax.set_ylim(-0.05, 1.05)
    ax.set_xlim(x_lo, x_hi)
    ax.set_xlabel('x')
    ax.set_ylabel('P(AE | x)')
    ax.set_title('S-shaped P(AE | x) with hard-mono fit (constant included, regularized)')
    ax.grid(True, alpha=0.3)
    ax.legend(loc='lower right', frameon=False)
    return ax

# ---------- Run fit ----------
theta_hat, res = fit_hard_mono_WITH_CONST_REG(X, y, rng=rng)
print("Optimization success:", res.success, "fval:", res.fun)
diag_report(theta_hat, X)

ax = plot_s_shape(theta_hat, X, y, rng=rng)
plt.show()




#CI Estimation



# Delta-method 
import numdifftools as nd  

# wrap scalar objective for numdifftools
def build_objective(X, y):
    def f(phi):
        return neg_post_phi_mono_WITH_CONST_REG(np.asarray(phi, float), X, y)
    return f

# compute Σ_φ (covariance in phi-space) at MAP using numdifftools.Hessian
phi_hat = res.x.copy()                    # MAP in raw-phi space 
f_obj = build_objective(X, y)             # scalar negative log-posterior
H = nd.Hessian(f_obj, method='central')(phi_hat)
Sigma_phi = invert_with_eigenfloor(H, floor=1e-6)


# 95% Wald band via Delta method
z = norm.ppf(0.975)  # 1.96 for 95%  ppf stands for percent point function — it’s the inverse CDF


def g_px_at_x(x):
    """Return g(φ) = P(AE | x, φ), so we can get ∇g(φ̂) via numdifftools.Gradient."""
    #Build a scalar function g(φ) = P(AE | x, φ) for a fixed x.
    #We return this function so numdifftools.Gradient can compute ∇g(φ̂).
    def g(phi):
        p, a, b, s, k, th = unpack_phi_mono(np.asarray(phi, float))
        L = (np.log(p) - np.log(1 - p)) + dE_full(x, a, b, s, k, th)
        return logistic(L)
    return g

# x-grid
x_lo = max(1e-6, float(X.min()) * 0.8)
x_hi = min(10.0, float(X.max()) * 1.2)
xg = np.linspace(x_lo, x_hi, 500)

p_hat = np.empty_like(xg)  # point estimate at MAP
p_lo  = np.empty_like(xg)  # lower 95%
p_hi  = np.empty_like(xg)  # upper 95%

for i, x in enumerate(xg):
    gx = g_px_at_x(x)
    # point estimate at MAP
    ph = gx(phi_hat)
    # gradient wrt φ at φ̂ via numdifftools.Gradient
    grad = nd.Gradient(gx, method='central')(phi_hat)  # shape (d,)
    # Delta-method variance on probability scale: var ≈ ∇g^T Σ_φ ∇g
    var = float(grad @ Sigma_phi @ grad)
    se  = np.sqrt(max(var, 0.0))

    p_hat[i] = ph
    p_lo[i]  = np.clip(ph - z * se, 0.0, 1.0)
    p_hi[i]  = np.clip(ph + z * se, 0.0, 1.0)

# plot
fig, ax = plt.subplots(figsize=(7,4.5 ))

# curve + band (sharp colors)
ax.plot(xg, p_hat, color="#000000", lw=2.2, label='P(AE|x) at MAP')
ax.fill_between(xg, p_lo, p_hi, facecolor="#1f77b4", alpha=0.18, label='95% Wald band (Delta)')
ax.plot(xg, p_lo, color="#1f77b4")
ax.plot(xg, p_hi, color="#1f77b4")

# overlay data with tiny vertical jitter
rng_plot = np.random.default_rng(999)
jit = (rng_plot.random(len(X)) - 0.5) * 0.06
ax.scatter(X[y==0], (y + jit)[y==0], s=22, alpha=0.55, color="#2ca02c", edgecolors='none', label='NC')
ax.scatter(X[y==1], (y + jit)[y==1], s=26, alpha=0.75, color="#ff7f0e", edgecolors='none', label='AE')
ax.set_ylim(-0.05, 1.05)
ax.set_xlabel('x')
ax.set_ylabel('P(AE | x)')
ax.set_title('Delta–method band using numdifftools Hessian/Gradient')
ax.grid(alpha=0.3)
ax.legend(loc='lower right')
plt.show()