qra/python/bayesian/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


# === 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


# 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))

    

    # 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_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



# 1- Delta-method 
# ===== Delta band + compact summaries (minimal) =====
import numpy as np
import matplotlib.pyplot as plt
import numdifftools as nd
from scipy.stats import norm

# 1) Covariance in raw-phi space at MAP
phi_hat   = res.x.copy()
f_obj     = lambda phi: neg_post_phi_mono_WITH_CONST_REG(np.asarray(phi, float), X, y)
H         = nd.Hessian(f_obj, method='central')(phi_hat)
Sigma_phi = np.linalg.pinv(0.5*(H + H.T))  # robust inverse

# 2) Delta band on P(AE|x) via numdifftools.Gradient
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)
p_lo  = np.empty_like(xg)
p_hi  = np.empty_like(xg)

for i, x in enumerate(xg):
    gx   = g_px(x)
    ph   = gx(phi_hat)
    grad = nd.Gradient(gx, method='central')(phi_hat)
    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)

# 3) Plot
fig, ax = plt.subplots(figsize=(7.2, 4.4), dpi=140)
ax.plot(xg, p_hat, lw=2.0, label='P(AE|x) @ MAP')
ax.fill_between(xg, p_lo, p_hi, alpha=0.20, label='95% Delta band')
rngp = np.random.default_rng(999); jit = (rngp.random(len(X)) - 0.5) * 0.06
ax.scatter(X[y==0], (y+jit)[y==0], s=22, alpha=0.55, edgecolors='none', label='NC')
ax.scatter(X[y==1], (y+jit)[y==1], s=26, alpha=0.75, edgecolors='none', label='AE')
ax.set_ylim(-0.05, 1.05); ax.set_xlabel('x'); ax.set_ylabel('P(AE | x)')
ax.grid(alpha=0.3); ax.legend(loc='lower right')
plt.tight_layout(); plt.show()



w = p_hi - p_lo
mask = (xg >= float(X.min())) & (xg <= float(X.max()))
print("\nBand width (95% pointwise): "
      f"overall mean {w.mean():.3f}, max {w.max():.3f}; "
      f"in-range mean {w[mask].mean():.3f}, max {w[mask].max():.3f}")


try:
    G = lambda phi: np.array(unpack_phi_mono(np.asarray(phi, float)), float)  # -> [p,a,b,s,k,theta]
    J = nd.Jacobian(G)(phi_hat)
    Sigma_theta = J @ Sigma_phi @ J.T
    se = np.sqrt(np.maximum(np.diag(Sigma_theta), 0.0))
    theta_hat_vec = G(phi_hat); names = ["p","a","b","s","k","theta"]
    print("\nParameter 95% CIs (Delta/Wald):")
    for nm, v, svi in zip(names, theta_hat_vec, se):
        print(f"  {nm:>6s} : {v:.6g}  [ {v - z*svi:.6g}, {v + z*svi:.6g} ]")
except Exception as e:
    print("(Parameter CI step skipped:", e, ")")