{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multilevel Modeling and Time-Series Forecasting\n",
    "\n",
    "In this lab you will fit multilevel (mixed-effects) models to nested behavioral data and compare them to a single-level regression that ignores the nesting structure. The dataset contains session-level response rates from eight participants, each observed across twenty sessions under varying reinforcement rates.\n",
    "\n",
    "**Learning objectives:**\n",
    "- Understand why ignoring nesting leads to biased inference\n",
    "- Fit and interpret mixed-effects models using `statsmodels`\n",
    "- Compare models using AIC and BIC\n",
    "- Visualize participant-level variation in intercepts and slopes\n",
    "\n",
    "**Required packages:** `pandas`, `numpy`, `matplotlib`, `statsmodels`, `scipy`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Run the cell below to import the libraries you will need."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import statsmodels.api as sm\n",
    "import statsmodels.formula.api as smf\n",
    "from scipy import stats\n",
    "\n",
    "plt.rcParams['figure.figsize'] = (10, 6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 1: Load and Explore the Data\n",
    "\n",
    "Load `nested_behavior_data.csv` into a pandas DataFrame. The dataset has four columns:\n",
    "\n",
    "| Column | Description |\n",
    "|---|---|\n",
    "| `participant_id` | Participant identifier (1-8) |\n",
    "| `session` | Session number (1-20) |\n",
    "| `reinforcement_rate` | Reinforcements delivered per minute |\n",
    "| `response_rate` | Responses per minute |\n",
    "\n",
    "After loading:\n",
    "1. Print the first few rows and the shape of the DataFrame.\n",
    "2. Compute summary statistics grouped by participant.\n",
    "3. Create a scatter plot of `reinforcement_rate` vs. `response_rate`, coloring points by participant."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "df = pd.read_csv('nested_behavior_data.csv')\n",
    "print(df.shape)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Summary statistics by participant\n",
    "summary = df.groupby('participant_id')[['reinforcement_rate', 'response_rate']].agg(['mean', 'std', 'min', 'max'])\n",
    "print(summary)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Scatter plot colored by participant\n",
    "fig, ax = plt.subplots()\n",
    "for pid, g in df.groupby('participant_id'):\n",
    "    ax.scatter(g['reinforcement_rate'], g['response_rate'], label=f'P{pid}', alpha=0.7)\n",
    "ax.set_xlabel('Reinforcement rate (per min)')\n",
    "ax.set_ylabel('Response rate (per min)')\n",
    "ax.set_title('Response vs. reinforcement rate, colored by participant')\n",
    "ax.legend(title='Participant', bbox_to_anchor=(1.02, 1), loc='upper left')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 2: Fit a Single-Level OLS Regression\n",
    "\n",
    "Fit an ordinary least squares (OLS) regression predicting `response_rate` from `reinforcement_rate`, pooling all participants together. This model ignores the nesting of observations within participants.\n",
    "\n",
    "1. Fit the model using `smf.ols()`.\n",
    "2. Print the summary.\n",
    "3. Plot the residuals vs. fitted values. Do you see any patterns that suggest the model is misspecified?\n",
    "4. Color the residual plot by participant. What do you notice?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "ols_result = smf.ols('response_rate ~ reinforcement_rate', data=df).fit()\n",
    "print(ols_result.summary())"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Residual plot colored by participant\n",
    "df['ols_fitted'] = ols_result.fittedvalues\n",
    "df['ols_resid'] = ols_result.resid\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "for pid, g in df.groupby('participant_id'):\n",
    "    ax.scatter(g['ols_fitted'], g['ols_resid'], label=f'P{pid}', alpha=0.7)\n",
    "ax.axhline(0, color='black', linewidth=1)\n",
    "ax.set_xlabel('Fitted values')\n",
    "ax.set_ylabel('Residuals')\n",
    "ax.set_title('OLS residuals vs. fitted, colored by participant')\n",
    "ax.legend(title='Participant', bbox_to_anchor=(1.02, 1), loc='upper left')\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Residuals cluster by participant -> observations within a participant are correlated.\n",
    "print('Mean residual by participant:')\n",
    "print(df.groupby('participant_id')['ols_resid'].mean())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question:** Why might the standard errors from the OLS model be misleading when the data are nested? Write your answer below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Observations within a participant are not independent -- each participant's residuals share a common offset (the per-participant mean residuals above are far from zero, ranging from about -8 to +10). OLS assumes independent errors, so it treats 160 correlated observations as if they were 160 independent ones. This overstates the effective sample size, shrinks the standard errors, and inflates the t-statistics, making effects look more precise and more significant than they really are."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 3: Fit a Random-Intercept Model\n",
    "\n",
    "Now fit a multilevel model that allows each participant to have their own intercept, while the slope of `reinforcement_rate` is fixed across participants. Use `smf.mixedlm()` from statsmodels.\n",
    "\n",
    "```python\n",
    "model = smf.mixedlm(\"response_rate ~ reinforcement_rate\",\n",
    "                     data=df,\n",
    "                     groups=df[\"participant_id\"])\n",
    "result = model.fit()\n",
    "```\n",
    "\n",
    "1. Fit the model and print the summary.\n",
    "2. What is the estimated variance of the random intercept?\n",
    "3. How does the fixed-effect slope compare to the OLS slope from Task 2?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Fit with maximum likelihood (reml=False) so AIC/BIC are comparable across models\n",
    "ri_model = smf.mixedlm('response_rate ~ reinforcement_rate', data=df, groups=df['participant_id'])\n",
    "ri_result = ri_model.fit(reml=False)\n",
    "print(ri_result.summary())\n",
    "\n",
    "print('\\nRandom-intercept variance (Group Var):', round(ri_result.cov_re.iloc[0, 0], 3))\n",
    "print('Mixed-model slope:', round(ri_result.fe_params['reinforcement_rate'], 3),\n",
    "      '| OLS slope:', round(ols_result.params['reinforcement_rate'], 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 4: Fit a Random-Intercept-and-Slope Model\n",
    "\n",
    "Extend the model so that each participant has both a unique intercept and a unique slope for `reinforcement_rate`. In statsmodels, you specify the random slope using the `re_formula` argument:\n",
    "\n",
    "```python\n",
    "model_rs = smf.mixedlm(\"response_rate ~ reinforcement_rate\",\n",
    "                        data=df,\n",
    "                        groups=df[\"participant_id\"],\n",
    "                        re_formula=\"~reinforcement_rate\")\n",
    "result_rs = model_rs.fit()\n",
    "```\n",
    "\n",
    "1. Fit the model and print the summary.\n",
    "2. Extract the random effects for each participant using `result_rs.random_effects`.\n",
    "3. For each participant, compute the participant-specific intercept and slope (fixed effect + random effect)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "rs_model = smf.mixedlm('response_rate ~ reinforcement_rate', data=df,\n",
    "                       groups=df['participant_id'], re_formula='~reinforcement_rate')\n",
    "rs_result = rs_model.fit(reml=False)\n",
    "print(rs_result.summary())"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Participant-specific intercepts and slopes = fixed effect + random effect\n",
    "fe_int = rs_result.fe_params['Intercept']\n",
    "fe_slope = rs_result.fe_params['reinforcement_rate']\n",
    "\n",
    "rows = []\n",
    "for pid, re in rs_result.random_effects.items():\n",
    "    rows.append({'participant_id': pid,\n",
    "                 'intercept': fe_int + re['Group'],\n",
    "                 'slope': fe_slope + re['reinforcement_rate']})\n",
    "participant_params = pd.DataFrame(rows).sort_values('participant_id').reset_index(drop=True)\n",
    "print(participant_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 5: Compare Models Using AIC and BIC\n",
    "\n",
    "Compare the three models (OLS, random-intercept, random-intercept-and-slope) using information criteria. Lower values indicate a better trade-off between fit and complexity.\n",
    "\n",
    "**Note:** statsmodels does not populate `.aic`/`.bic` on `MixedLMResults` (they return `nan`), so we compute them from the log-likelihood. All models are fit by maximum likelihood here so the values are comparable.\n",
    "\n",
    "Create a table summarizing the AIC and BIC for each model. Which model is preferred?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def info_criteria(result, n):\n",
    "    k = len(result.params) + 1  # +1 for the residual variance (scale)\n",
    "    aic = -2 * result.llf + 2 * k\n",
    "    bic = -2 * result.llf + k * np.log(n)\n",
    "    return aic, bic\n",
    "\n",
    "n = len(df)\n",
    "ri_aic, ri_bic = info_criteria(ri_result, n)\n",
    "rs_aic, rs_bic = info_criteria(rs_result, n)\n",
    "\n",
    "comparison = pd.DataFrame({\n",
    "    'model': ['OLS (pooled)', 'Random intercept', 'Random intercept + slope'],\n",
    "    'AIC': [ols_result.aic, ri_aic, rs_aic],\n",
    "    'BIC': [ols_result.bic, ri_bic, rs_bic],\n",
    "}).round(2)\n",
    "print(comparison)\n",
    "print('\\nPreferred by AIC:', comparison.loc[comparison['AIC'].idxmin(), 'model'])\n",
    "print('Preferred by BIC:', comparison.loc[comparison['BIC'].idxmin(), 'model'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question:** What do the AIC/BIC values tell you about the importance of accounting for individual differences in this dataset? Write your answer below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both criteria fall sharply as participant-level structure is added (AIC 1055 -> 821 -> 320; BIC 1061 -> 833 -> 338). The random-intercept-and-slope model is decisively preferred even under BIC's heavier penalty for extra parameters. Individual differences in *both* baseline responding and sensitivity to reinforcement are large, and a model that ignores them fits the data poorly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 6: Interpret Fixed and Random Effects\n",
    "\n",
    "Using the random-intercept-and-slope model:\n",
    "\n",
    "1. State in plain language what the fixed-effect intercept and slope represent.\n",
    "2. State what the random-effect variances represent.\n",
    "3. Which participants have the steepest and shallowest slopes? What might this mean behaviorally (e.g., sensitivity to reinforcement rate)?\n",
    "4. Is there a correlation between participant intercepts and slopes? If so, what does it mean?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "steep = participant_params.loc[participant_params['slope'].idxmax()]\n",
    "shallow = participant_params.loc[participant_params['slope'].idxmin()]\n",
    "print(f\"Steepest slope:  P{int(steep['participant_id'])} ({steep['slope']:.3f})\")\n",
    "print(f\"Shallowest slope: P{int(shallow['participant_id'])} ({shallow['slope']:.3f})\")\n",
    "\n",
    "r, p = stats.pearsonr(participant_params['intercept'], participant_params['slope'])\n",
    "print(f\"\\nIntercept-slope correlation: r = {r:.3f}, p = {p:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fixed-effect intercept (~8.4) is the average response rate predicted at zero reinforcement, and the fixed slope (~2.43) is the average increase in responses per minute for each additional reinforcer per minute. The random-effect variances describe how much participants deviate from those averages -- a large intercept variance means participants differ in baseline responding, and a non-zero slope variance means they differ in sensitivity to reinforcement. Here P8 has the steepest slope (~3.79, most sensitive) and P5 the shallowest (~0.99, nearly insensitive). The intercept-slope correlation is strongly negative (r = -0.97): participants with high baseline response rates tend to have shallower slopes, while low-baseline participants increase more steeply with reinforcement."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 7: Visualize Participant-Level Regression Lines\n",
    "\n",
    "Create a figure with:\n",
    "1. A scatter plot of the raw data, colored by participant.\n",
    "2. The OLS regression line (single line for all data) in black.\n",
    "3. The participant-specific regression lines from the random-intercept-and-slope model, each in a different color.\n",
    "\n",
    "This visualization should make it visually clear why a single regression line is inadequate."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "x_grid = np.linspace(df['reinforcement_rate'].min(), df['reinforcement_rate'].max(), 50)\n",
    "\n",
    "for _, row in participant_params.iterrows():\n",
    "    pid = int(row['participant_id'])\n",
    "    g = df[df['participant_id'] == pid]\n",
    "    pts = ax.scatter(g['reinforcement_rate'], g['response_rate'], alpha=0.5)\n",
    "    ax.plot(x_grid, row['intercept'] + row['slope'] * x_grid,\n",
    "            color=pts.get_facecolor()[0], linewidth=1.5)\n",
    "\n",
    "ax.plot(x_grid,\n",
    "        ols_result.params['Intercept'] + ols_result.params['reinforcement_rate'] * x_grid,\n",
    "        color='black', linewidth=3, linestyle='--', label='OLS (pooled)')\n",
    "ax.set_xlabel('Reinforcement rate (per min)')\n",
    "ax.set_ylabel('Response rate (per min)')\n",
    "ax.set_title('Participant-specific regression lines vs. pooled OLS')\n",
    "ax.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 8: Visualize Random Effects\n",
    "\n",
    "Create two plots:\n",
    "1. A caterpillar plot (forest plot) showing the random intercepts for each participant with confidence intervals.\n",
    "2. A caterpillar plot showing the random slopes for each participant.\n",
    "\n",
    "**Hint:** You can extract the random effects from `result_rs.random_effects` and approximate confidence intervals using the conditional standard errors."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Point estimates of the random effects (deviations from the fixed effect)\n",
    "re_int = {pid: re['Group'] for pid, re in rs_result.random_effects.items()}\n",
    "re_slope = {pid: re['reinforcement_rate'] for pid, re in rs_result.random_effects.items()}\n",
    "\n",
    "# Approximate SE bands from the estimated random-effect SDs (illustrative)\n",
    "se_int = np.sqrt(rs_result.cov_re.loc['Group', 'Group'])\n",
    "se_slope = np.sqrt(rs_result.cov_re.loc['reinforcement_rate', 'reinforcement_rate'])\n",
    "\n",
    "def caterpillar(ax, effects, se, title):\n",
    "    items = sorted(effects.items(), key=lambda kv: kv[1])\n",
    "    labels = [f'P{pid}' for pid, _ in items]\n",
    "    vals = [v for _, v in items]\n",
    "    y = np.arange(len(vals))\n",
    "    ax.errorbar(vals, y, xerr=1.96 * se, fmt='o', capsize=3)\n",
    "    ax.axvline(0, color='red', linestyle='--', linewidth=1)\n",
    "    ax.set_yticks(y); ax.set_yticklabels(labels)\n",
    "    ax.set_title(title); ax.set_xlabel('Deviation from fixed effect')\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))\n",
    "caterpillar(ax1, re_int, se_int, 'Random intercepts (~95% approx CI)')\n",
    "caterpillar(ax2, re_slope, se_slope, 'Random slopes (~95% approx CI)')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 9 (Optional): Time-Series Decomposition\n",
    "\n",
    "Select one participant and treat their session-by-session `response_rate` as a short time series.\n",
    "\n",
    "1. Plot the raw time series.\n",
    "2. Use `statsmodels.tsa.seasonal.seasonal_decompose` (with `model='additive'` and an appropriate `period`) to decompose the series into trend, seasonal, and residual components.\n",
    "3. Since our data may not have a true seasonal component over 20 sessions, focus on the trend and residual. What does the trend tell you about how this participant's behavior changed over sessions?\n",
    "\n",
    "**Note:** With only 20 observations, this decomposition is illustrative rather than definitive. The goal is to introduce the concept of decomposing temporal data."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from statsmodels.tsa.seasonal import seasonal_decompose\n",
    "\n",
    "p1 = df[df['participant_id'] == 1].sort_values('session').set_index('session')['response_rate']\n",
    "decomp = seasonal_decompose(p1, model='additive', period=4)\n",
    "\n",
    "fig = decomp.plot()\n",
    "fig.set_size_inches(10, 8)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print('Trend (start -> end):',\n",
    "      round(decomp.trend.dropna().iloc[0], 2), '->', round(decomp.trend.dropna().iloc[-1], 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reflection\n",
    "\n",
    "In a few sentences, address the following:\n",
    "\n",
    "1. Why is multilevel modeling particularly important for behavioral research where data are collected across multiple participants and sessions?\n",
    "2. How would you explain the difference between fixed effects and random effects to a colleague who has only used standard regression?\n",
    "3. What are the limitations of the models you fit today (e.g., assumptions about the error structure, linearity)?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Behavioral data are almost always nested (sessions within participants, participants within groups), and pooling ignores that structure -- biasing standard errors and masking individual differences that are often the phenomenon of interest. 2. Fixed effects are the average relationships that apply to the whole population; random effects capture how individual units depart from those averages, with the model estimating the *distribution* of those departures rather than a separate free parameter per unit. 3. These models still assume linearity, normally distributed random effects and residuals, and homoscedastic within-participant errors; with only 8 participants and 20 sessions each the variance components (and the time-series decomposition) are estimated imprecisely."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}