generate_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import pandas as pd\n",
    "from neo4j import ServiceUnavailable\n",
    "from CORD19_GraphOfDocs.neo4j_wrapper import Neo4jDatabase\n",
    "from CORD19_GraphOfDocs import select\n",
    "from sklearn.utils import shuffle\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def connect_to_the_database():\n",
    "    try:\n",
    "        database = Neo4jDatabase('bolt://localhost:7687', 'neo4j', '123')\n",
    "        # Neo4j server is unavailable.\n",
    "        # This client app cannot open a connection.\n",
    "    except ServiceUnavailable as error:\n",
    "        print('\\t* Neo4j database is unavailable.')\n",
    "        print('\\t* Please check the database connection before running this app.')\n",
    "        input('\\t* Press any key to exit the app...')\n",
    "        sys.exit(1)\n",
    "\n",
    "    return database\n",
    "\n",
    "def disconnect_from_the_database(database):\n",
    "    database.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "database = connect_to_the_database()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_duplicates(df):\n",
    "    df = df[~pd.DataFrame(np.sort(df[['node1','node2']].values, 1)).duplicated()]\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "846\n",
      "1544\n"
     ]
    }
   ],
   "source": [
    "number_of_samples = 1000\n",
    "train_positive_samples = select.get_positive_examples(database, limit=number_of_samples, train_set=True)\n",
    "train_negative_samples = select.get_negative_examples(database, limit=number_of_samples, train_set=True)\n",
    "\n",
    "test_positive_samples = select.get_positive_examples(database, limit=number_of_samples, train_set=False)\n",
    "test_negative_samples = select.get_negative_examples(database, limit=number_of_samples, train_set=False, min_hops=2, max_hops=2)\n",
    "\n",
    "columns = ['node1', 'node2', 'label']\n",
    "train_positive_samples_df = pd.DataFrame(train_positive_samples, columns=columns)\n",
    "train_positive_samples_df = remove_duplicates(train_positive_samples_df)\n",
    "train_negative_samples_df = pd.DataFrame(train_negative_samples, columns=columns)\n",
    "train_negative_samples_df = remove_duplicates(train_negative_samples_df)\n",
    "train_positive_samples_df = shuffle(train_positive_samples_df)\n",
    "train_negative_samples_df = shuffle(train_negative_samples_df)\n",
    "\n",
    "samples_num = min(len(train_positive_samples_df), len(train_negative_samples_df))\n",
    "train_positive_samples_df = train_positive_samples_df.iloc[:samples_num]\n",
    "train_negative_samples_df = train_negative_samples_df.iloc[:samples_num]\n",
    "\n",
    "test_positive_samples_df = pd.DataFrame(test_positive_samples, columns=columns)\n",
    "test_positive_samples_df = remove_duplicates(test_positive_samples_df)\n",
    "test_negative_samples_df = pd.DataFrame(test_negative_samples, columns=columns)\n",
    "test_negative_samples_df = remove_duplicates(test_negative_samples_df)\n",
    "test_positive_samples_df = shuffle(test_positive_samples_df)\n",
    "test_negative_samples_df = shuffle(test_negative_samples_df)\n",
    "\n",
    "samples_num = min(len(test_positive_samples_df), len(test_negative_samples_df))\n",
    "test_positive_samples_df = test_positive_samples_df.iloc[:samples_num]\n",
    "test_negative_samples_df = test_negative_samples_df.iloc[:samples_num]\n",
    "\n",
    "train_df = pd.concat([train_positive_samples_df, train_negative_samples_df])\n",
    "train_df = shuffle(train_df)\n",
    "test_df = pd.concat([test_positive_samples_df, test_negative_samples_df])\n",
    "test_df = shuffle(test_df)\n",
    "print(len(train_df))\n",
    "print(len(test_df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train df positive examples 423\n",
      "Train df negative examples 423\n",
      "Test df positive examples 772\n",
      "Test df negative examples 772\n"
     ]
    }
   ],
   "source": [
    "print(f'Train df positive examples {len(train_df[train_df[\"label\"] == 1])}')\n",
    "print(f'Train df negative examples {len(train_df[train_df[\"label\"] == 0])}')\n",
    "print(f'Test df positive examples {len(test_df[test_df[\"label\"] == 1])}')\n",
    "print(f'Test df negative examples {len(test_df[test_df[\"label\"] == 0])}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculate train df features\n",
      "Calculate test df features\n"
     ]
    }
   ],
   "source": [
    "def get_samples_dataframe_with_features(database, original_df, train_set):\n",
    "    pairs = original_df[['node1', 'node2', 'label']].values.tolist()\n",
    "    results = select.create_graph_features(database, pairs, train_set=train_set)\n",
    "    df = pd.DataFrame(results, columns=[\n",
    "        'node1',\n",
    "        'node2',\n",
    "        'adamic_adar',\n",
    "        'common_neighbors', \n",
    "        'preferential_attachment',\n",
    "        'total_neighbors',\n",
    "        'similarity',\n",
    "        'label',\n",
    "    ])\n",
    "    return df\n",
    "\n",
    "print('Calculate train df features')\n",
    "train_df = get_samples_dataframe_with_features(database, train_df, train_set=True)\n",
    "print('Calculate test df features')\n",
    "test_df = get_samples_dataframe_with_features(database, test_df, train_set=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.fillna({'similarity': 0}, inplace=True)\n",
    "test_df.fillna({'similarity': 0}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "disconnect_from_the_database(database)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.to_csv(f'train_{len(train_df)}.csv')\n",
    "test_df.to_csv(f'test_{len(test_df)}.csv')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}