{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulates a memory experiment (much like what we covered in class) and uses models to understand how people recognize words they've seen before\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "hit = yes to old item (correctly remember the word they saw)\n",
    "miss = no to old item (forgot they saw the word before)\n",
    "false alarm = yes to new (thought they saw a word they didnt)\n",
    "correct reject = No to new item (correctly said a new word wasnt shown)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory task results:\n",
      "Hits: 13/15 (0.87)\n",
      "Misses: 2/15 (0.13)\n",
      "False Alarms: 2/15 (0.13)\n",
      "Correct Rejections: 13/15 (0.87)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import stan\n",
    "import arviz as az\n",
    "\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "np.random.seed(2023)\n",
    "\n",
    "def simulate_memory_task(n_old, n_new, d_old=0.7, d_new=0.6, g=0.3):\n",
    "    \"\"\"Simulate data from a recognition memory task \"\"\"\n",
    "\n",
    "    detections = np.random.binomial(n=1, p=d_old, size=n_old) #detect old items\n",
    "    hits = sum(detections + (1-detections) * np.random.binomial(n=1, p=g, size=n_old)) #hit or guess\n",
    "    misses = n_old - hits\n",
    "\n",
    "    if d_new > 0:  # 2HT model\n",
    "        detections = np.random.binomial(n=1, p=d_new, size=n_new)\n",
    "        false_alarms = sum((1-detections) * np.random.binomial(n=1, p=g, size=n_new))\n",
    "    else:  # 1HT model\n",
    "        false_alarms = sum(np.random.binomial(n=1, p=g, size=n_new))\n",
    "    \n",
    "    correct_rejections = n_new - false_alarms\n",
    "    \n",
    "    return {\n",
    "        'hits': int(hits),\n",
    "        'misses': int(misses),\n",
    "        'false_alarms': int(false_alarms),\n",
    "        'correct_rejections': int(correct_rejections),\n",
    "        'n_old': n_old,\n",
    "        'n_new': n_new\n",
    "    }\n",
    "\n",
    "# Simulate data\n",
    "n_old = 15\n",
    "n_new = 15\n",
    "data = simulate_memory_task(n_old=n_old, n_new=n_new)\n",
    "\n",
    "print(\"Memory task results:\")\n",
    "print(f\"Hits: {data['hits']}/{data['n_old']} ({data['hits']/data['n_old']:.2f})\")\n",
    "print(f\"Misses: {data['misses']}/{data['n_old']} ({data['misses']/data['n_old']:.2f})\")\n",
    "print(f\"False Alarms: {data['false_alarms']}/{data['n_new']} ({data['false_alarms']/data['n_new']:.2f})\")\n",
    "print(f\"Correct Rejections: {data['correct_rejections']}/{data['n_new']} ({data['correct_rejections']/data['n_new']:.2f})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_1ht_code = \"\"\"\n",
    "data {\n",
    "  int<lower=0> n_old;      // Number of old items\n",
    "  int<lower=0> n_new;      // Number of new items\n",
    "  int<lower=0> hits;       // Number of hits\n",
    "  int<lower=0> misses;     // Number of misses\n",
    "  int<lower=0> false_alarms;       // Number of false alarms\n",
    "  int<lower=0> correct_rejections; // Number of correct rejections\n",
    "}\n",
    "\n",
    "parameters {\n",
    "  real<lower=0, upper=1> d;  // Detection parameter\n",
    "  real<lower=0, upper=1> g;  // Guessing parameter\n",
    "}\n",
    "\n",
    "model {\n",
    "  // Priors\n",
    "  d ~ beta(1, 1);  // Uninformative prior\n",
    "  g ~ beta(1, 1);  // Uninformative prior\n",
    "  \n",
    "  // Likelihood for old items\n",
    "  hits ~ binomial(n_old, d + (1-d) * g);\n",
    "  \n",
    "  // Likelihood for new items\n",
    "  false_alarms ~ binomial(n_new, g);\n",
    "}\n",
    "\n",
    "generated quantities {\n",
    "  // Predicted counts\n",
    "  int<lower=0> pred_hits = binomial_rng(n_old, d + (1-d) * g);\n",
    "  int<lower=0> pred_misses = n_old - pred_hits;\n",
    "  int<lower=0> pred_false_alarms = binomial_rng(n_new, g);\n",
    "  int<lower=0> pred_correct_rejections = n_new - pred_false_alarms;\n",
    "}\n",
    "\"\"\"\n",
    "model_2ht_code = \"\"\"\n",
    "data {\n",
    "  int<lower=0> n_old;      // Number of old items\n",
    "  int<lower=0> n_new;      // Number of new items\n",
    "  int<lower=0> hits;       // Number of hits\n",
    "  int<lower=0> misses;     // Number of misses\n",
    "  int<lower=0> false_alarms;       // Number of false alarms\n",
    "  int<lower=0> correct_rejections; // Number of correct rejections\n",
    "}\n",
    "\n",
    "parameters {\n",
    "  real<lower=0, upper=1> d_old;  // Detection parameter for old items\n",
    "  real<lower=0, upper=1> d_new;  // Detection parameter for new items\n",
    "  real<lower=0, upper=1> g;      // Guessing parameter\n",
    "}\n",
    "\n",
    "model {\n",
    "  // Priors\n",
    "  d_old ~ beta(1, 1);  // Uninformative prior\n",
    "  d_new ~ beta(1, 1);  // Uninformative prior\n",
    "  g ~ beta(1, 1);      // Uninformative prior\n",
    "  \n",
    "  // Likelihood for old items\n",
    "  hits ~ binomial(n_old, d_old + (1-d_old) * g);\n",
    "  \n",
    "  // Likelihood for new items\n",
    "  correct_rejections ~ binomial(n_new, d_new + (1-d_new) * (1-g));\n",
    "}\n",
    "\n",
    "generated quantities {\n",
    "  // Predicted counts\n",
    "  int<lower=0> pred_hits = binomial_rng(n_old, d_old + (1-d_old) * g);\n",
    "  int<lower=0> pred_misses = n_old - pred_hits;\n",
    "  int<lower=0> pred_false_alarms = binomial_rng(n_new, (1-d_new) * g);\n",
    "  int<lower=0> pred_correct_rejections = n_new - pred_false_alarms;\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Building: found in cache, done.Messages from stanc:\n",
      "Warning: The parameter g has 3 priors.\n",
      "Warning: The parameter d has 2 priors.\n",
      "Sampling:   0%\n",
      "Sampling:  25% (3000/12000)\n",
      "Sampling:  50% (6000/12000)\n",
      "Sampling:  75% (9000/12000)\n",
      "Sampling: 100% (12000/12000)\n",
      "Sampling: 100% (12000/12000), done.\n",
      "Messages received during sampling:\n",
      "  Gradient evaluation took 2.5e-05 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 0.25 seconds.\n",
      "  Adjust your expectations accordingly!\n",
      "  Gradient evaluation took 2.8e-05 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 0.28 seconds.\n",
      "  Adjust your expectations accordingly!\n",
      "  Gradient evaluation took 5e-06 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 0.05 seconds.\n",
      "  Adjust your expectations accordingly!\n",
      "  Gradient evaluation took 5e-06 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 0.05 seconds.\n",
      "  Adjust your expectations accordingly!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Building: found in cache, done.Messages from stanc:\n",
      "Warning: The parameter g has 3 priors.\n",
      "Warning: The parameter d_old has 2 priors.\n",
      "Warning: The parameter d_new has 2 priors.\n",
      "Sampling:   0%\n",
      "Sampling:  25% (3000/12000)\n",
      "Sampling:  50% (6000/12000)\n",
      "Sampling:  75% (9000/12000)\n",
      "Sampling: 100% (12000/12000)\n",
      "Sampling: 100% (12000/12000), done.\n",
      "Messages received during sampling:\n",
      "  Gradient evaluation took 9e-06 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 0.09 seconds.\n",
      "  Adjust your expectations accordingly!\n",
      "  Gradient evaluation took 9.1e-05 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 0.91 seconds.\n",
      "  Adjust your expectations accordingly!\n",
      "  Gradient evaluation took 8e-06 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 0.08 seconds.\n",
      "  Adjust your expectations accordingly!\n",
      "  Gradient evaluation took 1.2e-05 seconds\n",
      "  1000 transitions using 10 leapfrog steps per transition would take 0.12 seconds.\n",
      "  Adjust your expectations accordingly!\n"
     ]
    }
   ],
   "source": [
    "\n",
    "stan_data = {\n",
    "    'n_old': data['n_old'],\n",
    "    'n_new': data['n_new'],\n",
    "    'hits': data['hits'],\n",
    "    'misses': data['misses'],\n",
    "    'false_alarms': data['false_alarms'],\n",
    "    'correct_rejections': data['correct_rejections']\n",
    "}\n",
    "\n",
    "# Build and sample from the 1HT model\n",
    "model_1ht = stan.build(model_1ht_code, data=stan_data)\n",
    "fit_1ht = model_1ht.sample(num_chains=4, num_samples=2000, num_warmup=1000)\n",
    "\n",
    "# Build and sample from the 2HT model\n",
    "model_2ht = stan.build(model_2ht_code, data=stan_data)\n",
    "fit_2ht = model_2ht.sample(num_chains=4, num_samples=2000, num_warmup=1000)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 552608 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 562492 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 14 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "az.plot_trace(fit_1ht, figsize=(12, 8))\n",
    "plt.suptitle(\"1HT Model Trace Plots\", fontsize=16)\n",
    "plt.tight_layout()\n",
    "\n",
    "az.plot_trace(fit_2ht, figsize=(12, 8))\n",
    "plt.suptitle(\"2HT Model Trace Plots\", fontsize=16)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>mean</th>\n",
       "      <th>sd</th>\n",
       "      <th>hdi_3%</th>\n",
       "      <th>hdi_97%</th>\n",
       "      <th>mcse_mean</th>\n",
       "      <th>mcse_sd</th>\n",
       "      <th>ess_bulk</th>\n",
       "      <th>ess_tail</th>\n",
       "      <th>r_hat</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>d_old</th>\n",
       "      <td>0.621</td>\n",
       "      <td>0.230</td>\n",
       "      <td>0.172</td>\n",
       "      <td>0.964</td>\n",
       "      <td>0.004</td>\n",
       "      <td>0.003</td>\n",
       "      <td>3229.0</td>\n",
       "      <td>3084.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>d_new</th>\n",
       "      <td>0.604</td>\n",
       "      <td>0.236</td>\n",
       "      <td>0.136</td>\n",
       "      <td>0.953</td>\n",
       "      <td>0.004</td>\n",
       "      <td>0.003</td>\n",
       "      <td>3209.0</td>\n",
       "      <td>3121.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>g</th>\n",
       "      <td>0.491</td>\n",
       "      <td>0.238</td>\n",
       "      <td>0.095</td>\n",
       "      <td>0.886</td>\n",
       "      <td>0.005</td>\n",
       "      <td>0.002</td>\n",
       "      <td>2714.0</td>\n",
       "      <td>4377.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>pred_hits</th>\n",
       "      <td>12.645</td>\n",
       "      <td>1.832</td>\n",
       "      <td>9.000</td>\n",
       "      <td>15.000</td>\n",
       "      <td>0.021</td>\n",
       "      <td>0.017</td>\n",
       "      <td>7868.0</td>\n",
       "      <td>7605.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>pred_misses</th>\n",
       "      <td>2.355</td>\n",
       "      <td>1.832</td>\n",
       "      <td>0.000</td>\n",
       "      <td>6.000</td>\n",
       "      <td>0.021</td>\n",
       "      <td>0.017</td>\n",
       "      <td>7868.0</td>\n",
       "      <td>7262.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>pred_false_alarms</th>\n",
       "      <td>2.369</td>\n",
       "      <td>1.850</td>\n",
       "      <td>0.000</td>\n",
       "      <td>6.000</td>\n",
       "      <td>0.021</td>\n",
       "      <td>0.019</td>\n",
       "      <td>7477.0</td>\n",
       "      <td>6843.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>pred_correct_rejections</th>\n",
       "      <td>12.631</td>\n",
       "      <td>1.850</td>\n",
       "      <td>9.000</td>\n",
       "      <td>15.000</td>\n",
       "      <td>0.021</td>\n",
       "      <td>0.019</td>\n",
       "      <td>7477.0</td>\n",
       "      <td>6943.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                           mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  \\\n",
       "d_old                     0.621  0.230   0.172    0.964      0.004    0.003   \n",
       "d_new                     0.604  0.236   0.136    0.953      0.004    0.003   \n",
       "g                         0.491  0.238   0.095    0.886      0.005    0.002   \n",
       "pred_hits                12.645  1.832   9.000   15.000      0.021    0.017   \n",
       "pred_misses               2.355  1.832   0.000    6.000      0.021    0.017   \n",
       "pred_false_alarms         2.369  1.850   0.000    6.000      0.021    0.019   \n",
       "pred_correct_rejections  12.631  1.850   9.000   15.000      0.021    0.019   \n",
       "\n",
       "                         ess_bulk  ess_tail  r_hat  \n",
       "d_old                      3229.0    3084.0    1.0  \n",
       "d_new                      3209.0    3121.0    1.0  \n",
       "g                          2714.0    4377.0    1.0  \n",
       "pred_hits                  7868.0    7605.0    1.0  \n",
       "pred_misses                7868.0    7262.0    1.0  \n",
       "pred_false_alarms          7477.0    6843.0    1.0  \n",
       "pred_correct_rejections    7477.0    6943.0    1.0  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "az.summary(fit_1ht)\n",
    "az.summary(fit_2ht)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
