import numpy as np, pandas as pd, matplotlib.pyplot as plt, warnings
warnings.filterwarnings("ignore")
from scipy.optimize import minimize
from sklearn.linear_model import LogisticRegression, LinearRegression
BLUE="#2b6cb0"; RED="#c53030"; GREEN="#2f855a"; ORANGE="#dd6b20"; GREY="#a0aec0"; PURP="#6b46c1"
obs=pd.read_csv("lalonde_obs.csv"); exp=pd.read_csv("lalonde_exp.csv")
cov=["age","educ","black","hispan","married","nodegree","re74","re75"]
bench=exp.re78[exp.treat==1].mean()-exp.re78[exp.treat==0].mean()
W=obs.treat.values; Y=obs.re78.values; X=obs[cov].values.astype(float)
Xs=(X-X.mean(0))/X.std(0); Xt=Xs[W==1]; Xc=Xs[W==0]; Yt=Y[W==1]; Yc=Y[W==0]
target=Xt.mean(0)
def entropy_balance(Xc, target):                       # maximise entropy s.t. weighted means = target
    def obj(lam):
        a=-(Xc-target)@lam; M=a.max(); return M+np.log(np.exp(a-M).sum())   # convex dual
    lam=minimize(obj, np.zeros(Xc.shape[1]), method="BFGS", options={"maxiter":5000}).x
    a=-(Xc-target)@lam; a-=a.max(); w=np.exp(a); return w/w.sum()
w=entropy_balance(Xc, target)
att_eb=Yt.mean()-np.sum(w*Yc)
# report on the treated-group-sd scale, the convention the matching notebook uses
Xt_raw, Xc_raw = X[W==1], X[W==0]
tsd = Xt_raw.std(0)
smd_before=[np.abs(Xc_raw[:,j].mean()-Xt_raw[:,j].mean())/tsd[j] for j in range(len(cov))]
smd_after =[np.abs((w*Xc_raw[:,j]).sum()-Xt_raw[:,j].mean())/tsd[j] for j in range(len(cov))]
print(f"experimental benchmark ATT = ${bench:,.0f}")
print(f"entropy balancing ATT     = ${att_eb:,.0f}")
print(f"max |SMD| control vs treated: before {max(smd_before):.3f} -> after {max(smd_after):.1e}  (EXACT balance)")

rel = w*len(w)                                     # relative weight, 1 = uniform
ess = 1/np.sum(w**2)
print()
print(f"Now what that exactness cost. Effective control sample = {ess:.0f} of {len(Yc)} ({100*ess/len(Yc):.0f}% of the sample):")
print(f"  relative weights   min {rel.min():.3f}   median {np.median(rel):.3f}   max {rel.max():.2f}")
print(f"  controls carrying under a tenth of uniform weight: {(rel<0.1).sum()} of {len(rel)}")
print(f"  share of total weight held by the heaviest 25 controls: {100*np.sort(w)[-25:].sum():.0f}%")
fig,ax=plt.subplots(1,2,figsize=(13,4.4))
yv=np.arange(len(cov))
ax[0].scatter(smd_before,yv,color=RED,s=55,label="before (raw controls)",zorder=3)
ax[0].scatter(smd_after,yv,color=GREEN,s=55,label="after entropy balancing",zorder=3)
for y in yv: ax[0].plot([smd_before[y],smd_after[y]],[y,y],color=GREY,lw=.8)
ax[0].set_yticks(yv); ax[0].set_yticklabels(cov); ax[0].set_xlabel("|standardized mean gap| to treated"); ax[0].set_title("Entropy balancing: exact mean balance (green ≈ 0)"); ax[0].legend(fontsize=8)
ax[1].hist(w*len(w),bins=40,color=GREEN,alpha=.8); ax[1].axvline(1,color="k",ls="--",label="uniform weight")
ax[1].set_xlabel("relative control weight (1 = uniform)"); ax[1].set_ylabel("count"); ax[1].set_title(f"The price of exact balance: ESS {ess:.0f} of {len(Yc)}"); ax[1].legend(fontsize=8)
plt.tight_layout(); plt.show()
print()
print("Balance is EXACT by construction, with no propensity model and no fit-check-refit loop, and that is the")
print("genuine advantage of the method. But maximum entropy keeps the weights as uniform AS THE CONSTRAINTS")
print("ALLOW, and here the constraints are punishing because CPS respondents really do not look like NSW")
print(f"trainees. The typical control ends up under a fifth of uniform weight while a few carry up to {rel.max():.0f}x,")
print(f"and {100-100*ess/len(Yc):.0f}% of the effective sample is gone.")
print()
print("EXACT BALANCE IS NOT FREE -- IT IS PAID FOR IN VARIANCE. That trade is invisible if you report the")
print("balance table alone, and on a dataset with poor overlap it is exactly when the bill comes due.")

