import numpy as np
import pandas as pd
from scipy.optimize import minimize
from scipy import stats
import matplotlib.pyplot as plt

rng = np.random.default_rng(7)          # lecture number
FAST_MODE = True

np.set_printoptions(precision=4, suppress=True)

def generate_iv(T=500, seed=77):
    '''Listing 4-04's data generating process. True intercept 2, true price coefficient -0.5.'''
    st = np.random.RandomState(seed)
    st.uniform(size=T)                                   # kept so the stream matches
    kappa_t = 2 * st.normal(size=T)
    zeta_t = st.normal(size=T)
    p_t = kappa_t + zeta_t
    pihat_t = np.exp(2. - 0.5 * p_t + kappa_t) / (1 + np.exp(2. - 0.5 * p_t + kappa_t))
    return pihat_t, p_t, zeta_t, kappa_t


def iv(y_t, x_t_k, z_t_k):
    '''Just-identified instrumental variables, listing 4-04.'''
    return np.linalg.solve(z_t_k.T @ x_t_k, y_t @ z_t_k)


T = 500
pihat_t, p_t, zeta_t, kappa_t = generate_iv(T)
U_t = np.log(pihat_t / (1 - pihat_t))                    # Berry inversion, binomial case
regressors_t_k = np.stack((np.ones(T), p_t), axis=1)
instruments_t_k = np.stack((np.ones(T), zeta_t), axis=1)

print(f"{'':>28} {'intercept':>10} {'price':>10}")
print(f"{'truth':>28} {2.0:>10.4f} {-0.5:>10.4f}")
print(f"{'OLS (price is endogenous)':>28} "
      + ' '.join(f'{v:>10.4f}' for v in iv(U_t, regressors_t_k, regressors_t_k)))
print(f"{'IV with the cost shifter':>28} "
      + ' '.join(f'{v:>10.4f}' for v in iv(U_t, regressors_t_k, instruments_t_k)))
print(f'\ncorr(price, unobserved quality)   = {np.corrcoef(p_t, kappa_t)[0,1]:.3f}')
print(f'first-stage R^2 of the instrument = {np.corrcoef(p_t, zeta_t)[0,1]**2:.3f}')

                              intercept      price
                       truth     2.0000    -0.5000
   OLS (price is endogenous)     1.9388     0.2975
    IV with the cost shifter     1.9983    -0.5064

corr(price, unobserved quality)   = 0.892
first-stage R^2 of the instrument = 0.201

def generate_gmm(T=500, seed=77):
    '''Listing 4-05: two cost shifters, so the model is over-identified.'''
    st = np.random.RandomState(seed)
    st.uniform(size=T)
    kappa_t = 2 * st.normal(size=T)
    zeta_t = st.normal(size=T)
    zetaprime_t = st.normal(size=T)
    p_t = kappa_t + 0.3 * zeta_t + 0.7 * zetaprime_t
    pihat_t = np.exp(2. - 0.5 * p_t + kappa_t) / (1 + np.exp(2. - 0.5 * p_t + kappa_t))
    return pihat_t, p_t, zeta_t, zetaprime_t


def lambda_gmm(dep_t, reg_t_k, ins_t_l, W_l_l=None):
    '''IV-GMM.

    Two departures from listing 4-05. It reads its dependent variable from the enclosing
    scope rather than from its own argument; here it uses dep_t, as the signature promises.
    And it forms the T x T projector Z W Z'; the algebraically identical form below is
    O(T L^2) rather than O(T^2), which matters as soon as T is large.
    '''
    if W_l_l is None:
        W_l_l = np.linalg.inv(ins_t_l.T @ ins_t_l)
    A = reg_t_k.T @ ins_t_l @ W_l_l
    return np.linalg.solve(A @ (ins_t_l.T @ reg_t_k), A @ (ins_t_l.T @ dep_t))


def estimate_gmm(T_, seed):
    pihat, p, zeta, zetap = generate_gmm(T_, seed)
    U = np.log(pihat / (1 - pihat))
    reg = np.stack((np.ones(T_), p), axis=1)
    ins = np.stack((np.ones(T_), zeta, zetap), axis=1)
    l1 = lambda_gmm(U, reg, ins)
    resid = U - reg @ l1
    g = resid[:, None] * ins
    W2 = T_ * np.linalg.inv(g.T @ g)
    return l1, lambda_gmm(U, reg, ins, W2), np.corrcoef(p, zeta + zetap)[0, 1] ** 2


