combined_RF oct.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# build predictive model of genes -> STP\n",
    "\n",
    "\n",
    "1) Use ABA VISp,Hipp gene expression and STP data to fit RF for STP. \n",
    "\n",
    "2) test predictions - crossvalidation, test on hippocampal data, test on exc. neurons data?\n",
    "\n",
    "3) feature selection - compare datasets? - compare with Paul analysis\n",
    "\n",
    "4) stp factors genes - find better factors using interactions networks?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) fit RF ABA ge -> STP : Hippocampus, VISp\n",
    "\n",
    "a. select gene sets\n",
    "\n",
    "b. load gene expression data Hipp, Ent \n",
    "\n",
    "c. hippocampus STP data \n",
    "\n",
    "d. combine STP and gene expression data for Hippocampus\n",
    "\n",
    "e.load gene expression data VISp\n",
    "\n",
    "f.  coretex STP data\n",
    "\n",
    "g. combine STP and gene expression data for VISp cortex\n",
    "\n",
    "h. fit RF from sklit-learn\n",
    "\n",
    "f. test predictions :  cross validation - part of genes, part of STP labels, genes from different cortical areas, genes from hippocampus (Harris 2018)\n",
    "\n",
    "g. compare gene sets\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "do_laptop=1   # Laptop or desctop version\n",
    "\n",
    "\n",
    "do_add_scvi_latent_factors = 1\n",
    "do_add_scvi_imputed_genes = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# a. select gene sets - take Paul 2017 best pre-synaptic + CAMs\n",
    "# Paul 2017\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_prot_list(df_stp_ds, name):\n",
    "    pre_main=df_stp_ds.loc[:,['Protein','Gene rat']]\n",
    "    pre_main_nm = pre_main.loc[pre_main.loc[:,'Protein'].str.contains(':',na=False),'Protein']\n",
    "\n",
    "\n",
    "    pre_main_nm=pre_main_nm.reset_index()\n",
    "    i = pre_main_nm.loc[pre_main_nm.loc[:,'Protein']==name,:]\n",
    "    i1 = i.loc[:,'index']\n",
    "    i2 = pre_main_nm.loc[i.index[0]+1,'index']\n",
    "\n",
    "    Cas = df_stp_ds.loc[np.arange(i1.values[0]+1,i2),'Gene rat']\n",
    "    Cas=Cas.loc[Cas.loc[:].isna()==False]\n",
    "    Cas = Cas.str.split('\\W+',expand=True).iloc[:,0].values.tolist()\n",
    "    return Cas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d8 = '~/Documents/references/transcriptomes to STP/'\n",
    "if do_laptop:\n",
    "    d8 = '/Users/stepaniu/Documents/jan_2020/'\n",
    "\n",
    "df_stp_ds = pd.read_excel(d8+'Datasets_STP_hippocampus.xlsx', sheet_name='Sheet8')\n",
    "\n",
    "Cas = get_prot_list(df_stp_ds, 'Calcium sensors:')\n",
    "AZ2 = get_prot_list(df_stp_ds, 'active zone:')\n",
    "\n",
    "print(Cas)\n",
    "print(AZ2)\n",
    "\n",
    "\n",
    "d6=''\n",
    "if do_laptop:\n",
    "    d6 = '/Users/stepaniu/Documents/references/lists_of_proteins_2/Analysis_of_proteins/genetic_types/'\n",
    "\n",
    "fn6 = 'Paul_2017.xlsx'\n",
    "df_Paul_gs=pd.read_excel(d6 + fn6,header=2)\n",
    "\n",
    "d7 = '~/Documents/references/transcriptomes to STP/transcriptomes/Linnerson_2018/'\n",
    "fn7 = 'GO: synaptic vesicle cycle, mouse.xlsx'\n",
    "if do_laptop:\n",
    "    d7 = '/Users/stepaniu/Documents/jan_2020/'\n",
    "    fn7 = 'GO_ synaptic vesicle cycle, mouse.xlsx'\n",
    "df_go_priming = pd.read_excel(d7+fn7)\n",
    "SV_cycle=df_go_priming.loc[:,'Gene/product label'].value_counts().index.tolist()\n",
    "\n",
    "\n",
    "#print(df_Paul_gs.head())\n",
    "\n",
    "#select presynaptic :\n",
    "\n",
    "\n",
    "pre = ['Netrin-Unc5-Slit-Robo',\n",
    "'Nrxn+Nlgn+Nphx+Dystroglycan+Cbln',\n",
    "'neuronal_IgCAMs',\n",
    "'Contactins',\n",
    "'Semaphorin+Plexin',\n",
    "'LRR+Slitrk+Elfn+Lphn',\n",
    "'Eph+EphR',\n",
    "'Calcium channel regulatory subunits',\n",
    "'Calcium channel poreforming subunits',\n",
    "'Calcium channel gamma/TARP subunits',\n",
    "'Orphan GPCRs (based off HGNC)',\n",
    "'Metabotropic neurotransmitter receptors',\n",
    "'Neuropeptide receptors',\n",
    "'Trimeric G_protein_alpha',\n",
    "'Trimeric G_protein_beta',\n",
    "'Trimeric G_protein_gamma',\n",
    "'Rho-GEF',\n",
    "'Dock-GEF',\n",
    "'Dock + Rho-GEF',\n",
    "'Complexin-Syt',\n",
    "'SNAP-complex',\n",
    "'Syntaxin',\n",
    "'RIM-proteins/Active zone',\n",
    "'bHLH',\n",
    "'Nuclear receptors with C4 zinc fingers',\n",
    "'Homeo domain factors',\n",
    "'High-mobility group (HMG) domain factors',\n",
    "'C2H2 zinc finger factors',\n",
    "'bZIP',\n",
    "'ARID domain factors',\n",
    "'Runt domain factors',\n",
    "'Other C4 zinc finger-type',\n",
    "'C2HC zinc finger factors',\n",
    "'SMAD/NF-1 DNA-binding domain factors',\n",
    "'TATA-binding proteins',\n",
    "'Rel homology region (RHR) factors',\n",
    "'Tryptophan cluster factors',\n",
    "'MADS box factors',\n",
    "'Cold-shock domain factors',\n",
    "'CXXC zinc finger factors',\n",
    "'Leucine-rich repeat flightless-interacting proteins',\n",
    "'STAT domain factors',\n",
    "'Grainyhead domain factors'\n",
    "]\n",
    "\n",
    "#pre = ['Calcium channel regulatory subunits', 'Netrin-Unc5-Slit-Robo',\n",
    "#       'Nrxn+Nlgn+Nphx+Dystroglycan+Cbln', 'neuronal_IgCAMs', 'Contactins',\n",
    "#       'Metabotropic neurotransmitter receptors',      'Dock + Rho-GEF', 'Complexin-Syt']\n",
    "\n",
    "#pre = ['Syntaxin', 'RIM-proteins/Active zone', 'Dock-GEF','SNAP-complex' ]\n",
    "\n",
    "#pre = ['RIM-proteins/Active zone'] \n",
    "\n",
    "#pre = []\n",
    "\n",
    "\n",
    "Elfn = ['Elfn1', 'Elfn2']\n",
    "\n",
    "pregs = []\n",
    "for i in pre:\n",
    "    in_i = (df_Paul_gs.loc[:,'Gene Set']==i)\n",
    "    n_in_i = df_Paul_gs.loc[in_i,'Number of genes'].astype(int).values[0]\n",
    "    pri = df_Paul_gs.loc[in_i,df_Paul_gs.columns[4:4+n_in_i]]\n",
    "    pregs = pregs + pri.values.tolist()[0]\n",
    "    \n",
    "#pregs = pregs +  Elfn  \n",
    "#pregs = pregs + SV_cycle\n",
    "pregs = pregs +Cas + AZ2 \n",
    "#pregs = pregs +Cas \n",
    "\n",
    "## best gs\n",
    "#pregs = best_pre\n",
    "\n",
    "post = ['Netrin-Unc5-Slit-Robo',\n",
    "'Nrxn+Nlgn+Nphx+Dystroglycan+Cbln',\n",
    "'neuronal_IgCAMs',\n",
    "'Contactins',\n",
    "'Semaphorin+Plexin',\n",
    "'LRR+Slitrk+Elfn+Lphn',\n",
    "'Eph+EphR',\n",
    "'Calcium channel regulatory subunits',\n",
    "'Calcium channel poreforming subunits',\n",
    "'Calcium channel gamma/TARP subunits',\n",
    "'AMPA+TARP',\n",
    "'Orphan GPCRs (based off HGNC)',\n",
    "'Metabotropic neurotransmitter receptors',\n",
    "'Neuropeptide receptors']\n",
    "\n",
    "\n",
    "# post = ['Netrin-Unc5-Slit-Robo',\n",
    "# 'Nrxn+Nlgn+Nphx+Dystroglycan+Cbln',\n",
    "# 'neuronal_IgCAMs',\n",
    "# 'Contactins',\n",
    "# 'Semaphorin+Plexin',\n",
    "# 'LRR+Slitrk+Elfn+Lphn',\n",
    "# 'Eph+EphR']\n",
    "\n",
    "#post = ['Netrin-Unc5-Slit-Robo',\n",
    "#        'Nrxn+Nlgn+Nphx+Dystroglycan+Cbln', 'neuronal_IgCAMs', 'Contactins']\n",
    "\n",
    "\n",
    "#post = [ 'Nrxn+Nlgn+Nphx+Dystroglycan+Cbln']\n",
    "\n",
    "#post=[]\n",
    "\n",
    "\n",
    "postgs = []\n",
    "for i in post:\n",
    "    in_i = (df_Paul_gs.loc[:,'Gene Set']==i)\n",
    "    n_in_i = df_Paul_gs.loc[in_i,'Number of genes'].astype(int).values[0]\n",
    "    posti = df_Paul_gs.loc[in_i,df_Paul_gs.columns[4:4+n_in_i]]\n",
    "    postgs = postgs + posti.values.tolist()[0]\n",
    "    \n",
    "postgs = postgs +  Elfn    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pregs = pregs[0:14]\n",
    "postgs = postgs[0:9]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(pregs))\n",
    "print(len(postgs))\n",
    "postgs7='post__' + pd.DataFrame(postgs)\n",
    "pregs7='pre__' + pd.DataFrame(pregs)\n",
    "gs7 = pd.concat([pregs7, postgs7],axis=0).reset_index(drop=True)\n",
    "gs7.to_excel('gene_set_names.xlsx')\n",
    "gs7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## SET FLAG : do_select_best_genes_from_previous_RF_run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### fi2 = results of previous search\n",
    "#fi2.to_excel('best_genes_1691_RF.xlsx')\n",
    "#fi2.to_excel('best_genes_280_RF_2.xlsx')\n",
    "#fi2.to_excel('best_genes_280_RF_hip.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "do_select_best_genes_from_previous_RF_run=0\n",
    "if do_select_best_genes_from_previous_RF_run==1:\n",
    "    # fi2 = results of previous search\n",
    "    fi2  = pd.read_excel('best_genes_1691_RF.xlsx')\n",
    "    fi2  = pd.read_excel('best_genes_280_RF.xlsx')\n",
    "    fi2  = pd.read_excel('best_genes_280_RF.xlsx')\n",
    "    fi2  = pd.read_excel('best_genes_280_RF_hip.xlsx')\n",
    "\n",
    "    fi2 = pd.read_excel('best_genes_hipp.xlsx')\n",
    "    fi3 = pd.read_excel('best_genes_visp.xlsx')\n",
    "    #\n",
    "    nbest = 10\n",
    "    bestn = fi2.iloc[0:nbest,:]\n",
    "    #best_pre = bestn.index[bestn.loc[:,'compartment']=='pre'].values.tolist()\n",
    "    #best_post = bestn.index[bestn.loc[:,'compartment']=='post'].values.tolist()\n",
    "    \n",
    "    best_pre = bestn.loc[bestn.loc[:,'compartment']=='pre','Unnamed: 0'].values.tolist()\n",
    "    best_post = bestn.loc[bestn.loc[:,'compartment']=='post','Unnamed: 0'].values.tolist()\n",
    "    \n",
    "    print('best presynaptic genes')\n",
    "    print(best_pre)\n",
    "    print('best postsynaptic genes')\n",
    "    print(best_post)\n",
    "    \n",
    "    ## best gs\n",
    "    ## best gs\n",
    "    pregs = best_pre\n",
    "    postgs = best_post\n",
    "\n",
    "print('postsynaptic')\n",
    "print(postgs)\n",
    "print(len(postgs))\n",
    "\n",
    "print('presynaptic')\n",
    "print(pregs)\n",
    "print(len(pregs))\n",
    "print(len(postgs))\n",
    "\n",
    "pregs =np.unique(pregs).tolist()\n",
    "postgs =np.unique(postgs).tolist()\n",
    "\n",
    "import pickle\n",
    "RF_file = open('RF_290_hipp_gs.pickle', mode='wb')\n",
    "pickle.dump([pregs, postgs],RF_file)\n",
    "\n",
    "do_visp_features=0\n",
    "if do_visp_features==1:\n",
    "    #pickle_in = open('RF_290_hipp.pickle',\"rb\")\n",
    "    RF_file = open('RF_290_visp_gs.pickle', mode='rb')\n",
    "    \n",
    "    gs_visp = pickle.load(RF_file)\n",
    "    pregs =gs_visp[0] \n",
    "    postgs=gs_visp[1]\n",
    "    \n",
    "    \n",
    "##from mutual_info import entropy\n",
    "##e = entropy(X, k=1)\n",
    "#np.finfo(float).eps\n",
    "#print(len(pregs))\n",
    "#print(np.unique(pregs).shape)\n",
    "#print(len(postgs))\n",
    "#print(np.unique(postgs).shape)\n",
    "#len(pregs)+len(postgs)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# STP modeling functions: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PV->AAC ???\n",
    "Dt=50\n",
    "tF=10\n",
    "tD=410\n",
    "p0=0.23\n",
    "p=p0\n",
    "n = (1-p)\n",
    "n=1-(1-n)*np.exp(-Dt/tD)\n",
    "print(n)\n",
    "p=p+(1-p)*p0\n",
    "p=p0 +(p -p0)*np.exp((-Dt/tF))\n",
    "print(p)\n",
    "n*p/p0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CCK->CCK ???\n",
    "Dt=50\n",
    "tF=1542\n",
    "tD=115\n",
    "p0=0.11\n",
    "p=p0\n",
    "n = (1-p)\n",
    "n=1-(1-n)*np.exp(-Dt/tD)\n",
    "print(n)\n",
    "p=p+(1-p)*p0\n",
    "p=p0 +(p -p0)*np.exp((-Dt/tF))\n",
    "print(p)\n",
    "n*p/p0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PV->AAC ???\n",
    "Dt=50\n",
    "tF=1.6\n",
    "tD=930\n",
    "p0=0.26\n",
    "p=p0\n",
    "n = (1-p)\n",
    "n=1-(1-n)*np.exp(-Dt/tD)\n",
    "print(n)\n",
    "p=p+(1-p)*p0\n",
    "p=p0 +(p -p0)*np.exp((-Dt/tF))\n",
    "print(p)\n",
    "n*p/p0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ivy->PC ???\n",
    "Dt=50\n",
    "tF=62\n",
    "tD=144\n",
    "p0=0.32\n",
    "p=p0\n",
    "n = (1-p)\n",
    "n=1-(1-n)*np.exp(-Dt/tD)\n",
    "print(n)\n",
    "p=p+(1-p)*p0\n",
    "p=p0 +(p -p0)*np.exp((-Dt/tF))\n",
    "print(p)\n",
    "n*p/p0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CCK_DTI->PC ???\n",
    "Dt=50\n",
    "tF=14\n",
    "tD=185\n",
    "p0=0.15\n",
    "p=p0\n",
    "n = (1-p)\n",
    "n=1-(1-n)*np.exp(-Dt/tD)\n",
    "print(n)\n",
    "p=p+(1-p)*p0\n",
    "p=p0 +(p -p0)*np.exp((-Dt/tF))\n",
    "print(p)\n",
    "n*p/p0"
   ]
  },
  {
   "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": [
    "def  STP_sim(ge_data, T, init_state=None ):\n",
    "    # transform labels from TM to An:A1\n",
    "    #f = 20 # Hz\n",
    "    #N = 3\n",
    "    #T = np.arange(N)*1000/f\n",
    "\n",
    "    N    = len(T)\n",
    "    nc   = len(ge_data.index)\n",
    "\n",
    "    x_lower = np.array([1,       0.01,    1,      0.1,     1])\n",
    "    x_upper = np.array([10000,   1,      10000,    10,     1])\n",
    "    stp_ns =          ['tF',    'p0',    'tD',   'dp/p0', 'A']\n",
    "    for jj in range(len(stp_ns)):\n",
    "        nsj = stp_ns[jj]\n",
    "        ge_data.loc[:,nsj] = np.maximum(ge_data.loc[:,nsj].values,x_lower[jj])\n",
    "        ge_data.loc[:,nsj] = np.minimum(ge_data.loc[:,nsj].values,x_upper[jj])\n",
    "        if nsj=='dp/p0':\n",
    "            p0   = ge_data.loc[:,'p0'].values\n",
    "            dpp0 = ge_data.loc[:,'dp/p0'].values\n",
    "            dp = p0*dpp0\n",
    "            dp = np.minimum(dp,1)\n",
    "            ge_data.loc[:,nsj] = dp/(p0 +np.finfo(float).eps)\n",
    "\n",
    "    dpp0 = ge_data.loc[:,'dp/p0'].values\n",
    "    p0   = ge_data.loc[:,'p0'].values\n",
    "    tF   = ge_data.loc[:,'tF'].values\n",
    "    tD   = ge_data.loc[:,'tD'].values\n",
    "    A    = 1 + 0*ge_data.loc[:,'A'].values # simplify A\n",
    "\n",
    "    As = np.zeros((nc,N))\n",
    "    n = np.zeros((nc,))\n",
    "    p = np.zeros((nc,))\n",
    "    ns2 = np.zeros((nc,N))\n",
    "    ps2 = np.zeros((nc,N))\n",
    "\n",
    "    i=0\n",
    "    \n",
    "    if init_state is None :\n",
    "        n[:] = 1\n",
    "        p[:] = p0\n",
    "    else:\n",
    "        n = init_state[0]\n",
    "        p = init_state[1]\n",
    "\n",
    "\n",
    "    As[:,i] = A*n*p\n",
    "    \n",
    "    n = n*(1-p)\n",
    "    p = p + dp*(1-p)\n",
    "    \n",
    "    ns2[:,i]=n\n",
    "    ps2[:,i]=p\n",
    "\n",
    "    for i in range(1,N):\n",
    "        Dt=T[i]-T[i-1]\n",
    "        #n = 1 - (1 - (n -p*n))*np.exp((-Dt/tD).astype(float))\n",
    "        #p=p0 -(p0-(p + dpp0*p0*(1-p)))*np.exp((-Dt/tF).astype(float))\n",
    "        #As[:,i]=A*n*p\n",
    "        #ns2[:,i]=n\n",
    "        #ps2[:,i]=p\n",
    "\n",
    "        \n",
    "        n = 1 - (1 - n )*np.exp((-Dt/tD ).astype(float))\n",
    "        p=p0 +(p -p0)*np.exp((-Dt/tF).astype(float))\n",
    "            \n",
    "\n",
    "        As[:,i]=A*n*p\n",
    "       \n",
    "        n = n*(1-p)\n",
    "        p = p + dp*(1-p)\n",
    "        \n",
    "        ns2[:,i]=n\n",
    "        ps2[:,i]=p\n",
    "\n",
    "    \n",
    "\n",
    "    #aa = [As, ns2, ps2]\n",
    "    \n",
    "    return As, ns2, ps2, dpp0, p0, tF, tD, A"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "def  STP_sim(ge_data, T ):\n",
    "    # transform labels from TM to An:A1\n",
    "    #f = 20 # Hz\n",
    "    #N = 3\n",
    "    #T = np.arange(N)*1000/f\n",
    "\n",
    "    N    = len(T)\n",
    "    nc   = len(ge_data.index)\n",
    "\n",
    "    x_lower = np.array([1,       0.01,    1,      0.1,     1])\n",
    "    x_upper = np.array([10000,   1,      10000,    10,     1])\n",
    "    stp_ns =          ['tF',    'p0',    'tD',   'dp/p0', 'A']\n",
    "    for jj in range(len(stp_ns)):\n",
    "        nsj = stp_ns[jj]\n",
    "        ge_data.loc[:,nsj] = np.maximum(ge_data.loc[:,nsj].values,x_lower[jj])\n",
    "        ge_data.loc[:,nsj] = np.minimum(ge_data.loc[:,nsj].values,x_upper[jj])\n",
    "        if nsj=='dp/p0':\n",
    "            p0   = ge_data.loc[:,'p0'].values\n",
    "            dpp0 = ge_data.loc[:,'dp/p0'].values\n",
    "            dp = p0*dpp0\n",
    "            dp = np.minimum(dp,1)\n",
    "            ge_data.loc[:,nsj] = dp/(p0 +np.finfo(float).eps)\n",
    "\n",
    "    dpp0 = ge_data.loc[:,'dp/p0'].values\n",
    "    p0   = ge_data.loc[:,'p0'].values\n",
    "    tF   = ge_data.loc[:,'tF'].values\n",
    "    tD   = ge_data.loc[:,'tD'].values\n",
    "    A    = ge_data.loc[:,'A'].values\n",
    "\n",
    "    As = np.zeros((nc,N))\n",
    "    n = np.zeros((nc,))\n",
    "    p = np.zeros((nc,))\n",
    "    ns2 = np.zeros((nc,N))\n",
    "    ps2 = np.zeros((nc,N))\n",
    "\n",
    "    i=0\n",
    "    n[:] = 1\n",
    "    p[:] = p0\n",
    "    As[:,i] = A*n*p\n",
    "    ns2[:,i]=n\n",
    "    ps2[:,i]=p\n",
    "\n",
    "    for i in range(1,N):\n",
    "        Dt=T[i]-T[i-1]\n",
    "        n = 1 - (1 - (n -p*n))*np.exp((-Dt/tD).astype(float))\n",
    "        p=p0 -(p0-(p + dpp0*p0*(1-p)))*np.exp((-Dt/tF).astype(float))\n",
    "        As[:,i]=A*n*p\n",
    "        ns2[:,i]=n\n",
    "        ps2[:,i]=p\n",
    "\n",
    "    #aa = [As, ns2, ps2]\n",
    "    \n",
    "    return As, ns2, ps2, dpp0, p0, tF, tD, A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def  STP_sim2(x, T, init_state=None, model_type='tm5' ):\n",
    "\n",
    "    N    = len(T)\n",
    "    tF      = x[0] #.astype(float)\n",
    "    p00     = x[1]\n",
    "    tD      = x[2] #.astype(float)\n",
    "    dp      = x[3]\n",
    "    A       = 1 #x[4] # simplify A\n",
    "    \n",
    "    #breakpoint()\n",
    "    mod_fdr2=False\n",
    "    if model_type=='tm5_fdr2':  # should be :check freq. dependent recovery\n",
    "        tDmin     = x[5]\n",
    "        dd        = x[6]\n",
    "        t_FDR     = x[7]\n",
    "        mod_fdr2=True\n",
    "        tDmax  = tD\n",
    "        itDmin = 1/tDmin\n",
    "        itDmax = 1/tDmax\n",
    "        #breakpoint()\n",
    "        \n",
    "    mod_smr=False\n",
    "    if model_type=='tm5_smr':  # should be :check freq. dependent recovery\n",
    "        t_SMR   = x[8]\n",
    "        dp0     = x[9]\n",
    "        mod_smr=True\n",
    "        #p00  = p00\n",
    "\n",
    "    As = np.zeros((N,))\n",
    "    state = np.zeros((N,4))\n",
    "\n",
    "   \n",
    "    if init_state is None :\n",
    "        n = 1\n",
    "        p0=p00\n",
    "        p = p0\n",
    "        d = 0\n",
    "    else:\n",
    "        n = init_state[0]\n",
    "        p = init_state[1]\n",
    "        d = init_state[2]\n",
    "        p0= init_state[3]\n",
    "\n",
    "    \n",
    "    for i in range(0,N):\n",
    "        if i==0:\n",
    "            Dt = T[i]\n",
    "        else:\n",
    "            Dt = T[i]-T[i-1]\n",
    "        \n",
    "        if mod_fdr2:\n",
    "            d0=d\n",
    "            d = d*np.exp(-Dt/t_FDR) \n",
    "            n = 1 - (1 - n )*np.exp(-Dt*itDmax -(itDmin -itDmax)*t_FDR*(d0-d))\n",
    "        else:\n",
    "            #print(x,tD)\n",
    "            n = 1 - (1 - n )*np.exp(-Dt/tD )\n",
    "            \n",
    "        if mod_smr:\n",
    "            p01=p0\n",
    "            p0=p00 + (p0 -p00)*np.exp(-Dt/t_SMR)\n",
    "            p=p0 +(p -p01)*np.exp(-Dt/tF)\n",
    "        else:\n",
    "            p=p0 +(p -p0)*np.exp(-Dt/tF)\n",
    "            \n",
    "\n",
    "        As[i]=A*n*p\n",
    "       \n",
    "        n = n*(1-p)\n",
    "        p = p + dp*(1-p)\n",
    "        if mod_fdr2:\n",
    "            d  = d + dd*(1-d) \n",
    "        if mod_smr:\n",
    "            p0  = p0 - dp0*p0    \n",
    " \n",
    "        state[i] = [n,p,d,p0]\n",
    "\n",
    "    #return As, ns2, ps2, dpp0, p0, tF, tD, A\n",
    "    return As, state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def STP_sim_complex(ge_data,l_pre_post2,stp_aba_names):\n",
    "    # transform labels from TM to An:A1\n",
    "    #fs = [20, 50, 10] # Hz\n",
    "    fs = [20, 50, 10]\n",
    "    N = 5\n",
    "\n",
    "    #Trec = [250, 500, 1000]\n",
    "    Trec =[250, 1000]\n",
    "    DT0 = 25000\n",
    "    if l_pre_post2<ge_data.shape[0]:\n",
    "        xs  =ge_data.iloc[l_pre_post2:,:].loc[:,stp_aba_names].values\n",
    "        xs  =np.delete(xs, [5,6],axis=1)\n",
    "        As2=np.zeros((xs.shape[0],0))\n",
    "    \n",
    "    \n",
    "    As=np.zeros((l_pre_post2,0))\n",
    "    if 1:\n",
    "        raz=0\n",
    "        \n",
    "        for f in fs:\n",
    "            raz=raz+1\n",
    "            T = np.arange(N)*1000/f\n",
    "\n",
    "            if l_pre_post2>0:\n",
    "                if raz==1:\n",
    "                    print('model type = 4 parameters TM')\n",
    "                Asf, ns, ps, dpp0, p0, tF, tD, A = STP_sim(ge_data.iloc[0:l_pre_post2,:],T)\n",
    "                As = np.concatenate([As,Asf],axis=1)\n",
    "\n",
    "                for ri in range(len(Trec)):\n",
    "                    Asr, nsr, psr, dpp0, p0, tF, tD, A = STP_sim(ge_data.iloc[0:l_pre_post2,:],[Trec[ri]],init_state=[ns[:,-1],ps[:,-1]])\n",
    "                    As = np.concatenate([As,Asr],axis=1)\n",
    "\n",
    "            if l_pre_post2<ge_data.shape[0]:\n",
    "                if raz==1:\n",
    "                    print('model type =  TM 5 parameters , smr')\n",
    "                    \n",
    "                As2f = np.zeros((xs.shape[0],N+len(Trec)))\n",
    "                for i2 in range(xs.shape[0]):\n",
    "                    #as2, sts2 = STP_sim2(xs[i2,:],np.arange(8)*50,model_type = 'tm5') # preconditioning series???\n",
    "                    #as2, sts2 = STP_sim2(xs[i2,:], T+DT0, init_state=sts2[-1], model_type='tm5' ) \n",
    "\n",
    "                    as2, sts2 = STP_sim2(xs[i2,:],T,model_type = 'tm5_smr')\n",
    "\n",
    "                    As2f[i2,0:N] = as2\n",
    "                    for ri in range(len(Trec)):\n",
    "                        as2r, sts2r = STP_sim2(xs[i2,:], [Trec[ri]], init_state=sts2[-1], model_type='tm5_smr' ) \n",
    "                        As2f[i2,N+ri] = as2r\n",
    "\n",
    "                As2 = np.concatenate([As2,As2f],axis=1)\n",
    "\n",
    "        if l_pre_post2<ge_data.shape[0]:\n",
    "            As2 = As2/As2[:,0].reshape((-1,1))\n",
    "        else:\n",
    "            As2  = np.zeros((0,As.shape[1]))\n",
    "        \n",
    "        \n",
    "        if l_pre_post2>0:\n",
    "            As = As/As[:,0].reshape((-1,1))\n",
    "            As = np.concatenate([As,As2],axis=0)\n",
    "        else:\n",
    "            As =As2\n",
    "    else:\n",
    "        As, ns, ps, dpp0, p0, tF, tD, A = STP_sim(ge_data,T)\n",
    "\n",
    "\n",
    "    #import matplotlib.pyplot as plt    \n",
    "    #plt.plot(As2[400:410,:].transpose(),'o-')\n",
    "    \n",
    "    return As"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#df_stp.loc[:,'name'].str.contains('Jaing_2015_PC')\n",
    "df_stp=pd.read_excel(d5+fn5)\n",
    "#df_stp.loc[df_stp.loc[:,'stp_comment'].isna() ,'stp_comment']=''\n",
    "\n",
    "############   \n",
    "do_averagig_stp =0\n",
    "stp_TM_names=['tF', 'p0','tD','dp/p0','A']\n",
    "stp_TM_start = 187\n",
    "ncol = len(df_stp.columns)\n",
    "df_tm_name = pd.DataFrame(np.zeros(ncol).reshape((1,ncol)), columns =df_stp.columns )\n",
    "df_stp = pd.concat([df_tm_name, df_stp],axis=0).reset_index(drop=True)\n",
    "\n",
    "for i,nm in enumerate(stp_TM_names):\n",
    "    ii = stp_TM_start + np.arange(0,ncol-i-stp_TM_start,len(stp_TM_names))+i\n",
    "    df_stp.iloc[0,ii] = nm\n",
    "\n",
    "#df_stp.loc[df_stp.loc[:,'name'].str.contains('Jaing_2015_PC')!=True,:].reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## \n",
    "##\n",
    "## Hippocampus gene expression\n",
    "##\n",
    "## "
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "from scvi.dataset import DownloadableAnnDataset\n",
    "d4='~/Documents/references/transcriptomes to STP/transcriptomes/ABA/mouse_VISp_gene_expression_matrices_2018-06-14/'\n",
    "#df_aba_ad.write('mouse_VISp_2018-06-14_exon-matrix_GABAergic.h5ad')\n",
    "aba_vis_in = DownloadableAnnDataset(d4+\"data/mouse_VISp_2018-06-14_exon-matrix_GABAergic.h5ad\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scvi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "print(os.getcwd())\n",
    "#os.listdir()\n",
    "\n",
    "if do_laptop:\n",
    "    from scvi.dataset import LoomDataset, CsvDataset, Dataset10X, AnnDataset\n",
    "else:\n",
    "    from scvi.dataset import DownloadableAnnDataset\n",
    "\n",
    "d_aba  ='/home/stepaniu/Documents/references/transcriptomes to STP/transcriptomes/ABA_2019/'\n",
    "if do_laptop:\n",
    "    d_aba  ='/Users/stepaniu/Documents/jan_2020/ABA_2019/'\n",
    "\n",
    "df_aba_g = pd.read_csv(d_aba + 'medians.csv')\n",
    "df_aba_c = pd.read_csv(d_aba + 'sample_annotations.csv')\n",
    "\n",
    "\n",
    "d7 = '/home/stepaniu/Documents/references/transcriptomes to STP/transcriptomes/ABA_2019/'\n",
    "if do_laptop:\n",
    "    d7  ='/Users/stepaniu/Documents/jan_2020/ABA_2019/'\n",
    "\n",
    "if do_laptop:    \n",
    "    aba_vis_in = AnnDataset(\"ABA_2019_transcriptome.h5ad\", new_n_genes = 45768,\n",
    "                             save_path = d7) \n",
    "else:     \n",
    "    aba_vis_in = DownloadableAnnDataset(d7 +\"ABA_2019_transcriptome.h5ad\") \n",
    "\n",
    "obs = pd.read_excel(d7 +'ABA_2019_obs.xlsx') # samples\n",
    "var = pd.read_excel(d7 + 'ABA_2019_var.xlsx') # gene names\n",
    "#obs = aba_vis_in._obs # samples\n",
    "#var = aba_vis_in._var # gene names\n",
    "\n",
    "in_gaba = df_aba_c.loc[df_aba_c.loc[:,'sample_name'].isin(obs.loc[:,'samples']),['sample_name', 'class_label']]\n",
    "in_gaba = in_gaba.reset_index().set_index('sample_name',drop=True)\n",
    "in_gaba = in_gaba.loc[obs.loc[:,'samples'],:]\n",
    "in_gaba = in_gaba.reset_index().reset_index().set_index('sample_name',drop=True)\n",
    "in_gaba = in_gaba.loc[in_gaba.loc[:,'class_label']=='GABAergic',:].reset_index().set_index('level_0',drop=True)\n",
    "\n",
    "#aba_vis_in.gene_names = var.iloc[:,0].values\n",
    "aba_vis_in.gene_names = var.loc[:,'genes'].values\n",
    "vis_dat = aba_vis_in\n",
    "aba_vis_in = []\n",
    "\n",
    "in_ds = df_aba_c.loc[df_aba_c.loc[:,'sample_name'].isin(obs.loc[:,'samples']),['sample_name', 'class_label']]\n",
    "in_ds = in_ds.reset_index().set_index('sample_name',drop=True)\n",
    "in_ds = in_ds.loc[obs.loc[:,'samples'],:]\n",
    "in_ds = in_ds.reset_index().reset_index().set_index('sample_name',drop=True)\n",
    "#in_ds\n",
    "\n",
    "df_aba_vis_in_c = df_aba_c.loc[in_ds.loc[:,'index'],:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### add scVI latent factors and imputed genes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "do_replace_to_scVI_imputed=1      # replace raw gene expression data with their scvi imputation\n",
    "do_add_scvi_latent_factors=1      # add scvi latent factors data\n",
    "do_use_all_scvi_genes=0           # use all scvi genes expression in addition to pregs, postgs\n",
    "\n",
    "gs_name = '_gs5188' #'_gs1512' #'_gs5188' #'_gs1512' #'_gs15656' #'_gs5188' #'_gs1512' #'_gs219'\n",
    "if (gs_name == '_gs1512')|(gs_name == '_gs5188')|(gs_name == '_gs15656'):\n",
    "    nlf = 20\n",
    "elif gs_name == '_gs219':\n",
    "    nlf = 10\n",
    "    \n",
    "if (do_replace_to_scVI_imputed==1)|(do_add_scvi_latent_factors==1):\n",
    "    scVI_all = pd.read_hdf('scVI_latent_factors_and_imputation_aba2019'+gs_name+'.hdf',key='data') # this is 290 genes dataset\n",
    "    \n",
    "    scVI_all = scVI_all.set_index('sample_name') # scvi imputad genes and latent factors for hipp., enth. and vis. cortex\n",
    "\n",
    "    \n",
    "    gncolumns = scVI_all.columns[nlf+1:] # gene names columns in scvi data\n",
    "    lfcolumns = scVI_all.columns[:nlf]     # latent factors columns in scvi data\n",
    "\n",
    "    gnh = vis_dat.gene_names               # hippocampal vis_dat gene names\n",
    "    smh = vis_dat.obs                      # hippocampal vis_dat samples names\n",
    "\n",
    "    smh2 = smh.reset_index().set_index('samples')  # hippocampal samples  and their indexes in vis_dat X\n",
    "\n",
    "    gnscvi = pd.DataFrame(gncolumns)   # df from scvi genenames\n",
    "\n",
    "    gnh2 = pd.DataFrame(gnh).reset_index().set_index(0)   \n",
    "    ignh2 = gnh2.loc[gnscvi.loc[:,0],:].loc[:,'index'].values  # indeces of scvi gene names in vis_dat gene names\n",
    "\n",
    "    if (do_replace_to_scVI_imputed==1):\n",
    "        vis_dat._X[:,ignh2] = scVI_all.loc[smh2.index,gncolumns].values  # modified vis_dat matrix (imputed part of genes )\n",
    "    if (do_add_scvi_latent_factors==1):\n",
    "        scVI_lf_h = scVI_all.loc[smh2.index,lfcolumns]                   # latent factors for hippocampal samples\n",
    "        vis_dat._X = np.concatenate([vis_dat._X, scVI_lf_h.values],axis=1)\n",
    "        \n",
    "        lf_scvi_columns = lfcolumns.astype(str)\n",
    "        vis_dat.gene_names = np.concatenate([vis_dat.gene_names, lf_scvi_columns])\n",
    "        pregs = pregs + lf_scvi_columns.tolist()\n",
    "        postgs = postgs + lf_scvi_columns.tolist()\n",
    "    if (do_use_all_scvi_genes==1): \n",
    "        #pregs = list(set(pregs + gncolumns.tolist()))\n",
    "        #postgs = list(set(postgs + gncolumns.tolist()))\n",
    "        pregs = gncolumns.tolist()+ lf_scvi_columns.tolist()\n",
    "        postgs = gncolumns.tolist()+ lf_scvi_columns.tolist()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vis_dat.gene_names.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(pregs))\n",
    "len(postgs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##\n",
    "##\n",
    "##  HIPPOCAMPUS STP\n",
    "##\n",
    "##"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "do_laptop=0\n",
    "if do_laptop:\n",
    "    d5 = '~/Documents/references/lists_of_proteins_2/Analysis_of_proteins/genetic_types/transcriptomic_interneurons/additional cortical STP data/'\n",
    "else:\n",
    "    d5 = '/home/stepaniu/Documents/references/transcriptomes to STP/scVI_data/additional cortical STP data/'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled.xlsx'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled_UDF2.xlsx'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled_UDF3.xlsx'\n",
    "\n",
    "do_3pars_TM=0 # 3 or 4 parametric TM model?\n",
    "if do_3pars_TM:\n",
    "    fn5='STP_cortex_interneurons_Cre_labeled_UDF5_3pars.xlsx'\n",
    "else:\n",
    "    #fn5='STP_cortex_interneurons_Cre_labeled_UDF5_4pars.xlsx'\n",
    "    fn5='STP_hippocampus_all_UDF3.xlsx'\n",
    "\n",
    "df_stp=pd.read_excel(d5+fn5)\n",
    "\n",
    "ii=np.nonzero(df_stp.columns=='Var1_  1')[0][0]\n",
    "#tmp2 = df_stp.iloc[[9,13,14,15],ii:]\n",
    "tmp2 = df_stp.iloc[[13],ii:]\n",
    "for i in range(tmp2.shape[0]):\n",
    "    #tmp3 = tmp2.iloc[i,:].abs().values.reshape((-1,100)).T.ravel()\n",
    "    #tmp2.iloc[i,:]  = tmp3\n",
    "    tmp2.iloc[i,1::5]=0.24\n",
    "    \n",
    "#df_stp.iloc[[9,13,14,15],ii:]=tmp2   \n",
    "df_stp.iloc[[13],ii:]=tmp2   \n",
    "df_stp.to_excel(d5+fn5)\n",
    "df_stp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#do_laptop=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "#\n",
    "#  HIPPOCAMPUS STP\n",
    "#\n",
    "#\n",
    "import re\n",
    "import time\n",
    "t1 = time.time()\n",
    "\n",
    "\n",
    "# names of gene expression based synapses types\n",
    "\n",
    "glu_l23=[] #list(set(df_aba_vis_l23glu_c.loc[:,'cluster']))\n",
    "\n",
    "dsg = [['PC']+glu_l23,\n",
    "       ['Pvalb','Pvalb Tpbg','Pvalb Reln Itm2a','Pvalb Reln Tac1','Pvalb Sema3e Kank4'], \n",
    "       ['Sst','Sst Calb2 Pdlim5','Sst Calb2 Necab1', 'Sst Hpse Cbln4','Sst Hpse Sema3c','Sst Nr2f2 Necab1'],\n",
    "       ['Vip','Vip Lmo1 Myl1', 'Vip Rspo1 Itga4','Vip Ptprt Pkp2','Vip Rspo4 Rxfp1 Chat'],\n",
    "       ['Lamp5','Lamp5 Plch2 Dock5'],\n",
    "       ['MC','L23MC','Sst Calb2 Pdlim5'],\n",
    "        'L5MC']\n",
    "      \n",
    "\n",
    "#set(df_stp.loc[:,'synapse_type_2'])\n",
    "\n",
    "# names of electrophysiology based synapses types\n",
    "dse = [['PC', 'L23P', 'L23PC', 'L5P'],\n",
    "       ['Pvalb', 'PV', 'L23BC', 'L5BC'], \n",
    "       ['Sst', 'L23MC', 'L5MC', 'Sst', 'MC'], \n",
    "       ['Vip', 'VIP', 'BTC'],\n",
    "       ['NGC', 'L23NGC'],\n",
    "       ['L23MC' ],\n",
    "       ['L5MC']]\n",
    "\n",
    "\n",
    "# load stp data table\n",
    "#do_laptop=0\n",
    "if do_laptop:\n",
    "    #d5 = '~/Documents/references/lists_of_proteins_2/Analysis_of_proteins/genetic_types/transcriptomic_interneurons/additional cortical STP data/'\n",
    "    d5 = '/Users/stepaniu/Documents/jan_2020/'\n",
    "else:\n",
    "    d5 = '/home/stepaniu/Documents/references/transcriptomes to STP/scVI_data/additional cortical STP data/'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled.xlsx'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled_UDF2.xlsx'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled_UDF3.xlsx'\n",
    "\n",
    "do_3pars_TM=0 # 3 or 4 parametric TM model?\n",
    "if do_3pars_TM:\n",
    "    fn5='STP_cortex_interneurons_Cre_labeled_UDF5_3pars.xlsx'\n",
    "else:\n",
    "    #fn5='STP_cortex_interneurons_Cre_labeled_UDF5_4pars.xlsx'\n",
    "    fn5='STP_hippocampus_all_UDF3.xlsx'\n",
    "\n",
    "df_stp=pd.read_excel(d5+fn5)\n",
    "#df_stp.loc[df_stp.loc[:,'stp_comment'].isna() ,'stp_comment']=''\n",
    "\n",
    "############   \n",
    "do_averagig_stp =0\n",
    "stp_TM_names=['tF', 'p0','tD','dp/p0','A']\n",
    "stp_TM_start = np.nonzero(df_stp.columns=='references')[0][0]+1 # 187\n",
    "ncol = len(df_stp.columns)\n",
    "df_tm_name = pd.DataFrame(np.zeros(ncol).reshape((1,ncol)), columns =df_stp.columns )\n",
    "df_stp = pd.concat([df_tm_name, df_stp],axis=0).reset_index(drop=True)\n",
    "\n",
    "for i,nm in enumerate(stp_TM_names):\n",
    "    ii = stp_TM_start + np.arange(0,ncol-i-stp_TM_start,len(stp_TM_names))+i\n",
    "    df_stp.iloc[0,ii] = nm\n",
    "\n",
    "\n",
    "#print(df_stp.head())\n",
    "#############\n",
    "##Remove pure PPR data (Jaing_2015_PC)\n",
    "#df_stp = df_stp.loc[df_stp.loc[:,'name'].str.contains('Jaing_2015_PC')!=True,:].reset_index(drop=True) \n",
    "\n",
    "############\n",
    "\n",
    "\n",
    "#stp_TM_names_2=['tF', 'p0','tD','dp','A','A1','A2','tDmin','dd','t_FDR','t_SMR','dp0'] # tm5+smr\n",
    "##################\n",
    "#stp_2_default = np.ones(7)\n",
    "#for i in range(1,df_stp.shape[0]):\n",
    "#    stp_par = df_stp.iloc[i,stp_TM_start:].values.reshape((len(stp_TM_names),-1))\n",
    "    \n",
    "\n",
    "##################\n",
    "\n",
    "df_stp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hippocampus: allign gene expression and stp data (for selected gene sets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters:\n",
    "do_bootstrap=1\n",
    "do_median=0\n",
    "\n",
    "nsample_hipp = 200\n",
    "\n",
    "#for cell types diagrams : do_bootstrap=1\n",
    "#do_median=0\n",
    "#nsample_hipp = 400\n",
    "\n",
    "do_log='none'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# rules for cell types to aba clusters mapping:\n",
    "#1 dorsal : 'subregion_label'== 'anterior' , ventral: 'subregion_label'== posterior' ,  or by genes : see Cembrowski 2016\n",
    "#2 TA = ECIII (PCP4+) ->CA1 \n",
    "#3 pp = ECII (Reelin+) ->CA3, DG, CA2\n",
    "#4 TA = ECII (Wlf1+, Calb2+) ->CCK ppa in CA1 mostly\n",
    "#5 no cre lines annotations - use only old cck datatypes->clusters or corresponding cortical cre lines?\n",
    "#6 Ivy 1, 2 ?? - 2 clusters (Lhx6 1,2) and may be Lhx6 3 (no reelin), but also could be some Lamp5_5? - less reelin, more Nos1 \n",
    "#  use only only PPF Ivy cells? or try 2 subtypes = 2 clusters?? - dep=lhx6_4\n",
    "#7 NGF types 1,2? Price 2005 - correlate with age - select only facilitating (type 2)!\n",
    "#8 HIPP - Sst, pp-associated ~OLM??\n",
    "#9 HICAP - CCk - comissural assoc - ~ CCK Schffer comm : Sncg6 (Sncg/Ndnf HPF1)\n",
    "#10 Kohus: CCK DTI vs BC? - use Cbln4, Lgi2 ratio? (Favuzzi 2019) - Sncg2?, (Sncg6) - DTI vs Sncg/Ndnf HPF2,3,6, Serpinf1, Sncg7, Vip_14\n",
    "#11 CA1 PC: use subregions: anterior, mid-anterior  for dorsal\n",
    "\n",
    "\n",
    "\n",
    "# batches and gene sets\n",
    "pre_post = df_stp.loc[:,'synapse_type_2'].str.split(pat='->',expand=True)\n",
    "\n",
    "vbi = vis_dat.batch_indices.reshape((vis_dat.batch_indices.shape[0],))\n",
    "vbi=vbi>0\n",
    "\n",
    "df_gs  = pd.DataFrame( vis_dat.gene_names ) \n",
    "in_pregs = df_gs.loc[:,0].isin(pregs).values\n",
    "in_postgs = df_gs.loc[:,0].isin(postgs).values\n",
    "\n",
    "pregs2 = df_gs.loc[in_pregs,0].values\n",
    "postgs2 = df_gs.loc[in_postgs,0].values\n",
    "\n",
    "do_hippocampus = 1\n",
    "if do_hippocampus ==0:\n",
    "\n",
    "    do_2=1;\n",
    "    # select stp data for appropriate frequency and ages \n",
    "    stp_type = {'name':['Walker_2016_PV->MC L23',\n",
    "                   'Walker_2016_VIP->MC L23',\n",
    "                   'Yuste_2016_taus_PV->Sst L23 recalc',\n",
    "                   'Yuste_2016_taus_PV->VIP L23 recalc',\n",
    "                   'Yuste_2016_taus_Sst->VIP L23 recalc',\n",
    "                   'Yuste_2016_taus_VIP->Sst L23 recalc',\n",
    "                   'Yuste_2016_taus_VIP->VIP L23 recalc'],\n",
    "            'stp_freq': [20, 40, 50], 'stp_comment':  ['', ' P20-21', ' P20-22', ' P20-23', ' P25-31'],\n",
    "            'stp_data_type': df_stp.loc[0, :]!=0 }\n",
    "\n",
    "    if do_2==1:\n",
    "        stp_type = {'name':['Walker_2016_PV->MC L23',\n",
    "                    'Walker_2016_VIP->MC L23',\n",
    "                    'Yuste_2016_taus_PV->Sst L23 recalc',\n",
    "                    'Yuste_2016_taus_VIP->Sst L23 recalc'],\n",
    "            'stp_freq': [20, 40, 50], 'stp_comment':  ['', ' P20-21', ' P20-22', ' P20-23', ' P25-31'],\n",
    "            'stp_data_type': df_stp.loc[0, :]!=0 }\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    in_stp_data = df_stp.loc[:,'stp_freq'].isin([20])\n",
    "    in_stp_data= in_stp_data|(df_stp.loc[:,'stp_freq'].isin([40,50])&df_stp.loc[:,'name'].isin(stp_type['name']))\n",
    "    in_stp_data = in_stp_data&df_stp.loc[:,'stp_comment'].isin(pd.Series(stp_type['stp_comment']).str.strip())\n",
    "\n",
    "else:\n",
    "    \n",
    "    in_stp_data = pd.Series(np.ones(df_stp.shape[0])==1)\n",
    "    # filter stp dataset \n",
    "    in_stp_data.loc[df_stp.loc[:,'name'].isin(['Price_2005EC->NGF type1',\n",
    "                                               'Qin_2017EC->Ivy cell type1',\n",
    "                                               'Qin_2017CA3->Ivy cell type1'])]=False\n",
    "    #rem = df_stp.loc[:,'name'].str.contains('HIPP').fillna(value=True)\n",
    "    #in_stp_data.loc[rem]=False\n",
    "    \n",
    "    # 17–20 Sprague-Dawley rats?? - too high depression? - 0.39\n",
    "    rem = df_stp.loc[:,'name'].str.contains('Fuentealba_2008CA1->Ivy cell').fillna(value=True)\n",
    "    in_stp_data.loc[rem]=False    \n",
    "    \n",
    "# show selected stp dataset\n",
    "ddd=df_stp.loc[in_stp_data,['name','synapse_type_2', 'stp_freq','area']]\n",
    "print(ddd)\n",
    "print(in_stp_data.sum())\n",
    "\n",
    "synij = np.repeat(False,in_stp_data.shape[0])\n",
    "\n",
    "cols2 = ['name','perc']\n",
    "cl2post =pd.DataFrame([],columns= cols2)\n",
    "cl2pre =pd.DataFrame([],columns= cols2)\n",
    "\n",
    "\n",
    "# strip cluster names\n",
    "#list(set(df_aba_vis_in_c.loc[df_aba_vis_in_c.loc[:,'subclass'].str.contains(clns6),'cluster']))\n",
    "df_aba_vis_in_c.loc[:,'cluster_label']=df_aba_vis_in_c.loc[:,'cluster_label'].str.strip()\n",
    "#df_aba_vis_l23glu_c.loc[:,'cluster']=df_aba_vis_l23glu_c.loc[:,'cluster'].str.strip()\n",
    "\n",
    "do_simplfy_clusters = 1 # do not take into account minor subclasses\n",
    "Sst_NonMC ='Sst, Sst Calb2 Pdlim5-, Sst Chrna2 Glra3-, Sst Chrna2 Ptgdr-, Sst Myh8 Etv1-, Sst Tac2 Myh4-, Sst Nr2f2 Necab1-, Sst Chodl-,  Sst Calb2 Necab1-, Sst Tac1 Htr1d-, Sst Myh8 Fibin-';\n",
    "\n",
    "pregs3 = ['pre__'+s for s in pregs2.tolist()]\n",
    "postgs3 = ['post__'+s for s in postgs2.tolist()]\n",
    "ge_columns = ['cell_type2_pre', 'cell_type2_post',\n",
    "                                             'layer_pre','layer_post',\n",
    "                                             'cre_line_pre', 'cre_line_post'] +pregs3 +postgs3 + ['samples_pre', 'samples_post']\n",
    "\n",
    "stp_columns = stp_TM_names #list(set(df_stp.loc[0,:]).difference(set([0])))\n",
    "\n",
    "# output datasets\n",
    "stp_data = pd.DataFrame([], columns =['presynaptic', 'postsynaptic', 'stp_data'] )\n",
    "ge_data = pd.DataFrame([], columns =ge_columns+stp_columns+['index_ds'])\n",
    "\n",
    "\n",
    "for i in range(1,len(pre_post.index)): #range(1,len(pre_post.index)): # [1]: \n",
    "    \n",
    "    # syne3 - all available pre- and post- synaptic cells annotations : cre-lines, layers, morphologies\n",
    "    l = str.strip(df_stp.loc[i,'area'])\n",
    "    synel = re.split(';',l)\n",
    "    if len(synel)==1:\n",
    "        l1=l\n",
    "    else:\n",
    "        l  = str.strip(synel[1])\n",
    "        l1  = str.strip(synel[0])\n",
    "            \n",
    "    \n",
    "    if type(df_stp.loc[i,'synapse_type'])!=float:\n",
    "        #syne1 = re.split('\\;',df_stp.loc[i,'synapse_type']) no cre-line annotations for hippocampal data\n",
    "        syne1 = ''\n",
    "    else:\n",
    "        syne1 =''    \n",
    "    #syne2 = df_stp.loc[i,'synapse_type2'].str.split(pat='->',expand=True)\n",
    "    syne2 = re.split('->',df_stp.loc[i,'synapse_type_2'])\n",
    "    if len(syne2)==1:\n",
    "        pat ='→'\n",
    "        syne2 = re.split(pat,df_stp.loc[i,'synapse_type_2'])\n",
    "\n",
    "    syne3 ={'layer_pre':l1,  'cre_line_pre':'none', 'cell_type2_pre':str.strip(syne2[0]),'cluster_pre':[],\n",
    "            'layer_post':l, 'cre_line_post':'none','cell_type2_post':str.strip(syne2[1]),'cluster_post':[],  } \n",
    "\n",
    "    # add cre-lines when available\n",
    "    if len(syne1)>1:\n",
    "        syne3['cre_line_pre']  =str.strip(syne1[0])\n",
    "        syne3['cre_line_post'] =str.strip(syne1[1])\n",
    "    \n",
    "    # PV BC\n",
    "    if (syne3['cell_type2_post']=='PV')|(syne3['cell_type2_post']=='PVBC'):  \n",
    "        syne3['cluster_post']=[['Pvalb_3, Pvalb_10, Pvalb_4, Pvalb_6',100]] # modify!!! \n",
    "    if (syne3['cell_type2_pre']=='PV')|(syne3['cell_type2_pre']=='PVBC'):  \n",
    "        syne3['cluster_pre']=[['Pvalb_3, Pvalb_10, Pvalb_4, Pvalb_6',100]] # modify!!!   \n",
    "                     \n",
    "  \n",
    "    # HIPP -> HIPP Bartos\n",
    "    if syne3['cell_type2_post']=='HIPP':  \n",
    "        if syne3['layer_post'] == 'DG':\n",
    "            syne3['cluster_post']=[['Sst_15, Sst_24, Sst_13, Sst_6',100]] # modify!!!        \n",
    "    if syne3['cell_type2_pre']=='HIPP':  \n",
    "        if syne3['layer_pre'] == 'DG':\n",
    "            syne3['cluster_pre']=[['Sst_15, Sst_24, Sst_13, Sst_6',100]] # modify!!! \n",
    "            \n",
    "            \n",
    "    # HICAP -> HICAP Bartos\n",
    "    if syne3['cell_type2_post']=='HICAP':  \n",
    "        if syne3['layer_post'] == 'DG':\n",
    "            syne3['cluster_post']=[['Sncg6, Sncg/Ndnf HPF_1',100]] # modify!!!        \n",
    "    if syne3['cell_type2_pre']=='HICAP':  \n",
    "        if syne3['layer_pre'] == 'DG':\n",
    "            syne3['cluster_pre']=[['Sncg6, Sncg/Ndnf HPF_1',100]] # modify!!!      \n",
    "            \n",
    "            \n",
    "    # O-LM\n",
    "    if syne3['cell_type2_post']=='O-LM':  \n",
    "        if syne3['layer_post'] == 'CA1':\n",
    "            syne3['cluster_post']=[['Sst_15, Sst_24, Sst_13, Sst_6',100]] # modify!!!     \n",
    "            \n",
    "    # CCK BC\n",
    "    if syne3['cell_type2_post']=='CCKBC':  \n",
    "        syne3['cluster_post']=[['Sncg/Ndnf HPF_2, Sncg/Ndnf HPF_3, Sncg/Ndnf HPF_6, Serpinf1_1, Sncg_7, Vip_14',100]] # modify!!!    \n",
    "    if syne3['cell_type2_pre']=='CCKBC':  \n",
    "        syne3['cluster_pre']=[['Sncg/Ndnf HPF_2, Sncg/Ndnf HPF_3, Sncg/Ndnf HPF_6, Serpinf1_1, Sncg_7, Vip_14',100]] # modify!!!    \n",
    "\n",
    "    \n",
    "    # CCK DTI\n",
    "    if syne3['cell_type2_post']=='CCK_DTI':  \n",
    "        syne3['cluster_post']=[['Sncg/Ndnf HPF_1, Sncg6, Sncg2',100]] # modify!!!    \n",
    "    if syne3['cell_type2_pre']=='CCK_DTI':  \n",
    "        syne3['cluster_pre']=[['Sncg/Ndnf HPF_1, Sncg6, Sncg2',100]] # modify!!!    \n",
    "\n",
    "        \n",
    "        \n",
    "    # O-Bi\n",
    "    if syne3['cell_type2_post']=='O-Bi':  \n",
    "        if syne3['layer_post'] == 'CA1':\n",
    "            syne3['cluster_post']=[['Sst_27, Pvalb_1, Sst_25',100]] # modify!!!   \n",
    "     \n",
    "    # AAC\n",
    "    if syne3['cell_type2_pre']=='AAC':  \n",
    "        if syne3['layer_pre'] == 'CA3':\n",
    "            syne3['cluster_pre']=[['Pvalb_12, Pvalb_11',100]] # modify!!!   \n",
    "    if syne3['cell_type2_post']=='AAC':  \n",
    "        if syne3['layer_post'] == 'CA3':\n",
    "            syne3['cluster_post']=[['Pvalb_12, Pvalb_11',100]] # modify!!!         \n",
    "\n",
    "    # NGF    ? type2 vs type1? - remove type 1?  \n",
    "    if syne3['cell_type2_pre']=='NGF':  \n",
    "        syne3['cluster_pre']=[['Lamp5_2, Lamp5_3, Lamp5_4, Lamp5_5, Lamp5 Lhx6_2, Lamp5 Lhx6_3',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions       \n",
    "    if (syne3['cell_type2_post']=='NGF')|(syne3['cell_type2_post']=='NGF type1')|(syne3['cell_type2_post']=='NGF type2'):  \n",
    "        syne3['cluster_post']=[['Lamp5_2, Lamp5_3, Lamp5_4, Lamp5_5, Lamp5 Lhx6_2, Lamp5 Lhx6_3',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions   \n",
    "     \n",
    "    # Ivy cells   ? type2 vs type1? : CGE vs MGE? - no CGE is seen? or Lhs1 vs Lhs4? \n",
    "    if syne3['cell_type2_pre']=='Ivy cell':  \n",
    "        syne3['cluster_pre']=[['Lamp5 Lhx6_1, Lamp5 Lhx6_4',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions       \n",
    "    if (syne3['cell_type2_post']=='Ivy cell')|(syne3['cell_type2_post']=='Ivy cell type1')|(syne3['cell_type2_post']=='Ivy cell type2'):  \n",
    "        syne3['cluster_post']=[['Lamp5 Lhx6_1, Lamp5 Lhx6_4',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions   \n",
    "       \n",
    "    #IS3\n",
    "    if syne3['cell_type2_pre']=='IS3':  \n",
    "        syne3['cluster_pre']=[['Vip_6, Vip_12, Vip_15, Vip_17',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions       \n",
    "    \n",
    "    ###### EXCitatory\n",
    "    # CA1 PC        \n",
    "    if (syne3['cell_type2_pre']=='CA1'): \n",
    "        syne3['cell_type2_pre']=='PC'\n",
    "        if syne3['layer_pre'] == 'CA1':\n",
    "            syne3['cluster_pre']=[['CA1sp_7, CA1sp_8',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions       \n",
    "    if (syne3['cell_type2_post']=='CA1'):  \n",
    "        syne3['cell_type2_post']=='PC'\n",
    "        if syne3['layer_post'][0:3] == 'CA1':\n",
    "            syne3['cluster_post']=[['CA1sp_7, CA1sp_8',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions   \n",
    " \n",
    "            \n",
    "    # CA3 PC        \n",
    "    if (syne3['cell_type2_pre']=='CA3'): \n",
    "        syne3['cell_type2_pre']=='PC'\n",
    "        if syne3['layer_pre'] == 'CA3':\n",
    "            syne3['cluster_pre']=[['CA3sp_6, CA3sp_5, CA3sp_7, CA3sp_1',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions       \n",
    "    if (syne3['cell_type2_post']=='CA3'):  \n",
    "        syne3['cell_type2_post']=='PC'\n",
    "        if syne3['layer_post'] == 'CA3':\n",
    "            syne3['cluster_post']=[['CA3sp_6, CA3sp_5, CA3sp_7, CA3sp_1',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions   \n",
    "    \n",
    "    # CA3 PC ventral\n",
    "    if (syne3['cell_type2_pre']=='CA3')|(syne3['cell_type2_pre']=='PC'): \n",
    "        syne3['cell_type2_pre']=='PC'\n",
    "        if (syne3['layer_pre'] == 'CA3 ventral')&(syne3['layer_post'] == 'CA1 ventral'):\n",
    "            syne3['cluster_pre']=[['CA3sp_2, CA3sp_4',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions   \n",
    "\n",
    "    \n",
    "    \n",
    " \n",
    "    # ECII SS        \n",
    "    if (syne3['cell_type2_pre']=='EC'): \n",
    "        syne3['cell_type2_pre']=='PC'\n",
    "        if (syne3['layer_pre'] == 'EC')&((syne3['layer_post'] == 'DG')|(syne3['layer_post'] == 'CA3')):\n",
    "            syne3['cluster_pre']=[['L2/3 IT Ndst4 Endou_2',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions       \n",
    "\n",
    "\n",
    "    # ECIII PC        \n",
    "    if (syne3['cell_type2_pre']=='EC'):\n",
    "        syne3['cell_type2_pre']=='PC'\n",
    "        if (syne3['layer_pre'] == 'EC')&((syne3['layer_post'][0:3] == 'CA1')):\n",
    "            syne3['cluster_pre']=[['L2/3 IT Plch1_1, L2/3 IT Plch1_2',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions       \n",
    "\n",
    "            \n",
    "    # CA1 PC ventral       \n",
    "    if (syne3['cell_type2_pre']=='CA1 ventral'):  \n",
    "        syne3['cell_type2_pre']=='PC'\n",
    "        if syne3['layer_pre'] == 'CA1 ventral':\n",
    "            syne3['cluster_pre']=[['CA1sp_1, CA1sp_2, CA1sp_3',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions       \n",
    "    if (syne3['cell_type2_post']=='CA1 ventral'):  \n",
    "        syne3['cell_type2_post']=='PC'\n",
    "        if syne3['layer_post'] == 'CA1 ventral':\n",
    "            syne3['cluster_post']=[['CA1sp_1, CA1sp_2, CA1sp_3',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions   \n",
    "    \n",
    "    \n",
    "\n",
    "    # DG PC      \n",
    "    if (syne3['cell_type2_pre']=='DG'):\n",
    "        syne3['cell_type2_pre']=='PC'\n",
    "        #if syne3['layer_pre'] == 'DG':\n",
    "        syne3['cluster_pre']=[['DG_1, DG_2, DG_3',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions       \n",
    "    if (syne3['cell_type2_post']=='DG'): \n",
    "        syne3['cell_type2_post']=='PC'\n",
    "        #if syne3['layer_post'] == 'DG':\n",
    "        syne3['cluster_post']=[['DG_1, DG_2, DG_3',100]] # modify!!! - only ~35 cells? - selected by anterior and mid-anterior subregions   \n",
    "   \n",
    "    \n",
    "    \n",
    "    ###########  Select data:\n",
    "    ##synij = (pre_post.loc[:,0].isin(pree))&(pre_post.loc[:,1].isin(poste))\n",
    "    \n",
    "\n",
    "    synij[:] = False\n",
    "    synij[i] = in_stp_data[i]\n",
    "    \n",
    "    #print(syne3)\n",
    "    \n",
    "    if sum(synij)>0:\n",
    "        print('\\n\\n')\n",
    "        print('i = ',str(i))\n",
    "        print(syne3)\n",
    "        #if sum(in_stp_data&synij)>0:\n",
    "        #    in_stp_data = in_stp_data&synij\n",
    "\n",
    "        #add ge data\n",
    "        #'layer_pre':l,  'cre_line_pre':'', 'cell_type2_pre':syne2[0].str.strip(),'cluster_pre':[]\n",
    "        \n",
    "        # PC or GABAergic dataset?\n",
    "        batch = [ syne3['cell_type2_pre']=='PC',syne3['cell_type2_post']=='PC']\n",
    "        \n",
    "        # add ge data depending on available stp cell types labels:\n",
    "        #  select presynaptic cells\n",
    "        if 0: #batch[0]:\n",
    "            #in_preg  = df_aba_vis_l23glu_c.loc[:,'brain_subregion'].isin([syne3['layer_pre']]) #modify!!!\n",
    "            in_preg  = df_aba_vis_l23glu_c.loc[:,'brain_subregion']==syne3['layer_pre'] #modify!!!\n",
    "        else:\n",
    "            if len(syne3['cluster_pre'])==0:\n",
    "                # if len(syne3['cre_line_pre'])>0:\n",
    "                in_preg = df_aba_vis_in_c.loc[:,'full_genotype_label'].str.contains(syne3['cre_line_pre'])\n",
    "                ##in_preg = in_preg&df_aba_vis_in_c.loc[:,'brain_subregion'].isin([syne3['layer_pre']])\n",
    "                #in_preg = in_preg&df_aba_vis_in_c.loc[:,'brain_subregion'].str.contains(syne3['layer_pre'])\n",
    "            else:                  \n",
    "                ##in_preg = df_aba_vis_in_c.loc[:,'brain_subregion'].isin([syne3['layer_pre']])\n",
    "                in_preg = df_aba_vis_in_c.loc[:,'region_label'].str.contains(syne3['layer_pre'])\n",
    "                in_preg[:] = True\n",
    "                \n",
    "                cl = syne3['cluster_pre']\n",
    "                #cols2 = ['name','perc']\n",
    "                #cl2pre =pd.DataFrame([],columns= cols2)\n",
    "                in_preg0 = in_preg.copy()\n",
    "                in_preg[:] = False \n",
    "                for icl in range(len(cl)):\n",
    "                    clp = cl[icl][1]\n",
    "                    clns = re.split(',',cl[icl][0])\n",
    "                    cls3 = []\n",
    "                    cls3m=[]\n",
    "                    for icls in range(len(clns)):\n",
    "                        clns4 = str.strip(clns[icls])\n",
    "                        clns5 = re.split('\\_',clns4)\n",
    "                        if len(clns5)<2:\n",
    "                            clns6=clns5[0]\n",
    "                            cls3=cls3 + list(set(df_aba_vis_in_c.loc[df_aba_vis_in_c.loc[:,'subclass_label'].str.contains(clns6),'cluster_label']))\n",
    "                        else:\n",
    "                            clns4=str.strip(clns4)\n",
    "                            clns7=re.split('\\-',clns4)\n",
    "                            if len(clns7)>1:\n",
    "                                clns4=clns7[0]\n",
    "                                cls3m=cls3m+[str.strip(clns4)]\n",
    "                            else:\n",
    "                                cls3=cls3+[str.strip(clns4)]\n",
    "                    cls3 = list(set(cls3).difference(set(cls3m)))\n",
    "                    cl22=pd.DataFrame([cls3,[clp]],index=cols2).T\n",
    "                    cl2pre = cl2pre.append(cl22)\n",
    "                    \n",
    "                    in_preg = in_preg|(in_preg0&df_aba_vis_in_c.loc[:,'cluster_label'].isin(cls3))\n",
    "                    \n",
    "                if (len(syne3['cre_line_pre'])>0)&(syne3['cre_line_pre']!='none'):\n",
    "                    #in_preg = in_preg&df_aba_vis_in_c.loc[:,'driver_lines'].str.contains(syne3['cre_line_pre'])\n",
    "                    in_line=df_aba_vis_in_c.loc[:,'full_genotype_label'].str.contains(syne3['cre_line_pre'])\n",
    "                    print('is pre cre line')\n",
    "                    if in_line.sum()>0:\n",
    "                        in_preg = in_preg&in_line                        \n",
    "        \n",
    "        #  select postsynaptic cells\n",
    "        if 0: #batch[1]:\n",
    "            #in_postg = df_aba_vis_l23glu_c.loc[:,'brain_subregion'].isin([syne3['layer_post']]) #modify!!!\n",
    "            in_postg = df_aba_vis_l23glu_c.loc[:,'brain_subregion']==syne3['layer_post'] #modify!!!\n",
    "        else:    \n",
    "            if len(syne3['cluster_post'])==0:\n",
    "                # if len(syne3['cre_line_post'])>0:\n",
    "                in_postg = df_aba_vis_in_c.loc[:,'full_genotype_label'].str.contains(syne3['cre_line_post'])\n",
    "                #in_postg = in_postg&df_aba_vis_in_c.loc[:,'brain_subregion'].isin([syne3['layer_post']])\n",
    "                #in_postg = in_postg&df_aba_vis_in_c.loc[:,'brain_subregion_label'].str.contains(syne3['layer_post'])\n",
    "            else:    \n",
    "                ## under construction ...\n",
    "                ##in_postg = df_aba_vis_in_c.loc[:,'brain_subregion'].isin([syne3['layer_post']])\n",
    "                in_postg = df_aba_vis_in_c.loc[:,'region_label'].str.contains(syne3['layer_post'])\n",
    "                in_postg[:] = True\n",
    "                \n",
    "                cl = syne3['cluster_post']\n",
    "                #cols2 = ['name','perc']\n",
    "                #cl2post =pd.DataFrame([],columns= cols2)\n",
    "                in_postg0 = in_postg.copy()\n",
    "                in_postg[:] = False                 \n",
    "                for icl in range(len(cl)):\n",
    "                    clp = cl[icl][1]\n",
    "                    clns = re.split(',',cl[icl][0])\n",
    "                    cls3 = []\n",
    "                    cls3m=[]\n",
    "                    for icls in range(len(clns)):\n",
    "                        clns4 = str.strip(clns[icls])\n",
    "                        clns5 = re.split('\\_',clns4)\n",
    "                        if len(clns5)<2:\n",
    "                            clns6=clns5[0]\n",
    "                            cls3=cls3 + list(set(df_aba_vis_in_c.loc[df_aba_vis_in_c.loc[:,'subclass_label'].str.contains(clns6),'cluster_label']))\n",
    "                        else:\n",
    "                            clns4=str.strip(clns4)\n",
    "                            clns7=re.split('\\-',clns4)\n",
    "                            if len(clns7)>1:\n",
    "                                clns4=clns7[0]\n",
    "                                cls3m=cls3m+[str.strip(clns4)]\n",
    "                            else:\n",
    "                                cls3=cls3+[str.strip(clns4)]\n",
    "                    cls3 = list(set(cls3).difference(set(cls3m)))\n",
    "                    cl22=pd.DataFrame([cls3,[clp]],index=cols2).T\n",
    "                    cl2post = cl2post.append(cl22)\n",
    "                    \n",
    "                    #in_postg = in_postg&df_aba_vis_in_c.loc[:,'cluster'].isin(cls3)\n",
    "                    in_postg = in_postg|(in_postg0&df_aba_vis_in_c.loc[:,'cluster_label'].isin(cls3))\n",
    "                    \n",
    "                if (len(syne3['cre_line_post'])>0)&(syne3['cre_line_post']!='none'):\n",
    "                    in_line=df_aba_vis_in_c.loc[:,'full_genotype_label'].str.contains(syne3['cre_line_post'])\n",
    "                    print('is post cre line')\n",
    "                    if in_line.sum()>0:\n",
    "                        in_postg = in_postg&in_line\n",
    "            \n",
    "           \n",
    "            \n",
    "        # make equal size pre- and post- datasets\n",
    "        in_preg = in_preg&(in_preg.isna()==False)\n",
    "        in_postg = in_postg&(in_postg.isna()==False)\n",
    "        npre = in_preg.sum()\n",
    "        npost = in_postg.sum()\n",
    "        nsample = max([npost, npre] )\n",
    "        #nsample = min([npost, npre] ) # number of cells\n",
    "        \n",
    "        if (nsample<60)&(npre>0)&(npost>0):\n",
    "            nsample=60\n",
    "            \n",
    "        if (nsample>180)&(npre>0)&(npost>0):  # restrict nsample to avoid overrepresentation?\n",
    "            nsample=180    \n",
    "            \n",
    "        print('npre = '+ str(npre)+' , npost = ' + str(npost) + ' , nsample = '+str(nsample))    \n",
    "        \n",
    "        \n",
    "        ## select cell indices for gene expression sampling\n",
    "        #do_bootstrap=1\n",
    "        #if  do_bootstrap==0:\n",
    "        #    idx2 =   np.mod(np.arange(nsample),npre) \n",
    "        #    in_preg = np.nonzero(in_preg)[0][idx2]\n",
    "        #    idx2 =   np.mod(np.arange(nsample),npost) \n",
    "        #    in_postg = np.nonzero(in_postg)[0][idx2]\n",
    "        #else :\n",
    "        #    #random.uniform(low=0.0, high=1.0, size=None)\n",
    "        #    idx2 =  np.floor( np.random.uniform(0,npre, nsample) ).astype('int')\n",
    "        #    in_preg = np.nonzero(in_preg)[0][idx2]\n",
    "        #    idx2 =  np.floor( np.random.uniform(0,npost, nsample) ).astype('int')\n",
    "        #    in_postg = np.nonzero(in_postg)[0][idx2]\n",
    "            \n",
    "            \n",
    "        # select cell indices for gene expression sampling\n",
    "        \n",
    "                \n",
    "        ##\n",
    "        ##\n",
    "        ##  PARAMETERS OF GS AVERAGING\n",
    "        ##\n",
    "        ##\n",
    "        do_bootstrap=1\n",
    "        if  do_bootstrap==0:\n",
    "            \n",
    "            if (nsample<60)&(npre>0)&(npost>0):\n",
    "                nsample=60\n",
    "            \n",
    "            if (nsample>180)&(npre>0)&(npost>0):  # restrict nsample to avoid overrepresentation?\n",
    "                nsample=180 \n",
    "            \n",
    "            idx2 =   np.mod(np.arange(nsample),npre) \n",
    "            in_preg = np.nonzero(in_preg)[0][idx2]\n",
    "            idx2 =   np.mod(np.arange(nsample),npost) \n",
    "            in_postg = np.nonzero(in_postg)[0][idx2]\n",
    "            \n",
    "            in_preg_ = in_preg\n",
    "            in_postg_ = in_postg\n",
    "        else :\n",
    "            if (npre>0)&(npost>0):\n",
    "                nsample= nsample_hipp #= 200\n",
    "                \n",
    "            #random.uniform(low=0.0, high=1.0, size=None)\n",
    "            idx2pre =  np.floor( np.random.uniform(0,npre, (nsample,1)) ).astype('int')\n",
    "            in_preg_bs = np.nonzero(in_preg)[0][idx2pre]\n",
    "            idx2post =  np.floor( np.random.uniform(0,npost, (nsample,1)) ).astype('int')\n",
    "            in_postg_bs = np.nonzero(in_postg)[0][idx2post]\n",
    "            in_preg_ = in_preg_bs[:,0]\n",
    "            in_postg_ = in_postg_bs[:,0]    \n",
    "            \n",
    "\n",
    "\n",
    "        \n",
    "        sampl_post =df_aba_vis_in_c.iloc[in_postg_,:].loc[:,'sample_name']\n",
    "        sampl_pre =df_aba_vis_in_c.iloc[in_preg_,:].loc[:,'sample_name']\n",
    "        \n",
    "        if batch[0]:\n",
    "            #in_preg_cl = list(set(df_aba_vis_l23glu_c.iloc[in_preg,:].loc[:,'cluster']))\n",
    "            cl_pre = df_aba_vis_in_c.iloc[in_preg_,:].loc[:,'cluster_label']\n",
    "            in_preg_cl = cl_pre.value_counts()\n",
    "            \n",
    "        else:\n",
    "            #in_preg_cl = list(set(df_aba_vis_in_c.iloc[in_preg,:].loc[:,'cluster']))\n",
    "            cl_pre = df_aba_vis_in_c.iloc[in_preg_,:].loc[:,'cluster_label']\n",
    "            in_preg_cl = cl_pre.value_counts()\n",
    "            #in_preg_cl = df_aba_vis_in_c.iloc[in_preg_,:].loc[:,'cluster_label'].value_counts()\n",
    "        \n",
    "        if batch[1]:\n",
    "            #in_postg_cl = list(set(df_aba_vis_l23glu_c.iloc[in_postg,:].loc[:,'cluster']))\n",
    "            cl_post =df_aba_vis_in_c.iloc[in_postg_,:].loc[:,'cluster_label']\n",
    "            in_postg_cl = cl_post.value_counts()\n",
    "            #in_postg_cl = df_aba_vis_in_c.iloc[in_postg_,:].loc[:,'cluster_label'].value_counts()\n",
    "        else:        \n",
    "            #in_postg_cl = list(set(df_aba_vis_in_c.iloc[in_postg,:].loc[:,'cluster']))\n",
    "            \n",
    "            cl_post =df_aba_vis_in_c.iloc[in_postg_,:].loc[:,'cluster_label']\n",
    "            in_postg_cl = cl_post.value_counts()\n",
    "            #in_postg_cl = df_aba_vis_in_c.iloc[in_postg_,:].loc[:,'cluster_label'].value_counts()\n",
    "        in_preg_cl = in_preg_cl/in_preg_cl.sum()*100    \n",
    "        in_postg_cl = in_postg_cl/in_postg_cl.sum()*100    \n",
    "        print('presynaptic clusters')\n",
    "        print(in_preg_cl)\n",
    "        print('postsynaptic clusters')\n",
    "        print(in_postg_cl)\n",
    "\n",
    "        #if npost>npre:\n",
    "        #    idx = np.mod(np.arange(npost),npre)\n",
    "        #    #prege = prege[idx,:]\n",
    "        #    #postge = postge\n",
    "        #    prege = prege\n",
    "        #    postge = postge[0:npre,:]            \n",
    "        #elif npost<npre :\n",
    "        #    idx = np.mod(np.arange(npre),npost)\n",
    "        #    #prege = prege\n",
    "        #    #postge = postge[idx,:]\n",
    "        #    prege = prege[0:npost,:]\n",
    "        #    postge = postge         \n",
    "        \n",
    "        # take ge data    \n",
    "        #prege = vis_dat.X[:, in_pregs]  # presynaptic gene set expression\n",
    "        ##prege = prege[(vbi==batch[0]), :]\n",
    "        #prege = prege[in_preg,:]\n",
    "\n",
    "        #postge = vis_dat.X[:,  in_postgs] # postsynaptic gene set expression\n",
    "        ##postge = postge[(vbi==batch[1]),  :]\n",
    "        #postge = postge[in_postg,:]\n",
    "        \n",
    "        \n",
    "        \n",
    "        prege = vis_dat.X[:, in_pregs]  # presynaptic gene set expression\n",
    "        prege = prege[(vbi==batch[0]), :]\n",
    "        \n",
    "        postge = vis_dat.X[:,  in_postgs] # postsynaptic gene set expression\n",
    "        postge = postge[(vbi==batch[1]),  :]\n",
    "        \n",
    "        ##\n",
    "        ##\n",
    "        ##  PARAMETERS OF GS AVERAGING\n",
    "        ##\n",
    "        ##\n",
    "        #do_median=0      # do median of gene expression data\n",
    "        do_log='none'     # do logarithm transform of ge counts\n",
    "        \n",
    "        if  (do_bootstrap==0)|(do_median==0):\n",
    "            prege = prege[in_preg_,:]\n",
    "            postge = postge[in_postg_,:]\n",
    "            do_samples_columns=1\n",
    "        else:\n",
    "            do_samples_columns=0\n",
    "            nsample = nsample_hipp #200 \n",
    "            Nbs = nsample\n",
    "            ##random.uniform(low=0.0, high=1.0, size=None)\n",
    "            #nmax=5\n",
    "            nmax = npre #np.min([1000,npre])\n",
    "            idx2pre =  np.floor( np.random.uniform(0,npre, (nmax,Nbs)) ).astype('int')\n",
    "            in_preg_ = np.nonzero(in_preg)[0][np.arange(npre)]\n",
    "            nmax = npost #np.min([1000,npost])\n",
    "            idx2post =  np.floor( np.random.uniform(0,npost, (nmax,Nbs)) ).astype('int')\n",
    "            in_postg_ = np.nonzero(in_postg)[0][np.arange(npost)]\n",
    "\n",
    "                        # \n",
    "            # ??? transform ge to median bootstraped log ???\n",
    "            prege_ = prege[in_preg_,:]\n",
    "            postge_ = postge[in_postg_,:]\n",
    "\n",
    "            if do_log=='none':\n",
    "                 prege_ =  prege_\n",
    "            elif do_log=='log2':       \n",
    "                y = np.ma.masked_where(prege_ == 0, prege_)\n",
    "                prege_=np.ma.log2(y).filled(0)\n",
    "                y = np.ma.masked_where(postge_ == 0, postge_)\n",
    "                postge_=np.ma.log2(y).filled(0)\n",
    "            elif do_log=='log1p': \n",
    "                prege_ = np.log1p(prege_)\n",
    "                postge_ =  np.log1p(postge_)\n",
    "\n",
    "\n",
    "\n",
    "            npregs=prege_.shape[1]\n",
    "            ge_bs = np.zeros((Nbs,npregs+postge_.shape[1]))\n",
    "            for ibs in range(Nbs):\n",
    "                #ans=np.apply_along_axis(lambda v: np.median(v[v!=0]), 0, x)\n",
    "                #ans[np.isnan(ans)]=0.\n",
    "                xbs = prege_[idx2pre[:,ibs],:]\n",
    "                pregbs  = np.apply_along_axis(lambda v: np.median(v[v!=0]), 0, xbs)\n",
    "                pregbs[np.isnan(pregbs)]=0.\n",
    "                xbs = postge_[idx2post[:,ibs],:]\n",
    "                postgbs = np.apply_along_axis(lambda v: np.median(v[v!=0]), 0, xbs)\n",
    "                postgbs[np.isnan(postgbs)]=0.\n",
    "\n",
    "                ge_bs[ibs,0:npregs] =pregbs\n",
    "                ge_bs[ibs,npregs:]  =postgbs\n",
    "\n",
    "\n",
    "            #df_hr2 = np.log2(df_hr2d)\n",
    "            #df_hr2 = pd.concat([df_hr2,df_hr2s],axis=1) # log2 normalization\n",
    "            #df_hr2 = df_hr2.replace({0:np.nan})\n",
    "            #df_hr2 = np.log2(df_hr2)\n",
    "\n",
    "            #df_hr2 = pd.concat([df_hr.iloc[1:,0],df_hr2],axis=1)\n",
    "\n",
    "            #df_hr2.dtypes\n",
    "            #df_hr2.head()\n",
    "\n",
    "            #df_hr2 = df_hr2.replace({0:np.nan})\n",
    "            #df_hrmg=df_hr2.groupby(['Class']) # logarithm_2 of UMI counts\n",
    "            #df_hrmg=df_hrmg.median(numeric_only=True) # median of log2(UMI)\n",
    "\n",
    "            #df_hrmg0 = df_hrmg.iloc[:,0:-1]\n",
    "\n",
    "            #df_hrmsg = pd.concat([df_hrmg.iloc[:,-1]]*df_hrmg0.shape[1],axis=1)\n",
    "            #df_hrmsg.columns = df_hrmg.iloc[:,0:-1].columns\n",
    "\n",
    "            #df_hrmg = df_hrmg.iloc[:,0:-1] - df_hrmsg +6*np.log2(10)  # transform from log2( UMI counts) to log2(CPM)\n",
    "\n",
    "            # ??? transform ge to median bootstraped log ???\n",
    "\n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "              \n",
    "\n",
    "        ## add data to the final dataset    \n",
    "        #gnij = np.repeat(np.array([syne3['cell_type2_pre'], syne3['cell_type2_post'],\n",
    "        #                           syne3['layer_pre'], syne3['layer_post'],\n",
    "        #                           syne3['cre_line_pre'], syne3['cre_line_post']]).reshape((1,6)),nsample,axis=0 )\n",
    "        #gnij = pd.DataFrame(gnij, columns = ['cell_type2_pre', 'cell_type2_post',\n",
    "        #                                     'layer_pre','layer_post',\n",
    "        #                                     'cre_line_pre', 'cre_line_post'] )\n",
    "        #prege  = pd.DataFrame(prege,columns  = pregs3 )\n",
    "        #postge = pd.DataFrame(postge,columns = postgs3 )\n",
    "\n",
    "\n",
    "        #ge_dataij = pd.concat([gnij, prege, postge ], axis = 1)\n",
    "        \n",
    "        \n",
    "        # add data to the final dataset    \n",
    "        gnij = np.repeat(np.array([syne3['cell_type2_pre'], syne3['cell_type2_post'],\n",
    "                                   syne3['layer_pre'], syne3['layer_post'],\n",
    "                                   syne3['cre_line_pre'], syne3['cre_line_post']]).reshape((1,6)),nsample,axis=0 )\n",
    "        gnij = pd.DataFrame(gnij, columns = ['cell_type2_pre', 'cell_type2_post',\n",
    "                                             'layer_pre','layer_post',\n",
    "                                             'cre_line_pre', 'cre_line_post'] )\n",
    "        \n",
    "        if  (do_bootstrap==0)|(do_median==0):\n",
    "            prege  = pd.DataFrame(prege,columns  = pregs3 )\n",
    "            postge = pd.DataFrame(postge,columns = postgs3 )\n",
    "\n",
    "            samples_pre = pd.DataFrame(sampl_pre.values[:gnij.shape[0]],columns = ['samples_pre'] )\n",
    "            samples_post = pd.DataFrame(sampl_post.values[:gnij.shape[0]],columns = ['samples_post'] )\n",
    "            \n",
    "            ge_dataij = pd.concat([gnij, prege, postge, samples_pre, samples_post ], axis = 1)\n",
    "            \n",
    "            \n",
    "        else:\n",
    "            ge_bs = pd.DataFrame(ge_bs,columns = pregs3 + postgs3 )\n",
    "\n",
    "\n",
    "            ge_dataij = pd.concat([gnij, ge_bs ], axis = 1)\n",
    "\n",
    "       \n",
    "        \n",
    "        \n",
    "        # add stp data\n",
    "        # stp_dataij =  df_stp.loc[[0,i], stp_type['stp_data_type']!=0]\n",
    "        stp_dataij =  df_stp.loc[[0,i], df_stp.loc[0, :]!=0]\n",
    "        \n",
    "\n",
    "        i2 = np.trunc(np.arange(len(stp_dataij.columns))/len(stp_columns)).astype(int)\n",
    "        stp_dataij  = stp_dataij.T.set_index(i2).pivot(columns=0,values=i).loc[:,stp_columns] #ss.set_index('par',append=True).unstack('par')\n",
    "        #stp_dataij =  stp_dataij.T.pivot(index=0, values=synij).T\n",
    "        \n",
    "        idx = np.mod(np.arange(nsample),len(stp_dataij.index))\n",
    "        stp_dataij=stp_dataij.iloc[idx,:].reset_index(drop=True)  \n",
    "        \n",
    "        i_data_set = pd.DataFrame(i+np.zeros(len(stp_dataij.index)), index = stp_dataij.index, columns = ['index_ds'])\n",
    "        stp_dataij=pd.concat([stp_dataij, i_data_set],axis=1)\n",
    "        \n",
    "        #if do_averagig_stp:\n",
    "        #    #stp_dataij.loc[:,:] = stp_dataij.mean(axis=0).values\n",
    "        #    stp_dataij.loc[:,:] = stp_dataij.median(axis=0).values\n",
    "\n",
    "        \n",
    "        #stp_dataij = pd.DataFrame({'presynaptic':pree[0], 'postsynaptic': poste[0], \n",
    "        #                           'stp_data':stp_dataij.iloc[0,:].values})\n",
    "        #stp_data  = stp_data.append(stp_dataij) \n",
    "        \n",
    "        # names of ge and stp columns\n",
    "\n",
    "        \n",
    "        ge_dataij  = pd.concat([ge_dataij,stp_dataij] , axis=1)\n",
    "        \n",
    "        # add to all ge and stp dataframe\n",
    "        ge_data  = pd.concat([ge_data, ge_dataij] , axis=0)\n",
    "        \n",
    "#stp_data =  stp_data.reset_index(drop=True)  \n",
    "\n",
    "\n",
    "\n",
    "#stp_data.loc[stp_data.loc[:,'postsynaptic']=='NGC','postsynaptic']='Lamp5'\n",
    "#stp_data2 = ge_data.iloc[:,0:3].copy(deep=True)                                 # modify!\n",
    "#stp_data2.columns=['presynaptic', 'postsynaptic', 'ppr']\n",
    "#print(stp_data2.head())\n",
    "#for i in stp_data.index: #stp_data.index:\n",
    "#    pre=stp_data.loc[i,'presynaptic']\n",
    "#    post=stp_data.loc[i,'postsynaptic']\n",
    "#    stp_data2.loc[stp_data2.loc[:,'presynaptic'].isin([pre])&stp_data2.loc[:,'postsynaptic'].isin([post]),'ppr']=stp_data.loc[i,'stp_data'] \n",
    "    \n",
    "ge_data = ge_data.reset_index()    \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# transform labels from TM to An:A1\n",
    "#f = 20 # Hz\n",
    "#N = 5\n",
    "#T = np.arange(N)*1000/f\n",
    "\n",
    "\n",
    "#As, ns, ps, dpp0, p0, tF, tD, A = STP_sim(ge_data,T)\n",
    "As = STP_sim_complex(ge_data,ge_data.shape[0],[])\n",
    "N  = As.shape[1]\n",
    "\n",
    "#As =  As[:,1:]/np.tile(As[:,0],(N-1,1)).T\n",
    "As =  As[:,1:]/np.tile(As[:,0],(N-1,1)).T\n",
    "\n",
    "#As =  As[:,1:]/np.tile(As[:,0],(N-1,1)).T\n",
    "\n",
    "# add new stp parameters to ge_data dataset\n",
    "stp_columns2 = ['A'+str(s)+'/A1' for s in range(2,N+1)]\n",
    "stp_data2 = pd.DataFrame(As,columns = stp_columns2, index = ge_data.index)\n",
    "\n",
    "ge_data = pd.concat([ge_data, stp_data2 ],axis=1)\n",
    "\n",
    "if do_averagig_stp:\n",
    "     #stp_dataij.loc[:,:] = stp_dataij.mean(axis=0).values\n",
    "     #stp_dataij.loc[:,:] = stp_dataij.median(axis=0).values\n",
    "    all_i_ds =  ge_data.loc[:,'index_ds']                             \n",
    "    for i_ds in list(set(all_i_ds)):\n",
    "              ids2 =  ge_data.loc[:,'index_ds'].isin([i_ds])  \n",
    "              ge_data.loc[ids2,stp_columns2] = ge_data.loc[ids2,stp_columns2].median(axis=0).values                   \n",
    "                                   \n",
    "                                   \n",
    "## filter TM parameters\n",
    "do_filter_pars = 0\n",
    "\n",
    "if do_filter_pars!=0:\n",
    "    print((ge_data.loc[:,stp_columns ]<0).sum().sum())\n",
    "    print(len(ge_data.index))\n",
    "\n",
    "    nannot=6\n",
    "    annot_columns = ge_columns[0:nannot]\n",
    "    ge_data2 = ge_data.loc[:,annot_columns + stp_columns2].groupby(annot_columns).mean()\n",
    "    f,ax = plt.subplots(figsize=(5, 5))\n",
    "    ge_data2.T.plot(ax=ax)\n",
    "    plt.gca().legend('')\n",
    "    #plt.gca().title('An:A1 before filtering data')\n",
    "    ax.title.set_text('An:A1 before filtering data')\n",
    "    #ge_data2 = ge_data.loc[:,annot_columns + stp_columns2].groupby(annot_columns).mean()\n",
    "    print('mean of stp before filtering:')\n",
    "    print(ge_data2)\n",
    "\n",
    "    #i_filter_out = (ge_data.loc[:,'dp/p0']*ge_data.loc[:,'p0']+)|(ge_data.loc[:,'tD'])\n",
    "    #i_filter_out = i_filter_out&((ge_data.loc[:,'tF'])|(ge_data.loc[:,'tF']))\n",
    "    i_filter_out = np.sum(ps2>1,axis=1)|np.sum(ps2<0,axis=1)\n",
    "    i_filter_out = i_filter_out|(np.sum(ns2>1,axis=1)|np.sum(ns2<0,axis=1))\n",
    "    ge_data = ge_data.loc[i_filter_out==False ,:]\n",
    "\n",
    "    ge_data2 = ge_data.loc[:,annot_columns + stp_columns2].groupby(annot_columns).mean()\n",
    "    f1,ax1 = plt.subplots(figsize=(5, 5))\n",
    "    ge_data2.T.plot(ax=ax1)\n",
    "    plt.gca().legend('')\n",
    "    #plt.gca().title('An:A1 after filtering data')\n",
    "    ax1.title.set_text('An:A1 after filtering data')\n",
    "    #ge_data2 = ge_data.loc[:,annot_columns + stp_columns2].groupby(annot_columns).mean()\n",
    "\n",
    "    print('mean of stp after filtering:')\n",
    "    print(ge_data2)\n",
    "    \n",
    "#print('mean of stp after filtering:')\n",
    "#print(ge_data2)\n",
    "\n",
    "\n",
    "print(len(ge_data.index))\n",
    "#ge_data.loc[:,stp_columns ] =ge_data.loc[:,stp_columns ].abs()  \n",
    "\n",
    "\n",
    "#\n",
    "## normalize stp parameters - for TM parameters\n",
    "do_normalize_stp_pars = 0;\n",
    "if do_normalize_stp_pars!=0:\n",
    "    ge_data.loc[:,stp_columns  ] =np.log(ge_data.loc[:,stp_columns  ].values.astype(float))\n",
    "\n",
    "#  assign training and labels \n",
    "stp_columns1 = stp_columns \n",
    "stp_columns = stp_columns2 # select training labels names\n",
    "\n",
    "nannot = 7\n",
    "nstp =len(stp_columns)\n",
    "\n",
    "annot_columns = ge_columns[0:6]+['index_ds']\n",
    "if do_samples_columns==1:\n",
    "    nannot = 9\n",
    "    annot_columns = annot_columns + ['samples_pre', 'samples_post']\n",
    "\n",
    "## averaging of stp for each synapse type \n",
    "\n",
    "\n",
    "#X = ge_data.loc[:,ge_data.columns[0:]].values\n",
    "X = ge_data.loc[:,annot_columns + ge_columns[6:]].values\n",
    "y = ge_data.loc[:,annot_columns + stp_columns].values \n",
    "#y = stp_data2.loc[:,:].values\n",
    "\n",
    "## normalize stp parameters - for TM parameters\n",
    "#y[:,nannot: ] =np.log(y[:,nannot: ].astype(float))\n",
    "\n",
    "\n",
    "    \n",
    "\n",
    "# print elapsed time\n",
    "t2 = time.time()\n",
    "#print(regr.feature_importances_)\n",
    "print('Elapsed time: '+str(t2-t1))\n",
    "\n",
    "print(y.shape)\n",
    "X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ge_data_h = ge_data\n",
    "X_h = X\n",
    "y_h = y\n",
    "annot_columns_h = annot_columns\n",
    "ge_columns_h = ge_columns\n",
    "stp_columns_h = stp_columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ge_data.loc[0:10*0+5,['tF', 'p0','tD','dp']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ge_data_h.loc[:,annot_columns_h[-1]]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df_aba_vis_in_c.columns\n",
    "len(ge_columns_h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ge_data_h.shape)\n",
    "print(X_h.shape)\n",
    "print(y_h.shape)\n",
    "print(len(annot_columns_h))\n",
    "print(len(ge_columns_h))\n",
    "print(len(stp_columns_h))\n",
    "print(ge_columns_h[300:])\n",
    "print(ge_columns_h[0:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ge_data_h.to_excel('ge_data_h.xlsx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##\n",
    "##\n",
    "##  CORTEX gene expression\n",
    "##\n",
    "##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_path = '' #'data/'\n",
    "\n",
    "#d4='~/Documents/references/transcriptomes to STP/transcriptomes/ABA/mouse_VISp_gene_expression_matrices_2018-06-14/'\n",
    "if do_laptop:\n",
    "    d4 = '~/Documents/references/lists_of_proteins_2/Analysis_of_proteins/genetic_types/transcriptomic_interneurons/'\n",
    "else:\n",
    "    d4='~/Documents/references/transcriptomes to STP/transcriptomes/ABA/mouse_VISp_gene_expression_matrices_2018-06-14/'\n",
    "\n",
    "fn='mouse_VISp_2018-06-14_samples-columns.csv'\n",
    "df_aba_vis_c=pd.read_csv(d4+fn)\n",
    "print(df_aba_vis_c.columns)\n",
    "print(df_aba_vis_c.head())\n",
    "\n",
    "fn = 'mouse_VISp_2018-06-14_genes-rows.csv'\n",
    "df_aba_g=pd.read_csv(d4+fn)\n",
    "print(df_aba_g.head())\n",
    "print(len(df_aba_vis_c.index))\n",
    "\n",
    "in_glut_l23 = (df_aba_vis_c.loc[:,'class']=='Glutamatergic') #&(df_aba_vis_c.loc[:,'brain_subregion'].isin(['L2/3','L4','L5']))\n",
    "\n",
    "in_gaba = df_aba_vis_c.loc[:,'class']=='GABAergic'\n",
    "print(sum(in_gaba))\n",
    "df_aba_vis_in_c=df_aba_vis_c.loc[in_gaba,:]\n",
    "df_aba_vis_l23glu_c=df_aba_vis_c.loc[in_glut_l23,:]\n",
    "\n",
    "###\n",
    "print('loading gene expression...')\n",
    "\n",
    "d14 = '/home/stepaniu/Documents/references/transcriptomes to STP/scVI_data/'\n",
    "if do_laptop:\n",
    "    d14 = '/Users/stepaniu/Documents/jan_2020/'\n",
    "    from scvi.dataset import LoomDataset, CsvDataset, Dataset10X, AnnDataset\n",
    "else:\n",
    "    from scvi.dataset import DownloadableAnnDataset\n",
    "\n",
    "if do_laptop:    \n",
    "##df_aba_ad.write('mouse_VISp_2018-06-14_exon-matrix_GABAergic.h5ad')\n",
    "    aba_vis_in = AnnDataset(\"mouse_VISp_2018-06-14_exon-matrix_GABAergic.h5ad\", new_n_genes = 45768,\n",
    "                                   save_path = save_path) \n",
    "else:\n",
    "    aba_vis_in = DownloadableAnnDataset(d14+\"data/mouse_VISp_2018-06-14_exon-matrix_GABAergic.h5ad\") \n",
    "\n",
    "if do_laptop:\n",
    "    #aba_vis_l23glu = AnnDataset(\"mouse_VISp_2018-06-14_exon-matrix_L23_Glutamatergic_hda.h5ad\", new_n_genes = 45768,\n",
    "    #                               save_path = save_path) \n",
    "    #aba_vis_l23glu = AnnDataset(\"mouse_VISp_2018-06-14_exon-matrix_L2345_Glutamatergic_hda.h5ad\", new_n_genes = 45768,\n",
    "    #                               save_path = save_path) \n",
    "\n",
    "    aba_vis_l23glu = AnnDataset(\"mouse_VISp_2018-06-14_exon-matrix_L23456_Glutamatergic_hda.h5ad\", new_n_genes = 45768,\n",
    "                                   save_path = '') \n",
    "else:\n",
    "    aba_vis_l23glu = DownloadableAnnDataset(d14+\"mouse_VISp_2018-06-14_exon-matrix_L23456_Glutamatergic_hda.h5ad\") \n",
    "\n",
    "diff_order = sum(aba_vis_in.gene_names[0:].astype(int) - np.arange(len(aba_vis_in.gene_names)))\n",
    "print(diff_order)\n",
    "\n",
    "aba_vis_in.gene_names = df_aba_g.loc[:,'gene_symbol'].values\n",
    "\n",
    "aba_vis_l23glu.gene_names = df_aba_g.loc[:,'gene_symbol'].values\n",
    "\n",
    "\n",
    "\n",
    "print(len(aba_vis_in.gene_names))\n",
    "print(len(aba_vis_l23glu.gene_names))\n",
    "len(list(set(aba_vis_l23glu.gene_names).difference(set(aba_vis_in.gene_names))))\n",
    "print(len(df_aba_vis_l23glu_c.index))\n",
    "print(len(df_aba_vis_l23glu_c.index))\n",
    "print(len(df_aba_vis_in_c.index))\n",
    "\n",
    "from scvi.dataset.dataset import GeneExpressionDataset\n",
    "if do_laptop:\n",
    "    vis_dat = GeneExpressionDataset.concat_datasets(aba_vis_in,aba_vis_l23glu)    \n",
    "else:    \n",
    "    vis_dat  = GeneExpressionDataset()\n",
    "    vis_dat.populate_from_datasets([aba_vis_in,aba_vis_l23glu]) \n",
    "   \n",
    "\n",
    "print((vis_dat.batch_indices==0).sum())\n",
    "(vis_dat.batch_indices==1).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### replace some genes expression with scvi imputed ge, add scvi latent factors to data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# flags were defined for hippocampus:\n",
    "#do_replace_to_scVI_imputed=1      # replace raw gene expression data with their scvi imputation\n",
    "#do_add_scvi_latent_factors=1      # add scvi latent factors data\n",
    "#gs_name = '_gs219'\n",
    "if (do_replace_to_scVI_imputed==1)|(do_add_scvi_latent_factors==1):\n",
    "    scVI_all = pd.read_hdf('scVI_latent_factors_and_imputation_aba2019'+gs_name+'.hdf',key='data') # this is 290 genes dataset\n",
    "    \n",
    "    scVI_all = scVI_all.set_index('sample_name') # scvi imputed genes and latent factors for hipp., enth. and vis. cortex\n",
    "\n",
    "    # were defined before\n",
    "    #nlf = 10\n",
    "    #gncolumns = scVI_all.columns[nlf+1:-1] # gene names columns in scvi data\n",
    "    #lfcolumns = scVI_all.columns[:nlf]     # latent factors columns in scvi data\n",
    "\n",
    "    \n",
    "    smhi = aba_vis_in.obs\n",
    "    smhi = pd.DataFrame(np.arange(smhi.shape[0]).transpose().reshape((-1,1)),index = smhi.index.values, columns=['index'])\n",
    "    #smhi\n",
    "    smhi = smhi.reset_index().set_index('index')  # cortex interneurons samples and their indexes in ge-data\n",
    "    smhi.columns=['samples']\n",
    "    smh2i = smhi.reset_index().set_index('samples')  # cortex interneurons samples and their indexes in ge-data\n",
    "    #smh2i\n",
    "    \n",
    "    smhe = aba_vis_l23glu.obs\n",
    "\n",
    "    smhe = pd.DataFrame(smhe.reset_index(drop=True).loc[:,'sample_name'])\n",
    "    smhe.columns = ['samples']\n",
    "    #smhe\n",
    "\n",
    "    smh2e = smhe.reset_index().set_index('samples')  # cortex interneurons samples and their indexes in ge-data\n",
    "    #smh2e\n",
    "    \n",
    "    #smh = aba_vis_in.obs\n",
    "    smh2 = pd.concat([smh2i, smh2e+smh2i.shape[0]])\n",
    "    #smh2\n",
    "    \n",
    "    gnscvi = pd.DataFrame(gncolumns)   # df from scvi genenames\n",
    "\n",
    "    # assumption ??: the same gene names are in vis_dat, aba_vis_l23glu and aba_vis_in\n",
    "    gnh = vis_dat.gene_names               # cortex vis_dat gene names (the same names are in inh and exc data)\n",
    "    #smh = vis_dat.obs                     # cortex vis_dat samples names\n",
    "    gnh2 = pd.DataFrame(gnh).reset_index().set_index(0)   \n",
    "    ignh2 = gnh2.loc[gnscvi.loc[:,0],:].loc[:,'index'].values  # indices of scvi gene names in vis_dat gene names\n",
    "\n",
    "    gnhi = aba_vis_in.gene_names           # cortex aba_vis_in gene names (the same names are in inh and exc data)\n",
    "    #smh = vis_dat.obs                     # cortex aba_vis_in samples names\n",
    "    gnh2i = pd.DataFrame(gnhi).reset_index().set_index(0)   \n",
    "    ignh2i = gnh2i.loc[gnscvi.loc[:,0],:].loc[:,'index'].values  # indices of scvi gene names in aba_vis_in gene names\n",
    "\n",
    "    gnhe = aba_vis_l23glu.gene_names              # cortex aba_vis_l23glu gene names (the same names are in inh and exc data)\n",
    "    #smh = aba_vis_l23glu.obs                     # cortex aba_vis_l23glu samples names\n",
    "    gnh2e = pd.DataFrame(gnhe).reset_index().set_index(0)   \n",
    "    ignh2e = gnh2e.loc[gnscvi.loc[:,0],:].loc[:,'index'].values  # indices of scvi gene names in aba_vis_l23glu gene names\n",
    "    \n",
    "    \n",
    "    if (do_replace_to_scVI_imputed==1):\n",
    "        # correct this!\n",
    "        vis_dat._X[:,ignh2] = scVI_all.loc[smh2.index,gncolumns].values  # modified vis_dat matrix (imputed part of genes replaced)\n",
    "        aba_vis_l23glu._X[:,ignh2e] = scVI_all.loc[smh2e.index,gncolumns].values  # modified vis_dat matrix (imputed part of genes replaced)\n",
    "        aba_vis_in._X[:,ignh2i] = scVI_all.loc[smh2i.index,gncolumns].values  # modified vis_dat matrix (imputed part of genes replaced)\n",
    "    \n",
    "    \n",
    "    if (do_add_scvi_latent_factors==1):\n",
    "        scVI_lf_ci = scVI_all.loc[smh2i.index,lfcolumns]    # latent factors for hippocampal samples\n",
    "        scVI_lf_ce = scVI_all.loc[smh2e.index,lfcolumns]\n",
    "        scVI_lf_c = pd.concat([scVI_lf_ci,scVI_lf_ce])\n",
    "        vis_dat._X = np.concatenate([vis_dat._X, scVI_lf_c.values],axis=1)\n",
    "        # correct this!\n",
    "        \n",
    "        aba_vis_in._X = np.concatenate([aba_vis_in._X, scVI_lf_ci.values],axis=1)\n",
    "        aba_vis_l23glu._X = np.concatenate([aba_vis_l23glu._X, scVI_lf_ce.values],axis=1)\n",
    "        \n",
    "        #lf_scvi_columns = lfcolumns.astype(str)\n",
    "        vis_dat.gene_names = np.concatenate([vis_dat.gene_names, lf_scvi_columns])\n",
    "        # correct this!\n",
    "        \n",
    "        aba_vis_in.gene_names = np.concatenate([aba_vis_in.gene_names, lf_scvi_columns])\n",
    "        aba_vis_l23glu.gene_names = np.concatenate([aba_vis_l23glu.gene_names, lf_scvi_columns])\n",
    "        \n",
    "        #pregs = pregs + lf_scvi_columns.tolist()\n",
    "        #postgs = postgs + lf_scvi_columns.tolist()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(pregs))\n",
    "print(len(postgs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "i='Sst-IRES-Cre'\n",
    "li = df_aba_vis_in_c.loc[(df_aba_vis_in_c.loc[:,'driver_lines']==i),\n",
    "                             ['subclass','cluster','brain_subregion']]\n",
    "set(li.loc[:,'brain_subregion'])\n",
    "for reg in list(set(li.loc[:,'brain_subregion'])):\n",
    "    #reg = 'L4'\n",
    "    li = df_aba_vis_in_c.loc[(df_aba_vis_in_c.loc[:,'driver_lines']==i)&(df_aba_vis_in_c.loc[:,'brain_subregion']==reg),\n",
    "                             ['subclass','cluster','brain_subregion']]\n",
    "\n",
    "    f = plt.figure()\n",
    "    ax = f.add_axes()\n",
    "    plt.title(str(i)+' '+reg) \n",
    "    vc = li.loc[:,'cluster'].value_counts() \n",
    "\n",
    "    vc2 = vc/vc.sum()*100\n",
    "\n",
    "    if len(vc.index)>0: \n",
    "        vc2.plot(kind='bar',ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i='Pvalb-IRES-Cre'\n",
    "li = df_aba_vis_in_c.loc[(df_aba_vis_in_c.loc[:,'driver_lines']==i),\n",
    "                             ['subclass','cluster','brain_subregion']]\n",
    "set(li.loc[:,'brain_subregion'])\n",
    "for reg in list(set(li.loc[:,'brain_subregion'])):\n",
    "    #reg = 'L4'\n",
    "    li = df_aba_vis_in_c.loc[(df_aba_vis_in_c.loc[:,'driver_lines']==i)&(df_aba_vis_in_c.loc[:,'brain_subregion']==reg),\n",
    "                             ['subclass','cluster','brain_subregion']]\n",
    "\n",
    "    f = plt.figure()\n",
    "    ax = f.add_axes()\n",
    "    plt.title(str(i)+' '+reg) \n",
    "    vc = li.loc[:,'cluster'].value_counts() \n",
    "\n",
    "    vc2 = vc/vc.sum()*100\n",
    "\n",
    "    if len(vc.index)>0: \n",
    "        vc2.plot(kind='bar',ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##\n",
    "##\n",
    "##  CORTEX STP\n",
    "##\n",
    "##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "if do_laptop:\n",
    "    d5 = '~/Documents/references/lists_of_proteins_2/Analysis_of_proteins/genetic_types/transcriptomic_interneurons/additional cortical STP data/'\n",
    "else:\n",
    "    d5 = '/home/stepaniu/Documents/references/transcriptomes to STP/scVI_data/additional cortical STP data/'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled.xlsx'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled_UDF2.xlsx'\n",
    "fn5='STP_cortex_interneurons_Cre_labeled_UDF3.xlsx'\n",
    "fn5='STP_cortex_interneurons_Cre_labeled_UDF5_4pars.xlsx'\n",
    "\n",
    "df_stp=pd.read_excel(d5+fn5)\n",
    "df_stp.loc[df_stp.loc[:,'stp_comment'].isna() ,'stp_comment']=''\n",
    "\n",
    "############   \n",
    "do_averagig_stp =0\n",
    "stp_TM_names=['tF', 'p0','tD','dp/p0','A']\n",
    "stp_TM_start = 29\n",
    "ncol = len(df_stp.columns)\n",
    "df_tm_name = pd.DataFrame(np.zeros(ncol).reshape((1,ncol)), columns =df_stp.columns )\n",
    "df_stp = pd.concat([df_tm_name, df_stp],axis=0).reset_index(drop=True)\n",
    "\n",
    "for i,nm in enumerate(stp_TM_names):\n",
    "    ii = stp_TM_start + np.arange(0,ncol-i-stp_TM_start,len(stp_TM_names))+i\n",
    "    df_stp.iloc[0,ii] = nm\n",
    "\n",
    "#df_stp.loc[df_stp.loc[:,'name'].str.contains('Jaing_2015_PC')!=True,:].reset_index(drop=True)\n",
    "\n",
    "################# ADD ABA STP data\n",
    "do_add_aba_synphys=True\n",
    "do_add_aba_synphys_new=True\n",
    "\n",
    "#QTX.to_hdf('fit_aba_2019_A1_8_fited_pulses_syntypes_results','data')\n",
    "d_aba = '/home/stepaniu/Documents/references/transcriptomes to STP/'\n",
    "if do_laptop:\n",
    "    d_aba = '/Users/stepaniu/Documents/jan_2020/'\n",
    "\n",
    "# ACHTUNG!!! ACHTUNG!!!ACHTUNG!!!ACHTUNG!!!\n",
    "do_simple_TM_model = False  # simple tm vs tm+smr  \n",
    "# ACHTUNG!!! ACHTUNG!!!ACHTUNG!!!ACHTUNG!!!\n",
    "\n",
    "if do_simple_TM_model:    \n",
    "    STP_aba = pd.read_hdf(d_aba + 'fit_aba_2019_A1_8_fited_pulses_syntypes_results') # tm5\n",
    "    QTX2 = STP_aba.iloc[0:43,0:14]\n",
    "    QTX3 = STP_aba.iloc[43:,14:]\n",
    "    QTX3.index = QTX2.index\n",
    "    STP_aba = pd.concat([QTX2, QTX3],axis=1)\n",
    "else:\n",
    "    #STP_aba = pd.read_hdf(d_aba + 'fit_aba_2019_A1_8_fited_pulses_TM_8_syntypes_results') # tm5 + FDR\n",
    "    STP_aba = pd.read_hdf(d_aba + 'fit_aba_2019_mean_fited_pulses_TM_SMP0_100bs_syntypes_results') # tm5 + SMR\n",
    "    #STP_aba = pd.read_hdf(d_aba + 'fit_aba_2019_mean_fited_pulses_TM_SMP0_50bs_syntypes_results') \n",
    "\n",
    "STP_aba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_stp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "STP_aba.index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_stp_cells = STP_aba.loc[:,'comments'].str.split(pat='=', n=-1, expand=True)[1].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "f, ax =plt.subplots(figsize=(12, 5))\n",
    "#    plt.title(str(i0)+ '  '+str_syntype + '  ' +stp_columns2[3])\n",
    "#    plt.xlim((0.25,1.25))\n",
    "\n",
    "n_stp_cells.plot(kind = 'bar',ax=ax)\n",
    "n_stp_cells.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_stp_cells2 = n_stp_cells.reset_index()\n",
    "n_stp_cells3 = n_stp_cells2.loc[:,'index'].str.split(pat=';', n=-1, expand=True)\n",
    "n_stp_cells3[0] = n_stp_cells3[0].str.strip()\n",
    "n_stp_cells3[1] = n_stp_cells3[1].str.strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i2 = ['nr5a1 L2/3; nr5a1 L2/3', 'unknown L2/3; unknown L2/3',\n",
    "      'rorb L4; rorb L4', 'rbp4 L4; rbp4 L4',\n",
    "      'sim1 L5; sim1 L5', 'tlx3 L5; tlx3 L5', 'tlx3 L5; fam84b L5', \n",
    "      'ntsr1 L6; ntsr1 L6', \n",
    "       'unknown L2/3; pvalb L2/3','unknown L2/3; vip L2/3',\n",
    "       'nr5a1 L4; pvalb L4', 'nr5a1 L4; sst L4','nr5a1 L4; vip L4',\n",
    "       'sim1 L5; pvalb L5','sim1 L5; sst L5', 'tlx3 L5; sst L5', 'tlx3 L5; pvalb L2/3',\n",
    "       'ntsr1 L6; pvalb L6','ntsr1 L6; sst L6',\n",
    "        'pvalb L4; nr5a1 L4', 'pvalb L4; nr5a1 L5','pvalb L5; sim1 L5','pvalb L5; tlx3 L5',\n",
    "       'pvalb L6; ntsr1 L6',\n",
    "       'sst L2/3; sim1 L5','sst L4; nr5a1 L4','sst L5; tlx3 L5',\n",
    "       'pvalb L2/3; pvalb L2/3','pvalb L4; pvalb L2/3','pvalb L5; pvalb L5','pvalb,sst L5; pvalb L5',\n",
    "      'pvalb L6; pvalb L6','pvalb L2/3; sst L2/3','pvalb L2/3; vip L2/3', \n",
    "       'sst L6; pvalb L5','sst L2/3; sst L2/3','sst L5; sst L5','sst L6; sst L6','sst L2/3; vip L2/3','sst L4; vip L4',\n",
    "       'vip L2/3; sst L2/3','vip L4; sst L5','vip L2/3; vip L4',\n",
    "     ]\n",
    "\n",
    "n_stp_cells4 = n_stp_cells.loc[i2]\n",
    "n_stp_cells4 = n_stp_cells4.reset_index()\n",
    "is_uu = n_stp_cells4.loc[:,'index']=='unknown L2/3; unknown L2/3'\n",
    "n_stp_cells4.loc[is_uu,'index']='unknown L2/3; nr5a1 L2/3'\n",
    "n_stp_cells4 = n_stp_cells4.set_index('index')\n",
    "#%matplotlib inline\n",
    "#import matplotlib.pyplot as plt\n",
    "f, ax =plt.subplots(figsize=(15, 5))\n",
    "#    plt.title(str(i0)+ '  '+str_syntype + '  ' +stp_columns2[3])\n",
    "#    plt.xlim((0.25,1.25))\n",
    "\n",
    "n_stp_cells4.plot(kind = 'bar',ax=ax)\n",
    "\n",
    "plt.xticks(fontsize=16)\n",
    "plt.yticks(fontsize=16)\n",
    "plt.ylabel('Number of cells',fontsize=20)\n",
    "plt.xlabel('Synapse type',fontsize=20)\n",
    "plt.legend('')\n",
    "print(set(n_stp_cells.index).difference(set(i2)))\n",
    "n_stp_cells4.sum()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%matplotlib inline\n",
    "# transform labels from TM to An:A1\n",
    "fs = [20, 50, 10] # Hz\n",
    "N = 8\n",
    "\n",
    "\n",
    "\n",
    "#Trec = [125, 250, 500, 1000, 2000]\n",
    "Trec = [ 250, 500, 1000]\n",
    "DT0 = 25000\n",
    "#xs  =ge_data.iloc[l_pre_post2:,:].loc[:,stp_aba_names].values\n",
    "\n",
    "idx = 'vip L2/3; sst L2/3'\n",
    "xs  = STP_aba.loc[idx,'differential_evolution best x']\n",
    "\n",
    "xs  =np.delete(xs, [5,6],axis=1)\n",
    "As=np.zeros((l_pre_post2,0))\n",
    "As2=np.zeros((xs.shape[0],0))\n",
    "if do_add_aba_synphys:\n",
    "    \n",
    "    for f in fs:\n",
    "        T = np.arange(N)*1000/f\n",
    "        \n",
    "        Asf, ns, ps, dpp0, p0, tF, tD, A = STP_sim(ge_data.iloc[0:l_pre_post2,:],T)\n",
    "        As = np.concatenate([As,Asf],axis=1)\n",
    "        \n",
    "        for ri in range(len(Trec)):\n",
    "            Asr, nsr, psr, dpp0, p0, tF, tD, A = STP_sim(ge_data.iloc[0:l_pre_post2,:],[Trec[ri]],init_state=[ns[:,-1],ps[:,-1]])\n",
    "            As = np.concatenate([As,Asr],axis=1)\n",
    "        \n",
    "\n",
    "        As2f = np.zeros((xs.shape[0],N+len(Trec)))\n",
    "        for i2 in range(xs.shape[0]):\n",
    "            #as2, sts2 = STP_sim2(xs[i2,:],np.arange(8)*50,model_type = 'tm5') # preconditioning series???\n",
    "            #as2, sts2 = STP_sim2(xs[i2,:], T+DT0, init_state=sts2[-1], model_type='tm5' ) \n",
    "            \n",
    "            as2, sts2 = STP_sim2(xs[i2,:],T,model_type = 'tm5_smr')\n",
    "            \n",
    "            As2f[i2,0:N] = as2\n",
    "            for ri in range(len(Trec)):\n",
    "                as2r, sts2r = STP_sim2(xs[i2,:], [Trec[ri]], init_state=sts2[-1], model_type='tm5_smr' ) \n",
    "                As2f[i2,N+ri] = as2r\n",
    "                \n",
    "        As2 = np.concatenate([As2,As2f],axis=1)\n",
    "            \n",
    "    \n",
    "    As2 = As2/As2[:,0].reshape((-1,1))\n",
    "    As = As/As[:,0].reshape((-1,1))\n",
    "    As = np.concatenate([As,As2],axis=0)\n",
    "else:\n",
    "    As, ns, ps, dpp0, p0, tF, tD, A = STP_sim(ge_data,T)\n",
    "    \n",
    "\n",
    "import matplotlib.pyplot as plt    \n",
    "plt.plot(As2[:,:].transpose(),'o-')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "do_samples_columns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CORTEX : allign gene expression and stp data (for selected gene sets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# allign gene expression and stp data (for selected gene sets)\n",
    "import re\n",
    "import time\n",
    "t1 = time.time()\n",
    "\n",
    "\n",
    "\n",
    "# names of gene expression based synapses types\n",
    "\n",
    "glu_l23=list(set(df_aba_vis_l23glu_c.loc[:,'cluster']))\n",
    "\n",
    "dsg = [['PC']+glu_l23,\n",
    "       ['Pvalb','Pvalb Tpbg','Pvalb Reln Itm2a','Pvalb Reln Tac1','Pvalb Sema3e Kank4'], \n",
    "       ['Sst','Sst Calb2 Pdlim5','Sst Calb2 Necab1', 'Sst Hpse Cbln4','Sst Hpse Sema3c','Sst Nr2f2 Necab1'],\n",
    "       ['Vip','Vip Lmo1 Myl1', 'Vip Rspo1 Itga4','Vip Ptprt Pkp2','Vip Rspo4 Rxfp1 Chat'],\n",
    "       ['Lamp5','Lamp5 Plch2 Dock5'],\n",
    "       ['MC','L23MC','Sst Calb2 Pdlim5'],\n",
    "        'L5MC']\n",
    "      \n",
    "\n",
    "set(df_stp.loc[:,'synapse_type_2'])\n",
    "\n",
    "# names of electrophysiology based synapses types\n",
    "dse = [['PC', 'L23P', 'L23PC', 'L5P'],\n",
    "       ['Pvalb', 'PV', 'L23BC', 'L5BC'], \n",
    "       ['Sst', 'L23MC', 'L5MC', 'Sst', 'MC'], \n",
    "       ['Vip', 'VIP', 'BTC'],\n",
    "       ['NGC', 'L23NGC'],\n",
    "       ['L23MC' ],\n",
    "       ['L5MC']]\n",
    "\n",
    "\n",
    "\n",
    "# load stp data table\n",
    "if do_laptop:\n",
    "    d5 = '~/Documents/references/lists_of_proteins_2/Analysis_of_proteins/genetic_types/transcriptomic_interneurons/additional cortical STP data/'\n",
    "else:\n",
    "    d5 = '/home/stepaniu/Documents/references/transcriptomes to STP/scVI_data/additional cortical STP data/'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled.xlsx'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled_UDF2.xlsx'\n",
    "#fn5='STP_cortex_interneurons_Cre_labeled_UDF3.xlsx'\n",
    "\n",
    "do_3pars_TM=0 # 3 or 4 parametric TM model?\n",
    "if do_3pars_TM:\n",
    "    fn5='STP_cortex_interneurons_Cre_labeled_UDF5_3pars.xlsx'\n",
    "else:\n",
    "    fn5='STP_cortex_interneurons_Cre_labeled_UDF5_4pars.xlsx'\n",
    "\n",
    "df_stp=pd.read_excel(d5+fn5)\n",
    "df_stp.loc[df_stp.loc[:,'stp_comment'].isna() ,'stp_comment']=''\n",
    "\n",
    "# drop TC->\n",
    "TC_names = ['Miao_2016_fig1_TC->PC','Miao_2016_fig1_TC->PV']\n",
    "df_stp = df_stp.drop(index = df_stp.index[df_stp.loc[:,'name'].isin(TC_names)] )\n",
    "\n",
    "############   \n",
    "do_averagig_stp =0\n",
    "stp_TM_names=['tF', 'p0','tD','dp/p0','A']\n",
    "stp_TM_start = 29\n",
    "ncol = len(df_stp.columns)\n",
    "df_tm_name = pd.DataFrame(np.zeros(ncol).reshape((1,ncol)), columns =df_stp.columns )\n",
    "df_stp = pd.concat([df_tm_name, df_stp],axis=0).reset_index(drop=True)\n",
    "\n",
    "for i,nm in enumerate(stp_TM_names):\n",
    "    ii = stp_TM_start + np.arange(0,ncol-i-stp_TM_start,len(stp_TM_names))+i\n",
    "    df_stp.iloc[0,ii] = nm\n",
    "\n",
    "\n",
    "\n",
    "print(df_stp.head())\n",
    "#############\n",
    "##Remove pure PPR data (Jaing_2015_PC)\n",
    "#df_stp = df_stp.loc[df_stp.loc[:,'name'].str.contains('Jaing_2015_PC')!=True,:].reset_index(drop=True) \n",
    "\n",
    "############\n",
    "\n",
    "# batches and gene sets\n",
    "pre_post = df_stp.loc[:,'synapse_type_2'].str.split(pat='->',expand=True)\n",
    "\n",
    "if do_add_aba_synphys:\n",
    "    pre_post_aba = STP_aba.loc[:,'synapse_type_2'].str.split(pat=';',expand=True)\n",
    "\n",
    "vbi = vis_dat.batch_indices.reshape((vis_dat.batch_indices.shape[0],))\n",
    "vbi=vbi>0\n",
    "\n",
    "df_gs  = pd.DataFrame( vis_dat.gene_names ) \n",
    "in_pregs = df_gs.loc[:,0].isin(pregs).values\n",
    "in_postgs = df_gs.loc[:,0].isin(postgs).values\n",
    "\n",
    "pregs2 = df_gs.loc[in_pregs,0].values\n",
    "postgs2 = df_gs.loc[in_postgs,0].values\n",
    "\n",
    "\n",
    "do_2=1;\n",
    "# select stp data for appropriate frequency and ages \n",
    "stp_type = {'name':['Walker_2016_PV->MC L23',\n",
    "                    'Walker_2016_VIP->MC L23',\n",
    "                    'Yuste_2016_taus_PV->Sst L23 recalc',\n",
    "                    'Yuste_2016_taus_PV->VIP L23 recalc',\n",
    "                    'Yuste_2016_taus_Sst->VIP L23 recalc',\n",
    "                    'Yuste_2016_taus_VIP->Sst L23 recalc',\n",
    "                    'Yuste_2016_taus_VIP->VIP L23 recalc'],\n",
    "            'name_remove': [], \n",
    "            'stp_freq': [20, 40, 50], 'stp_comment':  ['', ' P20-21', ' P20-22', ' P20-23', ' P25-31'],\n",
    "            'stp_data_type': df_stp.loc[0, :]!=0 }\n",
    "\n",
    "if do_add_aba_synphys:\n",
    "    STP_aba_type = {'name':[],\n",
    "            'stp_freq': [], 'stp_comment':  [],\n",
    "            'stp_data_type': 'differential_evolution best x' }  # 'random_starts best x'\n",
    "    if do_simple_TM_model:\n",
    "        stp_aba_names = ['tF', 'p0','tD','dp','A','A1','A2'] # tm5\n",
    "    else:\n",
    "        stp_aba_names = ['tF', 'p0','tD','dp','A','A1','A2','tDmin','dd','t_FDR','t_SMR','dp0'] # tm5+smr\n",
    "    \n",
    "    # ACHTUNG!!! ACHTUNG!!! ACHTUNG!!!\n",
    "    do_shrink_stp_distribution = False # # shrink bootstrap distribution of stp parameters towords mean values - to test stp variability reduction\n",
    "    if do_shrink_stp_distribution==True:\n",
    "        shrink_stp_distribution_coef = 0*0.1\n",
    "    \n",
    "if do_2==1:\n",
    "    #stp_type = {'name':['Walker_2016_PV->MC L23',\n",
    "    #               'Walker_2016_VIP->MC L23',\n",
    "    #                'Yuste_2016_taus_PV->Sst L23 recalc',\n",
    "    #                'Yuste_2016_taus_VIP->Sst L23 recalc'],\n",
    "    #        'stp_freq': [20, 40, 50], 'stp_comment':  ['', ' P20-21', ' P20-22', ' P20-23', ' P25-31'],\n",
    "    #        'stp_data_type': df_stp.loc[0, :]!=0 }\n",
    "    \n",
    "    # without Yuste and\n",
    "    stp_type = {'name':['Walker_2016_PV->MC L23',\n",
    "                    'Walker_2016_VIP->MC L23'],\n",
    "             'name_remove': ['Jaing','Yuste'],  \n",
    "            'stp_freq': [20, 40, 50], 'stp_comment':  ['', ' P20-21', ' P20-22', ' P20-23', ' P25-31'],\n",
    "            'stp_data_type': df_stp.loc[0, :]!=0 }\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "in_stp_data = df_stp.loc[:,'stp_freq'].isin([20])\n",
    "in_stp_data= in_stp_data|(df_stp.loc[:,'stp_freq'].isin([40,50])&df_stp.loc[:,'name'].isin(stp_type['name']))\n",
    "in_stp_data = in_stp_data&df_stp.loc[:,'stp_comment'].isin(pd.Series(stp_type['stp_comment']).str.strip())\n",
    "if len(stp_type['name_remove'])>0:\n",
    "        remove = df_stp.loc[:,'name'].str.contains(stp_type['name_remove'][0])\n",
    "        for rmn in stp_type['name_remove']:\n",
    "             remove= remove|df_stp.loc[:,'name'].str.contains(rmn)\n",
    "        in_stp_data = in_stp_data&( remove!=True)\n",
    "        \n",
    "#show selected stp dataset\n",
    "ddd=df_stp.loc[in_stp_data,['name','synapse_type_2', 'stp_freq','area']]\n",
    "print(ddd)\n",
    "print(in_stp_data.sum())\n",
    "\n",
    "synij = np.repeat(False,in_stp_data.shape[0])\n",
    "if do_add_aba_synphys:\n",
    "    synij = np.repeat(False,in_stp_data.shape[0]+STP_aba.shape[0])\n",
    "\n",
    "cols2 = ['name','perc']\n",
    "cl2post =pd.DataFrame([],columns= cols2)\n",
    "cl2pre =pd.DataFrame([],columns= cols2)\n",
    "\n",
    "\n",
    "# strip cluster names\n",
    "#list(set(df_aba_vis_in_c.loc[df_aba_vis_in_c.loc[:,'subclass'].str.contains(clns6),'cluster']))\n",
    "df_aba_vis_in_c.loc[:,'cluster']=df_aba_vis_in_c.loc[:,'cluster'].str.strip()\n",
    "df_aba_vis_l23glu_c.loc[:,'cluster']=df_aba_vis_l23glu_c.loc[:,'cluster'].str.strip()\n",
    "\n",
    "do_simplfy_clusters = 1 # do not take into account minor subclasses\n",
    "Sst_NonMC ='Sst, Sst Calb2 Pdlim5-, Sst Chrna2 Glra3-, Sst Chrna2 Ptgdr-, Sst Myh8 Etv1-, Sst Tac2 Myh4-, Sst Nr2f2 Necab1-, Sst Chodl-,  Sst Calb2 Necab1-, Sst Tac1 Htr1d-, Sst Myh8 Fibin-';\n",
    "\n",
    "pregs3 = ['pre__'+s for s in pregs2.tolist()]\n",
    "postgs3 = ['post__'+s for s in postgs2.tolist()]\n",
    "ge_columns = ['cell_type2_pre', 'cell_type2_post',\n",
    "                                             'layer_pre','layer_post',\n",
    "                                             'cre_line_pre', 'cre_line_post'] +pregs3 +postgs3 + ['samples_pre', 'samples_post']\n",
    "\n",
    "stp_columns = stp_TM_names #list(set(df_stp.loc[0,:]).difference(set([0])))\n",
    "\n",
    "# output datasets\n",
    "stp_data = pd.DataFrame([], columns =['presynaptic', 'postsynaptic', 'stp_data'] )\n",
    "ge_data = pd.DataFrame([], columns =ge_columns+stp_columns+['index_ds'])\n",
    "\n",
    "if do_add_aba_synphys:\n",
    "    l_pre_post=len(pre_post.index)+len(pre_post_aba.index)\n",
    "else:\n",
    "    l_pre_post=len(pre_post.index)\n",
    "l_pre_post1 = len(pre_post.index)\n",
    "\n",
    "for i in range(1,l_pre_post): #range(1,len(pre_post.index)):# range(len(pre_post.index)): # [1]: #range(len(pre_post.index)):\n",
    "    if i==l_pre_post1:\n",
    "        l_pre_post2 = ge_data.shape[0]\n",
    "        \n",
    "    # syne3 - all available pre- and post- synaptic cells annotations : cre-lines, layers, morphologies\n",
    "    if i<l_pre_post1:\n",
    "        l = str.strip(df_stp.loc[i,'area'])\n",
    "        if type(df_stp.loc[i,'synapse_type'])!=float:\n",
    "            syne1 = re.split('\\;',df_stp.loc[i,'synapse_type'])\n",
    "        else:\n",
    "            syne1 =''    \n",
    "        #syne2 = df_stp.loc[i,'synapse_type2'].str.split(pat='->',expand=True)\n",
    "        syne2 = re.split('->',df_stp.loc[i,'synapse_type_2'])\n",
    "        if len(syne2)==1:\n",
    "            pat ='→'\n",
    "            syne2 = re.split(pat,df_stp.loc[i,'synapse_type_2'])\n",
    "\n",
    "        syne3 ={'layer_pre':l,  'cre_line_pre':'none', 'cell_type2_pre':str.strip(syne2[0]),'cluster_pre':[],\n",
    "                'layer_post':l, 'cre_line_post':'none','cell_type2_post':str.strip(syne2[1]),'cluster_post':[],  } \n",
    "\n",
    "        # add cre-lines when available\n",
    "        if len(syne1)>1:\n",
    "            syne3['cre_line_pre']  =str.strip(syne1[0])\n",
    "            syne3['cre_line_post'] =str.strip(syne1[1])\n",
    "\n",
    "        # introduce study-specific modifications to syne3 (using morphology+layer->clusters map)\n",
    "        if  df_stp.loc[i,'name'][0:5]=='Jaing':\n",
    "\n",
    "            # modify layer_post and reduce cell_type_post to mtype name\n",
    "            if syne3['cell_type2_post'][0:3]=='L23':\n",
    "                syne3['layer_post'] = 'L2/3'\n",
    "                syne3['cell_type2_post']=syne3['cell_type2_post'][3:]\n",
    "            elif syne3['cell_type2_post'][0:2]=='L5':\n",
    "                     syne3['layer_post'] = 'L5'\n",
    "                     syne3['cell_type2_post']=syne3['cell_type2_post'][2:]  \n",
    "\n",
    "            #  reduce cell_type_pre to mtype name\n",
    "            if syne3['cell_type2_pre'][0:3]=='L23':\n",
    "                syne3['layer_pre'] = 'L2/3'\n",
    "                syne3['cell_type2_pre']=syne3['cell_type2_pre'][3:]\n",
    "            elif syne3['cell_type2_pre'][0:2]=='L5':\n",
    "                     syne3['layer_pre'] = 'L5'\n",
    "                     syne3['cell_type2_pre']=syne3['cell_type2_pre'][2:]    \n",
    "\n",
    "            # for each mtype specify expected ABI gene custers and their expected % (based on Jiang 2015 cre-lines to morhology table)\n",
    "            if syne3['cell_type2_post']=='MC':\n",
    "                if syne3['layer_post'] == 'L2/3':\n",
    "                    syne3['cluster_post']=[['Sst Calb2 Pdlim5',100]]\n",
    "                if syne3['layer_post'] == 'L5':\n",
    "                    syne3['cre_line_post']='Chrna2'  \n",
    "                    syne3['cluster_post']=[['Sst',100]]\n",
    "\n",
    "            if syne3['cell_type2_post']=='BC':  \n",
    "                if syne3['layer_post'] == 'L2/3':\n",
    "                    if do_simplfy_clusters==0:\n",
    "                        syne3['cluster_post']=[['Pvalb, Pvalb Vipr2-',75],[Sst_NonMC,25]] # modify!!!    \n",
    "                    else:\n",
    "                        syne3['cluster_post']=[['Pvalb, Pvalb Vipr2-',100]] # modify!!!  \n",
    "\n",
    "                if syne3['layer_post'] == 'L5':\n",
    "                    if do_simplfy_clusters==0:\n",
    "                        syne3['cluster_post']=[['Pvalb, Pvalb Vipr2-',75],[Sst_NonMC,25]] # modify!!!    \n",
    "                    else:\n",
    "                        syne3['cluster_post']=[['Pvalb, Pvalb Vipr2-',100]] # modify!!!    \n",
    "            if syne3['cell_type2_post']=='BTC':\n",
    "                if do_simplfy_clusters==0:\n",
    "                    syne3['cluster_post']=[['Vip Lmo1 Myl1, Vip Ptprt Pkp2, Vip Rspo4 Rxfp1 Chat',90], # modify!!! \n",
    "                                         [Sst_NonMC,10]]\n",
    "                else:\n",
    "                    syne3['cluster_post']=[['Vip Lmo1 Myl1, Vip Ptprt Pkp2, Vip Rspo4 Rxfp1 Chat',100]]\n",
    "            if syne3['cell_type2_post']=='HEC':  # modify!!! \n",
    "                if do_simplfy_clusters==0:\n",
    "                    syne3['cluster_post']=[['Pvalb, Pvalb Vipr2-,Pvalb Tpbg-,Pvalb Reln Itm2a-,Pvalb Reln Tac1-',75],\n",
    "                                           [Sst_NonMC,25]]\n",
    "                else:\n",
    "                    syne3['cluster_post']=[['Pvalb, Pvalb Vipr2-,Pvalb Tpbg-,Pvalb Reln Itm2a-,Pvalb Reln Tac1-',100]]            \n",
    "            if syne3['cell_type2_post']=='NGC':            \n",
    "                  syne3['cluster_post']=[['Lamp5 Ntn1 Npy2r, Lamp5 Plch2 Dock5',100]] \n",
    "\n",
    "        if  df_stp.loc[i,'name'][0:5]=='Walke':    \n",
    "            if syne3['cell_type2_post']=='MC':\n",
    "                 syne3['cluster_post']=[['Sst Calb2 Pdlim5',100]]            \n",
    "\n",
    "        #if  df_stp.loc[i,'name'][0:5]=='Miao_':     \n",
    "\n",
    "        #if  df_stp.loc[i,'name'][0:5]=='Yuste':     \n",
    "    \n",
    "    # ABA synphys dataset\n",
    "    if i>=l_pre_post1:\n",
    "        #l = str.strip(STP_aba.loc[i1,'area'])\n",
    "\n",
    "            \n",
    "        i1 = STP_aba.index[i - l_pre_post1]\n",
    "          \n",
    "        syne0 = re.split('\\,',STP_aba.loc[i1,'synapse_type_2'])\n",
    "        syne1 = re.split('\\;',syne0[0])\n",
    "        \n",
    "        syne0 = re.split('\\,',STP_aba.loc[i1,'comments'])\n",
    "        syne1 = re.split('\\;',syne0[0])\n",
    "        syne12 = re.split('\\'',syne1[0])[1]\n",
    "        syne13 = re.split('\\'',syne1[1])[0]\n",
    "\n",
    "        syne_pre = re.split('\\s+',str.strip(syne12))\n",
    "        syne_post = re.split('\\s+',str.strip(syne13))\n",
    "        \n",
    "   \n",
    "        l1 = str.strip(syne_pre[1])\n",
    "        l2 = str.strip(syne_post[1])\n",
    "        \n",
    "        syne3 ={'layer_pre':[l1],  'cre_line_pre':'none', 'cell_type2_pre':'PC','cluster_pre':[],\n",
    "                'layer_post':[l2], 'cre_line_post':'none','cell_type2_post':'PC','cluster_post':[]  } \n",
    "\n",
    "        # add cre-lines when available\n",
    "        in_lines = pd.DataFrame(['vip','pvalb','sst'])\n",
    "        if str.strip(syne_pre[0])!='unknown':\n",
    "            syne3['cre_line_pre']   =str.strip(syne_pre[0])\n",
    "            if in_lines.isin([syne3['cre_line_pre']]).any()[0]:\n",
    "                syne3['cell_type2_pre'] = syne3['cre_line_pre'] \n",
    "        else:\n",
    "            syne3['cell_type2_pre']   = 'PC'\n",
    "            \n",
    "        if str.strip(syne_post[0])!='unknown':\n",
    "            syne3['cre_line_post']  =str.strip(syne_post[0])  \n",
    "            if in_lines.isin([syne3['cre_line_post']]).any()[0]:\n",
    "                syne3['cell_type2_post'] = syne3['cre_line_post'] \n",
    "        else:\n",
    "            syne3['cell_type2_post']   = 'PC'    \n",
    "            \n",
    "            \n",
    "        if syne3['cre_line_pre']=='tlx3':\n",
    "            syne3['layer_pre']=['L4-L6']\n",
    "                    \n",
    "        if syne3['cre_line_post']=='tlx3':\n",
    "            syne3['layer_post']=['L4-L6']\n",
    "        \n",
    "                    \n",
    "        if syne3['cre_line_pre']=='sim1':\n",
    "            syne3['layer_pre']=['L4-L5','L4-L6']\n",
    "        \n",
    "                    \n",
    "        if syne3['cre_line_post']=='sim1':\n",
    "            syne3['layer_post']=['L4-L5','L4-L6']\n",
    "            \n",
    "        if syne3['cre_line_post']=='vip':\n",
    "            if syne3['layer_post'][0]=='L2/3':\n",
    "                syne3['layer_post']=['L1-L2/3']\n",
    "        if syne3['cre_line_pre']=='vip':\n",
    "            if syne3['layer_pre'][0]=='L2/3':\n",
    "                syne3['layer_pre']=['L1-L2/3']\n",
    "                \n",
    "        if syne3['cre_line_post']=='sst':\n",
    "            if syne3['layer_post'][0]=='L2/3':\n",
    "                syne3['layer_post']=['L1-L2/3']\n",
    "        if syne3['cre_line_pre']=='sst':\n",
    "            if syne3['layer_pre'][0]=='L2/3':\n",
    "                syne3['layer_pre']=['L1-L2/3']        \n",
    " \n",
    "    \n",
    "    # # find sets of gene clusters corresponding to each pre- and post- neuron types pair from stp-table\n",
    "    # for ii in range(len(dse)): \n",
    "    #    if sum(np.array(dse[ii])==pre_post.loc[i,0])>0:\n",
    "    #        i1=ii\n",
    "    #    if sum(np.array(dse[ii])==pre_post.loc[i,1])>0:\n",
    "    #        j1=ii    \n",
    "    #\n",
    "    ## select names of stp-types\n",
    "    #pree = dse[i1]\n",
    "    #poste = dse[j1]\n",
    "    \n",
    "    ## select names of genes-types (ABA cluster)\n",
    "    #preg = dsg[i1]\n",
    "    #postg = dsg[j1]\n",
    "\n",
    "    ##synij = (pre_post.loc[:,0].isin(pree))&(pre_post.loc[:,1].isin(poste))\n",
    "    \n",
    "\n",
    "    synij[:] = False\n",
    "    \n",
    "    if i<l_pre_post1:\n",
    "        synij[i] = in_stp_data[i]\n",
    "    else:\n",
    "        synij[i] = True\n",
    "    \n",
    "    #print(syne3)\n",
    "    \n",
    "    if sum(synij)>0:\n",
    "        print('\\n\\n')\n",
    "        print('i = ',str(i))\n",
    "        print(syne3)\n",
    "        #if sum(in_stp_data&synij)>0:\n",
    "        #    in_stp_data = in_stp_data&synij\n",
    "\n",
    "        #add ge data\n",
    "        #'layer_pre':l,  'cre_line_pre':'', 'cell_type2_pre':syne2[0].str.strip(),'cluster_pre':[]\n",
    "        \n",
    "        # PC or GABAergic dataset?\n",
    "        batch = [ syne3['cell_type2_pre']=='PC',syne3['cell_type2_post']=='PC']\n",
    "        \n",
    "        # add ge data depending on available stp cell types labels:\n",
    "        #  select presynaptic cells\n",
    "        if batch[0]:\n",
    "\n",
    "            if i>=l_pre_post1:\n",
    "                in_preg  = df_aba_vis_l23glu_c.loc[:,'brain_subregion'].isin(syne3['layer_pre']) #modify!!!\n",
    "                #in_preg  = in_preg&(df_aba_vis_l23glu_c.loc[:,'driver_lines']==syne3['cre_line_pre']) #modify!!!\n",
    "                if syne3['cre_line_pre']!='none':\n",
    "                    in_preg  = in_preg&(df_aba_vis_l23glu_c.loc[:,'driver_lines'].str.lower().str.contains(syne3['cre_line_pre'])) #modify!!!\n",
    "                #set(df_aba_vis_l23glu_c.loc[:,'driver_lines'].str.lower().str.contains('nr5a1'))\n",
    "            else:\n",
    "                #in_preg  = df_aba_vis_l23glu_c.loc[:,'brain_subregion'].isin([syne3['layer_pre']]) #modify!!!\n",
    "                in_preg  = df_aba_vis_l23glu_c.loc[:,'brain_subregion'].isin([syne3['layer_pre']]) #modify!!!\n",
    "        else:\n",
    "            if len(syne3['cluster_pre'])==0:\n",
    "                # if len(syne3['cre_line_pre'])>0:\n",
    "                in_preg = df_aba_vis_in_c.loc[:,'driver_lines'].str.lower().str.contains(str.lower(syne3['cre_line_pre']))\n",
    "                #in_preg = in_preg&df_aba_vis_in_c.loc[:,'brain_subregion'].isin([syne3['layer_pre']])\n",
    "                if i>=l_pre_post1:\n",
    "                    in_preg = in_preg&df_aba_vis_in_c.loc[:,'brain_subregion'].isin(syne3['layer_pre'])\n",
    "                else:\n",
    "                    in_preg = in_preg&(df_aba_vis_in_c.loc[:,'brain_subregion'].str.contains(syne3['layer_pre']))\n",
    "                    \n",
    "                #if i==72:\n",
    "                #    breakpoint()\n",
    "                    \n",
    "                    \n",
    "            else:                  \n",
    "                #in_preg = df_aba_vis_in_c.loc[:,'brain_subregion'].isin([syne3['layer_pre']])\n",
    "                in_preg = df_aba_vis_in_c.loc[:,'brain_subregion'].str.contains(syne3['layer_pre'])\n",
    "                cl = syne3['cluster_pre']\n",
    "                #cols2 = ['name','perc']\n",
    "                #cl2pre =pd.DataFrame([],columns= cols2)\n",
    "                in_preg0 = in_preg.copy()\n",
    "                in_preg[:] = False \n",
    "                for icl in range(len(cl)):\n",
    "                    clp = cl[icl][1]\n",
    "                    clns = re.split(',',cl[icl][0])\n",
    "                    cls3 = []\n",
    "                    cls3m=[]\n",
    "                    for icls in range(len(clns)):\n",
    "                        clns4 = clns[icls]\n",
    "                        clns5 = re.split('\\s+',clns4)\n",
    "                        if len(clns5)<3:\n",
    "                            clns6=clns5[0]\n",
    "                            cls3=cls3 + list(set(df_aba_vis_in_c.loc[df_aba_vis_in_c.loc[:,'subclass'].str.contains(clns6),'cluster']))\n",
    "                        else:\n",
    "                            clns4=str.strip(clns4)\n",
    "                            clns7=re.split('\\-',clns4)\n",
    "                            if len(clns7)>1:\n",
    "                                clns4=clns7[0]\n",
    "                                cls3m=cls3m+[str.strip(clns4)]\n",
    "                            else:\n",
    "                                cls3=cls3+[str.strip(clns4)]\n",
    "                    cls3 = list(set(cls3).difference(set(cls3m)))\n",
    "                    cl22=pd.DataFrame([cls3,[clp]],index=cols2).T\n",
    "                    cl2pre = cl2pre.append(cl22)\n",
    "                    \n",
    "                    in_preg = in_preg|(in_preg0&df_aba_vis_in_c.loc[:,'cluster'].isin(cls3))\n",
    "                    \n",
    "                if len(syne3['cre_line_pre'])>0:\n",
    "                    #in_preg = in_preg&df_aba_vis_in_c.loc[:,'driver_lines'].str.contains(syne3['cre_line_pre'])\n",
    "                    in_line=df_aba_vis_in_c.loc[:,'driver_lines'].str.contains(syne3['cre_line_pre'])\n",
    "                    if in_line.sum()>0:\n",
    "                        in_preg = in_preg&in_line                        \n",
    "        \n",
    "        #  select postsynaptic cells\n",
    "        if batch[1]:\n",
    "\n",
    "            if i>=l_pre_post1:\n",
    "                in_postg  = df_aba_vis_l23glu_c.loc[:,'brain_subregion'].isin(syne3['layer_post']) #modify!!!\n",
    "                if syne3['cre_line_post']!='none':\n",
    "                    in_postg  = in_postg&(df_aba_vis_l23glu_c.loc[:,'driver_lines'].str.lower().str.contains(syne3['cre_line_post'])) #modify!!!\n",
    "            else:\n",
    "                #in_postg = df_aba_vis_l23glu_c.loc[:,'brain_subregion'].isin([syne3['layer_post']]) #modify!!!\n",
    "                in_postg = df_aba_vis_l23glu_c.loc[:,'brain_subregion']==syne3['layer_post'] #modify!!!\n",
    "        else:    \n",
    "            if len(syne3['cluster_post'])==0:\n",
    "                # if len(syne3['cre_line_post'])>0:\n",
    "                in_postg = df_aba_vis_in_c.loc[:,'driver_lines'].str.lower().str.contains(str.lower(syne3['cre_line_post']))\n",
    "                #in_postg = in_postg&df_aba_vis_in_c.loc[:,'brain_subregion'].isin([syne3['layer_post']])\n",
    "                if i>=l_pre_post1:\n",
    "                    in_postg = in_postg&df_aba_vis_in_c.loc[:,'brain_subregion'].isin(syne3['layer_post'])\n",
    "                else:\n",
    "                    in_postg = in_postg&df_aba_vis_in_c.loc[:,'brain_subregion'].str.contains(syne3['layer_post'])\n",
    "            else:    \n",
    "                ## under construction ...\n",
    "                #in_postg = df_aba_vis_in_c.loc[:,'brain_subregion'].isin([syne3['layer_post']])\n",
    "                in_postg = df_aba_vis_in_c.loc[:,'brain_subregion'].str.contains(syne3['layer_post'])\n",
    "                cl = syne3['cluster_post']\n",
    "                #cols2 = ['name','perc']\n",
    "                #cl2post =pd.DataFrame([],columns= cols2)\n",
    "                in_postg0 = in_postg.copy()\n",
    "                in_postg[:] = False                 \n",
    "                for icl in range(len(cl)):\n",
    "                    clp = cl[icl][1]\n",
    "                    clns = re.split(',',cl[icl][0])\n",
    "                    cls3 = []\n",
    "                    cls3m=[]\n",
    "                    for icls in range(len(clns)):\n",
    "                        clns4 = clns[icls]\n",
    "                        clns5 = re.split('\\s+',clns4)\n",
    "                        if len(clns5)<3:\n",
    "                            clns6=clns5[0]\n",
    "                            cls3=cls3 + list(set(df_aba_vis_in_c.loc[df_aba_vis_in_c.loc[:,'subclass'].str.contains(clns6),'cluster']))\n",
    "                        else:\n",
    "                            clns4=str.strip(clns4)\n",
    "                            clns7=re.split('\\-',clns4)\n",
    "                            if len(clns7)>1:\n",
    "                                clns4=clns7[0]\n",
    "                                cls3m=cls3m+[str.strip(clns4)]\n",
    "                            else:\n",
    "                                cls3=cls3+[str.strip(clns4)]\n",
    "                    cls3 = list(set(cls3).difference(set(cls3m)))\n",
    "                    cl22=pd.DataFrame([cls3,[clp]],index=cols2).T\n",
    "                    cl2post = cl2post.append(cl22)\n",
    "                    \n",
    "                    #in_postg = in_postg&df_aba_vis_in_c.loc[:,'cluster'].isin(cls3)\n",
    "                    in_postg = in_postg|(in_postg0&df_aba_vis_in_c.loc[:,'cluster'].isin(cls3))\n",
    "                    \n",
    "                if len(syne3['cre_line_post'])>0:\n",
    "                    in_line=df_aba_vis_in_c.loc[:,'driver_lines'].str.contains(syne3['cre_line_post'])\n",
    "                    if in_line.sum()>0:\n",
    "                        in_postg = in_postg&in_line\n",
    "            \n",
    "            \n",
    "        ## make equal size pre- and post- datasets\n",
    "        #in_preg = in_preg&(in_preg.isna()==False)\n",
    "        #in_postg = in_postg&(in_postg.isna()==False)\n",
    "        #npre = in_preg.sum()\n",
    "        #npost = in_postg.sum()\n",
    "        ##nsample = max([npost, npre] )\n",
    "        #nsample = min([npost, npre] ) # number of cells\n",
    "        #print('npre = '+ str(npre)+' , npost = ',str(npost))\n",
    "        #if (nsample<150)&(npre>0)&(npost>0):\n",
    "        #    nsample=150\n",
    "        \n",
    "        #idx2 =  np.mod(np.arange(nsample),npre)\n",
    "        #in_preg = np.nonzero(in_preg)[0][idx2]\n",
    "        #idx2 =  np.mod(np.arange(nsample),npost)\n",
    "        #in_postg = np.nonzero(in_postg)[0][idx2]\n",
    "        \n",
    "        \n",
    "        # make equal size pre- and post- datasets\n",
    "        in_preg = in_preg&(in_preg.isna()==False)\n",
    "        in_postg = in_postg&(in_postg.isna()==False)\n",
    "        npre = in_preg.sum()\n",
    "        npost = in_postg.sum()\n",
    "        nsample = max([npost, npre] )\n",
    "        #nsample = min([npost, npre] ) # number of cells\n",
    "        \n",
    "\n",
    "        \n",
    "   \n",
    "            \n",
    "        print('npre = '+ str(npre)+' , npost = ' + str(npost) + ' , nsample = '+str(nsample))    \n",
    "        \n",
    "        \n",
    "        # select cell indices for gene expression sampling\n",
    "        #do_bootstrap=1 # defined in hippocampus\n",
    "        if  do_bootstrap==0:\n",
    "            \n",
    "            if (nsample<60)&(npre>0)&(npost>0):\n",
    "                nsample=60\n",
    "            \n",
    "            if (nsample>180)&(npre>0)&(npost>0):  # restrict nsample to avoid overrepresentation?\n",
    "                nsample=180 \n",
    "            \n",
    "            idx2 =   np.mod(np.arange(nsample),npre) \n",
    "            in_preg = np.nonzero(in_preg)[0][idx2]\n",
    "            idx2 =   np.mod(np.arange(nsample),npost) \n",
    "            in_postg = np.nonzero(in_postg)[0][idx2]\n",
    "            \n",
    "            in_preg_ = in_preg\n",
    "            in_postg_ = in_postg\n",
    "        else :\n",
    "            if (npre>0)&(npost>0):\n",
    "                nsample=nsample_hipp #200 #defined in hippocampus\n",
    "            #random.uniform(low=0.0, high=1.0, size=None)\n",
    "            idx2pre =  np.floor( np.random.uniform(0,npre, (nsample,1)) ).astype('int')\n",
    "            in_preg_bs = np.nonzero(in_preg)[0][idx2pre]\n",
    "            idx2post =  np.floor( np.random.uniform(0,npost, (nsample,1)) ).astype('int')\n",
    "            in_postg_bs = np.nonzero(in_postg)[0][idx2post]\n",
    "            in_preg_ = in_preg_bs[:,0]\n",
    "            in_postg_ = in_postg_bs[:,0]\n",
    "        \n",
    "        \n",
    "\n",
    "        if batch[0]:\n",
    "            #in_preg_cl = list(set(df_aba_vis_l23glu_c.iloc[in_preg,:].loc[:,'cluster']))\n",
    "            in_preg_cl = df_aba_vis_l23glu_c.iloc[in_preg_,:].loc[:,'cluster'].value_counts()       \n",
    "            sampl_pre = df_aba_vis_l23glu_c.iloc[in_preg_,:].loc[:,'sample_name']\n",
    "        else:\n",
    "            #in_preg_cl = list(set(df_aba_vis_in_c.iloc[in_preg,:].loc[:,'cluster']))\n",
    "            in_preg_cl = df_aba_vis_in_c.iloc[in_preg_,:].loc[:,'cluster'].value_counts()\n",
    "            sampl_pre =df_aba_vis_in_c.iloc[in_preg_,:].loc[:,'sample_name']\n",
    "        \n",
    "        if batch[1]:\n",
    "            #in_postg_cl = list(set(df_aba_vis_l23glu_c.iloc[in_postg,:].loc[:,'cluster']))\n",
    "            in_postg_cl = df_aba_vis_l23glu_c.iloc[in_postg_,:].loc[:,'cluster'].value_counts()\n",
    "            sampl_post =df_aba_vis_l23glu_c.iloc[in_postg_,:].loc[:,'sample_name']\n",
    "        else:        \n",
    "            #in_postg_cl = list(set(df_aba_vis_in_c.iloc[in_postg,:].loc[:,'cluster']))\n",
    "            in_postg_cl = df_aba_vis_in_c.iloc[in_postg_,:].loc[:,'cluster'].value_counts()\n",
    "            sampl_post =df_aba_vis_in_c.iloc[in_postg_,:].loc[:,'sample_name']\n",
    "        in_preg_cl = in_preg_cl/in_preg_cl.sum()*100    \n",
    "        in_postg_cl = in_postg_cl/in_postg_cl.sum()*100    \n",
    "        print('presynaptic clusters')\n",
    "        print(in_preg_cl)\n",
    "        print('postsynaptic clusters')\n",
    "        print(in_postg_cl)\n",
    "        \n",
    "        \n",
    "\n",
    "        #if npost>npre:\n",
    "        #    idx = np.mod(np.arange(npost),npre)\n",
    "        #    #prege = prege[idx,:]\n",
    "        #    #postge = postge\n",
    "        #    prege = prege\n",
    "        #    postge = postge[0:npre,:]            \n",
    "        #elif npost<npre :\n",
    "        #    idx = np.mod(np.arange(npre),npost)\n",
    "        #    #prege = prege\n",
    "        #    #postge = postge[idx,:]\n",
    "        #    prege = prege[0:npost,:]\n",
    "        #    postge = postge         \n",
    "        \n",
    "        # take ge data  \n",
    "        \n",
    "        prege = vis_dat.X[:, in_pregs]  # presynaptic gene set expression\n",
    "        prege = prege[(vbi==batch[0]), :]\n",
    "        \n",
    "        postge = vis_dat.X[:,  in_postgs] # postsynaptic gene set expression\n",
    "        postge = postge[(vbi==batch[1]),  :]\n",
    "        \n",
    "        # use hippocampal settings\n",
    "        #do_median=1        # median over synapse types\n",
    "        #do_log = 'none'    # logarithmic transform of ge counts\n",
    "        \n",
    "        if  (do_bootstrap==0)|(do_median==0):\n",
    "            prege = prege[in_preg_,:]\n",
    "            postge = postge[in_postg_,:]\n",
    "            # do_samples_columns=1\n",
    "        else: \n",
    "            # do_samples_columns=0\n",
    "            #nsample = 100 # defined in hippocampus\n",
    "            Nbs = nsample_hipp\n",
    "            ##random.uniform(low=0.0, high=1.0, size=None)\n",
    "            #nmax=5\n",
    "            nmax = npre #np.min([1000,npre])\n",
    "            idx2pre =  np.floor( np.random.uniform(0,npre, (nmax,Nbs)) ).astype('int')\n",
    "            in_preg_ = np.nonzero(in_preg)[0][np.arange(npre)]\n",
    "            nmax = npost #np.min([1000,npost])\n",
    "            idx2post =  np.floor( np.random.uniform(0,npost, (nmax,Nbs)) ).astype('int')\n",
    "            in_postg_ = np.nonzero(in_postg)[0][np.arange(npost)]\n",
    "\n",
    "            # \n",
    "            # ??? transform ge to median bootstraped log ???\n",
    "            prege_ = prege[in_preg_,:]\n",
    "            postge_ = postge[in_postg_,:]\n",
    "            \n",
    "            if do_log=='none':\n",
    "                 prege_ =  prege_\n",
    "            elif do_log=='log2':       \n",
    "                y = np.ma.masked_where(prege_ == 0, prege_)\n",
    "                prege_=np.ma.log2(y).filled(0)\n",
    "                y = np.ma.masked_where(postge_ == 0, postge_)\n",
    "                postge_=np.ma.log2(y).filled(0)\n",
    "            elif do_log=='log1p': \n",
    "                prege_ = np.log1p(prege_)\n",
    "                postge_ =  np.log1p(postge_)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "            npregs=prege_.shape[1]\n",
    "            ge_bs = np.zeros((Nbs,npregs+postge_.shape[1]))\n",
    "            for ibs in range(Nbs):\n",
    "                #ans=np.apply_along_axis(lambda v: np.median(v[v!=0]), 0, x)\n",
    "                #ans[np.isnan(ans)]=0.\n",
    "                xbs = prege_[idx2pre[:,ibs],:]\n",
    "                pregbs  = np.apply_along_axis(lambda v: np.median(v[v!=0]), 0, xbs)\n",
    "                pregbs[np.isnan(pregbs)]=0.\n",
    "                xbs = postge_[idx2post[:,ibs],:]\n",
    "                postgbs = np.apply_along_axis(lambda v: np.median(v[v!=0]), 0, xbs)\n",
    "                postgbs[np.isnan(postgbs)]=0.\n",
    "\n",
    "                ge_bs[ibs,0:npregs] =pregbs\n",
    "                ge_bs[ibs,npregs:]  =postgbs\n",
    "\n",
    "\n",
    "            #df_hr2 = np.log2(df_hr2d)\n",
    "            #df_hr2 = pd.concat([df_hr2,df_hr2s],axis=1) # log2 normalization\n",
    "            #df_hr2 = df_hr2.replace({0:np.nan})\n",
    "            #df_hr2 = np.log2(df_hr2)\n",
    "\n",
    "            #df_hr2 = pd.concat([df_hr.iloc[1:,0],df_hr2],axis=1)\n",
    "\n",
    "            #df_hr2.dtypes\n",
    "            #df_hr2.head()\n",
    "\n",
    "            #df_hr2 = df_hr2.replace({0:np.nan})\n",
    "            #df_hrmg=df_hr2.groupby(['Class']) # logarithm_2 of UMI counts\n",
    "            #df_hrmg=df_hrmg.median(numeric_only=True) # median of log2(UMI)\n",
    "\n",
    "            #df_hrmg0 = df_hrmg.iloc[:,0:-1]\n",
    "\n",
    "            #df_hrmsg = pd.concat([df_hrmg.iloc[:,-1]]*df_hrmg0.shape[1],axis=1)\n",
    "            #df_hrmsg.columns = df_hrmg.iloc[:,0:-1].columns\n",
    "\n",
    "            #df_hrmg = df_hrmg.iloc[:,0:-1] - df_hrmsg +6*np.log2(10)  # transform from log2( UMI counts) to log2(CPM)\n",
    "\n",
    "            # ??? transform ge to median bootstraped log ???\n",
    "\n",
    "\n",
    "              \n",
    "\n",
    "        # add data to the final dataset    \n",
    "        gnij = np.repeat(np.array([syne3['cell_type2_pre'], syne3['cell_type2_post'],\n",
    "                                   syne3['layer_pre'], syne3['layer_post'],\n",
    "                                   syne3['cre_line_pre'], syne3['cre_line_post']]).reshape((1,6)),nsample,axis=0 )\n",
    "        gnij = pd.DataFrame(gnij, columns = ['cell_type2_pre', 'cell_type2_post',\n",
    "                                             'layer_pre','layer_post',\n",
    "                                             'cre_line_pre', 'cre_line_post'] )\n",
    "        \n",
    "        if  (do_bootstrap==0)|(do_median==0):\n",
    "            prege  = pd.DataFrame(prege,columns  = pregs3 )\n",
    "            postge = pd.DataFrame(postge,columns = postgs3 )\n",
    "\n",
    "\n",
    "            #ge_dataij = pd.concat([gnij, prege, postge ], axis = 1)\n",
    "            samples_pre = pd.DataFrame(sampl_pre.values[:gnij.shape[0]],columns = ['samples_pre'] )\n",
    "            samples_post = pd.DataFrame(sampl_post.values[:gnij.shape[0]],columns = ['samples_post'] )\n",
    "            \n",
    "            \n",
    "            ge_dataij = pd.concat([gnij, prege, postge, samples_pre, samples_post ], axis = 1)\n",
    "            \n",
    "            \n",
    "            \n",
    "        else:\n",
    "            ge_bs = pd.DataFrame(ge_bs,columns = pregs3 + postgs3 )\n",
    "\n",
    "\n",
    "            ge_dataij = pd.concat([gnij, ge_bs ], axis = 1)\n",
    "       \n",
    "\n",
    "        \n",
    "        \n",
    "        \n",
    "        # add stp data\n",
    "        if i<l_pre_post1:\n",
    "            stp_dataij =  df_stp.loc[[0,i], stp_type['stp_data_type']].T\n",
    "            stp_columns = stp_TM_names\n",
    "        else:\n",
    "            #stp_dataij  = STP_aba.loc[i1, STP_aba_type['stp_data_type']]\n",
    "            #stp_columns = np.array(stp_aba_names)\n",
    "            #nbs = int(len(stp_dataij)/len(stp_columns))\n",
    "            #stp_dataij  = pd.DataFrame(stp_dataij, index = np.repeat(stp_columns,nbs)).T\n",
    "            ##i2 = np.trunc(np.arange(len(stp_dataij.columns))/len(stp_columns)).astype(int)\n",
    "            #stp_dataij = stp_dataij.T.reset_index(drop=False)\n",
    "            #stp_dataij.columns =[0,i]\n",
    "            \n",
    "            stp_dataij  = STP_aba.loc[i1, STP_aba_type['stp_data_type']]\n",
    "            stp_columns = np.array(stp_aba_names)\n",
    "            #nbs = int(len(stp_dataij)/len(stp_columns))\n",
    "            nbs = int(stp_dataij.shape[0]) #/len(stp_columns))\n",
    "            stp_dataij  = pd.DataFrame(stp_dataij.reshape((-1,)),\n",
    "                                       index = np.tile(np.array(stp_aba_names),(nbs,))).T\n",
    "            \n",
    "            \n",
    "            if do_shrink_stp_distribution==True: # shrink bootstrap distribution of stp parameters towards mean values - to test stp variability reduction\n",
    "                print('shrink bootstrap distribution of stp parameters : ',shrink_stp_distribution_coef)\n",
    "                stp_dataij0 = stp_dataij.iloc[:,0:len(stp_aba_names)].values\n",
    "                #nbs00 = int(stp_dataij.shape[1]/len(stp_aba_names))\n",
    "                stp_dataij0  = pd.DataFrame(np.tile(stp_dataij0,[1,nbs]), columns = stp_dataij.columns)\n",
    "                stp_dataij = stp_dataij0 + (stp_dataij - stp_dataij0)*shrink_stp_distribution_coef\n",
    "            \n",
    "            \n",
    "            \n",
    "            #i2 = np.trunc(np.arange(len(stp_dataij.columns))/len(stp_columns)).astype(int)\n",
    "            stp_dataij = stp_dataij.T.reset_index(drop=False)\n",
    "            stp_dataij.columns =[0,i]\n",
    "            \n",
    "            \n",
    "        \n",
    "\n",
    "        i2 = np.trunc(np.arange(len(stp_dataij.index))/len(stp_columns)).astype(int)\n",
    "        stp_dataij  = stp_dataij.set_index(i2).pivot(columns=0,values=i).loc[:,stp_columns] #ss.set_index('par',append=True).unstack('par')\n",
    "        #stp_dataij =  stp_dataij.T.pivot(index=0, values=synij).T\n",
    "        \n",
    "        idx = np.mod(np.arange(nsample),len(stp_dataij.index))\n",
    "        stp_dataij=stp_dataij.iloc[idx,:].reset_index(drop=True)  \n",
    "        \n",
    "        i_data_set = pd.DataFrame(i+np.zeros(len(stp_dataij.index)), index = stp_dataij.index, columns = ['index_ds'])\n",
    "        stp_dataij=pd.concat([stp_dataij, i_data_set],axis=1)\n",
    "        \n",
    "        #if do_averagig_stp:\n",
    "        #    #stp_dataij.loc[:,:] = stp_dataij.mean(axis=0).values\n",
    "        #    stp_dataij.loc[:,:] = stp_dataij.median(axis=0).values\n",
    "\n",
    "        \n",
    "        #stp_dataij = pd.DataFrame({'presynaptic':pree[0], 'postsynaptic': poste[0], \n",
    "        #                           'stp_data':stp_dataij.iloc[0,:].values})\n",
    "        #stp_data  = stp_data.append(stp_dataij) \n",
    "        \n",
    "        # names of ge and stp columns\n",
    "\n",
    "        \n",
    "        ge_dataij  = pd.concat([ge_dataij,stp_dataij] , axis=1)\n",
    "        \n",
    "        # add to all ge and stp dataframe\n",
    "        ge_data  = pd.concat([ge_data, ge_dataij] , axis=0)\n",
    "        \n",
    "\n",
    "        \n",
    "#stp_data =  stp_data.reset_index(drop=True)  \n",
    "\n",
    "\n",
    "\n",
    "#stp_data.loc[stp_data.loc[:,'postsynaptic']=='NGC','postsynaptic']='Lamp5'\n",
    "#stp_data2 = ge_data.iloc[:,0:3].copy(deep=True)                                 # modify!\n",
    "#stp_data2.columns=['presynaptic', 'postsynaptic', 'ppr']\n",
    "#print(stp_data2.head())\n",
    "#for i in stp_data.index: #stp_data.index:\n",
    "#    pre=stp_data.loc[i,'presynaptic']\n",
    "#    post=stp_data.loc[i,'postsynaptic']\n",
    "#    stp_data2.loc[stp_data2.loc[:,'presynaptic'].isin([pre])&stp_data2.loc[:,'postsynaptic'].isin([post]),'ppr']=stp_data.loc[i,'stp_data'] \n",
    "    \n",
    "ge_data = ge_data.reset_index()    \n",
    "\n",
    " \n",
    "## transform labels from TM to An:A1\n",
    "As = STP_sim_complex(ge_data,l_pre_post2,stp_aba_names)\n",
    "\n",
    "N = As.shape[1]\n",
    "\n",
    "As =  As[:,1:]/As[:,0].reshape((-1,1))\n",
    "\n",
    "#As =  As[:,1:]/np.tile(As[:,0],(N-1,1)).T\n",
    "\n",
    "## add new stp parameters to ge_data dataset\n",
    "#stp_columns2 = ['A'+str(s)+'/A1' for s in range(1,N+1)]\n",
    "stp_columns2 = ['A'+str(s)+'/A1' for s in range(2,N+1)]\n",
    "stp_data2 = pd.DataFrame(As,columns = stp_columns2, index = ge_data.index)\n",
    "\n",
    "ge_data = pd.concat([ge_data, stp_data2 ],axis=1)\n",
    "\n",
    "if do_averagig_stp:\n",
    "     #stp_dataij.loc[:,:] = stp_dataij.mean(axis=0).values\n",
    "     #stp_dataij.loc[:,:] = stp_dataij.median(axis=0).values\n",
    "    all_i_ds =  ge_data.loc[:,'index_ds']                             \n",
    "    for i_ds in list(set(all_i_ds)):\n",
    "              ids2 =  ge_data.loc[:,'index_ds'].isin([i_ds])  \n",
    "              ge_data.loc[ids2,stp_columns2] = ge_data.loc[ids2,stp_columns2].median(axis=0).values                   \n",
    "                                   \n",
    "                                   \n",
    "## filter TM parameters\n",
    "do_filter_pars = 0\n",
    "\n",
    "if do_filter_pars!=0:\n",
    "    print((ge_data.loc[:,stp_columns ]<0).sum().sum())\n",
    "    print(len(ge_data.index))\n",
    "\n",
    "    nannot=6\n",
    "    annot_columns = ge_columns[0:nannot]\n",
    "    ge_data2 = ge_data.loc[:,annot_columns + stp_columns2].groupby(annot_columns).mean()\n",
    "    f,ax = plt.subplots(figsize=(5, 5))\n",
    "    ge_data2.T.plot(ax=ax)\n",
    "    plt.gca().legend('')\n",
    "    #plt.gca().title('An:A1 before filtering data')\n",
    "    ax.title.set_text('An:A1 before filtering data')\n",
    "    #ge_data2 = ge_data.loc[:,annot_columns + stp_columns2].groupby(annot_columns).mean()\n",
    "    print('mean of stp before filtering:')\n",
    "    print(ge_data2)\n",
    "\n",
    "    #i_filter_out = (ge_data.loc[:,'dp/p0']*ge_data.loc[:,'p0']+)|(ge_data.loc[:,'tD'])\n",
    "    #i_filter_out = i_filter_out&((ge_data.loc[:,'tF'])|(ge_data.loc[:,'tF']))\n",
    "    i_filter_out = np.sum(ps2>1,axis=1)|np.sum(ps2<0,axis=1)\n",
    "    i_filter_out = i_filter_out|(np.sum(ns2>1,axis=1)|np.sum(ns2<0,axis=1))\n",
    "    ge_data = ge_data.loc[i_filter_out==False ,:]\n",
    "\n",
    "    ge_data2 = ge_data.loc[:,annot_columns + stp_columns2].groupby(annot_columns).mean()\n",
    "    f1,ax1 = plt.subplots(figsize=(5, 5))\n",
    "    ge_data2.T.plot(ax=ax1)\n",
    "    plt.gca().legend('')\n",
    "    #plt.gca().title('An:A1 after filtering data')\n",
    "    ax1.title.set_text('An:A1 after filtering data')\n",
    "    #ge_data2 = ge_data.loc[:,annot_columns + stp_columns2].groupby(annot_columns).mean()\n",
    "\n",
    "    print('mean of stp after filtering:')\n",
    "    print(ge_data2)\n",
    "    \n",
    "#print('mean of stp after filtering:')\n",
    "#print(ge_data2)\n",
    "\n",
    "\n",
    "print(len(ge_data.index))\n",
    "#ge_data.loc[:,stp_columns ] =ge_data.loc[:,stp_columns ].abs()  \n",
    "\n",
    "\n",
    "#\n",
    "## normalize stp parameters - for TM parameters\n",
    "do_normalize_stp_pars = 0;\n",
    "if do_normalize_stp_pars!=0:\n",
    "    ge_data.loc[:,stp_columns  ] =np.log(ge_data.loc[:,stp_columns  ].values.astype(float))\n",
    "\n",
    "#  assign training and labels \n",
    "stp_columns1 = stp_columns \n",
    "stp_columns = stp_columns2 # select training labels names\n",
    "\n",
    "nannot = 7\n",
    "nstp =len(stp_columns)\n",
    "\n",
    "annot_columns = ge_columns[0:6]+['index_ds']\n",
    "if do_samples_columns==1:\n",
    "    nannot = 9\n",
    "    annot_columns = annot_columns + ['samples_pre', 'samples_post']    \n",
    "\n",
    "## averaging of stp for each synapse type \n",
    "\n",
    "\n",
    "#X = ge_data.loc[:,ge_data.columns[0:]].values\n",
    "X_c = ge_data.loc[:,annot_columns + ge_columns[6:]].values\n",
    "y_c = ge_data.loc[:,annot_columns + stp_columns].values \n",
    "#y = stp_data2.loc[:,:].values\n",
    "\n",
    "## normalize stp parameters - for TM parameters\n",
    "#y[:,nannot: ] =np.log(y[:,nannot: ].astype(float))\n",
    "\n",
    "\n",
    "    \n",
    "\n",
    "# print elapsed time\n",
    "t2 = time.time()\n",
    "#print(regr.feature_importances_)\n",
    "print('Elapsed time: '+str(t2-t1))\n",
    "\n",
    "print(y_c.shape)\n",
    "X_c.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(ge_data.columns[:]=='samples_pre').sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CORRECTED : \n",
    "there were errors in STP_cortex_interneurons_Cre_labeled_UDF5_4pars.xlsx table cre lines for Sst neurons assignment\n",
    "\n",
    "\n",
    "\n",
    "# ACHTUNG!!! - ERROR IN CLUSTER ASSIGNMENT!!!\n",
    "i =  22\n",
    "{'layer_pre': 'L4', 'cre_line_pre': 'PC', 'cell_type2_pre': 'PC', 'cluster_pre': [], 'layer_post': 'L4', 'cre_line_post': 'Pvalb-IRES-Cre', 'cell_type2_post': 'Sst', 'cluster_post': []}\n",
    "npre = 1034 , npost = 53 , nsample = 1034\n",
    "presynaptic clusters\n",
    "L4 IT VISp Rspo1            89.5\n",
    "L5 IT VISp Hsd11b1 Endou     8.5\n",
    "L5 IT VISp Whrn Tox2         0.5\n",
    "L5 NP VISp Trhr Met          0.5\n",
    "L2/3 IT VISp Agmat           0.5\n",
    "L5 IT VISp Batf3             0.5\n",
    "Name: cluster, dtype: float64\n",
    "postsynaptic clusters\n",
    "Pvalb Reln Itm2a     36.5\n",
    "Pvalb Reln Tac1      34.5\n",
    "Pvalb Tpbg           25.0\n",
    "Pvalb Calb1 Sst       3.0\n",
    "Pvalb Gpr149 Islr     1.0\n",
    "Name: cluster, dtype: float64\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i =  30\n",
    "{'layer_pre': 'L4', 'cre_line_pre': 'Pvalb-IRES-Cre', 'cell_type2_pre': 'PV', 'cluster_pre': [], 'layer_post': 'L4', 'cre_line_post': 'Sst-IRES-Cre', 'cell_type2_post': 'Sst', 'cluster_post': []}\n",
    "npre = 53 , npost = 16 , nsample = 53\n",
    "presynaptic clusters\n",
    "Pvalb Reln Itm2a     42.5\n",
    "Pvalb Reln Tac1      33.0\n",
    "Pvalb Tpbg           19.5\n",
    "Pvalb Gpr149 Islr     3.0\n",
    "Pvalb Calb1 Sst       2.0\n",
    "Name: cluster, dtype: float64\n",
    "postsynaptic clusters\n",
    "Sst Hpse Sema3c      26.5\n",
    "Sst Tac1 Htr1d       21.5\n",
    "Sst Calb2 Pdlim5     20.5\n",
    "Sst Chrna2 Ptgdr     18.0\n",
    "Sst Hpse Cbln4        7.5\n",
    "Pvalb Gpr149 Islr     6.0\n",
    "Name: cluster, dtype: float64\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ge_data_h.shape)\n",
    "print(X_h.shape)\n",
    "print(y_h.shape)\n",
    "print(len(annot_columns_h))\n",
    "print(len(ge_columns_h))\n",
    "print(len(stp_columns_h))\n",
    "print(ge_columns_h[300:])\n",
    "print(ge_columns_h[0:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ge_data.shape)\n",
    "print(X_c.shape)\n",
    "print(y_c.shape)\n",
    "print(len(annot_columns))\n",
    "print(len(ge_columns))\n",
    "print(len(stp_columns))\n",
    "print(ge_columns[300:])\n",
    "print(ge_columns[0:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ge_data.loc[11*200:11*200+5,['tF', 'p0','tD','dp']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "do_save_temporary=1\n",
    "if do_save_temporary==1:\n",
    "    t1=time.time()\n",
    "    print('save ge_data_h...')\n",
    "    #ge_data_h.to_excel('ge_data_h'+gs_name+'.xlsx')\n",
    "    ge_data_h.to_hdf('ge_data_h'+gs_name+'.hdf',key='data')\n",
    "    t2=time.time()\n",
    "    print('Elapsed time ',t2-t1, ' s')\n",
    "    print('save X_h...')\n",
    "    t1=time.time()\n",
    "    (pd.DataFrame(X_h)).to_hdf('X_h'+gs_name+'.hdf',key='data')\n",
    "    t2=time.time()\n",
    "    print('Elapsed time ',t2-t1, ' s')\n",
    "    print('save y_h...')\n",
    "    (pd.DataFrame(y_h)).to_hdf('y_h'+gs_name+'.hdf',key='data')\n",
    "    print('save annot_columns_h...')\n",
    "    (pd.DataFrame(annot_columns_h)).to_excel('annot_columns_h'+gs_name+'.xlsx')\n",
    "    print('save ge_columns_h...')\n",
    "    (pd.DataFrame(ge_columns_h)).to_excel('ge_columns_h'+gs_name+'.xlsx')\n",
    "    print('save stp_columns_h...')\n",
    "    (pd.DataFrame(stp_columns_h)).to_excel('stp_columns_h'+gs_name+'.xlsx')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "do_save_temporary=1\n",
    "if do_save_temporary==1:\n",
    "    print('save ge_data_c...')\n",
    "    t1=time.time()\n",
    "    #ge_data_h.to_excel('ge_data_h'+gs_name+'.xlsx')\n",
    "    ge_data.to_hdf('ge_data_c'+gs_name+'.hdf',key='data')\n",
    "    t2=time.time()\n",
    "    print('Elapsed time ',t2-t1, ' s')\n",
    "    print('save X_c...')\n",
    "    t1=time.time()\n",
    "    (pd.DataFrame(X_c)).to_hdf('X_c'+gs_name+'.hdf',key='data')\n",
    "    t2=time.time()\n",
    "    print('Elapsed time ',t2-t1, ' s')\n",
    "    print('save y_c...')\n",
    "    (pd.DataFrame(y_c)).to_hdf('y_c'+gs_name+'.hdf',key='data')\n",
    "    print('save annot_columns_c...')\n",
    "    (pd.DataFrame(annot_columns)).to_excel('annot_columns_c'+gs_name+'.xlsx')\n",
    "    print('save ge_columns_c...')\n",
    "    (pd.DataFrame(ge_columns)).to_excel('ge_columns_c'+gs_name+'.xlsx')\n",
    "    print('save stp_columns_c...')\n",
    "    (pd.DataFrame(stp_columns)).to_excel('stp_columns_c'+gs_name+'.xlsx')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "STP_aba_names = pd.DataFrame(STP_aba.index)\n",
    "STP_aba_names.to_excel('stp_aba_names.xlsx')\n",
    "ge_data.loc[np.arange(0,ge_data.index[-1],200),:].to_excel('ge_data_c_names'+gs_name+'.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SKIP NEXT SECTION : GRAPHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.cm as cm\n",
    "\n",
    "# plot gene expression\n",
    "print(ge_data.loc[[0,200,400,600,800,1000,1200],annot_columns])\n",
    "print(postgs)\n",
    "print(pregs)\n",
    "\n",
    "apostgs = ['post__'+p for p in postgs]\n",
    "apostgs = np.array(apostgs)[0:30]\n",
    "apostgs2 = np.array(postgs)[0:30]\n",
    "imda = ge_data.loc[1200,apostgs].map(float).values.reshape([1,-1])\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 10))\n",
    "\n",
    "c_orange = np.array([241,132,62])/255\n",
    "#c_blue = np.array([52,122,175])/255\n",
    "c_blue = np.array([65,43,93])/255\n",
    "#cols = np.array( [[1, 1, 1], [0, 0, 1], [1, 1, 0], [0, 1, 0]])\n",
    "cols = np.array( [[65,43,93], [241,132,62], [0, 255, 0]])/255\n",
    "\n",
    "\n",
    "im = ax.imshow(imda,cmap=cm.RdYlGn)\n",
    "Dn22=1\n",
    "# We want to show all ticks...\n",
    "xD = np.arange(0,apostgs.shape[0])\n",
    "yD = np.arange(0)\n",
    "ax.set_xticks(xD)\n",
    "ax.set_yticks(yD)\n",
    "# ... and label them with the respective list entries\n",
    "#ax.set_yticklabels(nms4)\n",
    "ax.set_xticklabels(apostgs2)\n",
    "#ax.set_yticklabels(vse_pre2[yD])\n",
    "#ax.tick_params(axis='x', colors='none')\n",
    "\n",
    "# Rotate the tick labels and set their alignment.\n",
    "plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "         rotation_mode=\"anchor\", fontsize=16)\n",
    "\n",
    "#axes.set_ylim(-1, nmm+1)\n",
    "\n",
    "ax.grid( b=False)\n",
    "\n",
    "apostgs = ['pre__'+p for p in pregs]\n",
    "apostgs = np.array(apostgs)[np.arange(0,140,5)]\n",
    "apostgs2 = np.array(pregs)[np.arange(0,140,5)]\n",
    "imda = ge_data.loc[1200,apostgs].map(float).values.reshape([1,-1])\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 10))\n",
    "\n",
    "c_orange = np.array([241,132,62])/255\n",
    "#c_blue = np.array([52,122,175])/255\n",
    "c_blue = np.array([65,43,93])/255\n",
    "#cols = np.array( [[1, 1, 1], [0, 0, 1], [1, 1, 0], [0, 1, 0]])\n",
    "cols = np.array( [[65,43,93], [241,132,62], [0, 255, 0]])/255\n",
    "\n",
    "\n",
    "im = ax.imshow(imda,cmap=cm.RdYlGn)\n",
    "Dn22=1\n",
    "# We want to show all ticks...\n",
    "xD = np.arange(0,apostgs.shape[0])\n",
    "yD = np.arange(0)\n",
    "ax.set_xticks(xD)\n",
    "ax.set_yticks(yD)\n",
    "# ... and label them with the respective list entries\n",
    "#ax.set_yticklabels(nms4)\n",
    "ax.set_xticklabels(apostgs2)\n",
    "#ax.set_yticklabels(vse_pre2[yD])\n",
    "#ax.tick_params(axis='x', colors='none')\n",
    "\n",
    "# Rotate the tick labels and set their alignment.\n",
    "plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "         rotation_mode=\"anchor\", fontsize=16)\n",
    "\n",
    "#axes.set_ylim(-1, nmm+1)\n",
    "\n",
    "ax.grid( b=False)\n",
    "\n"
   ]
  },
  {
   "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": [
    "import scipy as sci\n",
    "def epsp_sim(A, T_A,dt,tplus,tminus,Tmax):\n",
    "    \n",
    "    l, v = sci.linalg.eig(np.array([[-1/tplus,0],[1/tplus,-1/tminus]])) \n",
    "    l = np.real(l).reshape([-1,1])\n",
    "    vm = sci.linalg.inv(v)\n",
    "    #x = v*np.exp(l*t)*np.array([1,0]).reshape([-1,1])                  \n",
    "    X = np.array([0,0]).reshape([-1,1])\n",
    "\n",
    "    t=np.array([0])\n",
    "    for i in np.arange(1,T_A.size):\n",
    "        if T_A[i]-T_A[i-1]>Tmax+10*dt:\n",
    "            t2  =  np.arange(int(Tmax/dt))[np.newaxis]*dt  \n",
    "            t3  =  - 10*dt + np.arange(10)[np.newaxis]*dt   \n",
    "            \n",
    "            t5 = np.append(t2,t3+Tmax+10*dt ,axis=1)\n",
    "            t2 = np.append(t2,t3+T_A[i] - T_A[i-1],axis=1)\n",
    "\n",
    "        else: \n",
    "            t2  =  np.arange(int((T_A[i]-T_A[i-1])/dt))[np.newaxis]*dt \n",
    "            t5 = t2\n",
    "\n",
    "        #print(i)    \n",
    "        x0 = X[:,[-1]]   \n",
    "        x0[0,:] = x0[0,:]+A[i-1] \n",
    "        vvmx=(v@np.diag((vm@x0).ravel()))\n",
    "        x = vvmx@np.exp(l@t2)\n",
    "        X = np.append(X,x,axis=1)\n",
    "        #t = np.append(t,T_A[i-1] +t5)   \n",
    "        t = np.append(t,t[-1] +t5)   \n",
    "    \n",
    "    return X, t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "dt=0.1\n",
    "tplus=0.3\n",
    "tminus=10\n",
    "Tmax = 100\n",
    "A = np.array([0,2,3,4,5,0])\n",
    "T_A = np.array([0,200,250,300,550,650])\n",
    "\n",
    "X, t = epsp_sim(A, T_A,dt,tplus,tminus,Tmax)\n",
    "    \n",
    "f, ax =plt.subplots(figsize=(sx, sy))\n",
    "plt.plot(t,X.T)\n",
    "\n",
    "f, ax =plt.subplots(figsize=(sx, sy))\n",
    "plt.plot(X.T)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "l_pre_post23 = 2 #200 #int(200/10)\n",
    "ge_data_temp = ge_data.iloc[::l_pre_post23,:] #.loc[:,stp_aba_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ge_data_temp.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#annot_columns2 =np.array(annot_columns)\n",
    "annot_columns\n",
    "ge_data_temp.loc[:,annot_columns]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "dt=0.1\n",
    "import matplotlib.pyplot as plt \n",
    "%matplotlib inline\n",
    "# transform labels from TM to An:A1\n",
    "fs = [20 , 10, 40] # Hz\n",
    "N = 8\n",
    "\n",
    "Trec = [250, 1000]\n",
    "Trec=[]\n",
    "DT0 = 25000\n",
    "xs  =ge_data_temp.iloc[:,:].loc[:,stp_aba_names].values\n",
    "xs  =np.delete(xs, [5,6],axis=1)\n",
    "#As=np.zeros((l_pre_post2,0))\n",
    "As=np.zeros((xs.shape[0],0))\n",
    "As2=np.zeros((xs.shape[0],0))\n",
    "Tall_A = np.arange(1)\n",
    "if do_add_aba_synphys:\n",
    "    \n",
    "    for f in fs:\n",
    "        T = np.arange(N)*1000/f\n",
    "        Tall_A = np.append(Tall_A, T+DT0+Tall_A[-1])\n",
    "        \n",
    "        Asf, ns, ps, dpp0, p0, tF, tD, A = STP_sim(ge_data_temp.iloc[0:,:],T)\n",
    "        As = np.concatenate([As,Asf],axis=1)\n",
    "        \n",
    "        for ri in range(len(Trec)):\n",
    "            Asr, nsr, psr, dpp0, p0, tF, tD, A = STP_sim(ge_data_temp.iloc[0:,:],[Trec[ri]],init_state=[ns[:,-1],ps[:,-1]])\n",
    "            As = np.concatenate([As,Asr],axis=1)\n",
    "        Tall_A = np.append(Tall_A, np.array(Trec)+Tall_A[-1])\n",
    "        \n",
    "\n",
    "        As2f = np.zeros((xs.shape[0],N+len(Trec)))\n",
    "        for i2 in range(xs.shape[0]):\n",
    "            #as2, sts2 = STP_sim2(xs[i2,:],np.arange(8)*50,model_type = 'tm5') # preconditioning series???\n",
    "            #as2, sts2 = STP_sim2(xs[i2,:], T+DT0, init_state=sts2[-1], model_type='tm5' ) \n",
    "            \n",
    "            as2, sts2 = STP_sim2(xs[i2,:],T,model_type = 'tm5_smr')\n",
    "            \n",
    "            As2f[i2,0:N] = as2\n",
    "            for ri in range(len(Trec)):\n",
    "                as2r, sts2r = STP_sim2(xs[i2,:], [Trec[ri]], init_state=sts2[-1], model_type='tm5_smr' ) \n",
    "                As2f[i2,N+ri] = as2r\n",
    "                \n",
    "        As2 = np.concatenate([As2,As2f],axis=1)\n",
    "            \n",
    "    \n",
    "    As2 = As2/As2[:,0].reshape((-1,1))\n",
    "    As = As/As[:,0].reshape((-1,1))\n",
    "    Na1 = As.shape[1]\n",
    "    As = np.concatenate([As,As2],axis=0)\n",
    "else:\n",
    "    As, ns, ps, dpp0, p0, tF, tD, A = STP_sim(ge_data,T)\n",
    "    \n",
    "\n",
    "sx=10\n",
    "sy=5\n",
    "do_this=0\n",
    "if do_this==1:\n",
    "    for ii in range(As2.shape[0]):\n",
    "        f, ax =plt.subplots(figsize=(sx, sy))\n",
    "        plt.plot(As2[ii,:],'o-',linewidth=4, markersize=12)\n",
    "        #plt.ylim([0,1.2])\n",
    "        plt.xticks(fontsize=18)\n",
    "        plt.xticks(ticks=[0,5,10,15,20],fontsize=18)\n",
    "        plt.yticks(ticks=[0.5,1.0,1.5],fontsize=18)\n",
    "        plt.ylabel('An:A1',fontsize=20)\n",
    "        plt.xlabel('Stimulus number',fontsize=20)\n",
    "        # cell_type2_pre\tcell_type2_post\n",
    "        str_syn_type = ge_data_temp.iloc[ii,:].loc['cell_type2_pre']+'->'+ge_data_temp.iloc[ii,:].loc['cell_type2_post']\n",
    "        plt.title(str(ii)+' '+str_syn_type)\n",
    "        #plt.legend(fontsize=20)\n",
    "        #print(set(n_stp_cells.index).difference(set(i2)))\n",
    "\n",
    "do_this=0\n",
    "if do_this==1:\n",
    "    iis = [19,23,11,28,25,29,44,46,10]    \n",
    "    for ii in range(len(iis)): #(As.shape[0]): #(len(iis)):    \n",
    "        if ii<4:\n",
    "            cc=0  # sst\n",
    "            col = np.array([255, 0, 0])/255\n",
    "        if (ii>=4) and (ii<7):\n",
    "            cc=1  # pvalb\n",
    "            col = np.array([127, 96, 0])/255\n",
    "        if ii>=7:\n",
    "            cc=2  # vip \n",
    "            col = np.array([112, 48, 160])/255\n",
    "\n",
    "        f, ax =plt.subplots(figsize=(sx, sy))\n",
    "        dt=0.1\n",
    "        tplus=1\n",
    "        tminus=20\n",
    "        Tmax = 150\n",
    "        ##A = As[0,:Na1] #np.mean(As[:200,:Na1],axis=0)\n",
    "        #A = As[iis[ii],Na1-9:Na1]\n",
    "        #A = np.append([0],A)\n",
    "        #T_A = Tall_A[-9:-1]\n",
    "        #T_A = np.append(T_A,[T_A[-1]+2*Tmax])\n",
    "\n",
    "        A = As[iis[ii],0:Na1]\n",
    "        A = np.append([0],A)\n",
    "        T_A = Tall_A[0:-1]\n",
    "        T_A = np.append(T_A,[T_A[-1]+2*Tmax])\n",
    "\n",
    "        epsp, t = epsp_sim(A, T_A,dt,tplus,tminus,Tmax)\n",
    "        plt.plot(t[14800:], epsp[1,14800:].T,linewidth=4,color=col)\n",
    "\n",
    "        str_syn_type = ge_data_temp.iloc[iis[ii],:].loc['cell_type2_pre']+'->'+ge_data_temp.iloc[iis[ii],:].loc['cell_type2_post']\n",
    "        #plt.title(str(iis[ii])+' '+str_syn_type)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    iis = [19,23,11,28,25,29,44,46,10]    \n",
    "    for ii in range(len(iis)):    \n",
    "        if ii<4:\n",
    "            cc=0  # sst\n",
    "            col = np.array([255, 248, 82])/255\n",
    "        if (ii>=4) and (ii<7):\n",
    "            cc=1  # pvalb\n",
    "            col = np.array([0, 69, 245])/255\n",
    "        if ii>=7:\n",
    "            cc=2  # vip \n",
    "            col = np.array([115, 244, 64])/255\n",
    "\n",
    "        f, ax =plt.subplots(figsize=(sx, sy))\n",
    "        dt=0.1\n",
    "        tplus=1\n",
    "        tminus=20\n",
    "        Tmax = 150\n",
    "        ##A = As[0,:Na1] #np.mean(As[:200,:Na1],axis=0)\n",
    "        #A = As[iis[ii],Na1-9:Na1]\n",
    "        #A = np.append([0],A)\n",
    "        #T_A = Tall_A[-9:-1]\n",
    "        #T_A = np.append(T_A,[T_A[-1]+2*Tmax])\n",
    "\n",
    "        A = As[iis[ii],0:Na1]\n",
    "        A = np.append([0],A)\n",
    "        T_A = Tall_A[0:-1]\n",
    "        T_A = np.append(T_A,[T_A[-1]+2*Tmax])\n",
    "\n",
    "        epsp, t = epsp_sim(A, T_A,dt,tplus,tminus,Tmax)\n",
    "        plt.plot(t[14800:], epsp[1,14800:].T,linewidth=4,color=col)\n",
    "\n",
    "        str_syn_type = ge_data_temp.iloc[iis[ii],:].loc['cell_type2_pre']+'->'+ge_data_temp.iloc[iis[ii],:].loc['cell_type2_post']\n",
    "        #plt.title(str(iis[ii])+' '+str_syn_type)    \n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iis = [11,13] #, 14]\n",
    "dn2=100\n",
    "for ii in range(len(iis)):    \n",
    "    col = np.array([200, 200, 245])/255\n",
    "    \n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    dt=0.1\n",
    "    tplus=1\n",
    "    tminus=20\n",
    "    Tmax = 150\n",
    "    ##A = As[0,:Na1] #np.mean(As[:200,:Na1],axis=0)\n",
    "    #A = As[iis[ii],Na1-9:Na1]\n",
    "    #A = np.append([0],A)\n",
    "    #T_A = Tall_A[-9:-1]\n",
    "    #T_A = np.append(T_A,[T_A[-1]+2*Tmax])\n",
    "    Epsp =[]\n",
    "    for iii in range(dn2):\n",
    "        A = As[iis[ii]*dn2+iii,0:Na1]\n",
    "        A = np.append([0],A)\n",
    "        T_A = Tall_A[0:-1]\n",
    "        T_A = np.append(T_A,[T_A[-1]+2*Tmax])\n",
    "\n",
    "        epsp, t = epsp_sim(A, T_A,dt,tplus,tminus,Tmax)\n",
    "        Epsp = Epsp + [epsp[1,:].reshape([-1,1])]\n",
    "     \n",
    "    \n",
    "    amps3 = As[iis[ii]*dn2+np.arange(dn2),0:Na1][:,-8:-1]\n",
    "    ism1 = amps3[:,-1]>0.8\n",
    "    amps3 = amps3[:,:]\n",
    "    amps3 = amps3*(1 + np.random.normal(size=amps3.shape)*0.15)\n",
    "    #positions = t[14800] +3.7+25 + int(25)*np.arange(amps3.shape[1])\n",
    "    \n",
    "    Epsp = np.concatenate(Epsp, axis=1)\n",
    "    \n",
    "    Epsp = Epsp[:,ism1]\n",
    "    Epsp = (np.random.normal(size=Epsp.shape)*0.1+1)*Epsp\n",
    "    plt.plot(t[14800:], Epsp[14800:,:],linewidth=1,color=np.array([0, 69, 245])/255) \n",
    "    \n",
    "    fig,ax = plt.subplots(figsize=(sx, sy))\n",
    "    plt.plot(np.arange(amps3.shape[1])+1,amps3.transpose(),'-o',linewidth=1,color=col)\n",
    "    plt.boxplot(amps3, widths = 0.7)\n",
    "\n",
    "    ax.set_ylim([0,1.7])\n",
    "    \n",
    "#     pyplot.boxplot(x, notch=None, sym=None, vert=None, whis=None, positions=None, widths=None, \n",
    "#                    patch_artist=None, bootstrap=None, usermedians=None, conf_intervals=None, meanline=None, showmeans=None, \n",
    "#                    showcaps=None, showbox=None, showfliers=None, boxprops=None, labels=None, flierprops=None, medianprops=None,\n",
    "#                    meanprops=None, \n",
    "#                    capprops=None, whiskerprops=None, manage_ticks=True, autorange=False, zorder=None, *, data=None)\n",
    "    \n",
    "    str_syn_type = ge_data_temp.iloc[iis[ii]*dn2,:].loc['cell_type2_pre']+'->'+ge_data_temp.iloc[iis[ii],:].loc['cell_type2_post']\n",
    "    #plt.title(str(iis[ii])+' '+str_syn_type) \n",
    "    \n",
    "    fig,ax = plt.subplots(figsize=(sx, sy))\n",
    "    ge_data_temp.iloc[iis[ii]*dn2,:].loc['cell_type2_pre']\n",
    "    \n",
    "    plt.plot(np.arange(amps3.shape[1])+1,amps3.transpose(),'-o',linewidth=1,color=col)\n",
    "\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "amps3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "positions.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt \n",
    "%matplotlib inline\n",
    "# transform labels from TM to An:A1\n",
    "fs = [20 , 10, 40] # Hz\n",
    "N = 8\n",
    "\n",
    "Trec = [250, 1000]\n",
    "Trec=[]\n",
    "DT0 = 25000\n",
    "xs  =ge_data.iloc[l_pre_post2:,:].loc[:,stp_aba_names].values\n",
    "xs  =np.delete(xs, [5,6],axis=1)\n",
    "As=np.zeros((l_pre_post2,0))\n",
    "As2=np.zeros((xs.shape[0],0))\n",
    "Tall_A = np.arange(1)\n",
    "if do_add_aba_synphys:\n",
    "    \n",
    "    for f in fs:\n",
    "        T = np.arange(N)*1000/f\n",
    "        Tall_A = np.append(Tall_A, T+DT0+Tall_A[-1])\n",
    "        \n",
    "        Asf, ns, ps, dpp0, p0, tF, tD, A = STP_sim(ge_data.iloc[0:l_pre_post2,:],T)\n",
    "        As = np.concatenate([As,Asf],axis=1)\n",
    "        \n",
    "        for ri in range(len(Trec)):\n",
    "            Asr, nsr, psr, dpp0, p0, tF, tD, A = STP_sim(ge_data.iloc[0:l_pre_post2,:],[Trec[ri]],init_state=[ns[:,-1],ps[:,-1]])\n",
    "            As = np.concatenate([As,Asr],axis=1)\n",
    "        Tall_A = np.append(Tall_A, np.array(Trec)+Tall_A[-1])\n",
    "        \n",
    "\n",
    "        As2f = np.zeros((xs.shape[0],N+len(Trec)))\n",
    "        for i2 in range(xs.shape[0]):\n",
    "            #as2, sts2 = STP_sim2(xs[i2,:],np.arange(8)*50,model_type = 'tm5') # preconditioning series???\n",
    "            #as2, sts2 = STP_sim2(xs[i2,:], T+DT0, init_state=sts2[-1], model_type='tm5' ) \n",
    "            \n",
    "            as2, sts2 = STP_sim2(xs[i2,:],T,model_type = 'tm5_smr')\n",
    "            \n",
    "            As2f[i2,0:N] = as2\n",
    "            for ri in range(len(Trec)):\n",
    "                as2r, sts2r = STP_sim2(xs[i2,:], [Trec[ri]], init_state=sts2[-1], model_type='tm5_smr' ) \n",
    "                As2f[i2,N+ri] = as2r\n",
    "                \n",
    "        As2 = np.concatenate([As2,As2f],axis=1)\n",
    "            \n",
    "    \n",
    "    As2 = As2/As2[:,0].reshape((-1,1))\n",
    "    As = As/As[:,0].reshape((-1,1))\n",
    "    Na1 = As.shape[1]\n",
    "    As = np.concatenate([As,As2],axis=0)\n",
    "else:\n",
    "    As, ns, ps, dpp0, p0, tF, tD, A = STP_sim(ge_data,T)\n",
    "    \n",
    "\n",
    "sx=12\n",
    "sy=3\n",
    "f, ax =plt.subplots(figsize=(sx, sy))\n",
    "plt.plot(As2[0:1,:].transpose(),'o-',linewidth=4, markersize=12)\n",
    "plt.ylim([0,1.2])\n",
    "plt.xticks(fontsize=18)\n",
    "plt.xticks(ticks=[0,5,10,15,20],fontsize=18)\n",
    "plt.yticks(ticks=[0.5,1.0,1.5],fontsize=18)\n",
    "plt.ylabel('An:A1',fontsize=20)\n",
    "plt.xlabel('Stimulus number',fontsize=20)\n",
    "#plt.legend(fontsize=20)\n",
    "#print(set(n_stp_cells.index).difference(set(i2)))\n",
    "\n",
    "\n",
    "f, ax =plt.subplots(figsize=(sx, sy))\n",
    "dt=0.1\n",
    "tplus=1\n",
    "tminus=10\n",
    "Tmax = 300\n",
    "#A = As[0,:Na1] #np.mean(As[:200,:Na1],axis=0)\n",
    "A = As[0,:Na1]\n",
    "A = np.append([0],A)\n",
    "T_A = Tall_A\n",
    "T_A = np.append(T_A,[T_A[-1]+Tmax])\n",
    "\n",
    "epsp, t = epsp_sim(A, T_A,dt,tplus,tminus,Tmax)\n",
    "plt.plot(t, epsp[1,:].T,linewidth=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax =plt.subplots(figsize=(sx, sy))\n",
    "dt=0.1\n",
    "tplus=5\n",
    "tminus=50\n",
    "Tmax = 300\n",
    "\n",
    "DT = 200/6\n",
    "\n",
    "T_A = np.array([0]+(100+np.arange(9)*200/6).tolist()+[100+200/6*8+1000])\n",
    "\n",
    "i0=23 #42 #[18,23,42]\n",
    "ii = i0*200 + np.arange(200)\n",
    "A = np.mean(As[ii,:Na1],axis=0) #As[0,:Na1]#np.mean(As[:100,:Na1],axis=0)\n",
    "A = np.append([0],A)\n",
    "#T_A = Tall_A\n",
    "T_A = np.append(T_A,[T_A[-1]+400])\n",
    "\n",
    "epsp, t = epsp_sim(A, T_A,dt,tplus,tminus,Tmax)\n",
    "plt.plot(t, epsp[1,:].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## START FROM NEXT SECTION : Classify STP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Classify STP for cortex : 3 types by D/F ratio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "\n",
    "#\n",
    "#do_medians, do_log should be off!!!\n",
    "# Nsamples = 200 !!!\n",
    "\n",
    "\n",
    "## some distributions\n",
    "#stp_columns\n",
    "\n",
    "# f, ax =plt.subplots(figsize=(16, 7))\n",
    "# plt.plot(ge_data.iloc[3600:3800,:].loc[:,stp_columns].transpose(),'o-')\n",
    "# plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "#          rotation_mode=\"anchor\")\n",
    "\n",
    "#f, ax =plt.subplots(figsize=(16, 7))\n",
    "stp_columns2=np.array(stp_columns)\n",
    "annot_columns2 =np.array(annot_columns)\n",
    "sx=5\n",
    "sy=4\n",
    "\n",
    "\n",
    "i0=0\n",
    "Nbootstraps = 100\n",
    "\n",
    "skip_this=1\n",
    "\n",
    "Stp_type_p = []\n",
    "\n",
    "for i0 in range(0,53):#[42]: # [18,23,42] #range(0,53):\n",
    "    str_syntype = ge_data.loc[i0*Nbootstraps,annot_columns2[[0,2,1,3]]].astype(str).str.cat(others=['_','->','_','']).str.cat()\n",
    "\n",
    "    #print(ge_data.loc[i0*200,annot_columns[:-3]])\n",
    "\n",
    "    ii = i0*Nbootstraps + np.arange(Nbootstraps)\n",
    "    \n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    dat=ge_data.iloc[ii,:].loc[:,'tD'].map(float)\n",
    "    #dat=dat.loc[dat.loc[:]<800]\n",
    "    dat.loc[dat.loc[:]>500]=500\n",
    "    dat2=ge_data.iloc[ii,:].loc[:,'tF'].map(float)\n",
    "    #dat2=dat2.loc[dat2.loc[:]<5500]\n",
    "    dat2.loc[dat2.loc[:]>500]=500\n",
    "    plt.scatter(dat,dat2)\n",
    "    \n",
    "    #stp_ns =          ['tF',    'p0',    'tD',   'dp/p0', 'A']\n",
    "    dat_p0=ge_data.iloc[ii,:].loc[:,'p0'].map(float)\n",
    "    dat_dp=ge_data.iloc[ii,:].loc[:,'dp/p0'].map(float)\n",
    "    \n",
    "    #dat.hist(bins=25)\n",
    "    plt.title(str(i0)+ '  '+str_syntype + '  ' +'D->F')\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.xlabel('D, ms',fontsize=20)\n",
    "    plt.ylabel('F, ms',fontsize=20)\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    #ge_data.iloc[ii,:].loc[:,'tF'].map(float).hist(bins=25)\n",
    "    \n",
    "    dat3 = dat2/dat\n",
    "    dat3.hist(bins=25)\n",
    "    plt.title(str(i0)+ '  '+str_syntype + '  ' +'F/D')\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.xlabel('F/D',fontsize=20)\n",
    "    plt.ylabel('Frequency',fontsize=20)\n",
    "    nsam = dat3.shape[0]\n",
    "    #print(str(i0)+ '  '+str_syntype + '  ' +'F/D  :','<0.2 - ',(dat3<0.2).sum()/nsam,\n",
    "    #      '; >0.2, <2 - ',((dat3>0.2)&(dat3<2)).sum()/nsam, '; >2 - ',(dat3>2).sum()/nsam  )\n",
    "#     stp_type_p = [str(i0)+ '  '+str_syntype,\n",
    "#                   (dat3<0.4).sum()/nsam, ((dat3>0.4)&(dat3<2)).sum()/nsam, (dat3>2).sum()/nsam ,\n",
    "#                   (dat3.iloc[0]<0.4), (dat3.iloc[0]>0.4)&(dat3.iloc[0]<2), (dat3.iloc[0]>2),\n",
    "#                   np.median(dat3),dat3.iloc[0], dat3.values,\n",
    "#                   np.array([dat.iloc[0],dat2.iloc[0],dat_p0.iloc[0], dat_dp.iloc[0]])]\n",
    "    if i0<10:\n",
    "        stp_names = ['tD', 'tF','p0','dp/p0']\n",
    "    else:\n",
    "        # 'tF', 'p0','tD','dp','A','A1','A2','tDmin','dd','t_FDR','t_SMR','dp0'\n",
    "        stp_names = ['tD', 'tF','p0','dp','t_SMR','dp0']\n",
    "    stp_type_p = [str(i0)+ '  '+str_syntype,\n",
    "                  (dat3<0.4).sum()/nsam, ((dat3>0.4)&(dat3<2)).sum()/nsam, (dat3>2).sum()/nsam ,\n",
    "                  (dat3.iloc[0]<0.4), (dat3.iloc[0]>0.4)&(dat3.iloc[0]<2), (dat3.iloc[0]>2),\n",
    "                  np.median(dat3),dat3.iloc[0], dat3.values,\n",
    "                  np.array([dat.iloc[0],dat2.iloc[0],dat_p0.iloc[0], dat_dp.iloc[0]]),\n",
    "                  ge_data.iloc[ii[0],:].loc[stp_names].values]\n",
    "    Stp_type_p = Stp_type_p + [stp_type_p]\n",
    "    \n",
    "    \n",
    "    \n",
    "    print(str(i0)+ '  '+str_syntype + '  ' +'F/D  :',stp_type_p   )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### STP classes in hippocampus based on D/F ratio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "\n",
    "#\n",
    "#do_medians, do_log should be off!!!\n",
    "# Nsamples = 200 !!!\n",
    "\n",
    "\n",
    "## some distributions\n",
    "#stp_columns\n",
    "\n",
    "# f, ax =plt.subplots(figsize=(16, 7))\n",
    "# plt.plot(ge_data.iloc[3600:3800,:].loc[:,stp_columns].transpose(),'o-')\n",
    "# plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "#          rotation_mode=\"anchor\")\n",
    "\n",
    "#f, ax =plt.subplots(figsize=(16, 7))\n",
    "stp_columns2=np.array(stp_columns)\n",
    "annot_columns2 =np.array(annot_columns)\n",
    "sx=5\n",
    "sy=4\n",
    "i0=0\n",
    "skip_this=1\n",
    "\n",
    "Stp_type_p_h = []\n",
    "\n",
    "for i0 in range(0,40):#[42]: # [18,23,42] #range(0,53):\n",
    "    str_syntype = ge_data_h.loc[i0*Nbootstraps,annot_columns2[[0,2,1,3]]].astype(str).str.cat(others=['_','->','_','']).str.cat()\n",
    "\n",
    "    #print(ge_data.loc[i0*200,annot_columns[:-3]])\n",
    "\n",
    "    ii = i0*Nbootstraps + np.arange(Nbootstraps)\n",
    "    \n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    dat=ge_data_h.iloc[ii,:].loc[:,'tD'].map(float)\n",
    "    #dat=dat.loc[dat.loc[:]<800]\n",
    "    dat.loc[dat.loc[:]>500]=500\n",
    "    dat2=ge_data_h.iloc[ii,:].loc[:,'tF'].map(float)\n",
    "    #dat2=dat2.loc[dat2.loc[:]<5500]\n",
    "    dat2.loc[dat2.loc[:]>500]=500\n",
    "    plt.scatter(dat,dat2)\n",
    "    \n",
    "    #stp_ns =          ['tF',    'p0',    'tD',   'dp/p0', 'A']\n",
    "    dat_p0=ge_data.iloc[ii,:].loc[:,'p0'].map(float)\n",
    "    dat_dp=ge_data.iloc[ii,:].loc[:,'dp/p0'].map(float)\n",
    "    \n",
    "    #dat.hist(bins=25)\n",
    "    plt.title(str(i0)+ '  '+str_syntype + '  ' +'D->F')\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.xlabel('D, ms',fontsize=20)\n",
    "    plt.ylabel('F, ms',fontsize=20)\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    #ge_data.iloc[ii,:].loc[:,'tF'].map(float).hist(bins=25)\n",
    "    \n",
    "    dat3 = dat2/dat\n",
    "    dat3.hist(bins=25)\n",
    "    plt.title(str(i0)+ '  '+str_syntype + '  ' +'F/D')\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.xlabel('F/D',fontsize=20)\n",
    "    plt.ylabel('Frequency',fontsize=20)\n",
    "    nsam = dat3.shape[0]\n",
    "    #print(str(i0)+ '  '+str_syntype + '  ' +'F/D  :','<0.2 - ',(dat3<0.2).sum()/nsam,\n",
    "    #      '; >0.2, <2 - ',((dat3>0.2)&(dat3<2)).sum()/nsam, '; >2 - ',(dat3>2).sum()/nsam  )\n",
    "    #stp_type_p = [str(i0)+ '  '+str_syntype, (dat3<0.4).sum()/nsam, ((dat3>0.4)&(dat3<2)).sum()/nsam, (dat3>2).sum()/nsam, np.median(dat3), dat3.values ]\n",
    "    \n",
    "    stp_names = ['tD', 'tF','p0','dp/p0']\n",
    "    stp_type_p = [str(i0)+ '  '+str_syntype,\n",
    "                  (dat3<0.4).sum()/nsam, ((dat3>0.4)&(dat3<2)).sum()/nsam, (dat3>2).sum()/nsam ,\n",
    "                  (dat3.iloc[0]<0.4), (dat3.iloc[0]>0.4)&(dat3.iloc[0]<2), (dat3.iloc[0]>2),\n",
    "                  np.median(dat3),dat3.iloc[0], dat3.values,\n",
    "                  np.array([dat.iloc[0],dat2.iloc[0],dat_p0.iloc[0], dat_dp.iloc[0]]),\n",
    "                 ge_data.iloc[ii[0],:].loc[stp_names].values]\n",
    "    Stp_type_p_h = Stp_type_p_h + [stp_type_p]\n",
    "    \n",
    "    print(str(i0)+ '  '+str_syntype + '  ' +'F/D  :',stp_type_p   )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_Stp_type_p = pd.DataFrame(Stp_type_p + Stp_type_p_h, columns=['name','type_2','type_3','type_1',\n",
    "                                                                 'dtype_2','dtype_3','dtype_1',\n",
    "                                                                 'median','F_D_0','F_D','tm_DFUdU','tm_all'])\n",
    "df_Stp_type_p.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_Stp_type_p.to_excel('stp types probability cortex and hippocampus.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dat.iloc[0])\n",
    "print(dat2.iloc[0])\n",
    "print(dat_p0.iloc[0]) \n",
    "print(dat_dp.iloc[0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot EPSPs for different synaptic types (aka fig9b)"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_Stp_type_p = pd.read_excel('stp types probability cortex and hippocampus.xlsx')\n",
    "df_Stp_type_p['polarity'] = 'inh'\n",
    "\n",
    "is_PC_presynaptic = ge_data.loc[np.arange(df_Stp_type_p.shape[0])*Nbootstraps,annot_columns2[0]].astype(str).str.contains('PC').values\n",
    "df_Stp_type_p.loc[is_PC_presynaptic,'polarity']='ex'\n",
    "\n",
    "ex_types = ['CA1', 'CA3', 'EC','DG']\n",
    "for ex_type in ex_types:\n",
    "    is_PC_presynaptic = ge_data_h.loc[np.arange(df_Stp_type_p.shape[0])*Nbootstraps,annot_columns2[0]].astype(str).str.contains(ex_type).values\n",
    "    is_PC_presynaptic = np.nonzero(is_PC_presynaptic )[0]+53 \n",
    "    df_Stp_type_p.loc[is_PC_presynaptic,'polarity']='ex'    \n",
    "\n",
    "df_Stp_type_p.loc[53:73,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Q_TM_aba_synphys(x, amps,sig,Ts,DT0,model_type='tm5'):\n",
    "    #Q, A, A3 =  QA_TM_aba_synphys(x, amps,sig,Ts)\n",
    "    Q, A =  QA_TM_aba_synphys(x, amps,sig,Ts,DT0,model_type)\n",
    "    return Q    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def QA_TM_aba_synphys(x, amps,sig,Ts,DT0,model_type='tm5'):\n",
    "    #amps=args[0] \n",
    "    #sig=args[1] \n",
    "    #Ts=args[2]\n",
    "    \n",
    "    \n",
    "    fl=0\n",
    "    if fl==0:\n",
    "        ams = x[4:7]\n",
    "        #breakpoint()\n",
    "        if len(x)>7:\n",
    "            if model_type=='tm5_fdr2':\n",
    "                x2=np.copy(np.delete(x,[5,6]))\n",
    "            if model_type=='tm5_smr':\n",
    "                x2=np.copy(np.delete(x,[5,6]))    \n",
    "            \n",
    "            #x2[5:] = x[7:]\n",
    "        else:\n",
    "            x2=np.copy(x[0:5])\n",
    "    \n",
    "    n  = [0, 32, 44]\n",
    "    t2 = [6, 1,  1]\n",
    "    #T2  =[125, 250, 500, 1000, 2000, 4000]\n",
    "    eps = np.finfo(float).eps\n",
    "    z4 = np.arange(4)\n",
    "    z8 = np.arange(8)\n",
    "    A = np.zeros((amps.shape[0],))\n",
    "    #A3 = np.ones((amps.shape[0],))\n",
    "    #S3 = np.zeros((amps.shape[0],))\n",
    "    Q=0\n",
    "    for ii in range(3):\n",
    "        vv = z8 + n[ii]\n",
    "        A1e = amps[vv]  \n",
    "        S1e = sig[vv]\n",
    "        if fl==0:\n",
    "            x2[4] = ams[ii] # independent amplitudes for each stimulation frequency - to account for rundown etc.\n",
    "        \n",
    "        #vv_nonz = np.nonzero(np.abs(A1e[:])>eps*1e3)[0]\n",
    "        vv_z = np.abs(A1e[:])<eps*1e3\n",
    "        #print(vv_nonz)\n",
    "        #if len(vv_nonz)>0:\n",
    "        if sum(vv_z)!=len(A1e):\n",
    "            # preconditioning:\n",
    "            A1, states = STP_sim(x2, Ts[vv],model_type=model_type ) \n",
    "            np0 = states[-1]\n",
    "            \n",
    "            # first 8 stimuli responces:\n",
    "            #A1, states = STP_sim(x2, Ts[vv]+DT0, init_state=np0, model_type=model_type ) \n",
    "            #breakpoint()\n",
    "            A1[vv_z] = 0\n",
    "            A[vv] = A1\n",
    "            \n",
    "            #Q = Q + np.sum((A1e[vv_nonz]-A1[vv_nonz])**2/(S1e[vv_nonz]**2 + eps))\n",
    "            Q = Q + np.sum((A1e-A1)**2/(S1e**2 + eps))\n",
    "            # recovery responces\n",
    "            #np0 = [n12[-1], p12[-1], d12[-1]]\n",
    "            np0 = states[-1]\n",
    "            for i3 in range(t2[ii]):\n",
    "                vv2 = vv[-1]+1+z4+4*i3\n",
    "                A2e = amps[vv2] \n",
    "                S2e = sig[vv2] \n",
    "                ##A3[vv2] = np.median(A1e)\n",
    "                #vv_nonz2 = np.nonzero(np.abs(A2e[:])>eps*1e3)[0]\n",
    "                #print(vv_nonz2)\n",
    "                vv_z2 = np.abs(A2e[:])<eps*1e3\n",
    "                #if len(vv_nonz2)>0:\n",
    "                if sum(vv_z2)!=len(A2e):\n",
    "                    A2, states2 = STP_sim(x2, Ts[vv2], init_state=np0, model_type=model_type )\n",
    "                    \n",
    "                    A2[vv_z2] = 0\n",
    "                    \n",
    "                    A[vv2] = A2\n",
    "                    \n",
    "                    #breakpoint()\n",
    "                    #A2 = A2/np.median(A1)\n",
    "                    #A[vv2] = A2\n",
    "                    #breakpoint()\n",
    "                    \n",
    "                    #Q = Q + np.sum((A2e[vv_nonz2]-A2[vv_nonz2])**2/(S2e[vv_nonz2]**2 + eps))\n",
    "                    Q = Q + np.sum((A2e-A2)**2/(S2e**2 + eps))\n",
    "    \n",
    "    return Q, A #, A3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "x = [1.00000000e+01, 4.00000000e-01, 1.00000000e+02, 4.00000000e-01,\n",
    " 4.59968078e+00, 4.59968078e+00,4.59968078e+00, 1.00000000e+01,\n",
    " 5.00000000e-04, 1.00000000e+02]\n",
    "\n",
    "Ts = df.loc[1,'comments1']\n",
    "Ts = np.array(Ts[0][:].strip().split()).astype(float)\n",
    "print(Ts)\n",
    "\n",
    "vv=np.arange(len(As))\n",
    "npar=56\n",
    "amps = np.ones(npar)\n",
    "amps[vv]=As*1.05\n",
    "sig = amps/5\n",
    "DT0=2500000\n",
    "model_type = 'tm5'\n",
    "Q2,As2 = QA_TM_aba_synphys(x[0:7], amps,sig,Ts,DT0,model_type )\n",
    "plt.plot(As2)\n",
    "As2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#i=1\n",
    "def fit_TM_cell(i,df,par):\n",
    "    #breakpoint()\n",
    "    aba_case=False\n",
    "    if 'aba_case' in par:\n",
    "        aba_case = par['aba_case']\n",
    "        \n",
    "    model_type='tm5'\n",
    "    if 'model_type' in par:\n",
    "        model_type = par['model_type']   \n",
    "        \n",
    "    dont_plot=0\n",
    "    if 'dont_plot' in par:\n",
    "        dont_plot = par['dont_plot'] \n",
    "    \n",
    "    if par['fited_pulses']==True:\n",
    "        ampsa = df.loc[i,'stp_mean_fit_bootstrap']\n",
    "        siga = df.loc[i,'stp_sigma_fit_bootstrap']  \n",
    "        #if abs(ampsa[0])<1e-14:\n",
    "            #ampsa = df.loc[i,'stp_mean_bootstrap']\n",
    "            #siga = df.loc[i,'stp_sigma_bootstrap']\n",
    "    else:\n",
    "        ampsa = df.loc[i,'stp_mean_bootstrap']\n",
    "        siga = df.loc[i,'stp_sigma_bootstrap']\n",
    "        \n",
    "    Ts = df.loc[i,'comments1']\n",
    "    Ts = np.array(Ts[0][:].strip().split()).astype(float)\n",
    "    \n",
    "    #breakpoint()\n",
    "\n",
    "    #A3 = df.loc[i,'comments']\n",
    "    #A3i = np.array(A3[0][:].strip().split()).astype(float)\n",
    "    #A4 = np.ones((len(A3),len(A3i)))\n",
    "    #for ia3 in range(len(A3)):\n",
    "    #    A3i = np.array(A3[ia3][:].strip().split()).astype(float)\n",
    "    #    A4[ia3,:] = A3i\n",
    "    ##A4.shape \n",
    "    \n",
    "    \n",
    "    ibs  = par['ibs']\n",
    "    npar = 56\n",
    "    #amps = ampsa[ibs*npar + np.arange(npar)]\n",
    "    ampsa1 = ampsa.reshape((-1,npar))\n",
    "    amps = ampsa1[ibs,:].ravel()\n",
    "    \n",
    "    if aba_case:\n",
    "        # sig  = siga #[(ibs-1)*npar + np.arange(npar)]\n",
    "        #A3e  = A4[:,ibs ].ravel()\n",
    "\n",
    "        siga1 = siga.reshape((-1,npar))\n",
    "        sig = siga1[ibs,:].ravel()\n",
    "    else:\n",
    "        #sig = np.std(ampsa1,axis=0).ravel() # reestimate sigma\n",
    "        sig=pd.DataFrame(ampsa1).std(axis=0).values\n",
    "        \n",
    "    DT0 = 25000\n",
    "    #breakpoint()\n",
    "    # restrict set of fited data\n",
    "    vv = par['amplitudes_selected_for_TM_training']\n",
    "    #vv = np.arange(8)    # A 1:8\n",
    "    #vv = np.arange(npar) # all\n",
    "    #vv = np.concatenate([vv,vv+32, vv+44]) # without recovery\n",
    "    \n",
    "    amps0 = amps\n",
    "    amps = np.zeros(amps.shape)\n",
    "    amps[vv] = amps0[vv]\n",
    "\n",
    "         #  tF   p0     tD  dp         A            A1           A2      tDmin  dd   t_FDR t_SMR dp0 \n",
    "    x0 = [ 100,  0.1,  100, 0.1, amps[0]/0.1, amps[0]/0.1,  amps[0]/0.1, 10,   0.05, 100,  100,  0.02 ]\n",
    "    \n",
    "    x0 = np.array(x0)\n",
    "    \n",
    "    #dxx = np.array([1e3, 20, 1e3, 1e2,  20, 20, 20,  1e1, 20, 1e2 ])\n",
    "    dxx = np.array([1e3,  10, 1e3, 10,  20, 20, 20,  1e3, 20, 1e3, 1e3, 50 ])\n",
    "    \n",
    "    #breakpoint()\n",
    "    if model_type=='tm5':\n",
    "        x0 = x0[0:7]\n",
    "        dxx = dxx[0:7]\n",
    "        \n",
    "    if model_type=='tm4':\n",
    "        x0 = x0[0:7]\n",
    "        dxx = dxx[0:7]\n",
    "        \n",
    "    if model_type=='tm5_fdr2':\n",
    "        x0 = x0[0:10]\n",
    "        dxx = dxx[0:10]\n",
    "        \n",
    "    if model_type=='tm5_smr':\n",
    "        x0 = x0    \n",
    "        dxx = dxx\n",
    "    #print(x0)\n",
    "    Q0, A0 = QA_TM_aba_synphys(x0, amps, sig, Ts,DT0,model_type)\n",
    "    #print(x0)\n",
    "    print('initial Q: ',str(Q0))\n",
    "\n",
    "    #           tF      p0    tD        dp           A                A1                  A2              tDmin   dd   tD2\n",
    "    #xlower = [  0.1,   0.01,   0.1,     0.0001,     amps[0]/1*1e-1,   amps[0]/1*1e-1,     amps[0]/1*1e-1,  1,     1e-3, 1]\n",
    "    xlower = x0/dxx\n",
    "    xlower = np.array(xlower)\n",
    "    #xupper = [  1e5,   1,      1e5,     1,          amps[0]/0.01*1e1, amps[0]/0.01*1e1,   amps[0]/0.01*1e1, 100,  1,    1e4]\n",
    "    #xupper = np.array(xupper)\n",
    "    xupper = x0*dxx\n",
    "    bounds = [] #np.zeros(len(xlower))\n",
    "    for ibo in range(len(xlower)):\n",
    "        bounds = bounds + [(xlower[ibo],xupper[ibo])]\n",
    "    tt=[]\n",
    "    #import time\n",
    "    #import scipy.optimize as opt\n",
    "    t1 = time.time()\n",
    "    \n",
    "    #print(x0)\n",
    "    #print(A0)\n",
    "    #breakpoint()\n",
    "    \n",
    "    do_this=0\n",
    "    if do_this==1:\n",
    "        res1 = opt.shgo(Q_TM_aba_synphys, bounds , args=((amps, sig, Ts,DT0,model_type)), constraints=None, n=10,\n",
    "                            iters=1, callback=None, minimizer_kwargs=None, options=None, sampling_method='simplicial')\n",
    "        x_shgo = res1.x\n",
    "    else:\n",
    "        x_shgo = x0\n",
    "    \n",
    "    t2 = time.time()\n",
    "    tt = tt + [t2-t1]\n",
    "    print('shgo : elapsed time '+str(t2-t1)+'s ')\n",
    "\n",
    "    t1 = time.time()\n",
    "    do_good=True\n",
    "    if do_good:\n",
    "        popsize = 65 # best\n",
    "        maxiter=1000\n",
    "        tol=0.01\n",
    "    else:\n",
    "        popsize = 5  # fast\n",
    "        maxiter=10\n",
    "        tol=0.1\n",
    "        \n",
    "    #breakpoint()    \n",
    "    do_this=0\n",
    "    if do_this==1:\n",
    "        res2 = opt.differential_evolution(Q_TM_aba_synphys, bounds , args=((amps, sig, Ts,DT0,model_type)), strategy='best1bin',\n",
    "                           maxiter=maxiter, popsize=popsize, tol=tol, mutation=(0.5, 1), recombination=0.7, seed=None,\n",
    "                           callback=None, disp=False, polish=True, init='latinhypercube', atol=0,\n",
    "                           updating='immediate', workers=1)\n",
    "        x_de = res2.x\n",
    "    else:\n",
    "        x_de = x0  \n",
    "        \n",
    "    t2 = time.time()\n",
    "    tt = tt + [t2-t1]\n",
    "    print('differential_evolution: elapsed time '+str(t2-t1)+'s ')\n",
    "\n",
    "    t1 = time.time()\n",
    "    do_good=False\n",
    "    if do_good:\n",
    "        maxiter=1000 # good\n",
    "        accept=-5.0 \n",
    "        maxfun=10000000.0\n",
    "    else:\n",
    "        maxiter=10 # fast\n",
    "        maxfun=10.0\n",
    "        accept=-1.0\n",
    "    fl=0  \n",
    "    do_this=0\n",
    "    if do_this==1:\n",
    "        res3 = opt.dual_annealing(Q_TM_aba_synphys, bounds , args=((amps, sig, Ts,DT0,model_type)), maxiter=1000, \n",
    "                              local_search_options={},\n",
    "                              initial_temp=5230.0,restart_temp_ratio=2e-05, visit=2.62, \n",
    "                              accept=accept, maxfun=maxfun, seed=None,\n",
    "                              no_local_search=False, callback=None, x0=None)\n",
    "        x_da = res3.x\n",
    "    else:\n",
    "        x_da = x0 \n",
    "        \n",
    "    t2 = time.time()\n",
    "    Q, A = QA_TM_aba_synphys(x_da, amps, sig, Ts,DT0,model_type)\n",
    "    A_da=A\n",
    "    tt = tt + [t2-t1]\n",
    "    print('dual_annealing: elapsed time '+str(t2-t1)+'s ')\n",
    "\n",
    "    t1 = time.time()\n",
    "    do_this=0\n",
    "    if do_this==1:\n",
    "        res4 = opt.basinhopping(Q_TM_aba_synphys, x0 , niter=1, T=1.0, stepsize=1.0,\n",
    "                            minimizer_kwargs={'method':\"L-BFGS-B\", 'args':(amps, sig, Ts,DT0,model_type)}, \n",
    "                            take_step=None, accept_test=None, callback=None, interval=50, disp=False,\n",
    "                           niter_success=None, seed=None)\n",
    "        x_bh = res4.x\n",
    "    else:\n",
    "        x_bh = x0 \n",
    "        \n",
    "    t2 = time.time()\n",
    "    tt = tt + [t2-t1]\n",
    "    print('basinhopping: elapsed time '+str(t2-t1)+'s ')\n",
    "    #t1 = time.time()\n",
    "    #res4 = opt.basinhopping(Q_TM_aba_synphys, x0 , niter=100, T=1.0, stepsize=1.0,\n",
    "    #                        minimizer_kwargs={'method':\"L-BFGS-B\", 'args':(amps, sig, Ts,DT0)}, \n",
    "    #                        take_step=None, accept_test=None, callback=None, interval=50, disp=False,\n",
    "    #                        niter_success=None, seed=None)\n",
    "    t1 = time.time()\n",
    "    do_this=1\n",
    "    if do_this:\n",
    "        prob = pg.problem(jit_rosenbrock2(amps.reshape((-1,1)),sig.reshape((-1,1)),Ts,model_type,\n",
    "                                          np.array(xlower).reshape((-1,)),np.array(xupper).reshape((-1,))))\n",
    "        #pop = pg.population(prob=prob, size = 1000)\n",
    "        uda = pg.de1220(gen = 100, allowed_variants = [2,3,7,10,13,14,15,16], variant_adptv = 1, ftol = 1e-6, xtol = 1e-6, memory = False)\n",
    "        algo = pg.algorithm(uda)\n",
    "        archi = pg.archipelago(n=32,algo=algo, prob=prob, pop_size=150)\n",
    "        \n",
    "        #pop2 = algo.evolve(pop)\n",
    "        archi.evolve() \n",
    "        archi.wait()\n",
    "        \n",
    "        qq=np.array(archi.get_champions_f())\n",
    "        xx=np.array(archi.get_champions_x())\n",
    "        \n",
    "        x_sade = np.exp(xx[qq.argmin(),:])\n",
    "\n",
    "    else:\n",
    "        x_sade  = x0 \n",
    "    \n",
    "    t2 = time.time()\n",
    "    Q, A = QA_TM_aba_synphys(x_sade, amps, sig, Ts,DT0,model_type)\n",
    "    A_sade=A\n",
    "    print('self adaptive differential evolution: time ' + str(t2 - t1), '; Q ',str(Q)) \n",
    "    \n",
    "    \n",
    "    \n",
    "    t1 = time.time()\n",
    "    do_this=0\n",
    "    if do_this==1:\n",
    "        #res4 = opt.basinhopping(Q_TM_aba_synphys, x0 , niter=100, T=1.0, stepsize=1.0,\n",
    "        #                        minimizer_kwargs={'method':\"L-BFGS-B\", 'args':(amps, sig, Ts,DT0)}, \n",
    "        #                        take_step=None, accept_test=None, callback=None, interval=50, disp=False,\n",
    "        #                        niter_success=None, seed=None)\n",
    "        #nst=50 # good\n",
    "        nst = 10 #150 # better fit\n",
    "        Qs =  np.zeros(nst)\n",
    "        As =  np.zeros((nst,npar))\n",
    "        #vx = np.arange(5)  # fit 1 amplitude\n",
    "        vx = np.arange(len(x0)) # fit separate amplitude for each frequency\n",
    "        nx = len(x0)\n",
    "        xs =  np.zeros((nst,len(x0[vx]) ))\n",
    "        for ist in range(nst):\n",
    "\n",
    "            #x0i = np.exp(np.log(xlower) + np.random.rand(len(xlower))*(np.log(xupper)-np.log(xlower)))\n",
    "            x0i = x0*np.exp( 2*(np.random.rand(len(xlower))-0.5)*np.log(dxx)  )\n",
    "            x0i = x0i[vx]\n",
    "            #x0i[4:7] = x0i[4]\n",
    "            fl=1\n",
    "            #breakpoint()\n",
    "            resi=opt.minimize(Q_TM_aba_synphys, x0i, args=((amps, sig, Ts,DT0,model_type)), method=None, jac=None,\n",
    "                      hess=None, hessp=None, bounds=bounds[0:nx],\n",
    "                      constraints=(), tol=None, callback=None, \n",
    "                      options=None)\n",
    "            x=resi.x\n",
    "            Q,A = QA_TM_aba_synphys(x, amps, sig, Ts,DT0,model_type)\n",
    "            As[ist,:] = A\n",
    "            xs[ist,:] = x\n",
    "            Qs[ist] = Q\n",
    "\n",
    "            \n",
    "            Q_rs = np.min(Qs)\n",
    "            ia = np.nonzero(Q_rs==Qs)[0]\n",
    "            x_rs=xs[ia,:].ravel()\n",
    "            A_rs = As[ia,:].ravel()\n",
    "    else: \n",
    "        x_rs=x0\n",
    "        Q_rs = 1e5\n",
    "        A_rs = np.zeros((npar,))\n",
    "            \n",
    "    t2=time.time()\n",
    "    tt = tt + [t2-t1]\n",
    "    print('random_start: elapsed time '+str(t2-t1)+'s ')\n",
    "    Q2 = []\n",
    "    algs = ['shgo', 'differential_evolution','dual_annealing','basinhopping', 'self ajusting differential evolution',  'random_starts' ]\n",
    "    colors = ['y',   'm',                           'c',          'y',          'r',                                 'g' ]\n",
    "        \n",
    "    if dont_plot==0:\n",
    "        #%matplotlib\n",
    "        #import matplotlib.pyplot as plt\n",
    "        A3e=1\n",
    "        \n",
    "        plt.errorbar(np.arange(len(amps)), amps*A3e, yerr=sig*A3e, fmt =  'bo')\n",
    "        #plt.hold(True)\n",
    "\n",
    "\n",
    "        ii=0\n",
    "        x = x_shgo #res1.x\n",
    "        Q, A = QA_TM_aba_synphys(x, amps, sig, Ts,DT0,model_type)\n",
    "        Q2 = Q2 + [Q]\n",
    "        #plt.plot(np.arange(len(amps)), A, 'o-' ,color=colors[ii])\n",
    "        print(algs[ii]+ ' Q = '+str(Q)+' '+colors[ii])\n",
    "\n",
    "\n",
    "\n",
    "        ii=1\n",
    "        x = x_de #res2.x\n",
    "        x2=x\n",
    "        Q, A = QA_TM_aba_synphys(x, amps, sig, Ts,DT0,model_type)\n",
    "        Q2 = Q2 + [Q]\n",
    "        #plt.plot(np.arange(len(amps)), A, 'd-' ,color=colors[ii])\n",
    "        print(algs[ii]+ ' Q = '+str(Q)+' '+colors[ii])\n",
    "        ii=2\n",
    "        x = x_da #res3.x\n",
    "        #x3=x\n",
    "        Q, A = QA_TM_aba_synphys(x, amps, sig, Ts,DT0,model_type)\n",
    "        Q2 = Q2 + [Q]\n",
    "        #plt.plot(np.arange(len(amps)), A, 'x--',color=colors[ii] )\n",
    "        print(algs[ii]+ ' Q = '+str(Q)+' '+colors[ii])\n",
    "        ii=3\n",
    "        x = x_bh #res4.x\n",
    "        Q, A = QA_TM_aba_synphys(x, amps, sig, Ts,DT0,model_type)\n",
    "        Q2 = Q2 + [Q]\n",
    "        \n",
    "        #plt.plot(np.arange(len(amps)), A, '.--',color=colors[ii] )\n",
    "        print(algs[ii]+ ' Q = '+str(Q)+' '+colors[ii])\n",
    "\n",
    "        ii=4\n",
    "\n",
    "        x = x_sade #res4.x\n",
    "        Q, A = QA_TM_aba_synphys(x, amps, sig, Ts,DT0,model_type)\n",
    "        Q2 = Q2 + [Q]\n",
    "        plt.plot(np.arange(len(amps)), A, '.-',color='r')\n",
    "        print(algs[ii]+ ' Q = '+str(Q)+' '+colors[ii])\n",
    "\n",
    "        ii=5\n",
    "\n",
    "\n",
    "        #Q, A = QA_TM_aba_synphys(x, amps, sig, Ts,DT0)\n",
    "        Q=Q_rs \n",
    "        Q2 = Q2 + [Q]\n",
    "        #plt.plot(np.arange(len(amps)), A_rs, '.--',color=colors[ii] )\n",
    "        print(algs[ii]+ ' Q = '+str(Q)+' '+colors[ii])\n",
    "\n",
    "        #algs_short = ['differential_evolution',  'random_starts' ]\n",
    "        algs_short = ['random_starts' ]\n",
    "        #plt.legend(algs)\n",
    "        plt.legend(algs_short)\n",
    "\n",
    "        #fig = plt.gcf()\n",
    "        #set_position(plt.gca, [10, 10, 100, 500], which='both')\n",
    "        plt.xlim((min(vv)-1,max(vv)+2)) \n",
    "        #thismanager = plt.get_current_fig_manager()\n",
    "        #thismanager.window.SetPosition((500, 0))\n",
    "        #plt.hold(False)\n",
    "        plt.pause(0.5) \n",
    "    \n",
    "    #fig = plt.gcf()\n",
    "    AA = [A_sade,A_da,A_rs]\n",
    "    xx = [x_sade,x_da,x_rs]\n",
    "    return Q2, tt, xx, algs, colors, AA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy as sci\n",
    "def epsp_sim(A, T_A,dt,tplus,tminus,Tmax):\n",
    "    \n",
    "    l, v = sci.linalg.eig(np.array([[-1/tplus,0],[1/tplus,-1/tminus]])) \n",
    "    l = np.real(l).reshape([-1,1])\n",
    "    vm = sci.linalg.inv(v)\n",
    "    #x = v*np.exp(l*t)*np.array([1,0]).reshape([-1,1])                  \n",
    "    X = np.array([0,0]).reshape([-1,1])\n",
    "\n",
    "    t=np.array([0])\n",
    "    for i in np.arange(1,T_A.size):\n",
    "        if T_A[i]-T_A[i-1]>Tmax+10*dt:\n",
    "            t2  =  np.arange(int(Tmax/dt))[np.newaxis]*dt  \n",
    "            t3  =  - 10*dt + np.arange(10)[np.newaxis]*dt   \n",
    "            \n",
    "            t5 = np.append(t2,t3+Tmax+10*dt ,axis=1)\n",
    "            t2 = np.append(t2,t3+T_A[i] - T_A[i-1],axis=1)\n",
    "\n",
    "        else: \n",
    "            t2  =  np.arange(int((T_A[i]-T_A[i-1])/dt))[np.newaxis]*dt \n",
    "            t5 = t2\n",
    "\n",
    "        #print(i)    \n",
    "        x0 = X[:,[-1]]   \n",
    "        x0[0,:] = x0[0,:]+A[i-1] \n",
    "        vvmx=(v@np.diag((vm@x0).ravel()))\n",
    "        x = vvmx@np.exp(l@t2)\n",
    "        X = np.append(X,x,axis=1)\n",
    "        #t = np.append(t,T_A[i-1] +t5)   \n",
    "        t = np.append(t,t[-1] +t5)   \n",
    "    \n",
    "    return X, t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_fig9b(i0,ge_data, tplus=5, tminus=50, Nbootstraps=200, Dn=200, model_type='tm'):\n",
    "    dt=0.1\n",
    "    #tplus=5\n",
    "    #tminus=50\n",
    "    Tmax = 1000\n",
    "    DT = 200/6\n",
    "\n",
    "    #i0=23 #42 #[18,23,42]\n",
    "    ii = i0*Nbootstraps + np.arange(0,Nbootstraps,Dn)\n",
    "    #T_A = np.array((100+np.arange(8)*200/6).tolist()+[100+200/6*8+500]) #  from BBP cell paper figure\n",
    "    T_A = np.array((100+np.arange(9)*200/6).tolist()+[100+200/6*9+400]) # for classification task\n",
    "    \n",
    "    \n",
    "    if model_type=='tm':\n",
    "        stp_aba_names = ['tF', 'p0','tD','dp','A','A1','A2'] # tm5\n",
    "    if model_type=='tm_smr':\n",
    "        stp_aba_names = ['tF', 'p0','tD','dp','A','A1','A2','tDmin','dd','t_FDR','t_SMR','dp0'] # tm5+smr\n",
    "    \n",
    "    \n",
    "    \n",
    "    if model_type=='tm':\n",
    "        T_A0 = T_A-T_A[0]\n",
    "        As, ns, ps, dpp0, p0, tF, tD, A = STP_sim(ge_data.iloc[ii,:],T_A0)\n",
    "    if model_type=='tm_smr':\n",
    "        xs  =ge_data.iloc[ii,:].loc[:,stp_aba_names].values\n",
    "        xs  =np.delete(xs, [5,6],axis=1)\n",
    "\n",
    "        T_A0 = T_A-T_A[0]\n",
    "        As = np.zeros((xs.shape[0],T_A.shape[0]))\n",
    "        for i2 in range(xs.shape[0]):\n",
    "        \n",
    "            as2, sts2 = STP_sim2(xs[i2,:],T_A0,model_type='tm5_smr') \n",
    "            As[i2,0:as2.shape[0]] = as2\n",
    "\n",
    "\n",
    "    #A = np.mean(As,axis=0) #As[0,:Na1]#np.mean(As[:100,:Na1],axis=0)\n",
    "    #A = np.append([0],A)\n",
    "    ##T_A = Tall_A\n",
    "    \n",
    "    T_A =np.append([0],T_A)\n",
    "    T_A = np.append(T_A,[T_A[-1]+400])\n",
    "    \n",
    "    \n",
    "    Epsp =[]\n",
    "    for iii in range(As.shape[0]):\n",
    "        #A = As[iii,:]\n",
    "        A = As[iii,:]/As[iii,0]\n",
    "        A = np.append([0],A)\n",
    "        A = np.append(A,[0])\n",
    "\n",
    "        #print(T_A, A)\n",
    "        epsp, t = epsp_sim(A, T_A,dt,tplus,tminus,Tmax)\n",
    "        Epsp = Epsp + [epsp[1,:].reshape([-1,1])]\n",
    "     \n",
    "    \n",
    "    Epsp = np.concatenate(Epsp, axis=1)\n",
    "    \n",
    "    \n",
    "    \n",
    "    #plt.plot(t, epsp[1,:].T)\n",
    "    \n",
    "    return np.mean(Epsp,axis=1), t, As, T_A0, Epsp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Nbootstraps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Generate responses for each synapse type : 20Hz + 500ms recovery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simulate epsps for fig9a\n",
    "stp_columns2=np.array(stp_columns)\n",
    "annot_columns2 =np.array(annot_columns)\n",
    "sx=5\n",
    "sy=4\n",
    "i0=0\n",
    "skip_this=1\n",
    "eps = np.finfo(float).eps\n",
    "#Stp_type_p = []\n",
    "EPSP=[]\n",
    "Ams=[]\n",
    "Ams_s=[]\n",
    "Ams_m = []\n",
    "\n",
    "\n",
    "### ACHTUNG!!! ACHTUNG!!! ACHTUNG!!! ACHTUNG!!!\n",
    "#Nbootstraps = 100\n",
    "# ACHTUNG!!! ACHTUNG!!! ACHTUNG!!! ACHTUNG!!!\n",
    "\n",
    "\n",
    "for i0 in range(0,df_Stp_type_p.shape[0]):#[42]: # [18,23,42] #range(0,53):\n",
    "    if i0<53:\n",
    "        str_syntype = ge_data.loc[i0*Nbootstraps,\n",
    "                              annot_columns2[[0,2,1,3]]].astype(str).str.cat(others=['_','->','_','']).str.cat()\n",
    "    else:\n",
    "        str_syntype = ge_data_h.loc[(i0-53)*Nbootstraps,\n",
    "                              annot_columns2[[0,2,1,3]]].astype(str).str.cat(others=['_','->','_','']).str.cat()\n",
    "\n",
    "    #print(ge_data.loc[i0*200,annot_columns[:-3]])\n",
    "\n",
    "    #ii = i0*200 + np.arange(200)\n",
    "    \n",
    "    #epsp,t_epsp = plot_fig9b(i0,ge_data, tplus=2, tminus=20, Nbootstraps=200, Dn=200) # averaged all synapses\n",
    "    if i0<10:\n",
    "        model_type='tm'\n",
    "    elif i0<53:\n",
    "        model_type='tm_smr'\n",
    "    else:\n",
    "        model_type='tm'\n",
    "        \n",
    "    if i0<53:    \n",
    "        epsp,t_epsp,ams, T_A,epsp0 = plot_fig9b(i0,ge_data, tplus=3, tminus=25, Nbootstraps=Nbootstraps, Dn=5, model_type=model_type)  # only 1st response\n",
    "    else:\n",
    "        epsp,t_epsp,ams, T_A,epsp0 = plot_fig9b(i0-53,ge_data_h, tplus=3, tminus=25, Nbootstraps=Nbootstraps, Dn=5, model_type=model_type)  # only 1st response\n",
    "    \n",
    "    \n",
    "    EPSP = EPSP + [epsp0[:,0]]\n",
    "    Ams = Ams + [ams/ams[:,[0]]]\n",
    "    Ams_m = Ams_m + [np.mean(ams/ams[:,[0]],axis=0)]\n",
    "    \n",
    "    Ams_s = Ams_s + [np.std(ams/ams[:,[0]],axis=0)+10000*eps]\n",
    "    \n",
    "    print(str(i0)+ '  '+str_syntype + '  ' +'epsp')\n",
    "    \n",
    "    skip_this=1\n",
    "    if i0<53:\n",
    "        skip_this=0\n",
    "    if skip_this==0:\n",
    "        f, ax =plt.subplots(figsize=(sx, sy))\n",
    "\n",
    "        plt.plot(t_epsp,epsp0,color=[0.7,0.7,0.7])\n",
    "        plt.plot(t_epsp,epsp)\n",
    "        \n",
    "        plt.plot(t_epsp,epsp0[:,0],'g')\n",
    "\n",
    "        #dat.hist(bins=25)\n",
    "        stitle = str(i0)+ '  '+str_syntype + '  ' +'epsp '+np.array2string(df_Stp_type_p.loc[i0,\n",
    "                                ['polarity','tm_all','tm_DFUdU','dtype_1','dtype_2','dtype_3',\n",
    "                                 'type_1','type_2','type_3']].values)\n",
    "        \n",
    "        plt.title(stitle)\n",
    "        plt.xticks(fontsize=16)\n",
    "        plt.yticks(fontsize=16)\n",
    "        plt.xlabel('time, ms',fontsize=20)\n",
    "        plt.ylabel('EPSP, mV',fontsize=20)\n",
    "        \n",
    "        f, ax =plt.subplots(figsize=(sx, sy))\n",
    "        #plt.plot(T_A[1:-1],ams[0,:].transpose(),'og')\n",
    "        plt.plot(T_A,ams[0,:].transpose(),'og')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPSP2 = np.array(EPSP)\n",
    "\n",
    "Ams_m = np.array(Ams_m)\n",
    "Ams_s = np.array(Ams_s)\n",
    "print(Ams_m.shape)\n",
    "EPSP2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def  STP_sim2(x, T, init_state=None, model_type='tm5' ):\n",
    "\n",
    "    N    = len(T)\n",
    "    tF      = x[0] #.astype(float)\n",
    "    p00     = x[1]\n",
    "    tD      = x[2] #.astype(float)\n",
    "    dp      = x[3]\n",
    "    A       = 1 #x[4] # simplify A\n",
    "    \n",
    "    #breakpoint()\n",
    "    mod_fdr2=False\n",
    "    if model_type=='tm5_fdr2':  # should be :check freq. dependent recovery\n",
    "        tDmin     = x[5]\n",
    "        dd        = x[6]\n",
    "        t_FDR     = x[7]\n",
    "        mod_fdr2=True\n",
    "        tDmax  = tD\n",
    "        itDmin = 1/tDmin\n",
    "        itDmax = 1/tDmax\n",
    "        #breakpoint()\n",
    "        \n",
    "    mod_smr=False\n",
    "    if model_type=='tm5_smr':  # should be :check freq. dependent recovery\n",
    "        t_SMR   = x[8]\n",
    "        dp0     = x[9]\n",
    "        mod_smr=True\n",
    "        #p00  = p00\n",
    "\n",
    "    As = np.zeros((N,))\n",
    "    state = np.zeros((N,4))\n",
    "\n",
    "   \n",
    "    if init_state is None :\n",
    "        n = 1\n",
    "        p0=p00\n",
    "        p = p0\n",
    "        d = 0\n",
    "    else:\n",
    "        n = init_state[0]\n",
    "        p = init_state[1]\n",
    "        d = init_state[2]\n",
    "        p0= init_state[3]\n",
    "\n",
    "    \n",
    "    for i in range(0,N):\n",
    "        if i==0:\n",
    "            Dt = T[i]\n",
    "        else:\n",
    "            Dt = T[i]-T[i-1]\n",
    "        \n",
    "        if mod_fdr2:\n",
    "            d0=d\n",
    "            d = d*np.exp(-Dt/t_FDR) \n",
    "            n = 1 - (1 - n )*np.exp(-Dt*itDmax -(itDmin -itDmax)*t_FDR*(d0-d))\n",
    "        else:\n",
    "            #print(x,tD)\n",
    "            n = 1 - (1 - n )*np.exp(-Dt/tD )\n",
    "            \n",
    "        if mod_smr:\n",
    "            p01=p0\n",
    "            p0=p00 + (p0 -p00)*np.exp(-Dt/t_SMR)\n",
    "            p=p0 +(p -p01)*np.exp(-Dt/tF)\n",
    "        else:\n",
    "            p=p0 +(p -p0)*np.exp(-Dt/tF)\n",
    "            \n",
    "\n",
    "        As[i]=A*n*p\n",
    "       \n",
    "        n = n*(1-p)\n",
    "        p = p + dp*(1-p)\n",
    "        if mod_fdr2:\n",
    "            d  = d + dd*(1-d) \n",
    "        if mod_smr:\n",
    "            p0  = p0 - dp0*p0    \n",
    " \n",
    "        state[i] = [n,p,d,p0]\n",
    "\n",
    "    #return As, ns2, ps2, dpp0, p0, tF, tD, A\n",
    "    return As, state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Q_TM_simple(x, amps,sig,Ts,model_type='tm4'):\n",
    "    #Q, A, A3 =  QA_TM_aba_synphys(x, amps,sig,Ts)\n",
    "    Q, _ =  QA_TM_simple(x, amps,sig,Ts,model_type=model_type)\n",
    "    return Q  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def QA_TM_simple(x, amps,sig,T_A,model_type='tm4'):\n",
    "    eps = np.finfo(float).eps\n",
    "    \n",
    "    if model_type=='tm4':\n",
    "        #x = [tF, p0,tD,  dp]\n",
    "        x[3] = x[1] \n",
    "        amps2, sts2 = STP_sim2(x,T_A,model_type='tm5')\n",
    "    else:  \n",
    "        amps2, sts2 = STP_sim2(x,T_A,model_type=model_type)\n",
    "    \n",
    "    amps2 = amps2/amps2[0]\n",
    "    Q = np.sum((amps2-amps)**2/(sig**2 + eps))\n",
    "    return Q, amps2    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_STP_model(amps,sig,Ts,par,model_type = 'tm4'):\n",
    "\n",
    "    npar = amps.shape[1]\n",
    "\n",
    "\n",
    "         #  tF   p0     tD  dp         A            A1           A2      tDmin  dd   t_FDR t_SMR dp0 \n",
    "    x0 = [ 100,  0.1,  100, 0.1, amps[0]/0.1, amps[0]/0.1,  amps[0]/0.1, 10,   0.05, 100,  100,  0.02 ]\n",
    "    \n",
    "    x0 = np.array(x0)\n",
    "    \n",
    "    #dxx = np.array([1e3, 20, 1e3, 1e2,  20, 20, 20,  1e1, 20, 1e2 ])\n",
    "    dxx = np.array([1e3,  10, 1e3, 10,  20, 20, 20,  1e3, 20, 1e3, 1e3, 50 ])\n",
    "    \n",
    "    #breakpoint()\n",
    "    if model_type=='tm5':\n",
    "        x0 = x0[0:4]\n",
    "        dxx = dxx[0:4]\n",
    "\n",
    "        \n",
    "    if model_type=='tm4':\n",
    "        x0 = x0[0:4]\n",
    "        dxx = dxx[0:4]\n",
    "        \n",
    "\n",
    "    #print(x0)\n",
    "    Q0, A0 = QA_TM_simple(x0, amps, sig, Ts,model_type=model_type)\n",
    "    #print(x0)\n",
    "    print('initial Q: ',str(Q0))\n",
    "\n",
    "    #           tF      p0    tD        dp           A                A1                  A2              tDmin   dd   tD2\n",
    "    #xlower = [  0.1,   0.01,   0.1,     0.0001,     amps[0]/1*1e-1,   amps[0]/1*1e-1,     amps[0]/1*1e-1,  1,     1e-3, 1]\n",
    "    xlower = x0/dxx\n",
    "    xlower = np.array(xlower)\n",
    "    #xupper = [  1e5,   1,      1e5,     1,          amps[0]/0.01*1e1, amps[0]/0.01*1e1,   amps[0]/0.01*1e1, 100,  1,    1e4]\n",
    "    #xupper = np.array(xupper)\n",
    "    xupper = x0*dxx\n",
    "    bounds = [] #np.zeros(len(xlower))\n",
    "    for ibo in range(len(xlower)):\n",
    "        bounds = bounds + [(xlower[ibo],xupper[ibo])]\n",
    "    tt=[]\n",
    "    #import time\n",
    "    #import scipy.optimize as opt\n",
    "    t1 = time.time()\n",
    "\n",
    "    do_this=1\n",
    "    if do_this==1:\n",
    "        #res4 = opt.basinhopping(Q_TM_aba_synphys, x0 , niter=100, T=1.0, stepsize=1.0,\n",
    "        #                        minimizer_kwargs={'method':\"L-BFGS-B\", 'args':(amps, sig, Ts,DT0)}, \n",
    "        #                        take_step=None, accept_test=None, callback=None, interval=50, disp=False,\n",
    "        #                        niter_success=None, seed=None)\n",
    "        #nst=50 # good\n",
    "        nst = 100 #150 # better fit\n",
    "        Qs =  np.zeros(nst)\n",
    "        As =  np.zeros((nst,npar))\n",
    "        #vx = np.arange(5)  # fit 1 amplitude\n",
    "        vx = np.arange(len(x0)) # fit separate amplitude for each frequency\n",
    "        nx = len(x0)\n",
    "        xs =  np.zeros((nst,len(x0[vx]) ))\n",
    "        for ist in range(nst):\n",
    "\n",
    "            #x0i = np.exp(np.log(xlower) + np.random.rand(len(xlower))*(np.log(xupper)-np.log(xlower)))\n",
    "            x0i = x0*np.exp( 2*(np.random.rand(len(xlower))-0.5)*np.log(dxx)  )\n",
    "            x0i = x0i[vx]\n",
    "            #x0i[4:7] = x0i[4]\n",
    "            fl=1\n",
    "            #breakpoint()\n",
    "            resi=opt.minimize(Q_TM_simple, x0i, args=((amps, sig, Ts,model_type)), method=None, jac=None,\n",
    "                      hess=None, hessp=None, bounds=bounds[0:nx],\n",
    "                      constraints=(), tol=None, callback=None, \n",
    "                      options=None)\n",
    "            x=resi.x\n",
    "            if model_type=='tm4':\n",
    "                #x = [tF, p0,tD,  dp]\n",
    "                x[3] = x[1] \n",
    "            Q,A = QA_TM_simple(x, amps, sig, Ts,model_type=model_type)\n",
    "            As[ist,:] = A\n",
    "            xs[ist,:] = x\n",
    "            Qs[ist] = Q\n",
    "\n",
    "            \n",
    "            Q_rs = np.min(Qs)\n",
    "            ia = np.nonzero(Q_rs==Qs)[0]\n",
    "            x_rs=xs[ia,:].ravel()\n",
    "            A_rs = As[ia,:].ravel()\n",
    "    else: \n",
    "        x_rs=x0\n",
    "        Q_rs = 1e5\n",
    "        A_rs = np.zeros((npar,))\n",
    "            \n",
    "    t2=time.time()\n",
    "    tt = tt + [t2-t1]\n",
    "    print('random_start: elapsed time '+str(t2-t1)+'s ')\n",
    "    \n",
    "    return x_rs, Q_rs, A_rs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## refit from 4-parameters TM model (and SMR model) to 3 parameters TM, save new STP type to 'stp types probability cortex and hippocampus correct 33Hz.xlsx'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "import scipy.optimize as opt\n",
    "#from pygmo import \n",
    "#from pylab import *\n",
    "\n",
    "\n",
    "#  ACHTUNG!!! ACHTUNG!!! ACHTUNG!!! ACHTUNG!!!\n",
    "Ts = T_A\n",
    "#  ACHTUNG!!! ACHTUNG!!! ACHTUNG!!! ACHTUNG!!!\n",
    "\n",
    "t00 = time.time()\n",
    "Q3  = []\n",
    "TT  = []\n",
    "X   = []\n",
    "Amps2 = []\n",
    "Amps1 = []\n",
    "\n",
    "nalg = 1\n",
    "\n",
    "npar=EPSP2.shape[1]\n",
    "\n",
    "par = {'fited_pulses':True, 'amplitudes_selected_for_TM_training':[], 'aba_case':True, 'figure':0,\n",
    "       'model_type' : 'tm4' , 'ibs':0, 'x0': np.array([]) }\n",
    "\n",
    "par['model_type'] = 'tm4'  # extended TM model\n",
    "\n",
    "for i in range(EPSP2.shape[0]):# range(EPSP2.shape[0]): #range(0,len(df9.index)): #range(22,23): #range(len(df9.index)):\n",
    "    idx = i #df9.index[i]\n",
    "    print('\\n\\n'+str(i)+'  ')\n",
    "    #print('\\n\\n'+str(i)+'  '+df9.loc[idx,'comments']+'\\n')\n",
    "    #Ts  = df9.loc[idx,'comments1']\n",
    "    ##ampsa = df.loc[idx,'stp_mean_fit_bootstrap']\n",
    "    ##is_fited_pulses = (abs(ampsa[0])>1e-14)\n",
    "    \n",
    "    is_fited_pulses = True\n",
    "    if Ts.size!=0: #(len(Ts)!=0): #&((par['fited_pulses']==False)|is_fited_pulses):\n",
    "        fig, ax = plt.subplots(figsize=(10, 4))\n",
    "        par['figure']=fig\n",
    "    #try:\n",
    "        #par['fited_pulses']=True\n",
    "        #Q2, tt, x, algs, colors = fit_TM_cell(idx,df9,par)\n",
    "        amps = Ams_m[[i],:]\n",
    "        sig = Ams_s[[i],:] #*0 + Ams_m[[i],:]*0.1\n",
    "\n",
    "        i0=i\n",
    "        #stp_aba_names = ['tF', 'p0','tD','dp/p0']\n",
    "        #x0 = ge_data.iloc[i0*Nbootstraps,:].loc[stp_aba_names].values\n",
    "        #x0[3] = x0[3]*x0[1]\n",
    "        #Q0,amps0 = QA_TM_simple(x0, amps, sig, Ts,model_type='tm5') # ACHTUNG!!! - amps0 do not coincide with amps2???\n",
    "        #print('Q0 ',Q0,'amps0 ', amps0,'x0 ', x0)\n",
    "        \n",
    "        vv = np.arange(amps.shape[1]-1)\n",
    "        x,Q2,amps2 = fit_STP_model(amps[:,vv],sig[:,vv],Ts[vv],par,model_type = 'tm4')\n",
    "        \n",
    "        plt.plot(Ts,amps.ravel(),'-xb',markersize=12)\n",
    "        plt.plot(Ts[vv],amps2.ravel(),'-og')\n",
    "        #plt.plot(Ts,amps0.ravel(),'-dr')\n",
    "        \n",
    "        Amps2 = Amps2 + [amps2]\n",
    "        Amps1 = Amps1 + [amps]\n",
    "        \n",
    "        if i0<53:\n",
    "            str_syntype = ge_data.loc[i0*Nbootstraps,\n",
    "                              annot_columns2[[0,2,1,3]]].astype(str).str.cat(others=['_','->','_','']).str.cat()\n",
    "        else:\n",
    "            str_syntype = ge_data_h.loc[(i0-53)*Nbootstraps,\n",
    "                              annot_columns2[[0,2,1,3]]].astype(str).str.cat(others=['_','->','_','']).str.cat()\n",
    "       \n",
    "            \n",
    "        stitle = ' '+np.array2string(df_Stp_type_p.loc[i0,\n",
    "                                ['polarity','tm_all','tm_DFUdU','dtype_1','dtype_2','dtype_3',\n",
    "                                 'type_1','type_2','type_3']].values)\n",
    "        \n",
    "        df_Stp_type_p.loc[i0, 'tm_DFUdU'] = np.array2string(x[[2,0,1,3]])\n",
    "        tF_2_tD=x[0]/x[2]\n",
    "        df_Stp_type_p.loc[i0, ['dtype_2','dtype_3','dtype_1']]=np.array([(tF_2_tD<0.4),(tF_2_tD>0.4)&(tF_2_tD<2),(tF_2_tD>2)])\n",
    "    \n",
    "        stitle = str(i0)+ '  '+str_syntype + '  '+ np.array2string(x[0:4])+np.array2string(df_Stp_type_p.loc[i0,\n",
    "                                ['dtype_1','dtype_2','dtype_3']].values)\n",
    "        \n",
    "        plt.title(stitle)\n",
    "        \n",
    "        #plt.xticks(fontsize=16)\n",
    "        #plt.yticks(fontsize=16)\n",
    "        #plt.xlabel('time, ms',fontsize=20)\n",
    "        #plt.ylabel('EPSP, mV',fontsize=20)\n",
    "        plt.ylim([0, max([np.max(amps), np.max(amps2)])*1.2])\n",
    "        \n",
    "        nalg = 1 #len(tt)\n",
    "        tt=[]\n",
    "    #except:\n",
    "    else:\n",
    "        print(Ts)\n",
    "        Q2 = (10*np.zeros(nalg)).tolist()\n",
    "        tt = (0*np.zeros(nalg)).tolist()\n",
    "        x = (0*np.zeros(nalg)).tolist()\n",
    "        \n",
    "    Q3 = Q3 + [Q2]\n",
    "    TT = TT + [tt]\n",
    "    X = X + [x]\n",
    "    #df_Stp_type_p.loc[i0, 'tm_DFUdU'] = np.array2string(x[0,2,1,3])\n",
    "    #tF_2_tD=x[2]/x[0]\n",
    "    #df_Stp_type_p.loc[i0, ['dtype_2','dtype_3','dtype_1']]=np.array([(tF_2_tD<0.4),(tF_2_tD>0.4)&(tF_2_tD<2),(tF_2_tD>2)])\n",
    "    print(df_Stp_type_p.loc[i0, ['dtype_1','dtype_2','dtype_3']])\n",
    "    print('x tm4 ',np.array2string(x))\n",
    "    print(stitle)\n",
    "\n",
    "    \n",
    "Q3 = pd.DataFrame(Q3)  \n",
    "TT = pd.DataFrame(TT) \n",
    "X3 = pd.DataFrame(X)\n",
    "t10 = time.time()\n",
    "print('time : '+ str(t10-t00))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# # Save results\n",
    "\n",
    "df_Stp_type_p.to_excel('stp types probability cortex and hippocampus correct 33Hz.xlsx')\n",
    "\n",
    "# Q3.columns = ['random_starts'] #algs\n",
    "# algs2=algs\n",
    "# for ialg, alg in enumerate(algs2): \n",
    "#     algs2[ialg] = alg + ' , time'\n",
    "# TT.columns = algs2\n",
    "# QT = pd.concat([Q3,TT],axis=1)\n",
    "\n",
    "# X3 = pd.DataFrame(X)\n",
    "# #X3.columns = ['differential_evolution best x', 'dual_annealing best x', 'random_starts best x', 'shgo best x', 'basinhopping best x']\n",
    "# X3.columns = ['differential_evolution best x', 'dual_annealing best x', 'random_starts best x']\n",
    "\n",
    "\n",
    "# QTX_extended = pd.concat([df9,QT,X3],axis=1)\n",
    "\n",
    "# #QTX_extended.to_hdf('refit_aba_2019_A1_8_fited_pulses_TM_8_syntypes_results','data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "amps2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array2string(df_Stp_type_p.loc[i0,['polarity','type_1', 'type_2', 'type_3']].values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Amps1 = np.array(Amps1)\n",
    "Amps2 = np.array(Amps2)\n",
    "Amps2.shape"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# simulate epsps for fig9a\n",
    "stp_columns2=np.array(stp_columns)\n",
    "annot_columns2 =np.array(annot_columns)\n",
    "sx=5\n",
    "sy=4\n",
    "i0=0\n",
    "skip_this=1\n",
    "eps = np.finfo(float).eps\n",
    "#Stp_type_p = []\n",
    "EPSP=[]\n",
    "Ams=[]\n",
    "Ams_s=[]\n",
    "Ams_m = []\n",
    "\n",
    "\n",
    "### ACHTUNG!!! ACHTUNG!!! ACHTUNG!!! ACHTUNG!!!\n",
    "#Nbootstraps = 100\n",
    "# ACHTUNG!!! ACHTUNG!!! ACHTUNG!!! ACHTUNG!!!\n",
    "\n",
    "\n",
    "for i0 in range(0,df_Stp_type_p.shape[0]):#[42]: # [18,23,42] #range(0,53):\n",
    "    if i0<53:\n",
    "        str_syntype = ge_data.loc[i0*Nbootstraps,\n",
    "                              annot_columns2[[0,2,1,3]]].astype(str).str.cat(others=['_','->','_','']).str.cat()\n",
    "    else:\n",
    "        str_syntype = ge_data_h.loc[(i0-53)*Nbootstraps,\n",
    "                              annot_columns2[[0,2,1,3]]].astype(str).str.cat(others=['_','->','_','']).str.cat()\n",
    "\n",
    "    #print(ge_data.loc[i0*200,annot_columns[:-3]])\n",
    "\n",
    "    #ii = i0*200 + np.arange(200)\n",
    "    \n",
    "    #epsp,t_epsp = plot_fig9b(i0,ge_data, tplus=2, tminus=20, Nbootstraps=200, Dn=200) # averaged all synapses\n",
    "    if i0<10:\n",
    "        model_type='tm'\n",
    "    elif i0<53:\n",
    "        model_type='tm_smr'\n",
    "    else:\n",
    "        model_type='tm'\n",
    "        \n",
    "    if i0<53:    \n",
    "        epsp,t_epsp,ams, T_A,epsp0 = plot_fig9b(i0,ge_data, tplus=3, tminus=25, Nbootstraps=Nbootstraps, Dn=5, model_type=model_type)  # only 1st response\n",
    "    else:\n",
    "        epsp,t_epsp,ams, T_A,epsp0 = plot_fig9b(i0-53,ge_data_h, tplus=3, tminus=25, Nbootstraps=Nbootstraps, Dn=5, model_type=model_type)  # only 1st response\n",
    "    \n",
    "    \n",
    "    EPSP = EPSP + [epsp0[:,0]]\n",
    "    Ams = Ams + [ams/ams[:,[0]]]\n",
    "    Ams_m = Ams_m + [np.mean(ams/ams[:,[0]],axis=0)]\n",
    "    \n",
    "    Ams_s = Ams_s + [np.std(ams/ams[:,[0]],axis=0)+10000*eps]\n",
    "    \n",
    "    print(str(i0)+ '  '+str_syntype + '  ' +'epsp')\n",
    "    \n",
    "    skip_this=1\n",
    "    if i0<53:\n",
    "        skip_this=0\n",
    "    if skip_this==0:\n",
    "        f, ax =plt.subplots(figsize=(sx, sy))\n",
    "\n",
    "        plt.plot(t_epsp,epsp0,color=[0.7,0.7,0.7])\n",
    "        plt.plot(t_epsp,epsp)\n",
    "        \n",
    "        plt.plot(t_epsp,epsp0[:,0],'g')\n",
    "\n",
    "        #dat.hist(bins=25)\n",
    "        stitle = str(i0)+ '  '+str_syntype + '  ' +'epsp '+np.array2string(df_Stp_type_p.loc[i0,\n",
    "                                ['polarity','tm_all','tm_DFUdU','dtype_1','dtype_2','dtype_3',\n",
    "                                 'type_1','type_2','type_3']].values)\n",
    "        \n",
    "        plt.title(stitle)\n",
    "        plt.xticks(fontsize=16)\n",
    "        plt.yticks(fontsize=16)\n",
    "        plt.xlabel('time, ms',fontsize=20)\n",
    "        plt.ylabel('EPSP, mV',fontsize=20)\n",
    "        \n",
    "        f, ax =plt.subplots(figsize=(sx, sy))\n",
    "        #plt.plot(T_A[1:-1],ams[0,:].transpose(),'og')\n",
    "        plt.plot(T_A,ams[0,:].transpose(),'og')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Amps1 = np.squeeze(Amps1)\n",
    "Amps1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Reestimate STP type for each synapse type using STP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sx = 12\n",
    "sy=4\n",
    "\n",
    "do_draw_cortex = False # draw cortical or hippocampal synapses\n",
    "if do_draw_cortex:\n",
    "    remove = np.array( (53+np.arange(40)).tolist()+[4] )\n",
    "else:\n",
    "    remove = np.array( (0+np.arange(53)).tolist()+[] )\n",
    "    \n",
    "take = np.zeros(df_Stp_type_p.shape[0])==0\n",
    "take[remove]=False\n",
    "df_Stp_type_p2 = df_Stp_type_p.iloc[take,:]\n",
    "\n",
    "dt=0.1\n",
    "\n",
    "for pol in ['ex', 'inh']:\n",
    "    for stype in ['dtype_1', 'dtype_2', 'dtype_3']:\n",
    "        #df_pol_stype=df_Stp_type_p.loc[(df_Stp_type_p.loc[:,'polarity']==pol)&(df_Stp_type_p.loc[:,stype]>=0.7),:]\n",
    "        df_pol_stype=df_Stp_type_p2.loc[(df_Stp_type_p2.loc[:,'polarity']==pol)&(df_Stp_type_p2.loc[:,stype]==True),:]\n",
    "        if df_pol_stype.shape[0]>0:\n",
    "#             if do_draw_cortex:\n",
    "#                 idx = df_pol_stype.index[df_pol_stype.index[:]<53]\n",
    "#             else:\n",
    "#                 idx = df_pol_stype.index[df_pol_stype.index[:]>=53]\n",
    "            \n",
    "            idx = df_pol_stype.index\n",
    "        \n",
    "            amps2 = Amps2[idx,:].T # amplitudes 3_TM\n",
    "            amps1 = Amps1[idx,:].T # amplitudes 4-TM\n",
    "            \n",
    "            epsp = EPSP2[idx,:].T\n",
    "            am=np.max(epsp[0:int((100+15)/dt),:], axis=0)\n",
    "            epsp =epsp/np.tile(am,[epsp.shape[0],1])\n",
    "            \n",
    "            if (pol=='inh')&(stype=='dtype_2'):\n",
    "                take = np.zeros(epsp.shape[1])==0\n",
    "                min_amp=np.min(-epsp[np.arange(int(50/dt))+int((100+200/6*7)/dt),:],axis=0)\n",
    "                print(min_amp)\n",
    "                take[min_amp<-0.8] = False\n",
    "                epsp = epsp[:,take]\n",
    "                \n",
    "            if (pol=='ex')&(stype=='dtype_2'):\n",
    "                \n",
    "                take = np.zeros(epsp.shape[1])==0\n",
    "                min_amp=np.max(epsp[np.arange(int(50/dt))+int((100+200/6*7)/dt),:],axis=0)\n",
    "                print(min_amp)\n",
    "                take[min_amp>1.1]= False\n",
    "                epsp = epsp[:,take]    \n",
    "            \n",
    "            f, ax =plt.subplots(figsize=(sx, sy))\n",
    "\n",
    "            if pol=='ex':\n",
    "                print('ex', stype)\n",
    "                plt.plot(t_epsp,epsp,color='b')\n",
    "                plt.plot(t_epsp,np.median(epsp, axis=1),color='r')\n",
    "            else: \n",
    "                print('inh', stype)\n",
    "                plt.plot(t_epsp,-epsp,color='b')\n",
    "                plt.plot(t_epsp,-np.median(epsp, axis=1),color='r')\n",
    "            \n",
    "            #dat.hist(bins=25)\n",
    "            all_syn_names=df_pol_stype.loc[idx,'name'].str.cat(sep='; ')\n",
    "            \n",
    "            if pol=='ex':\n",
    "                npol='E'\n",
    "            else:\n",
    "                npol='I'\n",
    "            nstype = str.split(stype, '_')[1]        \n",
    "            plt.title(npol+nstype, fontsize=25)\n",
    "            plt.xticks(fontsize=16)\n",
    "            plt.yticks(fontsize=16)\n",
    "            plt.xlabel('time, ms',fontsize=20)\n",
    "            plt.ylabel('EPSP, mV',fontsize=20)\n",
    "            \n",
    "            print(npol+nstype + '  synaptic pair types: ' +all_syn_names+'\\n')\n",
    "            \n",
    "            \n",
    "            f, ax =plt.subplots(figsize=(sx, sy))\n",
    "            plt.plot(np.arange(amps1.shape[0]),amps1,color='b')\n",
    "            plt.plot(np.arange(amps1.shape[0]),np.median(amps1, axis=1),color='r')\n",
    "            \n",
    "            f, ax =plt.subplots(figsize=(sx, sy))\n",
    "            plt.plot(np.arange(amps2.shape[0]),amps2,color='b')\n",
    "            plt.plot(np.arange(amps2.shape[0]),np.median(amps2, axis=1),color='r')\n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examples of distributions of An/A1 and U,D,F parameters.  Cortex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "%matplotlib inline\n",
    "\n",
    "\n",
    "#\n",
    "#do_medians, do_log should be off!!!\n",
    "# Nsamples = 200 !!!\n",
    "\n",
    "\n",
    "## some distributions\n",
    "#stp_columns\n",
    "\n",
    "# f, ax =plt.subplots(figsize=(16, 7))\n",
    "# plt.plot(ge_data.iloc[3600:3800,:].loc[:,stp_columns].transpose(),'o-')\n",
    "# plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "#          rotation_mode=\"anchor\")\n",
    "\n",
    "#f, ax =plt.subplots(figsize=(16, 7))\n",
    "stp_columns2=np.array(stp_columns)\n",
    "annot_columns2 =np.array(annot_columns)\n",
    "sx=5\n",
    "sy=4\n",
    "i0=0\n",
    "\n",
    "for i0 in [42]: # [18,23,42] #range(0,53):\n",
    "    str_syntype = ge_data.loc[i0*200,annot_columns2[[0,2,1,3]]].astype(str).str.cat(others=['_','->','_','']).str.cat()\n",
    "\n",
    "    #print(ge_data.loc[i0*200,annot_columns[:-3]])\n",
    "\n",
    "    ii = i0*200 + np.arange(200)\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    ge_data.iloc[ii,:].loc[:,stp_columns2[0]].hist(bins=20)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +stp_columns2[0])\n",
    "    plt.xlim((0.25,1.25)) \n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.xlabel('A2/A1',fontsize=20)\n",
    "    plt.ylabel('Frequency',fontsize=20)\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    ge_data.iloc[ii,:].loc[:,stp_columns2[3]].hist(bins=20)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +stp_columns2[3])\n",
    "    plt.xlim((0.25,1.25))\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.xlabel('A5/A1',fontsize=20)\n",
    "    plt.ylabel('Frequency',fontsize=20)\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    ge_data.iloc[ii,:].loc[:,stp_columns2[4]].hist(bins=20)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +stp_columns2[4])\n",
    "    plt.xlim((0.25,1.25))\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.xlabel('Arecovery/A1, ',fontsize=20)\n",
    "    plt.ylabel('Frequency',fontsize=20)\n",
    "\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    dat=ge_data.iloc[ii,:].loc[:,'tD'].map(float)\n",
    "    dat=dat.loc[dat.loc[:]<800]\n",
    "    dat.hist(bins=25)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +'D')\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.xlabel('D, ms',fontsize=20)\n",
    "    plt.ylabel('Frequency',fontsize=20)\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    #ge_data.iloc[ii,:].loc[:,'tF'].map(float).hist(bins=25)\n",
    "    dat=ge_data.iloc[ii,:].loc[:,'tF'].map(float)\n",
    "    dat=dat.loc[dat.loc[:]<5500]\n",
    "    dat.hist(bins=25)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +'F')\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.xlabel('F, ms',fontsize=20)\n",
    "    plt.ylabel('Frequency',fontsize=20)\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    #ge_data.iloc[ii,:].loc[:,'p0'].map(float).hist(bins=25)\n",
    "    dat=ge_data.iloc[ii,:].loc[:,'p0'].map(float)\n",
    "    dat=dat.loc[dat.loc[:]<1]\n",
    "    dat.hist(bins=25)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +'U')\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.xlabel('U',fontsize=20)\n",
    "    plt.ylabel('Frequency',fontsize=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "#\n",
    "#do_medians, do_log should be off!!!\n",
    "#\n",
    "\n",
    "\n",
    "## some distributions\n",
    "#stp_columns\n",
    "\n",
    "# f, ax =plt.subplots(figsize=(16, 7))\n",
    "# plt.plot(ge_data.iloc[3600:3800,:].loc[:,stp_columns].transpose(),'o-')\n",
    "# plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "#          rotation_mode=\"anchor\")\n",
    "\n",
    "#f, ax =plt.subplots(figsize=(16, 7))\n",
    "stp_columns2=np.array(stp_columns)\n",
    "annot_columns2 =np.array(annot_columns)\n",
    "sx=5\n",
    "sy=4\n",
    "i0=0\n",
    "\n",
    "for i0 in [42]: # [18,23,42] #range(0,53):\n",
    "    str_syntype = ge_data.loc[i0*200,annot_columns2[[0,2,1,3]]].astype(str).str.cat(others=['_','->','_','']).str.cat()\n",
    "\n",
    "    #print(ge_data.loc[i0*200,annot_columns[:-3]])\n",
    "\n",
    "    ii = i0*200 + np.arange(200)\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    ge_data.iloc[ii,:].loc[:,stp_columns2[0]].hist(bins=20)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +stp_columns2[0])\n",
    "    plt.xlim((0.25,1.25))\n",
    "    plt.xlabel(stp_columns2[0],fontsize=16)\n",
    "    plt.ylabel('Frequency',fontsize=16)\n",
    "    plt.setp(ax.get_xticklabels(),fontsize=14)\n",
    "    plt.setp(ax.get_yticklabels(),fontsize=14)\n",
    "    ax.grid(b=False)\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    ge_data.iloc[ii,:].loc[:,stp_columns2[3]].hist(bins=20)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +stp_columns2[3])\n",
    "    plt.xlim((0.25,1.25))\n",
    "    plt.xlabel(stp_columns2[3],fontsize=16)\n",
    "    plt.ylabel('Frequency',fontsize=16)\n",
    "    plt.setp(ax.get_xticklabels(),fontsize=14)\n",
    "    plt.setp(ax.get_yticklabels(),fontsize=14)\n",
    "    ax.grid(b=False)\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    ge_data.iloc[ii,:].loc[:,stp_columns2[4]].hist(bins=20)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +stp_columns2[4])\n",
    "    plt.xlim((0.25,1.25))\n",
    "    plt.xlabel(stp_columns2[4],fontsize=16)\n",
    "    plt.ylabel('Frequency',fontsize=16)\n",
    "    plt.setp(ax.get_xticklabels(),fontsize=14)\n",
    "    plt.setp(ax.get_yticklabels(),fontsize=14)\n",
    "    ax.grid(b=False)\n",
    "    \n",
    "\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    dat=ge_data.iloc[ii,:].loc[:,'tD'].map(float)\n",
    "    dat=dat.loc[dat.loc[:]<800]\n",
    "    dat.hist(bins=25)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +'tD')\n",
    "    plt.xlabel('D',fontsize=16)\n",
    "    plt.ylabel('Frequency',fontsize=16)\n",
    "    plt.setp(ax.get_xticklabels(),fontsize=14)\n",
    "    plt.setp(ax.get_yticklabels(),fontsize=14)\n",
    "    ax.grid(b=False)\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    #ge_data.iloc[ii,:].loc[:,'tF'].map(float).hist(bins=25)\n",
    "    dat=ge_data.iloc[ii,:].loc[:,'tF'].map(float)\n",
    "    dat=dat.loc[dat.loc[:]<5500]\n",
    "    dat.hist(bins=25)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +'tF')\n",
    "    plt.xlabel('F',fontsize=16)\n",
    "    plt.ylabel('Frequency',fontsize=16)\n",
    "    plt.setp(ax.get_xticklabels(),fontsize=14)\n",
    "    plt.setp(ax.get_yticklabels(),fontsize=14)\n",
    "    ax.grid(b=False)\n",
    "\n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    #ge_data.iloc[ii,:].loc[:,'p0'].map(float).hist(bins=25)\n",
    "    dat=ge_data.iloc[ii,:].loc[:,'p0'].map(float)\n",
    "    dat=dat.loc[dat.loc[:]<1]\n",
    "    dat.hist(bins=25)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +'p0')\n",
    "    plt.xlabel('U',fontsize=16)\n",
    "    plt.ylabel('Frequency',fontsize=16)\n",
    "    plt.setp(ax.get_xticklabels(),fontsize=14)\n",
    "    plt.setp(ax.get_yticklabels(),fontsize=14)\n",
    "    ax.grid(b=False)\n",
    "    \n",
    "    f, ax =plt.subplots(figsize=(sx, sy))\n",
    "    ge_data.iloc[ii,:].loc[:,'p0'].map(float).hist(bins=25)\n",
    "    dat=ge_data.iloc[ii,:].loc[:,'p0'].map(float)\n",
    "    dat=dat.loc[dat.loc[:]<1]\n",
    "    dat.hist(bins=25)\n",
    "    #plt.title(str(i0)+ '  '+str_syntype + '  ' +'p0'\n",
    "    plt.xlabel('U',fontsize=16)\n",
    "    plt.ylabel('Frequency',fontsize=16)\n",
    "    plt.setp(ax.get_xticklabels(),fontsize=14)\n",
    "    plt.setp(ax.get_yticklabels(),fontsize=14)\n",
    "    plt.grid(b=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# library & dataset\n",
    "import seaborn as sns\n",
    "#df = sns.load_dataset('iris')\n",
    " \n",
    "# Basic 2D density plot\n",
    "sns.set_style(\"white\")\n",
    "#sns.kdeplot(df.sepal_width, df.sepal_length)\n",
    "#sns.plt.show()\n",
    " \n",
    "# Custom it with the same argument as 1D density plot\n",
    "dat=ge_data.iloc[ii,:].loc[:,'tD'].map(float)\n",
    "dat2=ge_data.iloc[ii,:].loc[:,'tF'].map(float)\n",
    "dat3=ge_data.iloc[ii,:].loc[:,'p0'].map(float)\n",
    "sns.kdeplot(np.log10(dat), np.log10(dat2), cmap=\"Reds\", shade=True, bw=.15) #, log_scale=True) #, log_scale=True\n",
    "\n",
    "plt.xlabel('Log(D)',fontsize=16)\n",
    "plt.ylabel('Log(F)',fontsize=16)\n",
    "plt.setp(plt.gca().get_xticklabels(),fontsize=14)\n",
    "plt.setp(plt.gca().get_yticklabels(),fontsize=14)\n",
    " \n",
    "# Some features are characteristic of 2D: color palette and wether or not color the lowest range\n",
    "#sns.kdeplot(df.sepal_width, df.sepal_length, cmap=\"Blues\", shade=True, shade_lowest=True, )\n",
    "##sns.plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Basic 2D density plot\n",
    "sns.set_style(\"white\")\n",
    "#sns.kdeplot(df.sepal_width, df.sepal_length)\n",
    "#sns.plt.show()\n",
    " \n",
    "# Custom it with the same argument as 1D density plot\n",
    "dat=ge_data.iloc[ii,:].loc[:,stp_columns2[0]].map(float)\n",
    "dat2=ge_data.iloc[ii,:].loc[:,stp_columns2[3]].map(float)\n",
    "dat3=ge_data.iloc[ii,:].loc[:,stp_columns2[4]].map(float)\n",
    "sns.kdeplot((dat), (dat2), cmap=\"Reds\", shade=True, bw=.05)\n",
    "plt.xlabel('A2:A1 20Hz',fontsize=16)\n",
    "plt.ylabel('A5:A1 20Hz',fontsize=16)\n",
    "plt.setp(plt.gca().get_xticklabels(),fontsize=14)\n",
    "plt.setp(plt.gca().get_yticklabels(),fontsize=14)\n",
    "#ge_data.iloc[ii,:].loc[:,stp_columns2[0]].hist(bins=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(__doc__)\n",
    "\n",
    "\n",
    "# Code source: Gaël Varoquaux\n",
    "#              Andreas Müller\n",
    "# Modified for documentation by Jaques Grobler\n",
    "# License: BSD 3 clause\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import ListedColormap\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.datasets import make_moons, make_circles, make_classification\n",
    "from sklearn.neural_network import MLPClassifier\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from sklearn.svm import SVC\n",
    "from sklearn.gaussian_process import GaussianProcessClassifier\n",
    "from sklearn.gaussian_process.kernels import RBF\n",
    "from sklearn.tree import DecisionTreeClassifier\n",
    "from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier\n",
    "from sklearn.naive_bayes import GaussianNB\n",
    "from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis\n",
    "\n",
    "h = .02  # step size in the mesh\n",
    "\n",
    "names = [\"Nearest Neighbors\", \"Linear SVM\", \"RBF SVM\", \"Gaussian Process\",\n",
    "         \"Decision Tree\", \"Random Forest\", \"Neural Net\", \"AdaBoost\",\n",
    "         \"Naive Bayes\", \"QDA\"]\n",
    "\n",
    "classifiers = [\n",
    "    KNeighborsClassifier(3),\n",
    "    SVC(kernel=\"linear\", C=0.025),\n",
    "    SVC(gamma=2, C=1),\n",
    "    GaussianProcessClassifier(1.0 * RBF(1.0)),\n",
    "    DecisionTreeClassifier(max_depth=5),\n",
    "    RandomForestClassifier(max_depth=5, n_estimators=10, max_features=1),\n",
    "    MLPClassifier(alpha=1, max_iter=1000),\n",
    "    AdaBoostClassifier(),\n",
    "    GaussianNB(),\n",
    "    QuadraticDiscriminantAnalysis()]\n",
    "\n",
    "classifiers = [\n",
    "    GaussianProcessClassifier(1.0 * RBF(1.0)),\n",
    "]\n",
    "\n",
    "X, y = make_classification(n_features=2, n_redundant=0, n_informative=2,\n",
    "                           random_state=1, n_clusters_per_class=1)\n",
    "rng = np.random.RandomState(2)\n",
    "X += 2 * rng.uniform(size=X.shape)\n",
    "linearly_separable = (X, y)\n",
    "\n",
    "datasets = [make_moons(noise=0.3, random_state=0),\n",
    "            make_circles(noise=0.2, factor=0.5, random_state=1),\n",
    "            linearly_separable\n",
    "            ]\n",
    "\n",
    "figure = plt.figure(figsize=(9, 9))\n",
    "i = 1\n",
    "# iterate over datasets\n",
    "for ds_cnt, ds in enumerate(datasets):\n",
    "    # preprocess dataset, split into training and test part\n",
    "    X, y = ds\n",
    "    X = StandardScaler().fit_transform(X)\n",
    "    X_train, X_test, y_train, y_test = \\\n",
    "        train_test_split(X, y, test_size=.4, random_state=42)\n",
    "\n",
    "    x_min, x_max = X[:, 0].min() - .5, X[:, 0].max() + .5\n",
    "    y_min, y_max = X[:, 1].min() - .5, X[:, 1].max() + .5\n",
    "    xx, yy = np.meshgrid(np.arange(x_min, x_max, h),\n",
    "                         np.arange(y_min, y_max, h))\n",
    "\n",
    "    # just plot the dataset first\n",
    "    cm = plt.cm.RdBu\n",
    "    cm_bright = ListedColormap(['#FF0000', '#0000FF'])\n",
    "    ax = plt.subplot(len(datasets), len(classifiers) + 1, i)\n",
    "    if ds_cnt == 0:\n",
    "        ax.set_title(\"Input data\")\n",
    "    # Plot the training points\n",
    "    ax.scatter(X_train[:, 0], X_train[:, 1], c=y_train, cmap=cm_bright,\n",
    "               edgecolors='k')\n",
    "    # Plot the testing points\n",
    "    ax.scatter(X_test[:, 0], X_test[:, 1], c=y_test, cmap=cm_bright, alpha=0.6,\n",
    "               edgecolors='k')\n",
    "    ax.set_xlim(xx.min(), xx.max())\n",
    "    ax.set_ylim(yy.min(), yy.max())\n",
    "    ax.set_xticks(())\n",
    "    ax.set_yticks(())\n",
    "    i += 1\n",
    "\n",
    "    # iterate over classifiers\n",
    "    for name, clf in zip(names, classifiers):\n",
    "        ax = plt.subplot(len(datasets), len(classifiers) + 1, i)\n",
    "        clf.fit(X_train, y_train)\n",
    "        score = clf.score(X_test, y_test)\n",
    "\n",
    "        # Plot the decision boundary. For that, we will assign a color to each\n",
    "        # point in the mesh [x_min, x_max]x[y_min, y_max].\n",
    "        if hasattr(clf, \"decision_function\"):\n",
    "            Z = clf.decision_function(np.c_[xx.ravel(), yy.ravel()])\n",
    "        else:\n",
    "            Z = clf.predict_proba(np.c_[xx.ravel(), yy.ravel()])[:, 1]\n",
    "\n",
    "        # Put the result into a color plot\n",
    "        Z = Z.reshape(xx.shape)\n",
    "        ax.contourf(xx, yy, Z, cmap=cm, alpha=.8)\n",
    "\n",
    "        # Plot the training points\n",
    "        ax.scatter(X_train[:, 0], X_train[:, 1], c=y_train, cmap=cm_bright,\n",
    "                   edgecolors='k')\n",
    "        # Plot the testing points\n",
    "        ax.scatter(X_test[:, 0], X_test[:, 1], c=y_test, cmap=cm_bright,\n",
    "                   edgecolors='k', alpha=0.6)\n",
    "\n",
    "        ax.set_xlim(xx.min(), xx.max())\n",
    "        ax.set_ylim(yy.min(), yy.max())\n",
    "        ax.set_xticks(())\n",
    "        ax.set_yticks(())\n",
    "        if ds_cnt == 0:\n",
    "            ax.set_title(name)\n",
    "        ax.text(xx.max() - .3, yy.min() + .3, ('%.2f' % score).lstrip('0'),\n",
    "                size=15, horizontalalignment='right')\n",
    "        i += 1\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Combine Cortex and Hipp. datasets, assign synapse type annotations, save to .hdf file "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ge_data_h.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ge_data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_c.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_h.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set(ge_data.columns).difference(set(ge_data_h.columns))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##\n",
    "##\n",
    "##  combine cortex and hippocampus datasets \n",
    "##\n",
    "##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ge_data_h = ge_data\n",
    "# X_h = X\n",
    "# y_h = y\n",
    "# annot_columns_h = annot_columns\n",
    "# ge_columns_h = ge_columns\n",
    "# stp_columns_h = stp_colum\n",
    "\n",
    "X = np.concatenate([X_c, X_h],axis=0)\n",
    "y = np.concatenate([y_c, y_h],axis=0)\n",
    "\n",
    "print(X_h.shape)\n",
    "print(X_c.shape)\n",
    "X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Xan = pd.DataFrame(X[0::200,0:6])\n",
    "Xan.to_excel('all_annotations.xlsx') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Xan2 = pd.DataFrame([[16+3,8+4,8+1,11+2],[9,7,6,12]],columns = ['In->In','In->Ex','Ex->In','In->In'],\n",
    "                    index=['VISp','Hippocampus']).T\n",
    "\n",
    "#import matplotlib.pyplot as plt\n",
    "f, ax =plt.subplots(figsize=(10, 5))\n",
    "#    plt.title(str(i0)+ '  '+str_syntype + '  ' +stp_columns2[3])\n",
    "#    plt.xlim((0.25,1.25))\n",
    "\n",
    "Xan2.plot(kind = 'bar',ax=ax)\n",
    "plt.ylim([0,20])\n",
    "plt.xticks(fontsize=18)\n",
    "plt.yticks(ticks=[0,5,10,15,20],fontsize=18)\n",
    "plt.ylabel('Number of synapse types',fontsize=20)\n",
    "#plt.xlabel('Synapse type',fontsize=20)\n",
    "plt.legend(fontsize=20)\n",
    "#print(set(n_stp_cells.index).difference(set(i2)))\n",
    "#n_stp_cells4.sum()\n",
    "\n",
    "\n",
    "\n",
    "Xan2.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#pd.DataFrame(X).to_excel('temp.xlsx')"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "all_types = []\n",
    "for i in ge_data_h.index:\n",
    "    all_type=''\n",
    "    for j in range(6):\n",
    "        all_type = all_type + str(ge_data_h.loc[i,ge_data_h.columns[j+1]])+' '\n",
    "    all_types = all_types  +[all_type]    \n",
    "\n",
    "all_types =  pd.DataFrame(all_types)\n",
    "subclasses_hipp1 = pd.DataFrame(pd.unique(all_types.loc[:,0].values))\n",
    "#subclasses =  pd.DataFrame(subclasses,columns=['pre', 'post'] )\n",
    "#subclasses_hipp = pd.concat([subclasses_hipp1, subclasses], axis=1)\n",
    "subclasses2_hipp = ['edg','eca','ee','Pvalb','Sst','Vip','Lamp5','Cck']\n",
    "subclasses3_hipp = [['DG'],\n",
    "                    ['CA1','CA2','CA3'],\n",
    "                    ['EC'],\n",
    "                    ['PV','Pvalb','PVBC','AAC'],\n",
    "                    ['Sst','HIPP','O-LM','O-Bi'],\n",
    "                    ['Vip','IS3'],\n",
    "                    ['Lamp5','Ivy cell','NGF'],\n",
    "                    ['Cck','CCKBC','HICAP','CCK_DTI']]\n",
    "subclasses_hipp1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'cell_type2_post', 'layer_post'   , 'cell_type2_pre', 'layer_pre'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "t1 = time.time()\n",
    "subclasses2_hipp = ['e_dg','e_ca','e_e','Pvalb_h','Sst_h','Vip','Lamp5_h','Cck_h']\n",
    "subclasses3_hipp = [['DG','CA3','CA3 ventral'],\n",
    "                    ['CA1','CA1 ventral'],\n",
    "                    ['EC'],\n",
    "                    ['PV','Pvalb','PVBC','AAC'],\n",
    "                    ['Sst','HIPP','O-LM','O-Bi'],\n",
    "                    ['Vip','IS3'],\n",
    "                    ['Lamp5','Ivy cell','Ivy cell type2','NGF','NGF type2'],\n",
    "                    ['Cck','CCKBC','HICAP','CCK_DTI']]\n",
    "\n",
    "\n",
    "ge_data_h_subcl = pd.DataFrame(np.zeros((ge_data_h.shape[0],4)),columns = ['name_long_pre','name_long_post',\n",
    "                                                                           'subclass_pre','subclass_post'])\n",
    "ge_data_h_subcl.index = ge_data_h.index\n",
    "precol = ['cell_type2_pre', 'layer_pre'] #ge_data_h.columns[[1,5]]\n",
    "postcol = ['cell_type2_post', 'layer_post' ] # ge_data_h.columns[[2,6]]\n",
    "\n",
    "\n",
    "all_type_1 = ge_data_h.loc[:,'cell_type2_pre']+ ' ' + ge_data_h.loc[:,'cre_line_pre'] +' '+ge_data_h.loc[:,'layer_pre']\n",
    "all_type_2 = ge_data_h.loc[:,'cell_type2_post']+ ' ' + ge_data_h.loc[:,'cre_line_post'] +' '+ge_data_h.loc[:,'layer_post']\n",
    "\n",
    "ge_data_h_subcl.loc[:,'name_long_pre'] = all_type_1\n",
    "ge_data_h_subcl.loc[:,'name_long_post'] = all_type_2\n",
    "\n",
    "for j in range(len(subclasses3_hipp)): \n",
    "    is_in_pre = ge_data_h.loc[:,precol[0]].isin(subclasses3_hipp[j])\n",
    "    is_in_post = ge_data_h.loc[:,postcol[0]].isin(subclasses3_hipp[j])\n",
    "    ge_data_h_subcl.loc[is_in_pre,'subclass_pre'] =subclasses2_hipp[j]\n",
    "    ge_data_h_subcl.loc[is_in_post,'subclass_post'] =subclasses2_hipp[j] \n",
    "    \n",
    "t2=time.time()\n",
    "print('elapsed time : '+str(t2-t1))\n",
    "ge_data_h_subcl    "
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "import time\n",
    "t1 = time.time()\n",
    "subclasses2_hipp = ['e_dg','e_ca','e_e','Pvalb_h','Sst_h','Vip','Lamp5_h','Cck_h']\n",
    "subclasses3_hipp = [['DG','CA3','CA3 ventral'],\n",
    "                    ['CA1','CA1 ventral'],\n",
    "                    ['EC'],\n",
    "                    ['PV','Pvalb','PVBC','AAC'],\n",
    "                    ['Sst','HIPP','O-LM','O-Bi'],\n",
    "                    ['Vip','IS3'],\n",
    "                    ['Lamp5','Ivy cell','Ivy cell type2','NGF','NGF type2'],\n",
    "                    ['Cck','CCKBC','HICAP','CCK_DTI']]\n",
    "\n",
    "\n",
    "ge_data_h_subcl = pd.DataFrame(np.zeros((ge_data_h.shape[0],4)),columns = ['name_long_pre','name_long_post',\n",
    "                                                                           'subclass_pre','subclass_post'])\n",
    "ge_data_h_subcl.index = ge_data_h.index\n",
    "precol = ['cell_type2_pre', 'layer_pre'] #ge_data_h.columns[[1,5]]\n",
    "postcol = ['cell_type2_post', 'layer_post' ] # ge_data_h.columns[[2,6]]\n",
    "\n",
    "all_type_1\n",
    "for i in ge_data_h.index:\n",
    "    all_type_1=''\n",
    "    all_type_2=''\n",
    "    for j in range(3):\n",
    "        all_type_1 = all_type_1+ str(ge_data_h.loc[i,ge_data_h.columns[2*j+1]])+' '\n",
    "    for j in range(3):\n",
    "        all_type_2 = all_type_2+ str(ge_data_h.loc[i,ge_data_h.columns[2*j+2]])+' '  \n",
    "    ge_data_h_subcl.loc[i,'name_long_pre'] = all_type_1\n",
    "    ge_data_h_subcl.loc[i,'name_long_post'] = all_type_2\n",
    "    \n",
    "    ge_data_hi = ge_data_h.loc[i,:]\n",
    "    for j in range(len(subclasses3_hipp)): \n",
    "        is_in_pre = ge_data_hi.loc[precol].isin(subclasses3_hipp[j]).sum()>0\n",
    "        is_in_post = ge_data_hi.loc[postcol].isin(subclasses3_hipp[j]).sum()>0\n",
    "        if is_in_pre:\n",
    "            ge_data_h_subcl.loc[i,'subclass_pre'] =subclasses2_hipp[j]\n",
    "        if is_in_post:\n",
    "            ge_data_h_subcl.loc[i,'subclass_post'] =subclasses2_hipp[j]    \n",
    "        \n",
    "    #all_types = all_types  +[all_type]    \n",
    "\n",
    "    \n",
    "    \n",
    "# all_types =  pd.DataFrame(all_types)   \n",
    "# subclasses = [['Pvalb', 'eh'],\n",
    "#               ['Pvalb', 'Pvalb'],\n",
    "#               ['eh','Cck'],\n",
    "#               ['ee','Lamp5'],\n",
    "#               ['eh','eh'],\n",
    "#               ['ee','Lamp5'],\n",
    "#               ['Cck','eh'],\n",
    "#               ['Pvalb','eh'],\n",
    "#               ['Lamp5', 'Lamp5'],\n",
    "#               ['eh','Pvalb'],\n",
    "#               ['eh','eh'],\n",
    "#               ['eh','eh'],\n",
    "#               ['Lamp5','eh'],\n",
    "#               ['Cck','eh'],\n",
    "#               ['eh','Pvalb'],\n",
    "#               ['Lamp5','eh'],\n",
    "#               ['Pvalb','Pvalb'],\n",
    "#               ['ee','Lamp5'],\n",
    "#               ['ee','Lamp5'],\n",
    "#               ['eh','eh'],\n",
    "#               ['eh','Cck'],\n",
    "#               ['ee','Cck'],\n",
    "#               ['ee','eh'],\n",
    "#               ['eh','Sst'],\n",
    "#               ['Vip','Sst'],\n",
    "#               ['ee','eh'],\n",
    "#               ['ee','Pvalb'],\n",
    "#               ['eh','Lamp5'],\n",
    "#               ['eh','Lamp5'],\n",
    "#               ['Cck','Cck'],\n",
    "#               ['Pvalb','eh'],\n",
    "#               ['eh', 'Sst']]\n",
    "\n",
    "# #subclasses_hipp1 = pd.DataFrame(list(set(all_types)) )\n",
    "# subclasses_hipp1 = pd.DataFrame(pd.unique(all_types.loc[:,0].values))\n",
    "# subclasses =  pd.DataFrame(subclasses,columns=['pre', 'post'] )\n",
    "# subclasses_hipp = pd.concat([subclasses_hipp1, subclasses], axis=1)\n",
    "\n",
    "\n",
    "# ge_data_h_subcl = pd.DataFrame(np.tile(all_types,(1,4)),columns=['name_long_pre','name_long_post',\n",
    "#                                                                  'subclass_pre','subclass_post'])\n",
    "# for i in range(subclasses_hipp.shape[0]):\n",
    "\n",
    "#     ge_data_h_subcl.loc[ge_data_h_subcl.loc[:,'name_long'].isin([subclasses_hipp.loc[i,0]]),'subclass_pre'] =subclasses_hipp.loc[i,'pre']\n",
    "#     ge_data_h_subcl.loc[ge_data_h_subcl.loc[:,'name_long'].isin([subclasses_hipp.loc[i,0]]),'subclass_post'] =subclasses_hipp.loc[i,'post']\n",
    "t2=time.time()\n",
    "print('elapsed time : '+str(t2-t1))\n",
    "ge_data_h_subcl    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_types = []\n",
    "all_types2 = []\n",
    "columns = ['cell_type2_pre','cell_type2_post','layer_pre','layer_post','cre_line_pre','cre_line_post']\n",
    "for i in ge_data.index:\n",
    "    all_type=''\n",
    "    all_type2=''\n",
    "    for j in range(3):\n",
    "        sj = str(ge_data.loc[i,columns[2*j]])\n",
    "        if sj=='MC':\n",
    "            sj='Sst'    \n",
    "        all_type = all_type  +' '+sj\n",
    "        \n",
    "    for j in range(3):\n",
    "        sj = str(ge_data.loc[i,columns[2*j+1]])\n",
    "        if sj=='MC':\n",
    "            sj='Sst'    \n",
    "        all_type2 = all_type2 +' '+sj    \n",
    "    all_types = all_types  +[str.lower(all_type)]  \n",
    "    all_types2 = all_types2  +[str.lower(all_type2)]   \n",
    "\n",
    "all_types =  pd.DataFrame(all_types, columns=['pre'])\n",
    "all_types2 =  pd.DataFrame(all_types2, columns=['post'])\n",
    "all_types = pd.concat([all_types,all_types2], axis=1)\n",
    "\n",
    "subclasses = pd.Series([ 'Sst','Pvalb','Vip','ev_l23','ev_l4','ev_l56','ev_l56'])\n",
    "subclasses2 = pd.Series(['sst','pvalb','vip','l2',    'l4',   'l5',    'l6'])\n",
    "subclasses2_visp =subclasses2   \n",
    "\n",
    "subclasses_visp_1 = pd.DataFrame(pd.unique(all_types.loc[:,'pre'].values))\n",
    "subclasses_visp_2 = pd.DataFrame(pd.unique(all_types.loc[:,'post'].values))\n",
    "# subclasses =  pd.DataFrame(subclasses,columns=['pre', 'post'] )\n",
    "# subclasses_visp = pd.concat([subclasses_visp1, subclasses], axis=1)\n",
    "\n",
    "ge_data_c_subcl = pd.DataFrame(np.tile(all_types,(1,2)))\n",
    "ge_data_c_subcl.columns=['name_long_pre','name_long_post','subclass_pre','subclass_post']\n",
    "for i in range(subclasses_visp_1.shape[0]):\n",
    "    type_name = str.lower(subclasses_visp_1.loc[i,0])\n",
    "#     is_in=[]\n",
    "#     for si in subclasses2:\n",
    "#         is_in = is_in +[ss.contains(si)]\n",
    "#     is_in=np.nonzero(np.array(is_in))[0][0] \n",
    "#     subcl_name = subclasses[is_in]\n",
    "    \n",
    "    for ii,si in enumerate(subclasses2):\n",
    "        if si in type_name:\n",
    "            subcl_name = subclasses[ii]\n",
    "            break\n",
    "    ge_data_c_subcl.loc[ge_data_c_subcl.loc[:,'name_long_pre']==type_name,'subclass_pre'] =subcl_name\n",
    "    \n",
    "for i in range(subclasses_visp_2.shape[0]):\n",
    "    type_name = str.lower(subclasses_visp_2.loc[i,0])\n",
    "    for ii,si in enumerate(subclasses2):\n",
    "        if si in type_name:\n",
    "            subcl_name = subclasses[ii]\n",
    "            break\n",
    "    ge_data_c_subcl.loc[ge_data_c_subcl.loc[:,'name_long_post']==type_name,'subclass_post'] =subcl_name    \n",
    "\n",
    "#     ge_data_c_subcl.loc[ge_data_c_subcl.loc[:,'name_long'].isin([subclasses_visp.loc[i,0]]),'subclass_pre'] =subclasses_visp.loc[i,'pre']\n",
    "#     ge_data_c_subcl.loc[ge_data_c_subcl.loc[:,'name_long'].isin([subclasses_visp.loc[i,0]]),'subclass_post'] =subclasses_visp.loc[i,'post']\n",
    "\n",
    "#ge_data_c_subcl    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ge_data_subcl = pd.concat([ge_data_c_subcl, ge_data_h_subcl],sort=False)\n",
    "subclasses2 =pd.concat([subclasses, pd.Series(subclasses2_hipp)])\n",
    "ge_data_subcl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "# TRAIN MODEL\n",
    "\n",
    "##\n",
    "##\n",
    "##  train RF \n",
    "##\n",
    "##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subclasses2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes = ['all','ex', 'inh', 'ex_ctx', 'ex_hipp','ex_ec','mge','cge','pvalb','vip','sst','cck','lamp5','ex_he',\n",
    "          'ex_l23','ex_l6','ex_l4','ca','dg','pvalb_h','pvalb_c','sst_h','sst_c']\n",
    "y_features = pd.DataFrame([['A2/A1_20Hz', 'A5/A1_20Hz', 'A2/A1_50Hz','A5/A1_50Hz','A250/A1','A1000/A1'],[0,4,8,11,5,6]],index=['names','index_y'])\n",
    "print(\"y features : \\n\", y_features,\"\\n\\n\")\n",
    "subclasses2 = pd.unique(subclasses2 )\n",
    "print(\"subclasses : \\n\",subclasses2,\"\\n\\n\") \n",
    "\n",
    "cl_in_subcl = np.array([[0,1,2,3,4,5,6,7,8,9,10,11,12],[3,4,5,6,7,8],\n",
    "                        [0,1,2,9,10,11,12],[3,4,5],[6,7],[8],[1,9,0,10],\n",
    "                        [2,11,12],[1,9],[2],[0,10],[12],[11],[6,7,8],[3],[4],[5],[7],[6],[9],[1],[10],[0]])\n",
    "classes = pd.DataFrame([np.array(classes), cl_in_subcl],index=['classes','subclasses'])\n",
    "classes = classes.T.set_index('classes').T\n",
    "classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_data(X,y,data_subclasses,subclasses):\n",
    "    #X_train = np.array([]).reshape((0,X.shape[1]))\n",
    "    #y_train = np.array([]).reshape((0,y.shape[1]))\n",
    "\n",
    "    in_subcl = (data_subclasses.loc[:,'subclass_post']!=False)\n",
    "    for i0 in range(len(subclasses[0])):\n",
    "        in_subcl = in_subcl|(data_subclasses.loc[:,'subclass_pre']==subclasses[0][i0])\n",
    "        \n",
    "    for i0 in range(len(subclasses[1])):        \n",
    "        #in_subcl = in_subcl|(ge_data_subcl.loc[:,'subclass_post']==subclasses2[sbi[i0]])\n",
    "        in_subcl = in_subcl|(data_subclasses.loc[:,'subclass_post']==subclasses[1][i0])\n",
    "        \n",
    "    in_subcl = np.nonzero(in_subcl.values)[0]\n",
    "    #X_train = np.concatenate([X_train,X[in_subcl]])   #np.delete(X, in_subcl,axis=0)\n",
    "    #y_train = np.concatenate([y_train,y[in_subcl]])   #np.delete(y, in_subcl,axis=0)\n",
    "    X_train = np.copy(X[in_subcl])   #np.delete(X, in_subcl,axis=0)\n",
    "    y_train = np.copy(y[in_subcl])   #np.delete(y, in_subcl,axis=0)\n",
    "    return X_train, y_train    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modf = pd.read_excel('synapses_types_tree.xlsx')\n",
    "modf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stp_columns3 = []\n",
    "stp_columns0 = ['A1',\n",
    "                'A2',\n",
    "                'A3',\n",
    "                'A4',\n",
    "                'A5']\n",
    "stpn_fr=[s+'_20Hz' for s in stp_columns0[1:]]\n",
    "stp_columns3 = stp_columns3 +stpn_fr+['A250_20Hz', 'A1000_20Hz']\n",
    "\n",
    "stpn_fr=[s+'_50Hz' for s in stp_columns0[0:]]\n",
    "stp_columns3 = stp_columns3 +stpn_fr+['A250_50Hz', 'A1000_50Hz']\n",
    "\n",
    "stpn_fr=[s+'_10Hz' for s in stp_columns0[0:]]\n",
    "stp_columns3 = stp_columns3 +stpn_fr+['A250_10Hz', 'A1000_10Hz']\n",
    "stp_columns3\n",
    "\n",
    "X_train = pd.DataFrame(np.copy(X),columns = annot_columns+ge_columns[6:])   #np.delete(X, in_subcl,axis=0)\n",
    "y_train = pd.DataFrame(np.copy(y),columns = annot_columns+stp_columns3)     #np.delete(X, in_subcl,axis=0)\n",
    "\n",
    "X_train = pd.concat([X_train,y_train.loc[:,stp_columns3]],axis=1)\n",
    "\n",
    "#X_train \n",
    "\n",
    "\n",
    "#'cla_cell_type_l2_pre'\n",
    "#'cla_cell_type_l2_post'\n",
    "#'cla_cell_type_l3_pre'\n",
    "#'cla_cell_type_l3_post'\n",
    "\n",
    "\n",
    "#classes \tall \tex \tinh \tex_ctx \tex_hipp \tex_ec \tmge \tcge \tpvalb \tvip \tsst \tcck \tlamp5\n",
    "\n",
    "modf = pd.read_excel('synapses_types_tree.xlsx')\n",
    "modf2 = pd.read_excel('synapses_types_tree.xlsx',sheet_name='Sheet2')\n",
    "modf3 = pd.read_excel('synapses_types_tree.xlsx',sheet_name='Sheet3')\n",
    "mod =[ {'name': 'all',    'classes':['all'],                                   'classes':'pre'},\n",
    "       {'name': 'ex_inh', 'classes':['ex','inh'],                              'side':'pre'},\n",
    "       {'name': 'ex_inh_post', 'classes':['ex','inh'],                         'side':'post'},\n",
    "       {'name': 'ex2_inh', 'classes':['ex_ctx','ex_ec','ex_hipp','inh'],       'side':'pre'},\n",
    "       {'name': 'ex_inh2', 'classes':['ex','mge','cge'],                       'side':'pre'},\n",
    "       {'name': 'ex_inh', 'classes':['ex','inh'],                              'side':'pre'},\n",
    "       {'name': 'ex2_inh', 'classes':['ex_ctx','ex_ec','ex_hipp','inh'],       'side':'pre'},\n",
    "       {'name': 'ex_inh3', 'classes':['ex','pvalb','sst','cge'],               'side':'pre'},\n",
    "       {'name': 'ex_inh4', 'classes':['ex','pvalb','sst','vip','cck','lamp5'], 'side':'pre'},\n",
    "     ]\n",
    "\n",
    "subcl_pre=ge_data_subcl.loc[:,'subclass_pre'].copy()\n",
    "subcl_pre.index=X_train.index\n",
    "\n",
    "subcl_post=ge_data_subcl.loc[:,'subclass_post'].copy()\n",
    "subcl_post.index=X_train.index\n",
    "\n",
    "# classes_columns = []\n",
    "# for md in mod:\n",
    "#     in_subcl = subcl_pre.copy()\n",
    "#     for icl, cl in enumerate(md['classes']):\n",
    "#         #in_subcl = in_subcl|(ge_data_subcl.loc[:,'subclass_pre'].isin(classes.loc['subclasses']))\n",
    "#         subcl = subclasses2[classes.loc['subclasses',cl]]\n",
    "#         in_subcl.loc[subcl_pre.isin(subcl)] = icl #icl\n",
    "#     in_subcl.name =  md['name']  \n",
    "#     classes_columns = classes_columns + [in_subcl.name]\n",
    "#     X_train = pd.concat([X_train,in_subcl],axis=1)    \n",
    " \n",
    "# #For modf case     \n",
    "# classes_columns = []\n",
    "# for i in modf.index:\n",
    "#     md = modf.loc[i,:]\n",
    "#     in_subcl = subcl_pre.copy()\n",
    "#     mdcl = np.char.strip(np.array(str.split(md['classes'],',')))\n",
    "#     for icl, cl in enumerate(mdcl):\n",
    "#         #in_subcl = in_subcl|(ge_data_subcl.loc[:,'subclass_pre'].isin(classes.loc['subclasses']))\n",
    "#         subcl = subclasses2[classes.loc['subclasses',cl]]\n",
    "#         in_subcl.loc[subcl_pre.isin(subcl)] = icl #icl\n",
    "#     in_subcl.name =  md['name']+'__pre'  \n",
    "#     classes_columns = classes_columns + [in_subcl.name]\n",
    "#     X_train = pd.concat([X_train,in_subcl],axis=1)  \n",
    "    \n",
    "#     in_subcl = subcl_post.copy()\n",
    "#     mdcl = np.char.strip(np.array(str.split(md['classes_post'],',')))\n",
    "#     for icl, cl in enumerate(mdcl):\n",
    "#         #in_subcl = in_subcl|(ge_data_subcl.loc[:,'subclass_pre'].isin(classes.loc['subclasses']))\n",
    "#         subcl = subclasses2[classes.loc['subclasses',cl]]\n",
    "#         in_subcl.loc[subcl_post.isin(subcl)] = icl #icl\n",
    "#     in_subcl.name =  md['name']+'__post'  \n",
    "#     classes_columns = classes_columns + [in_subcl.name]\n",
    "#     X_train = pd.concat([X_train,in_subcl],axis=1)\n",
    "\n",
    "## modf3 case - independent columns for pre and post cell types\n",
    "classes_columns = []\n",
    "for i in modf3.index:\n",
    "    md = modf3.loc[i,:]\n",
    "    in_subcl = subcl_pre.copy()\n",
    "    in_subcl.loc[:]=0 \n",
    "    mdcl = np.char.strip(np.array(str.split(md['classes'],',')))\n",
    "    for icl, cl in enumerate(mdcl):\n",
    "        #in_subcl = in_subcl|(ge_data_subcl.loc[:,'subclass_pre'].isin(classes.loc['subclasses']))\n",
    "        subcl = subclasses2[classes.loc['subclasses',cl]]\n",
    "        in_subcl.loc[subcl_pre.isin(subcl)] = cl #icl+1 #icl\n",
    "    in_subcl.name =  md['name']+'__pre'  \n",
    "    classes_columns = classes_columns + [in_subcl.name]\n",
    "    X_train = pd.concat([X_train,in_subcl],axis=1)  \n",
    "    \n",
    "    in_subcl = subcl_post.copy()\n",
    "    in_subcl.loc[:]=0\n",
    "    mdcl = np.char.strip(np.array(str.split(md['classes'],',')))\n",
    "    for icl, cl in enumerate(mdcl):\n",
    "        #in_subcl = in_subcl|(ge_data_subcl.loc[:,'subclass_pre'].isin(classes.loc['subclasses']))\n",
    "        subcl = subclasses2[classes.loc['subclasses',cl]]\n",
    "        #if cl!='all' \n",
    "        in_subcl.loc[subcl_post.isin(subcl)] = cl #icl+1 #icl\n",
    "    in_subcl.name =  md['name']+'__post'  \n",
    "    classes_columns = classes_columns + [in_subcl.name]\n",
    "    X_train = pd.concat([X_train,in_subcl],axis=1)\n",
    "\n",
    "# add all short subclasses  \n",
    "in_subcl = subcl_pre.copy()    \n",
    "in_subcl.name =  'vse'+'__pre'  \n",
    "classes_columns = classes_columns + [in_subcl.name]\n",
    "X_train = pd.concat([X_train,in_subcl],axis=1)\n",
    "    \n",
    "in_subcl = subcl_post.copy()    \n",
    "in_subcl.name =  'vse'+'__post'  \n",
    "classes_columns = classes_columns + [in_subcl.name]\n",
    "X_train = pd.concat([X_train,in_subcl],axis=1)\n",
    "         \n",
    "#for i in\n",
    "#mod1_ = mod1_ +  [ sm.OLS( y_train, X_train) ]\n",
    "#res = mod.fit()\n",
    "#print(res.summary())\n",
    "ge_columns_train      = ge_columns[6:]\n",
    "annot_columns_train   = annot_columns\n",
    "stp_columns_train     = stp_columns3\n",
    "classes_columns_train = classes_columns\n",
    "print(X_train.columns)\n",
    "X_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_subcl = subcl_post.copy()    \n",
    "in_subcl.name =  'vse'+'__post'  \n",
    "classes_columns = classes_columns + [in_subcl.name]\n",
    "X_train = pd.concat([X_train,in_subcl],axis=1)\n",
    "         \n",
    "#for i in\n",
    "#mod1_ = mod1_ +  [ sm.OLS( y_train, X_train) ]\n",
    "#res = mod.fit()\n",
    "#print(res.summary())\n",
    "ge_columns_train      = ge_columns[6:]\n",
    "annot_columns_train   = annot_columns\n",
    "stp_columns_train     = stp_columns3\n",
    "classes_columns_train = classes_columns\n",
    "print(X_train.columns)\n",
    "X_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#X_train0.to_excel('X_train0.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns  =[\n",
    "[ 'ge_columns_train',\n",
    "  'annot_columns_train',\n",
    "  'stp_columns_train',\n",
    "  'classes_columns_train',],\n",
    "[ ge_columns_train,\n",
    "  annot_columns_train,\n",
    "  stp_columns_train,\n",
    "  classes_columns_train,\n",
    "]]\n",
    "columns = pd.DataFrame(columns)\n",
    "columns=columns.T.set_index([0]).T\n",
    "columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train.columns[X_train.columns[:].duplicated()]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Remove duplicated columns ...')\n",
    "print(X_train.columns.shape)\n",
    "is_dup = np.nonzero(X_train.columns[:].duplicated())[0]\n",
    "is_dup0 = is_dup\n",
    "for i in is_dup:\n",
    "    X_train.columns = np.array(X_train.columns[:i].tolist()+ [X_train.columns[i]+'_2']+X_train.columns[i+1:].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "is_dup = np.nonzero(X_train.columns[:].duplicated())[0]\n",
    "print(is_dup0)\n",
    "is_dup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gs_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns.to_hdf('stp_to_ge_all_columns'+gs_name+'.hdf',key='data')\n",
    "X_train.to_hdf('stp_to_ge_all_data'+gs_name+'.hdf',key='data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes_columns_train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# START HERE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gs_name = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "columns  = pd.read_hdf('stp_to_ge_all_columns'+gs_name+'.hdf')\n",
    "ge_columns_train      = columns.loc[:,'ge_columns_train'].values[0]\n",
    "annot_columns_train   = columns.loc[:,'annot_columns_train'].values[0]\n",
    "stp_columns_train     = columns.loc[:,'stp_columns_train'].values[0]\n",
    "classes_columns_train = columns.loc[:,'classes_columns_train'].values[0]\n",
    "#print(X_train.columns)\n",
    "X_train0 = pd.read_hdf('stp_to_ge_all_data'+gs_name+'.hdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#classes \tall \tex \tinh \tex_ctx \tex_hipp \tex_ec \tmge \tcge \tpvalb \tvip \tsst \tcck \tlamp5\n",
    "\n",
    "modf = pd.read_excel('synapses_types_tree.xlsx')\n",
    "modf2 = pd.read_excel('synapses_types_tree.xlsx',sheet_name='Sheet2')\n",
    "modf3 = pd.read_excel('synapses_types_tree.xlsx',sheet_name='Sheet3')\n",
    "mod =[ {'name': 'all',    'classes':['all'],                                   'classes':'pre'},\n",
    "       {'name': 'ex_inh', 'classes':['ex','inh'],                              'side':'pre'},\n",
    "       {'name': 'ex_inh_post', 'classes':['ex','inh'],                         'side':'post'},\n",
    "       {'name': 'ex2_inh', 'classes':['ex_ctx','ex_ec','ex_hipp','inh'],       'side':'pre'},\n",
    "       {'name': 'ex_inh2', 'classes':['ex','mge','cge'],                       'side':'pre'},\n",
    "       {'name': 'ex_inh', 'classes':['ex','inh'],                              'side':'pre'},\n",
    "       {'name': 'ex2_inh', 'classes':['ex_ctx','ex_ec','ex_hipp','inh'],       'side':'pre'},\n",
    "       {'name': 'ex_inh3', 'classes':['ex','pvalb','sst','cge'],               'side':'pre'},\n",
    "       {'name': 'ex_inh4', 'classes':['ex','pvalb','sst','vip','cck','lamp5'], 'side':'pre'},\n",
    "     ]\n",
    "modf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modf3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decode_classes2(class2,modf3):\n",
    "    cla=[]\n",
    "    for cl2 in class2:\n",
    "        pref = str.rfind(cl2,'_pre')\n",
    "        if pref==-1:\n",
    "            pref = str.rfind(cl2,'_post')\n",
    "            if pref==-1:\n",
    "                spref='__pre'\n",
    "                \n",
    "            else:\n",
    "                cl2 = cl2[0:pref]\n",
    "                spref='__post'\n",
    "                \n",
    "        else:\n",
    "            spref='__pre'\n",
    "            cl2 = cl2[0:pref]\n",
    "        \n",
    "        #print(cl2)\n",
    "        md = modf3.loc[modf3.loc[:,'name_classes2']==cl2 ,'classes2'].iloc[0]  \n",
    "        #print(md)\n",
    "        mdcl = np.char.strip(np.array(str.split(md,',')))\n",
    "        for icl, cl in enumerate(mdcl):\n",
    "            #in_subcl = in_subcl|(ge_data_subcl.loc[:,'subclass_pre'].isin(classes.loc['subclasses']))\n",
    "            cla = cla + [cl + spref]\n",
    "    \n",
    "    return cla    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decode_classes2(['ex2_in_pre','ex2_in3_post'],modf3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sklearn as sk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Fit cell types hierarhy tree with linear models in leaves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_classes_tree(X2,y2,X2_cl,cli,model_type='lasso',nmin = 20, alpha=1, l1_ratio=0.5):\n",
    "    model = []\n",
    "    support =[]\n",
    "    ncl1 = np.max(X2_cl[:,cli[0]])+1 \n",
    "    ncl2 = np.max(X2_cl[:,cli[1]])+1 \n",
    "    for i1 in range(ncl1):\n",
    "        model1 = []\n",
    "        support1 = []\n",
    "        for i2 in range(ncl2):\n",
    "\n",
    "            is_in_cl = (X2_cl[:,cli[0]]==i1)&(X2_cl[:,cli[1]]==i2)\n",
    "            #           support2 = [np.sum(is_in_cl),\n",
    "            #           np.nonzero((X2_cl[:,cli[0]]==i1))[0],\n",
    "            #           np.nonzero((X2_cl[:,cli[1]]==i2))[0]]\n",
    "            support2 = np.sum(is_in_cl)\n",
    "            \n",
    "            support1 = support1 + [support2]\n",
    "            regr = 0\n",
    "            if support2>nmin:\n",
    "                y_train = y2[is_in_cl,:]\n",
    "                X_train = X2[is_in_cl,:]\n",
    "                if model_type=='ridge':\n",
    "                    regr = sk.linear_model.Ridge(alpha=0.50, fit_intercept=True, normalize=False,\n",
    "                                                 copy_X=True, max_iter=None, tol=0.001, solver='auto',\n",
    "                                                 random_state=None)\n",
    "                \n",
    "                elif model_type=='lasso':\n",
    "                    regr = sk.linear_model.Lasso(alpha=alpha, fit_intercept=True, normalize=False,\n",
    "                                                 precompute=False, copy_X=True,\n",
    "                                                 max_iter=1000, tol=0.0001, warm_start=False, positive=False, \n",
    "                                                 random_state=None, selection='cyclic')\n",
    "                elif model_type=='elastic_net':   \n",
    "                    regr = sk.linear_model.ElasticNet(alpha=alpha, l1_ratio=l1_ratio, fit_intercept=True,\n",
    "                                    normalize=False,\n",
    "                                    max_iter=1000, copy_X=True, tol=0.0001, warm_start=False,\n",
    "                                    random_state=None, selection='cyclic')\n",
    "                else:\n",
    "                    print(\" model types supported: ridge, lasso, elastic_net\")\n",
    "                \n",
    "                regr.fit(X_train, y_train)\n",
    "                #y_pred = regr.predict(X_test[:,nannot:]) \n",
    "\n",
    "\n",
    "            model1 = model1 + [regr]\n",
    "            \n",
    "        support = support + [np.array(support1)]\n",
    "        model = model + [np.array(model1)]    \n",
    "    return np.array(model), np.array(support)        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_classes_tree(model,X2,X2_cl,cli,nout=1,nmin = 20):\n",
    "    \n",
    "    ncl1 = np.max(X2_cl[:,cli[0]])+1 \n",
    "    ncl2 = np.max(X2_cl[:,cli[1]])+1 \n",
    "    y_pred = np.zeros((X2.shape[0],nout))\n",
    "    for i1 in range(ncl1):\n",
    "        #print(i1,\" f1\")\n",
    "        if len(model)>i1:\n",
    "            #print(i1,\" f2\")\n",
    "            model1 = model[i1]\n",
    "            for i2 in range(ncl2):\n",
    "\n",
    "                is_in_cl = (X2_cl[:,cli[0]]==i1)&(X2_cl[:,cli[1]]==i2)\n",
    "\n",
    "                support2 = np.sum(is_in_cl)\n",
    "                \n",
    "                #print(i1,i2,\" f3 \",support2)\n",
    "                #regr = 0\n",
    "                if support2>nmin:\n",
    "                    #print(i1,i2,\" f3 \",support2)\n",
    "                    if len(model1)>i2:\n",
    "                        #print(i1,i2,\" f4\")\n",
    "                        regr = model1[i2]\n",
    "                        #regr = sk.linear_model.MultiTaskElasticNet(alpha=alpha, l1_ratio=l1_ratio, fit_intercept=True,\n",
    "                        #                    normalize=False,\n",
    "                        #                    max_iter=1000, copy_X=True, tol=0.0001, warm_start=False,\n",
    "                        #                    random_state=None, selection='cyclic')\n",
    "                        #regr.fit(X_train, y_train)\n",
    "                        if regr!=0:\n",
    "                            X_test = X2[is_in_cl,:]\n",
    "                            y_pred[is_in_cl] = regr.predict(X_test).reshape((-1,1)) \n",
    "    \n",
    "    return y_pred       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Draw some interesting plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load important genes from iRF search\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "imp  = np.load('importance.npy')\n",
    "imp = pd.DataFrame(imp)\n",
    "imp = pd.concat([imp,pd.DataFrame(ge_columns_train )],axis=1)\n",
    "imp.columns = ['importance', 'genes']\n",
    "imp = imp.set_index('genes')\n",
    "\n",
    "imp0  = np.load('importance0.npy')\n",
    "imp0 = pd.DataFrame(imp0)\n",
    "imp0 = pd.concat([imp0,pd.DataFrame(ge_columns_train )],axis=1)\n",
    "imp0.columns = ['importance0', 'genes']\n",
    "imp0 = imp0.set_index('genes')\n",
    "\n",
    "\n",
    "# List of important genes : iRF\n",
    "imp = pd.concat([imp,imp0],axis=1)\n",
    "\n",
    "imp50 = imp.sort_values('importance',ascending=False).iloc[0:49,:]\n",
    "plt.plot(imp.sort_values('importance',ascending=False).loc[:,['importance','importance0']].values)\n",
    "imp50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes_columns_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# note: synapse type hierarhy in X(:,classes_columns_train) should be from modf!!!\n",
    "\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "t1=time.time()\n",
    "\n",
    "Dn = 5\n",
    "cla_n = classes_columns_train #['ex_inh']\n",
    "\n",
    "ge_n = imp50.index[0:50].tolist() #ge_columns_train #imp50.index[0:25].tolist()\n",
    "\n",
    "#ge_n =ge_columns_train\n",
    "X2 = X_train0.loc[:,ge_n ]\n",
    "\n",
    "do_PCA=False\n",
    "if do_PCA:\n",
    "    # PCA transformation\n",
    "    pca=PCA(n_components=30)\n",
    "    X2 =pca.fit_transform(X_train0.loc[:,ge_n ].values)\n",
    "    print(pca.explained_variance_ratio_)\n",
    "    print(pca.singular_values_) \n",
    "    ge_n = ['pca'+str(s) for s in range(X2.shape[1])]\n",
    "    X2=X2[0:X_train0.shape[0]:Dn,:]\n",
    "else:\n",
    "    #i_cl = np.nonzero(X2.columns.isin(cla_n))[0]\n",
    "    X2=X2.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "\n",
    "X2_cl = X_train0.loc[:,cla_n ]\n",
    "\n",
    "X2_cl=X2_cl.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "stp_n = ['A5_20Hz']\n",
    "#  'A2_20Hz',\n",
    "#  'A5_20Hz',\n",
    "#  'A250_20Hz',\n",
    "#  'A1000_20Hz',\n",
    "#  'A2_50Hz',\n",
    "#  'A2_10Hz',\n",
    "#  'A5_50Hz',\n",
    "#  'A5_10Hz',\n",
    "y2 =  X_train0.loc[:,stp_n].iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "##header = ge_n + cla_n + stp_n \n",
    "#header  = ge_n  + stp_n \n",
    "\n",
    "iy0=0\n",
    "mod_index = modf.index[[14]] #[3,9,13,16] [3] # modf.index #[0, 3, 9, 11, 13, 15]\n",
    "\n",
    "for ig,g in enumerate(ge_n):\n",
    "    #f, ax =plt.subplots(figsize=(16, 10))\n",
    "    ##f, ax = plt.figure()\n",
    "    ##ax = f.add_axes()\n",
    "    \n",
    "    #plt.title(stp_columns[iy]+\", model : \"+mdn+\", cv : \"+str(icv))\n",
    "    #plt.title(stp_columns[iy]+\", model : \"+mdn)\n",
    "    \n",
    "    #ax.set_title(stp_columns[i]+' '+', subclass out of bag: '+str(i0))\n",
    "    #yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "    #plt.plot(y2.ravel(),'ob')\n",
    "    #plt.plot(y_pred.ravel(),'xr')\n",
    "    for i in mod_index:\n",
    "        md = modf.loc[i,:]\n",
    "\n",
    "\n",
    "        mdn = md['name']\n",
    "        #print(mdn)\n",
    "\n",
    "        cla_n = np.array(classes_columns_train)\n",
    "\n",
    "        cli = np.nonzero(cla_n==mdn+'__pre')[0] + np.arange(2)\n",
    "\n",
    "        mdcl2 = np.char.strip(np.array(str.split(md['classes_post'],',')))\n",
    "        mdcl1 = np.char.strip(np.array(str.split(md['classes'],','))) \n",
    "        #print(mdcl1)\n",
    "        #print(mdcl2)\n",
    "        \n",
    "        f, ax =plt.subplots(figsize=(16, 10))\n",
    "        ##f, ax = plt.figure()\n",
    "        ##ax = f.add_axes()\n",
    "    \n",
    "        #plt.title(stp_columns[iy]+\", model : \"+mdn+\", cv : \"+str(icv))\n",
    "        plt.title(g+\", model : \"+mdn+\" classes pre : \"+str(mdcl1)+\" classes post : \"+str(mdcl2))\n",
    "        \n",
    "        plt.scatter(X2[:,ig],y2.ravel(),marker='o',c=X2_cl[:,cli[1]],s=100,alpha=0.5)\n",
    "        plt.scatter(X2[:,ig],y2.ravel(),marker='o',c=X2_cl[:,cli[0]],s=10)\n",
    "       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fit set of linear models for each level of cell types hierarhy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# note: synapse type hierarhy in X(:,classes_columns_train) should be from modf!!!\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "t1=time.time()\n",
    "\n",
    "Dn = 10\n",
    "cla_n = classes_columns_train #['ex_inh']\n",
    "ge_n = ge_columns_train #imp50.index[0:25].tolist()\n",
    "X2 = X_train0.loc[:,ge_n ]\n",
    "X2_cl = X_train0.loc[:,cla_n ]\n",
    "#i_cl = np.nonzero(X2.columns.isin(cla_n))[0]\n",
    "X2=X2.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "X2_cl=X2_cl.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "stp_n = ['A2_20Hz']\n",
    "#  'A2_20Hz',\n",
    "#  'A5_20Hz',\n",
    "#  'A250_20Hz',\n",
    "#  'A1000_20Hz',\n",
    "#  'A2_50Hz',\n",
    "#  'A2_10Hz',\n",
    "#  'A5_50Hz',\n",
    "#  'A5_10Hz',\n",
    "y2 =  X_train0.loc[:,stp_n].iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "##header = ge_n + cla_n + stp_n \n",
    "#header  = ge_n  + stp_n \n",
    "\n",
    "iy0=0\n",
    "mod_index = modf.index #[3] # modf.index #[0, 3, 9, 11, 13, 15]\n",
    "Y_pred = []\n",
    "for i in mod_index:\n",
    "    md = modf.loc[i,:]\n",
    "\n",
    "\n",
    "    mdn = md['name']\n",
    "    print(mdn)\n",
    "\n",
    "    cla_n = np.array(classes_columns_train)\n",
    "\n",
    "    cli = np.nonzero(cla_n==mdn+'__pre')[0] + np.arange(2)\n",
    "\n",
    "    mdcl2 = np.char.strip(np.array(str.split(md['classes_post'],',')))\n",
    "    mdcl1 = np.char.strip(np.array(str.split(md['classes'],','))) \n",
    "    print(mdcl1)\n",
    "    print(mdcl2)\n",
    "\n",
    "    y_pred = np.zeros((0,y2.shape[1]))\n",
    "    \n",
    "    ncv = 10\n",
    "    n_samp_cv = np.rint(X2.shape[0]/ncv)\n",
    "    samples_all = np.arange(X2.shape[0])\n",
    "    r2cv = np.zeros(ncv)\n",
    "    r4cv = np.zeros(ncv)\n",
    "    for icv in range(ncv): #range(ncv):\n",
    "        samples_test = (np.arange(n_samp_cv) + icv*n_samp_cv).astype(int)\n",
    "        samples_train = np.delete(np.copy(samples_all),samples_test)\n",
    "        X2train, y2train, X2train_cl = X2[samples_train,:], y2[samples_train,:], X2_cl[samples_train,:]\n",
    "        X2test, y2test, X2test_cl = X2[samples_test,:], y2[samples_test,:], X2_cl[samples_test,:]\n",
    "        \n",
    "        model, support = fit_classes_tree(X2train,y2train,X2train_cl,cli,model_type='elastic_net',\n",
    "                                          nmin = 20, alpha=0.2, l1_ratio=0.05)\n",
    "        #model, support = fit_classes_tree(X2,y2,X2_cl,cli,model_type='ridge',nmin = 20, alpha=1, l1_ratio=0.5)\n",
    "        y_pred_i = predict_classes_tree(model,X2test,X2test_cl,cli,nout=y2.shape[1],nmin = 2)\n",
    "        y_pred = np.concatenate([y_pred, y_pred_i],axis=0)\n",
    "        \n",
    "        iy=iy0\n",
    "        if np.var(y2[:,iy])!=0:\n",
    "            nonz = y_pred_i[:,iy]!=0\n",
    "            if np.sum(nonz)>0:\n",
    "                r2cv[icv]=1 - np.mean((y2test[nonz,iy] - y_pred_i[nonz,iy])**2)/np.var(y2[:,iy])\n",
    "                r4cv[icv]=1 - np.mean((y2test[nonz,iy] - y_pred_i[nonz,iy])**2)/np.var(y2test[nonz,iy])\n",
    "    \n",
    "    Y_pred = Y_pred + [y_pred]\n",
    "    \n",
    "    f, ax =plt.subplots(figsize=(16, 10))\n",
    "    ##f, ax = plt.figure()\n",
    "    ##ax = f.add_axes()\n",
    "    \n",
    "    #plt.title(stp_columns[iy]+\", model : \"+mdn+\", cv : \"+str(icv))\n",
    "    plt.title(stp_columns_train[iy]+\", model : \"+mdn)\n",
    "    \n",
    "    #ax.set_title(stp_columns[i]+' '+', subclass out of bag: '+str(i0))\n",
    "    #yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "    plt.plot(y2.ravel(),'ob')\n",
    "    plt.plot(y_pred.ravel(),'xr')\n",
    "    #plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n",
    "    iy=iy0\n",
    "    r2=0\n",
    "    r4=0\n",
    "    if np.var(y2[:,iy])!=0:\n",
    "        r2=1 - np.mean((y2[:,iy] - y_pred[:,iy])**2)/np.var(y2[:,iy])\n",
    "        r4=1 - np.var(y_pred[:,iy])/np.var(y2[:,iy])\n",
    "        \n",
    "    print(\"R**2 = \",r2,\"  Part of Var = \",r4,\"\\n\",\n",
    "          \"  R**2 cv mean = \", np.mean(r2cv),\"\\n  R**2 cv unnormed mean = \", np.mean(r4cv),\n",
    "          \"\\n  R**2 cv = \", r2cv,\"\\n  R**2 cv unnormed = \", r4cv,\"\\n\\n\")\n",
    "    \n",
    "    \n",
    "t2 =time.time()    \n",
    "print(\"Elapsed time \",t2-t1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Crossvalidation - kernel SVR on scVI space, RF with distance spliting penalty? - nonlocal predictors vs local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gs_name = '_gs1512' #'_gs15656' #'_gs1512' #'_gs15656' #'_gs1512'   #'_gs5188' #'_gs1512'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "do_laptop = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "columns  = pd.read_hdf('stp_to_ge_all_columns'+gs_name+'.hdf')\n",
    "ge_columns_train      = columns.loc[:,'ge_columns_train'].values[0]\n",
    "annot_columns_train   = columns.loc[:,'annot_columns_train'].values[0]\n",
    "stp_columns_train     = columns.loc[:,'stp_columns_train'].values[0]\n",
    "classes_columns_train = columns.loc[:,'classes_columns_train'].values[0]\n",
    "#print(X_train.columns)\n",
    "X_train0 = pd.read_hdf('stp_to_ge_all_data'+gs_name+'.hdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train0.shape\n",
    "print(X_train0.columns[0:20])\n",
    "print(X_train0.columns[200:300])\n",
    "X_train0.columns[300:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gs_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (gs_name=='_gs5188')|(gs_name=='_gs1512')|(gs_name=='_gs15656'):\n",
    "    nlf = 10   # 20 were in scvi model\n",
    "elif (gs_name=='_gs219'):\n",
    "    nlf = 10   # 10 were in scvi model\n",
    "lf_names = np.arange(nlf).astype(str)\n",
    "lf_names1 =  ['pre__'+s for s in lf_names.tolist()]\n",
    "lf_names2 =  ['post__'+s for s in lf_names.tolist()]\n",
    "\n",
    "lf_scvi_names = lf_names1 + lf_names2\n",
    "\n",
    "ge_columns_train2 = list(set(ge_columns_train).difference(set(lf_scvi_names)))\n",
    "lf_scvi_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### tSNE : cell types, STP space\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_aba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if do_laptop==False:\n",
    "    d_aba  ='/home/stepaniu/Documents/references/transcriptomes to STP/transcriptomes/ABA_2019/'\n",
    "else:\n",
    "    d_aba ='/Users/stepaniu/Documents/jan_2020/ABA_2019/'\n",
    "\n",
    "df_aba_g = pd.read_csv(d_aba + 'medians.csv')\n",
    "df_aba_c = pd.read_csv(d_aba + 'sample_annotations.csv')\n",
    "df_aba_tsne = pd.read_csv(d_aba + '2d_coordinates.csv')\n",
    "df_aba_tsne =df_aba_tsne.set_index('sample_name')\n",
    "df_aba_tsne "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#in_20190=df_aba_c.loc[:,'region_label'].isin(['VISp'])\n",
    "#set(df_aba_c.loc[in_20190&df_aba_c.loc[:,'class_label'].isin(['GABAergic', 'Glutamatergic']),'subclass_label'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#set(df_aba_vis_c.loc[df_aba_vis_c.loc[:,'class'].isin(['GABAergic', 'Glutamatergic']),'subclass'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#do_laptop=0\n",
    "if do_laptop:\n",
    "    d4 = '~/Documents/references/lists_of_proteins_2/Analysis_of_proteins/genetic_types/transcriptomic_interneurons/'\n",
    "else:\n",
    "    d4='~/Documents/references/transcriptomes to STP/transcriptomes/ABA/mouse_VISp_gene_expression_matrices_2018-06-14/'\n",
    "\n",
    "fn='mouse_VISp_2018-06-14_samples-columns.csv'\n",
    "df_aba_vis_c=pd.read_csv(d4+fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_aba_vis_c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "\n",
    "\n",
    "#\n",
    "#\n",
    "#  ALLIGNMET ABA 2018 and ABA 2019 samples\n",
    "#\n",
    "#\n",
    "\n",
    "\n",
    "do_alignment=0\n",
    "if do_alignment:\n",
    "    t1=time.time()\n",
    "    do_pre=0\n",
    "    if do_pre:\n",
    "        samps = sampr\n",
    "    else:\n",
    "        samps = sampo\n",
    "\n",
    "    if do_laptop:\n",
    "        d4 = '~/Documents/references/lists_of_proteins_2/Analysis_of_proteins/genetic_types/transcriptomic_interneurons/'\n",
    "    else:\n",
    "        d4='~/Documents/references/transcriptomes to STP/transcriptomes/ABA/mouse_VISp_gene_expression_matrices_2018-06-14/'\n",
    "\n",
    "    fn='mouse_VISp_2018-06-14_samples-columns.csv'\n",
    "    df_aba_vis_c=pd.read_csv(d4+fn)\n",
    "\n",
    "\n",
    "    #sm = [ 'F2S4_160825_057_G01']\n",
    "    #sm  =['F2S4_160331_001_B01']\n",
    "\n",
    "    #in_2019 = df_aba_c.loc[:,'subclass_label'].isin([df_2018_donor['subclass']])\n",
    "\n",
    "    in_20190=df_aba_c.loc[:,'region_label'].isin(['VISp'])\n",
    "\n",
    "    ll = []\n",
    "    nn = []\n",
    "    # index \tsample_name \tcluster \tsample_name_19 \tcluster_label\n",
    "    col1819=['index','sample_name','cluster','sample_name_19','cluster_label']\n",
    "    df_aba18_to_19 = pd.DataFrame([],columns = col1819)\n",
    "    for sm in samps: #.iloc[1:1000]: \n",
    "        if (df_aba18_to_19.loc[:,'sample_name']==sm).sum()==0:\n",
    "            df_2018_donor = df_aba_vis_c.loc[df_aba_vis_c.loc[:,'sample_name'].isin([sm]),['sample_name', 'donor','genotype','sex','subclass','class','cluster','brain_subregion']]\n",
    "            #df_2018_donor\n",
    "            #in_2019=in_2019&df_aba_c.loc[:,'cortical_layer_label'].isin(['L4'])\n",
    "            \n",
    "            #in_2019=in_2019&df_aba_c.loc[:,'external_donor_name_label'].isin(df_2018_donor['donor'])\n",
    "            in_2019=n_20190&df_aba_c.loc[:,'cortical_layer_label'].isin(df_2018_donor['brain_subregion'])\n",
    "            in_2019=in_2019&df_aba_c.loc[:,'full_genotype_label'].isin(df_2018_donor['genotype'])\n",
    "            in_2019=in_2019&df_aba_c.loc[:,'donor_sex_label'].isin(df_2018_donor['sex'])\n",
    "            in_2019=in_2019&df_aba_c.loc[:,'external_donor_name_label'].map(str).isin(df_2018_donor['donor'].map(str))\n",
    "            in_2019 = in_2019&df_aba_c.loc[:,'subclass_label'].isin(df_2018_donor['subclass'])\n",
    "            \n",
    "            ##df_aba_c.loc[in_2019,'sample_name']\n",
    "            col19 = ['sample_name', 'cortical_layer_label',\n",
    "                                             'full_genotype_label','external_donor_name_label',\n",
    "                                             'cluster_label','class_label','subclass_label',\n",
    "                                             'cell_type_alt_alias_label','donor_sex_label',\n",
    "                                             'cell_type_alias_label']\n",
    "            df_out19 = df_aba_c.loc[in_2019,col19]\n",
    "            col19[0] = 'sample_name_19'\n",
    "            df_out19.columns=col19\n",
    "\n",
    "            \n",
    "            ##in_2019=in_2019&df_aba_c.loc[:,'external_donor_name_label'].isin(df_2018_donor['donor'])\n",
    "            in_2018=df_aba_vis_c.loc[:,'brain_subregion'].isin(df_2018_donor['brain_subregion'])\n",
    "            in_2018=in_2018&df_aba_vis_c.loc[:,'genotype'].isin(df_2018_donor['genotype'])\n",
    "            in_2018=in_2018&df_aba_vis_c.loc[:,'sex'].isin(df_2018_donor['sex'])\n",
    "            in_2018=in_2018&df_aba_vis_c.loc[:,'donor'].isin(df_2018_donor['donor'])\n",
    "            in_2018 = in_2018&df_aba_vis_c.loc[:,'subclass'].isin(df_2018_donor['subclass'])\n",
    "            ##df_aba_c.loc[in_2019,'sample_name']\n",
    "            df_out18 = df_aba_vis_c.loc[in_2018,['sample_name', 'brain_subregion','genotype','donor','cluster','class','subclass','sex']]\n",
    "\n",
    "\n",
    "            #if df_out18.shape[0]==df_out19.shape[0]:\n",
    "            nmim = min([df_out18.shape[0],df_out19.shape[0]])\n",
    "            df1819 = pd.concat([df_out18.iloc[:nmim].loc[:,['sample_name','cluster']].reset_index(),\n",
    "                                df_out19.iloc[:nmim].loc[:,['sample_name_19','cluster_label']].reset_index(drop=True)],axis=1)\n",
    "            df_aba18_to_19 = df_aba18_to_19.append(df1819)\n",
    "\n",
    "            ll = ll + [df_out19.loc[:,'cortical_layer_label'].isin(df_2018_donor['brain_subregion']).sum()]\n",
    "            nn = nn + [df_out19.shape[0]]\n",
    "            #print(df_out18)\n",
    "            #print(df_out19)\n",
    "    t2 =time.time()\n",
    "    print('elapsed time '+str(t2-t1))\n",
    "    print(df_aba18_to_19.shape )\n",
    "    if do_pre:\n",
    "        df_aba18_to_19.to_excel('aba18_in_aba19_presynaptic.xlsx')\n",
    "    else:\n",
    "        df_aba18_to_19.to_excel('aba18_in_aba19_postsynaptic.xlsx')\n",
    "        \n",
    "    print(sampr.shape[0])    \n",
    "    df_aba18_to_19 \n",
    "    \n",
    "### filter\n",
    "df_aba18_to_19 = pd.read_excel('aba18_in_aba19_postsynaptic.xlsx')\n",
    "df_aba18_to_19_2 = df_aba18_to_19.loc[df_aba18_to_19.loc[:,'sample_name'].isin(sampo),:].copy()\n",
    "\n",
    "df_aba18_to_19 = pd.read_excel('aba18_in_aba19_presynaptic.xlsx')\n",
    "df_aba18_to_19_3 = df_aba18_to_19.loc[df_aba18_to_19.loc[:,'sample_name'].isin(sampr),:].copy()\n",
    "\n",
    "df_aba18_to_19_2 = df_aba18_to_19_2.append(df_aba18_to_19_3)\n",
    "\n",
    "\n",
    "df_aba18_to_19_2.to_excel('aba18_in_aba19_all_filtered.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_aba18_to_19_2 = pd.read_excel('aba18_in_aba19_all_filtered.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#[np.array([])]*cl18.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#to do !!! - clusters 18 to 19\n",
    "\n",
    "cl18 = set(df_aba18_to_19_2.loc[:,'cluster'])\n",
    "cl18=list(cl18)\n",
    "# stpn_test=[s+'_test' for s in stp_n]\n",
    "cl18_2 = np.array([cl18,[s+' counts' for s in cl18]])\n",
    "cl18_2= np.reshape(cl18_2.transpose(),(-1,))\n",
    "cl18=np.array(cl18)\n",
    "#df_cl18_to_cl19 = pd.DataFrame(np.zeros((5,cl18.shape[0]*2)),columns=cl18_2)\n",
    "df_cl18_to_cl19 = pd.DataFrame(['']*cl18.shape[0],index=cl18)\n",
    "for c18 in cl18:\n",
    "    c19=df_aba18_to_19_2.loc[df_aba18_to_19_2.loc[:,'cluster']==c18,'cluster_label'].value_counts()\n",
    "    print('\\n\\n')\n",
    "    print(c18)\n",
    "    print(c19)\n",
    "    #nmin=min([df_cl18_to_cl19.shape[0],c19.shape[0]])\n",
    "    nmin=min([1,c19.shape[0]])\n",
    "    \n",
    "    #df_cl18_to_cl19.loc[:,c18] = df_cl18_to_cl19.loc[:,c18].astype(str)\n",
    "    df_cl18_to_cl19.loc[c18,0] = c19.index[:nmin].tolist()[0] #.to_numpy()[0] #.tolist()\n",
    "    #df_cl18_to_cl19.iloc[0:nmin,:].loc[:,c18+' counts'] = c19[:nmin].values/c19.values.sum()\n",
    "    #df_cl18_to_cl19.loc[:nmin,c18] = c19[:nmin]/c19.sum()\n",
    "df_cl18_to_cl19"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load STP types based on F/D\n",
    "remove_st=[]\n",
    "\n",
    "df_Stp_type_p = pd.read_excel('stp types probability cortex and hippocampus.xlsx')\n",
    "#print('df_Stp_type_p.shape : ',df_Stp_type_p.shape, '  X_train0.shape : ', X_train0.shape[0]/N_bootstraps)\n",
    "all_samples2 = np.ones(df_Stp_type_p.shape[0])\n",
    "for st in remove_st: # remove bad samples\n",
    "    all_samples2[st] = 0\n",
    "df_Stp_type_p = df_Stp_type_p.iloc[all_samples2==1,:].reset_index()\n",
    "\n",
    "\n",
    "\n",
    "df8 = df_Stp_type_p.loc[:,'median'].copy()  # stp type based on median of f/d\n",
    "df8.name='type'\n",
    "df9=df8.copy()\n",
    "df9.loc[df8>2]=1\n",
    "df9.loc[df8<0.4]=2\n",
    "df9.loc[(df8>=0.4)&(df8<=2)]=3\n",
    "df_Stp_type_p = pd.concat([df_Stp_type_p,df9],axis=1)\n",
    "\n",
    "df_Stp_type_p.loc[:,'median'] = np.log(df_Stp_type_p.loc[:,'median'])\n",
    "\n",
    "\n",
    "df_Stp_type_p_c=df_Stp_type_p.loc[0:53,:]\n",
    "df_Stp_type_p_h=df_Stp_type_p.loc[53:,:]\n",
    "\n",
    "print('df_Stp_type_p.shape : ',df_Stp_type_p.shape)\n",
    "df_Stp_type_p.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_Stp_type_p_h.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_data22 = df_aba_c.loc[:,'subclass_label'].isin([scstpo])\n",
    "tsne1tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_1']\n",
    "tsne2tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_2']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train0"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "#   All synapses : cortex?\n",
    "#\n",
    "#\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "gray = np.array([200, 200, 200, 1])/255\n",
    "colas = np.array([[1,1,0],[0,0,1],[0,1,0]])\n",
    "for fff in range(4):\n",
    "    \n",
    "    plot_gaba = (fff==0)|(fff==2)\n",
    "\n",
    "    # in_data = df_aba_c.loc[:,'class_label'].isin(['GABAergic', 'Glutamatergic'])\n",
    "    # in_data = in_data&df_aba_c.loc[:,'region_label'].isin(['VISp','ENTl','ENTm','HIP'])\n",
    "    # tsne1 = df_aba_tsne.loc[ df_aba_c.loc[in_data, 'sample_name'], 'tsne_1']\n",
    "    # tsne2 = df_aba_tsne.loc[ df_aba_c.loc[in_data, 'sample_name'], 'tsne_2']\n",
    "    # c = df_aba_c.loc[in_data, 'subclass_color']\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 10), facecolor='w', edgecolor='k')\n",
    "    # plt.scatter(tsne1.values,tsne2.values,\n",
    "    #                 edgecolors='none', c=c, s=10, alpha=0.1)\n",
    "\n",
    "    # sbcls = set(df_aba_c.loc[in_data, 'subclass_label'])\n",
    "    # for tp in sbcls:\n",
    "    #     in_data2 = in_data&df_aba_c.loc[:,'subclass_label'].isin([tp])\n",
    "    #     tsne1tp = df_aba_tsne.loc[ df_aba_c.loc[in_data2, 'sample_name'], 'tsne_1']\n",
    "    #     tsne2tp = df_aba_tsne.loc[ df_aba_c.loc[in_data2, 'sample_name'], 'tsne_2']\n",
    "    #     plt.text(tsne1tp.median(),tsne2tp.median(),tp,fontsize=12)\n",
    "\n",
    "\n",
    "\n",
    "    X_train1 = X_train0.loc[X_train0.loc[:,'samples_pre'].isna()==False,:]\n",
    "    sampr = X_train1.loc[:,'samples_pre']\n",
    "\n",
    "    ##df_aba18_to_19 = pd.read_excel('aba18_in_aba19_presynaptic.xlsx')\n",
    "    #sampr19 = df_aba18_to_19_2.loc[df_aba18_to_19_2.loc[:,'sample_name'].isin(sampr),'sample_name_19']\n",
    "    df_aba18_to_19_3 = df_aba18_to_19_2.set_index('sample_name')\n",
    "    sampr19 = df_aba18_to_19_3.loc[sampr,'sample_name_19']\n",
    "    sampr190 = sampr19.copy()\n",
    "    sampr19 = sampr19.loc[sampr19.isna()==False ] \n",
    "\n",
    "\n",
    "\n",
    "    #sampr18 = df_aba18_to_19_2.loc[df_aba18_to_19_2.loc[:,'sample_name'].isin(sampr),'cluster']\n",
    "    #sampr19 = df_cl18_to_cl19[sampr18]\n",
    "    #sampr19 = \n",
    "    #### restart here!!!\n",
    "\n",
    "    tsne1pr = df_aba_tsne.loc[ sampr19, 'tsne_1']\n",
    "    tsne2pr = df_aba_tsne.loc[ sampr19, 'tsne_2']\n",
    "    plt.scatter(tsne1pr.values,tsne2pr.values,\n",
    "                    edgecolors='none', c='gray', s=15, alpha=0.4)\n",
    "    #X_train0.loc[:,annot_columns_train + list(ge_npr) + cla_n]\n",
    "\n",
    "\n",
    "    X_train2 = X_train0.loc[X_train0.loc[:,'samples_post'].isna()==False,:]\n",
    "    sampo = X_train2.loc[:,'samples_post']\n",
    "\n",
    "\n",
    "    ##df_aba18_to_19 = pd.read_excel('aba18_in_aba19_postsynaptic.xlsx')\n",
    "    #sampo19 = df_aba18_to_19_2.loc[df_aba18_to_19_2.loc[:,'sample_name'].isin(sampo),'sample_name_19']\n",
    "    sampo19 = df_aba18_to_19_3.loc[sampo,'sample_name_19']\n",
    "    sampo190 = sampo19.copy()\n",
    "    sampo19 = sampo19.loc[sampo19.isna()==False ] \n",
    "\n",
    "    tsne1po = df_aba_tsne.loc[ sampo19, 'tsne_1']\n",
    "    tsne2po = df_aba_tsne.loc[ sampo19, 'tsne_2']\n",
    "    plt.scatter(tsne1po.values,tsne2po.values,\n",
    "                    edgecolors='none', c='gray', s=15, alpha=0.4)\n",
    "\n",
    "\n",
    "    #### SYNAPSES\n",
    "    do_show_all_synapses=0\n",
    "    if do_show_all_synapses==1:\n",
    "        x = pd.concat([tsne1pr,tsne1po]).T\n",
    "        y = pd.concat([tsne2pr,tsne2po]).T\n",
    "        plt.plot(x.values,y.values, color=gray, linewidth=0.2, alpha=0.5)\n",
    "\n",
    "    # #### SYNAPSES2\n",
    "    #do_laptop=0\n",
    "    if do_laptop:\n",
    "        d4 = '~/Documents/references/lists_of_proteins_2/Analysis_of_proteins/genetic_types/transcriptomic_interneurons/'\n",
    "    else:\n",
    "        d4='~/Documents/references/transcriptomes to STP/transcriptomes/ABA/mouse_VISp_gene_expression_matrices_2018-06-14/'\n",
    "\n",
    "    fn='mouse_VISp_2018-06-14_samples-columns.csv'\n",
    "    df_aba_vis_c=pd.read_csv(d4+fn)\n",
    "\n",
    "    in_data = df_aba_vis_c.loc[:,'class'].isin(['GABAergic', 'Glutamatergic'])\n",
    "    #in_data = in_data&df_aba_c.loc[:,'region_label'].isin(['VISp','ENTl','ENTm','HIP'])\n",
    "\n",
    "    # df_aba_c2 = df_aba_c.loc[in_data,:].set_index('sample_name')\n",
    "    # df_aba_c3 = df_aba_c2.loc[sampr19,:]\n",
    "\n",
    "    # df_aba18_to_19_4 = df_aba18_to_19_2.set_index('sample_name_19')\n",
    "    # X_train0pr = X_train0.set_index('samples_pre').loc[:,'samples_post']\n",
    "    # X_train0pr = df_aba18_to_19_3.loc[X_train0pr,'sample_name_19']\n",
    "    # X_train0pr = X_train0pr.loc[df_aba18_to_19_4.loc[sampr19,'sample_name']]\n",
    "    # X_train0pr.index =  sampr19\n",
    "\n",
    "    df_aba_vis_c2 = df_aba_vis_c.set_index('sample_name')\n",
    "    df_aba_c2  =df_aba_c.set_index('cluster_label')\n",
    "    #L5 IT_1\n",
    "    #L5 IT ALM Cpa6 Gpr88\n",
    "\n",
    "\n",
    "    clu1 = df_aba_vis_c2.loc[sampr, 'cluster']\n",
    "    clu1 = clu1[clu1.isna()==False]\n",
    "\n",
    "    dif = list(set(clu1).difference(set(df_cl18_to_cl19.index)))\n",
    "    dif19 = [ 'L6 IT_5','L6 IT_5','Sst_16','Vip_8','Vip_9','Vip_11']\n",
    "    ddd=pd.DataFrame([dif19,dif]).T.set_index(1)\n",
    "    df_cl18_to_cl19_2 = pd.concat([df_cl18_to_cl19,ddd]) \n",
    "    df_cl18_to_cl19_2.loc['L5 IT ALM Cpa6 Gpr88',0]='L5 IT_1'\n",
    "    df_cl18_to_cl19_2.loc['Vip Igfbp4 Mab21l1',0]='Vip_3'\n",
    "    df_cl18_to_cl19_2.loc['Vip Igfbp6 Car10',0]='Vip_3'\n",
    "    df_cl18_to_cl19_2.loc['Vip Rspo1 Itga4',0]='Vip_12'\n",
    "    df_cl18_to_cl19_2.loc['L6 IT ALM Tgfb1',0]='L6 IT_5'\n",
    "    df_cl18_to_cl19_2.loc['Vip Crispld2 Htr2c',0]='Vip_8'\n",
    "    df_cl18_to_cl19_2.loc['Sst Crh 4930553C11Rik ',0]='Sst_20'\n",
    "    df_cl18_to_cl19_2.loc['L6b Hsd17b2',0]='L6b_5'\n",
    "    #Vip Crispld2 Kcne4\n",
    "    #Sst Crh 4930553C11Rik \n",
    "    #L6b Hsd17b2\n",
    " \n",
    "\n",
    "\n",
    "\n",
    "    clu1_19 = df_cl18_to_cl19_2.loc[clu1,:]\n",
    "\n",
    "\n",
    "    sbcls1 = set(df_aba_c.loc[df_aba_c.loc[:,'cluster_label'].isin(clu1_19.loc[:,0]), 'subclass_label'])\n",
    "    clu1_19_ =clu1_19.drop_duplicates()\n",
    "    sbs119=df_aba_c2.loc[clu1_19_.loc[:,0], 'subclass_label'].reset_index().drop_duplicates(['cluster_label']).set_index('cluster_label')\n",
    "    sbs119 = sbs119.loc[clu1_19.loc[:,0]]\n",
    "\n",
    "\n",
    "    clu2 = df_aba_vis_c2.loc[sampo, 'cluster']\n",
    "    clu2 = clu2[clu2.isna()==False]\n",
    "    clu2_19 = df_cl18_to_cl19_2.loc[clu2,:]\n",
    "    sbcls2 = set(df_aba_c.loc[df_aba_c.loc[:,'cluster_label'].isin(clu2_19.loc[:,0]), 'subclass_label'])\n",
    "\n",
    "\n",
    "    clu2_19_ =clu2_19.drop_duplicates()\n",
    "    sbs219=df_aba_c2.loc[clu2_19_.loc[:,0], 'subclass_label'].reset_index().drop_duplicates(['cluster_label']).set_index('cluster_label')\n",
    "    sbs219 = sbs219.loc[clu2_19.loc[:,0]]\n",
    "\n",
    "    c_sb_18_19 = pd.concat([sampr.loc[0:10599],sampo.loc[0:10599],clu1_19.reset_index(),clu2_19.reset_index(),\n",
    "                            sbs119.reset_index(),sbs219.reset_index()],axis=1)\n",
    "    c_sb_18_19.columns = ['samples_pre', 'samples_post' ,'cluster_pre','cluster_pre_19'\n",
    "                          ,'cluster_post','cluster_post_19',\n",
    "                          'cluster_label_', 'subclass_label_pre_19','cluster_label_p','subclass_label_post_19']\n",
    "\n",
    "    # df_aba_c2 = df_aba_c.loc[in_data,:].set_index('sample_name')\n",
    "    # df_aba_c3 = df_aba_c2.loc[sampr,:]\n",
    "    # X_train0pr = X_train0.set_index('samples_pre')\n",
    "\n",
    "    # in_data31 = in_data&df_aba_c.loc[:,'sample_name'].isin(sampr)\n",
    "    # sbcls1 = set(df_aba_c.loc[in_data31, 'subclass_label'])\n",
    "\n",
    "    # in_data32 = in_data&df_aba_c.loc[:,'sample_name'].isin(sampo)\n",
    "    # sbcls2 = set(df_aba_c.loc[in_data32, 'subclass_label'])\n",
    "\n",
    "    Nbootstrap = 200\n",
    "    Nst = int(c_sb_18_19.shape[0]/Nbootstrap)\n",
    "    for ii in range(Nst):\n",
    "        stii = c_sb_18_19.loc[ii*Nbootstrap:(ii+1)*Nbootstrap-1,:]\n",
    "        scstpr  = stii.loc[:,'subclass_label_pre_19'].value_counts().index[0]\n",
    "        scstpo  = stii.loc[:,'subclass_label_post_19'].value_counts().index[0]\n",
    "\n",
    "        in_data21 = df_aba_c.loc[:,'subclass_label'].isin([scstpr])\n",
    "        tsne1tp1 = df_aba_tsne.loc[ df_aba_c.loc[in_data21, 'sample_name'], 'tsne_1']\n",
    "        tsne2tp1 = df_aba_tsne.loc[ df_aba_c.loc[in_data21, 'sample_name'], 'tsne_2']\n",
    "\n",
    "        in_data22 = df_aba_c.loc[:,'subclass_label'].isin([scstpo])\n",
    "        tsne1tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_1']\n",
    "        tsne2tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_2']\n",
    "\n",
    "        if df_aba_c.loc[in_data21, 'class_label'].iloc[0]=='GABAergic':\n",
    "            co1='k'\n",
    "            lw1 = 5\n",
    "            is_gaba = True\n",
    "            \n",
    "        else:\n",
    "            co1='k'\n",
    "            lw1 = 5\n",
    "            is_gaba = False\n",
    "\n",
    "        if fff>1:\n",
    "            #col = colas[df_Stp_type_p_c.loc[ii,'type'].astype(int)-1,:]  \n",
    "            stype=df_Stp_type_p_c.loc[ii,'type'] #.astype(int)\n",
    "            if stype==1:\n",
    "                co1='y'\n",
    "            if stype==2:\n",
    "                co1='b'\n",
    "            if stype==3:\n",
    "                col='g'\n",
    "            #print(ii,df_Stp_type_p_c.loc[ii,:],col)\n",
    "        \n",
    "        x = np.array([tsne1tp1.median(),tsne1tp2.median()])\n",
    "        y = np.array([tsne2tp1.median(),tsne2tp2.median()])\n",
    "        if (is_gaba == plot_gaba):\n",
    "            #print(col)\n",
    "            plt.plot(x,y, color=co1, linewidth=lw1, alpha=0.2)\n",
    "            #plt.plot(x,y, color=co1, linewidth=lw1)\n",
    "            \n",
    "         \n",
    "        #plt.text(tsne1tp.median(),tsne2tp.median(),tp,fontsize=18)    \n",
    "\n",
    "\n",
    "    # for tp1 in sbcls1:\n",
    "    #     in_data21 = df_aba_c.loc[:,'subclass_label'].isin([tp1])\n",
    "    #     tsne1tp1 = df_aba_tsne.loc[ df_aba_c.loc[in_data21, 'sample_name'], 'tsne_1']\n",
    "    #     tsne2tp1 = df_aba_tsne.loc[ df_aba_c.loc[in_data21, 'sample_name'], 'tsne_2']\n",
    "    #     n1 = tsne1tp1.shape[0]\n",
    "\n",
    "    #     if df_aba_c.loc[in_data21, 'class_label'].iloc[0]=='GABAergic':\n",
    "    #         co1='g'\n",
    "    #         lw1 = 8\n",
    "    #     else:\n",
    "    #         co1='r'\n",
    "    #         lw1 = 4\n",
    "\n",
    "    #     df_aba_c31 = c_sb_18_19.loc[c_sb_18_19.loc[:,'subclass_label_pre_19'].isin([tp1]),:]\n",
    "\n",
    "    #     #df_aba_c31po = df_aba_c31.loc[:,'samples_post']\n",
    "    #     n30=df_aba_c31.shape[0]\n",
    "    #     #print(tp1 + '  '+str([n30, n1]))\n",
    "    #     if n30>0:\n",
    "    #         for tp2 in sbcls2:\n",
    "    #             in_data22 = df_aba_c.loc[:,'subclass_label'].isin([tp2])\n",
    "    #             tsne1tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_1']\n",
    "    #             tsne2tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_2']\n",
    "    #             n2 = tsne1tp2.shape[0]\n",
    "\n",
    "\n",
    "\n",
    "    #             #df_aba_c32 = df_aba_c2.loc[df_aba_c31po,:]\n",
    "    #             df_aba_c32tp2 = df_aba_c31.loc[df_aba_c31.loc[:,'subclass_label_post_19'].isin([tp2]),:]\n",
    "    #             n3 = df_aba_c32tp2.shape[0]\n",
    "    #             #print(tp1 + '  '+tp2+'  ' +str([n3, n30, n1, n2]))\n",
    "    #             if n3>0:\n",
    "    #                 x = np.array([tsne1tp1.median(),tsne1tp2.median()])\n",
    "    #                 y = np.array([tsne2tp1.median(),tsne2tp2.median()])\n",
    "    #                 plt.plot(x,y, color=co1, linewidth=lw1, alpha=0.2)\n",
    "    #                 #plt.text(tsne1tp.median(),tsne2tp.median(),tp,fontsize=18)\n",
    "\n",
    "\n",
    "    #### SYNAPSES3\n",
    "    #import plotly.graph_objects as go\n",
    "\n",
    "\n",
    "    do_synapses3 = 1\n",
    "    label =[]\n",
    "    x = []\n",
    "    y = []\n",
    "    source = []\n",
    "    target =[]\n",
    "    value = []\n",
    "    nl=0\n",
    "    if do_synapses3==1:\n",
    "        Nbootstrap = 200\n",
    "        Nst = int(c_sb_18_19.shape[0]/200)\n",
    "        for ii in range(Nst):\n",
    "            stii = c_sb_18_19.loc[ii*Nbootstrap:(ii+1)*Nbootstrap-1,:]\n",
    "            scstpr  = stii.loc[:,'subclass_label_pre_19'].value_counts().index[0]\n",
    "            scstpo  = stii.loc[:,'subclass_label_post_19'].value_counts().index[0]\n",
    "\n",
    "            in_data21 = df_aba_c.loc[:,'subclass_label'].isin([scstpr])\n",
    "            tsne1tp1 = df_aba_tsne.loc[ df_aba_c.loc[in_data21, 'sample_name'], 'tsne_1']\n",
    "            tsne2tp1 = df_aba_tsne.loc[ df_aba_c.loc[in_data21, 'sample_name'], 'tsne_2']\n",
    "\n",
    "            in_data22 = df_aba_c.loc[:,'subclass_label'].isin([scstpo])\n",
    "            tsne1tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_1']\n",
    "            tsne2tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_2']\n",
    "\n",
    "            if df_aba_c.loc[in_data21, 'class_label'].iloc[0]=='GABAergic':\n",
    "                co1='k'\n",
    "                lw1 = 8\n",
    "                is_gaba = True\n",
    "            else:\n",
    "                co1='k'\n",
    "                lw1 = 4\n",
    "                is_gaba = False\n",
    "\n",
    "            #x = np.array([tsne1tp1.median(),tsne1tp2.median()])\n",
    "            #y = np.array([tsne2tp1.median(),tsne2tp2.median()])\n",
    "            #plt.plot(x,y, color=co1, linewidth=lw1, alpha=0.4)\n",
    "            ##plt.text(tsne1tp.median(),tsne2tp.median(),tp,fontsize=18)  \n",
    "\n",
    "            tp1 = scstpr\n",
    "            tp2 = scstpo\n",
    "            if len(label)==0:\n",
    "\n",
    "                fl1=True\n",
    "            else:\n",
    "                fl1 = sum(np.array(label)==tp1)==0\n",
    "\n",
    "            if fl1:\n",
    "                nl=nl+1\n",
    "                nl1 = np.nonzero(np.array(label)==tp1)[0]\n",
    "                label = label + [tp1]\n",
    "                x = x + [tsne1tp1.median()]\n",
    "                y = y + [tsne2tp1.median()]\n",
    "\n",
    "\n",
    "            if sum(np.array(label)==tp2)==0:\n",
    "                nl=nl+1\n",
    "                nl2 = np.nonzero(np.array(label)==tp2)[0]\n",
    "                label = label + [tp2]\n",
    "                x = x + [ tsne1tp2.median()]\n",
    "                y = y + [ tsne2tp2.median()]\n",
    "\n",
    "            nl1 = np.nonzero(np.array(label)==tp1)[0]\n",
    "            nl2 = np.nonzero(np.array(label)==tp2)[0]\n",
    "\n",
    "            if x[nl1[0]]<=x[nl2[0]]:\n",
    "                source = source + [nl1[0]]\n",
    "                target = target + [nl2[0]]\n",
    "            else:\n",
    "                source = source + [nl2[0]]\n",
    "                target = target + [nl1[0]]\n",
    "            value = value + [0.5]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        x = np.array(x)\n",
    "        x = (x - x.min())\n",
    "        x = x/x.max()/1.1 + 0.01\n",
    "        x=x.tolist()\n",
    "        y = np.array(y)\n",
    "        y = (y - y.min())\n",
    "        y = y/y.max()/1.1 + 0.01\n",
    "        y=y.tolist()\n",
    "\n",
    "    #     fig = go.Figure()\n",
    "    #     fig.add_trace(go.Sankey(\n",
    "    #         arrangement = \"snap\",\n",
    "    #         node = dict(\n",
    "    #             label= label,\n",
    "    #              x= x,\n",
    "    #              y= y,\n",
    "    #              pad = 0.1,\n",
    "    #              thickness = 10,\n",
    "    #              line = dict(color = \"black\", width = 0.1)),  # 10 Pixels \n",
    "    #         link = dict(\n",
    "    #             source= source,\n",
    "    #             target= target,\n",
    "    #             value= value)))\n",
    "    #     fig.update_layout( font_size=20)\n",
    "    #     fig.show()\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    #### TSNE \n",
    "    in_data = df_aba_c.loc[:,'class_label'].isin(['GABAergic', 'Glutamatergic'])\n",
    "    in_data = in_data&df_aba_c.loc[:,'region_label'].isin(['VISp','ENTl','ENTm','HIP'])\n",
    "    tsne1 = df_aba_tsne.loc[ df_aba_c.loc[in_data, 'sample_name'], 'tsne_1']\n",
    "    tsne2 = df_aba_tsne.loc[ df_aba_c.loc[in_data, 'sample_name'], 'tsne_2']\n",
    "    c = df_aba_c.loc[in_data, 'subclass_color']\n",
    "    plt.scatter(tsne1.values,tsne2.values,\n",
    "                    edgecolors='none', c=c, s=7, alpha=0.1)\n",
    "\n",
    "    #### Subclasses names\n",
    "    in_data3 = in_data&df_aba_c.loc[:,'sample_name'].isin(sampr19.append(sampo19))\n",
    "    sbcls = set(df_aba_c.loc[in_data3, 'subclass_label'])\n",
    "    for tp in sbcls:\n",
    "        dy=0\n",
    "        if tp=='Meis2':\n",
    "            dy=-2\n",
    "        dx=0\n",
    "        if tp=='Pvalb':\n",
    "            dx=+7\n",
    "        if tp=='Sst':\n",
    "            dx=+2   \n",
    "        if tp=='L4 IT':\n",
    "            dx=+4     \n",
    "        in_data2 = in_data&df_aba_c.loc[:,'subclass_label'].isin([tp])\n",
    "        tsne1tp = df_aba_tsne.loc[ df_aba_c.loc[in_data2, 'sample_name'], 'tsne_1']\n",
    "        tsne2tp = df_aba_tsne.loc[ df_aba_c.loc[in_data2, 'sample_name'], 'tsne_2']\n",
    "        plt.text(tsne1tp.median()+dx,tsne2tp.median()+dy,tp,fontsize=22,c=np.array([50, 50, 50])/255)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    plt.axis('off')\n",
    "    #fig.set_tight_layout(True)\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "d_aba2020= '/Users/stepaniu/Documents/jan_2020/ABA_2020/'\n",
    "df_meta20 = pd.read_csv(d_aba2020+'metadata.csv')\n",
    "df_meta20.to_hdf('metadata_aba2020.hdf',key='data') #\n",
    "df_meta20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ABA 2020 10x from site\n",
    "\n",
    "d_aba2020= '/Users/stepaniu/Documents/jan_2020/ABA_2020/'\n",
    "df_tsne20 = pd.read_csv(d_aba2020+'tsne.csv')\n",
    "df_meta20=pd.read_hdf('metadata_aba2020.hdf')\n",
    "df_meta20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df_meta20.loc[:10,:].to_excel('meta_2020.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_meta20.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(set(df_meta20.loc[:,'platform_label']))\n",
    "print(df_meta20.shape)\n",
    "print(set(df_meta20.loc[:,'cortical_layer_label']))\n",
    "print(set(df_meta20.loc[:,'region_label']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ABA 2020 10x from U19 Zeng site http://data.nemoarchive.org/biccn/lab/zeng/transcriptome/scell/10X/processed/YaoHippo2020/\n",
    "# 1093036 cells CTX and HIPP without layers\n",
    "d_aba2020= '/Users/stepaniu/Documents/jan_2020/ABA_2020/Zeng_Tasic_2020_10x_CTX_HIPP/'\n",
    "df_tsne20x = pd.read_csv(d_aba2020+'CTX_Hip_umap.2d.10x.csv')\n",
    "#df_meta20=pd.read_hdf('metadata_aba2020.hdf')\n",
    "df_meta20x=pd.read_csv(d_aba2020 + 'CTX_Hip_anno_10x.v2.csv')\n",
    "df_meta20x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_meta20x.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(set(df_meta20x.loc[:,'platform_label']))\n",
    "print(df_meta20x.shape)\n",
    "print(set(df_meta20x.loc[:,'cortical_layer_label']))\n",
    "print(set(df_meta20x.loc[:,'region_label']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ABA 2020 Smart_seq from AIBS site - 74973 cells CTX and HIPP\n",
    "d_aba2020= '/Users/stepaniu/Documents/jan_2020/ABA_2020/Zeng_Tasic_YaoHippo2020_SmartSeq_CTX_HIPP/'\n",
    "df_tsne20ss = pd.read_csv(d_aba2020+'CTX_Hip_umap.2d.SS.csv')\n",
    "#df_meta20=pd.read_hdf('metadata_aba2020.hdf')\n",
    "df_meta20ss=pd.read_csv(d_aba2020 + 'CTX_Hip_anno_SS.v2.csv')\n",
    "df_meta20ss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_meta20ss.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(set(df_meta20ss.loc[:,'platform_label']))\n",
    "print(df_meta20ss.shape)\n",
    "print(set(df_meta20ss.loc[:,'cortical_layer_label']))\n",
    "print(set(df_meta20ss.loc[:,'region_label']))"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#X_train0.iloc[0:-1:200,:].loc[:,annot_columns].to_excel('temp.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "#   All synapses : hippocampus\n",
    "#\n",
    "#\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "gray = np.array([200, 200, 200, 1])/255\n",
    "\n",
    "for fff in range(4):\n",
    "    \n",
    "    plot_gaba = (fff==0)|(fff==2)\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 10), facecolor='w', edgecolor='k')\n",
    "\n",
    "    ### PRESYNAPTIC CELLS\n",
    "    X_train1 = X_train0.loc[X_train0.loc[:,'samples_pre'].isna()==False,:]\n",
    "    sampr = X_train1.loc[:,'samples_pre']\n",
    "\n",
    "    #df_aba18_to_19 = pd.read_excel('aba18_in_aba19_presynaptic.xlsx')\n",
    "    #sampr19 = df_aba18_to_19.loc[df_aba18_to_19.loc[:,'sample_name'].isin(sampr),'sample_name_19']\n",
    "\n",
    "    tsne1pr = df_aba_tsne.loc[ sampr, 'tsne_1']\n",
    "    tsne2pr = df_aba_tsne.loc[ sampr, 'tsne_2']\n",
    "    plt.scatter(tsne1pr.values,tsne2pr.values,\n",
    "                    edgecolors='none', c='gray', s=15, alpha=0.4)\n",
    "    #X_train0.loc[:,annot_columns_train + list(ge_npr) + cla_n]\n",
    "\n",
    "    ### POSTSYNAPTIC CELLS\n",
    "    X_train2 = X_train0.loc[X_train0.loc[:,'samples_post'].isna()==False,:]\n",
    "    sampo = X_train2.loc[:,'samples_post']\n",
    "\n",
    "    #df_aba18_to_19 = pd.read_excel('aba18_in_aba19_postsynaptic.xlsx')\n",
    "    #sampo19 = df_aba18_to_19.loc[df_aba18_to_19.loc[:,'sample_name'].isin(sampo),'sample_name_19']\n",
    "\n",
    "    tsne1po = df_aba_tsne.loc[ sampo, 'tsne_1']\n",
    "    tsne2po = df_aba_tsne.loc[ sampo, 'tsne_2']\n",
    "    plt.scatter(tsne1po.values,tsne2po.values,\n",
    "                    edgecolors='none', c='gray', s=4, alpha=0.4)\n",
    "\n",
    "\n",
    "    #### SYNAPSES\n",
    "    do_show_all_synapses=0\n",
    "    if do_show_all_synapses==1:\n",
    "        x = pd.concat([tsne1pr,tsne1po]).T\n",
    "        y = pd.concat([tsne2pr,tsne2po]).T\n",
    "        plt.plot(x.values,y.values, color=gray, linewidth=0.2, alpha=0.5)\n",
    "\n",
    "    #### SYNAPSES2\n",
    "    df_aba_c2 = df_aba_c.loc[in_data,:].set_index('sample_name')\n",
    "    df_aba_c3 = df_aba_c2.loc[sampr,:]\n",
    "    X_train0pr = X_train0.reset_index().set_index('samples_pre')\n",
    "\n",
    "    in_data31 = in_data&df_aba_c.loc[:,'sample_name'].isin(sampr)\n",
    "    sbcls1 = set(df_aba_c.loc[in_data31, 'subclass_label'])\n",
    "\n",
    "    in_data32 = in_data&df_aba_c.loc[:,'sample_name'].isin(sampo)\n",
    "    sbcls2 = set(df_aba_c.loc[in_data32, 'subclass_label'])\n",
    "    do_synapses2=1\n",
    "    if do_synapses2==1:\n",
    "        for tp1 in sbcls1:\n",
    "            in_data21 = in_data&df_aba_c.loc[:,'subclass_label'].isin([tp1])\n",
    "            tsne1tp1 = df_aba_tsne.loc[ df_aba_c.loc[in_data21, 'sample_name'], 'tsne_1']\n",
    "            tsne2tp1 = df_aba_tsne.loc[ df_aba_c.loc[in_data21, 'sample_name'], 'tsne_2']\n",
    "            n1 = tsne1tp1.shape[0]\n",
    "\n",
    "            if df_aba_c.loc[in_data21, 'class_label'].iloc[0]=='GABAergic':\n",
    "                co1='k'\n",
    "                lw1 = 5\n",
    "                is_gaba=True\n",
    "            else:\n",
    "                co1='k'\n",
    "                lw1 = 5\n",
    "                is_gaba=False\n",
    "                \n",
    "    \n",
    "\n",
    "            df_aba_c31 = df_aba_c3.loc[df_aba_c3.loc[:,'subclass_label'].isin([tp1]),:]\n",
    "            df_aba_c31po = X_train0pr.loc[df_aba_c31.index,'samples_post']\n",
    "            df_aba_c31poall = X_train0pr.loc[df_aba_c31.index,:]\n",
    "            n30=df_aba_c31po.shape[0]\n",
    "            #print(tp1 + '  '+str([n30, n1]))\n",
    "            if n30>0:\n",
    "                for tp2 in sbcls2:\n",
    "                    in_data22 = in_data&df_aba_c.loc[:,'subclass_label'].isin([tp2])\n",
    "                    tsne1tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_1']\n",
    "                    tsne2tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_2']\n",
    "                    n2 = tsne1tp2.shape[0]\n",
    "\n",
    "\n",
    "\n",
    "                    df_aba_c32 = df_aba_c2.loc[df_aba_c31po,:]\n",
    "                    df_aba_c32tp2 = df_aba_c32.loc[df_aba_c32.loc[:,'subclass_label'].isin([tp2]),:]\n",
    "                    n3 = df_aba_c32tp2.shape[0]\n",
    "                    #print(tp1 + '  '+tp2+'  ' +str([n3, n30, n1, n2]))\n",
    "                    #df_aba_c32tp2.index - all tp2 postsynaptic samples\n",
    "                    \n",
    "                    \n",
    "                    \n",
    "                    if n3>0:\n",
    "                        if fff>1:\n",
    "                            idx=df_aba_c31poall.loc[df_aba_c31po.isin(df_aba_c32tp2.index),'index']\n",
    "                            ii =int(idx.iloc[0]/Nbootstrap)\n",
    "                            #col = colas[df_Stp_type_p_c.loc[ii,'type'].astype(int)-1,:]  \n",
    "                            stype=df_Stp_type_p.loc[ii,'type'] #.astype(int)\n",
    "                            if stype==1:\n",
    "                                co1='y'\n",
    "                            if stype==2:\n",
    "                                co1='b'\n",
    "                            if stype==3:\n",
    "                                col='g'\n",
    "                            print(ii,df_Stp_type_p.loc[ii,:],col,tp1,tp2)\n",
    "                        \n",
    "                        x = np.array([tsne1tp1.median(),tsne1tp2.median()])\n",
    "                        y = np.array([tsne2tp1.median(),tsne2tp2.median()])\n",
    "                        if is_gaba==plot_gaba:\n",
    "                            plt.plot(x,y, color=co1, linewidth=lw1, alpha=0.2)\n",
    "                            #plt.plot(x,y, color=co1, linewidth=lw1)\n",
    "\n",
    "                        #plt.text(tsne1tp.median(),tsne2tp.median(),tp,fontsize=18)\n",
    "\n",
    "    #### SYNAPSES3\n",
    "    import plotly.graph_objects as go\n",
    "\n",
    "\n",
    "    do_synapses3 = 1\n",
    "    label =[]\n",
    "    x = []\n",
    "    y = []\n",
    "    source = []\n",
    "    target =[]\n",
    "    value = []\n",
    "    nl=0\n",
    "    if do_synapses3==1:\n",
    "        for tp1 in sbcls1:\n",
    "            in_data21 = in_data&df_aba_c.loc[:,'subclass_label'].isin([tp1])\n",
    "            tsne1tp1 = df_aba_tsne.loc[ df_aba_c.loc[in_data21, 'sample_name'], 'tsne_1']\n",
    "            tsne2tp1 = df_aba_tsne.loc[ df_aba_c.loc[in_data21, 'sample_name'], 'tsne_2']\n",
    "            n1 = tsne1tp1.shape[0]\n",
    "\n",
    "            if df_aba_c.loc[in_data21, 'class_label'].iloc[0]=='GABAergic':\n",
    "                co1='g'\n",
    "                lw1 = 8\n",
    "            else:\n",
    "                co1='r'\n",
    "                lw1 = 4\n",
    "                \n",
    "                \n",
    "\n",
    "            df_aba_c31 = df_aba_c3.loc[df_aba_c3.loc[:,'subclass_label'].isin([tp1]),:]\n",
    "            df_aba_c31po = X_train0pr.loc[df_aba_c31.index,'samples_post']\n",
    "            n30=df_aba_c31po.shape[0]\n",
    "            #print(tp1 + '  '+str([n30, n1]))\n",
    "            if n30>0:\n",
    "                for tp2 in sbcls2:\n",
    "                    in_data22 = in_data&df_aba_c.loc[:,'subclass_label'].isin([tp2])\n",
    "                    tsne1tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_1']\n",
    "                    tsne2tp2 = df_aba_tsne.loc[ df_aba_c.loc[in_data22, 'sample_name'], 'tsne_2']\n",
    "                    n2 = tsne1tp2.shape[0]\n",
    "\n",
    "\n",
    "\n",
    "                    df_aba_c32 = df_aba_c2.loc[df_aba_c31po,:]\n",
    "                    df_aba_c32tp2 = df_aba_c32.loc[df_aba_c32.loc[:,'subclass_label'].isin([tp2]),:]\n",
    "                    n3 = df_aba_c32tp2.shape[0]\n",
    "                    #print(tp1 + '  '+tp2+'  ' +str([n3, n30, n1, n2]))\n",
    "                    if n3>0:\n",
    "\n",
    "\n",
    "                        if len(label)==0:\n",
    "\n",
    "                            fl1=True\n",
    "                        else:\n",
    "                            fl1 = sum(np.array(label)==tp1)==0\n",
    "\n",
    "                        if fl1:\n",
    "                            nl=nl+1\n",
    "                            nl1 = np.nonzero(np.array(label)==tp1)[0]\n",
    "                            label = label + [tp1]\n",
    "                            x = x + [tsne1tp1.median()]\n",
    "                            y = y + [tsne2tp1.median()]\n",
    "\n",
    "\n",
    "                        if sum(np.array(label)==tp2)==0:\n",
    "                            nl=nl+1\n",
    "                            nl2 = np.nonzero(np.array(label)==tp2)[0]\n",
    "                            label = label + [tp2]\n",
    "                            x = x + [ tsne1tp2.median()]\n",
    "                            y = y + [ tsne2tp2.median()]\n",
    "\n",
    "                        nl1 = np.nonzero(np.array(label)==tp1)[0]\n",
    "                        nl2 = np.nonzero(np.array(label)==tp2)[0]\n",
    "\n",
    "                        if x[nl1[0]]<=x[nl2[0]]:\n",
    "                            source = source + [nl1[0]]\n",
    "                            target = target + [nl2[0]]\n",
    "                        else:\n",
    "                            source = source + [nl2[0]]\n",
    "                            target = target + [nl1[0]]\n",
    "                        value = value + [0.1]\n",
    "\n",
    "                        #x = np.array([tsne1tp1.median(),tsne1tp2.median()])\n",
    "                        #y = np.array([tsne2tp1.median(),tsne2tp2.median()])\n",
    "                        #plt.plot(x,y, color=co1, linewidth=lw1, alpha=0.2)\n",
    "                        ##plt.text(tsne1tp.median(),tsne2tp.median(),tp,fontsize=18)\n",
    "\n",
    "\n",
    "        x = np.array(x)\n",
    "        x = (x - x.min())\n",
    "        x = x/x.max()/1.0 + 0.01\n",
    "        x=x.tolist()\n",
    "        y = np.array(y)\n",
    "        y = (y - y.min())\n",
    "        y = y/y.max()/1.0 + 0.01\n",
    "        y=y.tolist()\n",
    "\n",
    "    #     fig = go.Figure()\n",
    "    #     fig.add_trace(go.Sankey(\n",
    "    #         arrangement = \"snap\",\n",
    "    #         node = dict(\n",
    "    #             label= label,\n",
    "    #              x= x,\n",
    "    #              y= y,\n",
    "    #              pad = 0.1,\n",
    "    #              thickness = 10,\n",
    "    #              line = dict(color = \"black\", width = 0.05)),  # 10 Pixels \n",
    "    #         link = dict(\n",
    "    #             source= source,\n",
    "    #             target= target,\n",
    "    #             value= value)))\n",
    "    #     fig.update_layout( font_size=20)\n",
    "    #     fig.show()\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    #### TSNE \n",
    "    in_data = df_aba_c.loc[:,'class_label'].isin(['GABAergic', 'Glutamatergic'])\n",
    "    in_data = in_data&df_aba_c.loc[:,'region_label'].isin(['VISp','ENTl','ENTm','HIP'])\n",
    "    tsne1 = df_aba_tsne.loc[ df_aba_c.loc[in_data, 'sample_name'], 'tsne_1']\n",
    "    tsne2 = df_aba_tsne.loc[ df_aba_c.loc[in_data, 'sample_name'], 'tsne_2']\n",
    "    c = df_aba_c.loc[in_data, 'subclass_color']\n",
    "    plt.scatter(tsne1.values,tsne2.values,\n",
    "                    edgecolors='none', c=c, s=7, alpha=0.1)\n",
    "\n",
    "    #### Subclasses names\n",
    "    in_data3 = in_data&df_aba_c.loc[:,'sample_name'].isin(sampr.append(sampo))\n",
    "    sbcls = set(df_aba_c.loc[in_data3, 'subclass_label'])\n",
    "    for tp in sbcls:\n",
    "        in_data2 = in_data&df_aba_c.loc[:,'subclass_label'].isin([tp])\n",
    "        tsne1tp = df_aba_tsne.loc[ df_aba_c.loc[in_data2, 'sample_name'], 'tsne_1']\n",
    "        tsne2tp = df_aba_tsne.loc[ df_aba_c.loc[in_data2, 'sample_name'], 'tsne_2']\n",
    "        plt.text(tsne1tp.median(),tsne2tp.median(),tp,fontsize=22)\n",
    "\n",
    "\n",
    "\n",
    "    plt.axis('off')\n",
    "    #fig.set_tight_layout(True)\n"
   ]
  },
  {
   "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": [
    "import plotly.graph_objects as go\n",
    "\n",
    "fig = go.Figure(go.Sankey(\n",
    "    arrangement = \"snap\",\n",
    "    node = {\n",
    "        \"label\": [\"A\", \"B\", \"C\", \"D\", \"E\", \"F\"],\n",
    "        \"x\": [0.2, 0.1, 0.5, 0.7, 0.3, 0.5],\n",
    "        \"y\": [0.7, 0.5, 0.2, 0.4, 0.2, 0.3],\n",
    "        'pad':10},  # 10 Pixels \n",
    "    link = {\n",
    "        \"source\": [0, 0, 1, 2, 5, 4, 3, 5],\n",
    "        \"target\": [5, 3, 4, 3, 0, 2, 2, 3],\n",
    "        \"value\": [1, 2, 1, 1, 1, 1, 1, 2]}))\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train0pr = X_train0.reset_index().set_index('samples_pre')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train0pr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " df_aba_c.loc[in_data22, 'class_label'].iloc[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sbcls2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yy=y_train[:,nannot:]\n",
    "y_u = UMAP(spread=2).fit_transform(yy)\n",
    "\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "y_u2 = TSNE(n_components=2).fit_transform(yy)\n",
    "\n",
    "fig = plt.figure(figsize=(10, 10), facecolor='w', edgecolor='k')\n",
    "plt.scatter(y_u[:, 0], y_u[:, 1],\n",
    "                edgecolors='none', c=\"#ff4400\", s=10)\n",
    "\n",
    "plt.axis('off')\n",
    "fig.set_tight_layout(True)\n",
    "\n",
    "fig1 = plt.figure(figsize=(10, 10), facecolor='w', edgecolor='k')\n",
    "plt.scatter(y_u2[:, 0], y_u2[:, 1],\n",
    "                edgecolors='none', c=\"#ff4400\", s=10)\n",
    "\n",
    "plt.axis('off')\n",
    "fig1.set_tight_layout(True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to do:\n",
    "#\n",
    "#\n",
    "# 1) visualize scVI latent space? - a. distances betwean synapses types b. paga? \n",
    "#\n",
    "#\n",
    "# 2) kernels - rbf will not be good?\n",
    "# train SVR in scVI space - a) factors b) imputed ge\n",
    "#\n",
    "# 3) RF with distance penalty?\n",
    "#\n",
    "# 4) test other classifiers in scVI space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DISTANCES\n",
    "from sklearn import metrics\n",
    "\n",
    "##lyamba_post = 0.5\n",
    "# ['all', 'ex_inh']\n",
    "Dn = 5\n",
    "Dn2 = int(100/Dn)\n",
    "#cla_n = ['ex_in_ex_in__pre',\n",
    "#         'ex_in_ex_in__post'] \n",
    "\n",
    "#cla_n = ['ex_in_ex_in__pre'] \n",
    "cla_n = ['vse__pre','vse__post']\n",
    "\n",
    "#ge_n = imp50.index[0:25].tolist()\n",
    "ge_n = lf_scvi_names #ge_columns_train2\n",
    "ge_n2=np.array(ge_n)\n",
    "ge_npr = ge_n2[np.char.rfind(ge_n2 , 'pre__')>-1]\n",
    "ge_nps = ge_n2[np.char.rfind(ge_n2 , 'post__')>-1]\n",
    "\n",
    "nannot = len(annot_columns_train)\n",
    "lge_n  = len(ge_n)\n",
    "lge_npr  = len(ge_npr)\n",
    "lge_nps  = len(ge_nps)\n",
    "\n",
    "\n",
    "X2pr = X_train0.loc[:,annot_columns_train + list(ge_npr) + cla_n]\n",
    "X2ps = X_train0.loc[:,annot_columns_train + list(ge_nps) + cla_n]\n",
    "\n",
    "if len(cla_n)>0:\n",
    "    i_clpr = np.nonzero(X2pr.columns.isin(cla_n))[0]  - nannot\n",
    "    i_clps = np.nonzero(X2pr.columns.isin(cla_n))[0]  - nannot\n",
    "    i_cl2pr = np.nonzero(X2pr.columns.isin(cla_n))[0]  \n",
    "    i_cl2ps = np.nonzero(X2ps.columns.isin(cla_n))[0]  \n",
    "\n",
    "    #X2.loc[:,cla_n] = X2.loc[:,cla_n]+1\n",
    "\n",
    "X2pr=X2pr.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "X2ps=X2ps.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "\n",
    "\n",
    "# select y\n",
    "stp_n = ['A2_20Hz']\n",
    "\n",
    "# 'A2_20Hz',\n",
    "#  'A3_20Hz',\n",
    "#  'A4_20Hz',\n",
    "#  'A5_20Hz',\n",
    "#  'A250_20Hz',\n",
    "#  'A1000_20Hz',\n",
    "#  'A1_50Hz',\n",
    "#  'A2_50Hz',\n",
    "#  'A3_50Hz',\n",
    "#  'A4_50Hz',\n",
    "#  'A5_50Hz',\n",
    "#  'A250_50Hz',\n",
    "#  'A1000_50Hz',\n",
    "#  'A1_10Hz',\n",
    "#  'A2_10Hz',\n",
    "#  'A3_10Hz',\n",
    "#  'A4_10Hz',\n",
    "#  'A5_10Hz',\n",
    "#  'A250_10Hz',\n",
    "#  'A1000_10Hz']\n",
    "y2 =  X_train0.loc[:,annot_columns_train+stp_n].iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "\n",
    "i_Xpr = np.arange(nannot,nannot+lge_npr)\n",
    "i_Xps = np.arange(nannot,nannot+lge_nps)\n",
    "\n",
    "#[‘braycurtis’, ‘canberra’, ‘chebyshev’, ‘correlation’, ‘dice’, ‘hamming’, ‘jaccard’,\n",
    "# ‘kulsinski’, ‘mahalanobis’, ‘minkowski’, ‘rogerstanimoto’, ‘russellrao’, ‘seuclidean’,\n",
    "# ‘sokalmichener’, ‘sokalsneath’, ‘sqeuclidean’, ‘yule’] \n",
    "\n",
    "Dpr = metrics.pairwise_distances(X2pr[:,i_Xpr], metric='sqeuclidean', n_jobs=-1, force_all_finite=True)\n",
    "Dps = metrics.pairwise_distances(X2ps[:,i_Xps], metric='sqeuclidean', n_jobs=-1, force_all_finite=True)\n",
    "\n",
    "X2cl = X_train0.loc[:,cla_n]\n",
    "X2cl = X2cl.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "D_mask = np.ones(Dpr.shape)\n",
    "#D_mask[]\n",
    "\n",
    "alp=16\n",
    "D = D_mask*np.sqrt(np.sqrt(alp*Dpr**4 + Dps**4))\n",
    "D = pd.DataFrame(D, columns=X2pr[:,i_cl2pr], index = X2pr[:,i_cl2pr] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 15))\n",
    "im = ax.imshow(-D)\n",
    "\n",
    "# We want to show all ticks...\n",
    "xD = np.arange(0,len(D.columns),Dn2)\n",
    "yD = np.arange(0,len(D.index),Dn2)\n",
    "ax.set_xticks(xD)\n",
    "ax.set_yticks(yD)\n",
    "# ... and label them with the respective list entries\n",
    "ax.set_xticklabels(D.columns[xD])\n",
    "ax.set_yticklabels(D.index[yD])\n",
    "\n",
    "# Rotate the tick labels and set their alignment.\n",
    "plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "         rotation_mode=\"anchor\")\n",
    "\n",
    "# Loop over data dimensions and create text annotations.\n",
    "#for i in range(len(hip2aba.index)):\n",
    "#    for j in range(len(hip2aba.columns)):\n",
    "#       text = ax.text(j, i, hip2aba[i, j],\n",
    "#                       ha=\"center\", va=\"center\", color=\"w\")\n",
    "\n",
    "ax.set_title(\"scVI latent factors based distances between synapses types\")\n",
    "fig.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### gs = '5188'\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 7))\n",
    "yD2 = np.arange(0,1600,400)\n",
    "yD2n = D.columns[yD2]\n",
    "plt.plot(np.arange(D.shape[0]), D.iloc[:,yD2],figure=fig)\n",
    "# We want to show all ticks...\n",
    "xD = np.arange(0,len(D.columns),Dn2)\n",
    "\n",
    "ax.set_xticks(xD)\n",
    "\n",
    "# ... and label them with the respective list entries\n",
    "ax.set_xticklabels(D.columns[xD])\n",
    "\n",
    "\n",
    "# Rotate the tick labels and set their alignment.\n",
    "plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "         rotation_mode=\"anchor\")\n",
    "\n",
    "ax.legend(yD2n)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### gs = '1512'\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 7))\n",
    "yD2 = np.arange(0,1600,400)\n",
    "yD2n = D.columns[yD2]\n",
    "plt.plot(np.arange(D.shape[0]), D.iloc[:,yD2],figure=fig)\n",
    "# We want to show all ticks...\n",
    "xD = np.arange(0,len(D.columns),Dn2)\n",
    "\n",
    "ax.set_xticks(xD)\n",
    "\n",
    "# ... and label them with the respective list entries\n",
    "ax.set_xticklabels(D.columns[xD])\n",
    "\n",
    "\n",
    "# Rotate the tick labels and set their alignment.\n",
    "plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "         rotation_mode=\"anchor\")\n",
    "\n",
    "ax.legend(yD2n)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### gs = '219'\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 7))\n",
    "yD2 = np.arange(0,1600,400)\n",
    "yD2n = D.columns[yD2]\n",
    "plt.plot(np.arange(D.shape[0]), D.iloc[:,yD2],figure=fig)\n",
    "# We want to show all ticks...\n",
    "xD = np.arange(0,len(D.columns),Dn2)\n",
    "\n",
    "ax.set_xticks(xD)\n",
    "\n",
    "# ... and label them with the respective list entries\n",
    "ax.set_xticklabels(D.columns[xD])\n",
    "\n",
    "\n",
    "# Rotate the tick labels and set their alignment.\n",
    "plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "         rotation_mode=\"anchor\")\n",
    "\n",
    "ax.legend(yD2n)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "cols=matplotlib.colors.ColorConverter.to_rgba_array(['b','r','#467821','0.4','m','y','purple',\n",
    "                                                'cyan','#B0E0E6','#8C0900','k','#30a2da','#6d904f','#ccebc4'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cols[2,:].reshape((1,-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### gs='15656'\n",
    "\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "# PCA transformation\n",
    "pca=PCA(n_components=10)\n",
    "#X2pr_pca =pca.fit_transform(X2pr[:,i_Xpr])\n",
    "\n",
    "ge_n = lf_scvi_names  #lf_scvi_names #ge_columns_train2\n",
    "ge_n2=np.array(ge_n)\n",
    "ge_npr = ge_n2[np.char.rfind(ge_n2 , 'pre__')>-1]\n",
    "\n",
    "X2 = X_train0.loc[:,annot_columns_train + list(ge_npr) + cla_n].values\n",
    "i_X2 = np.arange(nannot,nannot+len(ge_npr))\n",
    "X2_pca =pca.fit_transform(X2[:,i_X2])\n",
    "\n",
    "print(pca.explained_variance_ratio_)\n",
    "print(pca.singular_values_) \n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,1],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,2],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,3],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "### gs='5188'\n",
    "\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "# PCA transformation\n",
    "pca=PCA(n_components=20)\n",
    "#X2pr_pca =pca.fit_transform(X2pr[:,i_Xpr])\n",
    "\n",
    "ge_n = lf_scvi_names  #lf_scvi_names #ge_columns_train2\n",
    "ge_n2=np.array(ge_n)\n",
    "ge_npr = ge_n2[np.char.rfind(ge_n2 , 'pre__')>-1]\n",
    "\n",
    "X2 = X_train0.loc[:,annot_columns_train + list(ge_npr) + cla_n].values\n",
    "i_X2 = np.arange(nannot,nannot+len(ge_npr))\n",
    "X2_pca =pca.fit_transform(X2[:,i_X2])\n",
    "\n",
    "print(pca.explained_variance_ratio_)\n",
    "print(pca.singular_values_) \n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,1],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,2],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,3],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "### gs='1512'\n",
    "\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "# PCA transformation\n",
    "pca=PCA(n_components=10)\n",
    "#X2pr_pca =pca.fit_transform(X2pr[:,i_Xpr])\n",
    "\n",
    "ge_n = lf_scvi_names  #lf_scvi_names #ge_columns_train2\n",
    "ge_n2=np.array(ge_n)\n",
    "ge_npr = ge_n2[np.char.rfind(ge_n2 , 'pre__')>-1]\n",
    "\n",
    "X2 = X_train0.loc[:,annot_columns_train + list(ge_npr) + cla_n].values\n",
    "i_X2 = np.arange(nannot,nannot+len(ge_npr))\n",
    "X2_pca =pca.fit_transform(X2[:,i_X2])\n",
    "\n",
    "print(pca.explained_variance_ratio_)\n",
    "print(pca.singular_values_) \n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,1],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,2],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,3],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "\n",
    "# PCA transformation\n",
    "pca=PCA(n_components=10)\n",
    "#X2pr_pca =pca.fit_transform(X2pr[:,i_Xpr])\n",
    "\n",
    "ge_n = lf_scvi_names  #lf_scvi_names #ge_columns_train2\n",
    "ge_n2=np.array(ge_n)\n",
    "ge_npr = ge_n2[np.char.rfind(ge_n2 , 'pre__')>-1]\n",
    "\n",
    "X2 = X_train0.loc[:,annot_columns_train + list(ge_npr) + cla_n].values\n",
    "i_X2 = np.arange(nannot,nannot+len(ge_npr))\n",
    "X2_pca =pca.fit_transform(X2[:,i_X2])\n",
    "\n",
    "print(pca.explained_variance_ratio_)\n",
    "print(pca.singular_values_) \n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,1],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,2],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,3],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "\n",
    "# PCA transformation\n",
    "pca=PCA(n_components=10)\n",
    "#X2pr_pca =pca.fit_transform(X2pr[:,i_Xpr])\n",
    "\n",
    "ge_n = lf_scvi_names  #lf_scvi_names #ge_columns_train2\n",
    "ge_n2=np.array(ge_n)\n",
    "ge_npr = ge_n2[np.char.rfind(ge_n2 , 'pre__')>-1]\n",
    "\n",
    "X2 = X_train0.loc[:,annot_columns_train + list(ge_npr) + cla_n].values\n",
    "i_X2 = np.arange(nannot,nannot+len(ge_npr))\n",
    "X2_pca =pca.fit_transform(X2[:,i_X2])\n",
    "\n",
    "print(pca.explained_variance_ratio_)\n",
    "print(pca.singular_values_) \n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,1],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,2],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,-2]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,-2]==cii,0],X2_pca[X2[:,-2]==cii,3],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "\n",
    "# PCA transformation\n",
    "pca=PCA(n_components=10)\n",
    "#X2pr_pca =pca.fit_transform(X2pr[:,i_Xpr])\n",
    "\n",
    "ge_n = lf_scvi_names  #lf_scvi_names #ge_columns_train2\n",
    "ge_n2=np.array(ge_n)\n",
    "ge_npr = ge_n2[np.char.rfind(ge_n2 , 'post__')>-1]\n",
    "\n",
    "X2 = X_train0.loc[:,annot_columns_train + list(ge_npr) + cla_n].values\n",
    "i_X2 = np.arange(nannot,nannot+len(ge_npr))\n",
    "X2_pca =pca.fit_transform(X2[:,i_X2])\n",
    "\n",
    "print(pca.explained_variance_ratio_)\n",
    "print(pca.singular_values_) \n",
    "\n",
    "i_X_post = -1\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,i_X_post]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,i_X_post]==cii,0],X2_pca[X2[:,i_X_post]==cii,1],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,i_X_post]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,i_X_post]==cii,0],X2_pca[X2[:,i_X_post]==cii,2],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,i_X_post]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,i_X_post]==cii,0],X2_pca[X2[:,i_X_post]==cii,3],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "\n",
    "# PCA transformation\n",
    "pca=PCA(n_components=10)\n",
    "#X2pr_pca =pca.fit_transform(X2pr[:,i_Xpr])\n",
    "\n",
    "ge_n = lf_scvi_names  #lf_scvi_names #ge_columns_train2\n",
    "ge_n2=np.array(ge_n)\n",
    "ge_npr = ge_n2[np.char.rfind(ge_n2 , 'post__')>-1]\n",
    "\n",
    "X2 = X_train0.loc[:,annot_columns_train + list(ge_npr) + cla_n].values\n",
    "i_X2 = np.arange(nannot,nannot+len(ge_npr))\n",
    "X2_pca =pca.fit_transform(X2[:,i_X2])\n",
    "\n",
    "print(pca.explained_variance_ratio_)\n",
    "print(pca.singular_values_) \n",
    "\n",
    "i_X_post = -1\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,i_X_post]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,i_X_post]==cii,0],X2_pca[X2[:,i_X_post]==cii,1],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,i_X_post]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,i_X_post]==cii,0],X2_pca[X2[:,i_X_post]==cii,2],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    \n",
    "\n",
    "fig1, ax1 = plt.subplots(figsize=(15, 10))\n",
    "\n",
    "allcl0 = list(set(X2[:,i_X_post]))\n",
    "for ii,cii in enumerate(allcl0): \n",
    "    plt.scatter(X2_pca[X2[:,i_X_post]==cii,0],X2_pca[X2[:,i_X_post]==cii,3],c=cols[ii,:].reshape((1,-1)))\n",
    "plt.legend(allcl0)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# STOP HERE!!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scanpy as sc\n",
    "\n",
    "sc.settings.verbosity = 3\n",
    "sc.settings.set_figure_params(dpi=80, color_map='Greys')\n",
    "sc.logging.print_versions()\n",
    "results_file = 'X_scvi.h5ad'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import anndata\n",
    "adt = anndata.AnnData(X2[:,i_X2],obs = X2[:,i_X_post])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#adt.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ge_n = ge_columns_train2  #lf_scvi_names #ge_columns_train2\n",
    "ge_n2=np.array(ge_n)\n",
    "ge_npr1 = ge_n2[np.char.rfind(ge_n2 , 'post__')>-1]\n",
    "\n",
    "X21 = X_train0.loc[:,annot_columns_train + list(ge_npr1) + cla_n].values\n",
    "i_X21 = np.arange(nannot,nannot+len(ge_npr1))\n",
    "#X2_pca =pca.fit_transform(X2[:,i_X2])\n",
    "\n",
    "\n",
    "adt1 = anndata.AnnData(X21[:,i_X21],obs = X21[:,i_X_post])\n",
    "\n",
    "#sc.pp.filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scanpy.api as sca\n",
    "sca.pp.filter(adt1,min_cells=3)\n",
    "#sc.pp.filter(adt1,min_genes=10)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfX2pr = pd.DataFrame(X2pr)\n",
    "dfX2pr = dfX2pr.set_index([dfX2pr.columns[-1]])\n",
    "dfX2pr.to_excel('clustering_scvi_pre.xlsx')\n",
    "dfX2pr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.cluster.hierarchy import dendrogram\n",
    "from sklearn.cluster import AgglomerativeClustering\n",
    "\n",
    "def plot_dendrogram(model, **kwargs):\n",
    "    # Create linkage matrix and then plot the dendrogram\n",
    "\n",
    "    # create the counts of samples under each node\n",
    "    counts = np.zeros(model.children_.shape[0])\n",
    "    n_samples = len(model.labels_)\n",
    "    for i, merge in enumerate(model.children_):\n",
    "        current_count = 0\n",
    "        for child_idx in merge:\n",
    "            if child_idx < n_samples:\n",
    "                current_count += 1  # leaf node\n",
    "            else:\n",
    "                current_count += counts[child_idx - n_samples]\n",
    "        counts[i] = current_count\n",
    "\n",
    "    linkage_matrix = np.column_stack([model.children_, model.distances_,\n",
    "                                      counts]).astype(float)\n",
    "\n",
    "    # Plot the corresponding dendrogram\n",
    "    dendrogram(linkage_matrix, **kwargs)\n",
    "    \n",
    "    return  linkage_matrix\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "clust = AgglomerativeClustering(n_clusters=None, affinity='euclidean', memory=None, connectivity=None,\n",
    "                                           compute_full_tree='auto', linkage='ward', distance_threshold=1)\n",
    "\n",
    "clupr = clust.fit(X2pr[:,i_Xpr])\n",
    "clups = clust.fit(X2ps[:,i_Xps])\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 9))\n",
    "plt.title('Hierarchical Clustering Dendrogram')\n",
    "# plot the top three levels of the dendrogram\n",
    "lima = plot_dendrogram(clupr, truncate_mode='level', p=6,ax=ax)\n",
    "plt.xlabel(\"Number of points in node (or index of point if no parenthesis).\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lima.shape\n",
    "clupr.children_ \n",
    "clupr.distances_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 15))\n",
    "im = ax.imshow(-D)\n",
    "\n",
    "# We want to show all ticks...\n",
    "xD = np.arange(0,len(D.columns),Dn2)\n",
    "yD = np.arange(0,len(D.index),Dn2)\n",
    "ax.set_xticks(xD)\n",
    "ax.set_yticks(yD)\n",
    "# ... and label them with the respective list entries\n",
    "ax.set_xticklabels(D.columns[xD])\n",
    "ax.set_yticklabels(D.index[yD])\n",
    "\n",
    "# Rotate the tick labels and set their alignment.\n",
    "plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "         rotation_mode=\"anchor\")\n",
    "\n",
    "# Loop over data dimensions and create text annotations.\n",
    "#for i in range(len(hip2aba.index)):\n",
    "#    for j in range(len(hip2aba.columns)):\n",
    "#       text = ax.text(j, i, hip2aba[i, j],\n",
    "#                       ha=\"center\", va=\"center\", color=\"w\")\n",
    "\n",
    "ax.set_title(\"scVI latent factors based distances between synapses types\")\n",
    "fig.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(15, 7))\n",
    "yD2 = np.arange(0,1600,400)\n",
    "yD2n = D.columns[yD2]\n",
    "plt.plot(np.arange(D.shape[0]), D.iloc[:,yD2],figure=fig)\n",
    "# We want to show all ticks...\n",
    "xD = np.arange(0,len(D.columns),Dn2)\n",
    "\n",
    "ax.set_xticks(xD)\n",
    "\n",
    "# ... and label them with the respective list entries\n",
    "ax.set_xticklabels(D.columns[xD])\n",
    "\n",
    "\n",
    "# Rotate the tick labels and set their alignment.\n",
    "plt.setp(ax.get_xticklabels(), rotation=90, ha=\"right\",\n",
    "         rotation_mode=\"anchor\")\n",
    "\n",
    "ax.legend(yD2n)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "Q: 3 tymes: e_dg, e_dg ???\n",
    "Q:     NGF2 ???    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# START HERE!!!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### test pyMC3 model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.patches import Ellipse\n",
    "import numpy as np\n",
    "def plot_ellipse(ax,ms,ss):\n",
    "    #Seaborn color palette\n",
    "    colors = [(0.12156862745098039, 0.4666666666666667, 0.7058823529411765),\n",
    " (1.0, 0.4980392156862745, 0.054901960784313725),\n",
    " (0.17254901960784313, 0.6274509803921569, 0.17254901960784313),\n",
    " (0.8392156862745098, 0.15294117647058825, 0.1568627450980392),\n",
    " (0.5803921568627451, 0.403921568627451, 0.7411764705882353),\n",
    " (0.5490196078431373, 0.33725490196078434, 0.29411764705882354),\n",
    " (0.8901960784313725, 0.4666666666666667, 0.7607843137254902),\n",
    " (0.4980392156862745, 0.4980392156862745, 0.4980392156862745),\n",
    " (0.7372549019607844, 0.7411764705882353, 0.13333333333333333),\n",
    " (0.09019607843137255, 0.7450980392156863, 0.8117647058823529)]\n",
    "    k = ms.shape[0]\n",
    "    for i in range(k):\n",
    "        colr = colors[i]\n",
    "        var, U = np.linalg.eig(ss[i])\n",
    "        angle = 180. / np.pi * np.arccos(np.abs(U[0, 0])) \n",
    "        #Set the angle right\n",
    "        if (var[0]>var[1] and ss[i][1,0]<0) or (var[1]>var[0] and ss[i][1,0]>0):\n",
    "            angle = -angle\n",
    "        e = Ellipse(ms[i,:], 2 * np.sqrt(5.991 * var[0]),\n",
    "                    2 * np.sqrt(5.991 * var[1]),\n",
    "                    angle=angle)\n",
    "        e.set_alpha(0.25)\n",
    "        e.set_facecolor(colr)\n",
    "        e.set_zorder(10);\n",
    "        ax.add_artist(e);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymc3 as pm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from theano.tensor.nlinalg import matrix_inverse\n",
    "#from ellipse import plot_ellipse\n",
    "\n",
    "# Seed the rng for exact reproducible results\n",
    "seed = 68492\n",
    "np.random.seed(seed)\n",
    "\n",
    "## Generate some data for three groups.\n",
    "# Means, variances, covariances, and proportions\n",
    "mus = np.array([[-4,2],[0,1],[6,-2]]) \n",
    "variance1   = [1,.4,1.5]\n",
    "variance2   = [1,.8,5  ]\n",
    "covariances = [.5,0, -1]\n",
    "ps = np.array([0.2, 0.5, 0.3])\n",
    "D = mus[0].shape[0]\n",
    "# Total amount of data\n",
    "N = 1000\n",
    "# Number of groups\n",
    "K = 3\n",
    "# Form covariance matrix for each group\n",
    "sigmas = [np.array([[variance1[i],covariances[i]],[covariances[i],variance2[i]]]) for i in range(K)]\n",
    "# Form group assignments\n",
    "zs = np.array([np.random.multinomial(1, ps) for _ in range(N)]).T\n",
    "xs = [z[:, np.newaxis] * np.random.multivariate_normal(m, s, size=N)\n",
    "      for z, m, s in zip(zs, mus, sigmas)]\n",
    "# Stack data into single array\n",
    "data = np.sum(np.dstack(xs), axis=2)\n",
    "\n",
    "## Plot them nicely\n",
    "# Prepare subplots\n",
    "fig, ax = plt.subplots(figsize=(8, 6))\n",
    "# First, scatter\n",
    "plt.scatter(data[:, 0], data[:, 1], c='g', alpha=0.5)\n",
    "plt.scatter(mus[0, 0], mus[0, 1], c='r', s=100)\n",
    "plt.scatter(mus[1, 0], mus[1, 1], c='b', s=100)\n",
    "plt.scatter(mus[2, 0], mus[2, 1], c='y', s=100)\n",
    "# Then, ellipses\n",
    "plot_ellipse(ax,mus,sigmas)\n",
    "ax.axis('equal')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "## Build model and sample\n",
    "# Number of iterations for sampler\n",
    "draws = 2000\n",
    "# Prepare lists of starting points for mu to prevent label-switching problem\n",
    "testvals = [[-2,-2],[0,0],[2,2]]\n",
    "import time\n",
    "t1=time.time()\n",
    "# Model structure\n",
    "with pm.Model() as mvgmm:\n",
    "\n",
    "\n",
    "    \n",
    "    # Prior over component weights\n",
    "    p = pm.Dirichlet('p', a=np.array([1.]*K))\n",
    "\n",
    "    # Prior over component means\n",
    "    mus = [pm.MvNormal('mu_%d' % i,\n",
    "                        mu=pm.floatX(np.zeros(D)),\n",
    "                        tau=pm.floatX(0.1 * np.eye(D)),\n",
    "                        shape=(D,),\n",
    "                        testval=pm.floatX(testvals[i]))\n",
    "           for i in range(K)]\n",
    "\n",
    "    # Cholesky decomposed LKJ prior over component covariance matrices\n",
    "    packed_L = [pm.LKJCholeskyCov('packed_L_%d' % i,\n",
    "                                  n=D,\n",
    "                                  eta=2.,\n",
    "                                  sd_dist=pm.HalfCauchy.dist(1))\n",
    "                for i in range(K)]\n",
    "\n",
    "    # Unpack packed_L into full array\n",
    "    L = [pm.expand_packed_triangular(D, packed_L[i])\n",
    "         for i in range(K)]\n",
    "\n",
    "    # Convert L to sigma and tau for convenience\n",
    "    sigma = [pm.Deterministic('sigma_%d' % i ,L[i].dot(L[i].T))\n",
    "             for i in range(K)]\n",
    "    tau = [pm.Deterministic('tau_%d' % i,matrix_inverse(sigma[i]))\n",
    "           for i in range(K)]\n",
    "\n",
    "    print(\"# Specify the likelihood\")\n",
    "    # Specify the likelihood\n",
    "    mvnl = [pm.MvNormal.dist(mu=mus[i],chol=L[i])  \n",
    "           for i in range(K)]    \n",
    "    Y_obs = pm.Mixture('Y_obs',w=p, comp_dists=mvnl,observed=data)  \n",
    "\n",
    "    print(\"# Start the sampler!\")\n",
    "    # Start the sampler!\n",
    "    trace = pm.sample(draws, step=pm.NUTS(), chains=4)\n",
    "    \n",
    "    \n",
    "t2=time.time()\n",
    "print('elapsed time ',str(t2-t1))\n",
    "\n",
    "## Plot traces\n",
    "pm.traceplot(trace, var_names=['p', 'mu_0','mu_1','mu_2','tau_0','tau_1','tau_2','sigma_0','sigma_1','sigma_2'])\n",
    "plt.show()\n",
    "\n",
    "## Posterior Predictive Checks  \n",
    "# Obtain posterior samples\n",
    "pp = pm.sample_posterior_predictive(model=mvgmm, trace=trace)\n",
    "\n",
    "\n",
    "\n",
    "#f, ax =plt.subplots(figsize=(15, 5))\n",
    "f, ax =plt.subplots()\n",
    "# Plot original data\n",
    "plt.hist(data[:,0], bins=20, alpha=0.5)\n",
    "# Plot posterior predictives on top of that \n",
    "plt.hist(pp['Y_obs'][7000,:,0], bins=20, alpha=0.5)\n",
    "# Add legend and axes labels\n",
    "plt.legend(['Data','Predictions'])\n",
    "plt.xlabel('Simulated values')\n",
    "plt.ylabel('Frequencies')\n",
    "plt.show() \n",
    "\n",
    "f1, ax1 =plt.subplots()\n",
    "# Plot original data\n",
    "plt.hist(data[:,1], bins=20, alpha=0.5)\n",
    "# Plot posterior predictives on top of that \n",
    "plt.hist(pp['Y_obs'][7000,:,1], bins=20, alpha=0.5)\n",
    "# Add legend and axes labels\n",
    "plt.legend(['Data','Predictions'])\n",
    "plt.xlabel('Simulated values')\n",
    "plt.ylabel('Frequencies')\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "with pm.Model() as hierarchical_model:\n",
    "    # global model parameters\n",
    "    α_μ_tmp = pm.Normal('α_μ_tmp', mu=0., sd=100)\n",
    "    α_σ_tmp = pm.HalfNormal('α_σ_tmp', 5.)\n",
    "    β_μ = pm.Normal('β_μ', mu=0., sd=100)\n",
    "    β_σ = pm.HalfNormal('β_σ', 5.)\n",
    "\n",
    "    # train type specific model parameters\n",
    "    α_tmp = pm.Normal('α_tmp', mu=α_μ_tmp, sd=α_σ_tmp, shape=n_train_types)  \n",
    "    # Intercept for each train type, distributed around train type mean \n",
    "    β = pm.Normal('β', mu=β_μ, sd=β_σ, shape=n_train_types)\n",
    "    # Model error\n",
    "    eps = pm.HalfCauchy('eps', 5.)\n",
    "\n",
    "    fare_est = α_tmp[train_type_idx] + β[train_type_idx]*data.fare_encode.values\n",
    "\n",
    "    # Data likelihood\n",
    "    fare_like = pm.Normal('fare_like', mu=fare_est, sd=eps, observed=data.price)\n",
    "    \n",
    "with hierarchical_model:\n",
    "    hierarchical_trace = pm.sample(2000, tune=2000, target_accept=.9)\n",
    "    \n",
    "pm.traceplot(hierarchical_trace, var_names=['α_μ_tmp', 'β_μ', 'α_σ_tmp', 'β_σ', 'eps']);\n",
    "\n",
    "### GOOD MULTIDIMENSIONAL EXAMPLE\n",
    "### FROM https://github.com/junpenglao/GLMM-in-Python/blob/master/misc/pymc3_hierarchical_example.py\n",
    "\n",
    "model = pm.Model()\n",
    "with model:\n",
    "\n",
    "    # m_g ~ N(0, .1)\n",
    "    group_effects = pm.Normal(\n",
    "        \"group_effects\", 0, .1, shape=(n_group_predictors, n_predictors))\n",
    "    gp = pm.Normal(\"gp\", 0, .1, shape=(n_groups,1))\n",
    "    # gp = group_predictors\n",
    "    # sg ~ Uniform(.05, 10)\n",
    "    sg = pm.Uniform(\"sg\", .05, 10, testval=2.)\n",
    "    \n",
    "\n",
    "    # m ~ N(mg * pg, sg)\n",
    "    effects = pm.Normal(\"effects\",\n",
    "                     T.dot(gp, group_effects), sg ** -2,\n",
    "                     shape=(n_groups, n_predictors))\n",
    "\n",
    "    s = pm.Uniform(\"s\", .01, 10, shape=n_groups)\n",
    "\n",
    "    g = T.constant(group)\n",
    "\n",
    "    # y ~ Normal(m[g] * p, s)\n",
    "    mu_est = pm.Deterministic(\"mu_est\",T.sum(effects[g] * predictors, 1))\n",
    "    yd = pm.Normal('y',mu_est , s[g] ** -2, observed=y)\n",
    "\n",
    "    start = pm.find_MAP()\n",
    "    #h = find_hessian(start)\n",
    "\n",
    "    step = pm.NUTS(model.vars, scaling=start)\n",
    "\n",
    "with model:\n",
    "trace = pm.sample(3000, step, start)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gs_name = '_gs1512' #'_gs15656' #'_gs1512' #'_gs15656' #'_gs1512'   #'_gs5188' #'_gs1512'\n",
    "\n",
    "\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "columns  = pd.read_hdf('stp_to_ge_all_columns'+gs_name+'.hdf')\n",
    "ge_columns_train      = columns.loc[:,'ge_columns_train'].values[0]\n",
    "annot_columns_train   = columns.loc[:,'annot_columns_train'].values[0]\n",
    "stp_columns_train     = columns.loc[:,'stp_columns_train'].values[0]\n",
    "classes_columns_train = columns.loc[:,'classes_columns_train'].values[0]\n",
    "#print(X_train.columns)\n",
    "X_train0 = pd.read_hdf('stp_to_ge_all_data'+gs_name+'.hdf')\n",
    "\n",
    "X_train0.shape\n",
    "print(X_train0.columns[200:300])\n",
    "print(X_train0.columns[300:])\n",
    "\n",
    "\n",
    "\n",
    "print('gs_name - '+gs_name)\n",
    "\n",
    "if (gs_name=='_gs5188')|(gs_name=='_gs1512')|(gs_name=='_gs15656'):\n",
    "    nlf = 10   # 20 were in scvi model\n",
    "elif (gs_name=='_gs219'):\n",
    "    nlf = 10   # 10 were in scvi model\n",
    "lf_names = np.arange(nlf).astype(str)\n",
    "lf_names1 =  ['pre__'+s for s in lf_names.tolist()]\n",
    "lf_names2 =  ['post__'+s for s in lf_names.tolist()]\n",
    "\n",
    "lf_scvi_names = lf_names1 + lf_names2\n",
    "\n",
    "ge_columns_train2 = list(set(ge_columns_train).difference(set(lf_scvi_names)))\n",
    "lf_scvi_names\n",
    "\n",
    "#\n",
    "#\n",
    "#  genes to STP: Hierarchical Linear Model\n",
    "#\n",
    "#\n",
    "\n",
    "#### 4 models:\n",
    "#    A. 0->subclass_pre->+class_post\n",
    "#    B. 0->class_pre->subclass_pre->+x class_post\n",
    "#    C. 0->class_pre->subclass_pre->+subclass_post \n",
    "#    D. 0->class_pre->+x class_post\n",
    "\n",
    "H_Models = pd.DataFrame([           ['Ms1c2'  ,['subclass_pre', 'class_post']],\n",
    "                                    ['Mc1s1c2',['class_pre', 'subclass_pre', 'class_post']],\n",
    "                                    ['Mc1s1s2',['class_pre', 'subclass_pre', 'subclass_post']],\n",
    "                                    ['Mc1c2'  ,['class_pre', 'class_post']]],\n",
    "                                    columns = ['name','structure'])\n",
    "\n",
    "\n",
    "#### Load Data\n",
    "#### modify data: DO class_pre, class_post, subclass_pre, subclass_post columns\n",
    "subcl_names_pre = ['ex_ctx__pre', 'ex_ec__pre','ex_hipp__pre','sst__pre','pvalb__pre','cge__pre']\n",
    "subcl_names_post = ['ex_ctx__post', 'ex_ec__post','ex_hipp__post','sst__post','pvalb__post','cge__post']\n",
    "new_syntype_columns = pd.DataFrame([['class_pre',['ex__pre', 'inh__pre'], ['Glutamatergic', 'GABAergic']],\n",
    "                                    ['subclass_pre',subcl_names_pre, subcl_names_pre],\n",
    "                                    ['class_post',['ex__post', 'inh__post'], ['Glutamatergic', 'GABAergic']],\n",
    "                                    ['subclass_post',subcl_names_post, subcl_names_post]], columns = ['group','columns_X0','new_names'])\n",
    "\n",
    "for g in new_syntype_columns.index:\n",
    "    df_new = X_train0.loc[:,new_syntype_columns.loc[g,'columns_X0']].copy()\n",
    "    g20 = new_syntype_columns.loc[g,'columns_X0'][0]\n",
    "    for g2 in new_syntype_columns.loc[g,'columns_X0']:\n",
    "        df_new.loc[df_new.loc[:,g2]!=0, g20] = g2\n",
    "    df_new_g = df_new.loc[:,g20].copy()   \n",
    "    df_new_g.name = new_syntype_columns.loc[g,'group']\n",
    "    X_train0 = pd.concat([X_train0, df_new_g],axis=1)\n",
    "    classes_columns_train = classes_columns_train + [df_new_g.name]\n",
    "\n",
    "#X_train0.loc[:,classes_columns_train].to_excel('X_train0.xlsx')  \n",
    "\n",
    "\n",
    "\n",
    "####\n",
    "\n",
    "Dn = 50\n",
    "N_bootstraps = 100\n",
    "Dn2 = int(N_bootstraps/Dn) # 100 - should be a number of bootstraps per synapse type !!!\n",
    "Dn3 = int(X_train0.shape[0]/N_bootstraps) # number of synapse_types\n",
    "#cla_n = ['ex_in_ex_in__pre',\n",
    "#         'ex_in_ex_in__post'] \n",
    "#cla_n = ['ex_in_ex_in__pre'] \n",
    "#cla_n = []\n",
    "\n",
    "HM_name = 'Ms1c2'\n",
    "HM = H_Models.loc[H_Models.loc[:,'name']==HM_name, :]\n",
    "cla_n  = HM['structure'][0]\n",
    "\n",
    "#ge_n = imp50.index[0:25].tolist()\n",
    "ge_n = lf_scvi_names #ge_columns_train2\n",
    "\n",
    "nannot = len(annot_columns_train)\n",
    "lge_n  = len(ge_n)\n",
    "\n",
    "X2 = X_train0.loc[:,annot_columns_train + ge_n + cla_n]\n",
    "if len(cla_n)>0:\n",
    "    i_cl = np.nonzero(X2.columns.isin(cla_n))[0] - nannot\n",
    "    #X2.loc[:,cla_n] = X2.loc[:,cla_n]+1\n",
    "\n",
    "X2=X2.iloc[0:X_train0.shape[0]:Dn,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Index each of the unique variable values\n",
    "n_levels = len(cla_n)\n",
    "cla_n2 = np.array(cla_n)\n",
    "index = []\n",
    "index_df = []\n",
    "for lv in range(n_levels):\n",
    "    lv_index = X2.groupby(cla_n[0:lv+1] ).all().reset_index().reset_index()[['index']+cla_n[0:lv+1]]\n",
    "    index = index + [lv_index]\n",
    "    \n",
    "    lv_indexes_df = lv_index\n",
    "    if lv==1:\n",
    "        lv_indexes_df = pd.merge(index[lv-1], index[lv], how='inner',\n",
    "                                 on=cla_n[0:lv], suffixes=('_l'+str(lv-1), '_l'+str(lv)))\n",
    "    elif lv>1:\n",
    "        lv_indexes_df = pd.merge(lv_indexes_df, index[lv], how='inner', on=cla_n[0:lv])\n",
    "    index_df = index_df + [lv_indexes_df ]    \n",
    "        \n",
    "X3 = pd.merge(X2, lv_indexes_df, how='inner', on=cla_n[0:]).reset_index()\n",
    "\n",
    "import pymc3 as pm\n",
    "import theano.tensor as T\n",
    "\n",
    "indexes = []\n",
    "count = []\n",
    "for lv in range(n_levels):\n",
    "    if lv==0:\n",
    "        index_nm = 'index'\n",
    "    else:\n",
    "        index_nm = 'index_l'+str(lv-1)\n",
    "    lv_indexes = index_df[lv][index_nm].values\n",
    "    indexes = indexes + [lv_indexes]\n",
    "    count = count + [len(lv_indexes)]\n",
    "\n",
    "# degree_indexes = degree_index['index'].values\n",
    "# degree_count = len(degree_indexes)\n",
    "# degree_state_indexes = degree_state_indexes_df['index_d'].values\n",
    "# degree_state_count = len(degree_state_indexes)\n",
    "# degree_state_county_indexes = degree_state_county_indexes_df['index_ds'].values\n",
    "# degree_state_county_count = len(degree_state_county_indexes)\n",
    "\n",
    "with pm.Model() as mvgmm:\n",
    "    # Prior over component weights\n",
    "    p = pm.Dirichlet('p', a=np.array([1.]*K))\n",
    "\n",
    "    # Prior over component means\n",
    "    mus = [pm.MvNormal('mu_%d' % i,\n",
    "                        mu=pm.floatX(np.zeros(D)),\n",
    "                        tau=pm.floatX(0.1 * np.eye(D)),\n",
    "                        shape=(D,),\n",
    "                        testval=pm.floatX(testvals[i]))\n",
    "           for i in range(K)]\n",
    "\n",
    "    # Cholesky decomposed LKJ prior over component covariance matrices\n",
    "    packed_L = [pm.LKJCholeskyCov('packed_L_%d' % i,\n",
    "                                  n=D,\n",
    "                                  eta=2.,\n",
    "                                  sd_dist=pm.HalfCauchy.dist(1))\n",
    "                for i in range(K)]\n",
    "\n",
    "    # Unpack packed_L into full array\n",
    "    L = [pm.expand_packed_triangular(D, packed_L[i])\n",
    "         for i in range(K)]\n",
    "\n",
    "    # Convert L to sigma and tau for convenience\n",
    "    sigma = [pm.Deterministic('sigma_%d' % i ,L[i].dot(L[i].T))\n",
    "             for i in range(K)]\n",
    "    tau = [pm.Deterministic('tau_%d' % i,matrix_inverse(sigma[i]))\n",
    "           for i in range(K)]\n",
    "\n",
    "    print(\"# Specify the likelihood\")\n",
    "    # Specify the likelihood\n",
    "    mvnl = [pm.MvNormal.dist(mu=mus[i],chol=L[i])  \n",
    "           for i in range(K)]    \n",
    "    Y_obs = pm.Mixture('Y_obs',w=p, comp_dists=mvnl,observed=data)  \n",
    "\n",
    "    print(\"# Start the sampler!\")\n",
    "    # Start the sampler!\n",
    "    trace = pm.sample(draws, step=pm.NUTS(), chains=4)\n",
    "    \n",
    "\n",
    "with pm.Model() as model:\n",
    "    lv=0 # global\n",
    "    mu_0 = pm.MvNormal('mu_0',\n",
    "                        mu=pm.floatX(np.zeros(D)),\n",
    "                        sig=pm.floatX(2 * np.eye(D)),\n",
    "                        shape=(D,),\n",
    "                        testval=pm.floatX(testvals[i]))\n",
    "    mu_sig_0 = pm.Uniform('mu_sig_0', lower=0, upper=4)\n",
    "    b_0 = pm.MvNormal('b_0',\n",
    "                        mu=pm.floatX(np.zeros(D)),\n",
    "                        sig=pm.floatX(2 * np.eye(D)),\n",
    "                        shape=(D,),\n",
    "                        testval=pm.floatX(testvals[i]))\n",
    "    b_sig_0 = pm.Uniform('b_sig_0', lower=0, upper=2)\n",
    "    \n",
    "            \n",
    "    lv=1  #first level: e-e, e-i, i-i, i-e\n",
    "    # Prior over component means\n",
    "    mu_1 = pm.MvNormal('mu_1',\n",
    "                        mu=mu_0,\n",
    "                        sig=pm.floatX(2  * np.eye(D)) ,\n",
    "                        shape=(D,count[0]))\n",
    "    mu_sig_1 = pm.Uniform('mu_sig_1', lower=0, upper=3)\n",
    "    b_1 = pm.MvNormal('b_1',\n",
    "                        mu=pm.floatX(np.zeros(D)),\n",
    "                        sig=pm.floatX(2  * np.eye(D)),\n",
    "                        shape=(D,count[0]))\n",
    "    b_sig_1 = pm.Uniform('b_sig_1', lower=0, upper=2)       \n",
    "\n",
    "\n",
    "    \n",
    "    lv=2\n",
    "    # Prior over component means\n",
    "    mu_2 = pm.MvNormal('mu_%d' % i,\n",
    "                        mu=mu_1,\n",
    "                        sig=pm.floatX(mu_sig_1  * np.eye(D)) ,\n",
    "                        shape=(D,count[0]))\n",
    "    \n",
    "    # Cholesky decomposed LKJ prior over component covariance matrices\n",
    "    packed_L = [pm.LKJCholeskyCov('packed_L_%d' % i,\n",
    "                                  n=D,\n",
    "                                  eta=2.,\n",
    "                                  sd_dist=pm.HalfCauchy.dist(1))\n",
    "                for i in range(K)]\n",
    "\n",
    "    # Unpack packed_L into full array\n",
    "    L = [pm.expand_packed_triangular(D, packed_L[i])\n",
    "         for i in range(K)]\n",
    "\n",
    "    \n",
    "    # Convert L to sigma and tau for convenience\n",
    "    sigma = [pm.Deterministic('sigma_%d' % i ,L[i].dot(L[i].T))\n",
    "             for i in range(K)]\n",
    "    tau = [pm.Deterministic('tau_%d' % i,matrix_inverse(sigma[i]))\n",
    "           for i in range(K)]\n",
    "\n",
    "\n",
    "    sigma_unexplained = pm.HalfNormal('sigma_unexplained', sd=0.2) # unexplained variability\n",
    "    y_prediction = mu_2[indexed_salary_df['index'].values] * indexed_salary_df.month_index.values + degree_state_county_b[indexed_salary_df['index'].values]\n",
    "    \n",
    "    y_like = pm.Normal('y_likelihood', mu=y_prediction, sd=sigma_unexplained, observed=y3.loc[:,stp_n].values)\n",
    "    \n",
    "    # robust?\n",
    "    #y_like = pm.StudentT('y_like', nu=1, mu=0, sd=sigma_unexplained, observed=y_prediction - y3.loc[:,stp_n].values)\n",
    "    \n",
    "    global_m = pm.Normal('global_m', mu=6500, sd=100**2)\n",
    "    global_m_sd = pm.Uniform('global_m_sd', lower=0, upper=1000)\n",
    "    global_b = pm.Normal('global_b', mu=1000, sd=100**2)\n",
    "    global_b_sd = pm.Uniform('global_b_sd', lower=0, upper=1000)\n",
    "        \n",
    "    degree_m = pm.Normal('degree_m', mu=global_m, sd=global_m_sd, shape=degree_count)\n",
    "    degree_m_sd = pm.Uniform('degree_m_sd', lower=0, upper=1000, shape=degree_count)\n",
    "    degree_b = pm.Normal('degree_b', mu=global_b, sd=global_b_sd, shape=degree_count)\n",
    "    degree_b_sd = pm.Uniform('degree_b_sd', lower=0, upper=1000, shape=degree_count)\n",
    "    \n",
    "    degree_state_m = pm.Normal('degree_state_m', mu=degree_m[degree_state_indexes], sd=degree_m_sd[degree_state_indexes], shape=degree_state_count)\n",
    "    degree_state_m_sd = pm.Uniform('degree_state_m_sd', lower=0, upper=1000, shape=degree_state_count)\n",
    "    degree_state_b = pm.Normal('degree_state_b', mu=degree_b[degree_state_indexes], sd=degree_b_sd[degree_state_indexes], shape=degree_state_count)\n",
    "    degree_state_b_sd = pm.Uniform('degree_state_b_sd', lower=0, upper=1000, shape=degree_state_count)\n",
    "    \n",
    "    degree_state_county_m = pm.Normal('degree_state_county_m', mu=degree_state_m[degree_state_county_indexes], sd=degree_state_m_sd[degree_state_county_indexes], shape=degree_state_county_count)\n",
    "    degree_state_county_b = pm.Normal('degree_state_county_b', mu=degree_state_b[degree_state_county_indexes], sd=degree_state_b_sd[degree_state_county_indexes], shape=degree_state_county_count)\n",
    "    \n",
    "    error = pm.Uniform('error', lower=0, upper=10000)\n",
    "    \n",
    "    y_prediction = degree_state_county_m[indexed_salary_df['index'].values] * indexed_salary_df.month_index.values + degree_state_county_b[indexed_salary_df['index'].values]\n",
    "    \n",
    "    pm.StudentT('y_like', nu=1, mu=0, sd=error, observed=y_prediction - indexed_salary_df.salary.values)\n",
    "    \n",
    "    #model_trace = pm.sample(5000)\n",
    "    #trace = pm.sample(draws=DRAWS+BURNIN, chains=CHAINS, cores=CORES, tune=TUNE, random_seed=SEED, nuts_kwargs={'target_accept':0.95})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "#\n",
    "#  genes to STP: Hierarchical Linear Model\n",
    "#\n",
    "#\n",
    "\n",
    "#### 4 models:\n",
    "#    A. 0->subclass_pre->+class_post\n",
    "#    B. 0->class_pre->subclass_pre->+x class_post\n",
    "#    C. 0->class_pre->subclass_pre->+subclass_post \n",
    "#    D. 0->class_pre->+x class_post\n",
    "\n",
    "H_Models = pd.DataFrame([           ['Ms1c2'  ,['subclass_pre', 'class_post']],\n",
    "                                    ['Mc1s1c2',['class_pre', 'subclass_pre', 'class_post']],\n",
    "                                    ['Mc1s1s2',['class_pre', 'subclass_pre', 'subclass_post']],\n",
    "                                    ['Mc1c2'  ,['class_pre', 'class_post']]],\n",
    "                                    columns = ['name','structure'])\n",
    "\n",
    "\n",
    "#### Load Data\n",
    "#### modify data: DO class_pre, class_post, subclass_pre, subclass_post columns\n",
    "subcl_names_pre = ['ex_ctx__pre', 'ex_ec__pre','ex_hipp__pre','sst__pre','pvalb__pre','cge__pre']\n",
    "subcl_names_post = ['ex_ctx__post', 'ex_ec__post','ex_hipp__post','sst__post','pvalb__post','cge__post']\n",
    "new_syntype_columns = pd.DataFrame([['class_pre',['ex__pre', 'inh__pre'], ['Glutamatergic', 'GABAergic']],\n",
    "                                    ['subclass_pre',subcl_names_pre, subcl_names_pre],\n",
    "                                    ['class_post',['ex__post', 'inh__post'], ['Glutamatergic', 'GABAergic']],\n",
    "                                    ['subclass_post',subcl_names_post, subcl_names_post]], columns = ['group','columns_X0','new_names'])\n",
    "\n",
    "for g in new_syntype_columns.index:\n",
    "    df_new = X_train0.loc[:,new_syntype_columns.loc[g,'columns_X0']].copy()\n",
    "    g20 = new_syntype_columns.loc[g,'columns_X0'][0]\n",
    "    for g2 in new_syntype_columns.loc[g,'columns_X0']:\n",
    "        df_new.loc[df_new.loc[:,g2]!=0, g20] = g2\n",
    "    df_new_g = df_new.loc[:,g20].copy()   \n",
    "    df_new_g.name = new_syntype_columns.loc[g,'group']\n",
    "    X_train0 = pd.concat([X_train0, df_new_g],axis=1)\n",
    "    classes_columns_train = classes_columns_train + [df_new_g.name]\n",
    "\n",
    "#X_train0.loc[:,classes_columns_train].to_excel('X_train0.xlsx')  \n",
    "\n",
    "\n",
    "\n",
    "####\n",
    "\n",
    "Dn = 50\n",
    "N_bootstraps = 100\n",
    "Dn2 = int(N_bootstraps/Dn) # 100 - should be a number of bootstraps per synapse type !!!\n",
    "Dn3 = int(X_train0.shape[0]/N_bootstraps) # number of synapse_types\n",
    "#cla_n = ['ex_in_ex_in__pre',\n",
    "#         'ex_in_ex_in__post'] \n",
    "#cla_n = ['ex_in_ex_in__pre'] \n",
    "#cla_n = []\n",
    "\n",
    "HM_name = 'Ms1c2'\n",
    "HM = H_Models.loc[H_Models.loc[:,'name']==HM_name, :]\n",
    "cla_n  = HM['structure'][0]\n",
    "\n",
    "#ge_n = imp50.index[0:25].tolist()\n",
    "ge_n = lf_scvi_names #ge_columns_train2\n",
    "\n",
    "nannot = len(annot_columns_train)\n",
    "lge_n  = len(ge_n)\n",
    "\n",
    "X2 = X_train0.loc[:,annot_columns_train + ge_n + cla_n]\n",
    "if len(cla_n)>0:\n",
    "    i_cl = np.nonzero(X2.columns.isin(cla_n))[0] - nannot\n",
    "    #X2.loc[:,cla_n] = X2.loc[:,cla_n]+1\n",
    "\n",
    "X2=X2.iloc[0:X_train0.shape[0]:Dn,:]\n",
    "\n",
    "X2 =X2.values\n",
    "\n",
    "# select y\n",
    "stp_n = ['A2_20Hz'] #,'A5_20Hz','A250_20Hz']\n",
    "\n",
    "# 'A2_20Hz',\n",
    "#  'A3_20Hz',\n",
    "#  'A4_20Hz',\n",
    "#  'A5_20Hz',\n",
    "#  'A250_20Hz',\n",
    "#  'A1000_20Hz',\n",
    "#  'A1_50Hz',\n",
    "#  'A2_50Hz',\n",
    "#  'A3_50Hz',\n",
    "#  'A4_50Hz',\n",
    "#  'A5_50Hz',\n",
    "#  'A250_50Hz',\n",
    "#  'A1000_50Hz',\n",
    "#  'A1_10Hz',\n",
    "#  'A2_10Hz',\n",
    "#  'A3_10Hz',\n",
    "#  'A4_10Hz',\n",
    "#  'A5_10Hz',\n",
    "#  'A250_10Hz',\n",
    "#  'A1000_10Hz']\n",
    "y2 =  X_train0.loc[:,annot_columns_train+stp_n].iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "\n",
    "#header = ge_n + cla_n + stp_n \n",
    "header = ge_n  + stp_n \n",
    "annot_columns2 = annot_columns_train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyMC3 example : HLM 1D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import pymc3 as pm\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "\n",
    "\n",
    "# These are the \"true\" slopes and intercepts that I made up.\n",
    "set_parameters = (\n",
    "    (('no degree', 'CA', 'Riverside County'), (5000, 1250)),\n",
    "    (('no degree', 'IL', 'Cook County'), (6500, 1150)),\n",
    "    (('no degree', 'IL', 'Lake County'), (7000, 1350)),\n",
    "    \n",
    "    (('degree', 'CA', 'Riverside County'), (6000, 1250)),\n",
    "    (('degree', 'IL', 'Cook County'), (7500, 1150)),\n",
    "    (('degree', 'IL', 'Lake County'), (8000, 1350))\n",
    ")\n",
    "\n",
    "data_points = np.random.randint(1,24, size=len(set_parameters))\n",
    "\n",
    "# Go through each definition above and generate a fake time series\n",
    "\n",
    "\n",
    "rows = []\n",
    "for set_parameter, N in zip(set_parameters, data_points):\n",
    "  key, parameters = set_parameter\n",
    "  population_rows = [{\n",
    "      'degree': key[0],\n",
    "      'state': key[1],\n",
    "      'county': key[2],\n",
    "      'month_index': i,\n",
    "      'salary': i*parameters[0] + parameters[1] + np.random.normal(loc=0, scale=500)\n",
    "  } for i in range(N)]\n",
    "  rows += population_rows\n",
    "\n",
    "salary_df = pd.DataFrame(rows)\n",
    "\n",
    "#plt.scatter(salary_df.month_index, salary_df.salary, s=2);\n",
    "salary_df.to_excel('temp2.xlsx')\n",
    "#salary_df\n",
    "\n",
    "# Index each of the unique variable values\n",
    "degree_index = salary_df.groupby('degree').all().reset_index().reset_index()[['index', 'degree']]\n",
    "degree_state_index = salary_df.groupby(['degree', 'state']).all().reset_index().reset_index()[['index', 'degree', 'state']]\n",
    "degree_state_county_index = salary_df.groupby(['degree', 'state', 'county']).all().reset_index().reset_index()[['index', 'degree', 'state', 'county']]\n",
    "\n",
    "degree_state_indexes_df = pd.merge(degree_index, degree_state_index, how='inner', on='degree', suffixes=('_d', '_ds'))\n",
    "degree_state_county_indexes_df = pd.merge(degree_state_indexes_df, degree_state_county_index, how='inner', on=['degree', 'state'])\n",
    "indexed_salary_df = pd.merge(salary_df, degree_state_county_indexes_df, how='inner', on=['degree', 'state', 'county']).reset_index()\n",
    "\n",
    "\n",
    "import pymc3 as pm\n",
    "\n",
    "degree_indexes = degree_index['index'].values\n",
    "degree_count = len(degree_indexes)\n",
    "degree_state_indexes = degree_state_indexes_df['index_d'].values\n",
    "degree_state_count = len(degree_state_indexes)\n",
    "degree_state_county_indexes = degree_state_county_indexes_df['index_ds'].values\n",
    "degree_state_county_count = len(degree_state_county_indexes)\n",
    "\n",
    "with pm.Model() as model:\n",
    "    global_m = pm.Normal('global_m', mu=6500, sd=100**2)\n",
    "    global_m_sd = pm.Uniform('global_m_sd', lower=0, upper=1000)\n",
    "    global_b = pm.Normal('global_b', mu=1000, sd=100**2)\n",
    "    global_b_sd = pm.Uniform('global_b_sd', lower=0, upper=1000)\n",
    "    \n",
    "    degree_m = pm.Normal('degree_m', mu=global_m, sd=global_m_sd, shape=degree_count)\n",
    "    degree_m_sd = pm.Uniform('degree_m_sd', lower=0, upper=1000, shape=degree_count)\n",
    "    degree_b = pm.Normal('degree_b', mu=global_b, sd=global_b_sd, shape=degree_count)\n",
    "    degree_b_sd = pm.Uniform('degree_b_sd', lower=0, upper=1000, shape=degree_count)\n",
    "    \n",
    "    degree_state_m = pm.Normal('degree_state_m', mu=degree_m[degree_state_indexes], sd=degree_m_sd[degree_state_indexes], shape=degree_state_count)\n",
    "    degree_state_m_sd = pm.Uniform('degree_state_m_sd', lower=0, upper=1000, shape=degree_state_count)\n",
    "    degree_state_b = pm.Normal('degree_state_b', mu=degree_b[degree_state_indexes], sd=degree_b_sd[degree_state_indexes], shape=degree_state_count)\n",
    "    degree_state_b_sd = pm.Uniform('degree_state_b_sd', lower=0, upper=1000, shape=degree_state_count)\n",
    "    \n",
    "    degree_state_county_m = pm.Normal('degree_state_county_m', mu=degree_state_m[degree_state_county_indexes], sd=degree_state_m_sd[degree_state_county_indexes], shape=degree_state_county_count)\n",
    "    degree_state_county_b = pm.Normal('degree_state_county_b', mu=degree_state_b[degree_state_county_indexes], sd=degree_state_b_sd[degree_state_county_indexes], shape=degree_state_county_count)\n",
    "    \n",
    "    error = pm.Uniform('error', lower=0, upper=10000)\n",
    "    \n",
    "    y_prediction = degree_state_county_m[indexed_salary_df['index'].values] * indexed_salary_df.month_index.values + degree_state_county_b[indexed_salary_df['index'].values]\n",
    "    \n",
    "    pm.StudentT('y_like', nu=1, mu=0, sd=error, observed=y_prediction - indexed_salary_df.salary.values)\n",
    "    \n",
    "    model_trace = pm.sample(5000)"
   ]
  },
  {
   "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": {},
   "source": [
    "## CV of different type for all ML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import sklearn as sk\n",
    "\n",
    "def fit_classes_tree(X2,y2,X2_cl,cli,model_type='lasso',nmin = 20, alpha=1, l1_ratio=0.5):\n",
    "    model = []\n",
    "    support =[]\n",
    "    ncl1 = np.max(X2_cl[:,cli[0]])+1 \n",
    "    ncl2 = np.max(X2_cl[:,cli[1]])+1 \n",
    "    for i1 in range(ncl1):\n",
    "        model1 = []\n",
    "        support1 = []\n",
    "        for i2 in range(ncl2):\n",
    "\n",
    "            is_in_cl = (X2_cl[:,cli[0]]==i1)&(X2_cl[:,cli[1]]==i2)\n",
    "            #           support2 = [np.sum(is_in_cl),\n",
    "            #           np.nonzero((X2_cl[:,cli[0]]==i1))[0],\n",
    "            #           np.nonzero((X2_cl[:,cli[1]]==i2))[0]]\n",
    "            support2 = np.sum(is_in_cl)\n",
    "            \n",
    "            support1 = support1 + [support2]\n",
    "            regr = 0\n",
    "            if support2>nmin:\n",
    "                y_train = y2[is_in_cl,:]\n",
    "                X_train = X2[is_in_cl,:]\n",
    "                if model_type=='ridge':\n",
    "                    regr = sk.linear_model.Ridge(alpha=0.50, fit_intercept=True, normalize=False,\n",
    "                                                 copy_X=True, max_iter=None, tol=0.001, solver='auto',\n",
    "                                                 random_state=None)\n",
    "                \n",
    "                elif model_type=='lasso':\n",
    "                    regr = sk.linear_model.Lasso(alpha=alpha, fit_intercept=True, normalize=False,\n",
    "                                                 precompute=False, copy_X=True,\n",
    "                                                 max_iter=1000, tol=0.0001, warm_start=False, positive=False, \n",
    "                                                 random_state=None, selection='cyclic')\n",
    "                elif model_type=='elastic_net':   \n",
    "                    regr = sk.linear_model.ElasticNet(alpha=alpha, l1_ratio=l1_ratio, fit_intercept=True,\n",
    "                                    normalize=False,\n",
    "                                    max_iter=1000, copy_X=True, tol=0.0001, warm_start=False,\n",
    "                                    random_state=None, selection='cyclic')\n",
    "                else:\n",
    "                    print(\" model types supported: ridge, lasso, elastic_net\")\n",
    "                \n",
    "                regr.fit(X_train, y_train)\n",
    "                #y_pred = regr.predict(X_test[:,nannot:]) \n",
    "\n",
    "\n",
    "            model1 = model1 + [regr]\n",
    "            \n",
    "        support = support + [np.array(support1)]\n",
    "        model = model + [np.array(model1)]    \n",
    "    return np.array(model), np.array(support)        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_classes_tree(model,X2,X2_cl,cli,nout=1,nmin = 20):\n",
    "    \n",
    "    ncl1 = np.max(X2_cl[:,cli[0]])+1 \n",
    "    ncl2 = np.max(X2_cl[:,cli[1]])+1 \n",
    "    y_pred = np.zeros((X2.shape[0],nout))\n",
    "    for i1 in range(ncl1):\n",
    "        #print(i1,\" f1\")\n",
    "        if len(model)>i1:\n",
    "            #print(i1,\" f2\")\n",
    "            model1 = model[i1]\n",
    "            for i2 in range(ncl2):\n",
    "\n",
    "                is_in_cl = (X2_cl[:,cli[0]]==i1)&(X2_cl[:,cli[1]]==i2)\n",
    "\n",
    "                support2 = np.sum(is_in_cl)\n",
    "                \n",
    "                #print(i1,i2,\" f3 \",support2)\n",
    "                #regr = 0\n",
    "                if support2>nmin:\n",
    "                    #print(i1,i2,\" f3 \",support2)\n",
    "                    if len(model1)>i2:\n",
    "                        #print(i1,i2,\" f4\")\n",
    "                        regr = model1[i2]\n",
    "                        #regr = sk.linear_model.MultiTaskElasticNet(alpha=alpha, l1_ratio=l1_ratio, fit_intercept=True,\n",
    "                        #                    normalize=False,\n",
    "                        #                    max_iter=1000, copy_X=True, tol=0.0001, warm_start=False,\n",
    "                        #                    random_state=None, selection='cyclic')\n",
    "                        #regr.fit(X_train, y_train)\n",
    "                        if regr!=0:\n",
    "                            X_test = X2[is_in_cl,:]\n",
    "                            y_pred[is_in_cl] = regr.predict(X_test).reshape((-1,1)) \n",
    "    \n",
    "    return y_pred       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cv_random_groups(n_splits, groups_index, part ):\n",
    "    #  part % random syn-types level2 cv \n",
    "    #part = 10 # % of removed samples\n",
    "    #n_splits=10\n",
    "    \n",
    "    test_indexes = []\n",
    "    names = []\n",
    "    #cla_n = [cla_cell_type_l2]\n",
    "\n",
    "    #groups_index = X_train0.loc[:,cla_n]\n",
    "    #groups_index=groups_index.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "    all_gr = set(groups_index.reshape((-1,)))\n",
    "    all_gr = all_gr.difference(set([0]))\n",
    "    all_gr = list(all_gr)\n",
    "\n",
    "    #dn_syn_type = 100\n",
    "    nall = groups_index.shape[0]\n",
    "    all_index = np.arange(nall)\n",
    "    Dn3_out = int(nall*part/100)\n",
    "    idx0 = all_gr  #np.arange(len(all_gr))\n",
    "    #idx1 = np.arange(Dn2)\n",
    "    for icv in range(n_splits):\n",
    "        idx0_out = np.random.choice(idx0, size=len(all_gr), replace=False, p=None)\n",
    "        test_index = np.zeros((0,))\n",
    "        for io in idx0_out:\n",
    "            #train_index  = train_index + icv*dn + np.arange(dn)\n",
    "            idx1 = all_index[np.any(groups_index==io,axis=1)]\n",
    "            if len(idx1) + len(test_index)<Dn3_out:\n",
    "                test_index  = np.append(test_index, idx1)\n",
    "            else:\n",
    "                idx1 = np.random.choice(idx1, size=Dn3_out-len(test_index), replace=False, p=None)\n",
    "                test_index  = np.append(test_index, idx1)\n",
    "                break\n",
    "\n",
    "        test_indexes = test_indexes + [test_index]   \n",
    "        names = names + ['syn_types_l0_'+str(icv)+'  '+np.array2string(idx0_out)]\n",
    "        #df_cv.loc[icv,'train_index']=[train_index]\n",
    "        #df_cv.loc[icv,'name']='syn_types_l0_'+str(icv)+'  '+np.array2string(idx0_out)\n",
    "\n",
    "    df_cv1  = pd.DataFrame([], columns = ['test_index', 'name'])\n",
    "    return df_cv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gs_name = '_gs1512' #'_gs15656' #'_gs1512' #'_gs15656' #'_gs1512'   #'_gs5188' #'_gs1512'\n",
    "\n",
    "\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "columns  = pd.read_hdf('stp_to_ge_all_columns'+gs_name+'.hdf')\n",
    "ge_columns_train      = columns.loc[:,'ge_columns_train'].values[0]\n",
    "annot_columns_train   = columns.loc[:,'annot_columns_train'].values[0]\n",
    "stp_columns_train     = columns.loc[:,'stp_columns_train'].values[0]\n",
    "classes_columns_train = columns.loc[:,'classes_columns_train'].values[0]\n",
    "#print(X_train.columns)\n",
    "X_train0 = pd.read_hdf('stp_to_ge_all_data'+gs_name+'.hdf')\n",
    "\n",
    "X_train0.shape\n",
    "print(X_train0.columns[200:300])\n",
    "print(X_train0.columns[300:])\n",
    "\n",
    "\n",
    "\n",
    "print('gs_name - '+gs_name)\n",
    "\n",
    "if (gs_name=='_gs5188')|(gs_name=='_gs1512')|(gs_name=='_gs15656'):\n",
    "    nlf = 10   # 20 were in scvi model\n",
    "elif (gs_name=='_gs219'):\n",
    "    nlf = 10   # 10 were in scvi model\n",
    "lf_names = np.arange(nlf).astype(str)\n",
    "lf_names1 =  ['pre__'+s for s in lf_names.tolist()]\n",
    "lf_names2 =  ['post__'+s for s in lf_names.tolist()]\n",
    "\n",
    "lf_scvi_names = lf_names1 + lf_names2\n",
    "\n",
    "ge_columns_train2 = list(set(ge_columns_train).difference(set(lf_scvi_names)))\n",
    "lf_scvi_names\n",
    "\n",
    "#\n",
    "#\n",
    "#  genes to STP: Hierarchical Linear Model\n",
    "#\n",
    "#\n",
    "\n",
    "#### 4 models:\n",
    "#    A. 0->subclass_pre->+class_post\n",
    "#    B. 0->class_pre->subclass_pre->+x class_post\n",
    "#    C. 0->class_pre->subclass_pre->+subclass_post \n",
    "#    D. 0->class_pre->+x class_post\n",
    "\n",
    "H_Models = pd.DataFrame([           ['Ms1c2'  ,['subclass_pre', 'class_post']],\n",
    "                                    ['Mc1s1c2',['class_pre', 'subclass_pre', 'class_post']],\n",
    "                                    ['Mc1s1s2',['class_pre', 'subclass_pre', 'subclass_post']],\n",
    "                                    ['Mc1c2'  ,['class_pre', 'class_post']]],\n",
    "                                    columns = ['name','structure'])\n",
    "\n",
    "\n",
    "#### Load Data\n",
    "#### modify data: DO class_pre, class_post, subclass_pre, subclass_post columns\n",
    "subcl_names_pre = ['ex_ctx__pre', 'ex_ec__pre','ex_hipp__pre','sst__pre','pvalb__pre','cge__pre']\n",
    "subcl_names_post = ['ex_ctx__post', 'ex_ec__post','ex_hipp__post','sst__post','pvalb__post','cge__post']\n",
    "new_syntype_columns = pd.DataFrame([['class_pre',['ex__pre', 'inh__pre'], ['Glutamatergic', 'GABAergic']],\n",
    "                                    ['subclass_pre',subcl_names_pre, subcl_names_pre],\n",
    "                                    ['class_post',['ex__post', 'inh__post'], ['Glutamatergic', 'GABAergic']],\n",
    "                                    ['subclass_post',subcl_names_post, subcl_names_post]], columns = ['group','columns_X0','new_names'])\n",
    "\n",
    "for g in new_syntype_columns.index:\n",
    "    df_new = X_train0.loc[:,new_syntype_columns.loc[g,'columns_X0']].copy()\n",
    "    g20 = new_syntype_columns.loc[g,'columns_X0'][0]\n",
    "    for g2 in new_syntype_columns.loc[g,'columns_X0']:\n",
    "        df_new.loc[df_new.loc[:,g2]!=0, g20] = g2\n",
    "    df_new_g = df_new.loc[:,g20].copy()   \n",
    "    df_new_g.name = new_syntype_columns.loc[g,'group']\n",
    "    X_train0 = pd.concat([X_train0, df_new_g],axis=1)\n",
    "    classes_columns_train = classes_columns_train + [df_new_g.name]\n",
    "\n",
    "#X_train0.loc[:,classes_columns_train].to_excel('X_train0.xlsx')  \n",
    "\n",
    "\n",
    "\n",
    "####\n",
    "\n",
    "# Dn = 50\n",
    "# N_bootstraps = 100\n",
    "# Dn2 = int(N_bootstraps/Dn) # 100 - should be a number of bootstraps per synapse type !!!\n",
    "# Dn3 = int(X_train0.shape[0]/N_bootstraps) # number of synapse_types\n",
    "\n",
    "# HM_name = 'Ms1c2'\n",
    "# HM = H_Models.loc[H_Models.loc[:,'name']==HM_name, :]\n",
    "# cla_n  = HM['structure'][0]\n",
    "\n",
    "# #ge_n = imp50.index[0:25].tolist()\n",
    "# ge_n = lf_scvi_names #ge_columns_train2\n",
    "\n",
    "# nannot = len(annot_columns_train)\n",
    "# lge_n  = len(ge_n)\n",
    "\n",
    "# X2 = X_train0.loc[:,annot_columns_train + ge_n + cla_n]\n",
    "# if len(cla_n)>0:\n",
    "#     i_cl = np.nonzero(X2.columns.isin(cla_n))[0] - nannot\n",
    "#     #X2.loc[:,cla_n] = X2.loc[:,cla_n]+1\n",
    "\n",
    "# X2=X2.iloc[0:X_train0.shape[0]:Dn,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# note: synapse type hierarhy in X(:,classes_columns_train) should be from H_Models!!!\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "t1=time.time()\n",
    "\n",
    "Dn = 10\n",
    "N_bootstraps = 100\n",
    "Dn2 = int(N_bootstraps/Dn) # 100 - should be a number of bootstraps per synapse type !!!\n",
    "Dn3 = int(X_train0.shape[0]/N_bootstraps) # number of synapse_types\n",
    "\n",
    "####\n",
    "####   SELECT X\n",
    "####\n",
    "cla_n = classes_columns_train #['ex_inh']\n",
    "cla_n2 = pd.Series(cla_n)\n",
    "\n",
    "#ge_n = imp50.index[0:25].tolist()\n",
    "#ge_n = lf_scvi_names #ge_columns_train2\n",
    "ge_n = ge_columns_train #imp50.index[0:25].tolist()\n",
    "\n",
    "\n",
    "#X2 = X_train0.loc[:,annot_columns_train + ge_n + cla_n]\n",
    "X2 = X_train0.loc[:,ge_n ]\n",
    "X2_cl = X_train0.loc[:,cla_n ]\n",
    "X2_an = X_train0.loc[:,annot_columns_train ]\n",
    "#i_cl = np.nonzero(X2.columns.isin(cla_n))[0]\n",
    "nannot = len(annot_columns_train)\n",
    "lge_n  = len(ge_n)\n",
    "if len(cla_n)>0:\n",
    "    i_cl = np.nonzero(X2.columns.isin(cla_n))[0] - nannot\n",
    "    #X2.loc[:,cla_n] = X2.loc[:,cla_n]+1\n",
    "    \n",
    "X2=X2.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "X2_cl=X2_cl.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "X2_an=X2_an.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "\n",
    "\n",
    "stp_n = ['A2_20Hz']\n",
    "#  'A2_20Hz',\n",
    "#  'A5_20Hz',\n",
    "#  'A250_20Hz',\n",
    "#  'A1000_20Hz',\n",
    "#  'A2_50Hz',\n",
    "#  'A2_10Hz',\n",
    "#  'A5_50Hz',\n",
    "#  'A5_10Hz',\n",
    "y2 =  X_train0.loc[:,stp_n].iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "##header = ge_n + cla_n + stp_n \n",
    "#header  = ge_n  + stp_n \n",
    "\n",
    "iy0=0\n",
    "mod_index = modf.index #[3] # modf.index #[0, 3, 9, 11, 13, 15]\n",
    "Y_pred = []\n",
    "for i,mdn in enumerate(H_Models.loc[:,'name']):\n",
    "    \n",
    "    # HM_name = 'Ms1c2'\n",
    "    md = H_Models.loc[H_Models.loc[:,'name']==mdn, :]\n",
    "    \n",
    "\n",
    "    #md = modf.loc[i,:]\n",
    "\n",
    "\n",
    "    #mdn = md['name']\n",
    "    print(mdn)\n",
    "\n",
    "    cli = cla_n2.index[cla_n2.isin(md['structure'].values[0])].values\n",
    "\n",
    "    #mdcl2 = np.char.strip(np.array(str.split(md['classes_post'],',')))\n",
    "    #mdcl1 = np.char.strip(np.array(str.split(md['classes'],','))) \n",
    "    print(set(X2_cl[:,cli[0]]))\n",
    "    print(set(X2_cl[:,cli[1]]))\n",
    "    \n",
    "    X3_cl = np.copy(X2_cl)\n",
    "    for cli2 in cli:\n",
    "        n2n=pd.DataFrame(list(set(X2_cl[:,cli2]))).reset_index().set_index(0)\n",
    "        X3_cl[:,cli2] = n2n.loc[X2_cl[:,cli2]].values.ravel()\n",
    "\n",
    "    y_pred = np.zeros((0,y2.shape[1]))\n",
    "    \n",
    "    ncv = 10\n",
    "    n_samp_cv = np.rint(X2.shape[0]/ncv)\n",
    "    samples_all = np.arange(X2.shape[0])\n",
    "    r2cv = np.zeros(ncv)\n",
    "    r4cv = np.zeros(ncv)\n",
    "    for icv in range(ncv): #range(ncv):\n",
    "        samples_test = (np.arange(n_samp_cv) + icv*n_samp_cv).astype(int)\n",
    "        samples_train = np.delete(np.copy(samples_all),samples_test)\n",
    "        X2train, y2train, X2train_cl = X2[samples_train,:], y2[samples_train,:], X3_cl[samples_train,:]\n",
    "        X2test, y2test, X2test_cl = X2[samples_test,:], y2[samples_test,:], X3_cl[samples_test,:]\n",
    "        \n",
    "        model, support = fit_classes_tree(X2train,y2train,X2train_cl,cli,model_type='elastic_net',\n",
    "                                          nmin = 20, alpha=0.2, l1_ratio=0.05)\n",
    "        #model, support = fit_classes_tree(X2,y2,X2_cl,cli,model_type='ridge',nmin = 20, alpha=1, l1_ratio=0.5)\n",
    "        y_pred_i = predict_classes_tree(model,X2test,X2test_cl,cli,nout=y2.shape[1],nmin = 2)\n",
    "        y_pred = np.concatenate([y_pred, y_pred_i],axis=0)\n",
    "        \n",
    "        iy=iy0\n",
    "        if np.var(y2[:,iy])!=0:\n",
    "            nonz = y_pred_i[:,iy]!=0\n",
    "            if np.sum(nonz)>0:\n",
    "                r2cv[icv]=1 - np.mean((y2test[nonz,iy] - y_pred_i[nonz,iy])**2)/np.var(y2[:,iy])\n",
    "                r4cv[icv]=1 - np.mean((y2test[nonz,iy] - y_pred_i[nonz,iy])**2)/np.var(y2test[nonz,iy])\n",
    "    \n",
    "    Y_pred = Y_pred + [y_pred]\n",
    "    \n",
    "    f, ax =plt.subplots(figsize=(16, 10))\n",
    "    ##f, ax = plt.figure()\n",
    "    ##ax = f.add_axes()\n",
    "    \n",
    "    #plt.title(stp_columns[iy]+\", model : \"+mdn+\", cv : \"+str(icv))\n",
    "    plt.title(stp_columns_train[iy]+\", model : \"+mdn)\n",
    "    \n",
    "    #ax.set_title(stp_columns[i]+' '+', subclass out of bag: '+str(i0))\n",
    "    #yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "    plt.plot(y2.ravel(),'ob')\n",
    "    plt.plot(y_pred.ravel(),'xr')\n",
    "    #plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n",
    "    iy=iy0\n",
    "    r2=0\n",
    "    r4=0\n",
    "    if np.var(y2[:,iy])!=0:\n",
    "        r2=1 - np.mean((y2[:,iy] - y_pred[:,iy])**2)/np.var(y2[:,iy])\n",
    "        r4=1 - np.var(y_pred[:,iy])/np.var(y2[:,iy])\n",
    "        \n",
    "    print(\"R**2 = \",r2,\n",
    "          \"  Part of Var = \",r4,\"\\n\",\n",
    "          \"  R**2 cv mean = \", np.mean(r2cv),\n",
    "          \"\\n  R**2 cv unnormed mean = \", np.mean(r4cv),\n",
    "          \"\\n  R**2 cv = \", r2cv,\n",
    "          \"\\n  R**2 cv unnormed = \", r4cv,\n",
    "          \"\\n\\n\")\n",
    "    \n",
    "    \n",
    "t2 =time.time()    \n",
    "print(\"Elapsed time \",t2-t1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.datasets import make_regression\n",
    "\n",
    "from sklearn import preprocessing\n",
    "\n",
    "\n",
    "from sklearn.model_selection import KFold\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import sklearn as sk\n",
    "\n",
    "import time\n",
    "t1 = time.time()\n",
    "\n",
    "do_oob_subclasses=0\n",
    "\n",
    "#X2 = np.copy(X[range(0,X.shape[0],1)])\n",
    "#y2 = np.copy(y[range(0,X.shape[0],1)])\n",
    "#nsel = np.arange(1)\n",
    "#y2=y2[:,0:nannot+nsel.shape[0]]\n",
    "\n",
    "#import models.linear_regr as lir\n",
    "#import models.mean_regr as mir\n",
    "\n",
    "#from models.svm_regr import svm_regr\n",
    "#from models.DT_sklearn_regr import DT_sklearn_regr\n",
    "\n",
    "from importlib import reload\n",
    "#mtree = reload(mtree)\n",
    "#lir=reload(lir)\n",
    "#mir=reload(mir)\n",
    "\n",
    "# ====================\n",
    "# Settings\n",
    "# ====================\n",
    "mode = \"regr\"  # \"clf\" / \"regr\"\n",
    "save_model_tree = True  # save model tree?\n",
    "save_model_tree_predictions = True  # save model tree predictions/explanations?\n",
    "cross_validation = True  # cross-validate model tree?\n",
    "\n",
    "#ML_METHOD = 1\n",
    "\n",
    "# ====================\n",
    "# Load data\n",
    "# ====================\n",
    "\n",
    "\n",
    "# ['all', 'ex_inh']\n",
    "Dn = 5\n",
    "Dn2 = int(100/Dn) # 100 - should be a number of bootstraps per synapse type !!!\n",
    "Dn3 = int(X_train0.shape[0]/100) # number of synapse_types\n",
    "#cla_n = ['ex_in_ex_in__pre',\n",
    "#         'ex_in_ex_in__post'] \n",
    "\n",
    "#cla_n = ['ex_in_ex_in__pre'] \n",
    "cla_n = []\n",
    "\n",
    "#ge_n = imp50.index[0:25].tolist()\n",
    "ge_n = lf_scvi_names #ge_columns_train2\n",
    "\n",
    "nannot = len(annot_columns_train)\n",
    "lge_n  = len(ge_n)\n",
    "\n",
    "X2 = X_train0.loc[:,annot_columns_train + ge_n + cla_n]\n",
    "\n",
    "if len(cla_n)>0:\n",
    "    i_cl = np.nonzero(X2.columns.isin(cla_n))[0] - nannot\n",
    "\n",
    "    #X2.loc[:,cla_n] = X2.loc[:,cla_n]+1\n",
    "\n",
    "X2=X2.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "\n",
    "# select y\n",
    "stp_n = ['A2_20Hz'] #,'A5_20Hz','A250_20Hz']\n",
    "\n",
    "# 'A2_20Hz',\n",
    "#  'A3_20Hz',\n",
    "#  'A4_20Hz',\n",
    "#  'A5_20Hz',\n",
    "#  'A250_20Hz',\n",
    "#  'A1000_20Hz',\n",
    "#  'A1_50Hz',\n",
    "#  'A2_50Hz',\n",
    "#  'A3_50Hz',\n",
    "#  'A4_50Hz',\n",
    "#  'A5_50Hz',\n",
    "#  'A250_50Hz',\n",
    "#  'A1000_50Hz',\n",
    "#  'A1_10Hz',\n",
    "#  'A2_10Hz',\n",
    "#  'A3_10Hz',\n",
    "#  'A4_10Hz',\n",
    "#  'A5_10Hz',\n",
    "#  'A250_10Hz',\n",
    "#  'A1000_10Hz']\n",
    "y2 =  X_train0.loc[:,annot_columns_train+stp_n].iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "\n",
    "#header = ge_n + cla_n + stp_n \n",
    "header = ge_n  + stp_n \n",
    "annot_columns2 = annot_columns_train\n",
    "#\n",
    "#\n",
    "# CROSS-VALIDATION GROUPS\n",
    "#\n",
    "# cv groups\n",
    "all_index = np.arange(X2.shape[0])\n",
    "\n",
    "#\n",
    "#\n",
    "#\n",
    "#\n",
    "# 10-fold cv  : 10% of synapses in data collection order\n",
    "#\n",
    "#\n",
    "test_indexes = []\n",
    "names = []\n",
    "n_splits=10\n",
    "dn = int(X2.shape[0]/10)\n",
    "#df_cv  = pd.DataFrame(np.zeros((n_splits,2)), columns = ['train_index', 'name'])\n",
    "for icv in range(n_splits):\n",
    "    if icv==n_splits:\n",
    "        dn = X2.shape[0] - icv*dn\n",
    "    test_index  = icv*dn + np.arange(dn)\n",
    "    #df_cv.loc[icv,'train_index']=[train_index]\n",
    "    #df_cv.loc[icv,'name']=str(icv)\n",
    "    test_indexes = test_indexes + [test_index]\n",
    "    names = names + [str(icv)]\n",
    "    \n",
    "df_cv  = pd.DataFrame([test_indexes, names]).T\n",
    "df_cv.columns = ['test_index', 'name']\n",
    "\n",
    "\n",
    "#\n",
    "#\n",
    "#\n",
    "#\n",
    "#  10% random  synaptic-types :  cre-lines+subregions (level1)  cv \n",
    "#\n",
    "#\n",
    "part = 10 # % of removed samples\n",
    "n_splits=10\n",
    "test_indexes = []\n",
    "names = []\n",
    "#dn_syn_type = 100\n",
    "Dn3_out = int(Dn3*part/100)\n",
    "idx0 = np.arange(Dn3)\n",
    "idx1 = np.arange(Dn2)\n",
    "for icv in range(n_splits):\n",
    "    idx0_out = np.random.choice(idx0, size=Dn3_out, replace=False, p=None)\n",
    "    test_index = np.zeros((0,))\n",
    "    for io in idx0_out:\n",
    "        #train_index  = train_index + icv*dn + np.arange(dn)\n",
    "        test_index  = np.append(test_index, io*Dn2 + idx1)\n",
    "    test_indexes = test_indexes + [test_index]\n",
    "    names = names + ['syn_types_l0_'+str(icv)+'  '+np.array2string(idx0_out)]\n",
    "    #df_cv.loc[icv,'train_index']=[train_index]\n",
    "    #df_cv.loc[icv,'name']='syn_types_l0_'+str(icv)+'  '+np.array2string(idx0_out)\n",
    "\n",
    "#df_cv1  = pd.DataFrame([train_indexes, names], columns = ['train_index', 'name'])  \n",
    "df_cv1  = pd.DataFrame([test_indexes, names]).T\n",
    "df_cv1.columns = ['test_index', 'name']\n",
    "    \n",
    "df_cv = pd.concat([df_cv, df_cv1],axis=0)    \n",
    "\n",
    "\n",
    "#\n",
    "#\n",
    "#\n",
    "#\n",
    "#  10% random synaptic types : subclasses+regions (level2) cv \n",
    "#\n",
    "#\n",
    "part = 10 # % of removed samples\n",
    "n_splits=10\n",
    "cla_n = ['vse__pre', 'vse__post']\n",
    "groups_index = X_train0.loc[:,cla_n]\n",
    "groups_index=groups_index.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "\n",
    "df_cv1 = cv_random_groups(n_splits, groups_index, part ) \n",
    "\n",
    "df_cv = pd.concat([df_cv, df_cv1],axis=0)  \n",
    "\n",
    "\n",
    "#\n",
    "#\n",
    "#\n",
    "#\n",
    "#  10% random synaptic types : subclasses (level3) cv \n",
    "#\n",
    "#\n",
    "part = 10\n",
    "n_splits=10\n",
    "\n",
    "cla_n = ['ex_ctx__pre',\n",
    " 'ex_ctx__post',\n",
    " 'ex_ec__pre',\n",
    " 'ex_ec__post',\n",
    " 'ex_hipp__pre',\n",
    " 'ex_hipp__post',\n",
    " 'pvalb__pre',\n",
    " 'pvalb__post',\n",
    " 'sst__pre',\n",
    " 'sst__post',\n",
    " 'lamp5__pre',\n",
    " 'lamp5__post',\n",
    " 'vip__pre',\n",
    " 'vip__post',\n",
    " 'cck__pre',\n",
    " 'cck__post']\n",
    "#cla_n = ['cla_cell_type_l3_pre', 'cla_cell_type_l3_post']\n",
    "\n",
    "groups_index = X_train0.loc[:,cla_n]\n",
    "groups_index=groups_index.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "#all_gr = list(set(groups_index))\n",
    "\n",
    "df_cv1 = cv_random_groups(n_splits, groups_index, part ) \n",
    "\n",
    "df_cv = pd.concat([df_cv, df_cv1],axis=0) \n",
    "\n",
    "df_cv = df_cv.reset_index()\n",
    "\n",
    "\n",
    "df_cv2 = df_cv.loc[10:18,:]   ## select for current calculations\n",
    "#\n",
    "#\n",
    "# END CROSS-VALIDATION GROUPS\n",
    "#\n",
    "#\n",
    "\n",
    "if do_oob_subclasses==0:\n",
    "    ## select datasets for training and validation\n",
    "    #X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0) \n",
    "\n",
    "    #X, y = make_regression(n_features=4, n_informative=2,\n",
    "    #                        random_state=0, shuffle=False)\n",
    "    \n",
    "    #n_splits=10\n",
    "    #kf = KFold(n_splits=n_splits, shuffle=False)\n",
    "    i0=0\n",
    "    r6=[]\n",
    "    r7=[]\n",
    "    for ii in df_cv2.index:\n",
    "        i0=i0+1\n",
    "        #train_index, test_index in kf.split(X2):\n",
    "        test_index = df_cv2.loc[ii,'test_index'].astype(int)\n",
    "        train_index = np.delete(all_index.copy(), test_index)\n",
    "        \n",
    "        #i0=i0+1\n",
    "        print('\\n\\n '+ str(i0)+'   '+df_cv2.loc[ii,'name'])\n",
    "        #print(\"TRAIN:\", train_index,\"\\n\", \"TEST:\", test_index)\n",
    "        \n",
    "        X_train, X_test = X2[train_index], X2[test_index]\n",
    "        y_train, y_test = y2[train_index], y2[test_index]\n",
    "        \n",
    "        #X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0) \n",
    "\n",
    "        #\n",
    "        #\n",
    "        #  ML METHOD\n",
    "        #\n",
    "        #\n",
    "        met=0 #ML_METHOD\n",
    "        \n",
    "        if (met<4):\n",
    "            i_X = np.arange(nannot,nannot+lge_n)\n",
    "        else: # use cell types information\n",
    "            i_X = np.arange(nannot,X_train.shape[1])\n",
    "            \n",
    "            \n",
    "        if met==0:\n",
    "            regr = sk.neighbors.KNeighborsRegressor(n_neighbors=8, weights='uniform', \n",
    "                                       algorithm='auto', leaf_size=30, p=2,\n",
    "                                       metric='minkowski', metric_params=None,\n",
    "                                       n_jobs=-1)\n",
    "        \n",
    "        if met==1:\n",
    "            regr = RandomForestRegressor(random_state=2026,max_depth=7,min_samples_leaf=5,\n",
    "                                    n_estimators=200, oob_score=True, n_jobs=-1, max_samples=300)\n",
    "                    # RandomForestRegressor(bootstrap=True, criterion='mse', max_depth=2,\n",
    "            #           max_features='auto', max_leaf_nodes=None,\n",
    "            #           min_impurity_decrease=0.0, min_impurity_split=None,\n",
    "            #           min_samples_leaf=1, min_samples_split=2,\n",
    "            #           min_weight_fraction_leaf=0.0, n_estimators=100, n_jobs=None,\n",
    "            #           oob_score=False, random_state=0, verbose=0, warm_start=False)\n",
    "        elif met==2:\n",
    "            if i0==1:\n",
    "                regr = sk.linear_model.MultiTaskElasticNetCV(cv=5,n_jobs=-1)\n",
    "                regr.fit(X[:,nannot:], y[:,nannot:])\n",
    "                print(\"alpha \",regr.alpha_)\n",
    "                print(\"l1_ratio \",regr.l1_ratio_)\n",
    "        \n",
    "            regr = sk.linear_model.MultiTaskElasticNet(alpha=regr.alpha_, l1_ratio=regr.l1_ratio_,\n",
    "                                                       fit_intercept=True,\n",
    "                                                       normalize=False,\n",
    "                                 max_iter=1000, copy_X=True, tol=0.0001,\n",
    "                                                       warm_start=False, positive=False, \n",
    "                                                       random_state=None, selection='cyclic')\n",
    "        \n",
    "        elif met==3:\n",
    "            X_train[:,i_X] = preprocessing.scale(X_train[:,i_X])\n",
    "            y_train[:,nannot:] = preprocessing.scale(y_train[:,nannot:])\n",
    "            \n",
    "            X_test[:,i_X] = preprocessing.scale(X_test[:,i_X])\n",
    "            y_test[:,nannot:] = preprocessing.scale(y_test[:,nannot:])\n",
    "            \n",
    "            regr = sk.svm.SVR(kernel='poly', degree=3, gamma='scale',\n",
    "                    coef0=0.0, tol=0.001, C=1.0, epsilon=0.1, shrinking=True, \n",
    "                              cache_size=200, verbose=False, max_iter=-1)\n",
    "        \n",
    "        \n",
    "        elif met==4:\n",
    "            # Choose model\n",
    "            model = lir.linear_regr()\n",
    "            model.lambda_classes  = 0.1\n",
    "            model.class_index = i_cl\n",
    "\n",
    "            # Build model tree\n",
    "            model_tree = mtree.ModelTree(model, max_depth=4, min_samples_leaf=10,\n",
    "                                   search_type=\"adaptive\", n_search_grid=20)\n",
    "            \n",
    "        elif met==5:\n",
    "            # Choose model\n",
    "            #model = lir.linear_regr()\n",
    "            model = mir.mean_regr()\n",
    "            model.lambda_classes  = 0.0\n",
    "            model.class_index = i_cl\n",
    "\n",
    "            # Build model tree\n",
    "            model_tree = mtree.RFRegressor(n_estimators=50, max_depth=4,\n",
    "                 min_samples_split=2, min_samples_leaf=5, DTalgo='modeltree',\n",
    "                 model_leaf = model, n_search_grid = 25, n_search_type = \"adaptive\",  \n",
    "                 n_bootstrap=93*5, bootstrap_type='sequential', bootstrap_classes=None,n_bootstrap_groups=93)\n",
    "\n",
    "\n",
    "\n",
    "        if met==3:\n",
    "            #regr.fit(X_train[:,6:], y_train[:,0:].ravel())  \n",
    "            regr.fit(X_train[:,i_X], y_train[:,nannot:].ravel())  \n",
    "            y_pred = regr.predict(X_test[:,i_X]) \n",
    "            y_pred0 = regr.predict(X_train[:,i_X]) \n",
    "        elif met<4:\n",
    "            regr.fit(X_train[:,i_X], y_train[:,nannot:])\n",
    "            y_pred = regr.predict(X_test[:,i_X]) \n",
    "            y_pred0 = regr.predict(X_train[:,i_X]) \n",
    "        elif met==4:\n",
    "            # ====================\n",
    "            # Train model tree\n",
    "            # ====================\n",
    "            print(\"Training model tree with '{}'...\".format(model.__class__.__name__))\n",
    "            model_tree.fit(X_train[:,i_X], y_train[:,nannot:], verbose=True)\n",
    "            y_pred = model_tree.predict(X_test[:,i_X])\n",
    "            y_pred0 = model_tree.predict(X_train[:,i_X])\n",
    "        elif met==5:    \n",
    "            print(\"Training forest of model trees with '{}'...\".format(model.__class__.__name__))\n",
    "            model_tree.fit(X_train[:,i_X], y_train[:,nannot:], verbose=False)\n",
    "            y_pred = model_tree.predict(X_test[:,i_X]) \n",
    "            y_pred0 = model_tree.predict(X_train[:,i_X]) \n",
    "        \n",
    "        if (met==3)|(met==1)|(met==4)|(met==5)|(met==0):\n",
    "            y_pred = y_pred.reshape((-1,1))\n",
    "            y_pred0 = y_pred0.reshape((-1,1))\n",
    "            \n",
    "        \n",
    "                \n",
    "        stpn_test=[s+'_test' for s in stp_n]\n",
    "        stpn_pred=[s+'_pred' for s in stp_n]\n",
    "        \n",
    "        yy = np.concatenate((y_test, y_pred),axis=1)\n",
    "        #yy = pd.DataFrame(yy,columns=annot_columns_train +stpn_test[0:len(nsel)] +stpn_pred[0:len(nsel)])\n",
    "        yy = pd.DataFrame(yy,columns=annot_columns_train +stpn_test +stpn_pred)\n",
    "        yy.loc[:,annot_columns2] = yy.loc[:,annot_columns2].astype(str)\n",
    "        yy1 = yy.set_index(annot_columns2)\n",
    "        ###yy1=yy1.loc[:,stpn_test+stpn_pred] #/yy1.loc[:,'test']*100\n",
    "        yy1=yy1.astype(float).groupby(annot_columns2).mean()\n",
    "        \n",
    "        yy00 = np.concatenate((y_train, y_pred0),axis=1)\n",
    "        yy00 = pd.DataFrame(yy00,columns=annot_columns_train +stpn_test +stpn_pred)\n",
    "        yy00.loc[:,annot_columns2] = yy00.loc[:,annot_columns2].astype(str)\n",
    "        yy100 = yy00.set_index(annot_columns2)\n",
    "        ###yy1=yy1.loc[:,stpn_test+stpn_pred] #/yy1.loc[:,'test']*100\n",
    "        yy100=yy100.astype(float).groupby(annot_columns2).mean()\n",
    "        \n",
    "        #y_test2 = y_test[:,nannot:]\n",
    "        #y_train2 = y_train[:,nannot:]\n",
    "        r3=[]\n",
    "        r5=[]\n",
    "        r300=[]\n",
    "        r500=[]\n",
    "        for i in range(y_pred.shape[1]):\n",
    "            if 1: #(i==np.array([0,2,4])).sum()>0:\n",
    "                f, ax =plt.subplots(figsize=(15, 5))\n",
    "                ##f, ax = plt.figure()\n",
    "                ##ax = f.add_axes()\n",
    "                ##plt.title(stp_columns[i])\n",
    "                ax.set_title(stp_columns_train[i]+' '+', subclass out of bag: '+str(i0))\n",
    "                yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "                plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n",
    "                \n",
    "                \n",
    "                f0, ax0 =plt.subplots(figsize=(15, 5))\n",
    "                ##f, ax = plt.figure()\n",
    "                ##ax = f.add_axes()\n",
    "                ##plt.title(stp_columns[i])\n",
    "                ax.set_title(stp_columns_train[i]+' '+', training erros: '+str(i0))\n",
    "                yy100.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax0)\n",
    "                plt.xticks(np.arange(len(yy100.index)), yy100.index, rotation=90)\n",
    "                \n",
    "                f1, ax1 =plt.subplots(figsize=(15, 5))\n",
    "                ax.set_title(stp_columns_train[0]+' '+', subclass out of bag: '+str(i0))\n",
    "                plt.plot(np.concatenate([y_test[:,[nannot]],y_pred.reshape((-1,1))],axis=1))\n",
    "\n",
    "                f2, ax2 =plt.subplots(figsize=(15, 5))\n",
    "                ax.set_title(stp_columns_train[0]+' '+', subclass out of bag: '+str(i0))\n",
    "                plt.plot(np.concatenate([y_train[:,[nannot]],y_pred0.reshape((-1,1))],axis=1))\n",
    "\n",
    "\n",
    "            #ytest= yy1.loc[:,stpn_test[i]]  \n",
    "            #ypred=yy1.loc[:,stpn_pred[i]]\n",
    "            \n",
    "            ##r2=1 - (ytest.iloc[istr:] - ypred.iloc[istr:]).pow(2).mean()/ytest.iloc[istr:].var()\n",
    "            i2 = nannot+i\n",
    "            if np.var(y_test[:,i2])!=0:\n",
    "                r2=1 - np.mean((y_test[:,i2] - y_pred[:,i])**2)/np.var(y_test[:,i2])\n",
    "            else:\n",
    "                r2=0\n",
    "            r4=1 - np.mean((y_test[:,i2] - y_pred[:,i])**2)/np.var(y2[:,i2])\n",
    "            \n",
    "            r200=1 - np.mean((y_train[:,i2] - y_pred0[:,i])**2)/np.var(y_train[:,i2])\n",
    "            r400=1 - np.mean((y_train[:,i2] - y_pred0[:,i])**2)/np.var(y2[:,i2])    \n",
    "           # print(ytest.iloc[istr:].var())\n",
    "           # print((ytest.iloc[istr:] - ypred.iloc[istr:]).var())\n",
    "           # print(str(i) +' R^2 = '+str(r2))\n",
    "            r3 = r3 + [r2]\n",
    "            r5 = r5 + [r4]\n",
    "            r300 = r300 + [r200]\n",
    "            r500 = r500 + [r400]\n",
    "        print('subclass out of bag: '+str(i0))\n",
    "        \n",
    "        err =  y_test[:,[nannot]]-y_pred.reshape((-1,1))\n",
    "        print('mean error_test**2 = '+str(np.mean(err**2)))\n",
    "        err2 =  y_train[:,[nannot]]-y_pred0.reshape((-1,1))\n",
    "        print('mean error_train**2 = '+str(np.mean(err2**2)))\n",
    "        print('var all = '+str(np.mean(y2[:,[nannot]]**2) -np.mean(y2[:,[nannot]])**2))\n",
    "\n",
    "        \n",
    "        print('mean R^2 total = '+str(np.array(r5).mean()))\n",
    "        print('all R^2 total = '+str(np.array(r5)))\n",
    "        \n",
    "        print('mean R^2 train total = '+str(np.array(r500).mean()))\n",
    "        print('all R^2 train total = '+str(np.array(r500)))\n",
    "        \n",
    "        print('mean fitted var = '+str(np.array(r3).mean()))\n",
    "        print('all fitted var = '+str(np.array(r3)))\n",
    "\n",
    "        \n",
    "        if met==5:\n",
    "            tr=model_tree.models[0]\n",
    "            df = tree_to_df(tr.tree)\n",
    "            \n",
    "            X3, y3 = df_tree_to_data(df)\n",
    "            X3 = pd.DataFrame(X3,columns = ge_n + cla_n + ['leaves'])\n",
    "            y3 = pd.DataFrame(y3,columns = stp_n)\n",
    "            #X3.to_excel('X3.xlsx')\n",
    "            X4 = count_subclasses_in_leaves(X3,cla_n)\n",
    "            print(X4)\n",
    "\n",
    "        r6 = r6 + [np.array(r3).mean()]\n",
    "        r7 = r7 + [np.array(r5).mean()]\n",
    "        \n",
    "        if met==1:\n",
    "            print('out of bag R^2: '+str(regr.oob_score_))\n",
    "\n",
    "    rr = np.array([r6,r7])\n",
    "    rr = pd.DataFrame(rr)\n",
    "\n",
    "\n",
    "    \n",
    "t2 = time.time()\n",
    "#print(regr.feature_importances_)\n",
    "print('Elapsed time: '+str(t2-t1))\n",
    "#print(regr.predict([[0, 0, 0, 0]]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f1, ax1 =plt.subplots(figsize=(15, 5))\n",
    "ax.set_title(stp_columns_train[0]+' '+', subclass out of bag: '+str(i0))\n",
    "plt.plot(np.concatenate([y_test[:,[nannot]],y_pred.reshape((-1,1))],axis=1))\n",
    "\n",
    "\n",
    "#plt.plot(err,color='g')\n",
    "err =  y_test[:,[nannot]]-y_pred.reshape((-1,1))\n",
    "print('mean error_test**2 = '+str(np.mean(err**2)))\n",
    "err2 =  y_train[:,[nannot]]-y_pred0.reshape((-1,1))\n",
    "print('mean error_train**2 = '+str(np.mean(err2**2)))\n",
    "print('var all = '+str(np.mean(y2[:,[nannot]]**2) -np.mean(y2[:,[nannot]])**2))\n",
    "\n",
    "f2, ax2 =plt.subplots(figsize=(15, 5))\n",
    "ax.set_title(stp_columns_train[0]+' '+', subclass out of bag: '+str(i0))\n",
    "plt.plot(np.concatenate([y_train[:,[nannot]],y_pred0.reshape((-1,1))],axis=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Approach for testing leaves linearity for each gene aka Pavilidis 2019"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "predict_classes_tree(model,X2test,X2test_cl,cli,nout=y2.shape[1],nmin = 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(y_pred_i,'ro')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# F-test of regression models with and without subclass splitting (Pavilidis 2019)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import statsmodels.api as sm\n",
    "from statsmodels.formula.api import ols\n",
    "from statsmodels.graphics.api import interaction_plot, abline_plot\n",
    "from statsmodels.stats.anova import anova_lm\n",
    "\n",
    "formula = \"A2_20Hz ~ C(ex_in)  +  pre__Cplx1\"\n",
    "lm = ols(formula, data=X_train0[0:X_train0.shape[0]:10]).fit()\n",
    "#print(lm.summary())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ii=np.nonzero(np.array(ge_columns_train)=='pre__Cplx1')[0]\n",
    "ge_columns_train[ii[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(ge_columns_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(lm.__doc__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load important genes from iRF search\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "imp  = np.load('importance.npy')\n",
    "imp = pd.DataFrame(imp)\n",
    "imp = pd.concat([imp,pd.DataFrame(ge_columns_train)],axis=1)\n",
    "imp.columns = ['importance', 'genes']\n",
    "imp = imp.set_index('genes')\n",
    "\n",
    "imp0  = np.load('importance0.npy')\n",
    "imp0 = pd.DataFrame(imp0)\n",
    "imp0 = pd.concat([imp0,pd.DataFrame(ge_columns_train)],axis=1)\n",
    "imp0.columns = ['importance0', 'genes']\n",
    "imp0 = imp0.set_index('genes')\n",
    "\n",
    "\n",
    "# List of important genes : iRF\n",
    "imp = pd.concat([imp,imp0],axis=1)\n",
    "\n",
    "imp50 = imp.sort_values('importance',ascending=False).iloc[0:49,:]\n",
    "plt.plot(imp.sort_values('importance',ascending=False).loc[:,['importance','importance0']].values)\n",
    "imp50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imp50.index[0:15].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import pixiedust"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### train and test a model tree with linear leaves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pixiedust": {
     "displayParams": {}
    }
   },
   "outputs": [],
   "source": [
    "#%%pixie_debugger\n",
    "\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "d_mtree = \"/home/stepaniu/LearningX/advanced_ML/model_tree\"\n",
    "\n",
    "import sys\n",
    "# insert at 1, 0 is the script path (or '' in REPL)\n",
    "sys.path.insert(1, d_mtree)\n",
    "\n",
    "#from models.linear_regr import linear_regr\n",
    "import models.linear_regr as lir\n",
    "\n",
    "from models.svm_regr import svm_regr\n",
    "from models.DT_sklearn_regr import DT_sklearn_regr\n",
    "\n",
    "from models.modal_clf import modal_clf\n",
    "from models.DT_sklearn_clf import DT_sklearn_clf\n",
    "\n",
    "import os, pickle, csv\n",
    "\n",
    "import src.ModelTree as mtree\n",
    "from src.utils import load_csv_data, cross_validate\n",
    "\n",
    "\n",
    "from importlib import reload\n",
    "mtree = reload(mtree)\n",
    "lir=reload(lir)\n",
    "\n",
    "\n",
    "\n",
    "# ====================\n",
    "# Settings\n",
    "# ====================\n",
    "mode = \"regr\"  # \"clf\" / \"regr\"\n",
    "save_model_tree = True  # save model tree?\n",
    "save_model_tree_predictions = True  # save model tree predictions/explanations?\n",
    "cross_validation = True  # cross-validate model tree?\n",
    "\n",
    "# ====================\n",
    "# Load data\n",
    "# ====================\n",
    "#data_csv_data_filename = os.path.join(d_mtree,\"data\", \"data_clf.csv\")\n",
    "#X, y, header = load_csv_data(data_csv_data_filename, mode=mode, verbose=True)\n",
    "\n",
    "#X2 = np.copy(X[range(0,X.shape[0],10)])\n",
    "#y2 = np.copy(y[range(0,X.shape[0],10)])\n",
    "#nsel = np.arange(1)\n",
    "#y2=y2[:,0:nannot+nsel.shape[0]]\n",
    "\n",
    "\n",
    "# add unspliting class\n",
    "#cl1 = np.zeros((X.shape[0],1))\n",
    "#cl1[0:200]=1\n",
    "#cl1[1000:1200]=2\n",
    "\n",
    "\n",
    "#X2 = np.concatenate([X2,cl1],axis=1)\n",
    "#i_cl = [X2.shape[1]-1]\n",
    "\n",
    "# ['all', 'ex_inh']\n",
    "Dn = 10\n",
    "cla_n = ['ex_in_pre','ex_in_post']\n",
    "\n",
    "\n",
    "ge_n = imp50.index[0:25].tolist()\n",
    "X2 = X_train0.loc[:,ge_n + cla_n]\n",
    "i_cl = np.nonzero(X2.columns.isin(cla_n))[0]\n",
    "X2=X2.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "stp_n = ['A2_20Hz']\n",
    "#  'A2_20Hz',\n",
    "#  'A3_20Hz',\n",
    "#  'A4_20Hz',\n",
    "#  'A5_20Hz',\n",
    "#  'A250_20Hz',\n",
    "#  'A1000_20Hz',\n",
    "#  'A1_50Hz',\n",
    "#  'A2_50Hz',\n",
    "#  'A3_50Hz',\n",
    "#  'A4_50Hz',\n",
    "#  'A5_50Hz',\n",
    "#  'A250_50Hz',\n",
    "#  'A1000_50Hz',\n",
    "#  'A1_10Hz',\n",
    "#  'A2_10Hz',\n",
    "#  'A3_10Hz',\n",
    "#  'A4_10Hz',\n",
    "#  'A5_10Hz',\n",
    "#  'A250_10Hz',\n",
    "#  'A1000_10Hz']\n",
    "y2 =  X_train0.loc[:,stp_n].iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "#header = ge_n + cla_n + stp_n \n",
    "header = ge_n  + stp_n \n",
    "\n",
    "t1=time.time()\n",
    "\n",
    "# Choose model\n",
    "model = lir.linear_regr()\n",
    "model.lambda_classes  = 0.1\n",
    "model.class_index = i_cl\n",
    "\n",
    "# Build model tree\n",
    "model_tree = mtree.ModelTree(model, max_depth=3, min_samples_leaf=10,\n",
    "                       search_type=\"adaptive\", n_search_grid=25)\n",
    "\n",
    "#model_tree = mtree.ModelTree(model, max_depth=4, min_samples_leaf=10,\n",
    "#                       search_type=\"greedy\", n_search_grid=25)\n",
    "\n",
    "# ====================\n",
    "# Train model tree\n",
    "# ====================\n",
    "print(\"Training model tree with '{}'...\".format(model.__class__.__name__))\n",
    "model_tree.fit(X2, y2, verbose=True)\n",
    "\n",
    "t2=time.time() \n",
    "\n",
    "y2_pred = model_tree.predict(X2)\n",
    "\n",
    "t3=time.time() \n",
    "\n",
    "explanations = model_tree.explain(X2, header)\n",
    "loss = model_tree.loss(X2, y2, y2_pred)\n",
    "print(\" -> loss_train={:.6f}\\n\".format(loss))\n",
    "model_tree.export_graphviz(os.path.join(\"output\", \"model_tree\"), header,\n",
    "                           export_png=True, export_pdf=False)\n",
    "\n",
    "print(\"elapsed time : \",[t2-t1,t3-t2])\n",
    "\n",
    "# ====================\n",
    "# Save model tree results\n",
    "# ====================\n",
    "if save_model_tree:\n",
    "    model_tree_filename = os.path.join(\"output\", \"model_tree.p\")\n",
    "    print(\"Saving model tree to '{}'...\".format(model_tree_filename))\n",
    "    pickle.dump(model, open(model_tree_filename, 'wb'))\n",
    "\n",
    "if save_model_tree_predictions:\n",
    "    predictions_csv_filename = os.path.join(\"output\", \"model_tree_pred.csv\")\n",
    "    print(\"Saving mode tree predictions to '{}'\".format(predictions_csv_filename))\n",
    "    with open(predictions_csv_filename, \"w\") as f:\n",
    "        writer = csv.writer(f)\n",
    "        field_names = [\"x2\", \"y2\", \"y2_pred\", \"explanation\"]\n",
    "        writer.writerow(field_names)\n",
    "        for (x_i, y_i, y_pred_i, exp_i) in zip(X2, y2, y2_pred, explanations):\n",
    "            field_values = [x_i, y_i, y_pred_i, exp_i]\n",
    "            writer.writerow(field_values)\n",
    "\n",
    "# ====================\n",
    "# Cross-validate model tree\n",
    "# ====================\n",
    "cross_validation = 0\n",
    "if cross_validation:\n",
    "    t5 = time.time()\n",
    "    cross_validate(model_tree, X2, y2, kfold=5, seed=1)\n",
    "    t6 = time.time()\n",
    "    print(\"time for cv : \", t6-t5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test and train an RF of model_trees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X2.shape"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%%pixie_debugger\n",
    "\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "d_mtree = \"/home/stepaniu/LearningX/advanced_ML/model_tree\"\n",
    "\n",
    "import sys\n",
    "# insert at 1, 0 is the script path (or '' in REPL)\n",
    "sys.path.insert(1, d_mtree)\n",
    "\n",
    "#from models.linear_regr import linear_regr\n",
    "import models.linear_regr as lir\n",
    "\n",
    "from models.svm_regr import svm_regr\n",
    "from models.DT_sklearn_regr import DT_sklearn_regr\n",
    "\n",
    "from models.modal_clf import modal_clf\n",
    "from models.DT_sklearn_clf import DT_sklearn_clf\n",
    "\n",
    "import os, pickle, csv\n",
    "\n",
    "import src.ModelTree as mtree\n",
    "from src.utils import load_csv_data, cross_validate\n",
    "\n",
    "\n",
    "from importlib import reload\n",
    "mtree = reload(mtree)\n",
    "lir=reload(lir)\n",
    "\n",
    "\n",
    "\n",
    "# ====================\n",
    "# Settings\n",
    "# ====================\n",
    "mode = \"regr\"  # \"clf\" / \"regr\"\n",
    "save_model_tree = True  # save model tree?\n",
    "save_model_tree_predictions = True  # save model tree predictions/explanations?\n",
    "cross_validation = True  # cross-validate model tree?\n",
    "\n",
    "# ====================\n",
    "# Load data\n",
    "# ====================\n",
    "#data_csv_data_filename = os.path.join(d_mtree,\"data\", \"data_clf.csv\")\n",
    "#X, y, header = load_csv_data(data_csv_data_filename, mode=mode, verbose=True)\n",
    "\n",
    "#X2 = np.copy(X[range(0,X.shape[0],10)])\n",
    "#y2 = np.copy(y[range(0,X.shape[0],10)])\n",
    "#nsel = np.arange(1)\n",
    "#y2=y2[:,0:nannot+nsel.shape[0]]\n",
    "\n",
    "\n",
    "# add unspliting class\n",
    "#cl1 = np.zeros((X.shape[0],1))\n",
    "#cl1[0:200]=1\n",
    "#cl1[1000:1200]=2\n",
    "\n",
    "\n",
    "#X2 = np.concatenate([X2,cl1],axis=1)\n",
    "#i_cl = [X2.shape[1]-1]\n",
    "\n",
    "# ['all', 'ex_inh']\n",
    "Dn = 10\n",
    "cla_n = ['ex_in_pre','ex_in_post']\n",
    "cla_n = decode_classes2(cla_n,modf3)\n",
    "\n",
    "ge_n = imp50.index[0:50].tolist()\n",
    "X2 = X_train0.loc[:,ge_n + cla_n]\n",
    "i_cl = np.nonzero(X2.columns.isin(cla_n))[0]\n",
    "X2=X2.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "stp_n = ['A2_20Hz']\n",
    "#  'A2_20Hz',\n",
    "#  'A3_20Hz',\n",
    "#  'A4_20Hz',\n",
    "#  'A5_20Hz',\n",
    "#  'A250_20Hz',\n",
    "#  'A1000_20Hz',\n",
    "#  'A1_50Hz',\n",
    "#  'A2_50Hz',\n",
    "#  'A3_50Hz',\n",
    "#  'A4_50Hz',\n",
    "#  'A5_50Hz',\n",
    "#  'A250_50Hz',\n",
    "#  'A1000_50Hz',\n",
    "#  'A1_10Hz',\n",
    "#  'A2_10Hz',\n",
    "#  'A3_10Hz',\n",
    "#  'A4_10Hz',\n",
    "#  'A5_10Hz',\n",
    "#  'A250_10Hz',\n",
    "#  'A1000_10Hz']\n",
    "y2 =  X_train0.loc[:,stp_n].iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "#header = ge_n + cla_n + stp_n \n",
    "header = ge_n  + stp_n \n",
    "\n",
    "t1=time.time()\n",
    "\n",
    "# Choose model\n",
    "model = lir.linear_regr()\n",
    "model.lambda_classes  = 0.1\n",
    "model.class_index = i_cl\n",
    "\n",
    "# Build model tree RF\n",
    "#model_tree = mtree.ModelTree(model, max_depth=4, min_samples_leaf=10,\n",
    "#                       search_type=\"greedy\", n_search_grid=100)\n",
    "model_tree = mtree.RFRegressor(n_estimators=50, max_depth=4,\n",
    "                 min_samples_split=2, min_samples_leaf=10, DTalgo='modeltree',\n",
    "                 model_leaf = model, n_search_grid = 25, n_search_type = \"adaptive\",  \n",
    "                 n_bootstrap=100, bootstrap_type='random', bootstrap_classes=None)\n",
    "\n",
    "# ====================\n",
    "# Train model tree\n",
    "# ====================\n",
    "print(\"Training model tree with '{}'...\".format(model.__class__.__name__))\n",
    "model_tree.fit(X2, y2, verbose=False)\n",
    "\n",
    "t2=time.time() \n",
    "\n",
    "y2_pred = y2 #model_tree.predict(X2)\n",
    "\n",
    "t3=time.time() \n",
    "\n",
    "#explanations = model_tree.explain(X2, header)\n",
    "loss = 0 #model_tree.loss(X2, y2, y2_pred)\n",
    "print(\" -> loss_train={:.6f}\\n\".format(loss))\n",
    "\n",
    "\n",
    "#model_tree.export_graphviz(os.path.join(\"output\", \"model_tree\"), header,\n",
    "#                           export_png=True, export_pdf=False)\n",
    "\n",
    "print(\"elapsed time : \",[t2-t1,t3-t2])\n",
    "\n",
    "# ====================\n",
    "# Save model tree results\n",
    "# ====================\n",
    "if save_model_tree:\n",
    "    model_tree_filename = os.path.join(\"output\", \"model_tree.p\")\n",
    "    print(\"Saving model tree to '{}'...\".format(model_tree_filename))\n",
    "    pickle.dump(model, open(model_tree_filename, 'wb'))\n",
    "\n",
    "if save_model_tree_predictions:\n",
    "    predictions_csv_filename = os.path.join(\"output\", \"model_tree_pred.csv\")\n",
    "    print(\"Saving mode tree predictions to '{}'\".format(predictions_csv_filename))\n",
    "    with open(predictions_csv_filename, \"w\") as f:\n",
    "        writer = csv.writer(f)\n",
    "        field_names = [\"x2\", \"y2\", \"y2_pred\", \"explanation\"]\n",
    "        writer.writerow(field_names)\n",
    "        \n",
    "        # for (x_i, y_i, y_pred_i, exp_i) in zip(X2, y2, y2_pred, explanations):\n",
    "        for (x_i, y_i, y_pred_i, exp_i) in zip(X2, y2, y2_pred, y2_pred):\n",
    "            field_values = [x_i, y_i, y_pred_i, exp_i]\n",
    "            writer.writerow(field_values)\n",
    "\n",
    "# ====================\n",
    "# Cross-validate model tree\n",
    "# ====================\n",
    "cross_validation = 0\n",
    "if cross_validation:\n",
    "    t5 = time.time()\n",
    "    cross_validate(model_tree, X2, y2, kfold=5, seed=1)\n",
    "    t6 = time.time()\n",
    "    print(\"time for cv : \", t6-t5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cross_validation = 0\n",
    "if cross_validation:\n",
    "    t5 = time.time()\n",
    "    cross_validate(model_tree, X2, y2, kfold=5, seed=1)\n",
    "    t6 = time.time()\n",
    "    print(\"time for cv : \", t6-t5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.histogram([1,1,2,2,2,3,3,4,9])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_tree.models[0].__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tree_to_df(tr,node_code='0'):\n",
    "    columns = ['node_code','j_feature','threshold',\n",
    "               'n_samples','y','X','loss']\n",
    "    df = pd.DataFrame(np.zeros((1,7)),columns = columns)\n",
    "    \n",
    "    #if 'tree' in tr:\n",
    "    #        tr=tr.tree\n",
    "    df.loc[0,'loss'] = tr['loss']\n",
    "    df.loc[0,'j_feature'] = tr['j_feature']\n",
    "    df.loc[0,'threshold'] = tr['threshold']\n",
    "    df.loc[0,'n_samples'] = tr['n_samples']\n",
    "    df.loc[0,'node_code'] = node_code\n",
    "    if 'data' in tr:  \n",
    "        #print(tr['data'])\n",
    "        df.loc[0,'X'] = [tr['data'][0]]\n",
    "        df.loc[0,'y'] = [tr['data'][1]]\n",
    "    else:\n",
    "        dfl = tree_to_df(tr['children']['left'], node_code+'0')\n",
    "        dfr = tree_to_df(tr['children']['right'],node_code+'1')\n",
    "\n",
    "        df = pd.concat([df,dfl,dfr],axis=0).reset_index(drop=True)\n",
    "        \n",
    "    return df        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr=model_tree.models[0]\n",
    "df = tree_to_df(tr.tree)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def df_tree_to_data(df):\n",
    "    k=0\n",
    "    isX3=0\n",
    "    y3=0\n",
    "    for i in df.index:\n",
    "        if (np.isnan(df.loc[i,'j_feature'])):\n",
    "            Xi = df.loc[i,'X']\n",
    "            #print(Xi.shape)\n",
    "            y3i = df.loc[i,'y']\n",
    "            X3i  = np.concatenate([Xi,np.ones((Xi.shape[0],1))*i],axis=1)\n",
    "            if isX3==0:\n",
    "                X3=X3i\n",
    "                y3=y3i\n",
    "                isX3=1\n",
    "            else: \n",
    "                X3 = np.concatenate([X3,X3i],axis=0)\n",
    "                y3 = np.concatenate([y3,y3i],axis=0)\n",
    "\n",
    "    return X3, y3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def  count_subclasses_in_leaves(X3,cla_n):\n",
    "    lv=list(set(X3.loc[:,'leaves']))\n",
    "    nx3 = X3.shape[0]\n",
    "    ncl = int(X3.loc[:,cla_n].max().max()+1)\n",
    "    nclan=len(cla_n)\n",
    "\n",
    "    X4 = np.zeros((ncl*nclan,len(lv)))\n",
    "    nx4 = ncl*nclan\n",
    "    classes4 = (np.arange(nx4)/nx4*nclan).astype(int).reshape((nx4,1))\n",
    "    subclasses4 = np.mod(np.arange(nx4),nclan).reshape((nx4,1))\n",
    "    X4 = np.concatenate([classes4, subclasses4, X4],axis=1)\n",
    "    X4 = pd.DataFrame(X4, columns=['classes','subclasses']+lv)\n",
    "\n",
    "    for lvi in lv:\n",
    "        X3i = X3.loc[X3.loc[:,'leaves']==lvi,cla_n]\n",
    "        for icl,cla in enumerate(cla_n):\n",
    "            for cl in range(ncl):\n",
    "                X4.loc[cl+ncl*icl,lvi]=(X3i.loc[:,cla]==cl).sum()\n",
    "                \n",
    "    return X4            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gini_impurity(y):\n",
    "    p2 = 0.0\n",
    "    y_classes = list(set(y))\n",
    "    for c in y_classes:\n",
    "        p2 += (np.sum(y == c) / len(y))**2\n",
    "    loss = p2\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss(self, X, y, y_pred):\n",
    "    if self.lambda_classes==0:\n",
    "        l = mean_squared_error(y, y_pred)\n",
    "    else:\n",
    "        X_cl = X[:,self.class_index]\n",
    "        gi=0\n",
    "        for i in range(len(self.class_index)):\n",
    "            x_cl = X_cl[:,i]\n",
    "            gi = gi + 1-gini_impurity(x_cl[x_cl>0])  # x_cl==0 - should be subclasses not included in class i\n",
    "\n",
    "        gi=gi*np.sum(x_cl>0)/x_cl.shape[0]\n",
    "        l = mean_squared_error(y, y_pred) + self.lambda_classes*gi\n",
    "    return l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X3, y3 = df_tree_to_data(df)\n",
    "X3 = pd.DataFrame(X3,columns = ge_n + cla_n + ['leaves'])\n",
    "y3 = pd.DataFrame(y3,columns = stp_n)\n",
    "\n",
    "X4 = count_subclasses_in_leaves(X3,cla_n)\n",
    "    \n",
    "X3.to_excel('X3.xlsx')\n",
    "X4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare  X_train  with  X3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Crossvalidation - RF model_trees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "columns  = pd.read_hdf('stp_to_ge_all_columns.hdf')\n",
    "ge_columns_train      = columns.loc[:,'ge_columns_train'].values[0]\n",
    "annot_columns_train   = columns.loc[:,'annot_columns_train'].values[0]\n",
    "stp_columns_train     = columns.loc[:,'stp_columns_train'].values[0]\n",
    "classes_columns_train = columns.loc[:,'classes_columns_train'].values[0]\n",
    "#print(X_train.columns)\n",
    "X_train0 = pd.read_hdf('stp_to_ge_all_data.hdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes_columns_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train0.columns[316:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#X2 = X_train0.loc[:,annot_columns_train + ge_n + cla_n]\n",
    "X2 = X_train0.loc[:, annot_columns_train + cla_n ]\n",
    "X2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.datasets import make_regression\n",
    "\n",
    "\n",
    "from sklearn.model_selection import KFold\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import sklearn as sk\n",
    "\n",
    "import time\n",
    "t1 = time.time()\n",
    "\n",
    "do_oob_subclasses=0\n",
    "\n",
    "#X2 = np.copy(X[range(0,X.shape[0],1)])\n",
    "#y2 = np.copy(y[range(0,X.shape[0],1)])\n",
    "#nsel = np.arange(1)\n",
    "#y2=y2[:,0:nannot+nsel.shape[0]]\n",
    "\n",
    "import models.linear_regr as lir\n",
    "import models.mean_regr as mir\n",
    "\n",
    "from models.svm_regr import svm_regr\n",
    "from models.DT_sklearn_regr import DT_sklearn_regr\n",
    "\n",
    "from importlib import reload\n",
    "mtree = reload(mtree)\n",
    "lir=reload(lir)\n",
    "mir=reload(mir)\n",
    "\n",
    "# ====================\n",
    "# Settings\n",
    "# ====================\n",
    "mode = \"regr\"  # \"clf\" / \"regr\"\n",
    "save_model_tree = True  # save model tree?\n",
    "save_model_tree_predictions = True  # save model tree predictions/explanations?\n",
    "cross_validation = True  # cross-validate model tree?\n",
    "\n",
    "# ====================\n",
    "# Load data\n",
    "# ====================\n",
    "nannot = len(annot_columns_train)\n",
    "lge_n  = len(ge_n)\n",
    "\n",
    "# ['all', 'ex_inh']\n",
    "Dn = 1\n",
    "#cla_n = ['ex_in_ex_in__pre',\n",
    "#         'ex_in_ex_in__post'] \n",
    "cla_n = ['ex_in_ex_in__pre'] \n",
    "\n",
    "ge_n = imp50.index[0:25].tolist()\n",
    "X2 = X_train0.loc[:,annot_columns_train + ge_n + cla_n]\n",
    "i_cl = np.nonzero(X2.columns.isin(cla_n))[0] - nannot\n",
    "\n",
    "X2.loc[:,cla_n] = X2.loc[:,cla_n]+1\n",
    "\n",
    "X2=X2.iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "stp_n = ['A2_20Hz']\n",
    "\n",
    "# 'A2_20Hz',\n",
    "#  'A3_20Hz',\n",
    "#  'A4_20Hz',\n",
    "#  'A5_20Hz',\n",
    "#  'A250_20Hz',\n",
    "#  'A1000_20Hz',\n",
    "#  'A1_50Hz',\n",
    "#  'A2_50Hz',\n",
    "#  'A3_50Hz',\n",
    "#  'A4_50Hz',\n",
    "#  'A5_50Hz',\n",
    "#  'A250_50Hz',\n",
    "#  'A1000_50Hz',\n",
    "#  'A1_10Hz',\n",
    "#  'A2_10Hz',\n",
    "#  'A3_10Hz',\n",
    "#  'A4_10Hz',\n",
    "#  'A5_10Hz',\n",
    "#  'A250_10Hz',\n",
    "#  'A1000_10Hz']\n",
    "y2 =  X_train0.loc[:,annot_columns_train+stp_n].iloc[0:X_train0.shape[0]:Dn,:].values\n",
    "#header = ge_n + cla_n + stp_n \n",
    "header = ge_n  + stp_n \n",
    "annot_columns2 = annot_columns_train\n",
    " \n",
    "\n",
    "if do_oob_subclasses==0:\n",
    "    ## select datasets for training and validation\n",
    "    #X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0) \n",
    "\n",
    "    #X, y = make_regression(n_features=4, n_informative=2,\n",
    "    #                        random_state=0, shuffle=False)\n",
    "\n",
    "    kf = KFold(n_splits=10, shuffle=False)\n",
    "    i0=0\n",
    "    r6=[]\n",
    "    r7=[]\n",
    "    for train_index, test_index in kf.split(X2):\n",
    "        i0=i0+1\n",
    "        print('\\n\\n '+ str(i0))\n",
    "        #print(\"TRAIN:\", train_index,\"\\n\", \"TEST:\", test_index)\n",
    "        \n",
    "        X_train, X_test = X2[train_index], X2[test_index]\n",
    "        y_train, y_test = y2[train_index], y2[test_index]\n",
    "        \n",
    "        #X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0) \n",
    "\n",
    "        #\n",
    "        #\n",
    "        #  ML METHOD\n",
    "        #\n",
    "        #\n",
    "        met=5\n",
    "        \n",
    "        if (met<4):\n",
    "            i_X = np.arange(nannot,nannot+lge_n)\n",
    "        else:\n",
    "            i_X = np.arange(nannot,X_train.shape[1])\n",
    "        \n",
    "        if met==1:\n",
    "            regr = RandomForestRegressor(random_state=2026,max_depth=4,min_samples_leaf=10,\n",
    "                                    n_estimators=50, oob_score=True, n_jobs=-1, max_samples=200)\n",
    "                    # RandomForestRegressor(bootstrap=True, criterion='mse', max_depth=2,\n",
    "            #           max_features='auto', max_leaf_nodes=None,\n",
    "            #           min_impurity_decrease=0.0, min_impurity_split=None,\n",
    "            #           min_samples_leaf=1, min_samples_split=2,\n",
    "            #           min_weight_fraction_leaf=0.0, n_estimators=100, n_jobs=None,\n",
    "            #           oob_score=False, random_state=0, verbose=0, warm_start=False)\n",
    "        elif met==2:\n",
    "            if i0==1:\n",
    "                regr = sk.linear_model.MultiTaskElasticNetCV(cv=5,n_jobs=-1)\n",
    "                regr.fit(X[:,nannot:], y[:,nannot:])\n",
    "                print(\"alpha \",regr.alpha_)\n",
    "                print(\"l1_ratio \",regr.l1_ratio_)\n",
    "        \n",
    "            regr = sk.linear_model.MultiTaskElasticNet(alpha=regr.alpha_, l1_ratio=regr.l1_ratio_,\n",
    "                                                       fit_intercept=True,\n",
    "                                                       normalize=False,\n",
    "                                 max_iter=1000, copy_X=True, tol=0.0001,\n",
    "                                                       warm_start=False, positive=False, \n",
    "                                                       random_state=None, selection='cyclic')\n",
    "        \n",
    "        elif met==3:\n",
    "            regr = sk.svm.SVR(kernel='rbf', degree=3, gamma='scale',\n",
    "                    coef0=0.0, tol=0.001, C=3.0, epsilon=0.1, shrinking=True, \n",
    "                              cache_size=200, verbose=False, max_iter=-1)\n",
    "        \n",
    "        \n",
    "        elif met==4:\n",
    "            # Choose model\n",
    "            model = lir.linear_regr()\n",
    "            model.lambda_classes  = 0.1\n",
    "            model.class_index = i_cl\n",
    "\n",
    "            # Build model tree\n",
    "            model_tree = mtree.ModelTree(model, max_depth=4, min_samples_leaf=10,\n",
    "                                   search_type=\"adaptive\", n_search_grid=20)\n",
    "            \n",
    "        elif met==5:\n",
    "            # Choose model\n",
    "            #model = lir.linear_regr()\n",
    "            model = mir.mean_regr()\n",
    "            model.lambda_classes  = 0.0\n",
    "            model.class_index = i_cl\n",
    "\n",
    "            # Build model tree\n",
    "            model_tree = mtree.RFRegressor(n_estimators=50, max_depth=4,\n",
    "                 min_samples_split=2, min_samples_leaf=5, DTalgo='modeltree',\n",
    "                 model_leaf = model, n_search_grid = 25, n_search_type = \"adaptive\",  \n",
    "                 n_bootstrap=93*5, bootstrap_type='sequential', bootstrap_classes=None,n_bootstrap_groups=93)\n",
    "\n",
    "\n",
    "\n",
    "        if met==3:\n",
    "            #regr.fit(X_train[:,6:], y_train[:,0:].ravel())  \n",
    "            regr.fit(X_train[:,i_X], y_train[:,nannot:].ravel())  \n",
    "            y_pred = regr.predict(X_test[:,i_X]) \n",
    "        elif met<4:\n",
    "            regr.fit(X_train[:,i_X], y_train[:,nannot:])\n",
    "            y_pred = regr.predict(X_test[:,i_X]) \n",
    "        elif met==4:\n",
    "            # ====================\n",
    "            # Train model tree\n",
    "            # ====================\n",
    "            print(\"Training model tree with '{}'...\".format(model.__class__.__name__))\n",
    "            model_tree.fit(X_train[:,i_X], y_train[:,nannot:], verbose=True)\n",
    "            y_pred = model_tree.predict(X_test[:,i_X])\n",
    "        elif met==5:    \n",
    "            print(\"Training forest of model trees with '{}'...\".format(model.__class__.__name__))\n",
    "            model_tree.fit(X_train[:,i_X], y_train[:,nannot:], verbose=False)\n",
    "            y_pred = model_tree.predict(X_test[:,i_X]) \n",
    "        \n",
    "        if (met==3)|(met==1)|(met==4)|(met==5):\n",
    "            y_pred = y_pred.reshape((-1,1))\n",
    "        y_test2 = y_test[:,nannot:]\n",
    "        \n",
    "        \n",
    "        yy = np.concatenate((y_test, y_pred),axis=1)\n",
    "        stpn_test=[s+'_test' for s in stp_n]\n",
    "        stpn_pred=[s+'_pred' for s in stp_n]\n",
    "        #yy = pd.DataFrame(yy,columns=annot_columns_train +stpn_test[0:len(nsel)] +stpn_pred[0:len(nsel)])\n",
    "        \n",
    "        yy = pd.DataFrame(yy,columns=annot_columns_train +stpn_test +stpn_pred)\n",
    "        \n",
    "        yy.loc[:,annot_columns2] = yy.loc[:,annot_columns2].astype(str)\n",
    "        yy1 = yy.set_index(annot_columns2)\n",
    "\n",
    "        ###yy1=yy1.loc[:,stpn_test+stpn_pred] #/yy1.loc[:,'test']*100\n",
    "        yy1=yy1.astype(float).groupby(annot_columns2).mean()\n",
    "        r3=[]\n",
    "        r5=[]\n",
    "        for i in range(y_pred.shape[1]):\n",
    "            if (i==np.array([0,2,4])).sum()>0:\n",
    "                f, ax =plt.subplots(figsize=(16, 10))\n",
    "                ##f, ax = plt.figure()\n",
    "                ##ax = f.add_axes()\n",
    "                ##plt.title(stp_columns[i])\n",
    "                ax.set_title(stp_columns_train[i]+' '+', subclass out of bag: '+str(i0))\n",
    "                yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "                plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n",
    "\n",
    "\n",
    "            #ytest= yy1.loc[:,stpn_test[i]]  \n",
    "            #ypred=yy1.loc[:,stpn_pred[i]]\n",
    "            \n",
    "            ##r2=1 - (ytest.iloc[istr:] - ypred.iloc[istr:]).pow(2).mean()/ytest.iloc[istr:].var()\n",
    "            if np.var(y_test2[:,i])!=0:\n",
    "                r2=1 - np.mean((y_test2[:,i] - y_pred[:,i])**2)/np.var(y_test2[:,i])\n",
    "                r4=1 - np.mean((y_test2[:,i] - y_pred[:,i])**2)/np.var(y2[:,nannot+i])\n",
    "           # print(ytest.iloc[istr:].var())\n",
    "           # print((ytest.iloc[istr:] - ypred.iloc[istr:]).var())\n",
    "           # print(str(i) +' R^2 = '+str(r2))\n",
    "            r3 = r3 + [r2]\n",
    "            r5 = r5 + [r4]\n",
    "        print('subclass out of bag: '+str(i0))\n",
    "        print('mean R^2 = '+str(np.array(r3).mean()))\n",
    "        print('all R^2 = '+str(np.array(r3)))\n",
    "        print('mean R^2 total = '+str(np.array(r5).mean()))\n",
    "        print('all R^2 total = '+str(np.array(r5)))\n",
    "        \n",
    "        if met==5:\n",
    "            tr=model_tree.models[0]\n",
    "            df = tree_to_df(tr.tree)\n",
    "            \n",
    "            X3, y3 = df_tree_to_data(df)\n",
    "            X3 = pd.DataFrame(X3,columns = ge_n + cla_n + ['leaves'])\n",
    "            y3 = pd.DataFrame(y3,columns = stp_n)\n",
    "            #X3.to_excel('X3.xlsx')\n",
    "            X4 = count_subclasses_in_leaves(X3,cla_n)\n",
    "            print(X4)\n",
    "\n",
    "        r6 = r6 + [np.array(r3).mean()]\n",
    "        r7 = r7 + [np.array(r5).mean()]\n",
    "        \n",
    "        if met==1:\n",
    "            print('out of bag R^2: '+str(regr.oob_score_))\n",
    "\n",
    "    rr = np.array([r6,r7])\n",
    "    rr = pd.DataFrame(rr)\n",
    "\n",
    "\n",
    "    \n",
    "t2 = time.time()\n",
    "#print(regr.feature_importances_)\n",
    "print('Elapsed time: '+str(t2-t1))\n",
    "#print(regr.predict([[0, 0, 0, 0]]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#%matplotlib notebook\n",
    "#%matplotlib ipympl\n",
    "\n",
    "#import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import mpld3\n",
    "mpld3.enable_notebook()\n",
    "plt.rcParams['figure.figsize']=[14,7]\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "A=plt.imread('output/model_tree.png')\n",
    "#fig=plt.figure(figsize=(18, 15), facecolor='w', edgecolor='k')\n",
    "#f, ax =plt.subplots(figsize=(15, 10))\n",
    "plt.imshow(A)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%javascript\n",
    "IPython.OutputArea.prototype._should_scroll = function(lines){\n",
    "      return false;\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "\n",
    "%matplotlib widget\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (20,10)\n",
    "#plt.plot([1,2,3])\n",
    "\n",
    "A=plt.imread('output/model_tree.png')\n",
    "fig, ax = plt.subplots() #(figsize=(5,10))\n",
    "im = ax.imshow(A)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \n",
    "sb0 = np.arange(len(classes))\n",
    "mods  = []\n",
    "mods2  = []\n",
    "for i in range(classes.shape[1]):\n",
    "    sbi = classes.loc['subclasses',i]\n",
    "    print(sbi)\n",
    "        \n",
    "    subclasses3 = subclasses2[sbi]\n",
    "    X_train, y_train =  make_data(X,y,ge_data_subcl,[subclasses3,[]])   \n",
    "    regr = sk.linear_model.MultiTaskElasticNet(alpha=0.1, l1_ratio=0.5, fit_intercept=True,\n",
    "                                                   normalize=False,max_iter=1000, copy_X=True, tol=0.0001, \n",
    "                                        warm_start=False, random_state=None, selection='cyclic')\n",
    "    regr.fit(X_train[:,nannot:], y_train[:,nannot:])\n",
    "    \n",
    "    \n",
    "    X_train2, y_train2 =  make_data(X,y,ge_data_subcl,[[],subclasses3]) \n",
    "    regr2 = sk.linear_model.MultiTaskElasticNet(alpha=0.1, l1_ratio=0.5, fit_intercept=True,\n",
    "                                                   normalize=False,max_iter=1000, copy_X=True, tol=0.0001, \n",
    "                                        warm_start=False, random_state=None, selection='cyclic')\n",
    "    regr2.fit(X_train2[:,nannot:], y_train2[:,nannot:])\n",
    "    \n",
    "    mods = mods + [regr]\n",
    "    mods2 = mods2 + [regr2]\n",
    "    #X_test = X[in_subcl,:]\n",
    "    #y_test = y[in_subcl,:]"
   ]
  },
  {
   "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": [
    "print(X.shape)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "i0=1\n",
    "in_subcl = ge_data_subcl.loc[:,'subclass_pre']==subclasses2[i0]\n",
    "in_subcl = in_subcl|(ge_data_subcl.loc[:,'subclass_post']==subclasses2[i0])\n",
    "\n",
    "\n",
    "\n",
    "X_train = np.delete(X, np.nonzero(in_subcl)[0],axis=0)\n",
    "y_train = np.delete(y, np.nonzero(in_subcl)[0],axis=0)\n",
    "X_test = X[in_subcl,:]\n",
    "y_test = y[in_subcl,:]\n",
    "\n",
    "print(X_train.shape)\n",
    "print(X_test.shape)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#pd.DataFrame(y_test).to_excel('y_test.xlsx')\n",
    "pd.DataFrame(X_h).to_excel('xh.xlsx')\n",
    "pd.DataFrame(X).to_excel('x.xlsx')\n",
    "ge_data_subcl.to_excel('ge_subcl.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred.shape\n",
    "y_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # elnet a=0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # elastic net a=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # cv 300"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # cv 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # cv 8000"
   ]
  },
  {
   "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": [
    "import sklearn\n",
    "print(sklearn.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.datasets import make_regression\n",
    "\n",
    "\n",
    "from sklearn.model_selection import KFold\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import sklearn as sk\n",
    "\n",
    "import time\n",
    "t1 = time.time()\n",
    "\n",
    "do_oob_subclasses=0\n",
    "\n",
    "X2 = np.copy(X[range(0,X.shape[0],1)])\n",
    "y2 = np.copy(y[range(0,X.shape[0],1)])\n",
    "nsel = np.arange(1)\n",
    "y2=y2[:,0:nannot+nsel.shape[0]]\n",
    "\n",
    "if do_oob_subclasses==0:\n",
    "    ## select datasets for training and validation\n",
    "    #X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0) \n",
    "\n",
    "    #X, y = make_regression(n_features=4, n_informative=2,\n",
    "    #                        random_state=0, shuffle=False)\n",
    "\n",
    "    kf = KFold(n_splits=10, shuffle=False)\n",
    "    i0=0\n",
    "    r6=[]\n",
    "    r7=[]\n",
    "    for train_index, test_index in kf.split(X2):\n",
    "        i0=i0+1\n",
    "        print('\\n\\n '+ str(i0))\n",
    "        #print(\"TRAIN:\", train_index,\"\\n\", \"TEST:\", test_index)\n",
    "        \n",
    "        X_train, X_test = X2[train_index], X2[test_index]\n",
    "        y_train, y_test = y2[train_index], y2[test_index]\n",
    "        \n",
    "        #X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0) \n",
    "\n",
    "\n",
    "        met=1\n",
    "        \n",
    "        if met==1:\n",
    "            regr = RandomForestRegressor(random_state=2026,max_depth=8,min_samples_leaf=1,\n",
    "                                    n_estimators=2000, oob_score=True, n_jobs=-1, max_samples=500)\n",
    "                    # RandomForestRegressor(bootstrap=True, criterion='mse', max_depth=2,\n",
    "            #           max_features='auto', max_leaf_nodes=None,\n",
    "            #           min_impurity_decrease=0.0, min_impurity_split=None,\n",
    "            #           min_samples_leaf=1, min_samples_split=2,\n",
    "            #           min_weight_fraction_leaf=0.0, n_estimators=100, n_jobs=None,\n",
    "            #           oob_score=False, random_state=0, verbose=0, warm_start=False)\n",
    "        elif met==2:\n",
    "            if i0==1:\n",
    "                regr = sk.linear_model.MultiTaskElasticNetCV(cv=5,n_jobs=-1)\n",
    "                regr.fit(X[:,nannot:], y[:,nannot:])\n",
    "                print(\"alpha \",regr.alpha_)\n",
    "                print(\"l1_ratio \",regr.l1_ratio_)\n",
    "        \n",
    "            regr = sk.linear_model.MultiTaskElasticNet(alpha=regr.alpha_, l1_ratio=regr.l1_ratio_, fit_intercept=True,\n",
    "                                                       normalize=False,\n",
    "                                 max_iter=1000, copy_X=True, tol=0.0001, warm_start=False, positive=False, random_state=None, selection='cyclic')\n",
    "        \n",
    "        elif met==3:\n",
    "            regr = sk.svm.SVR(kernel='rbf', degree=3, gamma='scale',\n",
    "                    coef0=0.0, tol=0.001, C=3.0, epsilon=0.1, shrinking=True, cache_size=200, verbose=False, max_iter=-1)\n",
    "        \n",
    "        if met==3:\n",
    "            #regr.fit(X_train[:,6:], y_train[:,0:].ravel())  \n",
    "            regr.fit(X_train[:,nannot:], y_train[:,nannot:].ravel())  \n",
    "        else:\n",
    "            regr.fit(X_train[:,nannot:], y_train[:,nannot:])\n",
    "\n",
    "\n",
    "        y_pred = regr.predict(X_test[:,nannot:]) \n",
    "        if (met==3)|(met==1):\n",
    "            y_pred = y_pred.reshape((-1,1))\n",
    "        y_test2 = y_test[:,nannot:]\n",
    "        \n",
    "        \n",
    "        yy = np.concatenate((y_test, y_pred),axis=1)\n",
    "        stpn_test=[s+'_test' for s in stp_columns]\n",
    "        stpn_pred=[s+'_pred' for s in stp_columns]\n",
    "        yy = pd.DataFrame(yy,columns=annot_columns +stpn_test[0:len(nsel)] +stpn_pred[0:len(nsel)])\n",
    "        yy.loc[:,annot_columns2] = yy.loc[:,annot_columns2].astype(str)\n",
    "        yy1 = yy.set_index(annot_columns2)\n",
    "\n",
    "        ###yy1=yy1.loc[:,stpn_test+stpn_pred] #/yy1.loc[:,'test']*100\n",
    "        yy1=yy1.astype(float).groupby(annot_columns2).mean()\n",
    "        r3=[]\n",
    "        r5=[]\n",
    "        for i in range(y_pred.shape[1]):\n",
    "            if (i==np.array([0,2,4])).sum()>0:\n",
    "                f, ax =plt.subplots(figsize=(16, 10))\n",
    "                ##f, ax = plt.figure()\n",
    "                ##ax = f.add_axes()\n",
    "                ##plt.title(stp_columns[i])\n",
    "                ax.set_title(stp_columns[i]+' '+', subclass out of bag: '+str(i0))\n",
    "                yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "                plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n",
    "\n",
    "\n",
    "            #ytest= yy1.loc[:,stpn_test[i]]  \n",
    "            #ypred=yy1.loc[:,stpn_pred[i]]\n",
    "            \n",
    "            ##r2=1 - (ytest.iloc[istr:] - ypred.iloc[istr:]).pow(2).mean()/ytest.iloc[istr:].var()\n",
    "            if np.var(y_test2[:,i])!=0:\n",
    "                r2=1 - np.mean((y_test2[:,i] - y_pred[:,i])**2)/np.var(y_test2[:,i])\n",
    "                r4=1 - np.mean((y_test2[:,i] - y_pred[:,i])**2)/np.var(y[:,nannot+i])\n",
    "           # print(ytest.iloc[istr:].var())\n",
    "           # print((ytest.iloc[istr:] - ypred.iloc[istr:]).var())\n",
    "           # print(str(i) +' R^2 = '+str(r2))\n",
    "            r3 = r3 + [r2]\n",
    "            r5 = r5 + [r4]\n",
    "        print('subclass out of bag: '+str(i0))\n",
    "        print('mean R^2 = '+str(np.array(r3).mean()))\n",
    "        print('all R^2 = '+str(np.array(r3)))\n",
    "        print('mean R^2 total = '+str(np.array(r5).mean()))\n",
    "        print('all R^2 total = '+str(np.array(r5)))\n",
    "\n",
    "        r6 = r6 + [np.array(r3).mean()]\n",
    "        r7 = r7 + [np.array(r5).mean()]\n",
    "        \n",
    "        if met==1:\n",
    "            print('out of bag R^2: '+str(regr.oob_score_))\n",
    "\n",
    "    rr = np.array([r6,r7])\n",
    "    rr = pd.DataFrame(rr)\n",
    "else:\n",
    "    \n",
    "    oobs = []\n",
    "    errors = []\n",
    "    annot_columns2 = annot_columns\n",
    "    \n",
    "    r6=[]\n",
    "    r7=[]\n",
    "    for i0 in range(len(subclasses2)):\n",
    "        in_subcl = ge_data_subcl.loc[:,'subclass_pre']==subclasses2[i0]\n",
    "        in_subcl = in_subcl|(ge_data_subcl.loc[:,'subclass_post']==subclasses2[i0])\n",
    "        in_subcl = np.nonzero(in_subcl.values)[0]\n",
    "\n",
    "        X_train = np.delete(X, in_subcl,axis=0)\n",
    "        y_train = np.delete(y, in_subcl,axis=0)\n",
    "        #X_train = X[:,:]\n",
    "        #y_train = y[:,:]\n",
    "        \n",
    "        X_test = X[in_subcl,:]\n",
    "        y_test = y[in_subcl,:]\n",
    "        \n",
    "        do_this = 1\n",
    "\n",
    "        X_train, X_test0, y_train, y_test0 = train_test_split(X, y, test_size=0.25, random_state=0) \n",
    "        \n",
    "        print(X_train.shape)\n",
    "        print(X_test.shape)\n",
    "        \n",
    "        if do_this==1:\n",
    "            #if i0=0:\n",
    "            regr = RandomForestRegressor(random_state=2026,\n",
    "                                    n_estimators=500, oob_score=True, n_jobs=-1, max_samples=300)\n",
    "                \n",
    "        else:\n",
    "            regr = RandomForestRegressor(random_state=2026,\n",
    "                                    n_estimators=500, oob_score=True, n_jobs=-1)\n",
    "        \n",
    "        regr.fit(X_train[:,nannot:], y_train[:,nannot:])  \n",
    "        \n",
    "        y_pred = regr.predict(X_test[:,nannot:]) \n",
    "        y_test2 = y_test[:,nannot:]\n",
    "        \n",
    "        \n",
    "        yy = np.concatenate((y_test, y_pred),axis=1)\n",
    "        stpn_test=[s+'_test' for s in stp_columns]\n",
    "        stpn_pred=[s+'_pred' for s in stp_columns]\n",
    "        yy = pd.DataFrame(yy,columns=annot_columns +stpn_test +stpn_pred)\n",
    "        yy.loc[:,annot_columns2] = yy.loc[:,annot_columns2].astype(str)\n",
    "        yy1 = yy.set_index(annot_columns2)\n",
    "\n",
    "        ###yy1=yy1.loc[:,stpn_test+stpn_pred] #/yy1.loc[:,'test']*100\n",
    "        yy1=yy1.astype(float).groupby(annot_columns2).mean()\n",
    "        r3=[]\n",
    "        r5=[]\n",
    "        for i in range(y_pred.shape[1]):\n",
    "            if (i==np.array([0,1,9])).sum()>0:\n",
    "                f, ax =plt.subplots(figsize=(16, 10))\n",
    "                ##f, ax = plt.figure()\n",
    "                ##ax = f.add_axes()\n",
    "                ##plt.title(stp_columns[i])\n",
    "                ax.set_title(stp_columns[i]+' '+', subclass out of bag: '+subclasses2[i0])\n",
    "                yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "                plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n",
    "\n",
    "\n",
    "            #ytest= yy1.loc[:,stpn_test[i]]  \n",
    "            #ypred=yy1.loc[:,stpn_pred[i]]\n",
    "            \n",
    "            ##r2=1 - (ytest.iloc[istr:] - ypred.iloc[istr:]).pow(2).mean()/ytest.iloc[istr:].var()\n",
    "            if np.var(y_test2[:,i])!=0:\n",
    "                r2=1 - np.mean((y_test2[:,i] - y_pred[:,i])**2)/np.var(y_test2[:,i])\n",
    "                r4=1 - np.mean((y_test2[:,i] - y_pred[:,i])**2)/np.var(y[:,nannot+i])\n",
    "           # print(ytest.iloc[istr:].var())\n",
    "           # print((ytest.iloc[istr:] - ypred.iloc[istr:]).var())\n",
    "           # print(str(i) +' R^2 = '+str(r2))\n",
    "            r3 = r3 + [r2]\n",
    "            r5 = r5 + [r4]\n",
    "        print('subclass out of bag: '+subclasses2[i0])\n",
    "        print('mean R^2 = '+str(np.array(r3).mean()))\n",
    "        print('all R^2 = '+str(np.array(r3)))\n",
    "        print('mean R^2 total = '+str(np.array(r5).mean()))\n",
    "        print('all R^2 total = '+str(np.array(r5)))\n",
    "\n",
    "        r6 = r6 + [np.array(r3).mean()]\n",
    "        r7 = r7 + [np.array(r5).mean()]\n",
    "        oobs = oobs + [regr.oob_score_]\n",
    "    print('out of bag R^2: '+str(oobs))\n",
    "    \n",
    "    rr = np.array([subclasses2,r6,r7])\n",
    "    rr = pd.DataFrame(rr)\n",
    "    \n",
    "t2 = time.time()\n",
    "#print(regr.feature_importances_)\n",
    "print('Elapsed time: '+str(t2-t1))\n",
    "#print(regr.predict([[0, 0, 0, 0]]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TO DO : CHECK R^2 for ca and dg, ee cases!!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(regr.alpha_)\n",
    "print(regr.__dict__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # out of 10 cross val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # all data, 300 bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # all data, 1000 bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # cross val, 1000 bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # cross val, 300 bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr # cross val, all bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr  # cross val, 93 classes bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr  # all data, classes bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ca1=1.6\n",
    "ca2=2\n",
    "k=2.79\n",
    "# k=1.09\n",
    "\n",
    "cc1=(k**4 + ca1**4)/ca1**4*ca2**4/(k**4+ca2**4)\n",
    "\n",
    "k=1.09\n",
    "cc2=(k**4 + ca1**4)/ca1**4*ca2**4/(k**4+ca2**4)\n",
    "\n",
    "print(str([cc1,cc2,0.5*(cc1+cc2)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear regression model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Linear regression\n",
    "\n",
    "t1 = time.time()\n",
    "\n",
    "#regr = RandomForestRegressor(max_depth=15, random_state=0,\n",
    "#                            n_estimators=100)\n",
    "#regr = RandomForestRegressor(random_state=2026,\n",
    "#                            n_estimators=500)\n",
    "import matplotlib.pyplot as plt \n",
    "from sklearn import linear_model\n",
    "regr2 = linear_model.Lasso(alpha=0.01)\n",
    "regr2.fit(X_train[:,nannot:], y_train[:,nannot:])  \n",
    "#Lasso(alpha=0.1, copy_X=True, fit_intercept=True, max_iter=1000,\n",
    "#   normalize=False, positive=False, precompute=False, random_state=None,\n",
    "#   selection='cyclic', tol=0.0001, warm_start=False)\n",
    "\n",
    "\n",
    "#regr.fit(X_train[:,6:], y_train[:,0:].ravel())  \n",
    "#regr.fit(X_train[:,nannot:], y_train[:,nannot:])  \n",
    "# RandomForestRegressor(bootstrap=True, criterion='mse', max_depth=2,\n",
    "#           max_features='auto', max_leaf_nodes=None,\n",
    "#           min_impurity_decrease=0.0, min_impurity_split=None,\n",
    "#           min_samples_leaf=1, min_samples_split=2,\n",
    "#           min_weight_fraction_leaf=0.0, n_estimators=100, n_jobs=None,\n",
    "#           oob_score=False, random_state=0, verbose=0, warm_start=False)\n",
    "t2 = time.time()\n",
    "#print(regr.feature_importances_)\n",
    "print('Elapsed time: '+str(t2-t1))\n",
    "#print(regr.predict([[0, 0, 0, 0]]))\n",
    "\n",
    "# some analysis of results\n",
    "y_predl = regr2.predict(X_test[:,nannot:])  \n",
    "if len(y_predl.shape)==1:\n",
    "    y_predl=y_predl.reshape((len(y_predl),1))\n",
    "    \n",
    "yyl = np.concatenate((y_test, y_predl),axis=1)\n",
    "stpn_test=[s+'_test' for s in stp_columns]\n",
    "stpn_pred=[s+'_pred' for s in stp_columns]\n",
    "yyl = pd.DataFrame(yyl,columns=annot_columns +stpn_test +stpn_pred)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# make input and output stp labels dataset with errors\n",
    "stpn_error = [s+'_error' for s in stp_columns]\n",
    "yel = pd.DataFrame(yyl.loc[:,stpn_pred].values - yyl.loc[:,stpn_test].values, columns=stpn_error)\n",
    "yyl = pd.concat([yyl, yel.abs()], axis=1)\n",
    "\n",
    "#yy.set_index(['pre','post']).groupby()\n",
    "#yy.set_index(['pre','post']).groupby(['pre','post']).mean()\n",
    "\n",
    "# plot errors\n",
    "annot_columns2 = annot_columns\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(16, 10))\n",
    "\n",
    "#yy1l=yyl.set_index(annot_columns2)\n",
    "yyl.loc[:,annot_columns2] = yyl.loc[:,annot_columns2].astype(str)\n",
    "yy1l = yyl.set_index(annot_columns2)\n",
    "yy1l.loc[:,stpn_error]=yy1l.loc[:,stpn_error].values/(yy1l.loc[:,stpn_test].values+1e-15)*100\n",
    "yy1l=yy1l.astype(float).groupby(annot_columns2).mean()\n",
    "yy1l.loc[:,stpn_error].plot(ax=ax)\n",
    "plt.xticks(np.arange(len(yy1l.index)), yy1l.index, rotation=90)\n",
    "plt.title('errors of An/A1 predicted by lasso linear regrassion, %')\n",
    "plt.ylim((0,100))    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## plot test vs train STP data\n",
    "\n",
    "1. plot all parameters\n",
    "2. additional cross validation? - better estimate of generalization error?\n",
    "3. check dependence on rf parameters\n",
    "4. dependence on % of training data\n",
    "5. lists of most significant genes\n",
    "6. check gs assignment - plot some figures\n",
    "7. project on hipp-ds -> predict stp, compare\n",
    "8. an:a1 vs TM-par\n",
    "## \n",
    "## new (sep 2019):\n",
    "9. compare with linear fit\n",
    "10. iRFcompare small gs predictions\n",
    "11. compare with factors predicting neuron types, predict n-types using small ds\n",
    "12. R^2\n",
    "13. compare with list of STP molecular factors\n",
    "14. relative errors\n",
    "15. some graphs for particular factors?\n",
    "16. iRF\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#annot_columns +stpn_test +stpn_pred\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(X_train.shape)\n",
    "print(y_train.shape)\n",
    "y_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RF case\n",
    "y_pred = regr.predict(X_test[:,nannot:]) \n",
    "if len(y_pred.shape)==1:\n",
    "    y_pred=y_pred.reshape((len(y_pred),1))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "from sklearn import metrics\n",
    "\n",
    "#print('Mean Absolute Error:', metrics.mean_absolute_error(y_test[:,2:].ravel(), y_pred))  \n",
    "#print('Mean Squared Error:', metrics.mean_squared_error(y_test[:,2:].ravel(), y_pred))  \n",
    "#print('Root Mean Squared Error:', np.sqrt(metrics.mean_squared_error(y_test[:,2:].ravel(), y_pred))) \n",
    "\n",
    "print('Mean Absolute Error:', metrics.mean_absolute_error(y_test[:,nannot:], y_pred))  \n",
    "print('Mean Squared Error:', metrics.mean_squared_error(y_test[:,nannot:], y_pred))  \n",
    "print('Root Mean Squared Error:', np.sqrt(metrics.mean_squared_error(y_test[:,nannot:], y_pred))) \n",
    "\n",
    "\n",
    "#y_pred2 = y_pred.reshape((len(y_pred),nstp))\n",
    "yy = np.concatenate((y_test, y_pred),axis=1)\n",
    "stpn_test=[s+'_test' for s in stp_columns]\n",
    "stpn_pred=[s+'_pred' for s in stp_columns]\n",
    "yy = pd.DataFrame(yy,columns=annot_columns +stpn_test +stpn_pred)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# make input and output stp labels dataset with errors\n",
    "stpn_error = [s+'_error' for s in stp_columns]\n",
    "ye = pd.DataFrame(yy.loc[:,stpn_pred].values - yy.loc[:,stpn_test].values, columns=stpn_error)\n",
    "yy = pd.concat([yy, ye.abs()], axis=1)\n",
    "\n",
    "#yy.set_index(['pre','post']).groupby()\n",
    "#yy.set_index(['pre','post']).groupby(['pre','post']).mean()\n",
    "\n",
    "# plot errors\n",
    "annot_columns2 = annot_columns\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(16, 10))\n",
    "\n",
    "#yy1=yy.set_index(annot_columns2)\n",
    "yy.loc[:,annot_columns2] = yy.loc[:,annot_columns2].astype(str)\n",
    "yy1 = yy.set_index(annot_columns2)\n",
    "yy1.loc[:,stpn_error]=yy1.loc[:,stpn_error].values/(yy1.loc[:,stpn_test].values+1e-15)*100\n",
    "yy1=yy1.astype(float).groupby(annot_columns2).mean()\n",
    "yy1.loc[:,stpn_error].plot(ax=ax)\n",
    "plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n",
    "plt.title('errors of An/A1 predicted by RF, %')\n",
    "plt.ylim((0,100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yy.to_excel('all_tested_rf_280gs_with_PPR.xlsx')\n",
    "yy.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Linear model case\n",
    "\n",
    "#plot predicted and real (test) stp data\n",
    "#plt.figure()\n",
    "#yy1l=yyl.set_index(annot_columns2)\n",
    "\n",
    "#yyl.loc[:,annot_columns2] = yyl.loc[:,annot_columns2].astype(str)\n",
    "yy1l = yyl.set_index(annot_columns2)\n",
    "#yy1=yy1.loc[:,stpn_test+stpn_pred] #/yy1.loc[:,'test']*100\n",
    "yy1l=yy1l.astype(float).groupby(annot_columns2).mean()\n",
    "\n",
    "r3=[]\n",
    "for i in range(nstp):\n",
    "    f, ax =plt.subplots(figsize=(16, 10))\n",
    "    #f, ax = plt.figure()\n",
    "    #ax = f.add_axes()\n",
    "    #plt.title(stp_columns[i])\n",
    "    ax.set_title(stp_columns[i])\n",
    "    yy1l.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "    plt.xticks(np.arange(len(yy1l.index)), yy1l.index, rotation=90)\n",
    "    if i==0:\n",
    "        istr=0\n",
    "    else:\n",
    "        istr=0 #11 - if ppr included\n",
    "    \n",
    "    ytest= yy1l.loc[:,stpn_test[i]]  \n",
    "    ypred=yy1l.loc[:,stpn_pred[i]]\n",
    "    \n",
    "    if ytest.iloc[istr:].var()!=0:\n",
    "        r2=1 - (ytest.iloc[istr:] - ypred.iloc[istr:]).pow(2).mean()/(ytest.iloc[istr:].var() )\n",
    "    print(str(i) +' R^2 = '+str(r2))\n",
    "    r3 = r3 + [r2]\n",
    "\n",
    "print('mean R^2 = '+str(np.array(r3).mean()))\n",
    "#print(yy1.mean())\n",
    "yy1l.loc[:,stpn_test +stpn_pred]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# RF model case\n",
    "\n",
    "#plot predicted and real (test) stp data\n",
    "#plt.figure()\n",
    "\n",
    "#yy.loc[:,annot_columns2] = yy.loc[:,annot_columns2].astype(str)\n",
    "#yy1 = yy.set_index(annot_columns2)\n",
    "\n",
    "##yy1=yy1.loc[:,stpn_test+stpn_pred] #/yy1.loc[:,'test']*100\n",
    "yy1=yy1.astype(float).groupby(annot_columns2).mean()\n",
    "r3=[]\n",
    "for i in range(0,nstp):\n",
    "    f, ax =plt.subplots(figsize=(16, 10))\n",
    "    #f, ax = plt.figure()\n",
    "    #ax = f.add_axes()\n",
    "    #plt.title(stp_columns[i])\n",
    "    ax.set_title(stp_columns[i])\n",
    "    yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "    plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n",
    "    if i==0:\n",
    "        istr=0\n",
    "    else:\n",
    "        istr=0 # 11 - if ppr data excluded\n",
    "        \n",
    "    ytest= yy1.loc[:,stpn_test[i]]  \n",
    "    ypred=yy1.loc[:,stpn_pred[i]]\n",
    "    #r2=1 - (ytest.iloc[istr:] - ypred.iloc[istr:]).pow(2).mean()/ytest.iloc[istr:].var()\n",
    "    if ytest.iloc[istr:].var()!=0:\n",
    "        r2=1 - (ytest.iloc[istr:] - ypred.iloc[istr:]).var()/ytest.iloc[istr:].var()\n",
    "    print(ytest.iloc[istr:].var())\n",
    "    print((ytest.iloc[istr:] - ypred.iloc[istr:]).var())\n",
    "    print(str(i) +' R^2 = '+str(r2))\n",
    "    r3 = r3 + [r2]\n",
    "\n",
    "print('mean R^2 = '+str(np.array(r3).mean()))\n",
    "\n",
    "#f, ax =plt.subplots(figsize=(10, 10))\n",
    "#yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "\n",
    "#print(yy1.mean())\n",
    "yy1.loc[:,stpn_test +stpn_pred]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#yy1.to_excel('results_RF_10gs_without_PPR.xlsx')\n",
    "yy1.to_excel('results_RF_10gs_with_PPR.xlsx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Look at most important genes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#linear model case\n",
    "\n",
    "# set of most significant genes\n",
    " #fil = regr2.feature_importances_\n",
    "\n",
    "#plt.plot(fi)\n",
    "fil = regr2.coef_\n",
    "fil = np.absolute(fil)\n",
    "#fil = np.sum(fil,axis=0).transpose()\n",
    "fil = fil[0,:].transpose()\n",
    "fil.shape\n",
    "\n",
    "ifi = np.nonzero(fil)[0]\n",
    "fi2l = fil[ifi]\n",
    "ifi2 = np.argsort(fi2l)\n",
    "l=len(fi2l)\n",
    "#print(np.concatenate((ifi2.reshape((l,1)),fi2[ifi2].reshape((l,1))),axis=1))\n",
    "\n",
    "print(len(ge_columns))\n",
    "print(len(ge_data.columns))\n",
    "print(fil.shape)\n",
    "print(ge_columns[0:6])\n",
    "\n",
    "fi2l = pd.DataFrame(fil,columns=['importance'],index = ge_columns[6:])\n",
    "fi2l = pd.concat([fi2l, pd.DataFrame(np.repeat('post',len(fi2l.index)),columns = ['compartment'],index=fi2l.index)],axis=1)\n",
    "#fi2.loc[np.arange(len(fi2.index))<len(pregs),'compartment']='pre'\n",
    "fi2l.loc[fi2l.index.str.contains('pre'),'compartment']='pre'\n",
    "\n",
    "fi2l.index =  [id[5:] for id in fi2l.index ]\n",
    "\n",
    "fi2l=fi2l.sort_values('importance',ascending=False)\n",
    "best10 = fi2l.iloc[0:10,:]\n",
    "best20 = fi2l.iloc[0:20,:]\n",
    "best50 = fi2l.iloc[0:50,:]\n",
    "\n",
    "\n",
    "## linear case\n",
    "#fi2l.iloc[0:50,:] # A1:5, without averaging, 500 trees, with Yuste, 384gs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fi2l.iloc[0:50,:] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fi2l.iloc[0:50,:] # A1:5, without averaging, 500 trees,  50 best genes found from 280 gs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fi2l.iloc[0:50,:] # A1:5, without averaging, 500 trees, with Yuste, 20 best genes found from 280 gs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RF model case\n",
    "\n",
    "# set of most significant genes\n",
    "fi = regr.feature_importances_\n",
    "#plt.plot(fi)\n",
    "\n",
    "ifi = np.nonzero(fi)[0]\n",
    "fi2 = fi[ifi]\n",
    "ifi2 = np.argsort(fi2)\n",
    "l=len(fi2)\n",
    "#print(np.concatenate((ifi2.reshape((l,1)),fi2[ifi2].reshape((l,1))),axis=1))\n",
    "\n",
    "print(len(ge_columns))\n",
    "print(len(ge_data.columns))\n",
    "print(fi.shape)\n",
    "print(ge_columns[0:6])\n",
    "\n",
    "fi2 = pd.DataFrame(fi,columns=['importance'],index = ge_columns[6:])\n",
    "fi2 = pd.concat([fi2, pd.DataFrame(np.repeat('post',len(fi2.index)),columns = ['compartment'],index=fi2.index)],axis=1)\n",
    "#fi2.loc[np.arange(len(fi2.index))<len(pregs),'compartment']='pre'\n",
    "fi2.loc[fi2.index.str.contains('pre'),'compartment']='pre'\n",
    "\n",
    "fi2.index =  [id[5:] for id in fi2.index ]\n",
    "\n",
    "fi2=fi2.sort_values('importance',ascending=False)\n",
    "best10 = fi2.iloc[0:10,:]\n",
    "best20 = fi2.iloc[0:20,:]\n",
    "best50 = fi2.iloc[0:50,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fi2.iloc[0:50,:] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# STOP!!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fi2.to_excel('best_genes_hipp.xlsx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare with VISp trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "RF_file = open('RF_290_hipp.pickle', mode='wb')\n",
    "pickle.dump(regr,RF_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RF_file = open('RF_290_visp.pickle', mode='rb')\n",
    "regr_visp = pickle.load(RF_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RF case\n",
    "y_pred = regr_visp.predict(X_test[:,nannot:]) \n",
    "if len(y_pred.shape)==1:\n",
    "    y_pred=y_pred.reshape((len(y_pred),1))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "from sklearn import metrics\n",
    "\n",
    "#print('Mean Absolute Error:', metrics.mean_absolute_error(y_test[:,2:].ravel(), y_pred))  \n",
    "#print('Mean Squared Error:', metrics.mean_squared_error(y_test[:,2:].ravel(), y_pred))  \n",
    "#print('Root Mean Squared Error:', np.sqrt(metrics.mean_squared_error(y_test[:,2:].ravel(), y_pred))) \n",
    "\n",
    "print('Mean Absolute Error:', metrics.mean_absolute_error(y_test[:,nannot:], y_pred))  \n",
    "print('Mean Squared Error:', metrics.mean_squared_error(y_test[:,nannot:], y_pred))  \n",
    "print('Root Mean Squared Error:', np.sqrt(metrics.mean_squared_error(y_test[:,nannot:], y_pred))) \n",
    "\n",
    "\n",
    "#y_pred2 = y_pred.reshape((len(y_pred),nstp))\n",
    "yy = np.concatenate((y_test, y_pred),axis=1)\n",
    "stpn_test=[s+'_test' for s in stp_columns]\n",
    "stpn_pred=[s+'_pred' for s in stp_columns]\n",
    "yy = pd.DataFrame(yy,columns=annot_columns +stpn_test +stpn_pred)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# make input and output stp labels dataset with errors\n",
    "stpn_error = [s+'_error' for s in stp_columns]\n",
    "ye = pd.DataFrame(yy.loc[:,stpn_pred].values - yy.loc[:,stpn_test].values, columns=stpn_error)\n",
    "yy = pd.concat([yy, ye.abs()], axis=1)\n",
    "\n",
    "#yy.set_index(['pre','post']).groupby()\n",
    "#yy.set_index(['pre','post']).groupby(['pre','post']).mean()\n",
    "\n",
    "# plot errors\n",
    "annot_columns2 = annot_columns\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(12, 18))\n",
    "\n",
    "yy1=yy.set_index(annot_columns2)\n",
    "yy1.loc[:,stpn_error]=yy1.loc[:,stpn_error].values/(yy1.loc[:,stpn_test].values+1e-15)*100\n",
    "yy1=yy1.astype(float).groupby(annot_columns2).mean()\n",
    "yy1.loc[:,stpn_error].plot(ax=ax)\n",
    "plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# RF model case\n",
    "\n",
    "#plot predicted and real (test) stp data\n",
    "#plt.figure()\n",
    "yy1=yy.set_index(annot_columns2)\n",
    "#yy1=yy1.loc[:,stpn_test+stpn_pred] #/yy1.loc[:,'test']*100\n",
    "yy1=yy1.astype(float).groupby(annot_columns2).mean()\n",
    "r3=[]\n",
    "for i in range(nstp):\n",
    "    f, ax =plt.subplots(figsize=(10, 10))\n",
    "    #f, ax = plt.figure()\n",
    "    #ax = f.add_axes()\n",
    "    #plt.title(stp_columns[i])\n",
    "    ax.set_title(stp_columns[i])\n",
    "    yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "    plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n",
    "    if i==0:\n",
    "        istr=0\n",
    "    else:\n",
    "        istr=0 # 11 - if ppr data excluded\n",
    "        \n",
    "    ytest= yy1.loc[:,stpn_test[i]]  \n",
    "    ypred=yy1.loc[:,stpn_pred[i]]\n",
    "    #r2=1 - (ytest.iloc[istr:] - ypred.iloc[istr:]).pow(2).mean()/ytest.iloc[istr:].var()\n",
    "    r2=1 - (ytest.iloc[istr:] - ypred.iloc[istr:]).var()/ytest.iloc[istr:].var()\n",
    "    print(ytest.iloc[istr:].var())\n",
    "    print((ytest.iloc[istr:] - ypred.iloc[istr:]).var())\n",
    "    print(str(i) +' R^2 = '+str(r2))\n",
    "    r3 = r3 + [r2]\n",
    "\n",
    "print('mean R^2 = '+str(np.array(r3).mean()))\n",
    "\n",
    "f, ax =plt.subplots(figsize=(10, 10))\n",
    "yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "\n",
    "#print(yy1.mean())\n",
    "yy1.loc[:,stpn_test +stpn_pred]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fi2.iloc[0:50,:] # A1:5, without averaging, 500 trees,  10 from 280 gs R^2 = 0.86"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#fi2.to_excel('10_selected_from_280_gs.xlsx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# START!\n",
    "## iterative RF - search for stable interations of features\n",
    "\n",
    "## run python version of iRF :"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## # PYTHON iRF WORKS ONLY WITH BYNARY CLASSIFICATION DATA!!! - USE R iRF2.0 for multiclasses (or 1D regression)!!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import numpy as np\n",
    "from irf import irf_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(X_train.shape)\n",
    "print(X_test.shape)\n",
    "print(y_train.shape)\n",
    "print(y_test.shape)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "n_samples = 5000\n",
    "n_features = 500\n",
    "X_train = np.random.uniform(low=0, high=1, size=(n_samples, n_features))\n",
    "y_train = np.random.choice([0, 1], size=(n_samples,), p=[.5, .5])\n",
    "X_test = np.random.uniform(low=0, high=1, size=(n_samples, n_features))\n",
    "y_test = np.random.choice([0, 1], size=(n_samples,), p=[.5, .5])\n",
    "# The second feature (which is indexed by 1) is very important\n",
    "X_train[:, 1] = X_train[:, 1] + y_train\n",
    "X_test[:, 1] = X_test[:, 1] + y_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "rng = np.random.RandomState(10)  # deterministic random data\n",
    "a = np.hstack((rng.normal(size=1000),  rng.normal(loc=5, scale=2, size=1000)))\n",
    "print(a.shape)\n",
    "plt.hist(a, bins='auto')  # arguments are passed to np.histogram\n",
    "plt.title(\"Histogram with 'auto' bins\")\n",
    "plt.show()\n",
    "\n",
    "bins =[0.45, 0.75,1.0, 1.5, 5]\n",
    "y_train[:,nannot]\n",
    "in_bins =  np.histogram(y_train[:,nannot],10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import umap\n",
    "from sklearn.datasets import load_digits\n",
    "\n",
    "digits = load_digits()\n",
    "\n",
    "# embedding = umap.UMAP().fit_transform(digits.data)\n",
    "\n",
    "\n",
    "\n",
    "embedding = umap.UMAP(n_neighbors=5,\n",
    "                      min_dist=0.3,\n",
    "                      metric='correlation').fit_transform(digits.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import louvain\n",
    "import pyclustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyclustering.cluster import cluster_visualizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yy=y_train[:,nannot:]\n",
    "y_u = UMAP(spread=2).fit_transform(yy)\n",
    "\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "y_u2 = TSNE(n_components=2).fit_transform(yy)\n",
    "\n",
    "fig = plt.figure(figsize=(10, 10), facecolor='w', edgecolor='k')\n",
    "plt.scatter(y_u[:, 0], y_u[:, 1],\n",
    "                edgecolors='none', c=\"#ff4400\", s=10)\n",
    "\n",
    "plt.axis('off')\n",
    "fig.set_tight_layout(True)\n",
    "\n",
    "fig1 = plt.figure(figsize=(10, 10), facecolor='w', edgecolor='k')\n",
    "plt.scatter(y_u2[:, 0], y_u2[:, 1],\n",
    "                edgecolors='none', c=\"#ff4400\", s=10)\n",
    "\n",
    "plt.axis('off')\n",
    "fig1.set_tight_layout(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = [\"#991f1f\", \"#ff9999\", \"#ff4400\", \"#ff8800\", \"#664014\", \"#665c52\",\n",
    "          \"#cca300\", \"#f1ff33\", \"#b4cca3\", \"#0e6600\", \"#33ff4e\", \"#00ccbe\",\n",
    "          \"#0088ff\", \"#7aa6cc\", \"#293966\", \"#0000ff\", \"#9352cc\", \"#cca3c9\", \"#cc2996\"]\n",
    "\n",
    "fig = plt.figure(figsize=(15, 15), facecolor='w', edgecolor='k')\n",
    "for i, k in enumerate(list(set(df_aba_vis_in_c_s_vis.loc[:,'cluster_label'] ))):\n",
    "    i_clu = df_aba_vis_in_c_s_vis.index[df_aba_vis_in_c_s_vis.loc[:,'cluster_label'].isin([k])]\n",
    "    c =df_aba_vis_in_c_s_vis.loc[i_clu,'cluster_color'].iloc[0]\n",
    "    \n",
    "    plt.scatter(latent_u_vis[i_clu, 0], latent_u_vis[i_clu, 1], label=k,\n",
    "                edgecolors='none', c=c, s=15)\n",
    "    plt.legend(loc=9, borderaxespad=0, fontsize='small', markerscale=3)\n",
    "\n",
    "plt.axis('off')\n",
    "fig.set_tight_layout(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## run R iRF in notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0)\n",
    "np.save('X.npy',X[:,nannot:].astype(float))\n",
    "np.save('y.npy',y[:,nannot:].astype(float))\n",
    "\n",
    "np.save('X_train.npy',X_train[:,nannot:].astype(float))\n",
    "np.save('X_test.npy',X_test[:,nannot:].astype(float))\n",
    "np.save('y_train.npy',y_train[:,nannot:].astype(float))\n",
    "np.save('y_test.npy',y_test[:,nannot:].astype(float))\n",
    "\n",
    "#np.savetxt('X_train.csv',X_train[:,nannot:].astype(),fmt='%10.5f', delimiter=', ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rpy2.robjects import r, pandas2ri\n",
    "import rpy2.robjects.numpy2ri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rpy2.robjects.packages import importr\n",
    "irfr = importr('iRF')\n",
    "bs = importr('base')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#iRFr=r('iRF')\n",
    "iRFr = irfr.iRF\n",
    "X_test_r=X_test[:,nannot:]\n",
    "X_train_r=X_train[:,nannot:]\n",
    "y_train_r = y_train[:,nannot:]\n",
    "y_test_r = y_test[:,nannot:]\n",
    "\n",
    "import rpy2.robjects as ro\n",
    "import rpy2.robjects.numpy2ri\n",
    "rpy2.robjects.numpy2ri.activate()\n",
    "\n",
    "X_test_r2 = ro.r.matrix(X_test_r, nrow=X_test_r.shape[0], ncol=X_test_r.shape[1])\n",
    "X_train_r2 = ro.r.matrix(X_train_r, nrow=X_train_r.shape[0], ncol=X_train_r.shape[1])\n",
    "y_test_r2 = ro.r.matrix(y_test_r, nrow=y_test_r.shape[0], ncol=y_test_r.shape[1])\n",
    "y_train_r2 = ro.r.matrix(y_train_r, nrow=y_train_r.shape[0], ncol=y_train_r.shape[1])\n",
    "ro.r.assign(\"X_test_r\", X_test_r2 )\n",
    "ro.r.assign(\"X_train_r\", X_train_r2 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit = irfr.iRF(x=X_train_r2, y=y_train_r2, xtest=X_test_r2, ytest=y_test_r2 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test_r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PYTHON iRF WORKS ONLY WITH BYNARY CLASSIFICATION DATA!!! - USE R iRF2.0 for multiclasses (or 1D regression)!!!\n",
    "\n",
    "from irf import irf_utils\n",
    "import time\n",
    "t1 = time.time()\n",
    "\n",
    "all_rf_weights, all_K_iter_rf_data, \\\n",
    "    all_rf_bootstrap_output, all_rit_bootstrap_output, \\\n",
    "    stability_score = irf_utils.run_iRF(X_train=X_train[:,nannot:],\n",
    "                                        X_test=X_test[:,nannot:],\n",
    "                                        y_train=y_train[:,-1],\n",
    "                                        y_test=y_test[:,-1],\n",
    "                                        K=5,                          # number of iteration\n",
    "                                        n_estimators=20,              # number of trees in the forest\n",
    "                                        B=30,\n",
    "                                        random_state_classifier=2026, # random seed\n",
    "                                        propn_n_samples=.2,\n",
    "                                        bin_class_type=1,\n",
    "                                        M=20,\n",
    "                                        max_depth=5,\n",
    "                                        noisy_split=False,\n",
    "                                        num_splits=2,\n",
    "                                        n_estimators_bootstrap=5)\n",
    "\n",
    "\n",
    "t2 = time.time()\n",
    "#print(regr.feature_importances_)\n",
    "print('Elapsed time: '+str(t2-t1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit = iRFr(x=X_train_r, \n",
    "         y=y_train_r, \n",
    "         xtest=X_test_r, \n",
    "         ytest=y_test_r, \n",
    "         n.iter=5, \n",
    "         n.core=8,\n",
    "         select.iter = TRUE,\n",
    "         n.bootstrap=10\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## run iRF in R studio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "np.save('X.npy',X[:,nannot:].astype(float))\n",
    "np.save('y.npy',y[:,nannot:].astype(float))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### HERE clustering by xmeans_STP.ipynb  should be done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yc  = np.load('y_clusters.npy')\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, yc, test_size=0.25, random_state=0)\n",
    "print(X_train.shape)\n",
    "print(X_test.shape)\n",
    "print(set(yc[:,-1]))\n",
    "print( len(list( set(yc[:,-1]) )))\n",
    "print(yc.shape)\n",
    "\n",
    "np.save('Xtrain.npy',X_train[:,nannot:].astype(float))\n",
    "np.save('yctrain.npy',np.rint(y_train[:,-1]).astype(float))\n",
    "np.save('Xtest.npy',X_test[:,nannot:].astype(float))\n",
    "np.save('yctest.npy',np.rint(y_test[:,-1]).astype(float))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### HERE do iRF.Rmd comutations"
   ]
  },
  {
   "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": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "imp  = np.load('importance.npy')\n",
    "imp = pd.DataFrame(imp)\n",
    "imp = pd.concat([imp,pd.DataFrame(ge_columns[6:])],axis=1)\n",
    "imp.columns = ['importance', 'genes']\n",
    "imp = imp.set_index('genes')\n",
    "\n",
    "imp0  = np.load('importance0.npy')\n",
    "imp0 = pd.DataFrame(imp0)\n",
    "imp0 = pd.concat([imp0,pd.DataFrame(ge_columns[6:])],axis=1)\n",
    "imp0.columns = ['importance0', 'genes']\n",
    "imp0 = imp0.set_index('genes')\n",
    "\n",
    "\n",
    "# List of important genes : iRF\n",
    "imp = pd.concat([imp,imp0],axis=1)\n",
    "\n",
    "imp50 = imp.sort_values('importance',ascending=False).iloc[0:49,:]\n",
    "plt.plot(imp.sort_values('importance',ascending=False).loc[:,['importance','importance0']].values)\n",
    "imp50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#annot_columns + ge_columns[6:]\n",
    "imp500 = imp.sort_values('importance0',ascending=False).iloc[0:50,:]\n",
    "imp500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ge_columns[6:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### # train subclasses with best iRF subsets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ge_data_subcl"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "ge_columns_x = ge_columns[6:]\n",
    "ge_columns_x = pd.DataFrame(ge_columns_x)\n",
    "\n",
    "nsel=10\n",
    "print(imp50.index[0:nsel-1])\n",
    "selected_columns_x = ge_columns_x.isin(imp50.index[0:nsel-1])\n",
    "selected_columns_x =ge_columns_x[selected_columns_x.values].index.values \n",
    "print(selected_columns_x)\n",
    "selected_columns_x2 = np.concatenate([np.arange(nannot),selected_columns_x+ nannot])\n",
    "selected_columns_x2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.datasets import make_regression\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "import time\n",
    "t1 = time.time()\n",
    "\n",
    "ge_columns_x = ge_columns[6:]\n",
    "ge_columns_x = pd.DataFrame(ge_columns_x)\n",
    "\n",
    "nsel=25\n",
    "print(imp50.index[0:nsel-1])\n",
    "selected_columns_x = ge_columns_x.isin(imp50.index[0:nsel-1])\n",
    "selected_columns_x =ge_columns_x[selected_columns_x.values].index.values \n",
    "print(selected_columns_x)\n",
    "selected_columns_x2 = np.concatenate([np.arange(nannot),selected_columns_x+ nannot])\n",
    "X_selected = X[:,selected_columns_x2]\n",
    "\n",
    "do_oob_subclasses=1\n",
    "\n",
    "subclasses2 = pd.unique(subclasses2 )\n",
    "print(subclasses2) \n",
    "\n",
    "if do_oob_subclasses==1:\n",
    "\n",
    "    \n",
    "    oobs = []\n",
    "    errors = []\n",
    "    annot_columns2 = annot_columns\n",
    "    \n",
    "    r6=[]\n",
    "    r7=[]\n",
    "    for i0 in range(len(subclasses2)):\n",
    "        in_subcl = ge_data_subcl.loc[:,'subclass_pre']==subclasses2[i0]\n",
    "        in_subcl = in_subcl|(ge_data_subcl.loc[:,'subclass_post']==subclasses2[i0])\n",
    "        in_subcl = np.nonzero(in_subcl.values)[0]\n",
    "\n",
    "        X_train = np.delete(X_selected, in_subcl,axis=0)\n",
    "        y_train = np.delete(y, in_subcl,axis=0)\n",
    "        #X_train = X[:,:]\n",
    "        #y_train = y[:,:]\n",
    "        \n",
    "        X_test = X_selected[in_subcl,:]\n",
    "        y_test = y[in_subcl,:]\n",
    "        \n",
    "        \n",
    "\n",
    "        X_train, X_test0, y_train, y_test0 = train_test_split(X_selected, y, test_size=0.25, random_state=0) \n",
    "        \n",
    "        print(X_train.shape)\n",
    "        print(X_test.shape)\n",
    "        \n",
    "        regr = RandomForestRegressor(random_state=2026,\n",
    "                                    n_estimators=2500, oob_score=True, n_jobs=-1, max_samples=93)\n",
    "        regr.fit(X_train[:,nannot:], y_train[:,nannot:])  \n",
    "        \n",
    "        y_pred = regr.predict(X_test[:,nannot:]) \n",
    "        y_test2 = y_test[:,nannot:]\n",
    "        \n",
    "        \n",
    "        yy = np.concatenate((y_test, y_pred),axis=1)\n",
    "        stpn_test=[s+'_test' for s in stp_columns]\n",
    "        stpn_pred=[s+'_pred' for s in stp_columns]\n",
    "        yy = pd.DataFrame(yy,columns=annot_columns +stpn_test +stpn_pred)\n",
    "        yy.loc[:,annot_columns2] = yy.loc[:,annot_columns2].astype(str)\n",
    "        yy1 = yy.set_index(annot_columns2)\n",
    "\n",
    "        ###yy1=yy1.loc[:,stpn_test+stpn_pred] #/yy1.loc[:,'test']*100\n",
    "        yy1=yy1.astype(float).groupby(annot_columns2).mean()\n",
    "        r3=[]\n",
    "        r5=[]\n",
    "        for i in range(y_pred.shape[1]):\n",
    "            if (i==np.array([0,1,9])).sum()>0:\n",
    "                f, ax =plt.subplots(figsize=(16, 10))\n",
    "                ##f, ax = plt.figure()\n",
    "                ##ax = f.add_axes()\n",
    "                ##plt.title(stp_columns[i])\n",
    "                ax.set_title(stp_columns[i]+' '+', subclass out of bag: '+subclasses2[i0])\n",
    "                yy1.loc[:,[stpn_test[i],stpn_pred[i]]].plot(ax=ax)\n",
    "                plt.xticks(np.arange(len(yy1.index)), yy1.index, rotation=90)\n",
    "\n",
    "\n",
    "            #ytest= yy1.loc[:,stpn_test[i]]  \n",
    "            #ypred=yy1.loc[:,stpn_pred[i]]\n",
    "            \n",
    "            ##r2=1 - (ytest.iloc[istr:] - ypred.iloc[istr:]).pow(2).mean()/ytest.iloc[istr:].var()\n",
    "            if np.var(y_test2[:,i])!=0:\n",
    "                r2=1 - np.var(y_test2[:,i] - y_pred[:,i])/np.var(y_test2[:,i])\n",
    "                r4=1 - np.var(y_test2[:,i] - y_pred[:,i])/np.var(y[:,nannot+i])\n",
    "           # print(ytest.iloc[istr:].var())\n",
    "           # print((ytest.iloc[istr:] - ypred.iloc[istr:]).var())\n",
    "           # print(str(i) +' R^2 = '+str(r2))\n",
    "            r3 = r3 + [r2]\n",
    "            r5 = r5 + [r4]\n",
    "        print('subclass out of bag: '+subclasses2[i0])\n",
    "        print('mean R^2 = '+str(np.array(r3).mean()))\n",
    "        print('all R^2 = '+str(np.array(r3)))\n",
    "        print('mean R^2 total = '+str(np.array(r5).mean()))\n",
    "        print('all R^2 total = '+str(np.array(r5)))\n",
    "\n",
    "        r6 = r6 + [np.array(r3).mean()]\n",
    "        r7 = r7 + [np.array(r5).mean()]\n",
    "        oobs = oobs + [regr.oob_score_]\n",
    "    print('out of bag R^2: '+str(oobs))\n",
    "    \n",
    "    rr = np.array([subclasses2,r6,r7])\n",
    "    rr = pd.DataFrame(rr)\n",
    "    \n",
    "t2 = time.time()\n",
    "#print(regr.feature_importances_)\n",
    "print('Elapsed time: '+str(t2-t1))\n",
    "#print(regr.predict([[0, 0, 0, 0]]))\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import cross_val_score, GridSearchCV\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "\n",
    "def rfr_model(X, y):\n",
    "    # Perform Grid-Search\n",
    "    gsc = GridSearchCV(\n",
    "        estimator=RandomForestRegressor(),\n",
    "        param_grid={\n",
    "            'max_depth': range(3,7),\n",
    "            'n_estimators': (10, 50, 100, 1000),\n",
    "        },\n",
    "        cv=5, scoring='neg_mean_squared_error', verbose=0, n_jobs=-1)\n",
    "    \n",
    "    grid_result = gsc.fit(X, y)\n",
    "    best_params = grid_result.best_params_\n",
    "    \n",
    "    rfr = RandomForestRegressor(max_depth=best_params[\"max_depth\"], n_estimators=best_params[\"n_estimators\"],                               random_state=False, verbose=False)\n",
    "    # Perform K-Fold CV\n",
    "    scores = cross_val_score(rfr, X, y, cv=10, scoring='neg_mean_absolute_error')\n",
    "\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}