experimental benchmark ATT = $1,794
entropy balancing ATT     = $1,273
max |SMD| control vs treated: before 1.762 -> after 5.4e-06  (EXACT balance)

Now what that exactness cost. Effective control sample = 98 of 429 (23% of the sample):
  relative weights   min 0.019   median 0.183   max 9.42
  controls carrying under a tenth of uniform weight: 120 of 429
  share of total weight held by the heaviest 25 controls: 38%

Balance is EXACT by construction, with no propensity model and no fit-check-refit loop, and that is the
genuine advantage of the method. But maximum entropy keeps the weights as uniform AS THE CONSTRAINTS
ALLOW, and here the constraints are punishing because CPS respondents really do not look like NSW
trainees. The typical control ends up under a fifth of uniform weight while a few carry up to 9x,
and 77% of the effective sample is gone.

EXACT BALANCE IS NOT FREE -- IT IS PAID FOR IN VARIANCE. That trade is invisible if you report the
balance table alone, and on a dataset with poor overlap it is exactly when the bill comes due.

import statsmodels.api as sm
g=np.clip(LogisticRegression(penalty=None,max_iter=5000).fit(Xs,W).predict_proba(Xs)[:,1],.02,.98)
Qmod=LinearRegression().fit(np.column_stack([Xs,W]),Y)
Q1=Qmod.predict(np.column_stack([Xs,np.ones(len(W))])); Q0=Qmod.predict(np.column_stack([Xs,np.zeros(len(W))]))
Qobs=np.where(W==1,Q1,Q0); p1=W.mean()
H=(W-(1-W)*g/(1-g))/p1                                  # clever covariate for the ATT
eps=sm.OLS(Y-Qobs, H).fit().params[0]                   # targeting fluctuation
Q1s=Q1+eps*(1/p1); Q0s=Q0+eps*(-(g/(1-g))/p1)          # updated outcome model
att_tmle=np.mean((Q1s-Q0s)[W==1])
# initial (untargeted) plug-in for contrast
att_plugin=np.mean((Q1-Q0)[W==1])
print(f"initial plug-in (g-computation ATT) = ${att_plugin:,.0f}")
print(f"targeting fluctuation eps           = {eps:,.1f}")
print(f"TMLE ATT (after targeting)          = ${att_tmle:,.0f}")
print(f"experimental benchmark              = ${bench:,.0f}")
print("\nThe targeting step moves the plug-in estimate along the efficient-influence-function direction; TMLE is doubly robust")
print("(consistent if EITHER the outcome model or the propensity model is correct) and attains the efficiency bound.")
print()
print("Those are asymptotic guarantees, and this is one finite sample. Note which way targeting moved things:")
print(f"the plug-in sat at ${att_plugin:,.0f} and TMLE moved it to ${att_tmle:,.0f}, ${abs(att_tmle-att_plugin):,.0f} FURTHER from the")
print(f"${bench:,.0f} benchmark rather than closer. That is not evidence against TMLE -- an estimator with better")
print("asymptotic properties can land worse on any single dataset, and the benchmark is only visible here")
print("because someone ran the experiment. It is a reminder that efficiency is a claim about repeated")
print("sampling, not about this answer.")

initial plug-in (g-computation ATT) = $1,548
targeting fluctuation eps           = -36.3
TMLE ATT (after targeting)          = $1,199
experimental benchmark              = $1,794

The targeting step moves the plug-in estimate along the efficient-influence-function direction; TMLE is doubly robust
(consistent if EITHER the outcome model or the propensity model is correct) and attains the efficiency bound.

Those are asymptotic guarantees, and this is one finite sample. Note which way targeting moved things:
the plug-in sat at $1,548 and TMLE moved it to $1,199, $350 FURTHER from the
$1,794 benchmark rather than closer. That is not evidence against TMLE -- an estimator with better
asymptotic properties can land worse on any single dataset, and the benchmark is only visible here
because someone ran the experiment. It is a reminder that efficiency is a claim about repeated
sampling, not about this answer.