lam1_k, lam2_k, r2_first = estimate_gmm(T, 77)
print(f"{'':>28} {'intercept':>10} {'price':>10}")
print(f"{'truth':>28} {2.0:>10.4f} {-0.5:>10.4f}")
print(f"{'IV-GMM, one step':>28} " + ' '.join(f'{v:>10.4f}' for v in lam1_k))
print(f"{'IV-GMM, two step':>28} " + ' '.join(f'{v:>10.4f}' for v in lam2_k))
print(f'\nfirst-stage R^2 of the two cost shifters : {r2_first:.3f}')

# is the estimator centred on the truth, or biased? repeat over seeds.
nrep = 400 if not FAST_MODE else 200
print(f"\n{'T':>8} {'mean':>9} {'s.d.':>9} {'s.e. of mean':>14} "
      f"{'|mean - truth|':>16}")
for T_ in (500, 5000):
    est = np.array([estimate_gmm(T_, 1000 + r)[1][1] for r in range(nrep)])
    se = est.std(ddof=1) / np.sqrt(nrep)
    print(f'{T_:>8} {est.mean():>9.4f} {est.std(ddof=1):>9.4f} {se:>14.4f} '
          f'{abs(est.mean() + 0.5):>16.4f}')
    if T_ == 5000:
        assert abs(est.mean() + 0.5) < 4 * se, 'estimator should be centred on the truth'

                              intercept      price
                       truth     2.0000    -0.5000
            IV-GMM, one step     1.9981    -0.3511
            IV-GMM, two step     1.9945    -0.3462

first-stage R^2 of the two cost shifters : 0.129

       T      mean      s.d.   s.e. of mean   |mean - truth|
     500   -0.5082    0.1220         0.0086           0.0082

    5000   -0.4961    0.0365         0.0026           0.0039

from pathlib import Path


def load_nevo():
    '''Load Nevo's cereal data locally, then try two public mirrors.'''
    for parent in (Path.cwd(), *Path.cwd().parents):
        root = parent / 'mec_datasets' / 'demand_nevo'
        products = root / 'nevo_products.csv'
        consumers = root / 'nevo_agents.csv'
        if products.is_file() and consumers.is_file():
            return pd.read_csv(products), pd.read_csv(consumers), str(root)

    roots = ['https://raw.githubusercontent.com/math-econ-code/mec_datasets/main/demand_nevo/',
             'https://raw.githubusercontent.com/jeffgortmaker/pyblp/master/pyblp/data/']
    for root in roots:
        try:
            prod = pd.read_csv(root + 'nevo_products.csv')
            agents = pd.read_csv(root + 'nevo_agents.csv')
            return prod, agents, root
        except Exception:
            continue
    raise RuntimeError(
        'could not load Nevo data: no local mec_datasets copy was found '
        'and both public mirrors failed'
    )


prod, agents, root = load_nevo()
prod = prod.sort_values(['market_ids', 'brand_ids']).reset_index(drop=True)
agents = agents.sort_values(['market_ids']).reset_index(drop=True)

T = prod.market_ids.nunique()
Y = prod.groupby('market_ids').size().iloc[0]
I = agents.groupby('market_ids').size().iloc[0]
print(f'source: {root}')
print(f'markets T = {T}, brands J = {Y}, simulated consumers per market I = {I}')

pihat_t_y = prod.shares.to_numpy().reshape(T, Y)
p_t_y = prod.prices.to_numpy().reshape(T, Y)
sugar_t_y = prod.sugar.to_numpy().reshape(T, Y)
pihat0_t = 1 - pihat_t_y.sum(axis=1)
print(f'outside-good share: min {pihat0_t.min():.3f}, mean {pihat0_t.mean():.3f}, max {pihat0_t.max():.3f}')

U_logit = np.log(pihat_t_y) - np.log(pihat0_t)[:, None]      # Berry inversion, logit case
print(f'\ninverted U: min {U_logit.min():.2f}, mean {U_logit.mean():.2f}, '
      f'max {U_logit.max():.2f}')

