{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dynamical Systems Models: The Logistic ODE and Behavioral Acquisition\n",
    "\n",
    "In this lab you will work with the logistic ordinary differential equation (ODE) as a model of behavioral acquisition:\n",
    "\n",
    "$$\\frac{dx}{dt} = r \\cdot x \\cdot \\left(1 - \\frac{x}{K}\\right)$$\n",
    "\n",
    "where:\n",
    "- $x(t)$ is the response rate at time $t$\n",
    "- $r$ is the intrinsic growth rate (how quickly behavior accelerates early in acquisition)\n",
    "- $K$ is the carrying capacity (the asymptotic response rate)\n",
    "\n",
    "**Learning objectives:**\n",
    "- Implement numerical integration using Euler's method\n",
    "- Analyze equilibrium points and their stability\n",
    "- Construct phase portraits\n",
    "- Fit dynamical models to empirical data\n",
    "\n",
    "**Required packages:** `numpy`, `matplotlib`, `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 numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize import curve_fit\n",
    "\n",
    "plt.rcParams['figure.figsize'] = (10, 6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 1: Define the Logistic ODE and Implement Euler's Method\n",
    "\n",
    "First, define a Python function for the right-hand side of the logistic ODE:\n",
    "\n",
    "$$f(x) = r \\cdot x \\cdot \\left(1 - \\frac{x}{K}\\right)$$\n",
    "\n",
    "Then implement Euler's method to solve the ODE numerically. Recall that Euler's method updates the state as:\n",
    "\n",
    "$$x_{n+1} = x_n + \\Delta t \\cdot f(x_n)$$\n",
    "\n",
    "Write a function `euler_logistic(r, K, x0, dt, n_steps)` that returns arrays of time values and corresponding $x$ values.\n",
    "\n",
    "Use the following parameters for your first run: $r = 0.2$, $K = 45$, $x_0 = 1.5$, $\\Delta t = 0.1$, and enough steps to cover 30 time units."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def logistic_ode(x, r, K):\n",
    "    \"\"\"Right-hand side of the logistic ODE: dx/dt = r * x * (1 - x/K)\"\"\"\n",
    "    return r * x * (1 - x / K)\n",
    "\n",
    "\n",
    "def euler_logistic(r, K, x0, dt, n_steps):\n",
    "    \"\"\"Solve the logistic ODE using Euler's method.\"\"\"\n",
    "    t = np.zeros(n_steps + 1)\n",
    "    x = np.zeros(n_steps + 1)\n",
    "    x[0] = x0\n",
    "    for n in range(n_steps):\n",
    "        x[n + 1] = x[n] + dt * logistic_ode(x[n], r, K)\n",
    "        t[n + 1] = t[n] + dt\n",
    "    return t, x"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "r, K, x0, dt = 0.2, 45, 1.5, 0.1\n",
    "n_steps = int(30 / dt)\n",
    "t, x = euler_logistic(r, K, x0, dt, n_steps)\n",
    "\n",
    "plt.plot(t, x, lw=2)\n",
    "plt.axhline(K, color='gray', ls='--', label=f'K = {K}')\n",
    "plt.xlabel('Time')\n",
    "plt.ylabel('Response rate x(t)')\n",
    "plt.title(\"Logistic acquisition (Euler's method)\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 2: Explore the Effect of Parameters and Initial Conditions\n",
    "\n",
    "Create a figure with two subplots:\n",
    "\n",
    "**Subplot 1 - Varying initial conditions:** Fix $r = 0.2$ and $K = 45$. Plot the acquisition curve for $x_0 \\in \\{0.5, 1.5, 5.0, 15.0\\}$. All trajectories should converge to $K$, but they will take different amounts of time to get there.\n",
    "\n",
    "**Subplot 2 - Varying growth rate:** Fix $K = 45$ and $x_0 = 1.5$. Plot the acquisition curve for $r \\in \\{0.1, 0.2, 0.4, 0.8\\}$. How does $r$ affect the steepness and timing of the inflection point?\n",
    "\n",
    "Label your axes and include legends."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, (axA, axB) = plt.subplots(1, 2, figsize=(14, 5))\n",
    "\n",
    "for x0_i in [0.5, 1.5, 5.0, 15.0]:\n",
    "    t, x = euler_logistic(0.2, 45, x0_i, 0.1, 300)\n",
    "    axA.plot(t, x, label=f'x0 = {x0_i}')\n",
    "axA.axhline(45, color='gray', ls='--')\n",
    "axA.set_xlabel('Time'); axA.set_ylabel('x(t)')\n",
    "axA.set_title('Varying initial condition (r=0.2, K=45)'); axA.legend()\n",
    "\n",
    "for r_i in [0.1, 0.2, 0.4, 0.8]:\n",
    "    t, x = euler_logistic(r_i, 45, 1.5, 0.1, 300)\n",
    "    axB.plot(t, x, label=f'r = {r_i}')\n",
    "axB.axhline(45, color='gray', ls='--')\n",
    "axB.set_xlabel('Time'); axB.set_ylabel('x(t)')\n",
    "axB.set_title('Varying growth rate (K=45, x0=1.5)'); axB.legend()\n",
    "\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question:** In behavioral terms, what does $r$ represent? What does $K$ represent? How might these parameters relate to schedule of reinforcement and response ceiling effects?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Answer.** $r$ is the **acquisition rate** -- how quickly responding accelerates early in learning, before the ceiling is approached. Richer/denser reinforcement schedules (e.g., continuous reinforcement) tend to produce larger $r$. $K$ is the **asymptotic response rate** -- the ceiling set by the schedule, the response topography, and motor/temporal constraints. Once behavior reaches $K$, added practice yields little further increase (a ceiling effect)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 3: Find Equilibrium Points Analytically\n",
    "\n",
    "The equilibrium points (also called fixed points or steady states) are values of $x^*$ where $\\frac{dx}{dt} = 0$.\n",
    "\n",
    "Set $f(x^*) = r \\cdot x^* \\cdot (1 - x^*/K) = 0$ and solve for $x^*$.\n",
    "\n",
    "1. Write out the equilibrium points analytically in the markdown cell below.\n",
    "2. Verify your answer by evaluating `logistic_ode(x_star, r, K)` for each equilibrium and confirming the result is zero (or numerically very close to zero)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Equilibrium points.** $f(x^*) = r\\,x^*(1 - x^*/K) = 0$ requires $x^* = 0$ **or** $1 - x^*/K = 0 \\Rightarrow x^* = K$. So the two equilibria are $x^* = 0$ and $x^* = K = 45$, confirmed numerically below."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "r, K = 0.2, 45\n",
    "for x_star in [0.0, K]:\n",
    "    print(f\"f({x_star}) = {logistic_ode(x_star, r, K):.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 4: Linear Stability Analysis\n",
    "\n",
    "To determine whether each equilibrium is stable or unstable, compute the derivative of $f(x)$ with respect to $x$:\n",
    "\n",
    "$$f'(x) = \\frac{d}{dx}\\left[r \\cdot x \\cdot \\left(1 - \\frac{x}{K}\\right)\\right]$$\n",
    "\n",
    "Evaluate $f'(x^*)$ at each equilibrium point:\n",
    "- If $f'(x^*) < 0$, the equilibrium is **stable** (perturbations decay).\n",
    "- If $f'(x^*) > 0$, the equilibrium is **unstable** (perturbations grow).\n",
    "\n",
    "1. Derive $f'(x)$ analytically.\n",
    "2. Evaluate it at each equilibrium for $r = 0.2$ and $K = 45$.\n",
    "3. State which equilibrium is stable and which is unstable, and explain what this means for the acquisition of behavior."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Derivation.** $f(x) = rx - \\tfrac{r}{K}x^2$, so $f'(x) = r - \\tfrac{2r}{K}x = r\\left(1 - \\tfrac{2x}{K}\\right)$. At $x^*=0$: $f'(0) = r = 0.2 > 0$ -> **unstable**. At $x^*=K$: $f'(K) = r(1-2) = -r = -0.2 < 0$ -> **stable**. Behaviorally, zero responding is an unstable starting point (any responding grows), and the asymptote $K$ is the stable steady state acquisition settles into."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def logistic_ode_derivative(x, r, K):\n",
    "    \"\"\"f'(x) = r * (1 - 2x/K)\"\"\"\n",
    "    return r * (1 - 2 * x / K)\n",
    "\n",
    "r, K = 0.2, 45\n",
    "for x_star in [0.0, K]:\n",
    "    fp = logistic_ode_derivative(x_star, r, K)\n",
    "    kind = 'stable' if fp < 0 else 'unstable'\n",
    "    print(f\"f'({x_star}) = {fp:+.3f}  ->  {kind}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 5: Create a Phase Portrait\n",
    "\n",
    "A phase portrait plots $\\frac{dx}{dt}$ on the y-axis against $x$ on the x-axis. This shows at a glance how the rate of change depends on the current state.\n",
    "\n",
    "1. Create an array of $x$ values from $-2$ to $55$ (going slightly beyond $K$).\n",
    "2. Compute $f(x)$ for each value.\n",
    "3. Plot $f(x)$ vs. $x$.\n",
    "4. Mark the equilibrium points with dots.\n",
    "5. Add arrows along the x-axis indicating the direction of flow:\n",
    "   - Where $f(x) > 0$, $x$ is increasing (rightward arrows).\n",
    "   - Where $f(x) < 0$, $x$ is decreasing (leftward arrows).\n",
    "6. Draw a horizontal dashed line at $f(x) = 0$.\n",
    "\n",
    "**Hint:** Use `plt.annotate` with `arrowprops` or `plt.quiver` for the arrows."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "r, K = 0.2, 45\n",
    "xs = np.linspace(-2, 55, 400)\n",
    "fx = logistic_ode(xs, r, K)\n",
    "\n",
    "plt.plot(xs, fx, lw=2)\n",
    "plt.axhline(0, color='k', ls='--', lw=1)\n",
    "plt.plot([0, K], [0, 0], 'ro', ms=10, label='equilibria')\n",
    "\n",
    "# flow arrows along the x-axis\n",
    "for x_arrow in np.linspace(-1, 54, 18):\n",
    "    direction = np.sign(logistic_ode(x_arrow, r, K))\n",
    "    if direction != 0:\n",
    "        plt.annotate('', xy=(x_arrow + direction * 1.5, 0), xytext=(x_arrow, 0),\n",
    "                     arrowprops=dict(arrowstyle='->', color='green', lw=1.5))\n",
    "\n",
    "plt.xlabel('x (response rate)')\n",
    "plt.ylabel('dx/dt = f(x)')\n",
    "plt.title('Phase portrait of the logistic ODE')\n",
    "plt.legend(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question:** Looking at the phase portrait, explain in your own words why $x^* = K$ is a stable equilibrium and $x^* = 0$ is unstable (for $x > 0$). What happens if $x$ is slightly above $K$?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Answer.** For $0 < x < K$, $f(x) > 0$, so $x$ increases toward $K$; for $x > K$, $f(x) < 0$, so $x$ decreases back toward $K$. Flow on both sides points at $K$, making it stable. At $x = 0$ the flow points away (any small positive $x$ grows), making $0$ unstable. If $x$ is slightly above $K$, the negative rate pulls it back down to $K$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 6: The Analytical Solution\n",
    "\n",
    "The logistic ODE has a well-known analytical (closed-form) solution:\n",
    "\n",
    "$$x(t) = \\frac{K}{1 + \\left(\\frac{K - x_0}{x_0}\\right) e^{-rt}}$$\n",
    "\n",
    "1. Implement this as a Python function.\n",
    "2. Plot it alongside the Euler numerical solution from Task 1 to verify they agree.\n",
    "3. Try a larger time step (e.g., $\\Delta t = 1.0$) in Euler's method and compare. How does the step size affect accuracy?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def logistic_analytical(t, r, K, x0):\n",
    "    \"\"\"Analytical solution of the logistic ODE.\"\"\"\n",
    "    return K / (1 + ((K - x0) / x0) * np.exp(-r * t))"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "r, K, x0 = 0.2, 45, 1.5\n",
    "t_fine, x_euler = euler_logistic(r, K, x0, 0.1, 300)\n",
    "t_coarse, x_euler_coarse = euler_logistic(r, K, x0, 1.0, 30)\n",
    "x_exact = logistic_analytical(t_fine, r, K, x0)\n",
    "\n",
    "plt.plot(t_fine, x_exact, 'k-', lw=2, label='analytical')\n",
    "plt.plot(t_fine, x_euler, 'b--', label='Euler dt=0.1')\n",
    "plt.plot(t_coarse, x_euler_coarse, 'r:o', ms=3, label='Euler dt=1.0')\n",
    "plt.xlabel('Time'); plt.ylabel('x(t)')\n",
    "plt.title('Analytical vs. Euler (step-size effect)')\n",
    "plt.legend(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 7: Fit the Logistic Model to Empirical Acquisition Data\n",
    "\n",
    "Load `acquisition_data.csv`, which contains response rates across 30 sessions for a single participant learning a new operant response.\n",
    "\n",
    "1. Load the data and plot the raw acquisition curve.\n",
    "2. Use `scipy.optimize.curve_fit` to fit the analytical logistic solution to the data. You will need to define a wrapper function that takes `t` as the first argument and the parameters to fit ($r$, $K$, $x_0$) as subsequent arguments.\n",
    "3. Provide reasonable initial guesses: $r_0 = 0.3$, $K_0 = 40$, $x_{0,0} = 1.0$.\n",
    "4. Report the fitted parameter values and their standard errors.\n",
    "5. Plot the data and the fitted curve together."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "acq_data = pd.read_csv('acquisition_data.csv')\n",
    "\n",
    "plt.plot(acq_data['session'], acq_data['response_rate'], 'o')\n",
    "plt.xlabel('Session'); plt.ylabel('Response rate')\n",
    "plt.title('Observed acquisition data'); plt.show()\n",
    "acq_data.head()"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "sessions = acq_data['session'].values\n",
    "rates = acq_data['response_rate'].values\n",
    "\n",
    "def logistic_fit(t, r, K, x0):\n",
    "    return logistic_analytical(t, r, K, x0)\n",
    "\n",
    "popt, pcov = curve_fit(logistic_fit, sessions, rates,\n",
    "                       p0=[0.3, 40, 1.0], maxfev=10000)\n",
    "perr = np.sqrt(np.diag(pcov))"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "r_hat, K_hat, x0_hat = popt\n",
    "print(f\"r  = {r_hat:.4f} +/- {perr[0]:.4f}\")\n",
    "print(f\"K  = {K_hat:.4f} +/- {perr[1]:.4f}\")\n",
    "print(f\"x0 = {x0_hat:.4f} +/- {perr[2]:.4f}\")\n",
    "\n",
    "t_smooth = np.linspace(sessions.min(), sessions.max(), 200)\n",
    "plt.plot(sessions, rates, 'o', label='data')\n",
    "plt.plot(t_smooth, logistic_fit(t_smooth, *popt), 'r-', lw=2, label='logistic fit')\n",
    "plt.xlabel('Session'); plt.ylabel('Response rate')\n",
    "plt.title('Logistic fit to acquisition data'); plt.legend(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 8: Assess Model Fit\n",
    "\n",
    "1. Compute the residuals (observed - predicted) and plot them against session number.\n",
    "2. Compute $R^2$ for the fit.\n",
    "3. Are the residuals randomly scattered, or do you see systematic patterns? What might this suggest about the adequacy of the logistic model for these data?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "predicted = logistic_fit(sessions, *popt)\n",
    "residuals = rates - predicted\n",
    "\n",
    "ss_res = np.sum(residuals ** 2)\n",
    "ss_tot = np.sum((rates - rates.mean()) ** 2)\n",
    "r_squared = 1 - ss_res / ss_tot\n",
    "print(f\"R-squared = {r_squared:.4f}\")\n",
    "\n",
    "plt.axhline(0, color='gray', ls='--')\n",
    "plt.plot(sessions, residuals, 'o-')\n",
    "plt.xlabel('Session'); plt.ylabel('Residual (observed - predicted)')\n",
    "plt.title(f'Residuals (R^2 = {r_squared:.3f})'); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 9: Behavioral Interpretation\n",
    "\n",
    "Discuss the following in the cell below:\n",
    "\n",
    "1. What is the estimated carrying capacity $K$? How does this compare to what you would expect as a maximum response rate for an operant response?\n",
    "2. What is the estimated growth rate $r$? How might this parameter change under different reinforcement schedules (e.g., continuous vs. intermittent reinforcement)?\n",
    "3. What is the estimated initial response rate $x_0$? Does it make sense as a starting level of behavior?\n",
    "4. What are the limitations of using the logistic model for behavioral acquisition? Can you think of behavioral phenomena it would fail to capture (e.g., extinction, ratio strain, contrast effects)?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Discussion.** The fit gives $K \\approx 45$ responses/min (a plausible operant ceiling), $r \\approx 0.3$ (moderate acquisition speed -- expected to be larger under continuous reinforcement and smaller under lean intermittent schedules), and $x_0 \\approx 1$ (a low but non-zero starting rate, as expected before learning). Limitations: the logistic curve is monotonic and saturating, so it cannot capture extinction or decline after acquisition, ratio strain, behavioral contrast, spontaneous recovery, or non-monotonic dynamics. It also assumes a single fixed ceiling and symmetric approach."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reflection\n",
    "\n",
    "In a few sentences, address the following:\n",
    "\n",
    "1. How does thinking about behavior in terms of differential equations change your perspective compared to static curve-fitting?\n",
    "2. What is the advantage of having both a numerical method (Euler) and an analytical solution? When might you only have one or the other?\n",
    "3. How could you extend the logistic model to capture more complex behavioral dynamics (e.g., extinction following acquisition, or competing responses)?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Reflection.** Framing behavior as a differential equation shifts attention from the shape of a static curve to the *rule of change* -- how the current rate of responding drives its own future change -- which makes equilibria and stability natural objects of analysis. Having both Euler and the closed form lets you cross-check accuracy and trust the numerics on models that have no closed form; for most realistic (nonlinear, multi-variable) behavioral systems only a numerical method is available. The model could be extended with a time-varying or reduced $K$ to model extinction, or coupled logistic equations to model competing responses."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}