# reproduce ex1's NN-PS, IPW, AIPW for the comparison
ps=g; ti=np.where(W==1)[0]; ci=np.where(W==0)[0]
mps=ci[np.abs(ps[ti][:,None]-ps[ci][None,:]).argmin(1)]; att_nn=(Y[ti]-Y[mps]).mean()
keep=(ps>0.05)&(ps<0.95); wc=ps[keep]/(1-ps[keep]); Wk,Yk=W[keep],Y[keep]
att_ipw=Yk[Wk==1].mean()-np.sum(wc[Wk==0]*Yk[Wk==0])/np.sum(wc[Wk==0])
m0=LinearRegression().fit(X[W==0],Y[W==0]).predict(X); N1=(W==1).sum()
att_aipw=(np.sum((W==1)*(Y-m0))-np.sum((W==0)*(ps/(1-ps))*(Y-m0)))/N1
naive=Y[W==1].mean()-Y[W==0].mean()
res={"naive (CPS)":naive,"NN-propensity":att_nn,"IPW":att_ipw,"AIPW":att_aipw,"entropy\nbalancing":att_eb,"TMLE":att_tmle}
fig,ax=plt.subplots(figsize=(9,4.6))
nm=list(res); vv=[res[k] for k in nm]; cols=[GREY,BLUE,ORANGE,PURP,GREEN,RED]
ax.barh(nm,vv,color=cols); ax.axvline(bench,color="k",lw=2.5,ls="--"); ax.axvline(0,color="k",lw=.6)
lo_,hi_=min(vv+[0]),max(vv+[bench]); pad_=0.025*(hi_-lo_)
ax.set_xlim(lo_-7*pad_, hi_+7*pad_)
for i,v in enumerate(vv):
    ax.text(v+pad_ if v>=0 else v-pad_, i, f"${v:,.0f}", va="center",
            ha="left" if v>=0 else "right", fontsize=8)
ax.invert_yaxis(); ax.set_xlabel("estimated ATT (USD)")
ax.set_title(f"Modern balancing vs the matching-notebook estimators (dashed = ${bench:,.0f})")
plt.tight_layout(); plt.show()
print(f"  {'estimator':22} {'ATT':>9} {'vs benchmark':>14}")
for k_,v_ in res.items():
    print(f"  {k_.replace(chr(10),' '):22} ${v_:>8,.0f} {100*(v_-bench)/bench:>13.0f}%")
print(f"  {'experimental truth':22} ${bench:>8,.0f}")
print()
print("Entropy balancing and TMLE join IPW/AIPW in the same band, all far above the confounded naive gap.")
print("Modern balancing's edge is HOW it gets there -- exact or optimized balance and efficient")
print("doubly-robust inference -- and NOT a closer answer: every one of them sits about 30% below the")
print("benchmark that plain nearest-neighbour propensity matching hit almost exactly.")
print()
print(f"(AIPW reads ${res['AIPW']:,.0f} here against ${1226:,} on the matching page: this notebook reuses the CLIPPED")
print(" propensity score from the TMLE section, which moves the estimate by a dollar.)")

  estimator                    ATT   vs benchmark
  naive (CPS)            $    -635          -135%
  NN-propensity          $   1,792            -0%
  IPW                    $   1,316           -27%
  AIPW                   $   1,225           -32%
  entropy balancing      $   1,273           -29%
  TMLE                   $   1,199           -33%
  experimental truth     $   1,794

Entropy balancing and TMLE join IPW/AIPW in the same band, all far above the confounded naive gap.
Modern balancing's edge is HOW it gets there -- exact or optimized balance and efficient
doubly-robust inference -- and NOT a closer answer: every one of them sits about 30% below the
benchmark that plain nearest-neighbour propensity matching hit almost exactly.

(AIPW reads $1,225 here against $1,226 on the matching page: this notebook reuses the CLIPPED
 propensity score from the TMLE section, which moves the estimate by a dollar.)

Causal Inference II(c) — Modern Balancing: entropy balancing, CBPS, and TMLE¶

Optimize covariate balance directly, and estimate doubly-robustly and efficiently¶

1. Entropy balancing — exact covariate balance by construction¶

2. TMLE — doubly-robust and efficient, with a targeting step¶

3. Every method against the benchmark¶

4. Summary¶