<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Causal Inference of Lalonde
<table align="left">
  <td>
    <a target="_blank" href="https
Step1:  Data
The data are downloaded from
Dehejia and Wahba (1999)'s
data hosting website via functions
under ec.data.lalonde.
Experiment benchmark
The first column is the treatment assignment. The last column is the post-intervention outcome variable. The remaining 8 columns are pre-intervention covariates.
Step2:  We use two methods to estimate the average treatment effects on the treated.
Difference Estimator
Using a simple difference in means, we estimate that the program increased the average earnings of the treated by $1794.
Step3:  Using the estimator for the sampling variance (see formula 6.18 of Imbens and Rubin (2014)), we estimate the 95% $z$-interval of the difference as [479, 3109].
Step4:  Regression Estimator
The experimental treated and experimental control group are not perfectly balanced.
Step5:  Thus, we can fit a linear outcome regression model to "clean up" the imbalance. It yields an estimated average treatment effect on the treated of $1698 with a slightly more efficient 95% confidence interval [458, 2938].
Step6:  Observational controls
Step7:  We create two additional indicators umemployed1974 and umemployed1975 based on the earnings.
Step8:  The observational control group is significantly different from the experimental treatment group.
Step9:  The experimental treated subjects tend to be younger, with less education, black, single, without high school degree, and have smaller pre-intervention earnings.
Step10:  Empirical Calibration
Empirical calibration assigns weights to the observational control subjects such that the reweighted control group equates the treated group on moments of covariates. The weights are also sought to be as uniform as possible.
Prepare Balancing Constraints
There are 8 raw pre-intervention covariates. Simply balancing on their first moment might not be sufficient. We follow Hainmueller (2011)'s analysis and consider a total of 52 balancing constraints, which include single terms, selected one-way interactions, and selected squared terms. Note that balancing on all 52 transformed covariates might be excessive, we choose them simply to reproduce Hainmueller's results.
We use patsy's formula API to construct the covariate matrices, where "
Step11:  The full list of the 52 balancing constraints is shown below.
Step12:  Entropy Balancing
We first use the default ENTROPY objective which mimizes the Kullback-Leibler divergence between the balancing weights and uniform weights.
Step13:  It turns out that we need to relax the equality constraint a tiny bit to ensure convergence.
Step14:  Top 3 control observations account for 12% of total weight.
Step15:  The effective sample size of the control group is merely 103.4, a huge reduction from the original sample size 15992.
Step16:  We confirm that the means of the raw covariates are matched after weighting.
Step17:  The difference in means between the treated group and the reweighted control group yields an average treatment effect on the treated of $1571, which reproduces Hainmueller (2011)'s result.
Step18:  Under the homoskedastic assumption (i.e., equal unit-level conditional variance), we estimate the variance using formula prescribed in section 19.7 of Imbens and Rubin (2015).
Step19:  The 95% confidence interval is [55, 3088], wider than the reported [97, 3044] by Hainmueller (2011) (it is unclear what method was used there).
Step20:  Quadratic Balancing
We then use the QUADRATIC objective which minimizes the Euclidean distance between the balancing weights and the uniform weights.
Step21:  Due to being the explicit optimizing objective, the effective control sample size increases.
Step22:  The point estimate $1713 is getting closer to the experimental benchmark $1794.
Step23:  The 95% confidence interval becomes narrower due to increased effective control sample size.
Step24:  Imposing an Upper Bound on Weights
For both entropy balancing and quadratic balancing, one can optionally impose an upper bound on the weights to avoid extreme individual weights. For illustration, we impose an upper bound 0.015 on quadratic balancing weights, i.e., no single control individual accounts for 1.5% or more of total weights.
Step25:  We can confirm the resulting weights are indeed capped at 0.015.
Step26:  The cost of this additional constraint is a small loss in effective sample size.
Step27:  Note that this additional weights constraint is an integral part of the optimization, as opposed to a separate trimming step used with propensity score weighting type method. We don't sacrifice the balancing conditions.
Step28:  The point estimate is very close to the experimental estimate $1794.
<ASSISTANT_TASK:>
Python Code:
import numpy as np
import pandas as pd
import patsy
import seaborn as sns
import statsmodels.api as sm
# install and import ec
!pip install -q git+https://github.com/google/empirical_calibration
import empirical_calibration as ec
sns.set_style('whitegrid')
%config InlineBackend.figure_format='retina'
%precision 3
Explanation: Causal Inference of Lalonde
<table align="left">
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/google/empirical_calibration/blob/master/notebooks/causal_inference_lalonde.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/google/empirical_calibration/blob/master/notebooks/causal_inference_lalonde.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
</table>
Causal Inference of Lalonde
Background
Imports
Data
Experiment benchmark
Difference Estimator
Regression Estimator
Observational controls
Empirical Calibration
Prepare Balancing Constraints
Entropy Balancing
Quadratic Balancing
Imposing an Upper Bound on Weights
We illustrate the use of empirical calibration on the observational part of the LaLonde data to evaluate the effect of a job training program.
Background
The LaLonde (1986) data is a
canonical benchmark in the causal inference literature. It consists of three
groups for evaluating the effect of a large scale job training program — the
National Supported Work Demonstration (NSW):
An experimental treatment group with 185 observations.
An experimental control group with 260 observations.
An observational control group drawn from the Current Population Survey
    (CPS), with 15,992 observations.
