SAT_help_functions_old.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "356f5305",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0d84ed28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyddm.plot\n",
    "from pyddm import Model, set_N_cpus\n",
    "from pyddm.models import DriftConstant, DriftLinear, NoiseConstant,Bound,Drift, BoundConstant, OverlayNonDecision, ICPointSourceCenter,BoundCollapsingLinear\n",
    "from pyddm.functions import fit_adjust_model, display_model\n",
    "\n",
    "from pyddm import Fittable, Fitted, Sample\n",
    "from pyddm.models import LossRobustBIC, LossBIC, LossLikelihood, LossRobustLikelihood,LossSquaredError\n",
    "from pyddm.functions import fit_adjust_model\n",
    "\n",
    "import copy\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import seaborn as sns\n",
    "import scipy\n",
    "from statannot import add_stat_annotation\n",
    "import glob\n",
    "\n",
    "from sbi import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dbadcbd1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e7ede39",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_df_DDM2v2a2c2ter(df):\n",
    "    SAT = [copy.copy(e) for _ in range(len(df)) for e in [\"Speed\"]] + [copy.copy(e) for _ in range(len(df)) for e in [\"Accuracy\"]]\n",
    "    df_speed = df.loc[:, [\"Drift Rate Speed\", \"Boundary Separation Speed\", \"Non-decision Time Speed\", \"Collapse Speed\"]]\n",
    "    df_speed = df_speed.rename(columns={\"Drift Rate Speed\":\"Drift Rate\",\n",
    "                                        \"Boundary Separation Speed\": \"Boundary Separation\",\n",
    "                                        \"Non-decision Time Speed\": \"Non-decision Time\",\n",
    "                                        \"Collapse Speed\":\"Collapse Rate\"})\n",
    "\n",
    "    df_accuracy = df.loc[:, [\"Drift Rate Accuracy\", \"Boundary Separation Accuracy\", \"Non-decision Time Accuracy\", \"Collapse Accuracy\"]]\n",
    "    df_accuracy = df_accuracy.rename(columns={\"Drift Rate Accuracy\": \"Drift Rate\", \n",
    "                                              \"Boundary Separation Accuracy\": \"Boundary Separation\",\n",
    "                                              \"Non-decision Time Accuracy\": \"Non-decision Time\",\n",
    "                                              \"Collapse Accuracy\":\"Collapse Rate\"})\n",
    "\n",
    "    df_concat = pd.concat((df_speed, df_accuracy), ignore_index=True)\n",
    "    df_concat[\"SAT\"] = SAT\n",
    "\n",
    "    df_melted = pd.melt(df_concat, id_vars = [\"SAT\"], \n",
    "                            value_vars=['Drift Rate', 'Boundary Separation','Non-decision Time', 'Collapse Rate'],\n",
    "                            var_name = \"Parameter\", value_name = \"Value\")\n",
    "    return df_melted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0e1642f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DDM_result_sub_2v2c2ter(samp, v = [0,6], a = [0.5,8], c = [0,4], ter=[0.1,1],\n",
    "                                dx=.01, dt=.01, T_dur=5):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftConstantSpeedAcc(Vspeed=Fittable(minval=v[0], maxval=v[1]),\n",
    "                                                  Vacc=Fittable(minval=v[0], maxval=v[1])), \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundCollapsingLinearSAT(B      =   Fittable(minval=a[0], maxval=a[1]),\n",
    "                                                     cSpeed =   Fittable(minval=c[0], maxval=c[1]),\n",
    "                                                     cAcc   =   Fittable(minval=c[0], maxval=c[1])),\n",
    "                      \n",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "\n",
    "                      ),\n",
    "                      dx=dx, dt=dt, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit, fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    #print(model_fit.get_model_parameters())\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,7),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    #df['Participant']=i+1\n",
    "    #df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'drift': 'Drift Rate Speed',\n",
    "                            'drift': 'Drift Rate Accuracy',\n",
    "                            'B': 'Boundary Separation',\n",
    "                            'cSpeed': 'Collapse Speed',\n",
    "                            'cAcc':   'Collapse Accuracy',\n",
    "                            'tSpeed': 'Non-decision Time Speed',\n",
    "                            'tAcc': 'Non-decision Time Accuracy'})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28c64758",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DDM_result_sub_2v2a2c2ter(samp, v = [0,6], a = [0.5,8], c = [0,4], ter=[0.1,1],\n",
    "                                  dx=.01, dt=.01, T_dur=5):\n",
    "\n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftConstantSpeedAcc(Vspeed = Fittable(minval=v[0], maxval=v[1]),\n",
    "                                                  Vacc = Fittable(minval=v[0], maxval=v[1])), \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundCollapsingLinearSAT(BSpeed=  Fittable(minval=a[0], maxval=a[1]),\n",
    "                                                     BAcc =   Fittable(minval=a[0], maxval=a[1]),\n",
    "                                                     cSpeed = Fittable(minval=c[0], maxval=c[1]),\n",
    "                                                     cAcc =   Fittable(minval=c[0], maxval=c[1])),\n",
    "                      \n",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "\n",
    "                      ),\n",
    "                      dx=dx, dt=dt, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    #print(model_fit.get_model_parameters())\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,8),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    #df['Participant']=i+1\n",
    "    #df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'Vspeed': 'Drift Rate Speed',\n",
    "                            'Vacc': 'Drift Rate Accuracy',\n",
    "                            'BSpeed': 'Boundary Separation Speed',\n",
    "                            'BAcc':   'Boundary Separation Accuracy',\n",
    "                            'cSpeed': 'Collapse Speed',\n",
    "                            'cAcc':   'Collapse Accuracy',\n",
    "                            'tSpeed': 'Non-decision Time Speed',\n",
    "                            'tAcc': 'Non-decision Time Accuracy'})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48a0f8f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DDM_result_sub_2a2ter(samp, v = [0,6], a = [0.5,8], ter=[0.1,1], T_dur=3):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftConstant(drift=Fittable(minval=v[0], maxval=v[1])), \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundSpeedAcc(\n",
    "                          Bspeed=Fittable(minval=a[0], maxval=a[1]),\n",
    "                          Bacc=Fittable(minval=a[0], maxval=a[1])),\n",
    "                      \n",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "\n",
    "                      ),\n",
    "                      dx=.001, dt=.001, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    #print(model_fit.get_model_parameters())\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,5),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    #df['Participant']=i+1\n",
    "    #df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'drift': 'Drift Rate',\n",
    "                            'Bspeed': 'Boundary Separation Speed',\n",
    "                            'Bacc': 'Boundary Separation Accuracy',\n",
    "                            'tSpeed': 'Non-decision Time Speed',\n",
    "                            'tAcc': 'Non-decision Time Accuracy'})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b5aef99",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DDM_result_sub_2v2a2ter(samp, v = [0,6], a = [0.5,8], ter=[0.1,1], T_dur=5):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftConstantSpeedAcc(Vspeed = Fittable(minval=v[0], maxval=v[1]),\n",
    "                                                  Vacc = Fittable(minval=v[0], maxval=v[1])), \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundSpeedAcc(\n",
    "                          Bspeed=Fittable(minval=a[0], maxval=a[1]),\n",
    "                          Bacc=Fittable(minval=a[0], maxval=a[1])),\n",
    "                      \n",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "\n",
    "                      ),\n",
    "                      dx=.01, dt=.01, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    #print(model_fit.get_model_parameters())\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,6),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    #df['Participant']=i+1\n",
    "    #df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'Vspeed': 'Drift Rate Speed',\n",
    "                            'Vacc': 'Drift Rate Accuracy',\n",
    "                            'Bspeed': 'Boundary Separation Speed',\n",
    "                            'Bacc': 'Boundary Separation Accuracy',\n",
    "                            'tSpeed': 'Non-decision Time Speed',\n",
    "                            'tAcc': 'Non-decision Time Accuracy'})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebc235f8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faa80993",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e30046c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2a2ter(samp, v = [0,6], a = [0.5,8], k = [-10,10], ter=[0.1,1], T_dur=3):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftLinear(drift=Fittable(minval=v[0], maxval=v[1]),t = 0,\n",
    "                                        x=Fittable(minval=k[0], maxval=k[1])),                       \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundSpeedAcc(\n",
    "                          Bspeed=Fittable(minval=a[0], maxval=a[1]),\n",
    "                          Bacc=Fittable(minval=a[0], maxval=a[1])),\n",
    "                      \n",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "\n",
    "                      ),\n",
    "                      dx=.001, dt=.001, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    #print(model_fit.get_model_parameters())\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,6),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    #df['Participant']=i+1\n",
    "    #df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'drift': 'Drift Rate',\n",
    "                            'Bspeed': 'Boundary Separation Speed',\n",
    "                            'Bacc': 'Boundary Separation Accuracy',\n",
    "                            'K': \"Self-excitation\",\n",
    "                            'tSpeed': 'Non-decision Time Speed',\n",
    "                            'tAcc': 'Non-decision Time Accuracy'})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3ec99d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2k2ter(samp, v = [0,6], a = [0.5,8], k = [-10,10], ter=[0.1,1], T_dur=3):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftExcitationSpeedAcc(drift=Fittable(minval=v[0], maxval=v[1]),\n",
    "                                                    Kspeed=Fittable(minval=k[0], maxval=k[1]),\n",
    "                                                    Kacc=Fittable(minval=k[0], maxval=k[1])),                              \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundConstant(B=Fittable(minval=a[0], maxval=a[1])),\n",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                      ),\n",
    "                      dx=.001, dt=.001, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    #print(model_fit.get_model_parameters())\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,6),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    #df['Participant']=i+1\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'drift': 'Drift Rate',\n",
    "                            'B': 'Boundary Separation',\n",
    "                            'tSpeed': 'Non-decision Time Speed',\n",
    "                            'tAcc': 'Non-decision Time Accuracy',\n",
    "                            'Kspeed': \"Self-excitation Speed\",\n",
    "                            'Kacc': \"Self-excitation Accuracy\",\n",
    "                           })\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e43ea45",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56536892",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2v2a2ter(samp, v = [0,6], a = [0.5,8], k = [-10,10], ter=[0.1,1], T_dur=3):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftConstantSpeedAcc(Vspeed = Fittable(minval=v[0], maxval=v[1]),\n",
    "                                                  Vacc = Fittable(minval=v[0], maxval=v[1]),\n",
    "                                                  K = Fittable(minval=k[0], maxval=k[1])), \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundSpeedAcc(\n",
    "                          Bspeed=Fittable(minval=a[0], maxval=a[1]),\n",
    "                          Bacc=Fittable(minval=a[0], maxval=a[1])),\n",
    "                      \n",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "\n",
    "                      ),\n",
    "                      dx=.001, dt=.001, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    #print(model_fit.get_model_parameters())\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,7),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    #df['Participant']=i+1\n",
    "    #df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'Vspeed': 'Drift Rate Speed',\n",
    "                            'Vacc': 'Drift Rate Accuracy',\n",
    "                            'Bspeed': 'Boundary Separation Speed',\n",
    "                            'Bacc': 'Boundary Separation Accuracy',\n",
    "                            'K': \"Self-excitation\",\n",
    "                            'tSpeed': 'Non-decision Time Speed',\n",
    "                            'tAcc': 'Non-decision Time Accuracy'})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59dcf61f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2v2k2ter(samp, v = [0,6], a = [0.5,8], k = [-10,10], ter=[0.1,1], \n",
    "                                dx=.01, dt=.01, T_dur=5):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftBothSpeedAcc(Vspeed = Fittable(minval=v[0], maxval=v[1]),\n",
    "                                              Vacc = Fittable(minval=v[0], maxval=v[1]),\n",
    "                                              Kspeed=Fittable(minval=k[0], maxval=k[1]),\n",
    "                                              Kacc=Fittable(minval=k[0], maxval=k[1])), \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundConstant(B=Fittable(minval=a[0], maxval=a[1])),                      \n",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "\n",
    "                      ),\n",
    "                      dx=.01, dt=.01, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    #print(model_fit.get_model_parameters())\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,7),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    #df['Participant']=i+1\n",
    "    #df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'Vspeed': 'Drift Rate Speed',\n",
    "                            'Vacc': 'Drift Rate Accuracy',\n",
    "                            'B': 'Boundary Separation',\n",
    "                            'Kspeed': \"Self-excitation Speed\",\n",
    "                            'Kacc': \"Self-excitation Accuracy\",\n",
    "                            'tSpeed': 'Non-decision Time Speed',\n",
    "                            'tAcc': 'Non-decision Time Accuracy'})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eff1306f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_full(samp, v = [0,6], a = [0.5,8], k = [-10,10], ter=[0.1,1]\n",
    "                            ,dx=.01, dt=.01, T_dur=5):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftBothSpeedAcc(Vspeed = Fittable(minval=v[0], maxval=v[1]),\n",
    "                                              Vacc = Fittable(minval=v[0], maxval=v[1]),\n",
    "                                              Kspeed=Fittable(minval=k[0], maxval=k[1]),\n",
    "                                              Kacc=Fittable(minval=k[0], maxval=k[1])), \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundSpeedAcc(\n",
    "                          Bspeed=Fittable(minval=a[0], maxval=a[1]),\n",
    "                          Bacc=Fittable(minval=a[0], maxval=a[1])),\n",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "\n",
    "                      ),\n",
    "                      dx=dx, dt=dt, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    #print(model_fit.get_model_parameters())\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,8),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    #df['Participant']=i+1\n",
    "    #df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'Vspeed': 'Drift Rate Speed',\n",
    "                            'Vacc': 'Drift Rate Accuracy',\n",
    "                            'Bspeed': 'Boundary Separation Speed',\n",
    "                            'Bacc': 'Boundary Separation Accuracy',\n",
    "                            'Kspeed': \"Self-excitation Speed\",\n",
    "                            'Kacc': \"Self-excitation Accuracy\",\n",
    "                            'tSpeed': 'Non-decision Time Speed',\n",
    "                            'tAcc': 'Non-decision Time Accuracy'})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30d6fac3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DDM_result_sub(samp, SAT, v = [0,6], a = [0.5,8], ter=[0.1,1], T_dur=3):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                              drift=DriftConstant(drift=Fittable(minval=v[0], maxval=v[1])),\n",
    "                              noise=NoiseConstant(noise=1),\n",
    "                              bound=BoundConstant(B=Fittable(minval=a[0], maxval=a[1])),\n",
    "                              overlay=OverlayNonDecision(nondectime=Fittable(minval=ter[0], maxval=ter[1])),\n",
    "                              dx=.001, dt=.001, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,3),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    df['Participant']=i+1\n",
    "    df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'drift': 'Drift Rate',\n",
    "                            'B': 'Boundary Separation',\n",
    "                            'nondectime': \"Non-decision Time\"})\n",
    "\n",
    "    \n",
    "    df_all = pd.DataFrame(columns = ['Participant', 'drift', 'B', 'nondectime','SAT','BIC'])\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5e5d626",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2a(samp, v = [0,6], a = [0.5,8], k = [-10,10], ter=[0.1,1], T_dur=3):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftLinear(drift=Fittable(minval=v[0], maxval=v[1]),\n",
    "                                        t = 0, x=Fittable(minval=k[0], maxval=k[1])),                                noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundSpeedAcc(\n",
    "                          Bspeed=Fittable(minval=a[0], maxval=a[1]),\n",
    "                          Bacc=Fittable(minval=a[0], maxval=a[1])),\n",
    "                      \n",
    "                      overlay=OverlayNonDecision(nondectime=Fittable(minval=ter[0], maxval=ter[1])),\n",
    "                      dx=.001, dt=.001, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    #print(model_fit.get_model_parameters())\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,5),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    df['Participant']=i+1\n",
    "    #df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'drift': 'Drift Rate',\n",
    "                            'Bspeed': 'Boundary Separation Speed',\n",
    "                            'Bacc': 'Boundary Separation Accuracy',\n",
    "\n",
    "                            'nondectime': \"Non-decision Time\",\n",
    "                            'x': \"Self-excitation\"})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59b830c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2k(samp, v = [0,6], a = [0.5,8], k = [-10,10], ter=[0.1,1], T_dur=3):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftSpeedAcc(drift = Fittable(minval=v[0], maxval=v[1]),\n",
    "                                          Kspeed=Fittable(minval=k[0], maxval=k[1]),\n",
    "                                          Kacc=Fittable(minval=k[0], maxval=k[1])),                      \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundConstant(B=Fittable(minval=a[0], maxval=a[1])),\n",
    "                      overlay=OverlayNonDecision(nondectime=Fittable(minval=ter[0], maxval=ter[1])),\n",
    "                      dx=.001, dt=.001, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,5),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    df['Participant']=i+1\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'drift': 'Drift Rate',\n",
    "                            'B': 'Boundary Separation',\n",
    "                            'nondectime': \"Non-decision Time\",\n",
    "                            'Kspeed': 'Self-excitation Speed',\n",
    "                            'Kacc': 'Self-excitation Accuracy'\n",
    "                           })\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f025112e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub(samp, SAT, v = [0,6], a = [0.5,8], k = [-10,10], ter=[0.1,1], T_dur=3):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                              drift=DriftLinear(drift=Fittable(minval=v[0], maxval=v[1]),\n",
    "                                                t = 0, x=Fittable(minval=k[0], maxval=k[1])),                                noise=NoiseConstant(noise=1),\n",
    "                              bound=BoundConstant(B=Fittable(minval=a[0], maxval=a[1])),\n",
    "                              overlay=OverlayNonDecision(nondectime=Fittable(minval=ter[0], maxval=ter[1])),\n",
    "                              dx=.001, dt=.001, T_dur=T_dur)\n",
    "    #set_N_cpus(4)\n",
    "    fit_adjust_model(samp, model_fit,\n",
    "                     fitting_method=\"differential_evolution\",\n",
    "                     lossfunction=LossRobustBIC, verbose=False)\n",
    "    df = pd.DataFrame(data = np.array(model_fit.get_model_parameters()).reshape(1,4),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    df['Participant']=i+1\n",
    "    df['SAT'] = SAT\n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'drift': 'Drift Rate',\n",
    "                            'B': 'Boundary Separation',\n",
    "                            'nondectime': \"Non-decision Time\",\n",
    "                            'x': \"Self-excitation\"})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "239e9fe7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DDM_OUM_result(df_Dutilh2019, v, a, ter, k, T_dur=3):\n",
    "    df_all = pd.DataFrame(columns = ['Participant', 'Drift Rate','Boundary Separation','Non-decision Time','Effective Inhibition','SAT'],)\n",
    "\n",
    "    for i in range(len(df_Dutilh2019.index)):\n",
    "\n",
    "        print(i)\n",
    "        model = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftConstant(drift=df_Dutilh2019[\"Drift Rate\"][i]),   \n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundConstant(B = df_Dutilh2019[\"Boundary Separation\"][i]),\n",
    "                      overlay=OverlayNonDecision(nondectime = df_Dutilh2019[\"Non-decision Time\"][i]),\n",
    "                      dx=.001, dt=.001, T_dur=3)\n",
    "\n",
    "        sol = model.solve()\n",
    "        samp = sol.resample(1000)\n",
    "\n",
    "\n",
    "        model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftLinear(drift=Fittable(minval=v[0], maxval=v[1]),\n",
    "                                        t = 0,\n",
    "                                        x=Fittable(minval=k[0], maxval=k[1])),\n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundConstant(B=Fittable(minval=a[0], maxval=a[1])),\n",
    "                      overlay=OverlayNonDecision(nondectime=Fittable(minval=ter[0], maxval=ter[1])),\n",
    "                      dx=.001, dt=.001, T_dur=T_dur)\n",
    "\n",
    "        fit_adjust_model(samp, model_fit,\n",
    "                         fitting_method=\"differential_evolution\",\n",
    "                         lossfunction=LossRobustBIC, verbose=False)\n",
    "\n",
    "        #df_sub = pd.DataFrame(columns = ['Frequency', 'SAT', 'Participant', 'v','a','ter','k'],)\n",
    "\n",
    "        df_sub = pd.DataFrame({\n",
    "                               'Participant': df_Dutilh2019[\"Participant\"][i],\n",
    "                               'Drift Rate': np.array(model_fit.parameters()['drift']['drift']), \n",
    "                               'Boundary Separation': np.array(model_fit.parameters()['bound']['B']),\n",
    "                               'Non-decision Time': np.array(model_fit.parameters()['overlay']['nondectime']),\n",
    "                               'Effective Inhibition': np.array(model_fit.parameters()['drift']['x']),\n",
    "                               'SAT':df_Dutilh2019[\"SAT\"][i]\n",
    "                              }, index=[i])\n",
    "\n",
    "        df_all = df_all.append(df_sub)\n",
    "\n",
    "        print(df_sub)\n",
    "    return(df_all)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b84f4f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sim_fit_plot_all(par, par_fitted, par_labels, par_fitted_labels, s = 4, figsize=(10,10)):\n",
    "    \n",
    "    param_num = par.shape[1]\n",
    "    param_fitted_num = par_fitted.shape[1]\n",
    "    \n",
    "    fig, axes = plt.subplots(param_fitted_num, param_num, figsize=figsize)\n",
    "    \n",
    "    for i in range(param_num):\n",
    "        for j in range(param_fitted_num):\n",
    "            axes[j,i].scatter(par[:,i], par_fitted[:,j], s = s, facecolors='none', edgecolors='k')\n",
    "\n",
    "    for m in range(param_fitted_num):\n",
    "        axes[m,0].set_ylabel(par_fitted_labels[m],fontsize=18)#, rotation=0)\n",
    "    \n",
    "    for n in range(param_num):\n",
    "        #axes[0,n].set_title(par_labels[n])\n",
    "        axes[(param_fitted_num-1),n].set_xlabel(par_labels[n], fontsize=18)\n",
    "    fig.align_ylabels(axes[:, :])\n",
    "    fig.align_xlabels(axes[:, :])\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4541ded1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def standardize_result1(dataset, model):\n",
    "\n",
    "    df100_DDM = pd.read_csv(model + \"result\" + dataset + \".csv\")\n",
    "\n",
    "    ## calculate mean and std of each param for one dataset\n",
    "    df100_DDM[\"v_mean\"] = [np.mean(df100_DDM[\"Drift Rate\"])]*len(df100_DDM.index)\n",
    "    df100_DDM[\"v_sd\"] = [np.std(df100_DDM[\"Drift Rate\"])]*len(df100_DDM.index)\n",
    "\n",
    "    df100_DDM[\"a_mean\"] = [np.mean(df100_DDM[\"Boundary Separation\"])]*len(df100_DDM.index)\n",
    "    df100_DDM[\"a_sd\"] = [np.std(df100_DDM[\"Boundary Separation\"])]*len(df100_DDM.index)\n",
    "\n",
    "    df100_DDM[\"ter_mean\"] = [np.mean(df100_DDM[\"Non-decision Time\"])]*len(df100_DDM.index)\n",
    "    df100_DDM[\"ter_sd\"] = [np.std(df100_DDM[\"Non-decision Time\"])]*len(df100_DDM.index)\n",
    "\n",
    "    # standardize\n",
    "    df100_DDM[\"Drift Rate\"] = (df100_DDM[\"Drift Rate\"] - df100_DDM[\"v_mean\"])/df100_DDM[\"v_sd\"]\n",
    "    df100_DDM[\"Boundary Separation\"] = (df100_DDM[\"Boundary Separation\"] - df100_DDM[\"a_mean\"])/df100_DDM[\"a_sd\"]\n",
    "    df100_DDM[\"Non-decision Time\"] = (df100_DDM[\"Non-decision Time\"] - df100_DDM[\"ter_mean\"])/df100_DDM[\"ter_sd\"]\n",
    "    \n",
    "    if \"Effective Inhibition\" in list(df100_DDM.columns):\n",
    "        df100_DDM[\"k_mean\"] = [np.mean(df100_DDM[\"Effective Inhibition\"])]*len(df100_DDM.index)\n",
    "        df100_DDM[\"k_sd\"] = [np.std(df100_DDM[\"Effective Inhibition\"])]*len(df100_DDM.index)\n",
    "        df100_DDM[\"Effective Inhibition\"] = (df100_DDM[\"Effective Inhibition\"] - df100_DDM[\"k_mean\"])/df100_DDM[\"k_sd\"]\n",
    "\n",
    "    \n",
    "    return(df100_DDM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b35f6ca",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfd3ea08",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bca478fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_cohen_d(df_dif, model):\n",
    "\n",
    "    dataset_model = np.array(df_dif[['Dataset','Model']].loc[0])\n",
    "    mean = np.array(df_dif.groupby(['Dataset','Model'], as_index=False).agg(\"mean\"))[0,2:]\n",
    "    std = np.array(df_dif.groupby(['Dataset','Model'], as_index=False).agg(\"std\"))[0,2:]\n",
    "    d = mean/std\n",
    "    arr = np.concatenate((np.expand_dims(dataset_model, axis = 1),np.expand_dims(d, axis = 1)))\n",
    "    \n",
    "    if model == \"DDM\":\n",
    "        df = pd.DataFrame(arr.T, columns = [\"Dataset\", \"Model\", 'Drift Rate', 'Boundary Separation','Non-decision Time'])\n",
    "        df_melted = pd.melt(df, id_vars = [\"Dataset\", \"Model\"],\n",
    "                    value_vars=['Drift Rate', 'Boundary Separation','Non-decision Time'],\n",
    "                    var_name = \"Parameter\", value_name = \"Cohen's d\")\n",
    "    if model == \"OUM\" or model == \"DDMOUM\":\n",
    "        df = pd.DataFrame(arr.T, columns = [\"Dataset\", \"Model\", 'Drift Rate', 'Boundary Separation','Non-decision Time',\"Effective Inhibition\"])\n",
    "        df_melted = pd.melt(df, id_vars = [\"Dataset\", \"Model\"],\n",
    "                    value_vars=['Drift Rate', 'Boundary Separation','Non-decision Time', \"Effective Inhibition\"],\n",
    "                    var_name = \"Parameter\", value_name = \"Cohen's d\")\n",
    "    return(df_melted)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37fa7c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def standardize_result(df100_DDM, model):\n",
    "\n",
    "    ## calculate mean and std of each param for one dataset\n",
    "    df100_DDM[\"v_mean\"] = [np.mean(df100_DDM[\"Drift Rate\"])]*len(df100_DDM.index)\n",
    "    df100_DDM[\"v_sd\"] = [np.std(df100_DDM[\"Drift Rate\"])]*len(df100_DDM.index)\n",
    "\n",
    "    df100_DDM[\"a_mean\"] = [np.mean(df100_DDM[\"Boundary Separation\"])]*len(df100_DDM.index)\n",
    "    df100_DDM[\"a_sd\"] = [np.std(df100_DDM[\"Boundary Separation\"])]*len(df100_DDM.index)\n",
    "\n",
    "    df100_DDM[\"ter_mean\"] = [np.mean(df100_DDM[\"Non-decision Time\"])]*len(df100_DDM.index)\n",
    "    df100_DDM[\"ter_sd\"] = [np.std(df100_DDM[\"Non-decision Time\"])]*len(df100_DDM.index)\n",
    "\n",
    "    # standardize\n",
    "    df100_DDM[\"Drift Rate\"] = (df100_DDM[\"Drift Rate\"] - df100_DDM[\"v_mean\"])/df100_DDM[\"v_sd\"]\n",
    "    df100_DDM[\"Boundary Separation\"] = (df100_DDM[\"Boundary Separation\"] - df100_DDM[\"a_mean\"])/df100_DDM[\"a_sd\"]\n",
    "    df100_DDM[\"Non-decision Time\"] = (df100_DDM[\"Non-decision Time\"] - df100_DDM[\"ter_mean\"])/df100_DDM[\"ter_sd\"]\n",
    "    \n",
    "    if \"Effective Inhibition\" in list(df100_DDM.columns):\n",
    "        df100_DDM[\"k_mean\"] = [np.mean(df100_DDM[\"Effective Inhibition\"])]*len(df100_DDM.index)\n",
    "        df100_DDM[\"k_sd\"] = [np.std(df100_DDM[\"Effective Inhibition\"])]*len(df100_DDM.index)\n",
    "        df100_DDM[\"Effective Inhibition\"] = (df100_DDM[\"Effective Inhibition\"] - df100_DDM[\"k_mean\"])/df100_DDM[\"k_sd\"]\n",
    "    \n",
    "    if model == \"DDM\":\n",
    "        df_melted = pd.melt(df100_DDM, id_vars = [\"Dataset\", \"Model\"],\n",
    "                            value_vars=['Drift Rate', 'Boundary Separation','Non-decision Time'],\n",
    "                               var_name = \"Parameter\", value_name = \"Difference\")\n",
    "        #print(1)\n",
    "    if model == \"OUM\" or model == \"DDMOUM\":\n",
    "        df_melted = pd.melt(df100_DDM, id_vars = [\"Dataset\", \"Model\"],\n",
    "                            value_vars=['Drift Rate', 'Boundary Separation','Non-decision Time', \"Effective Inhibition\"],\n",
    "                               var_name = \"Parameter\", value_name = \"Difference\")\n",
    "        #print(2)\n",
    "    \n",
    "    return(df_melted)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ff399af9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_difference(dataset, df_SAT, model):\n",
    "    if model == \"DDM\":\n",
    "        S = np.array(df_SAT.loc[(df_SAT['SAT'] == \"Speed\")][[\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\"]])\n",
    "        A = np.array(df_SAT.loc[(df_SAT['SAT'] == \"Accuracy\")][[\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\"]])\n",
    "        dif = S-A\n",
    "        df_dif = pd.DataFrame(dif,columns = [\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\"])\n",
    "    elif model == \"OUM\" or model == \"DDMOUM\":\n",
    "        S = np.array(df_SAT.loc[(df_SAT['SAT'] == \"Speed\")][[\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\", \"Effective Inhibition\"]])\n",
    "        A = np.array(df_SAT.loc[(df_SAT['SAT'] == \"Accuracy\")][[\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\", \"Effective Inhibition\"]])\n",
    "        dif = S-A\n",
    "        df_dif = pd.DataFrame(dif,columns = [\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\", \"Effective Inhibition\"])\n",
    "    else:\n",
    "        print(\"Error: model not implemented\")\n",
    "    df_dif[\"Dataset\"] = [dataset]*len(df_dif)\n",
    "    df_dif[\"Model\"] = [model]*len(df_dif)\n",
    "    return(df_dif)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "384fe4e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_difference_DDM1(dataset, df_SAT, model):\n",
    "    S = np.array(df_SAT.loc[(df_SAT['SAT'] == \"Speed\")][[\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\"]])\n",
    "    A = np.array(df_SAT.loc[(df_SAT['SAT'] == \"Accuracy\")][[\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\"]])\n",
    "    dif = S-A\n",
    "    df_dif = pd.DataFrame(dif,columns = [\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\"])\n",
    "\n",
    "\n",
    "    df_dif_melted = pd.melt(df_dif, value_vars=['Drift Rate', 'Boundary Separation','Non-decision Time'],\n",
    "                           var_name = \"Parameter\", value_name = \"Difference\")\n",
    "\n",
    "    df_dif_melted[\"Dataset\"] = [dataset]*len(df_dif_melted)\n",
    "    df_dif_melted[\"Model\"] = [model]*len(df_dif_melted)\n",
    "    return(df_dif_melted)\n",
    "\n",
    "def get_difference_OUM1(dataset, df_SAT, model):\n",
    "    S = np.array(df_SAT.loc[(df_SAT['SAT'] == \"Speed\")][[\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\", \"Effective Inhibition\"]])\n",
    "    A = np.array(df_SAT.loc[(df_SAT['SAT'] == \"Accuracy\")][[\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\", \"Effective Inhibition\"]])\n",
    "    dif = S-A\n",
    "    df_dif = pd.DataFrame(dif,columns = [\"Drift Rate\", \"Boundary Separation\", \"Non-decision Time\", \"Effective Inhibition\"])\n",
    "    df_dif_melted = pd.melt(df_dif, value_vars=['Drift Rate', 'Boundary Separation','Non-decision Time', \"Effective Inhibition\"],\n",
    "                           var_name = \"Parameter\", value_name = \"Difference\")\n",
    "\n",
    "    df_dif_melted[\"Dataset\"] = [dataset]*len(df_dif_melted)\n",
    "    df_dif_melted[\"Model\"] = [model]*len(df_dif_melted)\n",
    "    return(df_dif_melted)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1def695f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_OUM(df_SAT_melted, comparisons_correction=None, figsize=(8,6)):\n",
    "\n",
    "    box_pairs=[\n",
    "       ((\"Drift Rate\", \"Speed\"), (\"Drift Rate\", \"Accuracy\")), \n",
    "       ((\"Boundary Separation\", \"Speed\"), (\"Boundary Separation\", \"Accuracy\")), \n",
    "       ((\"Non-decision Time\", \"Speed\"), (\"Non-decision Time\", \"Accuracy\")),\n",
    "       ((\"Effective Inhibition\", \"Speed\"), (\"Effective Inhibition\", \"Accuracy\"))\n",
    "      ]\n",
    "\n",
    "    plt.figure(figsize=figsize)\n",
    "    ax = sns.barplot(data=df_SAT_melted, x=\"Parameter\", y=\"Value\", hue=\"SAT\")\n",
    "    add_stat_annotation(ax, data=df_SAT_melted, x=\"Parameter\", y=\"Value\", hue=\"SAT\",\n",
    "                        box_pairs = box_pairs, text_format='star', loc='inside', verbose=2,\n",
    "                        comparisons_correction=comparisons_correction, test='t-test_paired')\n",
    "\n",
    "    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "\n",
    "    plt.show()\n",
    "    #plt.savefig('example_hue_layout.png', dpi=300, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d7cb8b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_DDM(df_SAT_melted, comparisons_correction=None, figsize=(8,6)):\n",
    "\n",
    "    box_pairs=[\n",
    "       ((\"Drift Rate\", \"Speed\"), (\"Drift Rate\", \"Accuracy\")), \n",
    "       ((\"Boundary Separation\", \"Speed\"), (\"Boundary Separation\", \"Accuracy\")), \n",
    "       ((\"Non-decision Time\", \"Speed\"), (\"Non-decision Time\", \"Accuracy\"))\n",
    "      ]\n",
    "\n",
    "    plt.figure(figsize=figsize)\n",
    "    ax = sns.barplot(data=df_SAT_melted, x=\"Parameter\", y=\"Value\", hue=\"SAT\")\n",
    "    add_stat_annotation(ax, data=df_SAT_melted, x=\"Parameter\", y=\"Value\", hue=\"SAT\",\n",
    "                        box_pairs = box_pairs, text_format='star', loc='inside', verbose=2,\n",
    "                        comparisons_correction=comparisons_correction, test='t-test_paired')\n",
    "\n",
    "    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "\n",
    "    plt.show()\n",
    "    #plt.savefig('example_hue_layout.png', dpi=300, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea70ab88",
   "metadata": {},
   "outputs": [],
   "source": [
    "def DDM_OUM_QP_plot(dataset, data_dir,sample_size, \n",
    "                   correct, corrects,\n",
    "                   cond, conds, rt,\n",
    "                    ax=None):\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    \n",
    "    quantiles = [.1, .3, .5, .7, .9]\n",
    "    ## behavioral data\n",
    "    RT_cor_s, RT_cor_a, RT_err_s, RT_err_a, nums = get_RT_quantile(data_dir = data_dir,sample_size= sample_size, \n",
    "                                                                   correct = correct, corrects = corrects,\n",
    "                                                                   cond = cond, conds = conds, rt = rt)\n",
    "    \n",
    "    ax.plot([nums[0]/(nums[0]+nums[2])]*5, RT_cor_s,\"o\", color = 'black', fillstyle = 'none')\n",
    "    ax.plot([nums[1]/(nums[1]+nums[3])]*5, RT_cor_a,\"o\", color = 'black', fillstyle = 'none')\n",
    "    ax.plot([nums[2]/(nums[0]+nums[2])]*5, RT_err_s,\"o\", color = 'black', fillstyle = 'none')\n",
    "    ax.plot([nums[3]/(nums[1]+nums[3])]*5, RT_err_a,\"o\", color = 'black', fillstyle = 'none')\n",
    "    \n",
    "    ## data generated from DDM\n",
    "    params_DDM = pd.read_csv(\"DDM2v2a2terresult\"+dataset+\".csv\")#.drop(labels=3, axis=0)\n",
    "    RT_cor_s, RT_cor_a, RT_err_s, RT_err_a, nums = get_RT_quantile_DDM(params_all = params_DDM,sample_size= sample_size, \n",
    "                                                                   correct = correct, corrects = corrects,\n",
    "                                                                   cond = cond, conds = conds, rt = rt)\n",
    "    \n",
    "    ax.plot([nums[0]/(nums[0]+nums[2])]*5, RT_cor_s,\"x\", color = '#008176', fillstyle = 'none')\n",
    "    ax.plot([nums[1]/(nums[1]+nums[3])]*5, RT_cor_a,\"x\", color = '#008176', fillstyle = 'none')\n",
    "    ax.plot([nums[2]/(nums[0]+nums[2])]*5, RT_err_s,\"x\", color = '#008176', fillstyle = 'none')\n",
    "    ax.plot([nums[3]/(nums[1]+nums[3])]*5, RT_err_a,\"x\", color = '#008176', fillstyle = 'none')\n",
    "    \n",
    "    ## data generated from OUM\n",
    "    params_OUM = pd.read_csv(\"OUM2v2k2terresult\"+dataset+\".csv\")#.drop(labels=3, axis=0)\n",
    "    RT_cor_s, RT_cor_a, RT_err_s, RT_err_a, nums = get_RT_quantile_OUM2v2k2ter(params_all = params_OUM,sample_size= sample_size, \n",
    "                                                                   correct = correct, corrects = corrects,\n",
    "                                                                   cond = cond, conds = conds, rt = rt)\n",
    "    \n",
    "    ax.plot([nums[0]/(nums[0]+nums[2])]*5, RT_cor_s,\"x\", color = '#c1272d', fillstyle = 'none')\n",
    "    ax.plot([nums[1]/(nums[1]+nums[3])]*5, RT_cor_a,\"x\", color = '#c1272d', fillstyle = 'none')\n",
    "    ax.plot([nums[2]/(nums[0]+nums[2])]*5, RT_err_s,\"x\", color = '#c1272d', fillstyle = 'none')\n",
    "    ax.plot([nums[3]/(nums[1]+nums[3])]*5, RT_err_a,\"x\", color = '#c1272d', fillstyle = 'none')\n",
    "    \n",
    "    # limits\n",
    "    ax.set_xlim(0,1)\n",
    "    ax.set_ylim(0,)\n",
    "        \n",
    "    ax.set_xlabel(\"Probability of Response\")\n",
    "    ax.set_ylabel(\"RT\")\n",
    "    #ax.legend(loc = 'upper right')\n",
    "    return(ax)\n",
    "#DDM_OUM_QP_plot(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1f3254c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def DDM_OUM_09quantile_plot(dataset, data_dir,sample_size, correct, corrects,\n",
    "                            cond, conds, rt,figsize=(8, 4),legend = False):\n",
    "    #if ax is None:\n",
    "    #    ax = plt.gca()\n",
    "    \n",
    "    figure, ax = plt.subplots(1,2,figsize=figsize)\n",
    "\n",
    "    ## data generated from DDM\n",
    "    params_DDM = pd.read_csv(\"DDM2v2a2terresult\"+dataset+\".csv\")#.drop(labels=3, axis=0)\n",
    "    RT_cor_s, RT_cor_a, RT_err_s, RT_err_a, nums = get_all_RT_quantile_DDM(params_all = params_DDM,sample_size= sample_size, \n",
    "                                                                   correct = correct, corrects = corrects,\n",
    "                                                                   cond = cond, conds = conds, rt = rt)\n",
    "    quantiles_DDM = pd.DataFrame({\"cor_s\":RT_cor_s[:,4],\"cor_a\":RT_cor_a[:,4], \"err_s\": RT_err_s[:,4],'err_a': RT_err_a[:,4]})\n",
    "\n",
    "    ## data generated from OUM\n",
    "    params_OUM = pd.read_csv(\"OUM2v2k2terresult\"+dataset+\".csv\")#.drop(labels=3, axis=0)\n",
    "    RT_cor_s, RT_cor_a, RT_err_s, RT_err_a, nums = get_all_RT_quantile_OUM2v2k2ter(params_all = params_OUM,sample_size= sample_size, \n",
    "                                                                   correct = correct, corrects = corrects,\n",
    "                                                                   cond = cond, conds = conds, rt = rt)\n",
    "    quantiles_OUM = pd.DataFrame({\"cor_s\":RT_cor_s[:,4],\"cor_a\":RT_cor_a[:,4], \"err_s\": RT_err_s[:,4],'err_a': RT_err_a[:,4]})\n",
    "    \n",
    "    ran = np.arange(0,max(quantiles_DDM.max(axis=None),quantiles_OUM.max(axis=None))+0.5,0.1)\n",
    "\n",
    "\n",
    "    ax[0].scatter(quantiles_DDM[\"cor_s\"],quantiles_OUM[\"cor_s\"], c = [[1,1,1]], edgecolors = \"#1F77B4\", alpha = 0.7, label = \"Speed\")\n",
    "    ax[0].scatter(quantiles_DDM[\"cor_a\"],quantiles_OUM[\"cor_a\"], c = [[1,1,1]], edgecolors = \"#FF7F0E\", alpha = 0.7, label = \"Accuracy\")\n",
    "    ax[0].set_xlabel(\"0.9 Quantile of DDM\")\n",
    "    ax[0].set_ylabel(\"0.9 Quantile of OUM\")\n",
    "    ax[0].plot(ran,ran,c = \"black\")\n",
    "    ax[0].set_title(\"Correct Choices\")\n",
    "\n",
    "    ax[1].scatter(quantiles_DDM[\"err_s\"],quantiles_OUM[\"err_s\"], c = [[1,1,1]], edgecolors = \"#1F77B4\", alpha = 0.7, label = \"Speed\")\n",
    "    ax[1].scatter(quantiles_DDM[\"err_a\"],quantiles_OUM[\"err_a\"], c = [[1,1,1]], edgecolors = \"#FF7F0E\", alpha = 0.7, label = \"Accuracy\")\n",
    "    ax[1].set_xlabel(\"0.9 Quantile of DDM\")\n",
    "    ax[1].set_title(\"Incorrect Choices\")\n",
    "    \n",
    "    \n",
    "    ax[1].plot(ran,ran,c = \"black\")\n",
    "\n",
    "    \n",
    "    #ax[1].set_ylabel(\"0.9 Quantile of OUM\")\n",
    "    #ax[1].set_yticklabels([])\n",
    "    if legend:\n",
    "        ax[1].legend(loc='upper right', bbox_to_anchor=(1.5, 1))\n",
    "    \n",
    "    return ax\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6db78dba",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28f6e1a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_all_RT_quantile_DDM(params_all, sample_size, correct, corrects, cond, conds, rt, quantiles = [.1, .3, .5, .7, .9]):\n",
    "\n",
    "    RT_cor_ss, RT_cor_as, RT_err_ss, RT_err_as, nums = np.zeros((sample_size,5)),np.zeros((sample_size,5)),np.zeros((sample_size,5)),np.zeros((sample_size,5)),np.zeros((sample_size,4))\n",
    "    for i in range(sample_size):\n",
    "        #print(i)\n",
    "        params = np.array(params_all)[i,1:7]\n",
    "        \n",
    "        \n",
    "        RT_cor_s, RT_cor_a, RT_err_s, RT_err_a, num = get_RT_quantile_DDM_sub(params = params)\n",
    "        #print(RT_err_a)\n",
    "        RT_cor_ss[i,:] = RT_cor_s\n",
    "        RT_cor_as[i,:] = RT_cor_a\n",
    "        RT_err_ss[i,:] = RT_err_s\n",
    "        RT_err_as[i,:] = RT_err_a\n",
    "        nums[i,:] = num\n",
    "    \n",
    "    return RT_cor_ss, RT_cor_as, RT_err_ss, RT_err_as, nums\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d9de53a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_all_RT_quantile_OUM2v2k2ter(params_all, sample_size, correct, corrects, cond, conds, rt, quantiles = [.1, .3, .5, .7, .9]):\n",
    "\n",
    "    RT_cor_ss, RT_cor_as, RT_err_ss, RT_err_as, nums = np.zeros((sample_size,5)),np.zeros((sample_size,5)),np.zeros((sample_size,5)),np.zeros((sample_size,5)),np.zeros((sample_size,4))\n",
    "    for i in range(sample_size):\n",
    "\n",
    "        params = np.array(params_all)[i,1:8]\n",
    "        RT_cor_s, RT_cor_a, RT_err_s, RT_err_a, num = get_RT_quantile_OUM2v2k2ter_sub(params = params)\n",
    "        #print(RT_err_a)\n",
    "        RT_cor_ss[i,:] = RT_cor_s\n",
    "        RT_cor_as[i,:] = RT_cor_a\n",
    "        RT_err_ss[i,:] = RT_err_s\n",
    "        RT_err_as[i,:] = RT_err_a\n",
    "        nums[i,:] = num\n",
    "    \n",
    "    return RT_cor_ss, RT_cor_as, RT_err_ss, RT_err_as, nums\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a813ade0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_all_RT_quantile(data_dir, sample_size, correct, corrects, cond, conds, rt, quantiles = [.1, .3, .5, .7, .9]):\n",
    "\n",
    "    RT_cor_ss, RT_cor_as, RT_err_ss, RT_err_as, nums = np.zeros((sample_size,5)),np.zeros((sample_size,5)),np.zeros((sample_size,5)),np.zeros((sample_size,5)),np.zeros((sample_size,4))\n",
    "    for i in range(sample_size):\n",
    "        #print(i)\n",
    "        path = os.path.join(data_dir+str(i+1)+\".csv\")\n",
    "        df = pd.read_csv(path, sep = \",\")\n",
    "        #print(df)\n",
    "        df = df[df[rt]<5]\n",
    "        RT_cor_s, RT_cor_a, RT_err_s, RT_err_a, num = get_RT_quantile_sub(\n",
    "            df = df, rt = rt,\n",
    "            correct = correct, corrects = corrects,\n",
    "            cond = cond, conds = conds)\n",
    "        #print(RT_cor_s)\n",
    "        RT_cor_ss[i,:] = RT_cor_s\n",
    "        RT_cor_as[i,:] = RT_cor_a\n",
    "        RT_err_ss[i,:] = RT_err_s\n",
    "        RT_err_as[i,:] = RT_err_a\n",
    "        nums[i,:] = num\n",
    "        \n",
    "        #print(RT_err_s)\n",
    "\n",
    "    \n",
    "    return RT_cor_ss, RT_cor_as, RT_err_ss, RT_err_as, nums\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afd62794",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "db1df8d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ks_data(data_dir, sample_size, correct, corrects, cond, conds, rt):\n",
    "\n",
    "    ks_cor_s, ks_cor_a, ks_err_s, ks_err_a = [],[],[],[]\n",
    "    for i in range(sample_size):\n",
    "        #print(i)\n",
    "        path = os.path.join(data_dir+str(i+1)+\".csv\")\n",
    "        df = pd.read_csv(path, sep = \",\")\n",
    "        #print(df)\n",
    "        df = df[df[rt]<5]\n",
    "        \n",
    "        df_s = df[df[cond] == conds[0]]\n",
    "        df_cor_s = df_s[df_s[correct]==corrects[0]]\n",
    "        df_err_s = df_s[df_s[correct]==corrects[1]]\n",
    "\n",
    "        df_a = df[df[cond] == conds[1]]\n",
    "        df_cor_a = df_a[df_a[correct]==corrects[0]]\n",
    "        df_err_a = df_a[df_a[correct]==corrects[1]]\n",
    "\n",
    "        \n",
    "        ks_cor_s.append(get_ks_stat(df_cor_s, eval_model = False, rt = rt))\n",
    "        ks_err_s.append(get_ks_stat(df_err_s, eval_model = False, rt = rt))\n",
    "        ks_cor_a.append(get_ks_stat(df_cor_a, eval_model = False, rt = rt))\n",
    "        ks_err_a.append(get_ks_stat(df_err_a, eval_model = False, rt = rt))\n",
    "        \n",
    "        ks_data = pd.DataFrame({\"cor_s\":ks_cor_s,\n",
    "                                \"cor_a\":ks_cor_a,\n",
    "                                \"err_s\": ks_err_s,\n",
    "                                'err_a': ks_err_a})\n",
    "    \n",
    "    return ks_data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fe14e5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ks_stat(df, eval_model = True, rt = \"RT\"):\n",
    "    #print(len(df[\"RT\"]))\n",
    "    if eval_model:\n",
    "        if len(df[rt])<1000:\n",
    "            return 0\n",
    "    else:\n",
    "        if len(df[rt])<10:\n",
    "            return 0\n",
    "    \n",
    "    normal_samp = np.random.normal(loc=np.mean(df[rt]), scale=np.std(df[rt]), size=len(df[rt]))\n",
    "    ks = scipy.stats.kstest(df[0:1000][rt], normal_samp)[0]\n",
    "    return ks\n",
    "\n",
    "\n",
    "def get_ks_DDM(params_all, sample_size, cond, conds, T_dur=5):\n",
    "\n",
    "    ks_all = pd.DataFrame()\n",
    "\n",
    "    for i in range(sample_size):\n",
    "        params = np.array(params_all)[i,1:7]\n",
    "        \n",
    "        Vspeed,Vacc, Bspeed,Bacc,tSpeed,tAcc = params[0],params[1],params[2],params[3],params[4],params[5]\n",
    "\n",
    "        model = Model(name='Simple model (fitted)',\n",
    "                          drift=DriftConstantSpeedAcc(Vspeed=Vspeed, Vacc=Vacc), \n",
    "                          noise=NoiseConstant(noise=1),\n",
    "                          bound=BoundSpeedAcc(Bspeed=Bspeed,Bacc=Bacc),\n",
    "                          overlay=NondecisionSpeedAcc(tSpeed=tSpeed,tAcc=tAcc),\n",
    "                          dx=.001, dt=.001, T_dur=T_dur)\n",
    "        sol_s = model.solve({cond:conds[0]})\n",
    "        samp_s = sol_s.resample(100000)\n",
    "        data_s = samp_s.to_pandas_dataframe(drop_undecided = True)\n",
    "\n",
    "        sol_a = model.solve({cond:conds[1]})\n",
    "        samp_a = sol_a.resample(100000)\n",
    "        data_a = samp_a.to_pandas_dataframe(drop_undecided = True)\n",
    "\n",
    "\n",
    "        ks_cor_s = get_ks_stat(data_s[data_s['choice']==1])\n",
    "        ks_err_s = get_ks_stat(data_s[data_s['choice']==0])\n",
    "        ks_cor_a = get_ks_stat(data_a[data_a['choice']==1])\n",
    "        ks_err_a = get_ks_stat(data_a[data_a['choice']==0])\n",
    "\n",
    "        ks = pd.DataFrame({\"cor_s\":[ks_cor_s], \"err_s\": [ks_err_s], \"cor_a\":[ks_cor_a],'err_a': [ks_err_a]})\n",
    "        ks[\"Participant\"] = i\n",
    "\n",
    "        ks_all = pd.concat((ks_all,ks))\n",
    "        \n",
    "    return ks_all\n",
    "\n",
    "def get_ks_OUM(params_all, sample_size, cond, conds, T_dur=5):\n",
    "\n",
    "    ks_all = pd.DataFrame()\n",
    "\n",
    "    for i in range(sample_size):\n",
    "        params = np.array(params_all)[i,1:8]\n",
    "        \n",
    "        Vspeed,Vacc, Kspeed,Kacc,B,tSpeed,tAcc = params[0],params[1],params[2],params[3],params[4],params[5],params[6]\n",
    "\n",
    "        model = Model(name='Simple model (fitted)',\n",
    "                          drift=DriftBothSpeedAcc(Vspeed=Vspeed, Vacc=Vacc, Kspeed=Kspeed, Kacc=Kacc),                              \n",
    "                          noise=NoiseConstant(noise=1),\n",
    "                          bound=BoundConstant(B=B),\n",
    "                          overlay=NondecisionSpeedAcc(tSpeed=tSpeed,tAcc=tAcc),\n",
    "                          dx=.001, dt=.001, T_dur=T_dur)\n",
    "        sol_s = model.solve({cond:conds[0]})\n",
    "        samp_s = sol_s.resample(100000)\n",
    "        data_s = samp_s.to_pandas_dataframe(drop_undecided = True)\n",
    "\n",
    "        sol_a = model.solve({cond:conds[1]})\n",
    "        samp_a = sol_a.resample(100000)\n",
    "        data_a = samp_a.to_pandas_dataframe(drop_undecided = True)\n",
    "\n",
    "\n",
    "        ks_cor_s = get_ks_stat(data_s[data_s['choice']==1])\n",
    "        ks_err_s = get_ks_stat(data_s[data_s['choice']==0])\n",
    "        ks_cor_a = get_ks_stat(data_a[data_a['choice']==1])\n",
    "        ks_err_a = get_ks_stat(data_a[data_a['choice']==0])\n",
    "\n",
    "        ks = pd.DataFrame({\"cor_s\":[ks_cor_s], \"err_s\": [ks_err_s], \"cor_a\":[ks_cor_a],'err_a': [ks_err_a]})\n",
    "        ks[\"Participant\"] = i\n",
    "\n",
    "        ks_all = pd.concat((ks_all,ks))\n",
    "        \n",
    "    return ks_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "057fed2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ks_DDM_matched(params_all, sample_size, cond, conds, RT = \"Time\", T_dur=5):\n",
    "\n",
    "    ks_all = pd.DataFrame()\n",
    "\n",
    "    for i in range(sample_size):\n",
    "        params = np.array(params_all)[i,1:7]\n",
    "        \n",
    "        Vspeed,Vacc, Bspeed,Bacc,tSpeed,tAcc = params[0],params[1],params[2],params[3],params[4],params[5]\n",
    "\n",
    "        model = Model(name='Simple model (fitted)',\n",
    "                          drift=DriftConstantSpeedAcc(Vspeed=Vspeed, Vacc=Vacc), \n",
    "                          noise=NoiseConstant(noise=1),\n",
    "                          bound=BoundSpeedAcc(Bspeed=Bspeed,Bacc=Bacc),\n",
    "                          overlay=NondecisionSpeedAcc(tSpeed=tSpeed,tAcc=tAcc),\n",
    "                          dx=.001, dt=.001, T_dur=T_dur)\n",
    "        sol_s = model.solve({cond:conds[0]})\n",
    "        samp_s = sol_s.resample(100000)\n",
    "        data_s = samp_s.to_pandas_dataframe(drop_undecided = True)\n",
    "\n",
    "        sol_a = model.solve({cond:conds[1]})\n",
    "        samp_a = sol_a.resample(100000)\n",
    "        data_a = samp_a.to_pandas_dataframe(drop_undecided = True)\n",
    "        \n",
    "        path = os.path.join(data_dir+str(i+1)+\".csv\")\n",
    "        df = pd.read_csv(path, sep = \",\")\n",
    "        #print(df)\n",
    "        df = df[df[RT]<5]\n",
    "        \n",
    "        df_s = df[df[cond] == conds[0]]\n",
    "        n_cor_s = len(df_s[df_s[correct]==corrects[0]])\n",
    "        n_err_s = len(df_s[df_s[correct]==corrects[1]])\n",
    "\n",
    "        df_a = df[df[cond] == conds[1]]\n",
    "        n_cor_a = len(df_a[df_a[correct]==corrects[0]])\n",
    "        n_err_a = len(df_a[df_a[correct]==corrects[1]])\n",
    "\n",
    "        #print(data_s[data_s['choice']==1][0:n_cor_s])\n",
    "        ks_cor_s = get_ks_stat(data_s[data_s['choice']==1][0:n_cor_s],eval_model = False)\n",
    "        ks_err_s = get_ks_stat(data_s[data_s['choice']==0][0:n_err_s],eval_model = False)\n",
    "        ks_cor_a = get_ks_stat(data_a[data_a['choice']==1][0:n_cor_a],eval_model = False)\n",
    "        ks_err_a = get_ks_stat(data_a[data_a['choice']==0][0:n_err_a],eval_model = False)\n",
    "\n",
    "        ks = pd.DataFrame({\"cor_s\":[ks_cor_s], \"err_s\": [ks_err_s], \"cor_a\":[ks_cor_a],'err_a': [ks_err_a]})\n",
    "        ks[\"Participant\"] = i\n",
    "\n",
    "        ks_all = pd.concat((ks_all,ks))\n",
    "        \n",
    "    return ks_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61cc01a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ks_OUM_matched(params_all, sample_size, cond, conds, RT = \"Time\", T_dur=5):\n",
    "\n",
    "    ks_all = pd.DataFrame()\n",
    "\n",
    "    for i in range(sample_size):\n",
    "        params = np.array(params_all)[i,1:8]\n",
    "        \n",
    "        Vspeed,Vacc, Kspeed,Kacc,B,tSpeed,tAcc = params[0],params[1],params[2],params[3],params[4],params[5],params[6]\n",
    "\n",
    "        model = Model(name='Simple model (fitted)',\n",
    "                          drift=DriftBothSpeedAcc(Vspeed=Vspeed, Vacc=Vacc, Kspeed=Kspeed, Kacc=Kacc),                              \n",
    "                          noise=NoiseConstant(noise=1),\n",
    "                          bound=BoundConstant(B=B),\n",
    "                          overlay=NondecisionSpeedAcc(tSpeed=tSpeed,tAcc=tAcc),\n",
    "                          dx=.001, dt=.001, T_dur=T_dur)\n",
    "        sol_s = model.solve({cond:conds[0]})\n",
    "        samp_s = sol_s.resample(100000)\n",
    "        data_s = samp_s.to_pandas_dataframe(drop_undecided = True)\n",
    "\n",
    "        sol_a = model.solve({cond:conds[1]})\n",
    "        samp_a = sol_a.resample(100000)\n",
    "        data_a = samp_a.to_pandas_dataframe(drop_undecided = True)\n",
    "        \n",
    "        path = os.path.join(data_dir+str(i+1)+\".csv\")\n",
    "        df = pd.read_csv(path, sep = \",\")\n",
    "        #print(df)\n",
    "        df = df[df[RT]<5]\n",
    "        \n",
    "        df_s = df[df[cond] == conds[0]]\n",
    "        n_cor_s = len(df_s[df_s[correct]==corrects[0]])\n",
    "        n_err_s = len(df_s[df_s[correct]==corrects[1]])\n",
    "\n",
    "        df_a = df[df[cond] == conds[1]]\n",
    "        n_cor_a = len(df_a[df_a[correct]==corrects[0]])\n",
    "        n_err_a = len(df_a[df_a[correct]==corrects[1]])\n",
    "\n",
    "        #print(data_s[data_s['choice']==1][0:n_cor_s])\n",
    "        ks_cor_s = get_ks_stat(data_s[data_s['choice']==1][0:n_cor_s],eval_model = False)\n",
    "        ks_err_s = get_ks_stat(data_s[data_s['choice']==0][0:n_err_s],eval_model = False)\n",
    "        ks_cor_a = get_ks_stat(data_a[data_a['choice']==1][0:n_cor_a],eval_model = False)\n",
    "        ks_err_a = get_ks_stat(data_a[data_a['choice']==0][0:n_err_a],eval_model = False)\n",
    "\n",
    "        ks = pd.DataFrame({\"cor_s\":[ks_cor_s], \"err_s\": [ks_err_s], \"cor_a\":[ks_cor_a],'err_a': [ks_err_a]})\n",
    "        ks[\"Participant\"] = i\n",
    "\n",
    "        ks_all = pd.concat((ks_all,ks))\n",
    "        \n",
    "    return ks_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62c00ec2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_normality_correct(ks_DDM,ks_OUM, ax=None):\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    \n",
    "    ax.scatter(ks_DDM[\"cor_s\"], ks_OUM[\"cor_s\"],c = [[1,1,1]], edgecolors = \"#1F77B4\", alpha = 0.7, label = \"Speed\")\n",
    "    ran = np.arange(0,max(ks_DDM[\"cor_s\"].max(axis=None),ks_OUM[\"cor_s\"].max(axis=None))+0.02,0.0001)\n",
    "    ax.plot(ran,ran,c = \"black\")\n",
    "    \n",
    "    ax.scatter(ks_DDM[\"cor_a\"], ks_OUM[\"cor_a\"],c = [[1,1,1]], edgecolors = \"#FF7F0E\", alpha = 0.7, label = \"Speed\")\n",
    "    ran = np.arange(0,max(ks_DDM[\"cor_a\"].max(axis=None),ks_OUM[\"cor_a\"].max(axis=None))+0.02,0.0001)\n",
    "    ax.plot(ran,ran,c = \"black\")\n",
    "    ax.set_title(\"Correct Choices\")\n",
    "    ax.set_xlabel(\"K-S Statistic of DDM\")\n",
    "    ax.set_ylabel(\"K-S Statistic of OUM\")\n",
    "    \n",
    "    return ax\n",
    "\n",
    "def plot_normality_error(ks_DDM,ks_OUM, ax=None):\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    \n",
    "    ax.scatter(ks_DDM[\"err_s\"], ks_OUM[\"err_s\"],c = [[1,1,1]], edgecolors = \"#1F77B4\", alpha = 0.7, label = \"Speed\")\n",
    "    ran = np.arange(0,max(ks_DDM[\"err_s\"].max(axis=None),ks_OUM[\"err_s\"].max(axis=None))+0.02,0.0001)\n",
    "    ax.plot(ran,ran,c = \"black\")\n",
    "    \n",
    "    ax.scatter(ks_DDM[\"err_a\"], ks_OUM[\"err_a\"],c = [[1,1,1]], edgecolors = \"#FF7F0E\", alpha = 0.7, label = \"Speed\")\n",
    "    ran = np.arange(0,max(ks_DDM[\"err_a\"].max(axis=None),ks_OUM[\"err_a\"].max(axis=None))+0.02,0.0001)\n",
    "    ax.plot(ran,ran,c = \"black\")\n",
    "    ax.set_title(\"Incorrect Choices\")\n",
    "    ax.set_xlabel(\"K-S Statistic of DDM\")\n",
    "    #ax.set_ylabel(\"K-S Statistic of OUM\")\n",
    "    \n",
    "    return ax"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}