Core.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Re-implmentation of AI Clinician Matlab Code in Python \n",
    "# Author: KyungJoong Kim (GIST, South Korea)\n",
    "# Date: 2020 June 2 \n",
    "# \n",
    "# This code is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;\n",
    "# without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE\n",
    "\n",
    "\n",
    "# Note \n",
    "#\n",
    "# K-Means in scikit-learn will produce differnt outcome with Matlab's original K-means \n",
    "# The random number generator will produce different random number sequence \n",
    "# \n",
    "# Todo: \n",
    "# eICU was not included in this re-implementation \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "## AI Clinician core code\n",
    "\n",
    "# (c) Matthieu Komorowski, Imperial College London 2015-2019\n",
    "# as seen in publication: https://www.nature.com/articles/s41591-018-0213-5\n",
    "\n",
    "# version 16 Feb 19\n",
    "# Builds 500 models using MIMIC-III training data\n",
    "# Records best candidate models along the way from off-policy policy evaluation on MIMIC-III validation data\n",
    "# Tests the best model on eRI data\n",
    "\n",
    "\n",
    "# TAKES:\n",
    "        # MIMICtable = m*59 table with raw values from MIMIC\n",
    "        # eICUtable = n*56 table with raw values from eICU\n",
    "        \n",
    "\n",
    "# GENERATES:\n",
    "        # MIMICraw = MIMIC RAW DATA m*47 array with columns in right order\n",
    "        # MIMICzs = MIMIC ZSCORED m*47 array with columns in right order, matching MIMICraw\n",
    "        # eICUraw = eICU RAW DATA n*47 array with columns in right order, matching MIMICraw\n",
    "        # eICUzs = eICU ZSCORED n*47 array with columns in right order, matching MIMICraw\n",
    "        # recqvi = summary statistics of all 500 models\n",
    "        # idxs = state membership of MIMIC test records, for all 500 models\n",
    "     \t# OA = optimal policy, for all 500 models\n",
    "        # allpols = detailed data about the best candidate models\n",
    "\n",
    "# This code is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;\n",
    "# without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE\n",
    "\n",
    "# Note: The size of the cohort will depend on which version of MIMIC-III is used.\n",
    "# The original cohort from the 2018 Nature Medicine publication was built using MIMIC-III v1.3.\n",
    "\n",
    "import pickle \n",
    "import numpy as np \n",
    "import pandas as pd \n",
    "from scipy.stats import zscore, rankdata\n",
    "import math \n",
    "import scipy.io as sio \n",
    "import datetime \n",
    "from scipy.stats.mstats import mquantiles\n",
    "from mdptoolbox.mdp import PolicyIteration\n",
    "from reinforcement_learning_mp import offpolicy_multiple_eval_010518\n",
    "from kmeans_mp import kmeans_with_multiple_runs \n",
    "from multiprocessing import freeze_support \n",
    "\n",
    "def my_zscore(x):\n",
    "    return zscore(x,ddof=1),np.mean(x,axis=0),np.std(x,axis=0,ddof=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######### Functions used in Reinforcement Learning ######## \n",
    "\n",
    "\n",
    "class PolicyIteration_with_Q(PolicyIteration):\n",
    "    def __init__(self, transitions, reward, discount, policy0=None,max_iter=1000, eval_type=0, skip_check=False):\n",
    "        # Python MDP toolbox from https://github.com/sawcordwell/pymdptoolbox\n",
    "        # In Matlab MDP Toolbox, P = (S, S, A), R = (S, A) \n",
    "        # In Python MDP Toolbox, P = (A, S, S), R= (S, A)\n",
    "\n",
    "        transitions = np.transpose(transitions,(2,0,1)).copy() # Change to Action First (A, S, S)\n",
    "        skip_check = True # To Avoid StochasticError: 'PyMDPToolbox - The transition probability matrix is not stochastic.'\n",
    "\n",
    "        PolicyIteration.__init__(self, transitions, reward, discount, policy0=None,max_iter=1000, eval_type=0, skip_check=skip_check)\n",
    "    \n",
    "    def _bellmanOperator_with_Q(self, V=None):\n",
    "        # Apply the Bellman operator on the value function.\n",
    "        #\n",
    "        # Updates the value function and the Vprev-improving policy.\n",
    "        #\n",
    "        # Returns: (policy, Q, value), tuple of new policy and its value\n",
    "        #\n",
    "        # If V hasn't been sent into the method, then we assume to be working\n",
    "        # on the objects V attribute\n",
    "        if V is None:\n",
    "            # this V should be a reference to the data rather than a copy\n",
    "            V = self.V\n",
    "        else:\n",
    "            # make sure the user supplied V is of the right shape\n",
    "            try:\n",
    "                assert V.shape in ((self.S,), (1, self.S)), \"V is not the \" \\\n",
    "                    \"right shape (Bellman operator).\"\n",
    "            except AttributeError:\n",
    "                raise TypeError(\"V must be a numpy array or matrix.\")\n",
    "        # Looping through each action the the Q-value matrix is calculated.\n",
    "        # P and V can be any object that supports indexing, so it is important\n",
    "        # that you know they define a valid MDP before calling the\n",
    "        # _bellmanOperator method. Otherwise the results will be meaningless.\n",
    "        Q = np.empty((self.A, self.S))\n",
    "        for aa in range(self.A):\n",
    "            Q[aa] = self.R[aa] + self.discount * self.P[aa].dot(V)\n",
    "        # Get the policy and value, for now it is being returned but...\n",
    "        # Which way is better?\n",
    "        # 1. Return, (policy, value)\n",
    "        return (Q.argmax(axis=0), Q, Q.max(axis=0))\n",
    "        # 2. update self.policy and self.V directly\n",
    "        # self.V = Q.max(axis=1)\n",
    "        # self.policy = Q.argmax(axis=1)\n",
    "    \n",
    "    def run(self):\n",
    "        # Run the policy iteration algorithm.\n",
    "        self._startRun()\n",
    "\n",
    "        while True:\n",
    "            self.iter += 1\n",
    "            # these _evalPolicy* functions will update the classes value\n",
    "            # attribute\n",
    "            if self.eval_type == \"matrix\":\n",
    "                self._evalPolicyMatrix()\n",
    "            elif self.eval_type == \"iterative\":\n",
    "                self._evalPolicyIterative()\n",
    "            # This should update the classes policy attribute but leave the\n",
    "            # value alone\n",
    "            policy_next, Q, null = self._bellmanOperator_with_Q()\n",
    "            del null\n",
    "            # calculate in how many places does the old policy disagree with\n",
    "            # the new policy\n",
    "            n_different = (policy_next != self.policy).sum()\n",
    "            # if verbose then continue printing a table\n",
    "            if self.verbose:\n",
    "                _printVerbosity(self.iter, n_different)\n",
    "            # Once the policy is unchanging of the maximum number of\n",
    "            # of iterations has been reached then stop\n",
    "            if n_different == 0:\n",
    "                if self.verbose:\n",
    "                    print(_MSG_STOP_UNCHANGING_POLICY)\n",
    "                break\n",
    "            elif self.iter == self.max_iter:\n",
    "                if self.verbose:\n",
    "                    print(_MSG_STOP_MAX_ITER)\n",
    "                break\n",
    "            elif self.iter > 20 and n_different <=5 : # This condition was added from the Nature Code \n",
    "                if self.verbose: \n",
    "                    print((_MSG_STOP))\n",
    "                break \n",
    "            else:\n",
    "                self.policy = policy_next\n",
    "\n",
    "        self._endRun()\n",
    "        \n",
    "        return Q \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == '__main__': \n",
    "    \n",
    "    freeze_support()\n",
    "    \n",
    "    # To ignore 'Runtime Warning: Invalid value encountered in greater' caused by NaN \n",
    "    np.warnings.filterwarnings('ignore')\n",
    "\n",
    "    # Load pickle \n",
    "    with open('step_4_start.pkl', 'rb') as file:\n",
    "        MIMICtable = pickle.load(file)\n",
    "    \n",
    "    #############################  MODEL PARAMETERS   #####################################\n",
    "\n",
    "    print('####  INITIALISATION  ####') \n",
    "\n",
    "    nr_reps=500               # nr of repetitions (total nr models) % 500 \n",
    "    nclustering=32            # how many times we do clustering (best solution will be chosen) % 32\n",
    "    prop=0.25                 # proportion of the data we sample for clustering\n",
    "    gamma=0.99                # gamma\n",
    "    transthres=5              # threshold for pruning the transition matrix\n",
    "    polkeep=1                 # count of saved policies\n",
    "    ncl=750                   # nr of states\n",
    "    nra=5                     # nr of actions (2 to 10)\n",
    "    ncv=5                     # nr of crossvalidation runs (each is 80% training / 20% test)\n",
    "    OA=np.full((752,nr_reps),np.nan)       # record of optimal actions\n",
    "    recqvi=np.full((nr_reps*2,30),np.nan)  # saves data about each model (1 row per model)\n",
    "    # allpols=[]  # saving best candidate models\n",
    "    \n",
    "    \n",
    "    # #################   Convert training data and compute conversion factors    ######################\n",
    "\n",
    "    # all 47 columns of interest\n",
    "    colbin = ['gender','mechvent','max_dose_vaso','re_admission'] \n",
    "    colnorm= ['age','Weight_kg','GCS','HR','SysBP','MeanBP','DiaBP','RR','Temp_C','FiO2_1',\\\n",
    "        'Potassium','Sodium','Chloride','Glucose','Magnesium','Calcium',\\\n",
    "        'Hb','WBC_count','Platelets_count','PTT','PT','Arterial_pH','paO2','paCO2',\\\n",
    "        'Arterial_BE','HCO3','Arterial_lactate','SOFA','SIRS','Shock_Index','PaO2_FiO2','cumulated_balance'] \n",
    "    collog=['SpO2','BUN','Creatinine','SGOT','SGPT','Total_bili','INR','input_total','input_4hourly','output_total','output_4hourly']\n",
    "\n",
    "    colbin=np.where(np.isin(MIMICtable.columns,colbin))[0]\n",
    "    colnorm=np.where(np.isin(MIMICtable.columns,colnorm))[0]\n",
    "    collog=np.where(np.isin(MIMICtable.columns,collog))[0]\n",
    "\n",
    "    # find patients who died in ICU during data collection period\n",
    "    # ii=MIMICtable.bloc==1&MIMICtable.died_within_48h_of_out_time==1& MIMICtable.delay_end_of_record_and_discharge_or_death<24;\n",
    "    # icustayidlist=MIMICtable.icustayid;\n",
    "    # ikeep=~ismember(icustayidlist,MIMICtable.icustayid(ii));\n",
    "    reformat5=MIMICtable.values.copy() \n",
    "    # reformat5=reformat5(ikeep,:);\n",
    "    icustayidlist=MIMICtable['icustayid']\n",
    "    icuuniqueids=np.unique(icustayidlist) # list of unique icustayids from MIMIC\n",
    "    idxs=np.full((icustayidlist.shape[0],nr_reps),np.nan) # record state membership test cohort\n",
    "\n",
    "    MIMICraw=MIMICtable.iloc[:, np.concatenate([colbin,colnorm,collog])] \n",
    "    MIMICraw=MIMICraw.values.copy()  # RAW values\n",
    "    MIMICzs=np.concatenate([reformat5[:, colbin]-0.5, zscore(reformat5[:,colnorm],ddof=1), zscore(np.log(0.1+reformat5[:, collog]),ddof=1)],axis=1)\n",
    "    MIMICzs[:,3]=np.log(MIMICzs[:,3]+0.6)   # MAX DOSE NORAD \n",
    "    MIMICzs[:,44]=2*MIMICzs[:,44]   # increase weight of this variable\n",
    "\n",
    "\n",
    "    # eICU section was not implemented \n",
    "    \n",
    "    # compute conversion factors using MIMIC data\n",
    "    a=MIMICraw[:, 0:3]-0.5 \n",
    "    b= np.log(MIMICraw[:,3]+0.1)\n",
    "    c,cmu,csigma = my_zscore(MIMICraw[:,4:36])\n",
    "    d,dmu,dsigma = my_zscore(np.log(0.1+MIMICraw[:,36:47]))\n",
    "    \n",
    "\n",
    "    ####################### Main LOOP ###########################\n",
    "    bestpol = 0 \n",
    "    \n",
    "    for modl in range(nr_reps):  # MAIN LOOP OVER ALL MODELS\n",
    "        N=icuuniqueids.size # total number of rows to choose from\n",
    "        grp=np.floor(ncv*np.random.rand(N,1)+1);  #list of 1 to 5 (20% of the data in each grp) -- this means that train/test MIMIC split are DIFFERENT in all the 500 models\n",
    "        crossval=1;\n",
    "        trainidx=icuuniqueids[np.where(grp!=crossval)[0]]\n",
    "        testidx=icuuniqueids[np.where(grp==crossval)[0]]\n",
    "        train=np.isin(icustayidlist,trainidx)\n",
    "        test=np.isin(icustayidlist,testidx)\n",
    "        X=MIMICzs[train,:]\n",
    "        Xtestmimic=MIMICzs[~train,:]\n",
    "        blocs=reformat5[train,0]\n",
    "        bloctestmimic=reformat5[~train,0]\n",
    "        ptid=reformat5[train,1]\n",
    "        ptidtestmimic=reformat5[~train,1] \n",
    "        outcome=9 #   HOSP _ MORTALITY = 7 / 90d MORTA = 9\n",
    "        Y90=reformat5[train,outcome];   \n",
    "\n",
    "\n",
    "        print('########################   MODEL NUMBER : ',modl)\n",
    "        print(datetime.datetime.now())\n",
    "\n",
    "        #######   find best clustering solution (lowest intracluster variability)  ####################\n",
    "        print('####  CLUSTERING  ####') # BY SAMPLING\n",
    "        N=X.shape[0] #total number of rows to choose from\n",
    "        sampl=X[np.where(np.floor(np.random.rand(N,1)+prop))[0],:]\n",
    "\n",
    "        C = kmeans_with_multiple_runs(ncl,10000,nclustering,sampl) \n",
    "        idx = C.predict(X)\n",
    "\n",
    "        ############################## CREATE ACTIONS  ########################\n",
    "        print('####  CREATE ACTIONS  ####') \n",
    "\n",
    "        nact=nra*nra\n",
    "\n",
    "        iol=MIMICtable.columns.get_loc('input_4hourly') \n",
    "        vcl=MIMICtable.columns.get_loc('max_dose_vaso') \n",
    "\n",
    "        a= reformat5[:,iol].copy()                   # IV fluid\n",
    "        a= rankdata(a[a>0])/a[a>0].shape[0]   # excludes zero fluid (will be action 1)\n",
    "\n",
    "        iof=np.floor((a+0.2499999999)*4)  #converts iv volume in 4 actions\n",
    "\n",
    "        a= reformat5[:,iol].copy() \n",
    "        a= np.where(a>0)[0]  # location of non-zero fluid in big matrix\n",
    "\n",
    "        io=np.ones((reformat5.shape[0],1))  # array of ones, by default     \n",
    "        io[a]=(iof+1).reshape(-1,1)   # where more than zero fluid given: save actual action\n",
    "        io = io.ravel() \n",
    "\n",
    "        vc=reformat5[:,vcl].copy() \n",
    "        vcr= rankdata(vc[vc!=0])/vc[vc!=0].size\n",
    "        vcr=np.floor((vcr+0.249999999999)*4)  # converts to 4 bins\n",
    "        vcr[vcr==0]=1\n",
    "        vc[vc!=0]=vcr+1 \n",
    "        vc[vc==0]=1\n",
    "\n",
    "        ma1 = np.array([np.median(reformat5[io==1,iol]),np.median(reformat5[io==2,iol]),np.median(reformat5[io==3,iol]), np.median(reformat5[io==4,iol]),np.median(reformat5[io==5,iol])]) # median dose of drug in all bins\n",
    "        ma2 = np.array([np.median(reformat5[vc==1,vcl]),np.median(reformat5[vc==2,vcl]),np.median(reformat5[vc==3,vcl]), np.median(reformat5[vc==4,vcl]),np.median(reformat5[vc==5,vcl])])\n",
    "\n",
    "        med = np.concatenate([io.reshape(-1,1),vc.reshape(-1,1)],axis=1)\n",
    "        uniqueValues,actionbloc = np.unique(med,axis=0,return_inverse=True)\n",
    "\n",
    "        actionbloctrain=actionbloc[train] \n",
    "\n",
    "        ma2Values = ma2[uniqueValues[:,1].astype('int64')-1].reshape(-1,1)\n",
    "        ma1Values = ma1[uniqueValues[:,0].astype('int64')-1].reshape(-1,1)\n",
    "\n",
    "        uniqueValuesdose = np.concatenate([ma2Values,ma1Values],axis=1) # median dose of each bin for all 25 actions \n",
    "\n",
    "        ####################################################################################################################################\n",
    "        print('####  CREATE QLDATA3  ####')\n",
    "\n",
    "        r=np.array([100, -100]).reshape(1,-1)\n",
    "        r2=r*(2*(1-Y90.reshape(-1,1))-1)\n",
    "        # because idx and actionbloctrain are index, it's equal to (Matlab's original value -1)\n",
    "        qldata=np.concatenate([blocs.reshape(-1,1), idx.reshape(-1,1), actionbloctrain.reshape(-1,1), Y90.reshape(-1,1), r2],axis=1)  # contains bloc / state / action / outcome&reward     \n",
    "        # 0 = died in Python, 1 = died in Matlab \n",
    "        qldata3=np.zeros((np.floor(qldata.shape[0]*1.2).astype('int64'),4))\n",
    "        c=-1\n",
    "        abss=np.array([ncl+1, ncl]) #absorbing states numbers # 751, 750 \n",
    "\n",
    "        for i in range(qldata.shape[0]-1):\n",
    "            c=c+1\n",
    "            qldata3[c,:]=qldata[i,0:4] \n",
    "            if(qldata[i+1,0]==1): #end of trace for this patient\n",
    "                c=c+1     \n",
    "                qldata3[c,:]=np.array([qldata[i,0]+1, abss[int(qldata[i,3])], -1, qldata[i,4]]) \n",
    "\n",
    "        qldata3=qldata3[:c+1,:]\n",
    "\n",
    "\n",
    "        # ###################################################################################################################################\n",
    "        print(\"####  CREATE TRANSITION MATRIX T(S'',S,A) ####\")\n",
    "        transitionr=np.zeros((ncl+2,ncl+2,nact))  #this is T(S',S,A)\n",
    "        sums0a0=np.zeros((ncl+2,nact)) \n",
    "\n",
    "        for i in range(qldata3.shape[0]-1):    \n",
    "            if (qldata3[i+1,0]!=1) : # if we are not in the last state for this patient = if there is a transition to make!\n",
    "                S0=int(qldata3[i,1]) \n",
    "                S1=int(qldata3[i+1,1])\n",
    "                acid= int(qldata3[i,2]) \n",
    "                transitionr[S1,S0,acid]=transitionr[S1,S0,acid]+1 \n",
    "                sums0a0[S0,acid]=sums0a0[S0,acid]+1\n",
    "\n",
    "        sums0a0[sums0a0<=transthres]=0  #delete rare transitions (those seen less than 5 times = bottom 50%!!)\n",
    "\n",
    "        for i in range(ncl+2): \n",
    "            for j in range(nact): \n",
    "                if sums0a0[i,j]==0: \n",
    "                    transitionr[:,i,j]=0; \n",
    "                else:\n",
    "                    transitionr[:,i,j]=transitionr[:,i,j]/sums0a0[i,j]\n",
    "\n",
    "\n",
    "        transitionr[np.isnan(transitionr)]=0  #replace NANs with zeros\n",
    "        transitionr[np.isinf(transitionr)]=0  #replace NANs with zeros\n",
    "\n",
    "        physpol=sums0a0/np.sum(sums0a0, axis=1).reshape(-1,1)    #physicians policy: what action was chosen in each state\n",
    "\n",
    "\n",
    "        print(\"####  CREATE TRANSITION MATRIX T(S,S'',A)  ####\")\n",
    "\n",
    "        transitionr2=np.zeros((ncl+2,ncl+2,nact))  # this is T(S,S',A)\n",
    "        sums0a0=np.zeros((ncl+2,nact))\n",
    "\n",
    "        for i in range(qldata3.shape[0]-1) : \n",
    "            if (qldata3[i+1,0]!=1) : # if we are not in the last state for this patient = if there is a transition to make!\n",
    "                S0=int(qldata3[i,1])\n",
    "                S1=int(qldata3[i+1,1])\n",
    "                acid= int(qldata3[i,2]) \n",
    "                transitionr2[S0,S1,acid]=transitionr2[S0,S1,acid]+1;  \n",
    "                sums0a0[S0,acid]=sums0a0[S0,acid]+1\n",
    "\n",
    "        sums0a0[sums0a0<=transthres]=0;  #delete rare transitions (those seen less than 5 times = bottom 50%!!) IQR = 2-17\n",
    "\n",
    "        for i in range(ncl+2): \n",
    "            for j in range(nact): \n",
    "                if sums0a0[i,j]==0:\n",
    "                    transitionr2[i,:,j]=0 \n",
    "                else: \n",
    "                    transitionr2[i,:,j]=transitionr2[i,:,j]/sums0a0[i,j]\n",
    "\n",
    "        transitionr2[np.isnan(transitionr2)]=0 #replace NANs with zeros\n",
    "        transitionr2[np.isinf(transitionr2)]=0 # replace infs with zeros\n",
    "\n",
    "        print('####  CREATE REWARD MATRIX  R(S,A) ####')\n",
    "        # CF sutton& barto bottom 1998 page 106. i compute R(S,A) from R(S'SA) and T(S'SA)\n",
    "        r3=np.zeros((ncl+2,ncl+2,nact))\n",
    "        r3[ncl,:,:]=-100\n",
    "        r3[ncl+1,:,:]=100\n",
    "        R=sum(transitionr*r3)\n",
    "        R=np.squeeze(R)   #remove 1 unused dimension\n",
    "\n",
    "\n",
    "        print('####  POLICY ITERATION   ####')\n",
    "\n",
    "\n",
    "        pi = PolicyIteration_with_Q(transitionr2, R, gamma, np.ones((ncl+2,1))) \n",
    "        Qon = np.transpose(pi.run()) \n",
    "        OptimalAction=np.argmax(Qon,axis=1).reshape(-1,1)  #deterministic \n",
    "        OA[:,modl]=OptimalAction.ravel() #save optimal actions\n",
    "\n",
    "\n",
    "        print('####  OFF-POLICY EVALUATION - MIMIC TRAIN SET ####')\n",
    "        # create new version of QLDATA3\n",
    "\n",
    "        r=np.array([100, -100]).reshape(1,-1)\n",
    "        r2=r*(2*(1-Y90.reshape(-1,1))-1)\n",
    "        # because idx and actionbloctrain are index, it's equal to (Matlab's original value -1)\n",
    "        qldata=np.concatenate([blocs.reshape(-1,1), idx.reshape(-1,1), actionbloctrain.reshape(-1,1), Y90.reshape(-1,1),np.zeros((idx.size,1)), r2[:,0].reshape(-1,1), ptid.reshape(-1,1) ],axis=1)   # contains bloc / state / action / outcome&reward     \n",
    "        # 0 = died in Python, 1 = died in Matlab \n",
    "        qldata3=np.zeros((np.floor(qldata.shape[0]*1.2).astype('int64'),8))\n",
    "        c=-1\n",
    "        abss=np.array([ncl+1, ncl]) #absorbing states numbers # 751, 750 \n",
    "\n",
    "        for i in range(qldata.shape[0]-1):\n",
    "            c=c+1\n",
    "            qldata3[c,:]=qldata[i,[0,1,2,4,6,6,6,6]] \n",
    "            if(qldata[i+1,0]==1): #end of trace for this patient\n",
    "                c=c+1     \n",
    "                qldata3[c,:]=np.array([qldata[i,0]+1, abss[int(qldata[i,3])], -1, qldata[i,5],0,0,-1,qldata[i,6]]) \n",
    "\n",
    "        qldata3=qldata3[:c+1,:]\n",
    "\n",
    "        #  add pi(s,a) and b(s,a)\n",
    "        p=0.01 #softening policies \n",
    "        softpi=physpol.copy() # behavior policy = clinicians' \n",
    "        for i in range(ncl): \n",
    "            ii=softpi[i,:]==0    \n",
    "            z=p/sum(ii)    \n",
    "            nz=p/sum(~ii)    \n",
    "            softpi[i,ii]=z;   \n",
    "            softpi[i,~ii]=softpi[i,~ii]-nz;\n",
    "\n",
    "        softb=np.abs(np.zeros((ncl+2,nact))-p/24) #\"optimal\" policy = target policy = evaluation policy \n",
    "\n",
    "        for i in range(ncl): \n",
    "            softb[i,OptimalAction[i]]=1-p\n",
    "\n",
    "\n",
    "        for i in range(qldata3.shape[0]):  # adding the probas of policies to qldata3\n",
    "            if qldata3[i,1]<ncl :\n",
    "                qldata3[i,4]=softpi[int(qldata3[i,1]),int(qldata3[i,2])]\n",
    "                qldata3[i,5]=softb[int(qldata3[i,1]),int(qldata3[i,2])]\n",
    "                qldata3[i,6]=OptimalAction[int(qldata3[i,1])]  #optimal action\n",
    "\n",
    "        qldata3train=qldata3.copy() \n",
    "\n",
    "\n",
    "        bootql,bootwis = offpolicy_multiple_eval_010518(qldata3,physpol, 0.99,1,6,750)\n",
    "\n",
    "        recqvi[modl,0]=modl\n",
    "        recqvi[modl,3]=np.nanmean(bootql)\n",
    "        recqvi[modl,4]=mquantiles(bootql,0.99, alphap=0.5, betap=0.5)[0]\n",
    "        recqvi[modl,5]=np.nanmean(bootwis) # we want this as high as possible\n",
    "        recqvi[modl,6]=mquantiles(bootwis,0.05, alphap=0.5, betap=0.5)[0] #we want this as high as possible\n",
    "\n",
    "\n",
    "        # testing on MIMIC-test\n",
    "        print('####  OFF-POLICY EVALUATION - MIMIC TEST SET ####')\n",
    "\n",
    "        # create new version of QLDATA3 with MIMIC TEST samples\n",
    "\n",
    "        idxtest = C.predict(Xtestmimic)\n",
    "\n",
    "        idxs[test,modl]=idxtest.ravel()  #important: record state membership of test cohort\n",
    "\n",
    "        actionbloctest=actionbloc[~train] \n",
    "        Y90test=reformat5[~train,outcome]\n",
    "\n",
    "        r=np.array([100, -100]).reshape(1,-1)\n",
    "        r2=r*(2*(1-Y90test.reshape(-1,1))-1)\n",
    "        # because idx and actionbloctrain are index, it's equal to (Matlab's original value -1)\n",
    "\n",
    "        qldata=np.concatenate([bloctestmimic.reshape(-1,1), idxtest.reshape(-1,1), actionbloctest.reshape(-1,1), Y90test.reshape(-1,1),np.zeros((idxtest.size,1)), r2[:,0].reshape(-1,1), ptidtestmimic.reshape(-1,1) ],axis=1)   # contains bloc / state / action / outcome&reward     \n",
    "        # 0 = died in Python, 1 = died in Matlab \n",
    "        qldata3=np.zeros((np.floor(qldata.shape[0]*1.2).astype('int64'),8))\n",
    "        c=-1\n",
    "        abss=np.array([ncl+1, ncl]) #absorbing states numbers # 751, 750 \n",
    "\n",
    "        for i in range(qldata.shape[0]-1):\n",
    "            c=c+1\n",
    "            qldata3[c,:]=qldata[i,[0,1,2,4,6,6,6,6]] \n",
    "            if(qldata[i+1,0]==1): #end of trace for this patient\n",
    "                c=c+1     \n",
    "                qldata3[c,:]=np.array([qldata[i,0]+1, abss[int(qldata[i,3])], -1, qldata[i,5],0,0,-1,qldata[i,6]]) \n",
    "\n",
    "        qldata3=qldata3[:c+1,:]\n",
    "\n",
    "        #  add pi(s,a) and b(s,a)\n",
    "        p=0.01 # small correction factor #softening policies \n",
    "        softpi=physpol.copy() # behavior policy = clinicians' \n",
    "        for i in range(ncl): \n",
    "            ii=softpi[i,:]==0    \n",
    "            z=p/sum(ii)    \n",
    "            nz=p/sum(~ii)    \n",
    "            softpi[i,ii]=z;   \n",
    "            softpi[i,~ii]=softpi[i,~ii]-nz;\n",
    "\n",
    "        softb=np.abs(np.zeros((ncl+2,nact))-p/24) #\"optimal\" policy = target policy = evaluation policy \n",
    "\n",
    "        for i in range(ncl): \n",
    "            softb[i,OptimalAction[i]]=1-p\n",
    "\n",
    "        for i in range(qldata3.shape[0]):  # adding the probas of policies to qldata3\n",
    "            if qldata3[i,1]<ncl :\n",
    "                qldata3[i,4]=softpi[int(qldata3[i,1]),int(qldata3[i,2])]\n",
    "                qldata3[i,5]=softb[int(qldata3[i,1]),int(qldata3[i,2])]\n",
    "                qldata3[i,6]=OptimalAction[int(qldata3[i,1])]  #optimal action\n",
    "\n",
    "        qldata3test=qldata3.copy() \n",
    "\n",
    "        bootmimictestql,bootmimictestwis = offpolicy_multiple_eval_010518(qldata3,physpol, 0.99,1,6,2000)\n",
    "\n",
    "        recqvi[modl,18]=mquantiles(bootmimictestql,0.95, alphap=0.5, betap=0.5)[0] #PHYSICIANS' 95% UB\n",
    "        recqvi[modl,19]=np.nanmean(bootmimictestql)\n",
    "        recqvi[modl,20]=mquantiles(bootmimictestql,0.99, alphap=0.5, betap=0.5)[0]\n",
    "        recqvi[modl,21]=np.nanmean(bootmimictestwis) \n",
    "        recqvi[modl,22]=mquantiles(bootmimictestwis,0.01, alphap=0.5, betap=0.5)[0] \n",
    "        recqvi[modl,23]=mquantiles(bootmimictestwis,0.05, alphap=0.5, betap=0.5)[0] #AI 95% LB, we want this as high as possible\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        if recqvi[modl,23] > 40: #saves time if policy is not good on MIMIC test: skips to next model\n",
    "            print('########################## eICU TEST SET #############################')\n",
    "            # eICU part was not implemented \n",
    "\n",
    "\n",
    "\n",
    "        # eICU testing was not included \n",
    "        if recqvi[modl,23]>0 : #  & recqvi(modl,14)>0   # if 95% LB is >0 : save the model (otherwise it's pointless)\n",
    "            print('####   GOOD MODEL FOUND - SAVING IT   ####' ) \n",
    "\n",
    "            # best pol \n",
    "            if(bestpol < recqvi[modl,23]):\n",
    "                print('Best policy was replaced => 95% LB is ',recqvi[modl,23])\n",
    "                bestpol = recqvi[modl,23] \n",
    "                \n",
    "                # save to pickle \n",
    "                with open('bestpol.pkl', 'wb') as file:\n",
    "                    pickle.dump(modl,file)\n",
    "                    pickle.dump(Qon,file)\n",
    "                    pickle.dump(physpol,file)\n",
    "                    pickle.dump(transitionr,file) \n",
    "                    pickle.dump(transitionr2,file)\n",
    "                    pickle.dump(R,file)\n",
    "                    pickle.dump(C,file)\n",
    "                    pickle.dump(train,file)\n",
    "                    pickle.dump(qldata3train,file)\n",
    "                    pickle.dump(qldata3test,file) \n",
    "\n",
    "\n",
    "    recqvi=recqvi[:modl+1,:]\n",
    "\n",
    "    # save to pickle for visualization \n",
    "    with open('step_5_start.pkl', 'wb') as file:\n",
    "        pickle.dump(MIMICzs,file)\n",
    "        pickle.dump(actionbloc,file)\n",
    "        pickle.dump(reformat5,file)\n",
    "        pickle.dump(recqvi,file)\n",
    "\n",
    "    # save recqvi in csv format \n",
    "    np.savetxt('recqvi.csv',recqvi,delimiter=',') \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}