{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f026829",
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing\n",
    "# multiprocessing.set_start_method('forkserver')\n",
    "# from concurrent.futures import ProcessPoolExecutor\n",
    "import controlpy\n",
    "import time\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import csv\n",
    "# import itertools\n",
    "# import scipy\n",
    "import numpy.linalg as LA\n",
    "import scipy as sc\n",
    "# from scipy import stats\n",
    "from os import listdir\n",
    "from os.path import join\n",
    "# from igraph import  *\n",
    "from scipy import stats\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import sparse\n",
    "import sys\n",
    "## from joblib import Parallel, delayed\n",
    "from joblib import Parallel, delayed\n",
    "# import tqdm\n",
    "from sklearn.model_selection import GridSearchCV,StratifiedKFold\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.model_selection import cross_val_score\n",
    "import pandas as pd\n",
    "from sklearn.utils import shuffle\n",
    "num_cores = multiprocessing.cpu_count()\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.utils import to_networkx\n",
    "from sklearn.metrics import accuracy_score\n",
    "import torch\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "# from rdkit import Chem\n",
    "import pickle\n",
    "from torch_geometric.datasets import TUDataset\n",
    "import matplotlib.colors as pltc\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87f7a236",
   "metadata": {},
   "outputs": [],
   "source": [
    "INCLUDE_GRM = True\n",
    "INCLUDE_INVERSE_GRM = False #-----\n",
    "INCLUDE_EIGEN_VALUES = True\n",
    "INCLUDE_MIN_MAX = True\n",
    "INCLUDE_MEAN = True\n",
    "INCLUDE_LAP_SPECT = True\n",
    "INCLUDE_LAP_SPECT_EXT = True\n",
    "INCLUDE_METRIC_DIMENSION = True\n",
    "INCLUDE_FEATURES = True\n",
    "LAPLACIAN_ENERGY = True\n",
    "ECC_SPECTRUM = True\n",
    "ECC_ENERGY = True\n",
    "WEINER_INDEX = True\n",
    "TRACE_DEG_SEQ = True\n",
    "INCLUDE_CYCLES = True\n",
    "\n",
    "DESCRIPTOR_SIZE_ONE_LEADER = 5\n",
    "LEADERS_SIZE_DET = 0.95\n",
    "\n",
    "def add_type_attribute(G, node_types):\n",
    "    type_mapping = {}\n",
    "    for index, n in enumerate(G.nodes()):\n",
    "        type_mapping[n] = node_types[index]\n",
    "    nx.set_node_attributes(G, type_mapping,\"type\")\n",
    "    return G\n",
    "\n",
    "def get_leaders_followers(node_types, elem,G):\n",
    "    G = add_type_attribute(G, node_types)\n",
    "    if elem is 'C':\n",
    "        leaders = [i for i,x in enumerate(node_types) if x !=elem]\n",
    "        followers = [i for i,x in enumerate(node_types) if x ==elem]\n",
    "\n",
    "    elif elem is 'VNode':\n",
    "        attrib = nx.get_node_attributes(G, \"type\")\n",
    "        leader_nodes = [x for x in G.nodes() if attrib[x]!='C']\n",
    "        neighbors = []\n",
    "        for n in leader_nodes:\n",
    "            nei = list(G.neighbors(n))\n",
    "            neighbors.extend(nei)\n",
    "        neighbors = list(set(neighbors))        \n",
    "        leaders = [i for i,x in enumerate(node_types) if x !='C']\n",
    "        leaders.extend(neighbors)\n",
    "        leaders = list(set(leaders))\n",
    "        followers = [i for i in range(len(list(G.nodes()))) if i not in leaders]\n",
    "        \n",
    "        \n",
    "    else:\n",
    "        leaders = [i for i,x in enumerate(node_types) if x ==elem]\n",
    "        followers = [i for i,x in enumerate(node_types) if x !=elem]\n",
    "    return leaders, followers\n",
    "\n",
    "def return_node_types(d,node_types):\n",
    "    g = to_networkx(d, to_undirected = True)\n",
    "    if node_types ==None:\n",
    "        x = d.x[:,0].numpy().astype(int).tolist()\n",
    "        return x,g\n",
    "    x = d.x\n",
    "    types_ = []\n",
    "    for xv in x:\n",
    "        type_index = (xv == 1.0).nonzero(as_tuple=True)[0].item()\n",
    "        types_.append(node_types.get(type_index))\n",
    "    return types_, g   \n",
    "\n",
    "def get_descriptor(d,node_types,ALL_ELEMENTS):\n",
    "    atom_types, G = return_node_types(d, node_types)\n",
    "    \n",
    "    \n",
    "    if nx.is_connected(G)== False:\n",
    "        G = add_vertex(G)\n",
    "        atom_types.append(\"VNode\")\n",
    "    if G.number_of_nodes()<10:\n",
    "        G,atom_types = add_multiple_copies(G,atom_types)\n",
    "    A1 = -(np.matrix(nx.laplacian_matrix(G,nodelist=sorted(G.nodes())).todense()))\n",
    "    n = A1.shape[0]\n",
    "    desc = {}\n",
    "    new_desc = []\n",
    "    total_types = len( ALL_ELEMENTS ) \n",
    "    desc_size = (total_types * DESCRIPTOR_SIZE_ONE_LEADER) + 2\n",
    "    \n",
    "    leaders = []\n",
    "    nodes = list(G.nodes())\n",
    "    \n",
    "    old_n = A1.shape[0]\n",
    "    num_l_idx = 0\n",
    "    \n",
    "    for elem in ALL_ELEMENTS:\n",
    "        traces = ranks = min_eigs = max_eigs = itraces = iranks = imin_eigs = imax_eigs = metric_dimension = 0\n",
    "        leaders, follows = get_leaders_followers(atom_types, elem, G)\n",
    "        if len(leaders) is 0 or (len(follows) is 0):\n",
    "            new_desc.extend( np.zeros(DESCRIPTOR_SIZE_ONE_LEADER,dtype=float) )\n",
    "            continue\n",
    "        keys = A1.shape\n",
    "        A2 = A1[follows, :]\n",
    "        A3 = A2[:, follows]\n",
    "        A = 1*A3\n",
    "        keys = A.shape\n",
    "        n = A.shape[0] ## number of followers\n",
    "        mask = np.ones(old_n, dtype=bool)\n",
    "        mask[leaders] = False\n",
    "        B = A1[mask,:] #removing corresponding rows of leaders from laplacian\n",
    "        B = B[:,leaders] #selecting corresponding columns of leaders from laplacian\n",
    "        A = np.mat(A)\n",
    "        B = np.mat(B)\n",
    "        if INCLUDE_GRM:\n",
    "            grm = controlpy.analysis.controllability_gramian(A,B)\n",
    "            ranks = LA.matrix_rank(grm)\n",
    "            traces = np.trace(grm) \n",
    "            if INCLUDE_EIGEN_VALUES :\n",
    "                w, v = LA.eig(grm)\n",
    "                a = np.real(w)\n",
    "                a[a == 0] = 0.0001\n",
    "                minval = np.min(a)\n",
    "                min_eigs = np.min( minval )\n",
    "                max_eigs = np.max( a ) \n",
    "            if INCLUDE_INVERSE_GRM :\n",
    "                try:\n",
    "                    grm = LA.pinv( grm ,hermitian=True)\n",
    "                except:\n",
    "                    print(\"Oops!\", sys.exc_info()[0], \"occurred.\")\n",
    "                    return None\n",
    "                itraces = np.trace(grm)\n",
    "                iranks =  LA.matrix_rank(grm)\n",
    "                if INCLUDE_EIGEN_VALUES :\n",
    "                    w, v = LA.eig(grm)\n",
    "                    a = np.real(w)\n",
    "                    a[a == 0] = 0.0001\n",
    "                    \n",
    "                    imin_eigs = np.min( a ) \n",
    "                    imax_eigs =  np.max( a )   \n",
    "\n",
    "            if INCLUDE_METRIC_DIMENSION:            \n",
    "                metric_dimension = compute_metric_dimension(G, leaders)\n",
    "\n",
    "        new_desc.append(traces)\n",
    "        new_desc.append(ranks)\n",
    "\n",
    "        if INCLUDE_EIGEN_VALUES :\n",
    "            new_desc.append(min_eigs)\n",
    "            new_desc.append(max_eigs)\n",
    "        if INCLUDE_INVERSE_GRM :   \n",
    "            new_desc.append(itraces)\n",
    "        if INCLUDE_EIGEN_VALUES and INCLUDE_INVERSE_GRM:\n",
    "            new_desc.append( imin_eigs)\n",
    "            new_desc.append(imax_eigs)\n",
    "\n",
    "        if INCLUDE_METRIC_DIMENSION:\n",
    "            new_desc.append(metric_dimension)\n",
    "\n",
    "    n,m = G.number_of_nodes(),G.number_of_edges()\n",
    "    new_desc.append(n)\n",
    "    new_desc.append(m)\n",
    "    new_desc.append(len(list(nx.biconnected_components(G))))\n",
    "    if INCLUDE_LAP_SPECT :\n",
    "        Ls = sorted(nx.laplacian_spectrum(G))\n",
    "        Ls2 = Ls.copy()\n",
    "        Ls2.extend( list(np.zeros(10) ) )\n",
    "        new_desc.append(Ls2[0])\n",
    "        new_desc.append(Ls2[1])\n",
    "        new_desc.append(Ls2[2])\n",
    "    if INCLUDE_LAP_SPECT_EXT :\n",
    "        new_desc.append(Ls2[3])\n",
    "        new_desc.append(Ls2[4])\n",
    "        \n",
    "        Ls2 = list(np.zeros(10) )\n",
    "        Ls2.extend( Ls )\n",
    "        \n",
    "        new_desc.append(Ls2[-1])\n",
    "        new_desc.append(Ls2[-2])\n",
    "        new_desc.append(Ls2[-3])\n",
    "        \n",
    "    if INCLUDE_CYCLES:\n",
    "        if n-m == 1:\n",
    "            new_desc.append(1)  \n",
    "        else:\n",
    "            new_desc.append(0)\n",
    "        if n-m == 0:\n",
    "            new_desc.append(1)\n",
    "        else:\n",
    "            new_desc.append(0)\n",
    "        if n-m == -1:\n",
    "            new_desc.append(1)\n",
    "        else:\n",
    "            new_desc.append(0)\n",
    "        if n-m < -1:\n",
    "            new_desc.append(1)\n",
    "        else:\n",
    "            new_desc.append(0)\n",
    "    if node_types==None:\n",
    "        features = list(np.mean(d.x.numpy(),axis= 0))\n",
    "        new_desc.extend(features)\n",
    "    \n",
    "    return new_desc\n",
    "\n",
    "def add_vertex(G):\n",
    "    G = nx.convert_node_labels_to_integers(G)\n",
    "    nodes = sorted(list(G.nodes()))\n",
    "    new_vertex = nodes[-1]+1 \n",
    "    G.add_node(new_vertex)\n",
    "    for n in G.nodes():\n",
    "        if n !=new_vertex:\n",
    "            G.add_edge(new_vertex, n)\n",
    "    return G\n",
    "def apply_RF_Grid(data_X,data_y):\n",
    "    num_cores = multiprocessing.cpu_count()\n",
    "    max_arr = []\n",
    "    estimator  = RandomForestClassifier(criterion='gini', max_depth=None, min_weight_fraction_leaf=0.0,\n",
    "                                            max_leaf_nodes=None, bootstrap=True, \n",
    "                                            oob_score=False, n_jobs=num_cores,verbose=0, warm_start=False,\n",
    "                                            class_weight=None)\n",
    "         \n",
    "    param_grid = {'n_estimators':[50,100,500], 'max_features':['sqrt'], \n",
    "                  'min_samples_split':[2,3,4,5,10], 'min_samples_leaf':[1,2,5]}\n",
    "    kf = StratifiedKFold(n_splits=10, random_state = 567, shuffle = True)\n",
    "    grid_rf    = GridSearchCV(estimator, param_grid, scoring='accuracy', n_jobs=num_cores, \n",
    "                 refit=True, cv=kf, verbose=1, pre_dispatch='n_jobs', \n",
    "                 error_score='raise')\n",
    "\n",
    "    grid_rf.fit(data_X, data_y)\n",
    "    \n",
    "    cv_results = pd.DataFrame(grid_rf.cv_results_)\n",
    "    std = cv_results.loc[grid_rf.best_index_].std_test_score\n",
    "    best_score = grid_rf.best_score_\n",
    "    return best_score, std\n",
    "\n",
    "\n",
    "def compute_metric_dimension(G, leaders):\n",
    "    nodes_list = list(G.nodes()) \n",
    "    distance_vec = []\n",
    "    for n in nodes_list:\n",
    "        v_ = [nx.shortest_path_length(G, source=n, target=l) for l in leaders]\n",
    "        distance_vec.append(v_)\n",
    "    distance_vec = np.array(distance_vec)\n",
    "    return np.unique(distance_vec, axis=0).shape[0]\n",
    "\n",
    "def add_multiple_copies(G,atom_types):\n",
    "    no_node = G.number_of_nodes()   \n",
    "    H = G.copy()\n",
    "    type_copy = atom_types.copy()\n",
    "    H = nx.convert_node_labels_to_integers(H, ordering='sorted')\n",
    "    val = sorted(G.nodes())\n",
    "    val0 = val[0]\n",
    "    while H.number_of_nodes()<10:\n",
    "        val1 = H.number_of_nodes()-1\n",
    "        H = nx.union(H, G, rename=('G1-', 'G2-'))\n",
    "        atom_types.extend(type_copy)\n",
    "        H.add_edge('G1-'+str(val1), 'G2-'+str(val0))\n",
    "        H = nx.convert_node_labels_to_integers(H, ordering='sorted')\n",
    "    return H, atom_types\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c5dd642",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "datasets = [\"PTC_MR\",\"ENZYMES\"]\n",
    "file = open(\"results/CTRL_bioinformatics_test.csv\",'a',newline = '')\n",
    "res_writer = csv.writer(file, delimiter=',',quotechar='|', quoting=csv.QUOTE_MINIMAL)\n",
    "res_writer.writerow([\"dataset\",\"CTRL_ACC\",\"std\",\"time\"])\n",
    "for f in datasets:\n",
    "    attr = False\n",
    "    if f==\"ENZYMES\":\n",
    "        graph_dataset = TUDataset(root = \"../data/TUDataset\", name=f,use_node_attr = True)\n",
    "        ALL_ELEMENTS = [0,1,2,3,4,5,6,8,9,10,11,12,13,14,15,16,17,18,19,20]\n",
    "        node_types = None\n",
    "    else:\n",
    "        graph_dataset = TUDataset(root = \"../data/TUDataset\", name=f,use_node_attr = False)\n",
    "        node_types = pickle.load(open(\"mapping_\"+f+\".pkl\",\"rb\"))\n",
    "        ALL_ELEMENTS = [ 'O', 'N', 'C', 'S', 'F', 'Cl', 'P', 'Br', 'I', 'Si', 'Se', 'As', 'Pt', \n",
    "                'B', 'Bi', 'Ga', 'Sn', 'Pb', 'Cu','VNode']\n",
    "    start = time.time()\n",
    "    control_des = Parallel(n_jobs=-1, verbose=2)(delayed(get_descriptor)(d, node_types, ALL_ELEMENTS) for d in graph_dataset)\n",
    "    end = time.time()\n",
    "    labels = [x.y.item() for x in graph_dataset]\n",
    "    np.save(\"data/CTRL_new_data/\"+f+\"_desc.npy\", control_des)\n",
    "    control_des = np.array(control_des)\n",
    "    \n",
    "    labels = np.array(labels)\n",
    "\n",
    "    acc, std = apply_RF_Grid(control_des, labels)\n",
    "    to_write = [f,acc,std,end-start]\n",
    "    res_writer.writerow(to_write)\n",
    "    file.flush()\n",
    "    print(\"results:\", to_write)\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "40e536c4c67f3855c32fb03afd43728ac25fb5ba75f4ca513917f227d6e650ce"
  },
  "kernelspec": {
   "display_name": "Python [conda env:rdkit]",
   "language": "python",
   "name": "conda-env-rdkit-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
