SAT_help_functions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5cf8b400",
   "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",
    "plt.rc('font',family='Arial')\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\n",
    "import scipy.stats as stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e4182e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DriftConstantSpeedAcc(Drift):\n",
    "    name = \"DriftConstantSpeedAcc\"\n",
    "    required_parameters = [\"Vspeed\",\"Vacc\"]\n",
    "    required_conditions = ['Cond']\n",
    "    def get_drift(self, x, t, conditions, **kwargs):\n",
    "        if conditions['Cond'] == 'Speed':\n",
    "            return self.Vspeed\n",
    "        elif conditions['Cond'] == 'Accuracy':\n",
    "            return self.Vacc\n",
    "        else: \n",
    "            print(\"Non-implemented SAT conditions\")\n",
    "\n",
    "class DriftBothSpeedAcc(Drift):\n",
    "    name = \"DriftBothSpeedAcc\"\n",
    "    required_parameters = [\"Vspeed\",\"Vacc\", \"Kspeed\", \"Kacc\"]\n",
    "    required_conditions = ['Cond']\n",
    "        \n",
    "    def get_drift(self, x, t, conditions, **kwargs):\n",
    "        if conditions['Cond'] == 'Speed':\n",
    "            return self.Vspeed + x*self.Kspeed \n",
    "        elif conditions['Cond'] == 'Accuracy':\n",
    "            return self.Vacc + x*self.Kacc\n",
    "        else: \n",
    "            print(\"Non-implemented SAT conditions\")\n",
    "\n",
    "class DriftMixSpeedAcc(Drift):\n",
    "    name = \"DriftMixSpeedAcc\"\n",
    "    required_parameters = [\"Vspeed\",\"Vacc\", \"K\"]\n",
    "    required_conditions = ['Cond']\n",
    "        \n",
    "    def get_drift(self, x, t, conditions, **kwargs):\n",
    "        if conditions['Cond'] == 'Speed':\n",
    "            return self.Vspeed + x*self.K \n",
    "        elif conditions['Cond'] == 'Accuracy':\n",
    "            return self.Vacc + x*self.K\n",
    "        else: \n",
    "            print(\"Non-implemented SAT conditions\")\n",
    "            \n",
    "class DriftKSpeedAcc(Drift):\n",
    "    name = \"DriftMixSpeedAcc\"\n",
    "    required_parameters = [\"V\", \"Kspeed\",\"Kacc\"]\n",
    "    required_conditions = ['Cond']\n",
    "        \n",
    "    def get_drift(self, x, t, conditions, **kwargs):\n",
    "        if conditions['Cond'] == 'Speed':\n",
    "            return self.V + x*self.Kspeed\n",
    "        elif conditions['Cond'] == 'Accuracy':\n",
    "            return self.V + x*self.Kacc\n",
    "        else: \n",
    "            print(\"Non-implemented SAT conditions\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e294ac5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BoundSpeedAcc(Bound):\n",
    "    name = \"BoundSpeedAcc\"\n",
    "    required_parameters = [\"Bspeed\", \"Bacc\"]\n",
    "    required_conditions = ['Cond']\n",
    "    def get_bound(self, conditions, *args, **kwargs):\n",
    "        assert self.Bacc > 0\n",
    "        assert self.Bspeed > 0\n",
    "        if conditions['Cond'] == 'Speed':\n",
    "            return self.Bspeed\n",
    "        elif conditions['Cond'] == 'Accuracy':\n",
    "            return self.Bacc\n",
    "        else:\n",
    "            print(\"Non-implemented SAT conditions\")\n",
    "            \n",
    "class BoundCollapsingWeibull(Bound):\n",
    "    name = \"Weibull CDF collapsing bounds\"\n",
    "    required_parameters = [\"aSpeed\", \"aprimeSpeed\", \"lamSpeed\", \"kSpeed\", \"aAcc\", \"aprimeAcc\", \"lamAcc\", \"kAcc\"]\n",
    "    required_conditions = ['Cond']\n",
    "    def get_bound(self, t,conditions, **kwargs):\n",
    "        if conditions['Cond'] == 'Speed':\n",
    "            l = self.lamSpeed\n",
    "            a = self.aSpeed\n",
    "            aprime = self.aprimeSpeed\n",
    "            k = self.kSpeed\n",
    "            \n",
    "            return a - (1 - np.exp(-(t/l)**k)) * (a - aprime)\n",
    "        elif conditions['Cond'] == 'Accuracy':\n",
    "            l = self.lamAcc\n",
    "            a = self.aAcc\n",
    "            aprime = self.aprimeAcc\n",
    "            k = self.kAcc    \n",
    "            return a - (1 - np.exp(-(t/l)**k)) * (a - aprime)\n",
    "        else: \n",
    "            print(\"Non-implemented SAT conditions\")\n",
    "            \n",
    "class BoundCollapsingLinearSAT(BoundCollapsingLinear):\n",
    "    name = \"BoundSpeedAcc\"\n",
    "    required_parameters = [\"BSpeed\", \"BAcc\", \"cSpeed\", 'cAcc']\n",
    "    required_conditions = ['Cond']\n",
    "        \n",
    "    def get_bound(self, t, conditions, **kwargs):\n",
    "        if conditions['Cond'] == 'Speed':\n",
    "            return max(self.BSpeed - self.cSpeed*t, 0.)\n",
    "        elif conditions['Cond'] == 'Accuracy':\n",
    "            return max(self.BAcc - self.cAcc*t, 0.)\n",
    "        else: \n",
    "            print(\"Non-implemented SAT conditions\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70b47356",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NondecisionSpeedAcc(OverlayNonDecision):\n",
    "    name = \"NondecisionSpeedAcc\"\n",
    "    required_parameters = [\"tSpeed\", 'tAcc']\n",
    "    required_conditions = ['Cond']\n",
    "        \n",
    "    def get_nondecision_time(self, conditions, **kwargs):\n",
    "        if conditions['Cond'] == 'Speed':\n",
    "            return self.tSpeed \n",
    "        elif conditions['Cond'] == 'Accuracy':\n",
    "            return self.tAcc\n",
    "        else: \n",
    "            print(\"Non-implemented SAT conditions\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b99d83ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DDM_result_sub_2v2a2ter(samp, v = [0,6], a = [0.5,6], ter=[0.1,1], 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=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",
    "                      ),dx=.01, dt=dt, T_dur=T_dur)\n",
    "    fit_adjust_model(samp, model_fit,fitting_method=\"differential_evolution\",lossfunction=LossRobustBIC, verbose=False)\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['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": "4d711499",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DDM_result_sub_2v2a2c2ter(samp, v = [0,6], a = [0.5,6], c = [0,4], ter=[0.1,1], 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",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                      ),dx=.01, dt=dt, T_dur=T_dur)\n",
    "    fit_adjust_model(samp, model_fit,fitting_method=\"differential_evolution\",lossfunction=LossRobustBIC, verbose=False)\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['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": "a468c2f9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a8f3981",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DDM_result_sub_2a2ter(samp, v = [0,6], a = [0.5,6], ter=[0.1,1], dx=.01, dt=.01, T_dur=5):\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",
    "                      overlay=NondecisionSpeedAcc(\n",
    "                          tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                          tAcc=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                      ),dx=.01, dt=dt, T_dur=T_dur)\n",
    "    fit_adjust_model(samp, model_fit,fitting_method=\"differential_evolution\",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",
    "    \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": "8889af45",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2v2k2ter(samp, v = [0,6], a = [0.5,6], k = [-15,15], ter=[0.1,1], T_dur=5, dt=.01):\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=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,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": 1,
   "id": "25f31119",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2a2ter(samp, v = [0,6], a = [0.5,6], k = [-15,15], ter=[0.1,1], T_dur=5, dt=.01):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftLinear(drift = Fittable(minval=v[0], maxval=v[1]),\n",
    "                                             x=Fittable(minval=k[0], maxval=k[1]),\n",
    "                                             t=0), \n",
    "\n",
    "                      noise=NoiseConstant(noise=1),\n",
    "                      bound=BoundSpeedAcc(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=.01, 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,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",
    "                            'x': \"Self-excitation\",\n",
    "                            'tSpeed': 'Non-decision Time Speed',\n",
    "                            'tAcc': 'Non-decision Time Accuracy'})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9402e91",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf36f127",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2k2ter(samp, v = [0,6], a = [0.5,6], k = [-15,15], ter=[0.1,1], T_dur=5, dt=.01):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftKSpeedAcc(V = 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=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,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={'V': 'Drift Rate',\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": "0a6ed0fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2v2a2ter(samp, v = [0,6], a = [0.5,6], k = [-15,15], ter=[0.1,1], T_dur=5, dt=.01):\n",
    "    \n",
    "    model_fit = Model(name='Simple model (fitted)',\n",
    "                      drift=DriftMixSpeedAcc(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",
    "                      dx=.01, dt=dt, T_dur=T_dur)\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,7),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "\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": "c9fdb72c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_OUM_result_sub_2v2a2k2ter(samp, v = [0,6], a = [0.5,6], k = [-15,15], ter=[0.1,1], T_dur=5, dt=.01):\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=.01, dt=dt, T_dur=T_dur)\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,8),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \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": "7cd99347",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DDM_result_sub_2v2a2Weilbull2ter(samp, v = [0,6], a = [0.5,6], c = [0,4], ter=[0.1,1],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=BoundCollapsingWeibull(aSpeed=Fittable(minval=1, maxval=4),\n",
    "                                                   aprimeSpeed=Fittable(minval=0, maxval=2),\n",
    "                                                   lamSpeed=Fittable(minval=0.01, maxval=2),\n",
    "                                                   kSpeed=Fittable(minval=0, maxval=5),\n",
    "                                                   aAcc=Fittable(minval=1, maxval=4),\n",
    "                                                   aprimeAcc=Fittable(minval=0, maxval=2),\n",
    "                                                   lamAcc=Fittable(minval=0.01, maxval=2),\n",
    "                                                   kAcc=Fittable(minval=0, maxval=5)),\n",
    "                      overlay=NondecisionSpeedAcc(tSpeed=Fittable(minval=ter[0], maxval=ter[1]),\n",
    "                                                  tAcc=Fittable(minval=ter[0], maxval=ter[1])),\n",
    "                      dx=0.01, dt=dt, T_dur=T_dur)\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,12),\n",
    "                      columns = model_fit.get_model_parameter_names())\n",
    "    \n",
    "    df['BIC'] = model_fit.get_fit_result().value()\n",
    "    df = df.rename(columns={'Vspeed': 'Drift Rate Speed',\n",
    "                            'Vacc': 'Drift Rate Accuracy',\n",
    "                            'aSpeed': 'Boundary Separation Speed',\n",
    "                            'aAcc':   '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": "d5490376",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sim_fit_plot_all(par, par_fitted, par_labels, par_fitted_labels, fontsize, 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",
    "            #axes[j,i].locator_params(axis='both', nbins=4)\n",
    "            if j!=(param_fitted_num-1):\n",
    "                axes[j,i].set_xticks([])\n",
    "            if i!=0:\n",
    "                axes[j,i].set_yticks([])\n",
    "\n",
    "    for m in range(param_fitted_num):\n",
    "        axes[m,0].set_ylabel(par_fitted_labels[m],fontsize=fontsize)#, rotation=0)\n",
    "\n",
    "    for n in range(param_num):\n",
    "        axes[(param_fitted_num-1),n].set_xlabel(par_labels[n], fontsize=fontsize)\n",
    "\n",
    "    fig.align_ylabels(axes[:, 0])\n",
    "    return fig, axes\n"
   ]
  }
 ],
 "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
}