notebooks/1 Data Collection (2023).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Collection (2023)\n",
    "\n",
    "This is the code we used for our 2023 data collection. It retrieves PubMed abstracts using EDirect. For 2021 data collection, see [Data Collection (2021) notebook](./1%20Data%20Collection%20(2021).ipynb).\n",
    "\n",
    "The code is divided into 3 parts:\n",
    "1. Load the EFO ontology and extract all task and construct names and queries.\n",
    "2. Retrieving articles from PubMed in MEDLINE format.\n",
    "3. Parse the MEDLINE files and from each extract the abstract, year, DOI, PMID, and journal title. The output dataset is stored in `data/pubmed/abstracts_2023.csv.gz`.\n",
    "\n",
    "## Requirements\n",
    "\n",
    "-  Before running the notebook, make sure you have EDirect installed. You can find instructions on how to install it here: https://www.ncbi.nlm.nih.gov/books/NBK179288/\n",
    "\n",
    "- You also need to set the `NCBI_API_KEY` environment variable to your NCBI API key. You can get one here: https://www.ncbi.nlm.nih.gov/account/settings/\n",
    "\n",
    "- This notebook uses task and construct names from the EFO ontology. Make sure you have the latest version of the ontology in the `data/ontologies/efo.owl` path. Download the ontology from here: https://huggingface.co/datasets/morteza/cogtext/blob/main/ontologies/efo.owl and put the OWL file in the `data/ontologies/` directory.\n",
    "\n",
    "- Activate the `cogtext` conda environment: `mamba activate cogtext`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 3\n",
    "\n",
    "import os\n",
    "import re\n",
    "import shlex\n",
    "import subprocess\n",
    "from pathlib import Path\n",
    "\n",
    "import pandas as pd\n",
    "import xmltodict\n",
    "from IPython.display import clear_output\n",
    "from joblib import Parallel, delayed\n",
    "from owlready2 import get_ontology\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# collect data for the following categories\n",
    "CATEGORIES = ['CognitiveTask', 'CognitiveConstruct']\n",
    "\n",
    "# load the EFO ontology\n",
    "ONTOLOGY = get_ontology('data/ontologies/efo.owl').load()\n",
    "\n",
    "OUTPUT_PATH = 'data/pubmed/abstracts_2023.csv.gz'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EF ontology contains 126 PubMed queries for CognitiveTasks.\n",
      "EF ontology contains 72 PubMed queries for CognitiveConstructs.\n"
     ]
    }
   ],
   "source": [
    "all_queries = {}\n",
    "\n",
    "for category in CATEGORIES:\n",
    "  category_queries = {e.name:str(e.pubmedQuery[0]).replace('\"', '\\\\\"')\n",
    "                      for e in ONTOLOGY[category].descendants()\n",
    "                      if len(e.pubmedQuery) > 0}\n",
    "  all_queries[category] = category_queries\n",
    "  print(f'EF ontology contains {len(category_queries)} PubMed queries for {category}s.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_medline_to_csv(category, subcategory, pubmed_file):\n",
    "\n",
    "    with open(pubmed_file, 'r') as f:\n",
    "        docs = f.read().replace('\\n      ',' ').split('\\n\\n')\n",
    "    \n",
    "    records = []\n",
    "\n",
    "    for doc in docs:\n",
    "    # for doc in docs:\n",
    "        doc = doc + '\\n'\n",
    "        if 'AB  - ' not in doc:\n",
    "            abstract = re.search(r'OAB\\s-\\s(.*)\\n', doc).group(1)\n",
    "        else:\n",
    "            abstract = re.search(r'AB\\s\\s-\\s(.*)\\n', doc).group(1)\n",
    "        if 'PT  - Book' in doc:\n",
    "            title = re.search(r'BTI\\s-\\s(.*)\\n', doc).group(1)\n",
    "            journal_iso = None\n",
    "            journal_title = None\n",
    "        else:\n",
    "            title = re.search(r'TI\\s\\s-\\s(.*)\\n', doc).group(1)\n",
    "            try:\n",
    "                journal_title = re.search(r'JT\\s\\s-\\s(.*)\\n', doc).group(1)\n",
    "                journal_iso = re.search(r'TA\\s\\s-\\s(.*)\\n', doc).group(1)\n",
    "            except:\n",
    "                journal_iso = None\n",
    "                journal_title = None                \n",
    "\n",
    "        pmid = re.search(r'PMID-\\s(.*)\\n', doc).group(1)\n",
    "        year = re.search(r'DP\\s\\s-\\s(\\d{4}).*\\n', doc).group(1)\n",
    "        if '[doi]' in doc:\n",
    "            try:\n",
    "                doi = re.search(r'(AID|LID)\\s-\\s(.*)\\[doi\\]', doc, re.DOTALL).group(2)\n",
    "            except:\n",
    "                print(doc)\n",
    "        else:\n",
    "            doi = None\n",
    "        records.append((pmid, doi, year, title, journal_title, journal_iso, abstract))\n",
    "    # print(a.group(0))\n",
    "\n",
    "    records = pd.DataFrame(records, columns=['pmid',\n",
    "                                             'doi', 'year', 'title', 'journal_title', 'journal_iso_abbreviation', 'abstract'])\n",
    "    records['category'] = category\n",
    "    records['subcategory'] = subcategory\n",
    "\n",
    "    records.to_csv(f'data/pubmed/{category}/{subcategory}.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done!\n"
     ]
    }
   ],
   "source": [
    "def pubmed_pipeline(category, subcategory, query, overwrite=True,\n",
    "                    edirect_dir='~/edirect'):\n",
    "    subcategory = subcategory.replace('/', '')\n",
    "    fname = Path('data/pubmed/.cache') / (subcategory + '.txt')\n",
    "    fname.parent.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    if overwrite or not fname.exists():\n",
    "        edirect_dir = os.path.expanduser(edirect_dir)\n",
    "\n",
    "        # create a temporary environment with edirect in the PATH\n",
    "        edirect_env = os.environ.copy()\n",
    "        edirect_env['PATH'] = f\"{edirect_dir}:{edirect_env['PATH']}\"\n",
    "\n",
    "        # define esearch and efetch commands\n",
    "        esearch_cmd = f'esearch -db pubmed -query \"{query}\" -pub abstract'\n",
    "        efetch_cmd = f'efetch -format medline'\n",
    "\n",
    "        # run esearch to retrieve the list of PMIDs\n",
    "        esearch = subprocess.Popen(shlex.split(esearch_cmd),\n",
    "                                   stdout=subprocess.PIPE, env=edirect_env)\n",
    "        out, err = esearch.communicate()\n",
    "\n",
    "        n_results = xmltodict.parse(out)['ENTREZ_DIRECT']['Count']\n",
    "        if n_results == '0':\n",
    "            print(f'[EDirect] No results for {subcategory}')\n",
    "            return\n",
    "\n",
    "        print(f'[EDirect] Fetching {n_results} articles for {subcategory}...')\n",
    "\n",
    "        # run efetch to retrieve the abstracts in medline format\n",
    "        efetch = subprocess.Popen(shlex.split(efetch_cmd),\n",
    "                                  stdin=subprocess.PIPE,\n",
    "                                  stdout=subprocess.PIPE, env=edirect_env)\n",
    "        out, err = efetch.communicate(input=out)\n",
    "\n",
    "        # store the results in a file\n",
    "        with open(fname, 'wb') as f:\n",
    "            f.write(out)\n",
    "\n",
    "    print(f'[EDirect] Finished {subcategory}')\n",
    "    try:\n",
    "        convert_medline_to_csv(category, subcategory, fname)\n",
    "    except Exception as e:\n",
    "        print(f'[CSV] Failed to convert {subcategory}.')\n",
    "        raise e\n",
    "\n",
    "# run the pipeline in parallel. Make sure n_jobs does not exceed the NCBI API limit.\n",
    "for category, queries in all_queries.items():\n",
    "\n",
    "    Path(f'data/pubmed/{category}/').mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    jobs = [delayed(pubmed_pipeline)(category, subcategory, query, overwrite=False)\n",
    "            for subcategory, query in queries.items()\n",
    "    ]\n",
    "\n",
    "    Parallel(n_jobs=7)(jobs)\n",
    "\n",
    "clear_output()\n",
    "print('Done!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "174it [00:28,  6.08it/s]\n"
     ]
    }
   ],
   "source": [
    "# aggregate all the data into a single file\n",
    "\n",
    "corpus_files = Path('data/pubmed/').glob('**/*.csv')\n",
    "\n",
    "dfs = []\n",
    "\n",
    "for fname in tqdm(corpus_files):\n",
    "\n",
    "  # find categories from the file name\n",
    "  cats = re.findall('.*/pubmed/(.*)/(.*)\\\\.csv', str(fname))\n",
    "\n",
    "  # ignore other csv files\n",
    "  if len(cats) == 0:\n",
    "    continue\n",
    "\n",
    "  category = cats[0][0]\n",
    "  subcategory = cats[0][1]\n",
    "\n",
    "  # load the data and add category and subcategory columns\n",
    "  df = pd.read_csv(fname)\n",
    "  df['category'] = category\n",
    "  df['subcategory'] = subcategory\n",
    "  dfs.append(df)\n",
    "\n",
    "# now aggregate all the data and store the compressed csv output (takes ~ 2min).\n",
    "pd.concat(dfs).to_csv(OUTPUT_PATH, index=False, compression='gzip')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((531748, 9), (635601, 9))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# report the number of abstracts per category\n",
    "pubmed2021 = pd.read_csv('data/pubmed/abstracts_2021.csv.gz')\n",
    "pubmed2023 = pd.read_csv('data/pubmed/abstracts_2023.csv.gz')\n",
    "pubmed2021.shape, pubmed2023.shape\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cogtext",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}