source: C:\Users\alfre\Dropbox\AGResearch\courses\m-e-c\mec_datasets\demand_nevo
markets T = 94, brands J = 24, simulated consumers per market I = 20
outside-good share: min 0.305, mean 0.524, max 0.815

inverted U: min -8.27, mean -3.85, max 0.20

y = U_logit.reshape(-1)
price = prod.prices.to_numpy()
Dbrand = pd.get_dummies(prod.brand_ids, prefix='b').to_numpy().astype(float)
instr = prod[[f'demand_instruments{k}' for k in range(20)]].to_numpy()
ones = np.ones((len(y), 1))
sugar = prod.sugar.to_numpy()[:, None]
mushy = prod.mushy.to_numpy()[:, None]


def tsls(y_, X_, Z_, W_=None):
    '''IV-GMM / two-stage least squares with weighting matrix W (default Z'Z inverse).'''
    if W_ is None:
        W_ = np.linalg.pinv(Z_.T @ Z_)
    A = X_.T @ Z_ @ W_ @ Z_.T
    return np.linalg.solve(A @ X_, A @ y_)


specs = [
    ('OLS, no controls',      np.column_stack([price[:, None], ones]), None),
    ('OLS, + sugar and mushy', np.column_stack([price[:, None], ones, sugar, mushy]), None),
    ('OLS, + brand dummies',  np.column_stack([price[:, None], Dbrand]), None),
    ('IV, no controls',       np.column_stack([price[:, None], ones]),
                              np.column_stack([instr, ones])),
    ('IV, + brand dummies',   np.column_stack([price[:, None], Dbrand]),
                              np.column_stack([instr, Dbrand])),
]
print(f"{'specification':>24} {'price coefficient':>19}")
for name, Xs, Zs in specs:
    b = np.linalg.lstsq(Xs, y, rcond=None)[0] if Zs is None else tsls(y, Xs, Zs)
    print(f'{name:>24} {b[0]:>19.3f}')

# first-stage strength of the twenty excluded instruments, given brand dummies
Zfull = np.column_stack([instr, Dbrand])
r_full = price - Zfull @ np.linalg.lstsq(Zfull, price, rcond=None)[0]
r_rest = price - Dbrand @ np.linalg.lstsq(Dbrand, price, rcond=None)[0]
Fstat = ((r_rest @ r_rest - r_full @ r_full) / 20) / (r_full @ r_full / (len(y) - Zfull.shape[1]))
print(f'\nfirst-stage F on the 20 excluded instruments (given brand dummies) = {Fstat:.0f}')

           specification   price coefficient
        OLS, no controls              -7.586
  OLS, + sugar and mushy             -10.120
    OLS, + brand dummies             -29.037
         IV, no controls              -8.686
     IV, + brand dummies             -30.186

first-stage F on the 20 excluded instruments (given brand dummies) = 3367

nodes = agents[['nodes0', 'nodes1', 'nodes2', 'nodes3']].to_numpy().reshape(T, I, 4)
income = agents.income.to_numpy().reshape(T, I)
weights = agents.weights.to_numpy().reshape(T, I)
nu_p, nu_s = nodes[:, :, 1], nodes[:, :, 2]              # price and sugar draws
print('agent weights sum to one in every market :',
      bool(np.allclose(weights.sum(axis=1), 1.0)))


def nu_of(tau):
    '''Heterogeneous part of utility: tau = (tau_1, tau_2, tau_3).'''
    t1_, t2_, t3_ = tau
    return (p_t_y[:, None, :] * (t1_ * nu_p[:, :, None] + t3_ * income[:, :, None])
            + sugar_t_y[:, None, :] * (t2_ * nu_s[:, :, None]))


def choice_probs(U_t_y, nu_t_i_y):
    a = U_t_y[:, None, :] + nu_t_i_y
    m = np.maximum(a.max(axis=2), 0.0)                   # log-sum-exp guard, outside good = 0
    e = np.exp(a - m[:, :, None])
    return e / (np.exp(-m) + e.sum(axis=2))[:, :, None]


def model_shares(U_t_y, nu_t_i_y):
    return (choice_probs(U_t_y, nu_t_i_y) * weights[:, :, None]).sum(axis=1)


