doc/examples/py_double_ml_multiway_cluster.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python: Cluster Robust Double Machine Learning\n",
    "\n",
    "This example demonstrates the cluster robust features of the [DoubleML](https://docs.doubleml.org/stable/index.html) package.\n",
    "The theoretical foundations are given in Chiang et al. 2021."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import ListedColormap\n",
    "import seaborn as sns\n",
    "\n",
    "from sklearn.model_selection import KFold, RepeatedKFold\n",
    "from sklearn.base import clone\n",
    "\n",
    "from sklearn.linear_model import LassoCV\n",
    "\n",
    "from doubleml import DoubleMLClusterData, DoubleMLData, DoubleMLPLIV\n",
    "\n",
    "from doubleml.datasets import make_pliv_multiway_cluster_CKMS2021"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two-Way Cluster Robust DML\n",
    "\n",
    "In a first part, we show how the two-way cluster roboust DML (Chiang et al. 2021) can be implemented with the [DoubleML](https://docs.doubleml.org/stable/index.html) package.\n",
    "Chiang et al. (2021) consider double-indexed data\n",
    "\n",
    "\\begin{equation}\n",
    "\\lbrace W_{ij}: i \\in \\lbrace 1, \\ldots, N \\rbrace, j \\in \\lbrace 1, \\ldots, M \\rbrace \\rbrace\n",
    "\\end{equation}\n",
    "\n",
    "and the partially linear IV regression model (PLIV)\n",
    "\n",
    "$$\\begin{aligned}\n",
    "Y_{ij} = D_{ij} \\theta_0 +  g_0(X_{ij}) + \\epsilon_{ij}, & &\\mathbb{E}(\\epsilon_{ij} | X_{ij}, Z_{ij}) = 0, \\\\\n",
    "Z_{ij} = m_0(X_{ij}) + v_{ij}, & &\\mathbb{E}(v_{ij} | X_{ij}) = 0.\n",
    "\\end{aligned}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulate multiway cluster data\n",
    "\n",
    "We use the PLIV data generating process described in Section 4.1 of Chiang et al. (2020).\n",
    "The DGP is defined as\n",
    "$$\\begin{aligned}\n",
    "Z_{ij} &= X_{ij}' \\xi_0 + V_{ij}, \\\\\n",
    "D_{ij} &= Z_{ij}' \\pi_{10} + X_{ij}' \\pi_{20} + v_{ij}, \\\\\n",
    "Y_{ij} &= D_{ij} \\theta + X_{ij}' \\zeta_0 + \\varepsilon_{ij},\n",
    "\\end{aligned}$$\n",
    "with\n",
    "$$\\begin{aligned}\n",
    "X_{ij} &= (1 - \\omega_1^X - \\omega_2^X) \\alpha_{ij}^X\n",
    "+ \\omega_1^X \\alpha_{i}^X + \\omega_2^X \\alpha_{j}^X, \\\\\n",
    "\\varepsilon_{ij} &= (1 - \\omega_1^\\varepsilon - \\omega_2^\\varepsilon) \\alpha_{ij}^\\varepsilon\n",
    "+ \\omega_1^\\varepsilon \\alpha_{i}^\\varepsilon + \\omega_2^\\varepsilon \\alpha_{j}^\\varepsilon, \\\\\n",
    "v_{ij} &= (1 - \\omega_1^v - \\omega_2^v) \\alpha_{ij}^v\n",
    "+ \\omega_1^v \\alpha_{i}^v + \\omega_2^v \\alpha_{j}^v, \\\\\n",
    "V_{ij} &= (1 - \\omega_1^V - \\omega_2^V) \\alpha_{ij}^V\n",
    "+ \\omega_1^V \\alpha_{i}^V + \\omega_2^V \\alpha_{j}^V,\n",
    "\\end{aligned}$$\n",
    "and $\\alpha_{ij}^X, \\alpha_{i}^X, \\alpha_{j}^X \\sim \\mathcal{N}(0, \\Sigma)$\n",
    "where  $\\Sigma$ is a $p_x \\times p_x$ matrix with entries\n",
    "$\\Sigma_{kj} = s_X^{|j-k|}$.\n",
    "Further\n",
    "$$\\begin{aligned}\n",
    "\\left(\\begin{matrix} \\alpha_{ij}^\\varepsilon \\\\ \\alpha_{ij}^v \\end{matrix}\\right),\n",
    "\\left(\\begin{matrix} \\alpha_{i}^\\varepsilon \\\\ \\alpha_{i}^v \\end{matrix}\\right),\n",
    "\\left(\\begin{matrix} \\alpha_{j}^\\varepsilon \\\\ \\alpha_{j}^v \\end{matrix}\\right)\n",
    "\\sim \\mathcal{N}\\left(0, \\left(\\begin{matrix} 1 & s_{\\varepsilon v} \\\\\n",
    "s_{\\varepsilon v} & 1 \\end{matrix} \\right) \\right)\n",
    "\\end{aligned}$$\n",
    "and $\\alpha_{ij}^V, \\alpha_{i}^V, \\alpha_{j}^V \\sim \\mathcal{N}(0, 1)$.\n",
    "\n",
    "Data from this DGP can be generated with the [make_pliv_multiway_cluster_CKMS2021()](https://docs.doubleml.org/stable/api/generated/doubleml.datasets.make_pliv_multiway_cluster_CKMS2021.html#doubleml.datasets.make_pliv_multiway_cluster_CKMS2021) function from [DoubleML](https://docs.doubleml.org/stable/index.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the simulation parameters\n",
    "N = 25  # number of observations (first dimension)\n",
    "M = 25  # number of observations (second dimension)\n",
    "dim_X = 100  # dimension of X\n",
    "np.random.seed(3141) # set seed\n",
    "\n",
    "obj_dml_data = make_pliv_multiway_cluster_CKMS2021(N, M, dim_X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data-Backend for Cluster Data\n",
    "The implemenation of cluster robust double machine learning is based on a special data-backend called [DoubleMLClusterData](https://docs.doubleml.org/stable/api/generated/doubleml.DoubleMLClusterData.html#doubleml.DoubleMLClusterData)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The simulated data is of type DoubleMLClusterData\n",
    "print(obj_dml_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The cluster variables are part of the DataFrame\n",
    "obj_dml_data.data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize the objects of class  DoubleMLPLIV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set machine learning methods for m, g & r\n",
    "learner = LassoCV()\n",
    "ml_g = clone(learner)\n",
    "ml_m = clone(learner)\n",
    "ml_r = clone(learner)\n",
    "\n",
    "# initialize the DoubleMLPLIV object\n",
    "dml_pliv_obj = DoubleMLPLIV(obj_dml_data,\n",
    "                            ml_g, ml_m, ml_r,\n",
    "                            n_folds=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dml_pliv_obj)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cluster Robust Cross Fitting\n",
    "A key element of cluster robust DML (Chiang et al. 2021) is a special sample splitting used for the cross-fitting.\n",
    "In case of two-way clustering, we assume $N$ clusters in the first dimension and $M$ clusters in the second dimension.\n",
    "\n",
    "For $K$-fold cross-fitting, Chiang et al. 2021 proposed to randomly partition $[N]:=\\{1,\\ldots,N\\}$ into $K$ subsets $\\{I_1, \\ldots, I_K\\}$ and $[M]:=\\{1,\\ldots,N\\}$ into $K$ subsets $\\{J_1, \\ldots, J_K\\}$.\n",
    "Effectively, one then considers $K^2$ folds.\n",
    "Basically for each $(k, \\ell) \\in \\{1, \\ldots, K\\} \\times \\{1, \\ldots, K\\}$, the nuisance functions are estimated for all double-indexed observations in $([N]\\setminus I_K) \\times ([M]\\setminus J_\\ell)$, i.e.,\n",
    "$$\n",
    "\\hat{\\eta}_{k\\ell} = \\hat{\\eta}\\left((W_{ij})_{(i,j)\\in ([N]\\setminus I_K) \\times ([M]\\setminus J_\\ell)}\\right)\n",
    "$$\n",
    "The causal parameter is then estimated as usual by solving a moment condition with a Neyman orthogonal score function.\n",
    "For two-way cluster robust double machine learning with algorithm DML2 this results in solving\n",
    "$$\n",
    "\\frac{1}{K^2} \\sum_{k=1}^{K} \\sum_{\\ell=1}^{K} \\frac{1}{|I_k| |J_\\ell|} \\sum_{(i,j) \\in I_K \\times J_\\ell}\n",
    "\\psi(W_{ij}, \\tilde{\\theta}_0, \\hat{\\eta}_{k\\ell}) = 0\n",
    "$$\n",
    "for $\\tilde{\\theta}_0$.\n",
    "Here $|I_k|$ denotes the cardinality, i.e., the number of clusters in the $k$-th fold for the first cluster variable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the sample splitting of the $N \\cdot M = 625$ observations into $K \\cdot K = 9$ folds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#discrete color scheme\n",
    "x = sns.color_palette(\"RdBu_r\", 7)\n",
    "cMap = ListedColormap([x[0], x[3], x[6]])\n",
    "plt.rcParams['figure.figsize'] = 15, 12\n",
    "sns.set(font_scale=1.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "smpls = dml_pliv_obj.smpls[0]\n",
    "n_folds_per_cluster = dml_pliv_obj._n_folds_per_cluster\n",
    "df = pd.DataFrame(np.zeros([N*M, n_folds_per_cluster*n_folds_per_cluster]))\n",
    "for i_split, this_split_ind in enumerate(smpls):\n",
    "    df.loc[this_split_ind[0], i_split] = -1.\n",
    "    df.loc[this_split_ind[1], i_split] = 1.\n",
    "\n",
    "ax = sns.heatmap(df, cmap=cMap);\n",
    "ax.invert_yaxis();\n",
    "ax.set_ylim([0, N*M]);\n",
    "ax.set_xlabel('Fold')\n",
    "ax.set_ylabel('Observation')\n",
    "colorbar = ax.collections[0].colorbar\n",
    "colorbar.set_ticks([-0.667, 0, 0.667])\n",
    "colorbar.set_ticklabels(['Nuisance', '', 'Score'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we visualize the sample splitting in terms of the cluster variables the sample splitting into $9$ folds $I_k \\times J_\\ell$ becomes clearer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "smpls_cluster = dml_pliv_obj.smpls_cluster[0]\n",
    "n_folds_per_cluster = dml_pliv_obj._n_folds_per_cluster\n",
    "for i_split in range(len(smpls_cluster)):\n",
    "    plt.subplot(n_folds_per_cluster, n_folds_per_cluster, i_split + 1)\n",
    "    df = pd.DataFrame(np.zeros([N*M, 1]),\n",
    "                  index = pd.MultiIndex.from_product([range(N), range(M)]),\n",
    "                  columns=['value'])\n",
    "\n",
    "    df.loc[pd.MultiIndex.from_product(smpls_cluster[i_split][0]), :] = -1.\n",
    "    df.loc[pd.MultiIndex.from_product(smpls_cluster[i_split][1]), :] = 1.\n",
    "\n",
    "    df_wide = df.reset_index().pivot(index=\"level_0\", columns=\"level_1\", values=\"value\")\n",
    "    df_wide.index.name=''\n",
    "    df_wide.columns.name=''\n",
    "\n",
    "    ax = sns.heatmap(df_wide, cmap=cMap);\n",
    "    ax.invert_yaxis();\n",
    "    ax.set_ylim([0, M]);\n",
    "    colorbar = ax.collections[0].colorbar\n",
    "    colorbar.set_ticks([-0.667, 0, 0.667])\n",
    "    \n",
    "    l = i_split % n_folds_per_cluster + 1\n",
    "    k = np.floor_divide(i_split, n_folds_per_cluster) + 1\n",
    "    title = f'Nuisance: $I_{{{k}}}^C \\\\times J_{{{l}}}^C$; Score: $I_{{{k}}} \\\\times J_{{{l}}}$'\n",
    "    ax.set_title(title)\n",
    "    if l == n_folds_per_cluster:\n",
    "        colorbar.set_ticklabels(['Nuisance', '', 'Score'])\n",
    "    else:\n",
    "        colorbar.set_ticklabels(['', '', ''])\n",
    "    if l == 1:\n",
    "        ax.set_ylabel('First Cluster Variable $k$')\n",
    "    if k == 3:\n",
    "        ax.set_xlabel('Second Cluster Variable $\\ell$')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cluster Robust Standard Errors\n",
    "In the abstract base class `DoubleML` the estimation of cluster robust standard errors is implemented for all supported double machine learning models.\n",
    "It is based on the assumption of a linear Neyman orthogonal score function.\n",
    "We use the notation $n \\wedge m := \\min\\{n,m\\}$.\n",
    "For the the asymptotic variance of\n",
    "$\\sqrt{\\underline{C}}(\\tilde{\\theta_0} - \\theta_0$ with\n",
    "$\\underline{C} := N \\wedge M$\n",
    "Chiang et al. 2021 then proposed the following estimator\n",
    "$$\n",
    "\\hat{\\sigma}^2 = \\hat{J}^{-1} \\hat{\\Gamma} \\hat{J}^{-1}\n",
    "$$\n",
    "where\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\hat{\\Gamma} = \\frac{1}{K^2} \\sum_{(k, \\ell) \\in[K]^2}\n",
    "\\Bigg[ \\frac{|I_k| \\wedge |J_\\ell|}{(|I_k||J_\\ell|)^2}\n",
    "\\bigg(&\\sum_{i \\in I_k} \\sum_{j \\in J_\\ell} \\sum_{j' \\in J_\\ell}\n",
    "\\psi(W_{ij}; \\tilde{\\theta}, \\hat{\\eta}_{k \\ell}) \\psi(W_{ij'}; \\tilde{\\theta}_0, \\hat{\\eta}_{k \\ell}) \\\\\n",
    "&+ \\sum_{i \\in I_k} \\sum_{i' \\in I_k} \\sum_{j \\in J_\\ell}\n",
    "\\psi(W_{ij}; \\tilde{\\theta}, \\hat{\\eta}_{k \\ell}) \\psi(W_{i'j}; \\tilde{\\theta}_0, \\hat{\\eta}_{k \\ell})\n",
    "\\bigg)\n",
    "\\Bigg]\n",
    "\\end{aligned}$$\n",
    "and\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\hat{J} = \\frac{1}{K^2} \\sum_{(k, \\ell) \\in[K]^2} \\frac{1}{|I_k||J_\\ell|}\n",
    "\\sum_{i \\in I_k} \\sum_{j \\in J_\\ell}\n",
    "\\psi_a(W_{ij}; \\tilde{\\theta}_0, \\hat{\\eta}_{k \\ell}).\n",
    "\\end{aligned}\n",
    "$$\n",
    "A $(1-\\alpha)$ confidence interval is then given by (Chiang et al. 2020)\n",
    "$$\\begin{aligned}\n",
    "\\left[\n",
    "\\tilde{\\theta} \\pm \\Phi^{-1}(1-\\alpha/2) \\sqrt{\\hat{\\sigma}^2 / \\underline{C}}\n",
    "\\right]\n",
    "\\end{aligned}\n",
    "$$\n",
    "with $\\underline{C} = N \\wedge M$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimate the PLIV model with cluster robust double machine learning\n",
    "dml_pliv_obj.fit()\n",
    "dml_pliv_obj.summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (One-Way) Cluster Robust Double Machine Learing\n",
    "\n",
    "We again use the PLIV data generating process described in Section 4.1 of Chiang et al. (2020).\n",
    "To obtain one-way clustered data, we set the following weights to zero\n",
    "$$\n",
    "\\omega_2^X = \\omega_2^\\varepsilon = \\omega_2^v = \\omega_2^V = 0.\n",
    "$$\n",
    "Again we can simulate this data with [make_pliv_multiway_cluster_CKMS2021()](https://docs.doubleml.org/stable/api/generated/doubleml.datasets.make_pliv_multiway_cluster_CKMS2021.html#doubleml.datasets.make_pliv_multiway_cluster_CKMS2021). To prepare the data-backend for one-way clustering, we only have to alter the `cluster_cols` to be `'cluster_var_i'`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obj_dml_data = make_pliv_multiway_cluster_CKMS2021(N, M, dim_X,\n",
    "                                                   omega_X=np.array([0.25, 0]),\n",
    "                                                   omega_epsilon=np.array([0.25, 0]),\n",
    "                                                   omega_v=np.array([0.25, 0]),\n",
    "                                                   omega_V=np.array([0.25, 0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obj_dml_data.cluster_cols = 'cluster_var_i'\n",
    "print(obj_dml_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set machine learning methods for m & g\n",
    "learner = LassoCV()\n",
    "ml_g = clone(learner)\n",
    "ml_m = clone(learner)\n",
    "ml_r = clone(learner)\n",
    "\n",
    "# initialize the DoubleMLPLIV object\n",
    "dml_pliv_obj = DoubleMLPLIV(obj_dml_data,\n",
    "                            ml_g, ml_m, ml_r,\n",
    "                            n_folds=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_pliv_obj.fit()\n",
    "dml_pliv_obj.summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Application\n",
    "As data application we revist the consumer demand example from Chiang et al. (2021).\n",
    "The U.S. automobile data of Berry, Levinsohn, and Pakes (1995) is obtained from the `R` package `hdm`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load and Process Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from rpy2.robjects.packages import importr, data\n",
    "from rpy2.robjects import r, pandas2ri\n",
    "pandas2ri.activate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hdm = importr('hdm')\n",
    "blp_data = data(hdm).fetch('BLP')['BLP'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_cols = ['hpwt', 'air', 'mpd', 'space']\n",
    "blp_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def construct_iv(blp_data, x_cols=['hpwt', 'air', 'mpd', 'space']):\n",
    "    n = blp_data.shape[0]\n",
    "    p = len(x_cols)\n",
    "    \n",
    "    firmid = blp_data['firm.id'].values\n",
    "    cdid = blp_data['cdid'].values\n",
    "    id_var = blp_data['id'].values\n",
    "    X = blp_data[x_cols]\n",
    "    \n",
    "    sum_other = pd.DataFrame(columns=['sum.other.' + var for var in x_cols],\n",
    "                             index=blp_data.index)\n",
    "    sum_rival = pd.DataFrame(columns=['sum.rival.' + var for var in x_cols],\n",
    "                             index=blp_data.index)\n",
    "    \n",
    "    for i in range(n):\n",
    "        other_ind = (firmid == firmid[i]) & (cdid == cdid[i]) & (id_var != id_var[i])\n",
    "        rival_ind = (firmid != firmid[i]) & (cdid == cdid[i])\n",
    "        for j in range(p):\n",
    "            sum_other.iloc[i, j] = X.iloc[:,j][other_ind].sum()\n",
    "            sum_rival.iloc[i, j] = X.iloc[:,j][rival_ind].sum()\n",
    "    \n",
    "    return pd.concat((sum_other, sum_rival), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iv_vars = construct_iv(blp_data, x_cols=['hpwt', 'air', 'mpd', 'space'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "poly = PolynomialFeatures(degree=3, include_bias=False)\n",
    "data_transf = poly.fit_transform(blp_data[x_cols])\n",
    "x_cols_poly = poly.get_feature_names(x_cols)\n",
    "data_transf = pd.DataFrame(data_transf, columns=x_cols_poly)\n",
    "data_transf.index = blp_data.index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sel_cols_chiang = list(np.setdiff1d(data_transf.columns,\n",
    "                                    ['hpwt air mpd', 'hpwt air space',\n",
    "                                     'hpwt mpd space', 'air mpd space']))\n",
    "sel_cols_chiang"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blp_data['log_p'] = np.log(blp_data['price'] + 11.761)\n",
    "\n",
    "y_col = 'y'\n",
    "d_col = 'log_p'\n",
    "cluster_cols = ['model.id', 'cdid']\n",
    "all_z_cols = ['sum.other.hpwt', 'sum.other.mpd', 'sum.other.space']\n",
    "z_col = all_z_cols[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_df = pd.concat((blp_data[[y_col] + [d_col] + cluster_cols],\n",
    "                    data_transf[sel_cols_chiang],\n",
    "                    iv_vars[all_z_cols]),\n",
    "                   axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize `DoubleMLClusterData` object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_data = DoubleMLClusterData(dml_df,\n",
    "                               y_col=y_col,\n",
    "                               d_cols=d_col,\n",
    "                               z_cols=z_col,\n",
    "                               cluster_cols=cluster_cols,\n",
    "                               x_cols=sel_cols_chiang)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dml_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner = LassoCV(max_iter=50000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_df = pd.DataFrame()\n",
    "n_rep = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Two-Way Clustering with Respect to Product and Market"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_data.z_cols = z_col\n",
    "dml_data.cluster_cols = ['model.id', 'cdid']\n",
    "dml_pliv = DoubleMLPLIV(dml_data,\n",
    "                        clone(learner), clone(learner), clone(learner),\n",
    "                        n_folds=2, n_rep=n_rep)\n",
    "dml_pliv.fit()\n",
    "res = dml_pliv.summary.reset_index(drop=True)\n",
    "res['z_col'] = dml_data.z_cols[0]\n",
    "res['clustering'] = 'two-way'\n",
    "res_df = res_df.append(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### One-Way Clustering with Respect to the Product"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_data.z_cols = z_col\n",
    "dml_data.cluster_cols = 'model.id'\n",
    "dml_pliv = DoubleMLPLIV(dml_data,\n",
    "                        clone(learner), clone(learner), clone(learner),\n",
    "                        n_folds=4, n_rep=n_rep)\n",
    "dml_pliv.fit()\n",
    "res = dml_pliv.summary.reset_index(drop=True)\n",
    "res['z_col'] = dml_data.z_cols[0]\n",
    "res['clustering'] = 'one-way-product'\n",
    "res_df = res_df.append(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### One-Way Clustering with Respect to the Market"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_data.z_cols = z_col\n",
    "dml_data.cluster_cols = 'cdid'\n",
    "dml_pliv = DoubleMLPLIV(dml_data,\n",
    "                        clone(learner), clone(learner), clone(learner),\n",
    "                        n_folds=4, n_rep=n_rep)\n",
    "dml_pliv.fit()\n",
    "res = dml_pliv.summary.reset_index(drop=True)\n",
    "res['z_col'] = dml_data.z_cols[0]\n",
    "res['clustering'] = 'one-way-market'\n",
    "res_df = res_df.append(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### No Clustering / Zero-Way Clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_data = DoubleMLData(dml_df,\n",
    "                        y_col=y_col,\n",
    "                        d_cols=d_col,\n",
    "                        z_cols=z_col,\n",
    "                        x_cols=sel_cols_chiang)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dml_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_data.z_cols = z_col\n",
    "dml_pliv = DoubleMLPLIV(dml_data,\n",
    "                        clone(learner), clone(learner), clone(learner),\n",
    "                        n_folds=4, n_rep=n_rep)\n",
    "dml_pliv.fit()\n",
    "res = dml_pliv.summary.reset_index(drop=True)\n",
    "res['z_col'] = dml_data.z_cols[0]\n",
    "res['clustering'] = 'zero-way'\n",
    "res_df = res_df.append(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Application Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "Berry, S., Levinsohn, J., and Pakes, A. (1995), “Automobile Prices in Market\n",
    "Equilibrium,” Econometrica: Journal of the Econometric Society, 63, 841-890.\n",
    "\n",
    "Chiang, H. D., Kato K., Ma, Y. and Sasaki, Y. (2021), Multiway Cluster Robust Double/Debiased Machine Learning, Journal of Business & Economic Statistics, doi: [10.1080/07350015.2021.1895815](https://doi.org/10.1080/07350015.2021.1895815), arXiv: [1909.03489](https://arxiv.org/abs/1909.03489)."
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}