# Partialling-Out Demonstration
# Module 3 - The Partialling-Out Interpretation
# Dataset: WAGE1

# Partialling-Out Demonstration
#
# Recover the education coefficient by first removing the part of education explained by experience and tenure.

# %% Cell 2
from pathlib import Path

DATASET = "WAGE1.DTA"
VARIABLES = ["lwage", "educ", "exper", "tenure"]
DATA_FOLDER = "https://drive.google.com/drive/folders/1_STdcydIcst-opcbwOKRFzUXsgxQgBoS?usp=sharing"

if not Path(DATASET).exists():
    raise FileNotFoundError(
        "Dataset file not installed yet\n"
        f"Dataset: WAGE1\n"
        f"Variables needed: {', '.join(VARIABLES)}\n"
        f"Course data folder: {DATA_FOLDER}\n"
        "Admin upload instruction: upload the dataset in Admin -> Datasets and make it available to this notebook."
    )

# %% Cell 3
import pandas as pd
import numpy as np

df = pd.read_stata(DATASET)[VARIABLES].dropna()

def residualize(series, controls):
    X = np.column_stack([np.ones(len(df)), df[controls].to_numpy()])
    beta = np.linalg.lstsq(X, series.to_numpy(), rcond=None)[0]
    return series.to_numpy() - X @ beta

educ_leftover = residualize(df["educ"], ["exper", "tenure"])
X_partial = np.column_stack([np.ones(len(df)), educ_leftover])
partial_beta = np.linalg.lstsq(X_partial, df["lwage"].to_numpy(), rcond=None)[0][1]

X_full = np.column_stack([np.ones(len(df)), df[["educ", "exper", "tenure"]].to_numpy()])
full_beta = np.linalg.lstsq(X_full, df["lwage"].to_numpy(), rcond=None)[0][1]

print("Partialling-out coefficient:", round(partial_beta, 4))
print("Full model education coefficient:", round(full_beta, 4))

# Interpretation
#
# The partialling-out coefficient uses only the part of education that is not explained by the controls.