The outcome variable is the post-intervention earnings in 1978. The following
pre-intervention covariates are available for all three groups.
Name         | Description
------------ | ---------------------------------------------------
age          | age.
education    | years of schooling.
black        | 1 if black, 0 otherwise.
hispanic     | 1 if hispanic, 0 otherwise.
married      | 1 if married, 0 otherwise.
nodegree     | 1 if not completed high school degree, 0 otherwise.
earnings1974 | pre-intervention earnings in 1974.
earnings1975 | pre-intervention earnings in 1975.
Imports
End of explanation
#@title experimentl treatment group
treated = ec.data.lalonde.experimental_treated()
treated.head()
#@title experimental control group
control = ec.data.lalonde.experimental_control()
control.head()
Explanation: Data
The data are downloaded from
Dehejia and Wahba (1999)'s
data hosting website via functions
under ec.data.lalonde.
Experiment benchmark
The first column is the treatment assignment. The last column is the post-intervention outcome variable. The remaining 8 columns are pre-intervention covariates.
End of explanation
diff = treated["earnings1978"].mean() - control["earnings1978"].mean()
diff
Explanation: We use two methods to estimate the average treatment effects on the treated.
Difference Estimator
Using a simple difference in means, we estimate that the program increased the average earnings of the treated by $1794.
End of explanation
sd = np.sqrt(
    treated["earnings1978"].var(ddof=1) / len(treated["earnings1978"]) +
    control["earnings1978"].var(ddof=1) / len(control["earnings1978"]))
sd
Explanation: Using the estimator for the sampling variance (see formula 6.18 of Imbens and Rubin (2014)), we estimate the 95% $z$-interval of the difference as [479, 3109].
End of explanation
df = pd.concat([treated, control])
print(df.groupby("treatment").mean().T)
Explanation: Regression Estimator
The experimental treated and experimental control group are not perfectly balanced.
End of explanation
lm = sm.OLS(df["earnings1978"], df.iloc[:, :-1])
print(lm.fit().summary())
Explanation: Thus, we can fit a linear outcome regression model to "clean up" the imbalance. It yields an estimated average treatment effect on the treated of $1698 with a slightly more efficient 95% confidence interval [458, 2938].
End of explanation
control = ec.data.lalonde.observational_control()
control.head()
Explanation: Observational controls
End of explanation
treated["unemployed1974"] = treated["earnings1974"] == 0
treated["unemployed1975"] = treated["earnings1975"] == 0
control["unemployed1974"] = control["earnings1974"] == 0
control["unemployed1975"] = control["earnings1975"] == 0
Explanation: We create two additional indicators umemployed1974 and umemployed1975 based on the earnings.
End of explanation
df = pd.concat([treated, control])
print(df.groupby("treatment").mean().T)
Explanation: The observational control group is significantly different from the experimental treatment group.
End of explanation
for name in df.columns[1:]:
  axes = df.hist(column=name, by=df["treatment"], figsize=(10, 2))
  axes[0].set_xlabel(name)
  axes[1].set_xlabel(name)
Explanation: The experimental treated subjects tend to be younger, with less education, black, single, without high school degree, and have smaller pre-intervention earnings.
End of explanation
formula = (
    "(age + education + black + hispanic + married + nodegree + earnings1974 +"
    " earnings1975 + unemployed1974 + unemployed1975)**2 + I(age**2) + "
    "I(education**2) - black:hispanic - education:nodegree - "
    "earnings1974:earnings1975 - earnings1974:unemployed1974 - "
    "earnings1975:unemployed1975"
)
treated2 = patsy.dmatrix(formula, treated, return_type="dataframe").iloc[:, 1:]
control2 = patsy.dmatrix(formula, control, return_type="dataframe").iloc[:, 1:]
Explanation: Empirical Calibration
Empirical calibration assigns weights to the observational control subjects such that the reweighted control group equates the treated group on moments of covariates. The weights are also sought to be as uniform as possible.
Prepare Balancing Constraints
There are 8 raw pre-intervention covariates. Simply balancing on their first moment might not be sufficient. We follow Hainmueller (2011)'s analysis and consider a total of 52 balancing constraints, which include single terms, selected one-way interactions, and selected squared terms. Note that balancing on all 52 transformed covariates might be excessive, we choose them simply to reproduce Hainmueller's results.
We use patsy's formula API to construct the covariate matrices, where ":" indicates interaction, and "-" means removes the term on the right.
End of explanation
for i, name in enumerate(control2.columns):
  print("%s: %s" % (i, name))