def blp_step(U_t_y, nu_t_i_y):
    '''One BLP contraction step.'''
    return U_t_y + np.log(pihat_t_y) - np.log(model_shares(U_t_y, nu_t_i_y))


def ipfp_step(U_t_y, nu_t_i_y):
    '''One IPFP / Sinkhorn update of the product-side potential, written independently.'''
    a = U_t_y[:, None, :] + nu_t_i_y
    m = np.maximum(a.max(axis=2), 0.0)
    D = np.exp(-m) + np.exp(a - m[:, :, None]).sum(axis=2)
    denom = (np.exp(nu_t_i_y - m[:, :, None]) / D[:, :, None]
             * weights[:, :, None]).sum(axis=1)
    return np.log(pihat_t_y) - np.log(denom)


tau_probe = np.array([1.0, 0.02, 5.0])
nu_probe = nu_of(tau_probe)
for shift in (0.0, 0.3, -0.7):
    d = U_logit + shift
    gap = np.abs(blp_step(d, nu_probe) - ipfp_step(d, nu_probe)).max()
    print(f'U shifted by {shift:>5}:  max |BLP step - IPFP step| = {gap:.2e}')
    assert gap < 1e-12

agent weights sum to one in every market : True
U shifted by   0.0:  max |BLP step - IPFP step| = 3.55e-15
U shifted by   0.3:  max |BLP step - IPFP step| = 1.78e-15
U shifted by  -0.7:  max |BLP step - IPFP step| = 1.78e-15

def contraction(nu_t_i_y, U0, tol=1e-13, maxit=10_000):
    '''BLP's inner loop. Returns U and the iteration count.'''
    d = U0.copy()
    for it in range(maxit):
        dn = blp_step(d, nu_t_i_y)
        if np.abs(dn - d).max() < tol:
            return dn, it + 1
        d = dn
    raise RuntimeError('contraction did not converge')


U_star, nit = contraction(nu_probe, U_logit)
resid = np.abs(model_shares(U_star, nu_probe) - pihat_t_y).max()
print(f'contraction: {nit} iterations to a tolerance of 1e-13')
print(f'share residual at the fixed point : {resid:.2e}')

# the modulus: the contraction factor observed along the path
d, gaps = U_logit.copy(), []
for _ in range(40):
    dn = blp_step(d, nu_probe)
    gaps.append(np.abs(dn - d).max())
    d = dn
ratios = np.array(gaps[1:]) / np.array(gaps[:-1])
print(f'observed contraction factor (last 10 steps) : {ratios[-10:].mean():.4f}')
print('a factor below one is what Franklin and Lorenz guarantee for IPFP.')
assert resid < 1e-12

contraction: 82 iterations to a tolerance of 1e-13
share residual at the fixed point : 1.16e-14
observed contraction factor (last 10 steps) : 0.7087
a factor below one is what Franklin and Lorenz guarantee for IPFP.

Phi_ty_k = np.column_stack([price[:, None], Dbrand])
zeta_ty_d = np.column_stack([instr, Dbrand])
nobs = len(y)
W_one = np.linalg.pinv(zeta_ty_d.T @ zeta_ty_d / nobs)


def gmm_objective(tau, W, full=False):
    U, it = contraction(nu_of(tau), U_logit)
    dv = U.reshape(-1)
    lambda_k = tsls(dv, Phi_ty_k, zeta_ty_d, W)
    kappa = dv - Phi_ty_k @ lambda_k
    gbar = zeta_ty_d.T @ kappa / nobs
    value = float(gbar @ W @ gbar)
    return (value, lambda_k, kappa, U, it) if full else value


starts = [[0.5, 0.0, 5.0], [1.5, 0.05, 10.0], [3.0, -0.05, 15.0]]
if FAST_MODE:
    starts = starts[:2]
print(f"{'start':>22} {'tau_1':>9} {'tau_2':>9} {'tau_3':>9} {'objective':>13} {'evals':>7}")
best = None
for st in starts:
    r = minimize(gmm_objective, np.array(st), args=(W_one,), method='Nelder-Mead',
                 options={'xatol': 1e-6, 'fatol': 1e-14, 'maxfev': 1200})
    print(f'{str(st):>22} ' + ' '.join(f'{v:>9.4f}' for v in r.x)
          + f' {r.fun:>13.6e} {r.nfev:>7}')
    if best is None or r.fun < best.fun:
        best = r
