{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "848e73da-e1c3-4617-b35c-eabefadc4856",
   "metadata": {},
   "source": [
    "# High School and Beyond (1982): Clustered Data\n",
    "\n",
    "One of the main uses for estimating equations historically has been to handle clustered data. This use was popularized by Liang & Zeger (1986). While `delicatessen` relies on estimating equations for other tasks, it can also be used to handle clustered data. This tutorial reviews how clustered observations can be handled using built-in `delicatessen` functionalities.\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7e27e2c9-3f56-42b4-ad80-51d7dcab963c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Versions\n",
      "NumPy:        2.3.5\n",
      "SciPy:        1.16.3\n",
      "pandas:       2.3.3\n",
      "Matplotlib:   3.10.8\n",
      "Delicatessen: 4.1\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import pandas as pd\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import delicatessen as deli\n",
    "from delicatessen import MEstimator\n",
    "from delicatessen.estimating_equations import ee_regression\n",
    "from delicatessen.utilities import aggregate_efuncs\n",
    "\n",
    "print(\"Versions\")\n",
    "print(\"NumPy:       \", np.__version__)\n",
    "print(\"SciPy:       \", sp.__version__)\n",
    "print(\"pandas:      \", pd.__version__)\n",
    "print(\"Matplotlib:  \", mpl.__version__)\n",
    "print(\"Delicatessen:\", deli.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae574b7b-9a75-45fc-b433-66794dc7a7d0",
   "metadata": {},
   "source": [
    "This tutorial uses data High School and Beyond study from 1982. This data set comes from the `mlmRev` R package. See that package for details. In this example, we will conduct a simple regression analysis on mathematical achievement scores of students by several individual and school level factors. Here, clustering is assumed to occur at the school-level and is considered an incidental feature (i.e., school-specific coefficients are not of interest). \n",
    "\n",
    "The following code loads this data (saved as a .csv file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f7aad6dd-24f9-4338-a1da-a714d1f161fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = pd.read_csv(\"data/hsb82.csv\")\n",
    "d['intercept'] = 1\n",
    "d['female'] = np.where(d['sx'] == 'Female', 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "89cb0f9c-59c8-40da-b6c9-715f614f8b05",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.asarray(d['mAch'])  # Math achievement score\n",
    "X = np.asarray(d[['intercept', 'female', 'cses']])\n",
    "g = np.asarray(d['school'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2879165c-093d-47dd-88d5-3555fb12fdb9",
   "metadata": {},
   "source": [
    "To begin, consider if we ignored the clustering. For this, we can fit a linear regression model using the built-in `ee_regression` function (as illustrated elsewhere)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dab359c8-7dc1-4d21-892f-8feba63febfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_i(theta):\n",
    "    return ee_regression(theta=theta, y=y, X=X, model='linear')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dafdc304-3722-4f77-bcc2-e09cfdfea240",
   "metadata": {},
   "outputs": [],
   "source": [
    "estr_i = MEstimator(psi_i, init=[10, 0, 0])\n",
    "estr_i.estimate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4ab068e7-4e6b-4c11-b1b3-12205b12e304",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==============================================================\n",
      "              Estimation Method: M-estimator\n",
      "--------------------------------------------------------------\n",
      "No. Observations:        7185 | No. Parameters:              3\n",
      "Solving algorithm:         lm | Max Iterations:           5000\n",
      "Solving tolerance:      1e-09 | Allow P-Inverse:             1\n",
      "Derivative Method:     approx | Deriv Approx:            1e-09\n",
      "Small N Correction:      None | Distribution:           Z-stat\n",
      "==============================================================\n",
      "   Theta   StdErr  Z-score      LCL      UCL  P-value  S-value \n",
      "--------------------------------------------------------------\n",
      "   12.01     0.11   114.12    11.80    12.21     0.00      inf \n",
      "    1.57     0.16     9.93     1.26     1.88     0.00    74.75 \n",
      "    2.14     0.12    17.71     1.90     2.38     0.00   230.81 \n",
      "==============================================================\n"
     ]
    }
   ],
   "source": [
    "estr_i.print_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88f3cce2-b597-4d9e-852a-89cc6b7fbe66",
   "metadata": {},
   "source": [
    "From this model, we see that being male and having a higher SES relative to your school's average SES had a positive association math achievement scores. \n",
    "\n",
    "## Clustering by School\n",
    "\n",
    "The previous results, particularly the inferential statistics (standard errors, Z-scores, confidence intervals, P-values, S-values), are all premised on that observations are independent. However, we might be skeptical of this assumption. In particular, those in the same school may be more similar than those from different schools. From a certain perspective, we can think about these correlated observations as contributing 'less than 1 unit's worth' of information to our model. We can use estimating equations and the sandwich variance to handle this challenge.\n",
    "\n",
    "To do this, we will essentially collapse the estimating functions from $n$, the number of units, to $m$, the number of schools. So, this changes our sample size (and thus all asymptotics will be based on $m$ and not $n$ anymore). How we collapse observations is determined by something called the 'working correlation matrix'. This matrix stipulates how observations are correlated (and is something we assume beforehand). The good news is that the sandwich variance is robust to misspecification of this working correlation matrix.\n",
    "\n",
    "Within `delicatessen`, the collapsing of estimating functions from $n$ to $m$ can be done by the `aggregate_efuncs` utility function. This function takes a given estimating function and adds together observations within the same cluster defined by the `group` argument. Note that this function only supports the 'independent' working correlation matrix. While this might be a known misspecification (in this and other clustering settings), this choice was made for several reasons: (1) this approach is more flexible and easily generalizes to arbitrary input estimating functions, (2) non-diagonal working correlation matrices rely on an additional assumption that may not hold and will produced biased *point* estimates when it doesn't. The second point is detailed further in Pepe & Anderson (1994) and Pan et al. (2000). The independent correlation matrix avoids this, so it doesn't rely on this assumption and won't produce biased point estimates. The downside of this choice is that the standard error estimate is not as efficient as could be (i.e., larger than needs be) when a non-independent working correlation matrix is specified and the prior assumption does hold. Despite this downside, the flexibility and robustness offered by this approach seems preferable. Therefore, only the independent working correlation matrix was made available.\n",
    "\n",
    "The following code uses the `aggregate_efuncs` function to condense the previous estimating functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0b17186e-b648-4851-8eb1-dca1ff67c924",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_s(theta):\n",
    "    return aggregate_efuncs(psi_i(theta), group=g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4f5ffd6a-cd33-4677-be7f-05156a862627",
   "metadata": {},
   "outputs": [],
   "source": [
    "estr_s = MEstimator(psi_s, init=[10, 0, 0])\n",
    "estr_s.estimate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "eaa8c13e-af26-4fa3-8054-8dda88963492",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==============================================================\n",
      "              Estimation Method: M-estimator\n",
      "--------------------------------------------------------------\n",
      "No. Observations:         160 | No. Parameters:              3\n",
      "Solving algorithm:         lm | Max Iterations:           5000\n",
      "Solving tolerance:      1e-09 | Allow P-Inverse:             1\n",
      "Derivative Method:     approx | Deriv Approx:            1e-09\n",
      "Small N Correction:      None | Distribution:           Z-stat\n",
      "==============================================================\n",
      "   Theta   StdErr  Z-score      LCL      UCL  P-value  S-value \n",
      "--------------------------------------------------------------\n",
      "   12.01     0.26    45.31    11.49    12.52     0.00      inf \n",
      "    1.57     0.31     5.03     0.96     2.19     0.00    20.96 \n",
      "    2.14     0.13    16.76     1.89     2.39     0.00   207.06 \n",
      "==============================================================\n"
     ]
    }
   ],
   "source": [
    "estr_s.print_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "affcf359-323c-43c6-97fb-d2bf340602f5",
   "metadata": {},
   "source": [
    "In the results output, we can see several changes to our results. First, in the meta-data we see the number of observations drop from 7185 to 160. This is because while there were 7185 students in the study, these students only came from 160 different schools. Second, we see the standard errors are substantially larger than the previous case. This then leads to differences in the Z-scores, confidence intervals, P-values, and S-values. These increased, as we would expect with clustered data where there is some correlation between observations. While our inferential results changed, the point estimates did not. Again, this is because of our selection of the independent working correlation matrix.\n",
    "\n",
    "That concludes this example of how to handle clustered data with `delicatessen`. While only show in the context of linear regression, the `aggregate_efuncs` function handles condensing any user-specified estimating functions.\n",
    "\n",
    "## References \n",
    "\n",
    "Liang KY, & Zeger L. (1986). Longitudinal data analysis using generalized linear models. *Biometrika*, 73(1), 13-22.\n",
    "\n",
    "Pan W, Louis TA, & Connett JE. (2000). A note on marginal linear regression with correlated response data. *The American Statistician*, 54(3), 191-195.\n",
    "\n",
    "Pepe SM & Anderson GL (1994). A cautionary note on inference for marginal regression models with longitudinal data and general correlated response data. *Communications in Statistics-Simulation and Computation*, 23, 939-951."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}