{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Week 7 Lab: How to Construct a Model\n",
    "\n",
    "This week is about building a quantitative model of a phenomenon from scratch:\n",
    "map the variables, write the governing equation, express it as recursive,\n",
    "difference, and differential forms, fit it to data, and run simulations to probe\n",
    "its behavior.\n",
    "\n",
    "This notebook works two complete examples end to end. Use them as templates for\n",
    "the **two environment-behavior relations of your own choosing** that the\n",
    "assignment asks you to model.\n",
    "\n",
    "- **Part A** - Lever pressing under a random-ratio (RR) schedule, using the\n",
    "  discrete recursion model $n_{t+1} = (\\alpha p - \\beta)\\, n_t$.\n",
    "- **Part B** - Salivary responding to a conditioned stimulus (CS), using the\n",
    "  acquisition model $R_{t+1} = R_t + \\gamma I\\,(1 - R_t)$.\n",
    "\n",
    "For each, we fit the model to data and/or simulate how the parameters shape the\n",
    "predicted behavior."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Import `numpy`, `matplotlib`, and `seaborn`, and define the shared simulation parameters (time grid and saturation strength) used later."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Shows equilibria and dynamics for:\n",
    "# 1) Continuous-time linear model: dn/dt = (\u03b1p - \u03b2) n\n",
    "# 2) Discrete-time recursion EXACTLY as specified: n_{t+1} = (\u03b1p - \u03b2) n_t\n",
    "# 3) Saturated (logistic-like) model: dn/dt = (\u03b1p - \u03b2) n - \u03b3 n^2 (finite equilibrium when \u03b1p>\u03b2)\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "# ---------- Parameters ----------\n",
    "n0 = 100\n",
    "t_end = 10.0\n",
    "dt = 0.01\n",
    "t = np.arange(0.0, t_end + dt, dt)\n",
    "\n",
    "# For saturated model\n",
    "gamma = 0.8  # saturation strength (choose >0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part A: Lever Pressing under a Random-Ratio Schedule\n",
    "\n",
    "The discrete recursion model is $n_{t+1} = (\\alpha p - \\beta)\\, n_t$, where $\\alpha$ amplifies responding per reinforcer, $\\beta$ is baseline decay, and $p$ is reinforcement probability. The single multiplier $(\\alpha p - \\beta)$ governs growth vs. decay.\n",
    "\n",
    "### Task A1: Simulate animal data\n",
    "\n",
    "Generate a noisy \"observed\" response trajectory from known parameters so we have something to fit."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# First, let's create some simulated \"animal data\" to demonstrate fitting\n",
    "np.random.seed(42)\n",
    "T_data = 20\n",
    "true_alpha = 1.8\n",
    "true_beta = 0.9\n",
    "true_p = 0.6\n",
    "n0_true = 150\n",
    "\n",
    "# Generate \"true\" trajectory\n",
    "factor_true = true_alpha * true_p - true_beta\n",
    "n_true = np.zeros(T_data + 1)\n",
    "n_true[0] = n0_true\n",
    "for k in range(T_data):\n",
    "    n_true[k+1] = factor_true * n_true[k]\n",
    "\n",
    "# Add measurement noise to simulate real data\n",
    "noise_level = 0.15\n",
    "observed_data = n_true + np.random.normal(0, noise_level * np.abs(n_true), len(n_true))\n",
    "observed_data = np.maximum(observed_data, 0)  # Ensure non-negative responses\n",
    "\n",
    "time_points = np.arange(0, T_data + 1)\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.plot(time_points, n_true, 'b-', linewidth=2, label='True model')\n",
    "plt.scatter(time_points, observed_data, color='red', s=50, alpha=0.7, label='Observed data')\n",
    "plt.title('Simulated Animal Response Data')\n",
    "plt.xlabel('Time (sessions)')\n",
    "plt.ylabel('Response count')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.3)\n",
    "sns.despine()\n",
    "plt.show()\n",
    "\n",
    "print(f\"True parameters: \u03b1={true_alpha}, \u03b2={true_beta}, p={true_p}\")\n",
    "print(f\"True factor (\u03b1p-\u03b2): {factor_true:.3f}\")\n",
    "print(f\"Data range: {observed_data.min():.1f} to {observed_data.max():.1f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task A2: Fit by least squares\n",
    "\n",
    "Define the model and a sum-of-squared-error objective, then minimize it with `scipy.optimize.minimize` (bounded). Recover $\\alpha, \\beta, p$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from scipy.optimize import minimize\n",
    "\n",
    "def discrete_recursion_model(params, time_points, n0):\n",
    "    \"\"\"\n",
    "    Generate predictions from the discrete recursion model\n",
    "    \n",
    "    Parameters:\n",
    "    - params: [alpha, beta, p] - model parameters to fit\n",
    "    - time_points: array of time points\n",
    "    - n0: initial value\n",
    "    \n",
    "    Returns:\n",
    "    - predicted values at each time point\n",
    "    \"\"\"\n",
    "    alpha, beta, p = params\n",
    "    factor = alpha * p - beta\n",
    "    \n",
    "    n_pred = np.zeros(len(time_points))\n",
    "    n_pred[0] = n0\n",
    "    \n",
    "    for k in range(len(time_points) - 1):\n",
    "        n_pred[k+1] = factor * n_pred[k]\n",
    "    \n",
    "    return n_pred\n",
    "\n",
    "def sum_squared_error(params, observed_data, time_points, n0):\n",
    "    \"\"\"Calculate sum of squared errors between observed and predicted\"\"\"\n",
    "    predicted = discrete_recursion_model(params, time_points, n0)\n",
    "    return np.sum((observed_data - predicted)**2)\n",
    "\n",
    "# Initial parameter guesses\n",
    "initial_guess = [1.5, 0.8, 0.5]  # [alpha, beta, p]\n",
    "\n",
    "# Parameter bounds (all should be positive)\n",
    "bounds = [(0.1, 5.0),   # alpha bounds\n",
    "          (0.1, 5.0),   # beta bounds  \n",
    "          (0.01, 0.99)] # p bounds (probability, so 0-1)\n",
    "\n",
    "# Fit the model\n",
    "result = minimize(sum_squared_error, \n",
    "                 initial_guess, \n",
    "                 args=(observed_data, time_points, observed_data[0]),\n",
    "                 bounds=bounds,\n",
    "                 method='L-BFGS-B')\n",
    "\n",
    "fitted_params = result.x\n",
    "alpha_fit, beta_fit, p_fit = fitted_params\n",
    "\n",
    "print(\"Least Squares Fitting Results:\")\n",
    "print(f\"Fitted \u03b1: {alpha_fit:.3f} (true: {true_alpha})\")\n",
    "print(f\"Fitted \u03b2: {beta_fit:.3f} (true: {true_beta})\")\n",
    "print(f\"Fitted p: {p_fit:.3f} (true: {true_p})\")\n",
    "print(f\"Fitted factor (\u03b1p-\u03b2): {alpha_fit*p_fit - beta_fit:.3f} (true: {factor_true:.3f})\")\n",
    "print(f\"Sum of squared errors: {result.fun:.2f}\")\n",
    "print(f\"Optimization successful: {result.success}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Plot the fitted model against the data\n",
    "fitted_predictions = discrete_recursion_model(fitted_params, time_points, observed_data[0])\n",
    "\n",
    "plt.figure(figsize=(12, 6))\n",
    "plt.scatter(time_points, observed_data, color='red', s=60, alpha=0.7, label='Observed data', zorder=3)\n",
    "plt.plot(time_points, n_true, 'b-', linewidth=2, label='True model', alpha=0.8)\n",
    "plt.plot(time_points, fitted_predictions, 'g--', linewidth=2, label='Fitted model', alpha=0.8)\n",
    "plt.title('Model Fitting Results: Discrete Recursion')\n",
    "plt.xlabel('Time (sessions)')\n",
    "plt.ylabel('Response count')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.3)\n",
    "sns.despine()\n",
    "plt.show()\n",
    "\n",
    "# Calculate R-squared\n",
    "ss_res = np.sum((observed_data - fitted_predictions)**2)\n",
    "ss_tot = np.sum((observed_data - np.mean(observed_data))**2)\n",
    "r_squared = 1 - (ss_res / ss_tot)\n",
    "print(f\"R-squared: {r_squared:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task A3: Fit by maximum likelihood (Poisson)\n",
    "\n",
    "For count data, assume a Poisson error distribution and minimize the negative log-likelihood instead."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from scipy.stats import poisson\n",
    "\n",
    "def negative_log_likelihood_poisson(params, observed_data, time_points, n0):\n",
    "    \"\"\"\n",
    "    Calculate negative log-likelihood assuming Poisson distribution\n",
    "    \"\"\"\n",
    "    predicted = discrete_recursion_model(params, time_points, n0)\n",
    "    \n",
    "    # Ensure predictions are positive for Poisson\n",
    "    predicted = np.maximum(predicted, 0.1)\n",
    "    \n",
    "    # Calculate negative log-likelihood\n",
    "    nll = 0\n",
    "    for obs, pred in zip(observed_data, predicted):\n",
    "        if obs >= 0:  # Only include valid observations\n",
    "            nll -= poisson.logpmf(int(obs), pred)\n",
    "    \n",
    "    return nll\n",
    "\n",
    "# Fit using maximum likelihood\n",
    "result_mle = minimize(negative_log_likelihood_poisson, \n",
    "                     initial_guess, \n",
    "                     args=(observed_data, time_points, observed_data[0]),\n",
    "                     bounds=bounds,\n",
    "                     method='L-BFGS-B')\n",
    "\n",
    "fitted_params_mle = result_mle.x\n",
    "alpha_mle, beta_mle, p_mle = fitted_params_mle\n",
    "\n",
    "print(\"Maximum Likelihood Estimation Results:\")\n",
    "print(f\"Fitted \u03b1: {alpha_mle:.3f} (true: {true_alpha})\")\n",
    "print(f\"Fitted \u03b2: {beta_mle:.3f} (true: {true_beta})\")\n",
    "print(f\"Fitted p: {p_mle:.3f} (true: {true_p})\")\n",
    "print(f\"Fitted factor (\u03b1p-\u03b2): {alpha_mle*p_mle - beta_mle:.3f} (true: {factor_true:.3f})\")\n",
    "print(f\"Negative log-likelihood: {result_mle.fun:.2f}\")\n",
    "print(f\"Optimization successful: {result_mle.success}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task A4: Evaluate goodness of fit\n",
    "\n",
    "Compare the least-squares and MLE fits using R-squared, RMSE, MAE, and AIC."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Compare both fitting methods\n",
    "fitted_predictions_mle = discrete_recursion_model(fitted_params_mle, time_points, observed_data[0])\n",
    "\n",
    "# Calculate multiple goodness-of-fit metrics\n",
    "def calculate_fit_metrics(observed, predicted):\n",
    "    \"\"\"Calculate various goodness-of-fit metrics\"\"\"\n",
    "    residuals = observed - predicted\n",
    "    \n",
    "    # R-squared\n",
    "    ss_res = np.sum(residuals**2)\n",
    "    ss_tot = np.sum((observed - np.mean(observed))**2)\n",
    "    r_squared = 1 - (ss_res / ss_tot)\n",
    "    \n",
    "    # Root Mean Square Error\n",
    "    rmse = np.sqrt(np.mean(residuals**2))\n",
    "    \n",
    "    # Mean Absolute Error\n",
    "    mae = np.mean(np.abs(residuals))\n",
    "    \n",
    "    # Akaike Information Criterion (AIC) - for model comparison\n",
    "    n = len(observed)\n",
    "    k = 3  # number of parameters (alpha, beta, p)\n",
    "    aic = n * np.log(ss_res/n) + 2*k\n",
    "    \n",
    "    return {\n",
    "        'R_squared': r_squared,\n",
    "        'RMSE': rmse,\n",
    "        'MAE': mae,\n",
    "        'AIC': aic\n",
    "    }\n",
    "\n",
    "# Calculate metrics for both methods\n",
    "metrics_ls = calculate_fit_metrics(observed_data, fitted_predictions)\n",
    "metrics_mle = calculate_fit_metrics(observed_data, fitted_predictions_mle)\n",
    "\n",
    "print(\"Model Evaluation Metrics:\")\n",
    "print(\"\\nLeast Squares Method:\")\n",
    "for metric, value in metrics_ls.items():\n",
    "    print(f\"  {metric}: {value:.4f}\")\n",
    "\n",
    "print(\"\\nMaximum Likelihood Method:\")\n",
    "for metric, value in metrics_mle.items():\n",
    "    print(f\"  {metric}: {value:.4f}\")\n",
    "\n",
    "# Compare both fits visually\n",
    "plt.figure(figsize=(14, 6))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.scatter(time_points, observed_data, color='red', s=60, alpha=0.7, label='Observed data')\n",
    "plt.plot(time_points, fitted_predictions, 'g--', linewidth=2, label='Least Squares fit')\n",
    "plt.plot(time_points, n_true, 'b-', linewidth=1, alpha=0.6, label='True model')\n",
    "plt.title(f'Least Squares Fit (R\u00b2 = {metrics_ls[\"R_squared\"]:.3f})')\n",
    "plt.xlabel('Time (sessions)')\n",
    "plt.ylabel('Response count')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.3)\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.scatter(time_points, observed_data, color='red', s=60, alpha=0.7, label='Observed data')\n",
    "plt.plot(time_points, fitted_predictions_mle, 'm--', linewidth=2, label='MLE fit')\n",
    "plt.plot(time_points, n_true, 'b-', linewidth=1, alpha=0.6, label='True model')\n",
    "plt.title(f'MLE Fit (R\u00b2 = {metrics_mle[\"R_squared\"]:.3f})')\n",
    "plt.xlabel('Time (sessions)')\n",
    "plt.ylabel('Response count')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "sns.despine()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task A5: Simulate the dynamics\n",
    "\n",
    "Now explore the model itself. Plot the discrete recursion for several values of $p$, then compare the continuous-time and saturated (logistic-like) versions of the same relation."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ---------- 1) Discrete-time recursion: n_{t+1} = (\u03b1p \u2212 \u03b2) n_t ----------\n",
    "T = 12\n",
    "steps = np.arange(0, T + 1)\n",
    "n0 = 100\n",
    "alpha = 2.0 # Twice as much as beta\n",
    "beta = 1.0\n",
    "p_values = [0.3, 0.5, 0.8]\n",
    "\n",
    "plt.figure(figsize=(9, 6))\n",
    "for p in p_values:\n",
    "    factor = alpha * p - beta  # multiplier for each step\n",
    "    n = np.zeros(T + 1)\n",
    "    n[0] = n0\n",
    "    for k in range(T):\n",
    "        n[k+1] = factor * n[k]\n",
    "    plt.plot(steps, n, marker='o', label=f\"p={p}  (\u03b1p-\u03b2={factor:.2f})\")\n",
    "plt.axhline(0.0, linestyle=\"--\")\n",
    "plt.title(\"Discrete recursion: n_{t+1} = (\u03b1p \u2212 \u03b2) n_t\")\n",
    "plt.xlabel(\"Step t\", fontsize=14, labelpad=12)\n",
    "plt.ylabel(\"n_t\", fontsize=14, labelpad=12)\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "sns.despine(top=True, right=True)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def plot_discrete_recursion(p_values, alpha, beta, n0, T):\n",
    "    steps = np.arange(T + 1)\n",
    "    plt.figure(figsize=(9, 6))\n",
    "    \n",
    "    for p in p_values:\n",
    "        factor = alpha * p - beta  # multiplier for each step\n",
    "        n = np.zeros(T + 1)\n",
    "        n[0] = n0\n",
    "        for k in range(T):\n",
    "            n[k+1] = factor * n[k]\n",
    "        plt.plot(steps, n, marker='o', label=f\"p={p}  (\u03b1p-\u03b2={factor:.2f})\")\n",
    "    \n",
    "    plt.axhline(0.0, linestyle=\"--\")\n",
    "    plt.title(r\"Discrete recursion: $n_{t+1} = (\\alpha p \u2212 \\beta) n_t$\")\n",
    "    plt.xlabel(\"Step t\", fontsize=14, labelpad=12)\n",
    "    plt.ylabel(r\"$n_t$\", fontsize=14, labelpad=12)\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    sns.despine(top=True, right=True)\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def plot_discrete_recursion_euler(p_values, alpha, beta, n0, T):\n",
    "    \"\"\"\n",
    "    Discrete-time Euler update of dn/dt = (\u03b1p \u2212 \u03b2) n with \u0394t=1:\n",
    "        n_{t+1} = (1 + (\u03b1p \u2212 \u03b2)) * n_t\n",
    "    \"\"\"\n",
    "    steps = np.arange(T + 1)\n",
    "    plt.figure(figsize=(9, 6))\n",
    "    for p in p_values:\n",
    "        r = alpha * p - beta\n",
    "        mult = 1.0 + r\n",
    "        n = np.zeros(T + 1)\n",
    "        n[0] = n0\n",
    "        for k in range(T):\n",
    "            n[k+1] = mult * n[k]\n",
    "        plt.plot(steps, n, marker='o', label=f\"p={p}  r={r:.2f}, multiplier={mult:.2f}\")\n",
    "    plt.axhline(0.0, linestyle=\"--\")\n",
    "    plt.title(r\"Discrete (Euler from ODE): $n_{t+1} = (1 + (\\alpha p - \\beta))\\,n_t$\")\n",
    "    plt.xlabel(\"Step t\", fontsize=14, labelpad=12)\n",
    "    plt.ylabel(r\"$n_t$\", fontsize=14, labelpad=12)\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    sns.despine(top=True, right=True)\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ---------- 1) Discrete-time recursion: n_{t+1} = (\u03b1p \u2212 \u03b2) n_t ----------\n",
    "plot_discrete_recursion(\n",
    "    p_values = [0.3, 0.5, 0.8], \n",
    "    alpha = 0.5, # Half as much as beta, \n",
    "    beta = 1.0, \n",
    "    n0 = 100, \n",
    "    T = 12\n",
    ")\n",
    "\n",
    "plot_discrete_recursion_euler(\n",
    "    p_values = [0.3, 0.5, 0.8], \n",
    "    alpha = 0.5, # Half as much as beta, \n",
    "    beta = 1.0, \n",
    "    n0 = 100, \n",
    "    T = 12\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ---------- 1) Discrete-time recursion: n_{t+1} = (\u03b1p \u2212 \u03b2) n_t ----------\n",
    "plot_discrete_recursion(\n",
    "    p_values = [0.3, 0.5, 0.8], \n",
    "    alpha = 2, \n",
    "    beta = 0.95, # Slightly less than 1\n",
    "    n0 = 100, \n",
    "    T = 12\n",
    ")\n",
    "\n",
    "plot_discrete_recursion_euler(\n",
    "    p_values = [0.3, 0.5, 0.8], \n",
    "    alpha = 2, \n",
    "    beta = 0.95, # Slightly less than 1\n",
    "    n0 = 100, \n",
    "    T = 12\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ---------- 1) Discrete-time recursion: n_{t+1} = (\u03b1p \u2212 \u03b2) n_t ----------\n",
    "plot_discrete_recursion(\n",
    "    p_values = [0.3, 0.5, 0.8], \n",
    "    alpha = 2, \n",
    "    beta = 0.25, \n",
    "    n0 = 100, \n",
    "    T = 12\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Continuous-time** linear form $dn/dt = (\\alpha p - \\beta)\\, n$ (analytical solution):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ---------- 2) Continuous-time linear: dn/dt = (\u03b1p - \u03b2) n ----------\n",
    "plt.figure(figsize=(9, 6))\n",
    "for p in p_values:\n",
    "    r = alpha * p - beta  # net growth rate\n",
    "    n = n0 * np.exp(r * t)  # analytical solution\n",
    "    plt.plot(t, n, label=f\"p={p}  (\u03b1p-\u03b2={r:.2f})\")\n",
    "    # Equilibria: n*=0 (only finite eq). Show baseline line at 0.\n",
    "plt.axhline(0.0, linestyle=\"--\")\n",
    "plt.title(\"Continuous-time: dn/dt = (\u03b1p \u2212 \u03b2) n\")\n",
    "plt.xlabel(\"Time t\", fontsize=14, labelpad=12)\n",
    "plt.ylabel(\"n(t)\", fontsize=14, labelpad=12)\n",
    "# plt.yscale(\"log\")\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "sns.despine(top=True, right=True)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Saturated** form $dn/dt = (\\alpha p - \\beta)\\, n - \\gamma n^2$ (finite equilibrium):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ---------- 3) Saturated model: dn/dt = (\u03b1p \u2212 \u03b2) n \u2212 \u03b3 n^2 ----------\n",
    "plt.figure(figsize=(9, 6))\n",
    "for p in p_values:\n",
    "    r = alpha * p - beta\n",
    "    # Numerical integration (Euler) for the nonlinear case\n",
    "    n = np.zeros_like(t)\n",
    "    n[0] = n0\n",
    "    for i in range(len(t) - 1):\n",
    "        dn = (r * n[i] - gamma * n[i] * n[i]) * dt\n",
    "        n[i+1] = max(0.0, n[i] + dn)  # enforce non-negativity for visualization\n",
    "    # Equilibrium: n* = max{0, r/gamma}\n",
    "    n_star = max(0.0, r / gamma)\n",
    "    plt.plot(t, n, label=f\"p={p}  n*={n_star:.2f}\")\n",
    "    plt.axhline(n_star, linestyle=\"--\")\n",
    "plt.title(\"Saturated multiplicative model: dn/dt = (\u03b1p \u2212 \u03b2) n \u2212 \u03b3 n\u00b2\")\n",
    "plt.xlabel(\"Time t\")\n",
    "plt.ylabel(\"n(t)\")\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part B: Salivary Responding to a Conditioned Stimulus\n",
    "\n",
    "Respondent acquisition follows $R_{t+1} = R_t + \\gamma I\\,(1 - R_t)$, where $I$ is CS intensity and $\\gamma$ is the learning rate. We show the same relation as a recursion, a difference equation (Euler with $\\Delta t \\neq 1$), and a differential equation (closed form), each while varying $I$ and $\\gamma$.\n",
    "\n",
    "### Common settings"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Compact demo: Respondent conditioning under varying CS intensity (I) and learning rate (gamma)\n",
    "# Three panels: recursion, difference (Euler with \u0394t\u22601), and differential (closed-form solution).\n",
    "# Each panel has two subplots: vary I (fixed gamma) and vary gamma (fixed I).\n",
    "\n",
    "# ------------------------ Common settings ------------------------\n",
    "R0 = 0.05                  # initial salivation magnitude\n",
    "I_vals = [0.2, 0.5, 0.8]   # CS intensities to compare\n",
    "gamma_vals = [0.2, 0.5, 0.8]  # learning rates to compare\n",
    "T_trials = 20         # number of discrete trials to simulate for recursion\n",
    "t_end = 10.0          # time horizon for differential/difference demonstrations\n",
    "dt_diff = 0.2         # \u0394t for the difference-equation Euler update (\u2260 1 to distinguish it)\n",
    "\n",
    "# Time arrays\n",
    "trials = np.arange(T_trials + 1)\n",
    "t = np.arange(0.0, t_end + 1e-9, 0.01)  # smooth time for closed form\n",
    "t_euler = np.arange(0.0, t_end + 1e-9, dt_diff)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task B1: Recursion\n",
    "\n",
    "$R_{t+1} = R_t + \\gamma I\\,(1 - R_t)$"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ------------------------ 1) Recursion ------------------------\n",
    "# R_{t+1} = R_t + \u03b3 I (1 - R_t)\n",
    "fig1, axes1 = plt.subplots(1, 2, figsize=(12, 4), sharey=True)\n",
    "fixed_gamma = 0.5\n",
    "fixed_I = 0.5\n",
    "\n",
    "# (a) Vary I at fixed gamma\n",
    "for I in I_vals:\n",
    "    R = np.zeros(T_trials + 1)\n",
    "    R[0] = R0\n",
    "    for k in range(T_trials):\n",
    "        R[k+1] = R[k] + fixed_gamma * I * (1.0 - R[k])\n",
    "    axes1[0].plot(trials, R, marker='o', label=f\"I={I}, \u03b3={fixed_gamma}\")\n",
    "axes1[0].set_title(\"Recursion: vary I (\u03b3 fixed)\")\n",
    "axes1[0].set_xlabel(\"Trial\")\n",
    "axes1[0].set_ylabel(\"R\")\n",
    "axes1[0].legend()\n",
    "\n",
    "# (b) Vary gamma at fixed I\n",
    "for gamma in gamma_vals:\n",
    "    R = np.zeros(T_trials + 1)\n",
    "    R[0] = R0\n",
    "    for k in range(T_trials):\n",
    "        R[k+1] = R[k] + gamma * fixed_I * (1.0 - R[k])\n",
    "    axes1[1].plot(trials, R, marker='o', label=f\"I={fixed_I}, \u03b3={gamma}\")\n",
    "axes1[1].set_title(\"Recursion: vary \u03b3 (I fixed)\")\n",
    "axes1[1].set_xlabel(\"Trial\")\n",
    "axes1[1].legend()\n",
    "\n",
    "fig1.suptitle(\"Respondent Conditioning \u2014 Discrete Recursion\", y=1.02, fontsize=12)\n",
    "fig1.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task B2: Difference equation (Euler)\n",
    "\n",
    "$\\Delta R_t = \\gamma I\\,(1 - R_t)\\,\\Delta t$"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ------------------------ 2) Difference equation ------------------------\n",
    "# \u0394R_t = \u03b3 I (1 - R_t) \u0394t  -> Euler update with \u0394t = dt_diff\n",
    "fig2, axes2 = plt.subplots(1, 2, figsize=(12, 4), sharey=True)\n",
    "\n",
    "# (a) Vary I at fixed gamma\n",
    "for I in I_vals:\n",
    "    R = np.zeros_like(t_euler)\n",
    "    R[0] = R0\n",
    "    for k in range(len(t_euler)-1):\n",
    "        dR = fixed_gamma * I * (1.0 - R[k]) * dt_diff\n",
    "        R[k+1] = R[k] + dR\n",
    "    axes2[0].plot(t_euler, R, marker='o', linestyle='-', label=f\"I={I}, \u03b3={fixed_gamma}\")\n",
    "axes2[0].set_title(f\"Difference (Euler) \u0394t={dt_diff}: vary I (\u03b3 fixed)\")\n",
    "axes2[0].set_xlabel(\"Time\")\n",
    "axes2[0].set_ylabel(\"R\")\n",
    "axes2[0].legend()\n",
    "\n",
    "# (b) Vary gamma at fixed I\n",
    "for gamma in gamma_vals:\n",
    "    R = np.zeros_like(t_euler)\n",
    "    R[0] = R0\n",
    "    for k in range(len(t_euler)-1):\n",
    "        dR = gamma * fixed_I * (1.0 - R[k]) * dt_diff\n",
    "        R[k+1] = R[k] + dR\n",
    "    axes2[1].plot(t_euler, R, marker='o', linestyle='-', label=f\"I={fixed_I}, \u03b3={gamma}\")\n",
    "axes2[1].set_title(f\"Difference (Euler) \u0394t={dt_diff}: vary \u03b3 (I fixed)\")\n",
    "axes2[1].set_xlabel(\"Time\")\n",
    "axes2[1].legend()\n",
    "\n",
    "fig2.suptitle(\"Respondent Conditioning \u2014 Difference Equation (Euler)\", y=1.02, fontsize=12)\n",
    "fig2.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task B3: Differential equation (closed form)\n",
    "\n",
    "$dR/dt = \\gamma I\\,(1 - R) \\;\\Rightarrow\\; R(t) = 1 - (1 - R_0)\\,e^{-\\gamma I t}$"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ------------------------ 3) Differential equation (closed-form) ------------------------\n",
    "# dR/dt = \u03b3 I (1 - R)  ->  R(t) = 1 - (1 - R0) exp(-\u03b3 I t)\n",
    "fig3, axes3 = plt.subplots(1, 2, figsize=(12, 4), sharey=True)\n",
    "\n",
    "# (a) Vary I at fixed gamma\n",
    "for I in I_vals:\n",
    "    R = 1.0 - (1.0 - R0) * np.exp(-fixed_gamma * I * t)\n",
    "    axes3[0].plot(t, R, label=f\"I={I}, \u03b3={fixed_gamma}\")\n",
    "axes3[0].set_title(\"Differential: vary I (\u03b3 fixed)\")\n",
    "axes3[0].set_xlabel(\"Time\")\n",
    "axes3[0].set_ylabel(\"R\")\n",
    "axes3[0].legend()\n",
    "\n",
    "# (b) Vary gamma at fixed I\n",
    "for gamma in gamma_vals:\n",
    "    R = 1.0 - (1.0 - R0) * np.exp(-gamma * fixed_I * t)\n",
    "    axes3[1].plot(t, R, label=f\"I={fixed_I}, \u03b3={gamma}\")\n",
    "axes3[1].set_title(\"Differential: vary \u03b3 (I fixed)\")\n",
    "axes3[1].set_xlabel(\"Time\")\n",
    "axes3[1].legend()\n",
    "\n",
    "fig3.suptitle(\"Respondent Conditioning \u2014 Differential Equation (Closed Form)\", y=1.02, fontsize=12)\n",
    "fig3.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task B4: Interactive 3D surfaces (Plotly)\n",
    "\n",
    "Visualize $R(I, \\gamma, t)$ as an animated surface over the $(I, \\gamma)$ grid for each of the three formulations."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Plotly 3D interactive meshes for respondent conditioning\n",
    "# We create three figures (recursion, difference/Euler, differential/closed-form).\n",
    "# Each figure shows a 3x3 grid over I (CS intensity) and gamma (learning rate), with\n",
    "# a time slider to scrub through trials/time. The surface z = R(I, gamma, t).\n",
    "\n",
    "import numpy as np\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "# ---------------- Common parameters ----------------\n",
    "R0 = 0.05\n",
    "I_vals = np.array([0.2, 0.5, 0.8])\n",
    "gamma_vals = np.array([0.2, 0.5, 0.8])\n",
    "\n",
    "# Build a 2D meshgrid for I, gamma\n",
    "I_grid, G_grid = np.meshgrid(I_vals, gamma_vals)  # shape (3,3)\n",
    "\n",
    "# ---------------- 1) Recursion ----------------\n",
    "# R_{t+1} = R_t + gamma * I * (1 - R_t), with trial index t = 0..T\n",
    "T_trials = 20\n",
    "\n",
    "def recursion_surface_at_trial(t_idx):\n",
    "    # Iterate t_idx steps for each (I, gamma) on the grid\n",
    "    R = np.full_like(I_grid, R0, dtype=float)\n",
    "    for _ in range(t_idx):\n",
    "        R = R + G_grid * I_grid * (1.0 - R)\n",
    "    return R\n",
    "\n",
    "frames_rec = []\n",
    "for k in range(T_trials + 1):\n",
    "    Z = recursion_surface_at_trial(k)\n",
    "    frames_rec.append(go.Frame(data=[go.Surface(z=Z, x=I_grid, y=G_grid, showscale=False)],\n",
    "                               name=f\"trial_{k}\"))\n",
    "\n",
    "fig_rec = go.Figure(\n",
    "    data=[go.Surface(z=recursion_surface_at_trial(0), x=I_grid, y=G_grid, showscale=True)],\n",
    "    layout=go.Layout(\n",
    "        height=600,\n",
    "        width=1000,\n",
    "        title=\"Respondent Conditioning \u2014 Recursion<br>z = R, x = I, y = \u03b3 (trial slider)\",\n",
    "        scene=dict(\n",
    "            xaxis_title=\"I (CS intensity)\",\n",
    "            yaxis_title=\"\u03b3 (learning rate)\",\n",
    "            zaxis_title=\"R (salivation)\",\n",
    "            zaxis=dict(range=[0, 1]),  # lock z limits\n",
    "            aspectmode='cube'\n",
    "        ),\n",
    "        updatemenus=[dict(type=\"buttons\",\n",
    "                          buttons=[dict(label=\"Play\",\n",
    "                                        method=\"animate\",\n",
    "                                        args=[None, dict(frame=dict(duration=200, redraw=True),\n",
    "                                                         fromcurrent=True, transition=dict(duration=0))]),\n",
    "                                   dict(label=\"Pause\",\n",
    "                                        method=\"animate\",\n",
    "                                        args=[[None], dict(frame=dict(duration=0, redraw=False),\n",
    "                                                           mode=\"immediate\")])])],\n",
    "        sliders=[dict(\n",
    "            steps=[dict(method=\"animate\",\n",
    "                        args=[[f\"trial_{k}\"],\n",
    "                              dict(mode=\"immediate\",\n",
    "                                   frame=dict(duration=0, redraw=True),\n",
    "                                   transition=dict(duration=0))],\n",
    "                        label=str(k))\n",
    "                   for k in range(T_trials + 1)],\n",
    "            currentvalue=dict(prefix=\"Trial: \")\n",
    "        )]\n",
    "    ),\n",
    "    frames=frames_rec\n",
    ")\n",
    "\n",
    "fig_rec.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ---------------- 2) Difference (Euler with \u0394t) ----------------\n",
    "dt = 0.2\n",
    "t_end = 10.0\n",
    "time_euler = np.arange(0.0, t_end + 1e-9, dt)\n",
    "\n",
    "def difference_surface_at_index(idx):\n",
    "    # Euler stepping up to idx for each grid point\n",
    "    R = np.full_like(I_grid, R0, dtype=float)\n",
    "    for _ in range(idx):\n",
    "        dR = G_grid * I_grid * (1.0 - R) * dt\n",
    "        R = R + dR\n",
    "    return R\n",
    "\n",
    "frames_diff = []\n",
    "for i in range(len(time_euler)):\n",
    "    Z = difference_surface_at_index(i)\n",
    "    frames_diff.append(go.Frame(data=[go.Surface(z=Z, x=I_grid, y=G_grid, showscale=False)],\n",
    "                                name=f\"tidx_{i}\"))\n",
    "\n",
    "fig_diff = go.Figure(\n",
    "    data=[go.Surface(z=difference_surface_at_index(0), x=I_grid, y=G_grid, showscale=True)],\n",
    "    layout=go.Layout(\n",
    "        height=600,\n",
    "        width=1000,\n",
    "        title=f\"Respondent Conditioning \u2014 Difference (Euler \u0394t={dt})<br>z = R, x = I, y = \u03b3 (time slider)\",\n",
    "        scene=dict(\n",
    "            xaxis_title=\"I (CS intensity)\",\n",
    "            yaxis_title=\"\u03b3 (learning rate)\",\n",
    "            zaxis_title=\"R (salivation)\",\n",
    "            zaxis=dict(range=[0, 1]),  # lock z limits\n",
    "            aspectmode='cube'\n",
    "        ),\n",
    "        updatemenus=[dict(type=\"buttons\",\n",
    "                          buttons=[dict(label=\"Play\",\n",
    "                                        method=\"animate\",\n",
    "                                        args=[None, dict(frame=dict(duration=120, redraw=True),\n",
    "                                                         fromcurrent=True, transition=dict(duration=0))]),\n",
    "                                   dict(label=\"Pause\",\n",
    "                                        method=\"animate\",\n",
    "                                        args=[[None], dict(frame=dict(duration=0, redraw=False),\n",
    "                                                           mode=\"immediate\")])])],\n",
    "        sliders=[dict(\n",
    "            steps=[dict(method=\"animate\",\n",
    "                        args=[[f\"tidx_{i}\"],\n",
    "                              dict(mode=\"immediate\",\n",
    "                                   frame=dict(duration=0, redraw=True),\n",
    "                                   transition=dict(duration=0))],\n",
    "                        label=f\"{time_euler[i]:.1f}\")\n",
    "                   for i in range(len(time_euler))],\n",
    "            currentvalue=dict(prefix=\"Time: \")\n",
    "        )]\n",
    "    ),\n",
    "    frames=frames_diff\n",
    ")\n",
    "\n",
    "fig_diff.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ---------------- 3) Differential (closed-form) ----------------\n",
    "# dR/dt = \u03b3 I (1 - R)  =>  R(t) = 1 - (1 - R0) exp(-\u03b3 I t)\n",
    "time_ct = np.linspace(0.0, 10.0, 51)\n",
    "\n",
    "def differential_surface_at_time(tval):\n",
    "    return 1.0 - (1.0 - R0) * np.exp(-G_grid * I_grid * tval)\n",
    "\n",
    "frames_ct = []\n",
    "for i, tval in enumerate(time_ct):\n",
    "    Z = differential_surface_at_time(tval)\n",
    "    frames_ct.append(go.Frame(data=[go.Surface(z=Z, x=I_grid, y=G_grid, showscale=False)],\n",
    "                              name=f\"t_{i}\"))\n",
    "\n",
    "fig_ct = go.Figure(\n",
    "    data=[go.Surface(z=differential_surface_at_time(0.0), x=I_grid, y=G_grid, showscale=True)],\n",
    "    layout=go.Layout(\n",
    "        height=600,\n",
    "        width=1000,\n",
    "        title=\"Respondent Conditioning \u2014 Differential (Closed Form)<br>z = R, x = I, y = \u03b3 (time slider)\",\n",
    "        scene=dict(\n",
    "            xaxis_title=\"I (CS intensity)\",\n",
    "            yaxis_title=\"\u03b3 (learning rate)\",\n",
    "            zaxis_title=\"R (salivation)\",\n",
    "            zaxis=dict(range=[0, 1]),  # lock z limits\n",
    "            aspectmode='cube'\n",
    "        ),\n",
    "        updatemenus=[dict(type=\"buttons\",\n",
    "                          buttons=[dict(label=\"Play\",\n",
    "                                        method=\"animate\",\n",
    "                                        args=[None, dict(frame=dict(duration=120, redraw=True),\n",
    "                                                         fromcurrent=True, transition=dict(duration=0))]),\n",
    "                                   dict(label=\"Pause\",\n",
    "                                        method=\"animate\",\n",
    "                                        args=[[None], dict(frame=dict(duration=0, redraw=False),\n",
    "                                                           mode=\"immediate\")])])],\n",
    "        sliders=[dict(\n",
    "            steps=[dict(method=\"animate\",\n",
    "                        args=[[f\"t_{i}\"],\n",
    "                              dict(mode=\"immediate\",\n",
    "                                   frame=dict(duration=0, redraw=True),\n",
    "                                   transition=dict(duration=0))],\n",
    "                        label=f\"{time_ct[i]:.1f}\")\n",
    "                   for i in range(len(time_ct))],\n",
    "            currentvalue=dict(prefix=\"Time: \")\n",
    "        )]\n",
    "    ),\n",
    "    frames=frames_ct\n",
    ")\n",
    "\n",
    "fig_ct.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrap-up\n",
    "\n",
    "You have taken two relations from a verbal description to a fitted, simulatable model. For your assignment, repeat this lifecycle for two relations of your own: draw the life-cycle diagram, write the basic equation and its recursive/difference/differential forms, fit it to a fake dataset, and simulate how the parameters change the predictions. Watch for predictions that are biologically impossible (negative or runaway response rates) - those are signals to revise the model."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}