print('\nall starts reach the same optimum: the outer objective is well behaved here,')
print('which is not guaranteed in general and is worth checking every time.')

                 start     tau_1     tau_2     tau_3     objective   evals

       [0.5, 0.0, 5.0]    1.2833   -0.0153   12.1525  4.250557e-02     246

     [1.5, 0.05, 10.0]    1.2833   -0.0153   12.1525  4.250557e-02     246

all starts reach the same optimum: the outer objective is well behaved here,
which is not guaranteed in general and is worth checking every time.

# efficient weighting matrix from the first-step residuals, then re-optimize
v1, lambda1_k, kappa1, d1, it1 = gmm_objective(best.x, W_one, full=True)
gmat = kappa1[:, None] * zeta_ty_d
W_two = np.linalg.pinv(gmat.T @ gmat / nobs)
r2 = minimize(gmm_objective, best.x, args=(W_two,), method='Nelder-Mead',
              options={'xatol': 1e-6, 'fatol': 1e-14, 'maxfev': 1200})
v2, lambda2_k, kappa2, d2, it2 = gmm_objective(r2.x, W_two, full=True)

# GMM standard errors on the linear parameters, conditional on tau
G = -(zeta_ty_d.T @ Phi_ty_k) / nobs
bread = np.linalg.pinv(G.T @ W_two @ G)
Vb = bread @ (G.T @ W_two @ np.linalg.pinv(W_two) @ W_two @ G) @ bread / nobs

print(f"{'':>26} {'estimate':>10} {'s.e.':>9}")
print(f"{'lambda_p (price)':>26} {lambda2_k[0]:>10.3f} {np.sqrt(Vb[0,0]):>9.3f}")
names = ['tau_1', 'tau_2', 'tau_price x income']
for k, nm in enumerate(names):
    print(f'{nm:>26} {r2.x[k]:>10.4f}')
print(f'\nGMM objective, two-step : {v2:.6e}   contraction iterations at the optimum: {it2}')
print(f'moment norm |Z^T kappa / n| : {np.abs(zeta_ty_d.T @ kappa2 / nobs).max():.2e}')
print('\nnote tau_2 is estimated near zero, and its sign is not identified')
print('(only |tau| enters the model), so it should be read as "no detectable')
print('unobserved heterogeneity in the taste for sugar, in this specification".')

                             estimate      s.e.
          lambda_p (price)    -36.026     0.930
                     tau_1     2.3067
                     tau_2    -0.0158
        tau_price x income    12.0097

GMM objective, two-step : 4.912201e-02   contraction iterations at the optimum: 98
moment norm |Z^T kappa / n| : 1.04e-01

note tau_2 is estimated near zero, and its sign is not identified
(only |tau| enters the model), so it should be read as "no detectable
unobserved heterogeneity in the taste for sugar, in this specification".

try:
    import pyblp
    pyblp.options.verbose = False
    HAVE_PYBLP = True
except ImportError:
    HAVE_PYBLP = False
    print('pyblp is not installed, so the external check is skipped.')
    print('Install it with `pip install pyblp` to run this comparison.')

