prepare_neo4j_import_csvs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neo4j CSV import preparation\n",
    "\n",
    "Daniel's import a single edge at a time using Py2neo is way too slow. I suspect that Neo4j can work much faster, because the CPU saturation during edge import is maybe only 30%, and therefore will be forced to move to CSV import because 1.2 million edges should take only a minute or two according to the online documentation (instead of 20 hours).\n",
    "\n",
    "## Steps\n",
    "\n",
    "1. Convert the hetio.graph format into a TSV file with the bare minimum of data needed to do the metapath calculations.\n",
    "2. Create node and edge TSV files.\n",
    "3. Batch import into neo4j\n",
    "4. Repeat for the permuted networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import collections\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "from glob import glob\n",
    "from tqdm import tqdm\n",
    "\n",
    "import hetio.hetnet\n",
    "import hetio.readwrite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph readers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def check_node_uniqueness(graph):\n",
    "    \"\"\"Check that all node identifiers are unique.\"\"\"\n",
    "    node_ids = []\n",
    "    for node in graph.get_nodes():\n",
    "        node_type, node_id = node.get_id()\n",
    "\n",
    "        node_ids.append(node_id)\n",
    "\n",
    "    node_ids = set(node_ids)\n",
    "\n",
    "    assert len(node_ids) == graph.n_nodes\n",
    "    \n",
    "    return node_ids\n",
    "\n",
    "def read_graph(path):\n",
    "    return hetio.readwrite.read_graph(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CSV Writer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def write_to_csv(fname, header, generator, sep=\",\"):\n",
    "    assert isinstance(header, list)\n",
    "    assert isinstance(generator, collections.Iterable)\n",
    "        \n",
    "    with open(fname, \"w\") as fout:\n",
    "        fout.write(\"{}\\n\".format(sep.join(header)))\n",
    "        \n",
    "        for vals in generator:\n",
    "            fout.write(\"{}\\n\".format(sep.join(str(v) for v in vals)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create CSV storage directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "csv_folder = os.path.join(\"data\", \"import_csvs\")\n",
    "if os.path.exists(csv_folder):\n",
    "    shutil.rmtree(csv_folder)\n",
    "    \n",
    "os.makedirs(csv_folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Node and edge generators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def gen_nodes(graph, show_progress=False):\n",
    "    def make_gen():\n",
    "        for node in graph.get_nodes():\n",
    "            node_type, node_id = node.get_id()\n",
    "\n",
    "            yield (node_id, node_id, '\"{}\"'.format(node.name), node_type)\n",
    "            \n",
    "    gen = make_gen()\n",
    "    if show_progress:\n",
    "        gen = tqdm(gen, total=graph.n_nodes, desc=\"Nodes\")\n",
    "    \n",
    "    return gen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def gen_edges(graph, node_ids, show_progress=False):\n",
    "    def as_type(metaedge):\n",
    "        rel_type = str(metaedge.kind).upper().replace(\" \", \"_\")\n",
    "        abbrev = metaedge.get_standard_abbrev()\n",
    "        return \"{}_{}\".format(rel_type, abbrev)\n",
    "    \n",
    "    def make_gen():\n",
    "        for edge in graph.get_edges(exclude_inverts=True):\n",
    "            source = edge.source.get_id()[1]\n",
    "            target = edge.target.get_id()[1]\n",
    "\n",
    "            edge_type = as_type(edge.metaedge)\n",
    "\n",
    "            assert source in node_ids and target in node_ids\n",
    "\n",
    "            yield (source, target, edge_type)\n",
    "            \n",
    "    gen = make_gen()\n",
    "    if show_progress:\n",
    "        gen = tqdm(gen, total=graph.n_edges, desc=\"Edges\")\n",
    "    \n",
    "    return gen"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph to CSV converter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def convert_to_csv(fname):\n",
    "    \"\"\"Assuming fnames match *.json.bz2\"\"\"\n",
    "    graph = read_graph(fname)\n",
    "    \n",
    "    node_ids = check_node_uniqueness(graph)\n",
    "    \n",
    "    out_name = os.path.basename(fname).split(\".\")[0]\n",
    "    out_path = os.path.join(\"data\", \"import_csvs\")\n",
    "\n",
    "    # write nodes to file\n",
    "    write_to_csv(\n",
    "        os.path.join(out_path, \"{}_nodes.csv\".format(out_name)),\n",
    "        [\":ID\", \"identifier:string\", \"name:string\", \":LABEL\"],\n",
    "        gen_nodes(graph, show_progress=True),\n",
    "    )\n",
    "\n",
    "    write_to_csv(\n",
    "        os.path.join(out_path, \"{}_edges.csv\".format(out_name)),\n",
    "        [\":START_ID\", \":END_ID\", \":TYPE\"],\n",
    "        gen_edges(graph, node_ids, show_progress=True)\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert networks to CSV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converting data/hetnet.json.bz2 to CSV\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Nodes: 100%|██████████| 47031/47031 [00:00<00:00, 126138.09it/s]\n",
      "Edges: 100%|██████████| 2194470/2194470 [00:20<00:00, 106031.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converting data/permuted/hetnet_perm-4.json.bz2 to CSV\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Nodes: 100%|██████████| 47031/47031 [00:00<00:00, 136276.62it/s]\n",
      "Edges: 100%|██████████| 2194470/2194470 [00:20<00:00, 106771.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converting data/permuted/hetnet_perm-5.json.bz2 to CSV\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Nodes: 100%|██████████| 47031/47031 [00:00<00:00, 120769.10it/s]\n",
      "Edges: 100%|██████████| 2194470/2194470 [00:20<00:00, 106765.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converting data/permuted/hetnet_perm-2.json.bz2 to CSV\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Nodes: 100%|██████████| 47031/47031 [00:00<00:00, 117344.63it/s]\n",
      "Edges: 100%|██████████| 2194470/2194470 [00:21<00:00, 103854.39it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converting data/permuted/hetnet_perm-1.json.bz2 to CSV\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Nodes: 100%|██████████| 47031/47031 [00:00<00:00, 137597.35it/s]\n",
      "Edges: 100%|██████████| 2194470/2194470 [00:20<00:00, 105673.70it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converting data/permuted/hetnet_perm-3.json.bz2 to CSV\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Nodes: 100%|██████████| 47031/47031 [00:00<00:00, 132561.45it/s]\n",
      "Edges: 100%|██████████| 2194470/2194470 [00:20<00:00, 107203.90it/s]\n"
     ]
    }
   ],
   "source": [
    "fnames = [\"data/hetnet.json.bz2\"] + glob(\"data/permuted/*.json.bz2\")\n",
    "\n",
    "for fname in fnames:\n",
    "    print(\"Converting {} to CSV\".format(fname))\n",
    "    convert_to_csv(fname)"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}