process_edgelist.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pickle\n",
    "from sklearn.preprocessing import LabelEncoder "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = os.path.join(os.getcwd(), 'wiki/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_graph(fname, data_dir=data_dir):\n",
    "    graph = []\n",
    "    \n",
    "    file_path = os.path.join(data_dir, '%s.txt' % (fname))\n",
    "    with open(file_path) as file:\n",
    "        for line in file:\n",
    "            line = line.strip().split('\\t')\n",
    "            graph.append(np.array([line[0], line[1]]))\n",
    "\n",
    "    return np.array(graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_graph(fname):\n",
    "    graph = load_graph(fname)\n",
    "    df_graph = pd.DataFrame(graph, columns=['cat', 'edge'])\n",
    "    \n",
    "    lb = LabelEncoder()\n",
    "    df_graph['cat_int'] = lb.fit_transform(graph[:, 0])\n",
    "    df_graph['edge_int'] = lb.fit_transform(graph[:, 1]) + len(np.unique(df_graph['cat_int']))\n",
    "    \n",
    "    df_graph[['cat_int', 'edge_int']].to_csv(os.path.join(data_dir, '%s.edgelist' % (fname)), header=False, index=False, sep=' ')\n",
    "    \n",
    "    int_to_cat = df_graph[['cat','cat_int']].drop_duplicates().set_index('cat_int').to_dict()['cat']\n",
    "    int_to_edge = df_graph[['edge','edge_int']].drop_duplicates().set_index('edge_int').to_dict()['edge']\n",
    "\n",
    "    pickle.dump(int_to_cat, open(os.path.join(data_dir, '%s_dict_cat.pkl' % (fname)), \"wb\"))\n",
    "    pickle.dump(int_to_edge, open(os.path.join(data_dir, '%s_dict_edge.pkl' % (fname)), \"wb\"))\n",
    "    \n",
    "    return df_graph, int_to_cat, int_to_edge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = ['cat_edges1', 'cat_edges_2', 'link_edges']\n",
    "for fname in fnames:\n",
    "    df_graph, int_to_cat, int_to_edge = convert_graph(fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "31739"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.max(list(int_to_edge.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_cifar_labels = pd.read_csv(os.path.join(data_dir, 'graph_to_cifar_labels.csv'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_to_cifar_labels = df_cifar_labels.set_index('wiki_graph').to_dict()['CIFAR']\n",
    "pickle.dump(graph_to_cifar_labels, open(os.path.join(data_dir, 'graph_to_cifar_labels.pkl'), \"wb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Aluminum_can': 'can',\n",
       " 'Arecaceae': 'palm_tree',\n",
       " 'Batoidea': 'ray',\n",
       " 'Chili pepper': 'sweet_pepper',\n",
       " 'Computer keyboard': 'keyboard',\n",
       " 'Light_fixture': 'lamp',\n",
       " 'Opossum': 'possum',\n",
       " 'Orange (fruit)': 'orange',\n",
       " 'Pinniped': 'seal',\n",
       " 'Plate (dishware)': 'plate',\n",
       " 'Table (furniture)': 'table',\n",
       " 'apples': 'apple',\n",
       " 'aquarium': 'aquarium_fish',\n",
       " 'baby': 'baby',\n",
       " 'bear': 'bear',\n",
       " 'beaver': 'beaver',\n",
       " 'bed': 'bed',\n",
       " 'bee': 'bee',\n",
       " 'beetle': 'beetle',\n",
       " 'bicycle': 'bicycle',\n",
       " 'bottles': 'bottle',\n",
       " 'bowls': 'bowl',\n",
       " 'boy': 'boy',\n",
       " 'bridge': 'bridge',\n",
       " 'bus': 'bus',\n",
       " 'butterfly': 'butterfly',\n",
       " 'camel': 'camel',\n",
       " 'castle': 'castle',\n",
       " 'caterpillar': 'caterpillar',\n",
       " 'cattle': 'cattle',\n",
       " 'chair': 'chair',\n",
       " 'chimpanzee': 'chimpanzee',\n",
       " 'clock': 'clock',\n",
       " 'cloud': 'cloud',\n",
       " 'cockroach': 'cockroach',\n",
       " 'couch': 'couch',\n",
       " 'crab': 'crab',\n",
       " 'crocodile': 'crocodile',\n",
       " 'cups': 'cup',\n",
       " 'dinosaur': 'dinosaur',\n",
       " 'dolphin': 'dolphin',\n",
       " 'elephant': 'elephant',\n",
       " 'flatfish': 'flatfish',\n",
       " 'forest': 'forest',\n",
       " 'fox': 'fox',\n",
       " 'girl': 'girl',\n",
       " 'hamster': 'hamster',\n",
       " 'house': 'house',\n",
       " 'kangaroo': 'kangaroo',\n",
       " 'lawn_mower': 'lawn_mower',\n",
       " 'leopard': 'leopard',\n",
       " 'lion': 'lion',\n",
       " 'lizard': 'lizard',\n",
       " 'lobster': 'lobster',\n",
       " 'man': 'man',\n",
       " 'maple': 'maple_tree',\n",
       " 'motorcycle': 'motorcycle',\n",
       " 'mountain': 'mountain',\n",
       " 'mouse': 'mouse',\n",
       " 'mushrooms': 'mushroom',\n",
       " 'oak': 'oak_tree',\n",
       " 'orchids': 'orchid',\n",
       " 'otter': 'otter',\n",
       " 'pears': 'pear',\n",
       " 'pickup_truck': 'pickup_truck',\n",
       " 'pine': 'pine_tree',\n",
       " 'plain': 'plain',\n",
       " 'poppy': 'poppy',\n",
       " 'porcupine': 'porcupine',\n",
       " 'rabbit': 'rabbit',\n",
       " 'raccoon': 'raccoon',\n",
       " 'road': 'road',\n",
       " 'rocket': 'rocket',\n",
       " 'roses': 'rose',\n",
       " 'sea': 'sea',\n",
       " 'shark': 'shark',\n",
       " 'shrew': 'shrew',\n",
       " 'skunk': 'skunk',\n",
       " 'skyscraper': 'skyscraper',\n",
       " 'snail': 'snail',\n",
       " 'snake': 'snake',\n",
       " 'spider': 'spider',\n",
       " 'squirrel': 'squirrel',\n",
       " 'streetcar': 'streetcar',\n",
       " 'sunflowers': 'sunflower',\n",
       " 'tank': 'tank',\n",
       " 'telephone': 'telephone',\n",
       " 'television': 'television',\n",
       " 'tiger': 'tiger',\n",
       " 'tractor': 'tractor',\n",
       " 'train': 'train',\n",
       " 'trout': 'trout',\n",
       " 'tulips': 'tulip',\n",
       " 'turtle': 'turtle',\n",
       " 'wardrobe': 'wardrobe',\n",
       " 'whale': 'whale',\n",
       " 'willow': 'willow_tree',\n",
       " 'wolf': 'wolf',\n",
       " 'woman': 'woman',\n",
       " 'worm': 'worm'}"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph_to_cifar_labels = pickle.load(open(os.path.join(data_dir, 'graph_to_cifar_labels.pkl'), \"rb\"))\n",
    "graph_to_cifar_labels"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}