if HAVE_PYBLP:
    # relabel the draws so pyblp's nodes0, nodes1 are this notebook's price and sugar draws
    agents_pyblp = pd.DataFrame({
        'market_ids': agents.market_ids, 'weights': agents.weights,
        'nodes0': agents.nodes1, 'nodes1': agents.nodes2, 'income': agents.income})

    problem = pyblp.Problem(
        (pyblp.Formulation('0 + prices', absorb='C(brand_ids)'),
         pyblp.Formulation('0 + prices + sugar')),
        prod, agent_formulation=pyblp.Formulation('0 + income'), agent_data=agents_pyblp)

    ref = problem.solve(sigma=np.diag([0.5, 0.001]),        # zeros are held fixed
                        pi=np.array([[5.0], [0.0]]),
                        method='2s',
                        optimization=pyblp.Optimization('l-bfgs-b', {'gtol': 1e-10}))

    mine = {'lambda_p (price)': lambda2_k[0], 'tau_1': abs(r2.x[0]),
            'tau_3': r2.x[2], 'tau_2': abs(r2.x[1])}
    theirs = {'lambda_p (price)': ref.beta[0, 0], 'tau_1': abs(ref.sigma[0, 0]),
              'tau_3': ref.pi[0, 0], 'tau_2': abs(ref.sigma[1, 1])}
    ses = {'lambda_p (price)': ref.beta_se[0, 0], 'tau_1': ref.sigma_se[0, 0],
           'tau_3': ref.pi_se[0, 0], 'tau_2': ref.sigma_se[1, 1]}

    print(f"{'parameter':>14} {'this notebook':>15} {'pyblp':>12} {'difference':>12} "
          f"{'s.e.':>9} {'diff / s.e.':>12}")
    for k in mine:
        d = abs(mine[k] - theirs[k])
        print(f'{k:>14} {mine[k]:>15.4f} {theirs[k]:>12.4f} {d:>12.4f} '
              f'{ses[k]:>9.4f} {d/ses[k]:>12.2f}')

    print(f'\npyblp version {pyblp.__version__}; '
          f'{ref.cumulative_objective_evaluations} objective evaluations '
          f'with analytic gradients')
    print(f'smallest eigenvalue of its reduced Hessian: '
          f'{np.linalg.eigvalsh(ref.reduced_hessian).min():.3e}')

    for k in ('lambda_p (price)', 'tau_1', 'tau_3'):
        assert abs(mine[k] - theirs[k]) < 0.5 * ses[k], f'{k} disagrees with pyblp'
    assert max(mine['tau_2'], theirs['tau_2']) < 0.05, \
        'both implementations should put tau_2 at essentially zero'

pyblp is not installed, so the external check is skipped.
Install it with `pip install pyblp` to run this comparison.

f_t_i_y = choice_probs(d2, nu_of(r2.x))
lambdap_t_i = lambda2_k[0] + r2.x[0] * nu_p + r2.x[2] * income
pi_fit = (f_t_i_y * weights[:, :, None]).sum(axis=1)

own = ((lambdap_t_i[:, :, None] * f_t_i_y * (1 - f_t_i_y) * weights[:, :, None]).sum(axis=1)
       * p_t_y / pi_fit)
own_logit = lambda2_k[0] * p_t_y * (1 - pi_fit)

print(f'own-price elasticity, random coefficients : mean {own.mean():>7.3f}, '
      f'median {np.median(own):>7.3f}')
print(f'own-price elasticity, plain logit         : mean {own_logit.mean():>7.3f}, '
      f'median {np.median(own_logit):>7.3f}')
print(f'products with inelastic demand (> -1)     : {100*(own > -1).mean():.1f}%')

own-price elasticity, random coefficients : mean  -3.717, median  -3.678
own-price elasticity, plain logit         : mean  -4.444, median  -4.374
products with inelastic demand (> -1)     : 0.0%

# Diversion ratios, pooled over every market and every ordered pair of brands.
# We compare the RC-logit diversion with the plain-logit ratio s_k/(1-s_j) evaluated at
# the SAME fitted shares, so the difference isolates the departure from IIA.
rows = []
for t in range(T):
    dsdp = -np.einsum('i,ij,ik->jk', weights[t] * lambdap_t_i[t], f_t_i_y[t], f_t_i_y[t])
    np.fill_diagonal(dsdp, (weights[t] * lambdap_t_i[t] * f_t_i_y[t].T
                            * (1 - f_t_i_y[t].T)).sum(axis=1))
    d_rc = dsdp / np.abs(np.diag(dsdp))[None, :]
    d_lg = np.tile(pi_fit[t][:, None], (1, Y)) / (1 - pi_fit[t])[None, :]
    for j in range(Y):
        for k in range(Y):
            if j != k:
                rows.append((d_rc[k, j], d_lg[k, j], p_t_y[t, j], p_t_y[t, k],
                             abs(sugar_t_y[t, k] - sugar_t_y[t, j])))
A = np.array(rows)
d_rc, d_lg, price_src, price_dst, sugar_gap = A.T
ratio = d_rc / d_lg

