{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Behavioral Momentum and Response Persistence\n",
    "\n",
    "In this lab, you will analyze resistance to change across reinforcement contexts, using three different *disruptors*. Behavioral momentum theory (Nevin, 1992; Nevin & Shahan, 2011) predicts that responding maintained by richer reinforcement is more resistant to disruption.\n",
    "\n",
    "The lab has two parts:\n",
    "\n",
    "- **Part 1 - Prefeeding.** Fit the simplified momentum equation to multiple-schedule data where prefeeding is the disruptor, and compare resistance to change across rich and lean components.\n",
    "- **Part 2 - Other disruptors.** Fit Nevin's resistance-to-extinction and alternative-reinforcement equations to data where extinction and alternative reinforcement are the disruptors.\n",
    "\n",
    "## Background (Part 1)\n",
    "\n",
    "Four subjects responded on a two-component multiple schedule with **rich** and **lean** reinforcement components. After baseline response rates were established, responding was disrupted by prefeeding at graded levels (0, 25, 50, 75, and 100g). The data record baseline and disrupted rates for each subject, component, and disruption level.\n",
    "\n",
    "The behavioral momentum equation describes the relation between disruption and proportional change in responding:\n",
    "\n",
    "$$\\log\\left(\\frac{B_x}{B_0}\\right) = \\frac{-x \\cdot c}{r \\cdot S}$$\n",
    "\n",
    "where $B_x$ is the response rate under disruption level $x$, $B_0$ is the baseline rate, $c$ is sensitivity to disruption, $r$ is the reinforcement rate, and $S$ captures the stimulus-reinforcer relation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 1: Import Libraries\n",
    "\n",
    "Import the libraries you will need for this lab: `pandas`, `numpy`, `matplotlib.pyplot`, and `scipy.optimize`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize import curve_fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 2: Load and Inspect the Data\n",
    "\n",
    "Load `momentum_data.csv` into a DataFrame. Examine the structure of the data. How many subjects are there? What are the unique disruption levels? Print summary statistics for the rich and lean components separately."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"momentum_data.csv\")\n",
    "\n",
    "print(\"Number of subjects:\", df['subject'].nunique())\n",
    "print(\"Disruption levels:\", sorted(df['disruption_level'].unique()))\n",
    "print(\"Components:\", df['component'].unique())\n",
    "\n",
    "for comp in ['rich', 'lean']:\n",
    "    print(f\"\\n--- {comp} component ---\")\n",
    "    print(df[df['component'] == comp][['baseline_rate', 'disrupted_rate']].describe())\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 3: Calculate Proportion of Baseline\n",
    "\n",
    "For each row in the dataset, calculate the proportion of baseline responding ($B_x / B_0$). Add this as a new column called `prop_baseline`. Then calculate $\\log_{10}(B_x / B_0)$ and store it in a column called `log_prop_baseline`.\n",
    "\n",
    "**Note:** At disruption level 0, the proportion should be 1.0 and the log proportion should be 0.0. Verify this is the case in your data."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "df['prop_baseline'] = df['disrupted_rate'] / df['baseline_rate']\n",
    "df['log_prop_baseline'] = np.log10(df['prop_baseline'])\n",
    "\n",
    "check = df[df['disruption_level'] == 0][['prop_baseline', 'log_prop_baseline']]\n",
    "print(\"At disruption level 0:\")\n",
    "print(\"  prop_baseline values:\", check['prop_baseline'].unique())\n",
    "print(\"  log_prop_baseline values:\", check['log_prop_baseline'].unique())\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 4: Visualize the Raw Disruption Functions\n",
    "\n",
    "Create a figure with two panels (subplots). In the left panel, plot proportion of baseline ($B_x / B_0$) as a function of disruption level, with separate lines for each subject. Use different colors or markers for rich vs. lean components. In the right panel, plot the same data on a log scale ($\\log_{10}(B_x / B_0)$ on the y-axis).\n",
    "\n",
    "Add appropriate axis labels, legends, and a title. What pattern do you notice regarding the rich vs. lean components?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, (axL, axR) = plt.subplots(1, 2, figsize=(14, 5))\n",
    "colors = {'rich': 'tab:blue', 'lean': 'tab:orange'}\n",
    "seen = set()\n",
    "\n",
    "for (subj, comp), g in df.groupby(['subject', 'component']):\n",
    "    g = g.sort_values('disruption_level')\n",
    "    label = comp if comp not in seen else None\n",
    "    seen.add(comp)\n",
    "    axL.plot(g['disruption_level'], g['prop_baseline'], marker='o',\n",
    "             color=colors[comp], alpha=0.6, label=label)\n",
    "    axR.plot(g['disruption_level'], g['log_prop_baseline'], marker='o',\n",
    "             color=colors[comp], alpha=0.6, label=label)\n",
    "\n",
    "axL.set_xlabel('Disruption Level (g prefeeding)'); axL.set_ylabel('Bx / B0')\n",
    "axL.set_title('Proportion of Baseline'); axL.legend()\n",
    "axR.set_xlabel('Disruption Level (g prefeeding)'); axR.set_ylabel('log10(Bx / B0)')\n",
    "axR.set_title('Log Proportion of Baseline'); axR.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 5: Compute Group Means\n",
    "\n",
    "Calculate the mean proportion of baseline and mean log proportion of baseline across subjects for each combination of component (rich/lean) and disruption level. Store these in a new DataFrame. This will be useful for fitting the momentum equation to averaged data."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "group_means = (df.groupby(['component', 'disruption_level'])\n",
    "                 [['prop_baseline', 'log_prop_baseline']]\n",
    "                 .mean()\n",
    "                 .reset_index())\n",
    "group_means"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 6: Define the Behavioral Momentum Equation\n",
    "\n",
    "Write a Python function that implements the behavioral momentum equation:\n",
    "\n",
    "$$\\log\\left(\\frac{B_x}{B_0}\\right) = \\frac{-x \\cdot c}{r \\cdot S}$$\n",
    "\n",
    "The function should take `x` (disruption level) as the independent variable and `c`, `r`, and `S` as parameters. It should return the predicted log proportion of baseline.\n",
    "\n",
    "**Hint:** Since $r$ and $S$ always appear as the product $r \\cdot S$, you may find it easier to define a combined parameter (e.g., `rS`) to avoid identifiability issues during fitting. Think about what this means for interpreting your results."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def momentum_model(x, c, rS):\n",
    "    \"\"\"Behavioral momentum: log10(Bx/B0) = -(c * x) / (r * S).\n",
    "\n",
    "    r and S enter only as the product r*S, so they are combined into a single\n",
    "    rS parameter. Only the ratio c/rS (the slope) is identifiable from one\n",
    "    disruption function -- see Task 7.\n",
    "    \"\"\"\n",
    "    return -(c * x) / rS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 7: Fit the Model to Each Component\n",
    "\n",
    "Using `scipy.optimize.curve_fit`, fit the behavioral momentum equation to the group-mean data for the **rich** component and the **lean** component separately.\n",
    "\n",
    "Report the best-fitting parameter values for each component. Which component shows a steeper disruption function (i.e., less resistance to change)?\n",
    "\n",
    "**Hints:**\n",
    "- Exclude the disruption level 0 data point from fitting (it is trivially 0 on the log scale).\n",
    "- Provide reasonable initial parameter guesses.\n",
    "- Consider whether parameter bounds are needed."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Only c/rS is identifiable, so we fix c = 1.0 and let rS absorb the slope.\n",
    "# A larger rS means a flatter function -> more resistance to change.\n",
    "C_FIXED = 1.0\n",
    "\n",
    "def fit_one(x, rS):\n",
    "    return momentum_model(x, C_FIXED, rS)\n",
    "\n",
    "fits = {}\n",
    "for comp in ['rich', 'lean']:\n",
    "    gm = group_means[(group_means['component'] == comp) &\n",
    "                     (group_means['disruption_level'] > 0)]\n",
    "    x = gm['disruption_level'].values.astype(float)\n",
    "    y = gm['log_prop_baseline'].values\n",
    "    popt, _ = curve_fit(fit_one, x, y, p0=[100.0], bounds=(1e-6, np.inf))\n",
    "    rS = popt[0]\n",
    "    fits[comp] = {'c': C_FIXED, 'rS': rS, 'slope': C_FIXED / rS}\n",
    "    print(f\"{comp:>4}: c={C_FIXED}, rS={rS:.2f}, slope=c/rS={C_FIXED/rS:.5f}\")\n",
    "\n",
    "steeper = min(fits, key=lambda c: fits[c]['rS'])\n",
    "print(f\"\\nSteeper (less resistant) component: {steeper}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 8: Evaluate Model Fit\n",
    "\n",
    "For each component, calculate the following goodness-of-fit measures:\n",
    "\n",
    "1. **R-squared ($R^2$):** proportion of variance accounted for.\n",
    "2. **Root Mean Squared Error (RMSE):** average magnitude of the residuals.\n",
    "\n",
    "How well does the behavioral momentum equation describe the disruption data? Are the fits comparable across components?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "for comp in ['rich', 'lean']:\n",
    "    gm = group_means[(group_means['component'] == comp) &\n",
    "                     (group_means['disruption_level'] > 0)]\n",
    "    x = gm['disruption_level'].values.astype(float)\n",
    "    y = gm['log_prop_baseline'].values\n",
    "    pred = momentum_model(x, fits[comp]['c'], fits[comp]['rS'])\n",
    "    ss_res = np.sum((y - pred) ** 2)\n",
    "    ss_tot = np.sum((y - np.mean(y)) ** 2)\n",
    "    r2 = 1 - ss_res / ss_tot\n",
    "    rmse = np.sqrt(np.mean((y - pred) ** 2))\n",
    "    fits[comp]['r2'] = r2\n",
    "    fits[comp]['rmse'] = rmse\n",
    "    print(f\"{comp:>4}: R2={r2:.3f}, RMSE={rmse:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 9: Plot Model Predictions Against Data\n",
    "\n",
    "Create a publication-quality figure showing:\n",
    "- The group-mean log proportion of baseline data points for each component (use distinct markers for rich and lean).\n",
    "- The fitted model predictions as smooth curves overlaid on the data.\n",
    "- A legend identifying the components.\n",
    "- Appropriate axis labels (\"Disruption Level (g prefeeding)\" and \"log(Bx/B0)\").\n",
    "\n",
    "Generate the smooth curves by evaluating the model function at many x-values between 0 and 100."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(9, 6))\n",
    "markers = {'rich': 'o', 'lean': 's'}\n",
    "colors = {'rich': 'tab:blue', 'lean': 'tab:orange'}\n",
    "xx = np.linspace(0, 100, 200)\n",
    "\n",
    "for comp in ['rich', 'lean']:\n",
    "    gm = group_means[group_means['component'] == comp]\n",
    "    ax.scatter(gm['disruption_level'], gm['log_prop_baseline'],\n",
    "               marker=markers[comp], color=colors[comp], s=70,\n",
    "               label=f\"{comp} (data)\", zorder=3)\n",
    "    yy = momentum_model(xx, fits[comp]['c'], fits[comp]['rS'])\n",
    "    ax.plot(xx, yy, color=colors[comp],\n",
    "            label=f\"{comp} (fit, rS={fits[comp]['rS']:.1f})\")\n",
    "\n",
    "ax.set_xlabel('Disruption Level (g prefeeding)')\n",
    "ax.set_ylabel('log(Bx/B0)')\n",
    "ax.set_title('Behavioral Momentum Model: Predictions vs. Data')\n",
    "ax.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 10: Interpretation and Discussion (Part 1)\n",
    "\n",
    "In a markdown cell below, address the following:\n",
    "\n",
    "1. How do the fitted parameter values differ between the rich and lean components? What do these differences mean in terms of behavioral momentum theory?\n",
    "2. Is resistance to change greater in the rich component, as predicted by the theory? How do you quantify this from your model fits?\n",
    "3. What are the limitations of using the simplified momentum equation with prefeeding as the disruptor?\n",
    "4. If you combined the parameter $r \\cdot S$ into a single value, what additional information would you need to separate them? Why might this matter for applied predictions?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Other Disruptors\n",
    "\n",
    "### Task 11: Resistance to Extinction\n",
    "\n",
    "Prefeeding is only one way to disrupt responding. Load `behavioral_momentum_extinction_data.csv`, where **extinction** is the disruptor across sessions. Implement Nevin's resistance-to-extinction equation\n",
    "\n",
    "$$\\frac{B_t}{B_0} = 10^{\\,-t\\,(c + d\\,r)/\\sqrt{r}}$$\n",
    "\n",
    "and use `curve_fit` to estimate `c`, `d`, `r`, and `B0` for each condition. Report the fitted parameters and $R^2$, and overlay the fitted curves on the data. (Bounds: $c, d > 0$; $r > 0.1$; $B_0$ a plausible baseline rate.)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "ext = pd.read_csv('behavioral_momentum_extinction_data.csv')\n",
    "\n",
    "def momentum_extinction(t, c, d, r, B0):\n",
    "    \"\"\"Resistance to extinction: Bt/B0 = 10^(-t * (c + d*r) / sqrt(r)).\"\"\"\n",
    "    return B0 * 10 ** ((-t * (c + d * r)) / (r ** 0.5))\n",
    "\n",
    "ext_results = []\n",
    "for cond in ext['condition'].unique():\n",
    "    cd = ext[ext['condition'] == cond]\n",
    "    t = cd['session'].values.astype(float)\n",
    "    B = cd['response_rate'].values\n",
    "    popt, _ = curve_fit(momentum_extinction, t, B, p0=[0.2, 0.1, 1.0, 40.0],\n",
    "                        bounds=([0, 0, 0.1, 0], [1, 1, 5, 100]))\n",
    "    pred = momentum_extinction(t, *popt)\n",
    "    r2 = 1 - np.sum((B - pred) ** 2) / np.sum((B - np.mean(B)) ** 2)\n",
    "    ext_results.append({'condition': cond, 'c': popt[0], 'd': popt[1],\n",
    "                        'r': popt[2], 'B0': popt[3], 'r2': r2})\n",
    "    print(f\"{cond}: c={popt[0]:.3f}, d={popt[1]:.3f}, r={popt[2]:.3f}, \"\n",
    "          f\"B0={popt[3]:.1f}, R2={r2:.3f}\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(9, 5))\n",
    "for row in ext_results:\n",
    "    cd = ext[ext['condition'] == row['condition']].sort_values('session')\n",
    "    ax.scatter(cd['session'], cd['response_rate'], s=30)\n",
    "    tt = np.linspace(cd['session'].min(), cd['session'].max(), 100)\n",
    "    ax.plot(tt, momentum_extinction(tt, row['c'], row['d'], row['r'], row['B0']),\n",
    "            label=row['condition'])\n",
    "ax.set_xlabel('Extinction session'); ax.set_ylabel('Response rate')\n",
    "ax.set_title('Resistance to Extinction'); ax.legend()\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 12: Disruption by Alternative Reinforcement\n",
    "\n",
    "Load `behavioral_momentum_alternative_data.csv`, where responding is disrupted by **alternative reinforcement** delivered at rate $R_a$. Implement\n",
    "\n",
    "$$\\frac{B_x}{B_0} = 10^{\\,-p\\,R_a/\\sqrt{r + R_a}}$$\n",
    "\n",
    "and fit `p`, `r`, and `B0` per condition. Overlay the fitted curves, and quantify how much responding is suppressed from baseline to the highest alternative-reinforcement rate."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "alt = pd.read_csv('behavioral_momentum_alternative_data.csv')\n",
    "\n",
    "def momentum_alternative(Ra, p, r, B0):\n",
    "    \"\"\"Disruption by alternative reinforcement: Bx/B0 = 10^(-p*Ra / sqrt(r + Ra)).\"\"\"\n",
    "    return B0 * 10 ** ((-p * Ra) / ((r + Ra) ** 0.5))\n",
    "\n",
    "alt_results = []\n",
    "for cond in alt['condition'].unique():\n",
    "    cd = alt[alt['condition'] == cond]\n",
    "    Ra = cd['Ra'].values.astype(float)\n",
    "    B = cd['response_rate'].values\n",
    "    popt, _ = curve_fit(momentum_alternative, Ra, B, p0=[0.3, 1.0, 50.0],\n",
    "                        bounds=([0, 0.1, 0], [2, 5, 100]))\n",
    "    pred = momentum_alternative(Ra, *popt)\n",
    "    r2 = 1 - np.sum((B - pred) ** 2) / np.sum((B - np.mean(B)) ** 2)\n",
    "    baseline = cd[cd['Ra'] == 0]['response_rate'].iloc[0]\n",
    "    final = cd[cd['Ra'] == cd['Ra'].max()]['response_rate'].iloc[0]\n",
    "    suppression = (baseline - final) / baseline * 100\n",
    "    alt_results.append({'condition': cond, 'p': popt[0], 'r': popt[1],\n",
    "                        'B0': popt[2], 'r2': r2, 'suppression_pct': suppression})\n",
    "    print(f\"{cond}: p={popt[0]:.3f}, r={popt[1]:.3f}, B0={popt[2]:.1f}, \"\n",
    "          f\"R2={r2:.3f}, suppression={suppression:.1f}%\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(9, 5))\n",
    "for row in alt_results:\n",
    "    cd = alt[alt['condition'] == row['condition']].sort_values('Ra')\n",
    "    ax.scatter(cd['Ra'], cd['response_rate'], s=30)\n",
    "    rr = np.linspace(cd['Ra'].min(), cd['Ra'].max(), 100)\n",
    "    ax.plot(rr, momentum_alternative(rr, row['p'], row['r'], row['B0']),\n",
    "            label=row['condition'])\n",
    "ax.set_xlabel('Alternative reinforcement rate (Ra)'); ax.set_ylabel('Response rate')\n",
    "ax.set_title('Disruption by Alternative Reinforcement'); ax.legend()\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrap-up\n",
    "\n",
    "Across all three disruptors (prefeeding, extinction, alternative reinforcement), what is the common signature of behavioral momentum? Compare how each disruptor enters its equation. Which parameters index *sensitivity to disruption* versus the *baseline mass* of behavior, and how do the recovered values line up with the theory's prediction that richer reinforcement yields greater resistance to change?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}