Binomial mixed models: likelihood, latent effects and integration

from pathlib import Path
import json,sys
import numpy as np
import pandas as pd

from scipy.special import expit,logsumexp
from scipy.optimize import brentq
from scipy.integrate import quad
from scipy.linalg import cho_factor,cho_solve
from IPython.display import display,Math
ROOT=next(p for p in [Path.cwd(),*Path.cwd().parents] if (p/'Cargo.toml').exists())
import sys
sys.path.insert(0,str(ROOT/'docs/figures'))
from bridge import render_figure
sys.path.insert(0,str(ROOT/'docs/verifications'))
from compare_logistic_mixed import evaluator
fixture=json.loads((ROOT/'tests/fixtures/logistic-mixed-r.json').read_text())
report=json.loads((ROOT/'docs/verifications/logistic-mixed-results.json').read_text())
c=next(c for c in fixture['cases'] if c['name']=='cbpp_61')
fit=next(r['PHP'] for r in report['results'] if r['name']==c['name'])
raw=pd.read_csv(ROOT/'docs/data/cbpp.csv');display(raw.head(8))
x=np.array(c['x']);y=np.array(c['y']);m=np.array(c['weights']);s=np.round(y*m);groups=np.array(c['groups'])
beta=np.array(fit['coefficients']);tau=fit['random_intercept_sd']
def matrix(a,name):
    a=np.atleast_2d(a)
    body=r' \\ '.join(' & '.join(f'{v:.6g}' for v in row) for row in a)
    display(Math(name+r'=\begin{bmatrix}'+body+r'\end{bmatrix}'))
ids=np.flatnonzero(groups==groups[0]);display(raw.iloc[ids])
matrix(x[ids],r'X_{\mathrm{first\ herd}}')
matrix(beta[:,None],r'\hat\beta')
matrix((x[ids]@beta)[:,None],r'X_g\hat\beta')
print('X shape:',x.shape,'group count:',len(set(groups)))

e=x[ids]@beta;yg=y[ids];mg=m[ids]
def kernel(u):
    eta=e+tau*u
    return -.5*u*u-np.sum(mg*(yg*np.logaddexp(0,-eta)+(1-yg)*np.logaddexp(0,eta)))
mode=brentq(lambda u:tau*np.sum(mg*(yg-expit(e+tau*u)))-u,-20,20,xtol=1e-13)
p=expit(e+tau*mode);h=1+tau*tau*np.sum(mg*p*(1-p))
display(pd.DataFrame(dict(period=raw.iloc[ids].period,successes=s[ids],trials=mg,
                         fixed_log_odds=e,conditional_log_odds=e+tau*mode,
                         probability=p,score_contribution=s[ids]-mg*p,
                         information_contribution=mg*p*(1-p))))
print('Conditional standardized mode:',mode)
print('Random intercept b=tau*u:',tau*mode)
print('Score at mode:',tau*np.sum(mg*(yg-p))-mode)
print('Posterior precision:',h)
assert np.isclose(tau*mode,fit['random_effects'][0],atol=1e-7)

z,w=np.polynomial.hermite_e.hermegauss(9);w/=np.sqrt(2*np.pi)
u=mode+z/np.sqrt(h)
display(pd.DataFrame(dict(normal_node=z,normal_weight=w,adaptive_node=u,
                         log_integrand=[kernel(v) for v in u],
                         corrected_log_term=[np.log(weight)+kernel(v)+node*node/2 for node,weight,v in zip(z,w,u)])))
print('Laplace group log integral, without binomial constant:',kernel(mode)-.5*np.log(h))
print('9-node group log integral, without binomial constant:',
      logsumexp([np.log(weight)+kernel(v)+node*node/2 for node,weight,v in zip(z,w,u)])-.5*np.log(h))

likelihood=evaluator(c);theta=np.r_[beta,tau]
def information(step):
    delta=step*(1+abs(theta));n=len(theta);H=np.zeros((n,n))
    objective=lambda t:-likelihood(t[:-1],t[-1],61)[0]
    f0=objective(theta)
    for i in range(n):
        a=theta.copy();b=theta.copy();a[i]+=delta[i];b[i]-=delta[i]
        H[i,i]=(objective(a)+objective(b)-2*f0)/delta[i]**2
        for j in range(i):
            pp=theta.copy();pm=theta.copy();mp=theta.copy();mm=theta.copy()
            pp[i]+=delta[i];pp[j]+=delta[j];pm[i]+=delta[i];pm[j]-=delta[j]
            mp[i]-=delta[i];mp[j]+=delta[j];mm[i]-=delta[i];mm[j]-=delta[j]
            H[i,j]=H[j,i]=(objective(pp)-objective(pm)-objective(mp)+objective(mm))/(4*delta[i]*delta[j])
    return H
H=(4*information(5e-4)-information(1e-3))/3
cov=cho_solve(cho_factor(H),np.eye(len(theta)))[:-1,:-1]
matrix(H,r'H_{\beta,\tau}')
matrix(cov,r'\widehat{\mathrm{Var}}(\hat\beta)')
print('Standard errors:',np.sqrt(cov.diagonal()))
assert np.allclose(cov,fit['covariance'],atol=1e-6,rtol=1e-4)

rare=[r for r in report['results'] if r['dataset']=='rare_events']
display(pd.DataFrame([dict(nodes=r['quadrature_points'],intercept=r['PHP']['coefficients'][0],
                          group_SD=r['PHP']['random_intercept_sd'],
                          approximate_log_likelihood=r['PHP']['log_likelihood'],
                          independently_integrated_log_likelihood=r['scipy_integrated_log_likelihood'],
                          approximation_error=r['approximation_log_likelihood_error']) for r in rare]))
# Compare posterior shape with the Gaussian approximation in a small group.
rc=next(c for c in fixture['cases'] if c['name']=='rare_events_61')
rf=next(r['PHP'] for r in rare if r['quadrature_points']==61)
rx=np.array(rc['x']);ry=np.array(rc['y']);rg=np.array(rc['groups']);rtau=rf['random_intercept_sd'];re=rx@np.array(rf['coefficients'])
choices=[]
for group in dict.fromkeys(rg):
    ids=rg==group;ee=re[ids];yy=ry[ids]
    k=lambda u,ee=ee,yy=yy:-.5*u*u-np.sum(yy*np.logaddexp(0,-ee-rtau*u)+(1-yy)*np.logaddexp(0,ee+rtau*u))
    uu=brentq(lambda u:rtau*np.sum(yy-expit(ee+rtau*u))-u,-30,30)
    hh=1+rtau*rtau*np.sum(expit(ee+rtau*uu)*expit(-ee-rtau*uu))
    norm=quad(lambda u:np.exp(k(u)-k(uu)),-np.inf,np.inf,epsabs=1e-11)[0]
    error=k(uu)-.5*np.log(hh)-(k(uu)+np.log(norm)-.5*np.log(2*np.pi))
    choices.append((abs(error),group,k,uu,hh,norm))
_,group,k,uu,hh,norm=max(choices,key=lambda v:v[0])
u=np.linspace(-4,4,400)
actual=np.array([np.exp(k(v)-k(uu))/norm for v in u]);gaussian=np.sqrt(hh/(2*np.pi))*np.exp(-.5*hh*(u-uu)**2)
render_figure('posterior',dict(u=u.tolist(),actual=actual.tolist(),gaussian=gaussian.tolist(),
                             mode=float(uu),group=int(group)))