Explanation: The full list of the 52 balancing constraints is shown below.
End of explanation
weights, success = ec.calibrate(covariates=control2.values,
                                target_covariates=treated2.values)
success
Explanation: Entropy Balancing
We first use the default ENTROPY objective which mimizes the Kullback-Leibler divergence between the balancing weights and uniform weights.
End of explanation
weights, success = ec.calibrate(covariates=control2.values,
                                target_covariates=treated2.values,
                                l2_norm=1e-5)
success
Explanation: It turns out that we need to relax the equality constraint a tiny bit to ensure convergence.
End of explanation
pd.Series(weights).plot(style='o', figsize=(12, 3))
Explanation: Top 3 control observations account for 12% of total weight.
End of explanation
effective_control_size = 1.0 / np.square(weights).sum()
effective_control_size
Explanation: The effective sample size of the control group is merely 103.4, a huge reduction from the original sample size 15992.
End of explanation
pd.concat([treated.mean(axis=0),
           pd.Series(np.matmul(control.astype('float64').T, weights),
                     index=treated.columns)], axis=1)
Explanation: We confirm that the means of the raw covariates are matched after weighting.
End of explanation
diff = treated["earnings1978"].mean() - np.sum(control["earnings1978"] * weights)
diff
Explanation: The difference in means between the treated group and the reweighted control group yields an average treatment effect on the treated of $1571, which reproduces Hainmueller (2011)'s result.
End of explanation
sd = np.sqrt(treated["earnings1978"].var(ddof=1) *
             (1 / len(treated["earnings1978"]) + 1 / effective_control_size))
sd
Explanation: Under the homoskedastic assumption (i.e., equal unit-level conditional variance), we estimate the variance using formula prescribed in section 19.7 of Imbens and Rubin (2015).
End of explanation
(diff - 1.96 * sd, diff + 1.96 * sd)
Explanation: The 95% confidence interval is [55, 3088], wider than the reported [97, 3044] by Hainmueller (2011) (it is unclear what method was used there).
End of explanation
weights, success = ec.calibrate(
    covariates=control2.values,
    target_covariates=treated2.values,
    objective=ec.Objective.QUADRATIC,
    l2_norm=1e-2)
success
Explanation: Quadratic Balancing
We then use the QUADRATIC objective which minimizes the Euclidean distance between the balancing weights and the uniform weights.
End of explanation
effective_control_size = 1.0 / np.square(weights).sum()
effective_control_size
pd.Series(weights).plot(style='o', figsize=(12, 3))
Explanation: Due to being the explicit optimizing objective, the effective control sample size increases.
End of explanation
diff = treated["earnings1978"].mean() - np.sum(control["earnings1978"] * weights)
diff
Explanation: The point estimate $1713 is getting closer to the experimental benchmark $1794.
End of explanation
sd = np.sqrt(treated["earnings1978"].var(ddof=1) *
             (1 / len(treated["earnings1978"]) + 1 / effective_control_size))
sd
(diff - 1.96 * sd, diff + 1.96 * sd)
Explanation: The 95% confidence interval becomes narrower due to increased effective control sample size.
End of explanation
weights, success = ec.calibrate(
    covariates=control2.values,
    target_covariates=treated2.values,
    objective=ec.Objective.QUADRATIC,
    max_weight=0.015,
    l2_norm=1e-2)
success
Explanation: Imposing an Upper Bound on Weights
For both entropy balancing and quadratic balancing, one can optionally impose an upper bound on the weights to avoid extreme individual weights. For illustration, we impose an upper bound 0.015 on quadratic balancing weights, i.e., no single control individual accounts for 1.5% or more of total weights.
End of explanation
pd.Series(weights).plot(style='o', figsize=(12, 3))
Explanation: We can confirm the resulting weights are indeed capped at 0.015.
End of explanation
effective_control_size = 1.0 / np.square(weights).sum()
effective_control_size
Explanation: The cost of this additional constraint is a small loss in effective sample size.
End of explanation
pd.concat([treated.mean(axis=0),
           pd.Series(np.matmul(control.astype('float64').T, weights),
                     index=treated.columns)], axis=1)
Explanation: Note that this additional weights constraint is an integral part of the optimization, as opposed to a separate trimming step used with propensity score weighting type method. We don't sacrifice the balancing conditions.
End of explanation
diff = treated["earnings1978"].mean() - np.sum(control["earnings1978"] * weights)
diff
Explanation: The point estimate is very close to the experimental estimate $1794.
End of explanation
<END_TASK>