Matching V_EE estimates from Sorkin estimator back to firm IDs about pytwoway HOT 10 CLOSED

import numpy as np
import pandas as pd
import bipartitepandas as bpd
import pytwoway as tw

n_workers = 750
n_firms = 5
n_periods = 2
## Firm and worker effects ##
psi = np.random.normal(size=n_firms)
exp_psi = np.exp(psi)
exp_psi /= exp_psi.sum()
psi = np.log(exp_psi)
## Simulate data ##
i_sim = np.repeat(np.arange(n_workers), n_periods)
t_sim = np.tile(np.arange(n_periods), n_workers)
j_sim = np.zeros((n_workers, n_periods), dtype=int)
y_sim = np.zeros(n_workers * n_periods)

## Simulate firms ##
for i in range(n_workers):
    # Simulate first firm randomly
    j1 = np.random.choice(range(n_firms))
    j_sim[i, 0] = j1
    j_prev = j1
    prev_match_draw = psi[j1] + np.random.gumbel()
    for t in range(1, n_periods):
        # Now worker moves to next firm with probability exp(jt) / (exp(jt) + exp(j1))
        jt = np.random.choice(list(range(j_prev)) + list(range(j_prev + 1, n_firms)))
        match_draw = psi[jt] + np.random.gumbel()
        if np.random.uniform() < exp_psi[jt] / (exp_psi[jt] + exp_psi[j_prev]):
#         if match_draw >  prev_match_draw:
            # Move!
            j_sim[i, t] = jt
            j_prev = jt
            prev_match_draw = match_draw
        else:
            j_sim[i, t] = j_prev

j_sim = j_sim.flatten()

## Construct BipartitePandas DataFrame ##
default = True
if default:
    clean_params = bpd.clean_params({'connectedness': 'strongly_connected', 'verbose': False})
    sim_data = bpd.BipartiteDataFrame(i=i_sim, j=j_sim, y=y_sim, t=t_sim, track_id_changes=True).clean(clean_params)
    sim_data = sim_data.collapse(is_sorted=True, copy=False)
    sim_es = sim_data.to_eventstudy(is_sorted=True, copy=False)
else:
    ## Construct BipartitePandas DataFrame ##
    clean_params = bpd.clean_params({'connectedness': 'strongly_connected', 'drop_returns': 'returns', 'verbose': False})
    sim_data = bpd.BipartiteDataFrame(i=i_sim, j=j_sim, y=y_sim, t=t_sim, track_id_changes=True).clean(clean_params)
    sim_data = sim_data.collapse(is_sorted=True, copy=False)
    sim_es = sim_data.to_permutedeventstudy(is_sorted=True, copy=False)
sim_es['w'] = sim_es['w2'] / sim_es['w1']

sorkin_sim = tw.SorkinEstimator()
sorkin_sim.fit(sim_es)

M0 = sim_es.groupby(['j1', 'j2'])['w'].sum().unstack(fill_value=0).to_numpy().T
S0inv = np.diag(1 / M0.sum(axis=0))
print(np.linalg.eig(S0inv @ M0)[0][0])
evec = np.real(np.linalg.eig(S0inv @ M0)[1][:, 0])
evec /= evec.sum()
np.log(evec)

print(np.corrcoef(sorkin_sim.V_EE, psi)[0, 1])
print(np.corrcoef(np.log(evec), psi)[0, 1])

dataraw=Data.get('i j year  wages ', valuelabel=True,missingval=np.nan)
df=pd.DataFrame(dataraw,columns=['i', 'j',  't', 'y'])

bdf=bpd.BipartiteDataFrame(df, track_id_changes=True)
bdf=bdf.clean(clean_params)
bdf=bdf.collapse(is_sorted=True, copy=False)

fe_estimator=tw.FEEstimator(bdf,fe_params)
fe_estimator.fit()
fe_estimator.summary

bdf['akm']=bdf['psi_hat'].astype('float')
bdf=bdf.add_column(col_name='akm', dtype='float', how_collapse='mean', long_es_split=True)

bdf=bdf.to_eventstudy(is_sorted=True, copy=False)

akm_series=bdf.groupby('j1').first()['akm1']

sorkin_estimator = tw.SorkinEstimator()
sorkin_estimator.fit(bdf)

np.corrcoef(akm_series, sorkin_estimator.V_EE)[0,1]

bdf2 = bdf.to_long()
bdf2['v_ee_values'] = sorkin_estimator.V_EE[bdf2['j']]
df_export = bdf2.groupby('j')['v_ee_values'].first().reset_index()