# Ceteris Lab downloadable Python script
# Course: Fundamentals of Python for Financial Econometrics

import numpy as np
from sklearn.linear_model import LogisticRegression
rng=np.random.default_rng(28); n=1500; x=rng.normal(size=(n,3)); ps=1/(1+np.exp(-(x[:,0]-.5*x[:,1]))); d=rng.binomial(1,ps); y=2*d+x[:,0]+rng.normal(size=n)
logit=LogisticRegression().fit(x,d); phat=logit.predict_proba(x)[:,1]
w=d/phat+(1-d)/(1-phat)
print('IPW ATE',np.average(d*y/phat)-np.average((1-d)*y/(1-phat)))