baseline_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "udPUmw0nyW0x"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pulse2percept.datasets import load_perezfornos2012\n",
    "import time\n",
    "from pulse2percept.models import TemporalModel\n",
    "from pulse2percept.stimuli.pulse_trains import BiphasicPulseTrain\n",
    "from scipy.optimize import minimize, differential_evolution\n",
    "from datetime import datetime\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = load_perezfornos2012()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(data['time_series'][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0zfHyux0WmRg",
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fixed parameters - condition dependent\n",
    "class Args():\n",
    "    EAR = 20 # electrode activation rate \n",
    "    exper_dur = 30 # experiment duration\n",
    "    stim_dur = 10 # stimulus duration\n",
    "    interpolation_factor = 2 # number of frames sampled every second\n",
    "    \n",
    "args = Args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rUbabzToSGf2"
   },
   "outputs": [],
   "source": [
    "# fit parameters\n",
    "PER=2 # persistence duration (usually between 0 and 8 seconds)\n",
    "PFO=10  # perceptual fading onset (usually between 0 and 10 seconds)\n",
    "PFD=10 # perceptual fading duration (usually between 0.5 and 60 seconds)\n",
    "brightness_scale = 10\n",
    "\n",
    "def brightness_equ(PER, PFO, PFD, brightness_scale):\n",
    "    # fixed parameters\n",
    "    SigmaX=.4 # phosphene dimension in x axis\n",
    "    SigmaY=.4 # phosphene dimension in y axis\n",
    "    # X0=24  # distance between two adjecent phosphenes in x axis\n",
    "    # Y0=24  # distance between two adjecent phosphenes in y axis\n",
    "    ElectrodeArraySize=1 # number of electrodes (60(val=1), 256(2), 961(3), or 2916(4) electrodes)\n",
    "    ElectrodeArrayStructure=1 # electrode array structure (square(val = 1) or hexagonal(2)).\n",
    "\n",
    "    va = 0\n",
    "    vb = 0.333\n",
    "    vc = 0.667\n",
    "    vd = 1\n",
    "    # MMr = 6  # number of rows\n",
    "    # MMc = 10  # number of columns;\n",
    "    # C = 240  # number of pixel columns\n",
    "    # R = 144 # number of pixel rows, choose: phosphene size 0.4; X0 and Y0 = 24\n",
    "\n",
    "    SigmaX=.4 # phosphene dimension in x axis\n",
    "    SigmaY=.4 # phosphene dimension in y axis\n",
    "    MMr = 3  # number of rows\n",
    "    MMc = 3  # number of columns;\n",
    "    X0 = 24  # distance between two adjecent phosphenes in x axis\n",
    "    Y0 = 24  # distance between two adjecent phosphenes in y axis\n",
    "    C = MMc*Y0  # number of pixel columns\n",
    "    R = MMr*X0  # number of pixel rows, choose: phosphene size 0.4; X0 and Y0 = 24\n",
    "\n",
    "    DFR = 30 # display frame rate (usually 30 or 60 FPS)\n",
    "\n",
    "    v2 = 0\n",
    "    v3 = 0\n",
    "    count = 0\n",
    "\n",
    "    EAR = args.EAR  # electrode activation rate (usually 6 or 10 Hz)\n",
    "    exper_dur = args.exper_dur  # experiment duration\n",
    "    stim_dur = args.stim_dur  # stimulus duration\n",
    "\n",
    "    PER_Factor = 3.9/(PER*DFR) # 3.9 was set for 2 precent threshold, i.e. the phosphene falls to 2 precent of its original intensity;\n",
    "    PF_Factor = PFO*EAR\n",
    "    PFD_Factor = 3.9/(PFD*DFR)\n",
    "\n",
    "    # Creating meshgrid\n",
    "    x = np.arange(0, C) \n",
    "    y = np.arange(0, R)  \n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    # Initialization\n",
    "    z = np.zeros((R, C))\n",
    "    Frame = np.zeros((R, C))\n",
    "    FormerFrame = np.zeros((R, C))\n",
    "    PFC = np.zeros((MMr, MMc))\n",
    "    Theta = np.zeros((MMr, MMc))\n",
    "    a = np.zeros((MMr, MMc))\n",
    "    b = np.zeros((MMr, MMc))\n",
    "    c = np.zeros((MMr, MMc))\n",
    "    sigma_x = np.zeros((MMr, MMc))\n",
    "    sigma_y = np.zeros((MMr, MMc))\n",
    "\n",
    "    count = 0\n",
    "    for s in range(MMr):\n",
    "        for t in range(MMc):\n",
    "            Theta[s, t] = 0\n",
    "            sigma_x[s, t] = SigmaX\n",
    "            sigma_y[s, t] = SigmaY\n",
    "            a[s, t] = (np.cos(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_x[s, t]**2) + \\\n",
    "                      (np.sin(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_y[s, t]**2)\n",
    "            b[s, t] = -(np.sin(np.deg2rad(2 * Theta[s, t]))) / (4 * sigma_x[s, t]**2) + \\\n",
    "                      (np.sin(np.deg2rad(2 * Theta[s, t]))) / (4 * sigma_y[s, t]**2)\n",
    "            c[s, t] = (np.sin(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_x[s, t]**2) + \\\n",
    "                      (np.cos(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_y[s, t]**2)\n",
    "\n",
    "    I_frame = [1]*(args.stim_dur*args.EAR)*MMr*MMc  # the brightness of the video is always 1\n",
    "    # video frame rate is the same as the electrode firing rate\n",
    "    \n",
    "    I_frame_0 = [0]*(args.exper_dur*args.EAR-args.stim_dur*args.EAR)*MMr*MMc  # no visual input after stim_dur\n",
    "    I_frame = np.append(np.array(I_frame), np.array(I_frame_0))\n",
    "    I_frame = np.array(I_frame).reshape(args.exper_dur*args.EAR, MMr, MMc)\n",
    "\n",
    "    results = []\n",
    "\n",
    "    for frame in range(exper_dur*args.EAR):\n",
    "        I = I_frame[frame]\n",
    "        # I = I / I.max()\n",
    "        MeanMatrix = np.zeros((MMr, MMc))\n",
    "        s = 0\n",
    "        t = 0\n",
    "        v = 0\n",
    "\n",
    "        start = int(np.floor((I.shape[0] % MMr) / 2) + np.floor(I.shape[0] / MMr) * MMr) - 1\n",
    "        stop = int(1 + np.floor((I.shape[0] % MMr) / 2)) - 1\n",
    "        step = -int(np.floor(I.shape[0] / MMr))\n",
    "        for m in range(start, stop, step):\n",
    "            for n in range(0, I.shape[1]//MMc*MMc - 1, I.shape[1] // MMc):\n",
    "                for j in range(m, m - (I.shape[0] // MMr), -1):\n",
    "                    # for k in range(n, n + (I.shape[1] // MMc)-2):\n",
    "                    for k in range(n, n - (I.shape[1] // MMc), -1):\n",
    "                        v += I[j, k]\n",
    "                v /= ((I.shape[0] // MMr) * (I.shape[1] // MMc))\n",
    "\n",
    "                if 0 <= v <= 0.25:\n",
    "                    v = va\n",
    "                elif 0.25 < v <= 0.5:\n",
    "                    v = vb\n",
    "                elif 0.5 < v <= 0.75:\n",
    "                    v = vc\n",
    "                else:\n",
    "                    v = vd\n",
    "                MeanMatrix[s, t] = v\n",
    "                v = 0\n",
    "                t += 1\n",
    "            t = 0\n",
    "            s += 1\n",
    "\n",
    "        getting_all_subframe = []\n",
    "\n",
    "        SubFrame_range = int(DFR) // int(EAR) + 1\n",
    "        if SubFrame_range <= 2: SubFrame_range = 2\n",
    "        for SubFrame in range(1, SubFrame_range):\n",
    "            for s in range(MMr):\n",
    "                for t in range(MMc):\n",
    "                    adjusted_t = t + 1\n",
    "                    adjusted_s = s + 1\n",
    "                    z += MeanMatrix[s, t] * np.exp(\n",
    "                        # -0.015 * (1 / (3 * MeanMatrix[s, t] + 1)) *\n",
    "                        -0.015 * (1 / (3 * MeanMatrix[s, t] + 1)) *\n",
    "                        # (a[s, t] * (X - X0 * t - v2 + X0 // 2) ** 2 +\n",
    "                        (a[s, t] * (X - X0 * adjusted_t - v2 + X0 // 2) ** 2 +\n",
    "                        # 2 * b[s, t] * (X - X0 * t - v2 + X0 // 2) * (Y - Y0 * s + Y0 // 2) +\n",
    "                        2 * b[s, t] * (X - X0 * adjusted_t - v2 + X0 // 2) * (Y - Y0 * adjusted_s + Y0 // 2) +\n",
    "                        # c[s, t] * (Y - Y0 * s + Y0 // 2) ** 2))\n",
    "                        c[s, t] * (Y - Y0 * adjusted_s + Y0 // 2) ** 2))\n",
    "\n",
    "                v2 -= (X0 / 2) * v3\n",
    "                v3 *= -1\n",
    "\n",
    "            for s in range(R, 0, -X0):\n",
    "                for t in range(0, C, Y0):\n",
    "                    if SubFrame == 1:\n",
    "                        condition1 = (np.mean(z[s-X0:s-1, t-1:t+Y0-2]) == 0 and np.mean(FormerFrame[s-X0:s-1, t-1:t+Y0-2]) == 0)\n",
    "                        condition2 = (np.mean(z[s-X0:s-1, t-1:t+Y0-2]) < np.mean(FormerFrame[s-X0:s-1, t:t+Y0-2]))\n",
    "                        condition3 = (MeanMatrix[int(s/X0)-1, t // Y0] == 0)\n",
    "                        condition4 = (PFC[int(s/X0)-1, t // Y0] < PF_Factor)\n",
    "                        if condition1 or condition2 or condition3 and condition4:\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PER_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                        elif np.mean(z[s-X0:s-1, t-1:t+Y0-2]) >= np.mean(FormerFrame[s-X0:s-1, t-1:t+Y0-2]) and condition4:\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = z[s-X0:s-1, t-1:t+Y0-2]\n",
    "                            PFC[s//X0-1, t//Y0] += 1\n",
    "                        elif PFC[s//X0-1, t//Y0] >= PF_Factor:\n",
    "                            # Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-3 * PER_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PFD_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                    else:\n",
    "                        Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PER_Factor) * Frame[s-X0:s-1, t-1:t+Y0-2]\n",
    "            z = np.zeros((R,C))\n",
    "            count = count + 1\n",
    "\n",
    "            FormerFrame = Frame\n",
    "            getting_all_subframe.append(Frame.max())\n",
    "        results.append(max(getting_all_subframe))\n",
    "\n",
    "    output =  [x * brightness_scale for x in results]\n",
    "\n",
    "    output_ = []\n",
    "    combined_frame = args.EAR/args.interpolation_factor\n",
    "    for i in range(args.exper_dur*args.interpolation_factor):\n",
    "        output_.append(max(output[int(i*combined_frame):int((i+1)*combined_frame)]))\n",
    "    return output_\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "XK-Hbc6URsmj",
    "outputId": "9fe7ccf9-3df6-43e6-d47a-bdf2ad08cb8f",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "def objective_function(params, gt_curves_train):\n",
    "    # persistence time, fading time onset, fading time duration, brightness scale\n",
    "    PER, PFO, PFD, brightness_scale = params\n",
    "    \n",
    "    total_mse = 0\n",
    "    y_predicted = brightness_equ(PER, PFO, PFD, brightness_scale)\n",
    "    for y in gt_curves_train: total_mse += np.mean((np.array(y_predicted) - y) ** 2)\n",
    "        \n",
    "    print(total_mse/len(gt_curves_train))\n",
    "    \n",
    "    return total_mse/len(gt_curves_train)\n",
    "\n",
    "\n",
    "for cond in range(5):\n",
    "    for subj in range(9):\n",
    "        current_time = datetime.now()\n",
    "        formatted_time = current_time.strftime(\"%Y-%m-%d %H:%M:%S\")\n",
    "\n",
    "        i = cond*9 + subj\n",
    "\n",
    "        if i < -1 : pass\n",
    "        else:\n",
    "            idx_lst = list(range(cond*9 , (cond+1)*9))\n",
    "            idx_lst.remove(i)\n",
    "            \n",
    "            if i <= 8: \n",
    "                args.EAR = 20\n",
    "                args.exper_dur = 30\n",
    "                args.stim_dur = 10\n",
    "            elif i <= 17: \n",
    "                args.EAR = 5\n",
    "                args.exper_dur = 30\n",
    "                args.stim_dur = 10\n",
    "            elif i <= 26:\n",
    "                args.EAR = 60\n",
    "                args.exper_dur = 30\n",
    "                args.stim_dur = 10\n",
    "                \n",
    "            elif i<=35: \n",
    "                args.EAR = 20\n",
    "                args.exper_dur = 10\n",
    "                args.stim_dur = 1\n",
    "            else: \n",
    "                args.EAR = 20\n",
    "                args.exper_dur = 70\n",
    "                args.stim_dur = 60\n",
    "                \n",
    "            # curve = np.nan_to_num(data['time_series'][i], 0)[4:args.exper_dur+4]\n",
    "            # ground_truth = curve  \n",
    "            \n",
    "            gt_curves_train = [np.nan_to_num(data['time_series'][idx], 0)[4:args.exper_dur*args.interpolation_factor+4] for idx in idx_lst]\n",
    "            gt_curves_val = np.nan_to_num(data['time_series'][i], 0)[4:args.exper_dur*args.interpolation_factor+4]\n",
    "    \n",
    "            # Initial guesses for parameters\n",
    "            initial_params = [2, 10, 10, 10]\n",
    "            # Run the optimization, passing the ground truth as an argument\n",
    "            result = minimize(objective_function, initial_params, args=(gt_curves_train,), method='Nelder-Mead')\n",
    "            # Extract optimized parameters\n",
    "            optimized_params = result.x\n",
    "            print(\"Optimized Parameters:\", optimized_params)\n",
    "            \n",
    "            file_path = '/hdd/yuchen/baseline_test.txt'  \n",
    "            line_to_add = '' \n",
    "            line_to_add += str(i) + ': '\n",
    "            line_to_add += str(list(optimized_params)) + '\\n' + formatted_time + '\\n'\n",
    "            with open(file_path, 'a') as file:\n",
    "                file.write(line_to_add)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## train per panel 0.25 interp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fixed parameters - condition dependent\n",
    "class Args():\n",
    "    EAR = 20 # electrode activation rate \n",
    "    exper_dur = 30 # experiment duration\n",
    "    stim_dur = 10 # stimulus duration\n",
    "    interpolation_factor = 4 # number of frames sampled every second\n",
    "    \n",
    "args = Args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fit parameters\n",
    "PER=2 # persistence duration (usually between 0 and 8 seconds)\n",
    "PFO=10  # perceptual fading onset (usually between 0 and 10 seconds)\n",
    "PFD=10 # perceptual fading duration (usually between 0.5 and 60 seconds)\n",
    "brightness_scale = 10\n",
    "\n",
    "def brightness_equ(PER, PFO, PFD, brightness_scale):\n",
    "    SigmaX=.4 # phosphene dimension in x axis\n",
    "    SigmaY=.4 # phosphene dimension in y axis\n",
    "    ElectrodeArraySize=1 # number of electrodes (60(val=1), 256(2), 961(3), or 2916(4) electrodes)\n",
    "    ElectrodeArrayStructure=1 # electrode array structure (square(val = 1) or hexagonal(2)).\n",
    "\n",
    "    va = 0\n",
    "    vb = 0.333\n",
    "    vc = 0.667\n",
    "    vd = 1\n",
    "\n",
    "    SigmaX=.4 # phosphene dimension in x axis\n",
    "    SigmaY=.4 # phosphene dimension in y axis\n",
    "    MMr = 2  # number of rows\n",
    "    MMc = 3  # number of columns;\n",
    "    X0 = 20  # distance between two adjecent phosphenes in x axis\n",
    "    Y0 = 20  # distance between two adjecent phosphenes in y axis\n",
    "    C = MMc*Y0  # number of pixel columns\n",
    "    R = MMr*X0  # number of pixel rows, choose: phosphene size 0.4; X0 and Y0 = 24\n",
    "\n",
    "    DFR = 30 # display frame rate (usually 30 or 60 FPS)\n",
    "\n",
    "    v2 = 0\n",
    "    v3 = 0\n",
    "    count = 0\n",
    "\n",
    "    EAR = args.EAR  # electrode activation rate (usually 6 or 10 Hz)\n",
    "    exper_dur = args.exper_dur  # experiment duration\n",
    "    stim_dur = args.stim_dur  # stimulus duration\n",
    "\n",
    "    PER_Factor = 3.9/(PER*DFR) # 3.9 was set for 2 precent threshold, i.e. the phosphene falls to 2 precent of its original intensity;\n",
    "    PF_Factor = PFO*EAR\n",
    "    PFD_Factor = 3.9/(PFD*DFR)\n",
    "\n",
    "    # Creating meshgrid\n",
    "    x = np.arange(0, C) \n",
    "    y = np.arange(0, R)  \n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    # Initialization\n",
    "    z = np.zeros((R, C))\n",
    "    Frame = np.zeros((R, C))\n",
    "    FormerFrame = np.zeros((R, C))\n",
    "    PFC = np.zeros((MMr, MMc))\n",
    "    Theta = np.zeros((MMr, MMc))\n",
    "    a = np.zeros((MMr, MMc))\n",
    "    b = np.zeros((MMr, MMc))\n",
    "    c = np.zeros((MMr, MMc))\n",
    "    sigma_x = np.zeros((MMr, MMc))\n",
    "    sigma_y = np.zeros((MMr, MMc))\n",
    "\n",
    "    count = 0\n",
    "    for s in range(MMr):\n",
    "        for t in range(MMc):\n",
    "            Theta[s, t] = 0\n",
    "            sigma_x[s, t] = SigmaX\n",
    "            sigma_y[s, t] = SigmaY\n",
    "            a[s, t] = (np.cos(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_x[s, t]**2) + \\\n",
    "                      (np.sin(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_y[s, t]**2)\n",
    "            b[s, t] = -(np.sin(np.deg2rad(2 * Theta[s, t]))) / (4 * sigma_x[s, t]**2) + \\\n",
    "                      (np.sin(np.deg2rad(2 * Theta[s, t]))) / (4 * sigma_y[s, t]**2)\n",
    "            c[s, t] = (np.sin(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_x[s, t]**2) + \\\n",
    "                      (np.cos(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_y[s, t]**2)\n",
    "\n",
    "    I_frame = [1]*(args.stim_dur*args.EAR)*MMr*MMc  # the brightness of the video is always 1\n",
    "    # video frame rate is the same as the electrode firing rate\n",
    "    \n",
    "    I_frame_0 = [0]*(args.exper_dur*args.EAR-args.stim_dur*args.EAR)*MMr*MMc  # no visual input after stim_dur\n",
    "    I_frame = np.append(np.array(I_frame), np.array(I_frame_0))\n",
    "    I_frame = np.array(I_frame).reshape(args.exper_dur*args.EAR, MMr, MMc)\n",
    "\n",
    "    results = []\n",
    "\n",
    "    for frame in range(exper_dur*args.EAR):\n",
    "        I = I_frame[frame]\n",
    "        # I = I / I.max()\n",
    "        MeanMatrix = np.zeros((MMr, MMc))\n",
    "        s = 0\n",
    "        t = 0\n",
    "        v = 0\n",
    "\n",
    "        start = int(np.floor((I.shape[0] % MMr) / 2) + np.floor(I.shape[0] / MMr) * MMr) - 1\n",
    "        stop = int(1 + np.floor((I.shape[0] % MMr) / 2)) - 1\n",
    "        step = -int(np.floor(I.shape[0] / MMr))\n",
    "        for m in range(start, stop, step):\n",
    "            for n in range(0, I.shape[1]//MMc*MMc - 1, I.shape[1] // MMc):\n",
    "                for j in range(m, m - (I.shape[0] // MMr), -1):\n",
    "                    # for k in range(n, n + (I.shape[1] // MMc)-2):\n",
    "                    for k in range(n, n - (I.shape[1] // MMc), -1):\n",
    "                        v += I[j, k]\n",
    "                v /= ((I.shape[0] // MMr) * (I.shape[1] // MMc))\n",
    "\n",
    "                if 0 <= v <= 0.25:\n",
    "                    v = va\n",
    "                elif 0.25 < v <= 0.5:\n",
    "                    v = vb\n",
    "                elif 0.5 < v <= 0.75:\n",
    "                    v = vc\n",
    "                else:\n",
    "                    v = vd\n",
    "                MeanMatrix[s, t] = v\n",
    "                v = 0\n",
    "                t += 1\n",
    "            t = 0\n",
    "            s += 1\n",
    "\n",
    "        getting_all_subframe = []\n",
    "\n",
    "        SubFrame_range = int(DFR) // int(EAR) + 1\n",
    "        if SubFrame_range <= 2: SubFrame_range = 2\n",
    "        for SubFrame in range(1, SubFrame_range):\n",
    "            for s in range(MMr):\n",
    "                for t in range(MMc):\n",
    "                    adjusted_t = t + 1\n",
    "                    adjusted_s = s + 1\n",
    "                    z += MeanMatrix[s, t] * np.exp(\n",
    "                        # -0.015 * (1 / (3 * MeanMatrix[s, t] + 1)) *\n",
    "                        -0.015 * (1 / (3 * MeanMatrix[s, t] + 1)) *\n",
    "                        # (a[s, t] * (X - X0 * t - v2 + X0 // 2) ** 2 +\n",
    "                        (a[s, t] * (X - X0 * adjusted_t - v2 + X0 // 2) ** 2 +\n",
    "                        # 2 * b[s, t] * (X - X0 * t - v2 + X0 // 2) * (Y - Y0 * s + Y0 // 2) +\n",
    "                        2 * b[s, t] * (X - X0 * adjusted_t - v2 + X0 // 2) * (Y - Y0 * adjusted_s + Y0 // 2) +\n",
    "                        # c[s, t] * (Y - Y0 * s + Y0 // 2) ** 2))\n",
    "                        c[s, t] * (Y - Y0 * adjusted_s + Y0 // 2) ** 2))\n",
    "\n",
    "                v2 -= (X0 / 2) * v3\n",
    "                v3 *= -1\n",
    "\n",
    "            for s in range(R, 0, -X0):\n",
    "                for t in range(0, C, Y0):\n",
    "                    if SubFrame == 1:\n",
    "                        condition1 = (np.mean(z[s-X0:s-1, t-1:t+Y0-2]) == 0 and np.mean(FormerFrame[s-X0:s-1, t-1:t+Y0-2]) == 0)\n",
    "                        condition2 = (np.mean(z[s-X0:s-1, t-1:t+Y0-2]) < np.mean(FormerFrame[s-X0:s-1, t:t+Y0-2]))\n",
    "                        condition3 = (MeanMatrix[int(s/X0)-1, t // Y0] == 0)\n",
    "                        condition4 = (PFC[int(s/X0)-1, t // Y0] < PF_Factor)\n",
    "                        if condition1 or condition2 or condition3 and condition4:\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PER_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                        elif np.mean(z[s-X0:s-1, t-1:t+Y0-2]) >= np.mean(FormerFrame[s-X0:s-1, t-1:t+Y0-2]) and condition4:\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = z[s-X0:s-1, t-1:t+Y0-2]\n",
    "                            PFC[s//X0-1, t//Y0] += 1\n",
    "                        elif PFC[s//X0-1, t//Y0] >= PF_Factor:\n",
    "                            # Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-3 * PER_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PFD_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                    else:\n",
    "                        Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PER_Factor) * Frame[s-X0:s-1, t-1:t+Y0-2]\n",
    "            z = np.zeros((R,C))\n",
    "            count = count + 1\n",
    "\n",
    "            FormerFrame = Frame\n",
    "            getting_all_subframe.append(Frame.max())\n",
    "        results.append(max(getting_all_subframe))\n",
    "\n",
    "    output =  [x * brightness_scale for x in results]\n",
    "\n",
    "    output_ = []\n",
    "    combined_frame = args.EAR/args.interpolation_factor\n",
    "    for i in range(args.exper_dur*args.interpolation_factor):\n",
    "        output_.append(max(output[int(i*combined_frame):int((i+1)*combined_frame)]))\n",
    "    return output_\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "data = pd.read_csv('/hdd/yuchen/perez-fornos-2012.csv')  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "args.interpolation_factor = 4\n",
    "\n",
    "def objective_function(params, gt_curves_val):\n",
    "    # persistence time, fading time onset, fading time duration, brightness scale\n",
    "    PER, PFO, PFD, brightness_scale = params\n",
    "    \n",
    "    total_mse = 0\n",
    "    y_predicted = brightness_equ(PER, PFO, PFD, brightness_scale)\n",
    "    mse = np.mean((np.array(y_predicted) - gt_curves_val) ** 2)\n",
    "    print(mse)\n",
    "    return mse\n",
    "\n",
    "\n",
    "for cond in range(5):\n",
    "    for subj in range(9):\n",
    "        current_time = datetime.now()\n",
    "        formatted_time = current_time.strftime(\"%Y-%m-%d %H:%M:%S\")\n",
    "\n",
    "        i = cond*9 + subj\n",
    "\n",
    "        if i < -1 : pass\n",
    "        else:\n",
    "            idx_lst = list(range(cond*9 , (cond+1)*9))\n",
    "            idx_lst.remove(i)\n",
    "            \n",
    "            if i <= 8: \n",
    "                args.EAR = 20\n",
    "                args.exper_dur = 30\n",
    "                args.stim_dur = 10\n",
    "            elif i <= 17: \n",
    "                args.EAR = 5\n",
    "                args.exper_dur = 30\n",
    "                args.stim_dur = 10\n",
    "            elif i <= 26:\n",
    "                args.EAR = 60\n",
    "                args.exper_dur = 30\n",
    "                args.stim_dur = 10\n",
    "                \n",
    "            elif i<=35: \n",
    "                args.EAR = 20\n",
    "                args.exper_dur = 10\n",
    "                args.stim_dur = 1\n",
    "            else: \n",
    "                args.EAR = 20\n",
    "                args.exper_dur = 70\n",
    "                args.stim_dur = 60\n",
    "                \n",
    "            # gt_curves_val = np.nan_to_num(data['time_series'][i], 0)[4:args.exper_dur*args.interpolation_factor+4]\n",
    "            gt_curves_val = np.nan_to_num(data.iloc[i, 6:args.exper_dur*args.interpolation_factor+6].to_numpy().astype('float64'), 0)  \n",
    "            \n",
    "    \n",
    "            # Initial guesses for parameters\n",
    "            initial_params = [2, 10, 10, 10]\n",
    "            # Run the optimization, passing the ground truth as an argument\n",
    "            result = minimize(objective_function, initial_params, args=(gt_curves_val,), method='Nelder-Mead')\n",
    "            # Extract optimized parameters\n",
    "            optimized_params = result.x\n",
    "            print(\"Optimized Parameters:\", optimized_params)\n",
    "            \n",
    "            file_path = '/hdd/yuchen/baseline_perpanel-025interp.txt'  \n",
    "            line_to_add = '' \n",
    "            line_to_add += str(i) + ': '\n",
    "            line_to_add += str(list(optimized_params)) + '\\n' + formatted_time + '\\n'\n",
    "            with open(file_path, 'a') as file:\n",
    "                file.write(line_to_add)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## train per-condition 0.25 interp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fixed parameters - condition dependent\n",
    "class Args():\n",
    "    EAR = 20 # electrode activation rate \n",
    "    exper_dur = 30 # experiment duration\n",
    "    stim_dur = 10 # stimulus duration\n",
    "    interpolation_factor = 2 # number of frames sampled every second\n",
    "    \n",
    "args = Args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "data = pd.read_csv('perez-fornos-2012.csv')  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fit parameters\n",
    "PER=2 # persistence duration (usually between 0 and 8 seconds)\n",
    "PFO=10  # perceptual fading onset (usually between 0 and 10 seconds)\n",
    "PFD=10 # perceptual fading duration (usually between 0.5 and 60 seconds)\n",
    "brightness_scale = 10\n",
    "\n",
    "def brightness_equ(PER, PFO, PFD, brightness_scale):\n",
    "    SigmaX=.4 # phosphene dimension in x axis\n",
    "    SigmaY=.4 # phosphene dimension in y axis\n",
    "    ElectrodeArraySize=1 # number of electrodes (60(val=1), 256(2), 961(3), or 2916(4) electrodes)\n",
    "    ElectrodeArrayStructure=1 # electrode array structure (square(val = 1) or hexagonal(2)).\n",
    "\n",
    "    va = 0\n",
    "    vb = 0.333\n",
    "    vc = 0.667\n",
    "    vd = 1\n",
    "\n",
    "    SigmaX=.4 # phosphene dimension in x axis\n",
    "    SigmaY=.4 # phosphene dimension in y axis\n",
    "    MMr = 2  # number of rows\n",
    "    MMc = 3  # number of columns;\n",
    "    X0 = 20  # distance between two adjecent phosphenes in x axis\n",
    "    Y0 = 20  # distance between two adjecent phosphenes in y axis\n",
    "    C = MMc*Y0  # number of pixel columns\n",
    "    R = MMr*X0  # number of pixel rows, choose: phosphene size 0.4; X0 and Y0 = 24\n",
    "\n",
    "    DFR = 30 # display frame rate (usually 30 or 60 FPS)\n",
    "\n",
    "    v2 = 0\n",
    "    v3 = 0\n",
    "    count = 0\n",
    "\n",
    "    EAR = args.EAR  # electrode activation rate (usually 6 or 10 Hz)\n",
    "    exper_dur = args.exper_dur  # experiment duration\n",
    "    stim_dur = args.stim_dur  # stimulus duration\n",
    "\n",
    "    PER_Factor = 3.9/(PER*DFR) # 3.9 was set for 2 precent threshold, i.e. the phosphene falls to 2 precent of its original intensity;\n",
    "    PF_Factor = PFO*EAR\n",
    "    PFD_Factor = 3.9/(PFD*DFR)\n",
    "\n",
    "    # Creating meshgrid\n",
    "    x = np.arange(0, C) \n",
    "    y = np.arange(0, R)  \n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    # Initialization\n",
    "    z = np.zeros((R, C))\n",
    "    Frame = np.zeros((R, C))\n",
    "    FormerFrame = np.zeros((R, C))\n",
    "    PFC = np.zeros((MMr, MMc))\n",
    "    Theta = np.zeros((MMr, MMc))\n",
    "    a = np.zeros((MMr, MMc))\n",
    "    b = np.zeros((MMr, MMc))\n",
    "    c = np.zeros((MMr, MMc))\n",
    "    sigma_x = np.zeros((MMr, MMc))\n",
    "    sigma_y = np.zeros((MMr, MMc))\n",
    "\n",
    "    count = 0\n",
    "    for s in range(MMr):\n",
    "        for t in range(MMc):\n",
    "            Theta[s, t] = 0\n",
    "            sigma_x[s, t] = SigmaX\n",
    "            sigma_y[s, t] = SigmaY\n",
    "            a[s, t] = (np.cos(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_x[s, t]**2) + \\\n",
    "                      (np.sin(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_y[s, t]**2)\n",
    "            b[s, t] = -(np.sin(np.deg2rad(2 * Theta[s, t]))) / (4 * sigma_x[s, t]**2) + \\\n",
    "                      (np.sin(np.deg2rad(2 * Theta[s, t]))) / (4 * sigma_y[s, t]**2)\n",
    "            c[s, t] = (np.sin(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_x[s, t]**2) + \\\n",
    "                      (np.cos(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_y[s, t]**2)\n",
    "\n",
    "    I_frame = [1]*(args.stim_dur*args.EAR)*MMr*MMc  # the brightness of the video is always 1\n",
    "    # video frame rate is the same as the electrode firing rate\n",
    "    \n",
    "    I_frame_0 = [0]*(args.exper_dur*args.EAR-args.stim_dur*args.EAR)*MMr*MMc  # no visual input after stim_dur\n",
    "    I_frame = np.append(np.array(I_frame), np.array(I_frame_0))\n",
    "    I_frame = np.array(I_frame).reshape(args.exper_dur*args.EAR, MMr, MMc)\n",
    "\n",
    "    results = []\n",
    "\n",
    "    for frame in range(exper_dur*args.EAR):\n",
    "        I = I_frame[frame]\n",
    "        # I = I / I.max()\n",
    "        MeanMatrix = np.zeros((MMr, MMc))\n",
    "        s = 0\n",
    "        t = 0\n",
    "        v = 0\n",
    "\n",
    "        start = int(np.floor((I.shape[0] % MMr) / 2) + np.floor(I.shape[0] / MMr) * MMr) - 1\n",
    "        stop = int(1 + np.floor((I.shape[0] % MMr) / 2)) - 1\n",
    "        step = -int(np.floor(I.shape[0] / MMr))\n",
    "        for m in range(start, stop, step):\n",
    "            for n in range(0, I.shape[1]//MMc*MMc - 1, I.shape[1] // MMc):\n",
    "                for j in range(m, m - (I.shape[0] // MMr), -1):\n",
    "                    # for k in range(n, n + (I.shape[1] // MMc)-2):\n",
    "                    for k in range(n, n - (I.shape[1] // MMc), -1):\n",
    "                        v += I[j, k]\n",
    "                v /= ((I.shape[0] // MMr) * (I.shape[1] // MMc))\n",
    "\n",
    "                if 0 <= v <= 0.25:\n",
    "                    v = va\n",
    "                elif 0.25 < v <= 0.5:\n",
    "                    v = vb\n",
    "                elif 0.5 < v <= 0.75:\n",
    "                    v = vc\n",
    "                else:\n",
    "                    v = vd\n",
    "                MeanMatrix[s, t] = v\n",
    "                v = 0\n",
    "                t += 1\n",
    "            t = 0\n",
    "            s += 1\n",
    "\n",
    "        getting_all_subframe = []\n",
    "\n",
    "        SubFrame_range = int(DFR) // int(EAR) + 1\n",
    "        if SubFrame_range <= 2: SubFrame_range = 2\n",
    "        for SubFrame in range(1, SubFrame_range):\n",
    "            for s in range(MMr):\n",
    "                for t in range(MMc):\n",
    "                    adjusted_t = t + 1\n",
    "                    adjusted_s = s + 1\n",
    "                    z += MeanMatrix[s, t] * np.exp(\n",
    "                        # -0.015 * (1 / (3 * MeanMatrix[s, t] + 1)) *\n",
    "                        -0.015 * (1 / (3 * MeanMatrix[s, t] + 1)) *\n",
    "                        # (a[s, t] * (X - X0 * t - v2 + X0 // 2) ** 2 +\n",
    "                        (a[s, t] * (X - X0 * adjusted_t - v2 + X0 // 2) ** 2 +\n",
    "                        # 2 * b[s, t] * (X - X0 * t - v2 + X0 // 2) * (Y - Y0 * s + Y0 // 2) +\n",
    "                        2 * b[s, t] * (X - X0 * adjusted_t - v2 + X0 // 2) * (Y - Y0 * adjusted_s + Y0 // 2) +\n",
    "                        # c[s, t] * (Y - Y0 * s + Y0 // 2) ** 2))\n",
    "                        c[s, t] * (Y - Y0 * adjusted_s + Y0 // 2) ** 2))\n",
    "\n",
    "                v2 -= (X0 / 2) * v3\n",
    "                v3 *= -1\n",
    "\n",
    "            for s in range(R, 0, -X0):\n",
    "                for t in range(0, C, Y0):\n",
    "                    if SubFrame == 1:\n",
    "                        condition1 = (np.mean(z[s-X0:s-1, t-1:t+Y0-2]) == 0 and np.mean(FormerFrame[s-X0:s-1, t-1:t+Y0-2]) == 0)\n",
    "                        condition2 = (np.mean(z[s-X0:s-1, t-1:t+Y0-2]) < np.mean(FormerFrame[s-X0:s-1, t:t+Y0-2]))\n",
    "                        condition3 = (MeanMatrix[int(s/X0)-1, t // Y0] == 0)\n",
    "                        condition4 = (PFC[int(s/X0)-1, t // Y0] < PF_Factor)\n",
    "                        if condition1 or condition2 or condition3 and condition4:\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PER_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                        elif np.mean(z[s-X0:s-1, t-1:t+Y0-2]) >= np.mean(FormerFrame[s-X0:s-1, t-1:t+Y0-2]) and condition4:\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = z[s-X0:s-1, t-1:t+Y0-2]\n",
    "                            PFC[s//X0-1, t//Y0] += 1\n",
    "                        elif PFC[s//X0-1, t//Y0] >= PF_Factor:\n",
    "                            # Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-3 * PER_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PFD_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                    else:\n",
    "                        Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PER_Factor) * Frame[s-X0:s-1, t-1:t+Y0-2]\n",
    "            z = np.zeros((R,C))\n",
    "            count = count + 1\n",
    "\n",
    "            FormerFrame = Frame\n",
    "            getting_all_subframe.append(Frame.max())\n",
    "        results.append(max(getting_all_subframe))\n",
    "\n",
    "    output =  [x * brightness_scale for x in results]\n",
    "\n",
    "    output_ = []\n",
    "    combined_frame = args.EAR/args.interpolation_factor\n",
    "    for i in range(args.exper_dur*args.interpolation_factor):\n",
    "        output_.append(max(output[int(i*combined_frame):int((i+1)*combined_frame)]))\n",
    "    return output_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stim_info = {}\n",
    "# EAR, exper_dur, stim_dur\n",
    "for i in range(9): stim_info[i] = [20, 30, 10]\n",
    "for i in range(9,18): stim_info[i] = [5, 30, 10]\n",
    "for i in range(18,27): stim_info[i] = [60, 30, 10]\n",
    "for i in range(27,36): stim_info[i] = [20, 10, 1]\n",
    "for i in range(36,45): stim_info[i] = [20, 70, 60]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "args.interpolation_factor = 4\n",
    "\n",
    "def objective_function(params, gt_curves_train, stim_lst):\n",
    "    # persistence time, fading time onset, fading time duration, brightness scale\n",
    "    PER, PFO, PFD, brightness_scale = params\n",
    "    \n",
    "    total_mse = 0\n",
    "    for gt_idx in range(len(gt_curves_train)):\n",
    "        gt_y = gt_curves_train[gt_idx]\n",
    "        stim = stim_lst[gt_idx]\n",
    "\n",
    "        args.EAR = stim[0]\n",
    "        args.exper_dur = stim[1]\n",
    "        args.stim_dur = stim[2]\n",
    "\n",
    "        y_predicted = brightness_equ(PER, PFO, PFD, brightness_scale)\n",
    "        total_mse += np.mean((np.array(y_predicted) - gt_y) ** 2)\n",
    "    # print(total_mse/4)\n",
    "    return total_mse/4\n",
    "\n",
    "for subj in range(0, 9):\n",
    "    for cond in range(5):\n",
    "\n",
    "        i = cond*9 + subj\n",
    "        print(i)\n",
    "\n",
    "        if i != 22 : pass\n",
    "        else:\n",
    "            idx_lst = [0*9+subj, 1*9+subj, 2*9+subj, 3*9+subj, 4*9+subj]\n",
    "            idx_lst.remove(i)\n",
    "            \n",
    "            args.EAR = stim_info[i][0]\n",
    "            args.exper_dur = stim_info[i][1]\n",
    "            args.stim_dur = stim_info[i][2]\n",
    "\n",
    "            gt_curves_train = [np.nan_to_num(data.iloc[idx, 6:stim_info[idx][1]*args.interpolation_factor+6].to_numpy().astype('float64'), 0) for idx in idx_lst] \n",
    "            gt_curves_val = np.nan_to_num(data.iloc[i, 6:args.exper_dur*args.interpolation_factor+6].to_numpy().astype('float64'), 0)  \n",
    "\n",
    "            stim_lst = [stim_info[idx] for idx in idx_lst]\n",
    "            \n",
    "            # Initial guesses for parameters\n",
    "            initial_params = [2, 10, 10, 10]\n",
    "            # Run the optimization, passing the ground truth as an argument\n",
    "            result = minimize(objective_function, initial_params, args=(gt_curves_train,stim_lst), method='Nelder-Mead')\n",
    "            # Extract optimized parameters\n",
    "            optimized_params = result.x\n",
    "            print(\"Optimized Parameters:\", optimized_params)\n",
    "\n",
    "            current_time = datetime.now()\n",
    "            formatted_time = current_time.strftime(\"%Y-%m-%d %H:%M:%S\")\n",
    "        \n",
    "            \n",
    "            file_path = 'baseline_percondition_025interp.txt'  \n",
    "            line_to_add = '' \n",
    "            line_to_add += str(i) + ': '\n",
    "            line_to_add += str(list(optimized_params)) + '\\n' + formatted_time + '\\n'\n",
    "            with open(file_path, 'a') as file:\n",
    "                file.write(line_to_add)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## train per-subject 0.25 interp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "data = pd.read_csv('perez-fornos-2012.csv')  \n",
    "\n",
    "# fixed parameters - condition dependent\n",
    "class Args():\n",
    "    EAR = 20 # electrode activation rate \n",
    "    exper_dur = 30 # experiment duration\n",
    "    stim_dur = 10 # stimulus duration\n",
    "    interpolation_factor = 2 # number of frames sampled every second\n",
    "    \n",
    "args = Args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fit parameters\n",
    "PER=2 # persistence duration (usually between 0 and 8 seconds)\n",
    "PFO=10  # perceptual fading onset (usually between 0 and 10 seconds)\n",
    "PFD=10 # perceptual fading duration (usually between 0.5 and 60 seconds)\n",
    "brightness_scale = 10\n",
    "\n",
    "def brightness_equ(PER, PFO, PFD, brightness_scale):\n",
    "    SigmaX=.4 # phosphene dimension in x axis\n",
    "    SigmaY=.4 # phosphene dimension in y axis\n",
    "    ElectrodeArraySize=1 # number of electrodes (60(val=1), 256(2), 961(3), or 2916(4) electrodes)\n",
    "    ElectrodeArrayStructure=1 # electrode array structure (square(val = 1) or hexagonal(2)).\n",
    "\n",
    "    va = 0\n",
    "    vb = 0.333\n",
    "    vc = 0.667\n",
    "    vd = 1\n",
    "\n",
    "    SigmaX=.4 # phosphene dimension in x axis\n",
    "    SigmaY=.4 # phosphene dimension in y axis\n",
    "    MMr = 2  # number of rows\n",
    "    MMc = 3  # number of columns;\n",
    "    X0 = 20  # distance between two adjecent phosphenes in x axis\n",
    "    Y0 = 20  # distance between two adjecent phosphenes in y axis\n",
    "    C = MMc*Y0  # number of pixel columns\n",
    "    R = MMr*X0  # number of pixel rows, choose: phosphene size 0.4; X0 and Y0 = 24\n",
    "\n",
    "    DFR = 30 # display frame rate (usually 30 or 60 FPS)\n",
    "\n",
    "    v2 = 0\n",
    "    v3 = 0\n",
    "    count = 0\n",
    "\n",
    "    EAR = args.EAR  # electrode activation rate (usually 6 or 10 Hz)\n",
    "    exper_dur = args.exper_dur  # experiment duration\n",
    "    stim_dur = args.stim_dur  # stimulus duration\n",
    "\n",
    "    PER_Factor = 3.9/(PER*DFR) # 3.9 was set for 2 precent threshold, i.e. the phosphene falls to 2 precent of its original intensity;\n",
    "    PF_Factor = PFO*EAR\n",
    "    PFD_Factor = 3.9/(PFD*DFR)\n",
    "\n",
    "    # Creating meshgrid\n",
    "    x = np.arange(0, C) \n",
    "    y = np.arange(0, R)  \n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    # Initialization\n",
    "    z = np.zeros((R, C))\n",
    "    Frame = np.zeros((R, C))\n",
    "    FormerFrame = np.zeros((R, C))\n",
    "    PFC = np.zeros((MMr, MMc))\n",
    "    Theta = np.zeros((MMr, MMc))\n",
    "    a = np.zeros((MMr, MMc))\n",
    "    b = np.zeros((MMr, MMc))\n",
    "    c = np.zeros((MMr, MMc))\n",
    "    sigma_x = np.zeros((MMr, MMc))\n",
    "    sigma_y = np.zeros((MMr, MMc))\n",
    "\n",
    "    count = 0\n",
    "    for s in range(MMr):\n",
    "        for t in range(MMc):\n",
    "            Theta[s, t] = 0\n",
    "            sigma_x[s, t] = SigmaX\n",
    "            sigma_y[s, t] = SigmaY\n",
    "            a[s, t] = (np.cos(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_x[s, t]**2) + \\\n",
    "                      (np.sin(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_y[s, t]**2)\n",
    "            b[s, t] = -(np.sin(np.deg2rad(2 * Theta[s, t]))) / (4 * sigma_x[s, t]**2) + \\\n",
    "                      (np.sin(np.deg2rad(2 * Theta[s, t]))) / (4 * sigma_y[s, t]**2)\n",
    "            c[s, t] = (np.sin(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_x[s, t]**2) + \\\n",
    "                      (np.cos(np.deg2rad(Theta[s, t]))**2) / (2 * sigma_y[s, t]**2)\n",
    "\n",
    "    I_frame = [1]*(args.stim_dur*args.EAR)*MMr*MMc  # the brightness of the video is always 1\n",
    "    # video frame rate is the same as the electrode firing rate\n",
    "    \n",
    "    I_frame_0 = [0]*(args.exper_dur*args.EAR-args.stim_dur*args.EAR)*MMr*MMc  # no visual input after stim_dur\n",
    "    I_frame = np.append(np.array(I_frame), np.array(I_frame_0))\n",
    "    I_frame = np.array(I_frame).reshape(args.exper_dur*args.EAR, MMr, MMc)\n",
    "\n",
    "    results = []\n",
    "\n",
    "    for frame in range(exper_dur*args.EAR):\n",
    "        I = I_frame[frame]\n",
    "        # I = I / I.max()\n",
    "        MeanMatrix = np.zeros((MMr, MMc))\n",
    "        s = 0\n",
    "        t = 0\n",
    "        v = 0\n",
    "\n",
    "        start = int(np.floor((I.shape[0] % MMr) / 2) + np.floor(I.shape[0] / MMr) * MMr) - 1\n",
    "        stop = int(1 + np.floor((I.shape[0] % MMr) / 2)) - 1\n",
    "        step = -int(np.floor(I.shape[0] / MMr))\n",
    "        for m in range(start, stop, step):\n",
    "            for n in range(0, I.shape[1]//MMc*MMc - 1, I.shape[1] // MMc):\n",
    "                for j in range(m, m - (I.shape[0] // MMr), -1):\n",
    "                    # for k in range(n, n + (I.shape[1] // MMc)-2):\n",
    "                    for k in range(n, n - (I.shape[1] // MMc), -1):\n",
    "                        v += I[j, k]\n",
    "                v /= ((I.shape[0] // MMr) * (I.shape[1] // MMc))\n",
    "\n",
    "                if 0 <= v <= 0.25:\n",
    "                    v = va\n",
    "                elif 0.25 < v <= 0.5:\n",
    "                    v = vb\n",
    "                elif 0.5 < v <= 0.75:\n",
    "                    v = vc\n",
    "                else:\n",
    "                    v = vd\n",
    "                MeanMatrix[s, t] = v\n",
    "                v = 0\n",
    "                t += 1\n",
    "            t = 0\n",
    "            s += 1\n",
    "\n",
    "        getting_all_subframe = []\n",
    "\n",
    "        SubFrame_range = int(DFR) // int(EAR) + 1\n",
    "        if SubFrame_range <= 2: SubFrame_range = 2\n",
    "        for SubFrame in range(1, SubFrame_range):\n",
    "            for s in range(MMr):\n",
    "                for t in range(MMc):\n",
    "                    adjusted_t = t + 1\n",
    "                    adjusted_s = s + 1\n",
    "                    z += MeanMatrix[s, t] * np.exp(\n",
    "                        # -0.015 * (1 / (3 * MeanMatrix[s, t] + 1)) *\n",
    "                        -0.015 * (1 / (3 * MeanMatrix[s, t] + 1)) *\n",
    "                        # (a[s, t] * (X - X0 * t - v2 + X0 // 2) ** 2 +\n",
    "                        (a[s, t] * (X - X0 * adjusted_t - v2 + X0 // 2) ** 2 +\n",
    "                        # 2 * b[s, t] * (X - X0 * t - v2 + X0 // 2) * (Y - Y0 * s + Y0 // 2) +\n",
    "                        2 * b[s, t] * (X - X0 * adjusted_t - v2 + X0 // 2) * (Y - Y0 * adjusted_s + Y0 // 2) +\n",
    "                        # c[s, t] * (Y - Y0 * s + Y0 // 2) ** 2))\n",
    "                        c[s, t] * (Y - Y0 * adjusted_s + Y0 // 2) ** 2))\n",
    "\n",
    "                v2 -= (X0 / 2) * v3\n",
    "                v3 *= -1\n",
    "\n",
    "            for s in range(R, 0, -X0):\n",
    "                for t in range(0, C, Y0):\n",
    "                    if SubFrame == 1:\n",
    "                        condition1 = (np.mean(z[s-X0:s-1, t-1:t+Y0-2]) == 0 and np.mean(FormerFrame[s-X0:s-1, t-1:t+Y0-2]) == 0)\n",
    "                        condition2 = (np.mean(z[s-X0:s-1, t-1:t+Y0-2]) < np.mean(FormerFrame[s-X0:s-1, t:t+Y0-2]))\n",
    "                        condition3 = (MeanMatrix[int(s/X0)-1, t // Y0] == 0)\n",
    "                        condition4 = (PFC[int(s/X0)-1, t // Y0] < PF_Factor)\n",
    "                        if condition1 or condition2 or condition3 and condition4:\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PER_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                        elif np.mean(z[s-X0:s-1, t-1:t+Y0-2]) >= np.mean(FormerFrame[s-X0:s-1, t-1:t+Y0-2]) and condition4:\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = z[s-X0:s-1, t-1:t+Y0-2]\n",
    "                            PFC[s//X0-1, t//Y0] += 1\n",
    "                        elif PFC[s//X0-1, t//Y0] >= PF_Factor:\n",
    "                            # Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-3 * PER_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                            Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PFD_Factor) * FormerFrame[s-X0:s-1, t-1:t+Y0-2]\n",
    "                    else:\n",
    "                        Frame[s-X0:s-1, t-1:t+Y0-2] = np.exp(-PER_Factor) * Frame[s-X0:s-1, t-1:t+Y0-2]\n",
    "            z = np.zeros((R,C))\n",
    "            count = count + 1\n",
    "\n",
    "            FormerFrame = Frame\n",
    "            getting_all_subframe.append(Frame.max())\n",
    "        results.append(max(getting_all_subframe))\n",
    "\n",
    "    output =  [x * brightness_scale for x in results]\n",
    "\n",
    "    output_ = []\n",
    "    combined_frame = args.EAR/args.interpolation_factor\n",
    "    for i in range(args.exper_dur*args.interpolation_factor):\n",
    "        output_.append(max(output[int(i*combined_frame):int((i+1)*combined_frame)]))\n",
    "    return output_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stim_info = {}\n",
    "# EAR, exper_dur, stim_dur\n",
    "for i in range(9): stim_info[i] = [20, 30, 10]\n",
    "for i in range(9,18): stim_info[i] = [5, 30, 10]\n",
    "for i in range(18,27): stim_info[i] = [60, 30, 10]\n",
    "for i in range(27,36): stim_info[i] = [20, 10, 1]\n",
    "for i in range(36,45): stim_info[i] = [20, 70, 60]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "args.interpolation_factor = 4\n",
    "\n",
    "def objective_function(params, gt_curves_train, stim_lst):\n",
    "    # persistence time, fading time onset, fading time duration, brightness scale\n",
    "    PER, PFO, PFD, brightness_scale = params\n",
    "    y_predicted = brightness_equ(PER, PFO, PFD, brightness_scale)\n",
    "    \n",
    "    \n",
    "    total_mse = 0\n",
    "    for gt_idx in range(len(gt_curves_train)):\n",
    "        gt_y = gt_curves_train[gt_idx]\n",
    "        # stim = stim_lst[gt_idx]\n",
    "        # args.EAR = stim[0]\n",
    "        # args.exper_dur = stim[1]\n",
    "        # args.stim_dur = stim[2]\n",
    "\n",
    "        total_mse += np.mean((np.array(y_predicted) - gt_y) ** 2)\n",
    "    # print(total_mse/len(gt_curves_train))\n",
    "    return total_mse/len(gt_curves_train)\n",
    "\n",
    "for subj in range(9):\n",
    "    for cond in range(5):\n",
    "\n",
    "        i = cond*9 + subj\n",
    "        print(i)\n",
    "\n",
    "        if i in [0,9,18,27] : pass\n",
    "        else:\n",
    "            idx_lst = [cond*9 + sub for sub in range(9)]\n",
    "            idx_lst.remove(i)\n",
    "            \n",
    "            args.EAR = stim_info[i][0]\n",
    "            args.exper_dur = stim_info[i][1]\n",
    "            args.stim_dur = stim_info[i][2]\n",
    "\n",
    "            gt_curves_train = [np.nan_to_num(data.iloc[idx, 6:stim_info[idx][1]*args.interpolation_factor+6].to_numpy().astype('float64'), 0) for idx in idx_lst] \n",
    "            gt_curves_val = np.nan_to_num(data.iloc[i, 6:args.exper_dur*args.interpolation_factor+6].to_numpy().astype('float64'), 0)  \n",
    "\n",
    "            stim_lst = [stim_info[idx] for idx in idx_lst]\n",
    "            \n",
    "            # Initial guesses for parameters\n",
    "            initial_params = [2, 10, 10, 10]\n",
    "            # Run the optimization, passing the ground truth as an argument\n",
    "            result = minimize(objective_function, initial_params, args=(gt_curves_train,stim_lst), method='Nelder-Mead')\n",
    "            # Extract optimized parameters\n",
    "            optimized_params = result.x\n",
    "            print(\"Optimized Parameters:\", optimized_params)\n",
    "\n",
    "            current_time = datetime.now()\n",
    "            formatted_time = current_time.strftime(\"%Y-%m-%d %H:%M:%S\")\n",
    "            \n",
    "            file_path = 'baseline_persubject_025interp.txt'  \n",
    "            line_to_add = '' \n",
    "            line_to_add += str(i) + ': '\n",
    "            line_to_add += str(list(optimized_params)) + '\\n' + formatted_time + '\\n'\n",
    "            with open(file_path, 'a') as file:\n",
    "                file.write(line_to_add)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "top5HkC-WjYQ"
   ],
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}