print(f'ordered pairs pooled over all markets : {len(A):,}')
print(f'mean plain-logit diversion            : {d_lg.mean():.5f}')
print(f'mean |RC - logit| diversion           : {np.abs(d_rc - d_lg).mean():.5f}')
print(f'mean |RC / logit - 1|                 : {np.abs(ratio - 1).mean():.3f}'
      '   <- the departure from IIA, in relative terms')

print(f"\n{'correlation of RC/logit - 1 with':>34} {'':>8}")
for nm, v in [('price of the brand losing sales', price_src),
              ('price of the brand gaining sales', price_dst),
              ('|price gap| between the two', np.abs(price_dst - price_src)),
              ('|sugar gap| between the two', sugar_gap)]:
    print(f'{nm:>34} {np.corrcoef(ratio - 1, v)[0,1]:>+8.3f}')

q = np.quantile((price_src + price_dst) / 2, [0, .25, .5, .75, 1.])
print(f"\n{'mean price of the pair':>26} {'mean RC / logit diversion':>27}")
for a_, b_ in zip(q[:-1], q[1:]):
    m = ((price_src + price_dst) / 2 >= a_) & ((price_src + price_dst) / 2 <= b_)
    print(f'   [{a_:.3f}, {b_:.3f}]{"":>10} {ratio[m].mean():>21.3f}')
assert np.abs(ratio - 1).mean() > 0.05

ordered pairs pooled over all markets : 51,888
mean plain-logit diversion            : 0.02023
mean |RC - logit| diversion           : 0.00323
mean |RC / logit - 1|                 : 0.190   <- the departure from IIA, in relative terms

  correlation of RC/logit - 1 with         
   price of the brand losing sales   +0.600
  price of the brand gaining sales   +0.226
       |price gap| between the two   -0.088
       |sugar gap| between the two   -0.031

    mean price of the pair   mean RC / logit diversion
   [0.057, 0.112]                           1.074
   [0.112, 0.125]                           1.135
   [0.125, 0.139]                           1.192
   [0.139, 0.212]                           1.328

fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))

sig_grid = np.linspace(0.0, 4.0, 17 if not FAST_MODE else 13)
prof = [gmm_objective(np.array([sg, r2.x[1], r2.x[2]]), W_two) for sg in sig_grid]
ax[0].plot(sig_grid, prof, 'o-', color='C0', ms=4)
ax[0].axvline(r2.x[0], color='C3', ls='--', lw=1.2,
              label=rf'$\hat\tau_1={r2.x[0]:.2f}$')
ax[0].set_xlabel(r'$\tau_1$'); ax[0].set_ylabel('GMM objective')
ax[0].set_title('Outer objective, profiled')
ax[0].legend(fontsize=9); ax[0].grid(alpha=.25)

pair_price = (price_src + price_dst) / 2
edges = np.quantile(pair_price, np.linspace(0, 1, 11))
ctr = 0.5 * (edges[:-1] + edges[1:])
means = [ratio[(pair_price >= a_) & (pair_price <= b_)].mean()
         for a_, b_ in zip(edges[:-1], edges[1:])]
ax[1].plot(ctr, means, 'o-', color='C0', ms=5)
ax[1].axhline(1.0, color='0.4', ls='--', lw=1.2, label='plain logit (IIA)')
ax[1].set_xlabel('mean price of the pair of brands')
ax[1].set_ylabel('RC diversion / plain-logit diversion')
ax[1].set_title('The departure from IIA rises with price')
ax[1].legend(fontsize=9); ax[1].grid(alpha=.25)

plt.tight_layout(); plt.show()

Endogeneity, IV-GMM, and BLP¶

Alfred Galichon (NYU)¶

'math+econ+code' masterclass series: discrete choice¶

With python code examples¶

Learning objectives¶

References¶

Motivation: the analyst does not see everything the consumer sees¶

1. Endogeneity in a controlled experiment¶

2. Nevo's cereal data, and the logit benchmark¶

3. Random coefficients, and BLP's contraction¶

The contraction is an IPFP half-step¶

4. The nested fixed point estimator¶

Validation against a reference implementation¶

Reading the estimates¶

5. What the random coefficients buy¶

Summing up¶

Exercises¶

Run notes¶