IDM_helper_functions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "90eaed2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "#os.chdir('/Users/wangjiashun/Documents/GitHub/sbi') \n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "53c165dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "from typing import Optional\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3bcbf4ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from sbi import utils as utils\n",
    "from sbi import analysis as analysis\n",
    "from sbi.inference.base import infer\n",
    "\n",
    "from sbi.neural_nets.mnle import MixedDensityEstimator, CategoricalNet\n",
    "from pyknos.nflows import flows, transforms\n",
    "from pyknos.nflows import distributions as distributions_\n",
    "\n",
    "\n",
    "from sbi.inference import MNLE\n",
    "from pyro.distributions import InverseGamma\n",
    "from torch.distributions import Beta, Binomial, Gamma\n",
    "from sbi.utils import MultipleIndependent\n",
    "\n",
    "from sbi.inference import MCMCPosterior, VIPosterior, RejectionPosterior\n",
    "from sbi.utils.torchutils import atleast_2d\n",
    "\n",
    "from sbi.utils import mcmc_transform\n",
    "from sbi.inference.potentials.base_potential import BasePotential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "13857fe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyddm import Model\n",
    "from pyddm.models import DriftConstant, DriftLinear, NoiseConstant, BoundConstant, OverlayNonDecision, ICPointSourceCenter\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52c091f3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2e0b298e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def energy(y1,y2,W_pos=52500,W_neg=8400,B1=0,B2=0,theta=51450):\n",
    "    # no multiple lines!!!\n",
    "    E = -W_pos*y1**2 - W_pos*y2**2 + W_neg*y1*y2 - (B1-theta)*y1 - (B2-theta)*y2\n",
    "    return E\n",
    "\n",
    "\n",
    "def dEdy1(y1, y2, W_pos, W_neg, B1, theta):\n",
    "    result = -W_pos*2*y1+W_neg*y2-(B1-theta)\n",
    "    return result\n",
    "\n",
    "\n",
    "def dEdy2(y1,y2, W_pos, W_neg, B2, theta):\n",
    "    result = -W_pos*2*y2+W_neg*y1-(B2-theta)\n",
    "    return result\n",
    "\n",
    "\n",
    "def free_energy(y1,y2,N=2000,W_pos=52500,W_neg=8400,B1=0,B2=0,beta=1/24,theta=51450):\n",
    "    sum_y = -N/2*(y1*np.log(y1)+(1-y1)*np.log(1-y1) + y2*np.log(y2)+(1-y2)*np.log(1-y2)) \n",
    "    # no multiple lines!!!\n",
    "    F = -W_pos*y1**2 - W_pos*y2**2 + W_neg*y1*y2 - (B1-theta)*y1 - (B2-theta)*y2 - beta**(-1)*sum_y\n",
    "    return F\n",
    "\n",
    "\n",
    "def dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N):\n",
    "    #print(log(y1/(1-y1)))\n",
    "    result = -W_pos*2*y1+W_neg*y2-(B1-theta)+beta**(-1)*0.5*N*(np.log(y1/(1-y1)))\n",
    "    return result\n",
    "\n",
    "\n",
    "def dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N):\n",
    "    result = -W_pos*2*y2+W_neg*y1-(B2-theta)+beta**(-1)*0.5*N*(np.log(y2/(1-y2)))\n",
    "    return result\n",
    "\n",
    "def clamp(n, minn=0.0001, maxn=0.9999):\n",
    "    return(max(min(maxn, n), minn))\n",
    "\n",
    "def get_state(y1,y2,h): # 1 for option 1, 0 for option 2, 100 for undecided\n",
    "    overlap = ((y1>(1-h))&(y2>(1-h)))&((y1<h)&(y2<h))\n",
    "    if overlap:\n",
    "        return 100\n",
    "    if(y1>(1-h))&(y2<h):\n",
    "        state = 1 # select 1\n",
    "    elif(y1<h)&(y2>(1-h)):\n",
    "        state = 0 # select 2\n",
    "    else:\n",
    "        state = 100 # undecided\n",
    "    \n",
    "    return state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2edb823d",
   "metadata": {},
   "outputs": [],
   "source": [
    "## get the probability for all states \n",
    "def all_states_prob(N = 1000, beta = 1/24): # here N is the number of neurons in one population\n",
    "    prob = np.zeros((N, N))\n",
    "    # first calculate Z\n",
    "    for i in range(N-1):\n",
    "        for j in range(N-1):\n",
    "            y1, y2 = (i+1)/N, (j+1)/N\n",
    "            prob[i,j] = np.exp(-beta*free_energy(y1, y2))\n",
    "\n",
    "    Z = np.sum(prob)\n",
    "    prob_scaled = prob/Z # divided by Z to make it a valid pmf\n",
    "    \n",
    "    return prob_scaled\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "30827926",
   "metadata": {},
   "outputs": [],
   "source": [
    "#def continue_sim(y1, y2, h):\n",
    "#    overlap = ((y1>(1-h))&(y2>(1-h)))&((y1<h)&(y2<h)) # in overlapped region?\n",
    "#    if overlap:\n",
    "#        return True\n",
    "#    else:\n",
    "#        return ((y1<(1-h))&(y2<(1-h)))|((y1>h)&(y2>h))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3665808e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "52385ac3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_one_trial(parameters, N = 2000, W_pos = 52500, W_neg = 8400, Bns = 2500, Bs = 1000, \n",
    "                        beta = 1/24, theta = 51450, D = 0.05, sigma = 0.01, \n",
    "                        dt = 0.001, init_loc = (0.3,0.3), mode = \"coarse-grained\"):\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3,\n",
    "    C = parameters[0]\n",
    "    h = parameters[1]\n",
    "    Ter = parameters[2]\n",
    "    \n",
    "    B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "    y1, y2 = init_loc[0], init_loc[1]\n",
    "    n_iter = 1\n",
    "    loc_y1, loc_y2 = [], []\n",
    "    \n",
    "    if mode == \"Euler-Maruyama\":\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            dW1 = np.random.normal()*sqrt(dt) \n",
    "            dW2 = np.random.normal()*sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "    \n",
    "    elif mode == \"coarse-grained\":\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            y_step = np.random.multivariate_normal(mean = [0,0], cov = [[sigma**2,0],[0,sigma**2]])\n",
    "            y_candidate1, y_candidate2 = clamp(y1+y_step[0]), clamp(y2+y_step[1])\n",
    "            df = free_energy(y_candidate1,y_candidate2,B1=B1,B2=B2) - free_energy(y1,y2,B1=B1,B2=B2)\n",
    "            # Matropolis acceptance rule\n",
    "            if df<=0: # accept the move\n",
    "                y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            else: # accept the move with a probability\n",
    "                p_accept = np.exp(-beta*df)\n",
    "                if p_accept >= np.random.uniform(1):\n",
    "                    y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "    else:\n",
    "            print(\"Wrong mode!\")\n",
    "    RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "    R = get_state(y1 = y1, y2 = y2, h = h)-1 # 0 -> choose option 1, 1 -> choose option 2\n",
    "    print(RT)\n",
    "    print(R)\n",
    "    return torch.cat((RT, R), dim=1)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29834228",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28dd6810",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1677e47e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_state_DDM(y,a):\n",
    "    if y > a:\n",
    "        state = 0  # 0 for option 1\n",
    "    elif y < 0:\n",
    "        state = 1  # 1 for option 2\n",
    "    else:\n",
    "        state = 2  # 2 for not finished\n",
    "    return state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "535fc5c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#parameters = np.array([[1,1,0.3], [1,1,0.3]])\n",
    "\n",
    "#simul_DDM(parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c5f84216",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_DDM(parameters, dt = 0.001, c = 0.01):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        v = parameters[i,0]\n",
    "        a = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        \n",
    "        init_loc = a*0.5\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y = init_loc\n",
    "\n",
    "        while(get_state_DDM(y,a)==2):\n",
    "            dW = np.random.normal()*np.sqrt(dt) \n",
    "            dy = v*dt + c*dW\n",
    "            y = y + dy\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "        \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_DDM(y, a) # different from IDM\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "73f5db76",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_DDM_five_param(parameters):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        v = parameters[i,0]\n",
    "        a = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        dt = parameters[i,3]\n",
    "        c = parameters[i,4]\n",
    "        \n",
    "        init_loc = a*0.5\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y = init_loc\n",
    "\n",
    "        while(get_state_DDM(y,a)==2):\n",
    "            dW = np.random.normal()*np.sqrt(dt) \n",
    "            dy = v*dt + c*dW\n",
    "            y = y + dy\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "        \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_DDM(y, a) # different from IDM\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "423d25de",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_EM(parameters, N = 2000, W_pos = 52500, W_neg = 8400, Bns = 2500, Bs = 500, \n",
    "                        beta = 1/24, theta = 51450,dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        D = parameters[i,3]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt) \n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "152f28ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_EM3(parameters, N = 2000, W_pos = 52500, W_neg = 8400, Bns = 2500, Bs = 1000, \n",
    "                        beta = 1/24, theta = 51450,D = 0.05, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt) \n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "\n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aed9a0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_EM1(parameters, C = 0.1, h = 0.4, Ter = 0, D = 0.05, N = 2000, W_pos = 52500, Bns = 0, Bs = 2500, \n",
    "                  beta = 1/24, theta = 51450, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    \n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "\n",
    "        \n",
    "        #W_pos = parameters[i,0]\n",
    "        W_neg = parameters[i,0]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt) \n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = 0.999, 0.001\n",
    "                else: \n",
    "                    y1, y2 = 0.001, 0.999\n",
    "                break\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "        \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f387c873",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_Inh_D(parameters, C = 0.2, h = 0.4, Ter = 0.5, N = 2000, Bns = 0, Bs = 2500, beta = 1/24, \n",
    "                  W_pos = 52500, theta = 51450, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    \n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "\n",
    "        \n",
    "        W_neg = parameters[i,0]\n",
    "        D = parameters[i,1]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt) \n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = 0.999, 0.001\n",
    "                else: \n",
    "                    y1, y2 = 0.001, 0.999\n",
    "                break\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "        \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ceaa8d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_Inh_Exc(parameters, C = 0.1, h = 0.4, Ter = 0, D = 0.05, N = 2000, Bns = 0, Bs = 2500, beta = 1/24, \n",
    "                  theta = 51450, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    \n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "\n",
    "        \n",
    "        W_pos = parameters[i,0]\n",
    "        W_neg = parameters[i,1]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt) \n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = 0.999, 0.001\n",
    "                else: \n",
    "                    y1, y2 = 0.001, 0.999\n",
    "                break\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "        \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "335b7a16",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_C_Inh_Exc(parameters, h = 0.4, Ter = 0.5, D = 0.05, N = 2000, Bns = 0, Bs = 2500, beta = 1/24, \n",
    "                  theta = 51450, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    \n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "\n",
    "        C = parameters[i,0]\n",
    "        W_pos = parameters[i,1]\n",
    "        W_neg = parameters[i,2]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns # 3000, 2000\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt) \n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = 0.999, 0.001\n",
    "                else: \n",
    "                    y1, y2 = 0.001, 0.999\n",
    "                break\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "        \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9969e4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_Inh_Exc(parameters, C = 0.2, h = 0.4, Ter = 0.5, D = 0.05, N = 2000, Bns = 0, Bs = 2500, beta = 1/24, \n",
    "                  theta = 51450, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    \n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "\n",
    "        \n",
    "        W_pos = parameters[i,0]\n",
    "        W_neg = parameters[i,1]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns # 3000, 2000\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt) \n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = 0.999, 0.001\n",
    "                else: \n",
    "                    y1, y2 = 0.001, 0.999\n",
    "                break\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "        \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d652ce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_EM5_Exc(parameters, N = 2000, Bns = 2500, Bs = 1000, beta = 1/24, \n",
    "                  theta = 51450, W_neg = 8400, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    \n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        D = parameters[i,3]\n",
    "        W_pos = parameters[i,4]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt) \n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = 0.999, 0.001\n",
    "                else: \n",
    "                    y1, y2 = 0.001, 0.999\n",
    "                break\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "        \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d0638e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_EM5_Inh(parameters, N = 2000, Bns = 2500, Bs = 1000, beta = 1/24, \n",
    "                  theta = 51450, W_pos = 52500, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    \n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        D = parameters[i,3]\n",
    "        W_neg = parameters[i,4]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt) \n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = 0.999, 0.001\n",
    "                else: \n",
    "                    y1, y2 = 0.001, 0.999\n",
    "                break\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "        \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce1f8537",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9209bd17",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13e5107e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe4cfdc9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "48d6df06",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_EM6_trajectory(C = 0.5, h = 0.3, Ter = 0.3, D = 0.05, \n",
    "                             W_pos = 52500, W_neg = 8400, N = 2000, Bns = 2000, Bs = 500, \n",
    "                             beta = 1/24, theta = 51450, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    \n",
    "\n",
    "    B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "    y1, y2 = init_loc[0], init_loc[1]\n",
    "    n_iter = 1\n",
    "    loc_y1, loc_y2 = [y1], [y2]\n",
    "\n",
    "    while get_state(y1=y1, y2=y2, h=h)==100:\n",
    "        dW1 = np.random.normal()*np.sqrt(dt) \n",
    "        dW2 = np.random.normal()*np.sqrt(dt)\n",
    "        dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "        dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "        y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "        y2 = clamp(y2 + dy2)\n",
    "        loc_y1.append(y1)\n",
    "        loc_y2.append(y2)\n",
    "        n_iter = n_iter + 1\n",
    "\n",
    "    return loc_y1, loc_y2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dd942b7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_EM6(parameters, N = 2000, Bns = 2000, Bs = 500, beta = 1/24, \n",
    "                  theta = 51450, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    \n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        D = parameters[i,3]   #D = 0.05\n",
    "        W_pos = parameters[i,4]\n",
    "        W_neg = parameters[i,5]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==100:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt) \n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                break\n",
    "\n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)   # 1 for option 1, 0 for option 2, 100 for undecided\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "        \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3b1f3ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_CG_four_param(parameters, N = 2000, W_pos = 52500, W_neg = 8400, Bns = 2500, Bs = 1000, \n",
    "                        beta = 1/24, theta = 51450, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        sigma = parameters[i,3]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            y_step = np.random.multivariate_normal(mean = [0,0], cov = [[sigma**2,0],[0,sigma**2]])\n",
    "            y_candidate1, y_candidate2 = clamp(y1+y_step[0]), clamp(y2+y_step[1])\n",
    "            df = free_energy(y_candidate1,y_candidate2,B1=B1,B2=B2) - free_energy(y1,y2,B1=B1,B2=B2)\n",
    "            # Matropolis acceptance rule\n",
    "            if df<=0: # accept the move\n",
    "                y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            else: # accept the move with a probability\n",
    "                p_accept = np.exp(-beta*df)\n",
    "                if p_accept >= np.random.uniform(1):\n",
    "                    y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = 0.999, 0.001\n",
    "                else: \n",
    "                    y1, y2 = 0.001, 0.999\n",
    "                break\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)\n",
    "\n",
    "def simul_IDM_CG_five_param(parameters, N = 2000, W_pos = 52500, W_neg = 8400, Bns = 2500, Bs = 1000, \n",
    "                        beta = 1/24, theta = 51450, init_loc = (0.3,0.3)):\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        sigma = parameters[i,3]\n",
    "        dt = parameters[i,4]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            y_step = np.random.multivariate_normal(mean = [0,0], cov = [[sigma**2,0],[0,sigma**2]])\n",
    "            y_candidate1, y_candidate2 = clamp(y1+y_step[0]), clamp(y2+y_step[1])\n",
    "            df = free_energy(y_candidate1,y_candidate2,B1=B1,B2=B2) - free_energy(y1,y2,B1=B1,B2=B2)\n",
    "            # Matropolis acceptance rule\n",
    "            if df<=0: # accept the move\n",
    "                y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            else: # accept the move with a probability\n",
    "                p_accept = np.exp(-beta*df)\n",
    "                if p_accept >= np.random.uniform(1):\n",
    "                    y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "266e3ee4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_CG_inh_pos(parameters, N = 2000, W_pos = 52500, Bns = 2500, Bs = 1000, \n",
    "                        beta = 1/24, theta = 51450, dt = 0.001, sigma = 0.01, init_loc = (0.3,0.3)):\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        W_neg = parameters[i,2]\n",
    "        W_pos = parameters[i,3]\n",
    "        #Ter = parameters[i,3]\n",
    "        #sigma = parameters[i,3]\n",
    "        Ter = 0.5\n",
    "        \n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            y_step = np.random.multivariate_normal(mean = [0,0], cov = [[sigma**2,0],[0,sigma**2]])\n",
    "            y_candidate1, y_candidate2 = clamp(y1+y_step[0]), clamp(y2+y_step[1])\n",
    "            df = free_energy(y_candidate1,y_candidate2,B1=B1,B2=B2,W_neg=W_neg, W_pos = W_pos) - free_energy(y1,y2,B1=B1,B2=B2,W_neg=W_neg, W_pos = W_pos)\n",
    "            # Matropolis acceptance rule\n",
    "            if df<=0: # accept the move\n",
    "                y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            else: # accept the move with a probability\n",
    "                p_accept = np.exp(-beta*df)\n",
    "                if p_accept >= np.random.uniform(1):\n",
    "                    y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if (n_iter*dt + Ter)>3:\n",
    "                break\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        \n",
    "        if get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            R = np.random.choice([0,1]) ## random choose when longer than 3s\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "        #print(parameters[i,:])\n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd867e35",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_CG_inh(parameters, N = 2000, W_pos = 52500, Bns = 2500, Bs = 1000, \n",
    "                        beta = 1/24, theta = 51450, dt = 0.001, sigma = 0.01, init_loc = (0.3,0.3)):\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        W_neg = parameters[i,2]\n",
    "        #Ter = parameters[i,3]\n",
    "        #sigma = parameters[i,3]\n",
    "        Ter = 0.5\n",
    "        \n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            y_step = np.random.multivariate_normal(mean = [0,0], cov = [[sigma**2,0],[0,sigma**2]])\n",
    "            y_candidate1, y_candidate2 = clamp(y1+y_step[0]), clamp(y2+y_step[1])\n",
    "            df = free_energy(y_candidate1,y_candidate2,B1=B1,B2=B2,W_neg=W_neg) - free_energy(y1,y2,B1=B1,B2=B2,W_neg=W_neg)\n",
    "            # Matropolis acceptance rule\n",
    "            if df<=0: # accept the move\n",
    "                y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            else: # accept the move with a probability\n",
    "                p_accept = np.exp(-beta*df)\n",
    "                if p_accept >= np.random.uniform(1):\n",
    "                    y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if (n_iter*dt + Ter)>3:\n",
    "                break\n",
    "\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        \n",
    "        if get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            R = np.random.choice([0,1]) ## random choose when longer than 3s\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "        #print(parameters[i,:])\n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc137813",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97bc9fe9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_CG(parameters, N = 2000, W_pos = 52500, W_neg = 8400, Bns = 2500, Bs = 1000, \n",
    "                        beta = 1/24, theta = 51450, dt = 0.001, sigma = 0.002, init_loc = (0.3,0.3)):\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        #sigma = parameters[i,3]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "            y_step = np.random.multivariate_normal(mean = [0,0], cov = [[sigma**2,0],[0,sigma**2]])\n",
    "            y_candidate1, y_candidate2 = clamp(y1+y_step[0]), clamp(y2+y_step[1])\n",
    "            df = free_energy(y_candidate1,y_candidate2,B1=B1,B2=B2) - free_energy(y1,y2,B1=B1,B2=B2)\n",
    "            # Matropolis acceptance rule\n",
    "            if df<=0: # accept the move\n",
    "                y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            else: # accept the move with a probability\n",
    "                p_accept = np.exp(-beta*df)\n",
    "                if p_accept >= np.random.uniform(1):\n",
    "                    y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = 0.999, 0.001\n",
    "                else: \n",
    "                    y1, y2 = 0.001, 0.999\n",
    "                break\n",
    "                \n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc2a5b1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_more_param(parameters, N = 2000, W_pos = 52500, W_neg = 8400, Bns = 2500, Bs = 1000, \n",
    "                        beta = 1/24, theta = 51450, D = 0.05, \n",
    "                        dt = 0.001, init_loc = (0.3,0.3), mode = \"coarse-grained\"):\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        sigma = parameters[i,3]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        if mode == \"Euler-Maruyama\":\n",
    "            while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "                dW1 = np.random.normal()*sqrt(dt) \n",
    "                dW2 = np.random.normal()*sqrt(dt)\n",
    "                dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + sqrt(2*D)*dW1\n",
    "                dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + sqrt(2*D)*dW2\n",
    "                y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "                y2 = clamp(y2 + dy2)\n",
    "                loc_y1.append(y1)\n",
    "                loc_y2.append(y2)\n",
    "                n_iter = n_iter + 1\n",
    "\n",
    "        elif mode == \"coarse-grained\":\n",
    "            while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "                y_step = np.random.multivariate_normal(mean = [0,0], cov = [[sigma**2,0],[0,sigma**2]])\n",
    "                y_candidate1, y_candidate2 = clamp(y1+y_step[0]), clamp(y2+y_step[1])\n",
    "                df = free_energy(y_candidate1,y_candidate2,B1=B1,B2=B2) - free_energy(y1,y2,B1=B1,B2=B2)\n",
    "                # Matropolis acceptance rule\n",
    "                if df<=0: # accept the move\n",
    "                    y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "                else: # accept the move with a probability\n",
    "                    p_accept = np.exp(-beta*df)\n",
    "                    if p_accept >= np.random.uniform(1):\n",
    "                        y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "                loc_y1.append(y1)\n",
    "                loc_y2.append(y2)\n",
    "                n_iter = n_iter + 1\n",
    "        else:\n",
    "                print(\"Wrong mode!\")\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "3f78dba8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_mult_trial(parameters, N = 2000, W_pos = 52500, W_neg = 8400, Bns = 2500, Bs = 1000, \n",
    "                        beta = 1/24, theta = 51450, D = 0.05, sigma = 0.01, \n",
    "                        dt = 0.001, init_loc = (0.3,0.3), mode = \"coarse-grained\"):\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        if mode == \"Euler-Maruyama\":\n",
    "            while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "                dW1 = np.random.normal()*sqrt(dt) \n",
    "                dW2 = np.random.normal()*sqrt(dt)\n",
    "                dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + sqrt(2*D)*dW1\n",
    "                dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + sqrt(2*D)*dW2\n",
    "                y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "                y2 = clamp(y2 + dy2)\n",
    "                loc_y1.append(y1)\n",
    "                loc_y2.append(y2)\n",
    "                n_iter = n_iter + 1\n",
    "\n",
    "        elif mode == \"coarse-grained\":\n",
    "            while get_state(y1=y1, y2=y2, h=h)==0:\n",
    "                y_step = np.random.multivariate_normal(mean = [0,0], cov = [[sigma**2,0],[0,sigma**2]])\n",
    "                y_candidate1, y_candidate2 = clamp(y1+y_step[0]), clamp(y2+y_step[1])\n",
    "                df = free_energy(y_candidate1,y_candidate2,B1=B1,B2=B2) - free_energy(y1,y2,B1=B1,B2=B2)\n",
    "                # Matropolis acceptance rule\n",
    "                if df<=0: # accept the move\n",
    "                    y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "                else: # accept the move with a probability\n",
    "                    p_accept = np.exp(-beta*df)\n",
    "                    if p_accept >= np.random.uniform(1):\n",
    "                        y1, y2 = y_candidate1, y_candidate2 # clamp to [0,1]\n",
    "                loc_y1.append(y1)\n",
    "                loc_y2.append(y2)\n",
    "                n_iter = n_iter + 1\n",
    "        else:\n",
    "                print(\"Wrong mode!\")\n",
    "        RT = clamp(n_iter*dt + Ter, 0, 3)\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)-1\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c180d436",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot them in one pairplot as contours (obtained via KDE on the samples).\n",
    "from sbi.analysis import pairplot\n",
    "\n",
    "#parameters_true = [[0.2,0.4,0.3]]\n",
    "def pair_plot(prior, mnle_samples, parameters_o, figsize = (5,5), contour = 0.7, labels=['C', 'h','Ter']):\n",
    "    fig, ax = pairplot(\n",
    "        [\n",
    "            prior.sample((1000,)),\n",
    "            mnle_samples,\n",
    "        ],\n",
    "        points=parameters_o,\n",
    "        diag=\"kde\",\n",
    "        upper=\"contour\",\n",
    "        kde_offdiag=dict(bins=50),\n",
    "        kde_diag=dict(bins=100),\n",
    "        contour_offdiag=dict(levels=[contour]),\n",
    "        points_colors=[\"k\"],\n",
    "        points_offdiag=dict(marker=\"*\", markersize=10),\n",
    "        labels=labels,\n",
    "        figsize=figsize\n",
    "    )\n",
    "\n",
    "\n",
    "    plt.legend(\n",
    "    [\"Prior\", \"MNLE\", 'True Parameters'], fontsize=10, bbox_to_anchor=(-3, 1), loc=2\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "8af17d88",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sim_fit_plot(Vs, As, Ters, par_fitted, sim_num = 30):\n",
    "\n",
    "    fig, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, figsize=(10,10))\n",
    "\n",
    "    ax1.scatter(Vs[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "    ax2.scatter(As[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "    ax3.scatter(Ters[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "    ax4.scatter(Vs[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "    ax5.scatter(As[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "    ax6.scatter(Ters[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "    ax7.scatter(Vs[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "    ax8.scatter(As[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "    ax9.scatter(Ters[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "\n",
    "    ax1.set_title(\"v\")\n",
    "    ax2.set_title(\"a\")\n",
    "    ax3.set_title(\"ter\")\n",
    "\n",
    "    ax1.set_ylabel(\"C\", rotation=0)\n",
    "    ax4.set_ylabel(\"h\", rotation=0)\n",
    "    ax7.set_ylabel(\"ter\", rotation=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f09aafd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sim_fit_plot_4(Vs, As, Ters, cs, par_fitted, sim_num = 30):\n",
    "\n",
    "    fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8), (ax9, ax10, ax11, ax12), (ax13, ax14, ax15, ax16)) = plt.subplots(4, 4, figsize=(10,10))\n",
    "\n",
    "    ax1.scatter(Vs[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "    ax2.scatter(As[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "    ax3.scatter(Ters[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "    ax4.scatter(cs[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "\n",
    "    ax5.scatter(Vs[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "    ax6.scatter(As[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "    ax7.scatter(Ters[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "    ax8.scatter(cs[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "\n",
    "    ax9.scatter(Vs[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "    ax10.scatter(As[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "    ax11.scatter(Ters[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "    ax12.scatter(cs[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "    \n",
    "    \n",
    "    ax13.scatter(Vs[0:sim_num], par_fitted[0:sim_num,3], s = 2)\n",
    "    ax14.scatter(As[0:sim_num], par_fitted[0:sim_num,3], s = 2)\n",
    "    ax15.scatter(Ters[0:sim_num], par_fitted[0:sim_num,3], s = 2)\n",
    "    ax16.scatter(cs[0:sim_num], par_fitted[0:sim_num,3], s = 2)\n",
    "\n",
    "    ax1.set_title(\"v\")\n",
    "    ax2.set_title(\"a\")\n",
    "    ax3.set_title(\"ter\")\n",
    "    ax4.set_title(\"c\")\n",
    "\n",
    "    ax1.set_ylabel(\"C\", rotation=0)\n",
    "    ax5.set_ylabel(\"h\", rotation=0)\n",
    "    ax9.set_ylabel(\"ter\", rotation=0)\n",
    "    ax13.set_ylabel(\"sigma\", rotation=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dc8faeb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10ee277a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeca3a9b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "f830b887",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def simul_DDM_mult_trial(v, a, ter, n_trials):\n",
    "    model = Model(name='Simple model',\n",
    "              drift=DriftConstant(drift=v),\n",
    "              bound=BoundConstant(B=a),\n",
    "              overlay=OverlayNonDecision(nondectime=ter),\n",
    "              noise=NoiseConstant(noise=1),\n",
    "              dx=.001, dt=.001, T_dur=5)\n",
    "    sol = model.solve()\n",
    "\n",
    "    # generate pseudo-random data\n",
    "    samp = sol.resample(n_trials)\n",
    "\n",
    "\n",
    "    len_corr = len(samp.corr)\n",
    "    len_err = len(samp.err)\n",
    "    DDM_sim_data = np.zeros((len_corr+len_err, 2))\n",
    "    DDM_sim_data[:len_corr,0] = samp.corr\n",
    "    DDM_sim_data[len_corr:len_corr+len_err,0] = samp.err\n",
    "\n",
    "    # coded as correct/error, like IDM\n",
    "    DDM_sim_data[:len_corr,1] = 0 # correct\n",
    "    DDM_sim_data[len_corr:len_corr+len_err,1] = 1 # error\n",
    "    return(torch.tensor(DDM_sim_data))\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a8ed621",
   "metadata": {},
   "outputs": [],
   "source": [
    "#from pyddm import Model\n",
    "#from pyddm.models import DriftConstant, NoiseConstant, BoundConstant, OverlayNonDecision, ICPointSourceCenter\n",
    "#from pyddm.functions import fit_adjust_model, display_model\n",
    "\n",
    "#Vs = np.random.uniform(low=0, high=3, size=30)\n",
    "#As = np.random.uniform(low=0.4, high=0.9, size=30)\n",
    "#Ters = np.random.uniform(low=0.1, high=0.7, size=30)\n",
    "\n",
    "#n_trials = 100\n",
    "#num_samples = 100\n",
    "#par_fitted = np.zeros((30,3))\n",
    "#for i in range(len(Vs)):\n",
    "#    print(i)\n",
    "#    v,a,ter = Vs[i], As[i], Ters[i]\n",
    "#    DDM_o = simul_DDM_mult_trial(v, a, ter, n_trials) # simulate DDM data\n",
    "    \n",
    "#    mnle_samples = mnle_posterior.sample((num_samples,), x=x_o) # fit with IDM\n",
    "#    mnle_posterior.set_default_x(DDM_o)\n",
    "#    mnle_map = mnle_posterior.map()\n",
    "    \n",
    "#    par_fitted[i,:] = np.array(mnle_map)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b29e223",
   "metadata": {},
   "outputs": [],
   "source": [
    "### simulate LCA from my code\n",
    "def get_state_LCA(y,a):\n",
    "    if y > a:\n",
    "        state = 0  # 0 for option 1\n",
    "    elif y < -a:\n",
    "        state = 1  # 1 for option 2\n",
    "    else:\n",
    "        state = 2  # 2 for not finished\n",
    "    return state\n",
    "\n",
    "def simul_LCA(parameters, init_loc = 0, k = 0.05, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        I = parameters[i,0]\n",
    "        a = parameters[i,1]\n",
    "        inh = parameters[i,2]\n",
    "        #Ter = parameters[i,3]\n",
    "        Ter = 0.5\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y = init_loc\n",
    "\n",
    "        while(get_state_LCA(y,a)==2):\n",
    "            dW = np.random.normal()*np.sqrt(dt) \n",
    "            dy = (I - k*y + inh*y)*dt + c*dW*np.sqrt(2)\n",
    "            y = y + dy\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "        \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_LCA(y, a) # different from IDM\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)\n",
    "\n",
    "def simul_LCA_four_param(parameters, init_loc = 0, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        I = parameters[i,0]\n",
    "        a = parameters[i,1]\n",
    "        inh = parameters[i,2]\n",
    "        k = parameters[i,3]\n",
    "        Ter = 0.5\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y = init_loc\n",
    "\n",
    "        while(get_state_LCA(y,a)==2):\n",
    "            dW = np.random.normal()*np.sqrt(dt) \n",
    "            dy = (I - k*y + inh*y)*dt + c*dW*np.sqrt(2)\n",
    "            y = y + dy\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "        \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_LCA(y, a) # different from IDM\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)\n",
    "\n",
    "def simul_LCA_all_param(parameters, init_loc = 0, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        I = parameters[i,0]\n",
    "        a = parameters[i,1]\n",
    "        inh = parameters[i,2]\n",
    "        Ter = parameters[i,3]\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y = init_loc\n",
    "\n",
    "        while(get_state_LCA(y,a)==2):\n",
    "            dW = np.random.normal()*np.sqrt(dt) \n",
    "            dy = (I + inh*y)*dt + c*dW*np.sqrt(2)\n",
    "            y = y + dy\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "        \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_LCA(y, a) # different from IDM\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "    \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0734fc8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_state_LCA_full(y1,y2,a):  ## can we also add the second threshold for Andrew's model?\n",
    "    #print(y1,y2,a)\n",
    "    if (y1 > a) and (y2 < a):\n",
    "        state = 0  # 0 for option 1\n",
    "    elif (y2 > a) and (y1 < a):\n",
    "        state = 1  # 1 for option 2\n",
    "    elif (y1 > a) and (y2 > a):        \n",
    "        if y1 > y2:\n",
    "            state = 0\n",
    "        if y1 <= y2:\n",
    "            state = 1\n",
    "    else:\n",
    "        state = 2  # 2 for not finished\n",
    "    return state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09579c60",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_state_LCA_full_3AFT(y1,y2,y3,a):  ## can we also add the second threshold for Andrew's model?\n",
    "    #print(y1,y2,a)\n",
    "    if (y1 > a) and ((y2 < a) and (y3 < a)):\n",
    "        state = 0  # 0 for option 1\n",
    "    elif (y2 > a) and ((y1 < a) and (y3 < a)):\n",
    "        state = 1  # 1 for option 2\n",
    "    elif (y3 > a) and (y1 < a) and (y2 < a):\n",
    "        state = 2\n",
    "    else:\n",
    "        state = 3  # 3 for not finished\n",
    "    return state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cbd18bb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "642eb9bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_state_LCA_full_3AFT1(y1,y2,y3,a):  ## can we also add the second threshold for Andrew's model?\n",
    "    #print(y1,y2,a)\n",
    "    if max(y1,y2,y3) > a:\n",
    "        index = [y1,y2,y3].index(max(y1,y2,y3))\n",
    "        return index\n",
    "    else:\n",
    "        return 3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "589ab558",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_LCA4_3AFT(parameters, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        dI = parameters[i,0]\n",
    "        #I = parameters[i,1]\n",
    "        I = 1\n",
    "        inh = parameters[i,1]\n",
    "        k = parameters[i,2]\n",
    "        a = parameters[i,3]\n",
    "        Ter = 0\n",
    "        \n",
    "        I1 = I + dI\n",
    "        I2 = I\n",
    "        I3 = I\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y1,y2,y3 = 0,0,0 # initial location\n",
    "        y1s,y2s,y3s = [],[],[]\n",
    "        while(get_state_LCA_full_3AFT1(y1,y2,y3,a)==3):\n",
    "            dW1, dW2, dW3 = np.random.normal()*np.sqrt(dt), np.random.normal()*np.sqrt(dt),np.random.normal()*np.sqrt(dt)  \n",
    "\n",
    "            dy1 = (I1 - k*y1 - inh*y2 - inh*y3)*dt + c*dW1\n",
    "            dy2 = (I2 - k*y2 - inh*y1 - inh*y3)*dt + c*dW2\n",
    "            dy3 = (I3 - k*y3 - inh*y1 - inh*y2)*dt + c*dW3\n",
    " \n",
    "            y1 = clamp(y1 + dy1, 0, 100)\n",
    "            y2 = clamp(y2 + dy2, 0, 100)\n",
    "            y3 = clamp(y3 + dy3, 0, 100)\n",
    "            y1s.append(y1)\n",
    "            y2s.append(y2)\n",
    "            y3s.append(y3)\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()<(1/3):\n",
    "                    y1, y2, y3 = a+0.01, a-0.01, a-0.01\n",
    "                elif np.random.random()<(1/2):\n",
    "                    y1, y2, y3 = a-0.01, a+0.01, a-0.01\n",
    "                else: \n",
    "                    y1, y2, y3 = a-0.01, a-0.01, a+0.01\n",
    "                break  \n",
    "            \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_LCA_full_3AFT(y1,y2,y3,a) # different from IDM\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "      \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "        \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1) #,y1s,y2s,y3s, n_iter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a81236d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_LCA5_3AFT(parameters, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        dI = parameters[i,0]\n",
    "        #I = parameters[i,1]\n",
    "        I = 1\n",
    "        inh = parameters[i,1]\n",
    "        k = parameters[i,2]\n",
    "        a = parameters[i,3]\n",
    "        Ter = parameters[i,4]\n",
    "        \n",
    "        I1 = I + dI\n",
    "        I2 = I\n",
    "        I3 = I\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y1,y2,y3 = 0,0,0 # initial location\n",
    "        y1s,y2s,y3s = [],[],[]\n",
    "        while(get_state_LCA_full_3AFT1(y1,y2,y3,a)==3):\n",
    "            dW1, dW2, dW3 = np.random.normal()*np.sqrt(dt), np.random.normal()*np.sqrt(dt),np.random.normal()*np.sqrt(dt)  \n",
    "\n",
    "            dy1 = (I1 - k*y1 - inh*y2 - inh*y3)*dt + c*dW1\n",
    "            dy2 = (I2 - k*y2 - inh*y1 - inh*y3)*dt + c*dW2\n",
    "            dy3 = (I3 - k*y3 - inh*y1 - inh*y2)*dt + c*dW3\n",
    " \n",
    "            y1 = clamp(y1 + dy1, 0, 100)\n",
    "            y2 = clamp(y2 + dy2, 0, 100)\n",
    "            y3 = clamp(y3 + dy3, 0, 100)\n",
    "            y1s.append(y1)\n",
    "            y2s.append(y2)\n",
    "            y3s.append(y3)\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()<(1/3):\n",
    "                    y1, y2, y3 = a+0.01, a-0.01, a-0.01\n",
    "                elif np.random.random()<(1/2):\n",
    "                    y1, y2, y3 = a-0.01, a+0.01, a-0.01\n",
    "                else: \n",
    "                    y1, y2, y3 = a-0.01, a-0.01, a+0.01\n",
    "                break  \n",
    "            \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_LCA_full_3AFT(y1,y2,y3,a) # different from IDM\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "      \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "        \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1) #,y1s,y2s,y3s, n_iter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5a68d04",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_LCA_full_3AFT(parameters, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        dI = parameters[i,0]\n",
    "        I = parameters[i,1]\n",
    "        inh = parameters[i,2]\n",
    "        k = parameters[i,3]\n",
    "        a = parameters[i,4]\n",
    "        Ter = parameters[i,5]\n",
    "        \n",
    "        I1 = I + dI\n",
    "        I2 = I\n",
    "        I3 = I\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y1,y2,y3 = 0,0,0 # initial location\n",
    "\n",
    "        while(get_state_LCA_full_3AFT(y1,y2,y3,a)==3):\n",
    "            dW1, dW2, dW3 = np.random.normal()*np.sqrt(dt), np.random.normal()*np.sqrt(dt),np.random.normal()*np.sqrt(dt)  \n",
    "\n",
    "            dy1 = (I1 - k*y1 - inh*y2 - inh*y3)*dt + c*dW1\n",
    "            dy2 = (I2 - k*y2 - inh*y1 - inh*y3)*dt + c*dW2\n",
    "            dy3 = (I3 - k*y3 - inh*y1 - inh*y2)*dt + c*dW3\n",
    " \n",
    "            y1 = clamp(y1 + dy1, 0, 100)\n",
    "            y2 = clamp(y2 + dy2, 0, 100)\n",
    "            y3 = clamp(y3 + dy3, 0, 100)\n",
    "            \n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()<(1/3):\n",
    "                    y1, y2, y3 = a+0.01, a-0.01, a-0.01\n",
    "                elif np.random.random()<(1/2):\n",
    "                    y1, y2, y3 = a-0.01, a+0.01, a-0.01\n",
    "                else: \n",
    "                    y1, y2, y3 = a-0.01, a-0.01, a+0.01\n",
    "                break  \n",
    "            \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_LCA_full_3AFT(y1,y2,y3,a) # different from IDM\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "      \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "            \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adbb8b3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_LCA_full5(parameters, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        I = 1\n",
    "        dI = parameters[i,0]\n",
    "        a = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        inh = parameters[i,3]\n",
    "        k = parameters[i,4]\n",
    "        \n",
    "        I1 = I + dI\n",
    "        I2 = I\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y1,y2 = 0,0 # initial location\n",
    "\n",
    "        while(get_state_LCA_full(y1,y2,a)==2):\n",
    "            dW1, dW2 = np.random.normal()*np.sqrt(dt), np.random.normal()*np.sqrt(dt) \n",
    "\n",
    "            dy1 = (I1 - k*y1 - inh*y2)*dt + c*dW1\n",
    "            dy2 = (I2 - k*y2 - inh*y1)*dt + c*dW2\n",
    "\n",
    "            \n",
    "            y1 = clamp(y1 + dy1, 0, 100)\n",
    "            y2 = clamp(y2 + dy2, 0, 100)\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = a+0.01, a-0.01\n",
    "                else: \n",
    "                    y1, y2 = a-0.01, a+0.01\n",
    "                break  \n",
    "            \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_LCA_full(y1,y2,a) # different from IDM\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "      \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "            \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "471b5c36",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_LCA_full(parameters, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        dI = parameters[i,0]\n",
    "        I = parameters[i,1]\n",
    "        inh = parameters[i,2]\n",
    "        k = parameters[i,3]\n",
    "        a = parameters[i,4]\n",
    "        Ter = parameters[i,5]\n",
    "        \n",
    "        I1 = I + dI\n",
    "        I2 = I\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y1,y2 = 0,0 # initial location\n",
    "\n",
    "        while(get_state_LCA_full(y1,y2,a)==2):\n",
    "            dW1, dW2 = np.random.normal()*np.sqrt(dt), np.random.normal()*np.sqrt(dt) \n",
    "\n",
    "            dy1 = (I1 - k*y1 - inh*y2)*dt + c*dW1\n",
    "            dy2 = (I2 - k*y2 - inh*y1)*dt + c*dW2\n",
    "\n",
    "            \n",
    "            y1 = clamp(y1 + dy1, 0, 100)\n",
    "            y2 = clamp(y2 + dy2, 0, 100)\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = a+0.01, a-0.01\n",
    "                else: \n",
    "                    y1, y2 = a-0.01, a+0.01\n",
    "                break  \n",
    "            \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_LCA_full(y1,y2,a) # different from IDM\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "      \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "            \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "58323483",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sim0\n",
      "tensor(0.7900)\n",
      "tensor([[1.4690, 1.0000],\n",
      "        [1.7440, 0.0000],\n",
      "        [3.4000, 1.0000],\n",
      "        [1.9320, 0.0000],\n",
      "        [3.3490, 0.0000],\n",
      "        [0.8570, 0.0000],\n",
      "        [3.1590, 0.0000],\n",
      "        [3.4000, 1.0000],\n",
      "        [1.1220, 0.0000],\n",
      "        [2.7230, 0.0000]])\n"
     ]
    }
   ],
   "source": [
    "#x = simul_LCA_full_simp(np.tile([0.5542, 1.4183, 0.2164], (100,1)))\n",
    "#print((100-sum(x[:,1]))/100)\n",
    "#print(x[0:10,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adc43c36",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_OU_analytical(parameters, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        lam = parameters[i,0]\n",
    "        y = parameters[i,1]\n",
    "        theta = parameters[i,2]\n",
    "\n",
    "        n_iter = 1\n",
    "        #y = 0 # initial location\n",
    "\n",
    "        while((y<a) and (y>-a)):\n",
    "            dW = np.random.normal()*np.sqrt(dt)\n",
    "            dy = lam*(theta-y)*dt + c*dW\n",
    "            y = y + dy\n",
    "\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y = a+0.01\n",
    "                else:\n",
    "                    y = -a-0.01\n",
    "                break\n",
    "            \n",
    "        RT = n_iter*dt + Ter\n",
    "        #print(y)\n",
    "        if y > a:\n",
    "            R = 0\n",
    "        if y <-a:\n",
    "            R = 1     # different from IDM\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "      \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "            \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6718241",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_Leaky_accum(parameters, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        dI = parameters[i,0]\n",
    "        I = 1\n",
    "        k = parameters[i,1]\n",
    "        a = parameters[i,2]\n",
    "        Ter = parameters[i,3]\n",
    "        \n",
    "        I1 = I + dI\n",
    "        I2 = I\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y1,y2 = 0,0 # initial location\n",
    "\n",
    "        while(get_state_LCA_full(y1,y2,a)==2):\n",
    "            dW1, dW2 = np.random.normal()*np.sqrt(dt), np.random.normal()*np.sqrt(dt) \n",
    "\n",
    "            dy1 = (I1 + k*y1)*dt + c*dW1\n",
    "            dy2 = (I2 + k*y2)*dt + c*dW2\n",
    "\n",
    "            y1 = clamp(y1 + dy1, 0, 100)\n",
    "            y2 = clamp(y2 + dy2, 0, 100)\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y1, y2 = a+0.01, a-0.01\n",
    "                else: \n",
    "                    y1, y2 = a-0.01, a+0.01\n",
    "                break  \n",
    "            \n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state_LCA_full(y1,y2,a) # different from IDM\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "      \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "            \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad63293e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def g(lam0, t0, z0, b0):\n",
    "    \n",
    "    t = lam0*t0\n",
    "    z = np.sqrt(lam0)*(z0-theta0)\n",
    "    b = np.sqrt(lam0)*(b0-theta0)\n",
    "    \n",
    "    numerator = (z-b)*exp(b*(z-b) - 0.5*e^(-t)(z-b)^2/np.sinh(t) + t/2)\n",
    "    denominator = np.sqrt(2*np.pi*np.sinh(t)**3)\n",
    "    \n",
    "    return numerator/denominator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2edb2129",
   "metadata": {},
   "outputs": [],
   "source": [
    "def G(lam0, t0, z0, b0):\n",
    "    \n",
    "    t = lam0*t0\n",
    "    z = np.sqrt(lam0)*(z0-theta0)\n",
    "    b = np.sqrt(lam0)*(b0-theta0)\n",
    "    \n",
    "    result = 2*np.exp(b*(z-b))*np.erf(e^(-t)(z-b)/np.sinh(t)**0.5)\n",
    "    \n",
    "    return G(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad1490ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_LCA_full_simp_bound(parameters, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        dI = parameters[i,0]\n",
    "        a = parameters[i,1]\n",
    "        bk = parameters[i,2]\n",
    "        Ter = parameters[i,3]\n",
    "        #Ter = 0.4\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y = a/2 # initial location\n",
    "\n",
    "        while((y<a) and (y>0)):\n",
    "            dW = np.random.normal()*np.sqrt(dt)\n",
    "            dy = (dI - bk*y)*dt + c*dW\n",
    "            y = y + dy\n",
    "\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y = a+0.01\n",
    "                else:\n",
    "                    y = -0.01\n",
    "                break\n",
    "            \n",
    "        RT = n_iter*dt + Ter\n",
    "        #print(y)\n",
    "        if y > a:\n",
    "            R = 0\n",
    "        if y < 0:\n",
    "            R = 1     # different from IDM\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "      \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "            \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0405ac95",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_LCA_full_simp(parameters, dt = 0.001, c = 1):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        dI = parameters[i,0]\n",
    "        a = parameters[i,1]\n",
    "        bk = parameters[i,2]\n",
    "        Ter = parameters[i,3]\n",
    "        #Ter = 0.4\n",
    "        \n",
    "        n_iter = 1\n",
    "        #location = init_loc\n",
    "        y = 0 # initial location\n",
    "\n",
    "        while((y<a) and (y>-a)):\n",
    "            dW = np.random.normal()*np.sqrt(dt)\n",
    "            dy = (dI - bk*y)*dt + c*dW\n",
    "            y = y + dy\n",
    "\n",
    "            #location = c(location, y)\n",
    "            n_iter = n_iter + 1\n",
    "            \n",
    "            if n_iter >=3000:\n",
    "                if np.random.random()>0.5:\n",
    "                    y = a+0.01\n",
    "                else:\n",
    "                    y = -a-0.01\n",
    "                break\n",
    "            \n",
    "        RT = n_iter*dt + Ter\n",
    "        #print(y)\n",
    "        if y > a:\n",
    "            R = 0\n",
    "        if y <-a:\n",
    "            R = 1     # different from IDM\n",
    "        \n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "        \n",
    "      \n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "            \n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c47c24d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_LCA1(parameters, ter = 0.5):\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        v = parameters[i,0]\n",
    "        a = parameters[i,1]\n",
    "        leakage = parameters[i,2]\n",
    "        \n",
    "        model = Model(name='Simple model',\n",
    "              drift=DriftLinear(drift=v, t = 0, x=leakage),\n",
    "              bound=BoundConstant(B=a),\n",
    "              overlay=OverlayNonDecision(nondectime=ter),\n",
    "              noise=NoiseConstant(noise=1),\n",
    "              dx=.001, dt=.001, T_dur=5)\n",
    "        sol = model.solve()\n",
    "\n",
    "        # generate pseudo-random data\n",
    "        samp = sol.resample(1)\n",
    "\n",
    "        if len(samp.corr) == 0:\n",
    "            RT = samp.err[0]\n",
    "            R = 1\n",
    "        if len(samp.err) == 0:\n",
    "            RT = samp.corr[0]\n",
    "            R = 0\n",
    "\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "\n",
    "\n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adc602d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sim_fit_plot_IDM_LCA(Cs, hs, Excs, Inhs, par_fitted, sim_num = 30):\n",
    "\n",
    "    fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8), (ax9, ax10, ax11, ax12)) = plt.subplots(3, 4, figsize=(10,10))\n",
    "\n",
    "    ax1.scatter(Cs[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "    ax2.scatter(hs[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "    ax3.scatter(Excs[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "    ax4.scatter(Inhs[0:sim_num], par_fitted[0:sim_num,0], s = 2)\n",
    "\n",
    "    ax5.scatter(Cs[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "    ax6.scatter(hs[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "    ax7.scatter(Excs[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "    ax8.scatter(Inhs[0:sim_num], par_fitted[0:sim_num,1], s = 2)\n",
    "\n",
    "    ax9.scatter(Cs[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "    ax10.scatter(hs[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "    ax11.scatter(Excs[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "    ax12.scatter(Inhs[0:sim_num], par_fitted[0:sim_num,2], s = 2)\n",
    "    \n",
    "\n",
    "\n",
    "    ax1.set_title(\"C\")\n",
    "    ax2.set_title(\"h\")\n",
    "    ax3.set_title(\"Exc\")\n",
    "    ax4.set_title(\"Inh\")\n",
    "\n",
    "    ax1.set_ylabel(\"I\", rotation=0)\n",
    "    ax5.set_ylabel(\"A\", rotation=0)\n",
    "    ax9.set_ylabel(\"Inh\", rotation=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06c86f18",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sim_fit_plot_5(parameters, par_fitted, sim_num = 30):\n",
    "\n",
    "    fig, axes = plt.subplots(6, 6, figsize=(10,10))\n",
    "    for i in range(6):\n",
    "        for j in range(6):\n",
    "            axes[i,j].scatter(parameters[0:sim_num,i], par_fitted[0:sim_num,j], s = 4)\n",
    "\n",
    "            axes[0,0].set_ylabel(\"I1\")\n",
    "            axes[1,0].set_ylabel(\"I2\")\n",
    "            axes[2,0].set_ylabel(\"a\")\n",
    "            axes[3,0].set_ylabel(\"inh\")\n",
    "            axes[4,0].set_ylabel(\"k\")\n",
    "            axes[5,0].set_ylabel(\"Ter\")\n",
    "\n",
    "            axes[0,0].set_title(\"I1\")\n",
    "            axes[0,1].set_title(\"I2\")\n",
    "            axes[0,2].set_title(\"a\")\n",
    "            axes[0,3].set_title(\"inh\")\n",
    "            axes[0,4].set_title(\"k\")\n",
    "            axes[0,5].set_title(\"Ter\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fb377cd",
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06222a6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def DDM_IDM(prior_DDM, prior_IDM, c = 1, n_trials = 250, sim_num = 30):\n",
    "    param_DDM = prior_DDM.sample((sim_num,))\n",
    "    par_fitted = np.zeros((sim_num,param_IDM.shape[1]))\n",
    "    for i in range(sim_num):\n",
    "        print(i)\n",
    "        DDM_o = simul_DDM(np.tile(param_DDM[i,:], (n_trials,1)), c = c) # simulate IDM data 0.1**0.5\n",
    "\n",
    "        mcmc_parameters = dict(num_chains=50,thin=10,warmup_steps=50,init_strategy=\"proposal\")\n",
    "        mnle_posterior = trainer.build_posterior(mcmc_method=\"slice_np_vectorized\", \n",
    "                                                 mcmc_parameters=mcmc_parameters, \n",
    "                                                 prior = prior_IDM)\n",
    "\n",
    "        mnle_posterior.set_default_x(DDM_o)\n",
    "        mnle_map = mnle_posterior.map()\n",
    "\n",
    "        par_fitted[i,:] = np.array(mnle_map)\n",
    "    return param_DDM, par_fitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9112453a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def IDM_IDM_EM6(prior_IDM, n_trials = 200, sim_num = 30):\n",
    "    param_IDM = prior_IDM.sample((sim_num,))\n",
    "    par_fitted = np.zeros((sim_num,param_IDM.shape[1]))\n",
    "    for i in range(sim_num):\n",
    "        print(i)\n",
    "        IDM_o = simul_IDM_EM6(np.tile(param_IDM[i,:], (n_trials,1))) # \n",
    "\n",
    "        mcmc_parameters = dict(num_chains=50,thin=10,warmup_steps=50,init_strategy=\"proposal\")\n",
    "        mnle_posterior = trainer.build_posterior(mcmc_method=\"slice_np_vectorized\", \n",
    "                                                 mcmc_parameters=mcmc_parameters, \n",
    "                                                 prior = prior_IDM)\n",
    "\n",
    "        mnle_posterior.set_default_x(IDM_o)\n",
    "        mnle_map = mnle_posterior.map()\n",
    "\n",
    "        par_fitted[i,:] = np.array(mnle_map)\n",
    "    return param_IDM, par_fitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "738532c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def IDM_IDM_EM5_Exc(prior_IDM, n_trials = 200, sim_num = 30):\n",
    "    param_IDM = prior_IDM.sample((sim_num,))\n",
    "    par_fitted = np.zeros((sim_num,param_IDM.shape[1]))\n",
    "    for i in range(sim_num):\n",
    "        print(i)\n",
    "        IDM_o = simul_IDM_EM5_Exc(np.tile(param_IDM[i,:], (n_trials,1))) # \n",
    "\n",
    "        mcmc_parameters = dict(num_chains=50,thin=10,warmup_steps=50,init_strategy=\"proposal\")\n",
    "        mnle_posterior = trainer.build_posterior(mcmc_method=\"slice_np_vectorized\", \n",
    "                                                 mcmc_parameters=mcmc_parameters, \n",
    "                                                 prior = prior_IDM)\n",
    "\n",
    "        mnle_posterior.set_default_x(IDM_o)\n",
    "        mnle_map = mnle_posterior.map()\n",
    "\n",
    "        par_fitted[i,:] = np.array(mnle_map)\n",
    "    return param_IDM, par_fitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "861e5935",
   "metadata": {},
   "outputs": [],
   "source": [
    "def IDM_IDM_EM5_Exc(prior_IDM, n_trials = 200, sim_num = 30):\n",
    "    param_IDM = prior_IDM.sample((sim_num,))\n",
    "    par_fitted = np.zeros((sim_num,param_IDM.shape[1]))\n",
    "    for i in range(sim_num):\n",
    "        print(i)\n",
    "        IDM_o = simul_IDM_EM5_Exc(np.tile(param_IDM[i,:], (n_trials,1))) # \n",
    "\n",
    "        mcmc_parameters = dict(num_chains=50,thin=10,warmup_steps=50,init_strategy=\"proposal\")\n",
    "        mnle_posterior = trainer.build_posterior(mcmc_method=\"slice_np_vectorized\", \n",
    "                                                 mcmc_parameters=mcmc_parameters, \n",
    "                                                 prior = prior_IDM)\n",
    "\n",
    "        mnle_posterior.set_default_x(IDM_o)\n",
    "        mnle_map = mnle_posterior.map()\n",
    "\n",
    "        par_fitted[i,:] = np.array(mnle_map)\n",
    "    return param_IDM, par_fitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30637137",
   "metadata": {},
   "outputs": [],
   "source": [
    "def IDM_IDM_EM5_Inh(prior_IDM, n_trials = 200, sim_num = 30):\n",
    "    param_IDM = prior_IDM.sample((sim_num,))\n",
    "    par_fitted = np.zeros((sim_num,param_IDM.shape[1]))\n",
    "    for i in range(sim_num):\n",
    "        print(i)\n",
    "        IDM_o = simul_IDM_EM5_Inh(np.tile(param_IDM[i,:], (n_trials,1))) # \n",
    "\n",
    "        mcmc_parameters = dict(num_chains=50,thin=10,warmup_steps=50,init_strategy=\"proposal\")\n",
    "        mnle_posterior = trainer.build_posterior(mcmc_method=\"slice_np_vectorized\", \n",
    "                                                 mcmc_parameters=mcmc_parameters, \n",
    "                                                 prior = prior_IDM)\n",
    "\n",
    "        mnle_posterior.set_default_x(IDM_o)\n",
    "        mnle_map = mnle_posterior.map()\n",
    "\n",
    "        par_fitted[i,:] = np.array(mnle_map)\n",
    "    return param_IDM, par_fitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bff914a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def IDM_IDM_EM2(prior_IDM, n_trials = 200, sim_num = 30):\n",
    "    param_IDM = prior_IDM.sample((sim_num,))\n",
    "    par_fitted = np.zeros((sim_num,param_IDM.shape[1]))\n",
    "    for i in range(sim_num):\n",
    "        print(i)\n",
    "        IDM_o = simul_IDM_Inh_Exc(np.tile(param_IDM[i,:], (n_trials,1))) # \n",
    "\n",
    "        #mcmc_parameters = dict(num_chains=50,thin=10,warmup_steps=50,init_strategy=\"proposal\")\n",
    "        mnle_posterior = trainer.build_posterior(prior = prior_IDM)\n",
    "\n",
    "        #mcmc_method=\"slice_np_vectorized\", \n",
    "        #mcmc_parameters=mcmc_parameters, \n",
    "        \n",
    "        mnle_posterior.set_default_x(IDM_o)\n",
    "        mnle_map = mnle_posterior.map()\n",
    "        print(param_IDM[i,:])\n",
    "        print(mnle_map)\n",
    "        par_fitted[i,:] = np.array(mnle_map)\n",
    "    return param_IDM, par_fitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a6482f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def LCA_IDM_EM6(prior_LCA, prior_IDM, n_trials = 200, sim_num = 30):\n",
    "    param_LCA = prior_LCA.sample((sim_num,))\n",
    "    param_IDM = prior_IDM.sample((sim_num,))\n",
    "\n",
    "    par_fitted = np.zeros((sim_num,param_IDM.shape[1]))\n",
    "    for i in range(sim_num):\n",
    "        print(i)\n",
    "        LCA_o = simul_LCA_full(np.tile(param_LCA[i,:], (n_trials,1))) # simulate IDM data 0.1**0.5\n",
    "\n",
    "        mcmc_parameters = dict(num_chains=50,thin=10,warmup_steps=50,init_strategy=\"proposal\")\n",
    "        mnle_posterior = trainer.build_posterior(mcmc_method=\"slice_np_vectorized\", \n",
    "                                                 mcmc_parameters=mcmc_parameters, \n",
    "                                                 prior = prior_IDM)\n",
    "\n",
    "        mnle_posterior.set_default_x(LCA_o)\n",
    "        mnle_map = mnle_posterior.map()\n",
    "\n",
    "        par_fitted[i,:] = np.array(mnle_map)\n",
    "    return param_IDM, par_fitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "04102b3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#x = [8374, 8209, 8534, 8374, 8272, 8471, 8532, 8442, 8417, 8486, 8496, 8239, 8299, 8441, 8364, 8463]\n",
    "#y = [8296, 8284, 8530, 8471, 8415, 8531, 8527, 8599, 8524, 8246, 8315, 8574, 8285, 8279, 8244, 8258]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae4f726f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simul_IDM_EM4_D_bs700(parameters, N = 2000, Bns = 2000, Bs = 700, beta = 1/24,\n",
    "                  theta = 51450, dt = 0.001, init_loc = (0.3,0.3)):\n",
    "\n",
    "    #C = 0.1, h = 0.4, Ter = 0.3\n",
    "    n_sim = np.shape(parameters)[0]\n",
    "    RTs, Rs = [], []\n",
    "    for i in range(n_sim):\n",
    "        C = parameters[i,0]\n",
    "        h = parameters[i,1]\n",
    "        Ter = parameters[i,2]\n",
    "        D = parameters[i,3]\n",
    "        W_pos = 52500\n",
    "        W_neg = 8400\n",
    "\n",
    "        B1, B2 = Bs*(1+C) + Bns, Bs*(1-C) + Bns\n",
    "        y1, y2 = init_loc[0], init_loc[1]\n",
    "        n_iter = 1\n",
    "        loc_y1, loc_y2 = [], []\n",
    "\n",
    "        while get_state(y1=y1, y2=y2, h=h)==100:\n",
    "            dW1 = np.random.normal()*np.sqrt(dt)\n",
    "            dW2 = np.random.normal()*np.sqrt(dt)\n",
    "            dy1 = -beta*D*dFdy1(y1,y2, W_pos, W_neg, B1, theta, beta, N)*dt + np.sqrt(2*D)*dW1\n",
    "            dy2 = -beta*D*dFdy2(y1,y2, W_pos, W_neg, B2, theta, beta, N)*dt + np.sqrt(2*D)*dW2\n",
    "            y1 = clamp(y1 + dy1) # clamp to [0,1]\n",
    "            y2 = clamp(y2 + dy2)\n",
    "            loc_y1.append(y1)\n",
    "            loc_y2.append(y2)\n",
    "            n_iter = n_iter + 1\n",
    "\n",
    "            if n_iter >=3000:\n",
    "                break\n",
    "\n",
    "        RT = n_iter*dt + Ter\n",
    "        R = get_state(y1 = y1, y2 = y2, h = h)   # 1 for option 1, 0 for option 2, 100 for undecided\n",
    "        RTs.append(RT)\n",
    "        Rs.append(R)\n",
    "\n",
    "\n",
    "        if i%1000 ==0:\n",
    "            print(\"sim\"+str(i))\n",
    "\n",
    "    RTs = torch.reshape(torch.Tensor(RTs), [n_sim,1])\n",
    "    Rs = torch.reshape(torch.Tensor(Rs), [n_sim,1])\n",
    "    return torch.cat((RTs,Rs),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08d33915",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b77303d8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5691e6d0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aca953d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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
}