Possibly working Nickel bias correction

sylvia_nickell_bias.py

from linearmodels import PanelOLS
import numpy as np
import pandas as pd
from itertools import product
import seaborn as sns
import matplotlib.pyplot as plt

def generate_data(N, T, rho, sigma_a = 5, sigma_eps=1, burn_in = 100):
    # dumb for-loop to generate data, not efficient but easy to understand
    data = []
    for i in range(N):
        a = sigma_a * np.random.normal(0, 1)
        lag_y = 0  # initial value
        for t in range(T + burn_in):
            y = a + rho * lag_y + sigma_eps*np.random.normal(0, 1)

            if t >= burn_in:
                data.append({
                    'y': y,
                    'ylag': lag_y,
                    'entity': i,
                    'time': t - burn_in
                })
            
            lag_y = y

    df = pd.DataFrame(data)
    return df



def get_moments(alpha, df):
    df = df.copy()

    T = len(np.unique(df["time"]))

    df["e"] = df["y"] - alpha * df["ylag"]

    df["ebar"] = df.groupby("entity")["e"].transform("mean")
    df["ylagbar"] = df.groupby("entity")["ylag"].transform("mean")
    
    # Correction term
    b = -1/((1 - alpha) * T) * (1 - (1-alpha**T)/((1-alpha) * T))
    
    df["sigma_term"] = (df["e"] - df["ebar"]) * df["e"]
    sigma2 = 1/(T-1) * df.groupby("entity")["sigma_term"].transform("sum")
    moment = (df["ylag"] - df["ylagbar"]) * df["e"] -  b * sigma2

    return np.mean(moment)



def objective(alpha, df):
    moments = get_moments(alpha, df)
    return (np.mean(moments))**2


np.random.seed(1)

N = 1000
T = 5
rho = .3
df = generate_data(N, T, rho=rho)

alphas = np.linspace(0, .95, 101)
results = np.zeros_like(alphas)
for i, alpha in enumerate(alphas):
    results[i] = objective(alpha, df)

alphabc = alphas[np.argmin(results)]
print("Minimum via GMM")
print(alphabc)
plt.plot(alphas, results)

print("\n")
print("Compare to OLS")
model = PanelOLS.from_formula('y ~ ylag + EntityEffects', data=df.set_index(['entity', 'time']))
alphaols = model.fit().params.iloc[0]
alphaols