pubmad.utils

  0import logging
  1from transformers import AutoTokenizer, AutoModelForSequenceClassification
  2import torch
  3from typing import List, Tuple
  4import requests
  5import networkx as nx
  6from itertools import product
  7from pubmad.types import Article, Entity
  8import matplotlib.pyplot as plt
  9import os
 10from nltk.tokenize import sent_tokenize
 11import pickle
 12from pathlib import Path
 13import time
 14from Bio import Entrez, Medline
 15from pyvis.network import Network
 16import re
 17import random
 18import numpy as np
 19import IPython
 20from math import ceil
 21
 22import nltk
 23nltk.download('punkt', quiet=True)
 24
 25# Number of sequences per batch
 26MAX_SEQ_PER_BATCH = 2
 27
 28print("Loading models...")
 29chemprot_model = AutoModelForSequenceClassification.from_pretrained(
 30    "pier297/autotrain-chemprot-re-838426740")
 31chemprot_tokenizer = AutoTokenizer.from_pretrained(
 32    "pier297/autotrain-chemprot-re-838426740")
 33
 34rel_tokenizer = AutoTokenizer.from_pretrained(
 35    "JacopoBandoni/BioBertRelationGenesDiseases")
 36
 37rel_model = AutoModelForSequenceClassification.from_pretrained(
 38    "JacopoBandoni/BioBertRelationGenesDiseases")
 39print('Finished loading models.')
 40
 41chemprot_model.eval()
 42rel_model.eval()
 43
 44device = 'cuda' if torch.cuda.is_available() else 'cpu'
 45rel_model = rel_model.to(device)
 46chemprot_model = chemprot_model.to(device)
 47
 48# logging.disable(logging.INFO) # disable INFO and DEBUG logging everywhere
 49# or
 50# disable WARNING, INFO and DEBUG logging everywhere
 51logging.disable(logging.WARNING)
 52
 53
 54def download_articles(title: str, start_year: int, end_year: int, max_results: int = 100, author: str = '', sort_by: str = 'relevance') -> List[Article]:
 55    """
 56    Download articles from PubMed using BioPython library.
 57    Args:
 58        query (str): The query to search for.
 59        start_year (int): The start year to search for.
 60        end_year (int): The end year to search for.
 61        max_results (int): The maximum number of results to return.
 62        author (str): The author to search for, leave empty to search for all authors.
 63        sort_by (str): The sort order to use. Can be 'relevance' to retrieve the most relevant results, or 'date' to sort by publication date. Defaults to 'relevance'.
 64    Returns:
 65        List[Article] A list of articles.
 66    """
 67    current_path = Path(os.getcwd()) / 'cache'
 68
 69    if not os.path.exists(current_path):
 70        os.mkdir(current_path)
 71
 72    file_name = '{}_{}_{}_{}.txt'.format(
 73        title, start_year, end_year, max_results)
 74    if os.path.exists(current_path / file_name):
 75        print("Loading from cache...")
 76        # TODO: clear_cache is not implemented yet
 77        with open(current_path / file_name, 'rb') as f:
 78            return pickle.load(f)
 79
 80    Entrez.email = 'pubmadbiosearch@gmail.com'
 81
 82    if author != "":
 83        query = '(' + title + ') AND (' + author + '[Author])'
 84    else:
 85        query = title
 86
 87    handle = Entrez.esearch(db='pubmed',
 88                            sort=sort_by,
 89                            retmax=max_results,
 90                            retmode='xml',
 91                            datetype='pdat',
 92                            mindate=str(start_year),
 93                            maxdate=str(end_year),
 94                            term=query)
 95    results = Entrez.read(handle)
 96    id_list = results['IdList']
 97    ids = ','.join(id_list)
 98    handle = Entrez.efetch(db='pubmed', rettype="medline",
 99                           id=ids, retmode='text')
100    results = Medline.parse(handle)
101    results = list(results)
102    articles = []
103    for article in results:
104        abstract = article.get("AB", "?")
105        title = article.get("TI", "?")
106        if abstract != "?" and title != "?":
107            if abstract.find("This article has been withdrawn at the request of the author(s) and/or editor") == -1 and \
108               abstract.find("An amendment to this paper has been published and can be accessed via a link at the top of the paper") == -1 and \
109               abstract.find("This corrects the article") == -1 and abstract.find("A Correction to this paaper has been published") == -1:
110                articles.append(Article(title=title, abstract=abstract,
111                                        pmid=article['PMID'], full_text='', publication_data=None))
112
113    print("Found {} articles".format(len(articles)))
114
115    # save articles to pickle file
116    with open(current_path / file_name, 'wb') as f:
117        pickle.dump(articles, f)
118    return articles
119
120
121def extract_entities_pmids(pmids: List[str]) -> List[List[Entity]]:
122    """Extract entities from one or more articles identified by the pmids provided
123
124    Args:
125        pmids (List[str]): list of pmids to extract entities from
126
127    Returns:
128        List[List[Entity]]: list of entities for each article
129    """
130
131    if len(pmids) == 0:
132        return []
133
134    # Query BERN2 server for entities.
135    extracted = False
136    bern_result = None
137
138    i = 0
139
140    while not extracted and i < 3:
141        try:
142            bern_result = query_pmid(pmids)
143            extracted = True
144        except Exception as e:
145            print("Error querying BERN2 server")
146            print("Sleep for 3 seconds and try again to not get banned by BERN2")
147            time.sleep(3)
148            i += 1
149            continue
150
151    entities = []
152
153    # Parse the result
154    for i, article in enumerate(bern_result):
155        entities_per_article = []
156        for entity in article["annotations"]:
157            if (entity['obj'] == 'disease' or entity['obj'] == 'gene' or entity['obj'] == 'drug'):
158                new_entity = Entity(mesh_id=entity['id'], mention=entity['mention'], type=entity['obj'],
159                                    prob=entity['prob'], span_begin=entity['span']['begin'], span_end=entity['span']['end'], pmid=pmids[i])
160                entities_per_article.append(new_entity)
161        entities.append(entities_per_article)
162
163    return entities
164
165
166def extract_entities(article: Article, source: str = 'abstract') -> List[Entity]:
167    """
168    Extract entities from an article using BERN2 (online).
169
170    Args:
171        article (Article): The article to extract entities from. 
172        source (str): The source to extract entities from. Can be 'abstract' or 'full_text'. Defaults to 'abstract'.
173
174    Returns:
175        List[Entity]: A list of entities.
176    """
177    # We call bern2 server for ner
178    text = ''
179    if source == 'abstract':
180        text = article.abstract
181    elif source == 'full_text':
182        text = article.full_text
183    else:
184        raise ValueError("Invalid source: {}".format(source))
185
186    # Query BERN2 server for entities.
187    extracted = False
188    bern_result = None
189
190    i = 0
191
192    while not extracted and i < 3:
193        try:
194            bern_result = query_plain(text)
195            extracted = True
196        except Exception as e:
197            print("Error querying BERN2 server")
198            print("Sleep for 3 seconds and try again to not get banned by BERN2")
199            time.sleep(3)
200            i += 1
201            continue
202
203    if i == 3:
204        return []
205
206    # Parse the result.
207    annotations = bern_result['annotations']
208    entities = []
209    for entity in annotations:
210        if (entity['obj'] == 'disease' or entity['obj'] == 'gene' or entity['obj'] == 'drug'):
211            pmid = article.pmid
212            if type(article.pmid) == list:
213                pmid = article.pmid[0]
214            new_entity = Entity(mesh_id=entity['id'], mention=entity['mention'], type=entity['obj'],
215                                prob=entity['prob'], span_begin=entity['span']['begin'], span_end=entity['span']['end'], pmid=pmid)
216            entities.append(new_entity)
217
218    return entities
219
220
221def display_graph(graph: nx.Graph, hide_isolated_nodes: bool = True, show: bool = True):
222    """
223    Display a graph.
224
225    Args:
226        graph (nx.Graph): The graph to display.
227        hide_isolated_nodes (bool): Whether to hide isolated nodes. Defaults to True.
228        show (bool): Whether to show the graph. Defaults to True. Otherwise it shows the window without stopping the program.
229    """
230    plt.figure()
231    if hide_isolated_nodes:
232        graph.remove_nodes_from(list(nx.isolates(graph)))
233
234    gene_nodes = [n for n, d in graph.nodes(data=True) if d['type'] == 'gene']
235    disease_nodes = [n for n, d in graph.nodes(
236        data=True) if d['type'] == 'disease']
237    other_nodes = [n for n, d in graph.nodes(
238        data=True) if d['type'] != 'gene' and d['type'] != 'disease']
239
240    pos = nx.spring_layout(graph, 100)
241
242    weights = nx.get_edge_attributes(graph, 'weight')
243
244    nx.draw_networkx_nodes(graph, pos, nodelist=gene_nodes,
245                           node_color='r', node_size=100, alpha=0.8, label='gene')
246    nx.draw_networkx_nodes(graph, pos, nodelist=disease_nodes,
247                           node_color='b', node_size=100, alpha=0.8, label='disease')
248    nx.draw_networkx_nodes(graph, pos, nodelist=other_nodes,
249                           node_color='g', node_size=100, alpha=0.8, label='drug')
250
251    nx.draw_networkx_labels(graph, pos, labels={
252                            n: d['mention'] for n, d in graph.nodes(data=True)}, font_size=10)
253
254    nx.draw_networkx_edges(graph, pos, width=1.0, alpha=0.5)
255    plt.axis('off')
256    plt.legend()
257    if show:
258        plt.show()
259    else:
260        plt.draw()
261
262
263def extract_naive_relations(entities: List[Entity]) -> List[Tuple[Entity, Entity, float]]:
264    '''
265    Extract relations from a list of entities.
266
267    Args:
268        entities (List[Entity]): A list of entities.
269
270    Returns:
271        List[Tuple[Entity, Entity, float]]: A list of relations.
272    '''
273    # Connect all the entities to each other.
274    relations = []
275    for entity1, entity2 in product(entities, entities):
276        if entity1.mesh_id != entity2.mesh_id:
277            # 1 meaning that the relation is present.
278            relations.append((entity1, entity2, 1))
279    return relations
280
281
282def extract_biobert_relations(article: Article, source: str = 'abstract', clear_cache: bool = False) -> Tuple[List[Entity], List[Tuple[Entity, Entity]], bool]:
283    """
284    Extract the entities from the article using BERN2 (online).
285    Then it calls BioBERT to extract relations between genes and diseases.
286
287    Args:
288        article (Article): The article to extract entities from.
289        source (str): The source to extract entities from. Can be 'abstract' or 'full_text'. Defaults to 'abstract'.
290        clear_cache (bool): Whether to clear the BioBERT cache. Defaults to False.
291
292    Returns:
293        Tuple[List[Entity], List[Tuple[Entity, Entity]], bool]: A tuple of entities, relations and whether the cache was used.
294
295    """
296    file_name = str(article.pmid)[:128] + '.txt'
297    file_name = file_name.replace('\n', '_')
298    path = Path(os.getcwd()) / 'cache'
299
300    if not os.path.exists(path / 'entities'):
301        # create directory
302        os.mkdir(path / 'entities')
303
304    if not os.path.exists(path / 'relations'):
305        # create directory
306        os.mkdir(path / 'relations')
307
308    if os.path.exists(path / 'entities' / file_name) and os.path.exists(path / 'relations' / file_name):
309        if clear_cache:
310            print("Clearing entities and relations cache for PMID {}".format(
311                article.pmid))
312            os.remove(path / 'entities' / file_name)
313            os.remove(path / 'relations' / file_name)
314        else:
315            with open(path / 'entities' / file_name, 'rb') as f:
316                entities = pickle.load(f)
317            with open(path / 'relations' / file_name, 'rb') as f:
318                relations = pickle.load(f)
319            return entities, relations, True
320
321    text = ''
322    if source == 'abstract':
323        text = article.abstract
324    elif source == 'full_text':
325        text = article.full_text
326    else:
327        raise ValueError("Invalid source: {}".format(source))
328
329    entities = extract_entities(article)
330
331    span_sentences = tokenize_into_sentences(article, source)
332
333    # divide entities in gene and disease entities
334    gene_entities = []
335    disease_entities = []
336    drug_entities = []
337    for entity in entities:
338        if entity.type == "disease":
339            disease_entities.append(entity)
340        elif entity.type == "gene":
341            gene_entities.append(entity)
342        elif entity.type == "drug":
343            drug_entities.append(entity)
344
345    relations = []
346    chemprot_batch = []
347    biobert_batch = []
348    for gene_idx, gene_entity in enumerate(gene_entities):
349        # extract gene-drug relations using fine-tuned bert on chemprot
350        for drug_idx, drug_entity in enumerate(drug_entities):
351            # find the sentence that contains the gene and disease
352            sentence_index_gene = find_entity(gene_entity, span_sentences)
353            sentence_index_drug = find_entity(drug_entity, span_sentences)
354
355            masked_text = ''
356            if gene_entity.span_begin < drug_entity.span_begin:
357                # wrap the gene_entity in << >>
358                # and the drug with [[ ]]
359                masked_text = text[span_sentences[sentence_index_gene][0]:gene_entity.span_begin] + "<< " + text[gene_entity.span_begin:gene_entity.span_end] + " >>" + \
360                    text[gene_entity.span_end:drug_entity.span_begin] + "[[ " + text[drug_entity.span_begin:drug_entity.span_end] + \
361                    " ]]" + \
362                    text[drug_entity.span_end:span_sentences[sentence_index_drug][1]]
363            else:
364                # wrap the drug_entity in << >>
365                # and the gene with [[ ]]
366                masked_text = text[span_sentences[sentence_index_drug][0]:drug_entity.span_begin] + "<< " + text[drug_entity.span_begin:drug_entity.span_end] + " >>" + \
367                    text[drug_entity.span_end:gene_entity.span_begin] + "[[ " + text[gene_entity.span_begin:gene_entity.span_end] + \
368                    " ]]" + \
369                    text[gene_entity.span_end:span_sentences[sentence_index_gene][1]]
370
371            chemprot_batch.append(
372                {'gene_idx': gene_idx, 'drug_idx': drug_idx, 'masked_text': masked_text})
373
374        # extract gene-disease relations using biobert
375        for disease_idx, disease_entity in enumerate(disease_entities):
376            # find the sentence that contains the gene and disease
377            sentence_index_gene = find_entity(gene_entity, span_sentences)
378            sentence_index_disease = find_entity(
379                disease_entity, span_sentences)
380
381            masked_text = ''
382            if gene_entity.span_begin < disease_entity.span_begin:
383                masked_text = text[span_sentences[sentence_index_gene][0]:gene_entity.span_begin] + "@GENE$" + text[gene_entity.span_end:
384                                                                                                                    disease_entity.span_begin] + "@DISEASE$" + text[disease_entity.span_end:span_sentences[sentence_index_disease][1]]
385            else:
386                masked_text = text[span_sentences[sentence_index_disease][0]:disease_entity.span_begin] + "@DISEASE$" + \
387                    text[disease_entity.span_end:gene_entity.span_begin] + "@GENE$" + \
388                    text[gene_entity.span_end:span_sentences[sentence_index_gene][1]]
389
390            biobert_batch.append(
391                {'gene_idx': gene_idx, 'disease_idx': disease_idx, 'masked_text': masked_text})
392
393    # Predict the relations using biobert
394    if len(biobert_batch) > 0:
395        for batch_idx in range(0, len(biobert_batch), MAX_SEQ_PER_BATCH):
396            batch = biobert_batch[batch_idx:batch_idx + MAX_SEQ_PER_BATCH]
397            masked_texts = [x['masked_text'] for x in batch]
398            tok_texts = rel_tokenizer(
399                masked_texts, max_length=512, padding=True, truncation=True, return_tensors='pt').to(device)
400            outputs = rel_model(**tok_texts)
401            class_logits = outputs["logits"].detach().cpu().numpy()
402            # len x 2
403            # apply softmax to get the probabilities
404            class_probs = np.exp(
405                class_logits) / np.sum(np.exp(class_logits), axis=1, keepdims=True)
406            for i in range(len(class_logits)):
407                if class_probs[i][0] > 0.5:
408                    j = i + batch_idx
409                    gene = gene_entities[biobert_batch[j]['gene_idx']]
410                    disease = disease_entities[biobert_batch[j]['disease_idx']]
411                    relations.append((gene, disease, float(class_probs[i][0])))
412
413            # free gpu memory
414            del tok_texts
415
416    # # Predict the relations using chemprot
417    if len(chemprot_batch) > 0:
418        for batch_idx in range(0, len(chemprot_batch), MAX_SEQ_PER_BATCH):
419            batch = chemprot_batch[batch_idx:batch_idx + MAX_SEQ_PER_BATCH]
420            masked_texts = [x['masked_text'] for x in batch]
421            tok_texts = chemprot_tokenizer(
422                masked_texts, max_length=512, padding=True, truncation=True, return_tensors='pt').to(device)
423            outputs = chemprot_model(**tok_texts)
424            class_logits = outputs["logits"].detach().cpu().numpy()
425            # len x 13
426            # apply softmax to get the probabilities
427            class_probs = np.exp(
428                class_logits) / np.sum(np.exp(class_logits), axis=1, keepdims=True)
429            for i in range(len(class_logits)):
430                # if 1 class is above 0.5, then we consider it as a relation
431                max_p = np.max(class_probs[i])
432                if max_p > 0.5:
433                    j = i + batch_idx
434                    gene = gene_entities[chemprot_batch[j]['gene_idx']]
435                    drug = drug_entities[chemprot_batch[j]['drug_idx']]
436                    relations.append((gene, drug, float(max_p)))
437
438            # free gpu memory
439            del tok_texts
440
441    # save entities and relations to file with pickle
442    with open(path / 'entities' / file_name, 'wb') as f:
443        pickle.dump(entities, f)
444    with open(path / 'relations' / file_name, 'wb') as f:
445        pickle.dump(relations, f)
446
447    return entities, relations, False
448
449
450def query_plain(text, url="http://bern2.korea.ac.kr/plain"):
451    """
452    Query the BERN2 server for plain text.
453    Args:
454        text (str): the text to be annotated
455        url (str): the url of the BERN2 server.
456    """
457    # Limit the text at 5000 characters
458    if len(text) > 5000:
459        text = text[:5000]
460
461    result = requests.post(url, json={'text': text})
462
463    if result.status_code != 200:
464        print("Error: {}".format(result.status_code))
465        raise Exception("Error: {}".format(result.status_code))
466
467    return result.json()
468
469
470def query_pmid(pmids, url="http://bern2.korea.ac.kr/pubmed"):
471    """query the BERN2 server for pmids"""
472    return requests.get(url + "/" + ",".join(pmids)).json()
473
474
475def tokenize_into_sentences(article: Article, source: str = 'abstract') -> List[List[int]]:
476    """tokenize the text into sentences returning the index of begin and end of each sentence
477
478    Args:
479        article (Article): article to tokenize.
480        source (str): The source of the text.
481
482    Returns:
483        List: A list of couples containing the begin and end of each sentence.
484    """
485    text = ''
486    if source == 'abstract':
487        text = article.abstract
488    elif source == 'full_text':
489        text = article.full_text
490    else:
491        raise ValueError("Invalid source: {}".format(source))
492
493    sentences = sent_tokenize(text)
494    span_sentences = []
495    # for each sentence, put the start and end index of each sentence
496    start = 0
497    for sentence in sentences:
498        span_sentences.append([start, start + len(sentence)])
499        start += len(sentence) + 1
500    return span_sentences
501
502
503def find_entity(entity: Entity, span_sentences: List[List[int]]) -> int:
504    """find in which sentence the entity is
505
506    Args:
507        entity (Entity): entity to find
508        span_sentences (List): A list of couples containing the index of the begin and end of each sentence.
509
510    Returns:
511        int: index of the sentence containing the entity
512    """
513    for i in range(len(span_sentences)):
514        if span_sentences[i][0] <= entity.span_begin and entity.span_end <= span_sentences[i][1]:
515            return i
516    return -1
517
518
519def add_title_node(G, n, d):
520    result_html = f"""
521  <h3>{d['mention']}</h3>
522  """
523    # articles:
524    #  <ul>
525    # """
526    #pmids = re.split("[,\\n]", d['pmid'])
527    # for pmid in pmids:
528    #  result_html += f"<li>{pmid}</li>"
529    #result_html += "</ul>"
530    return result_html
531
532
533def add_title_edge(G, edge):
534    # find intersection between pmids of the two nodes
535    node1 = dict(G.nodes(data=True))[edge[0]]
536    node2 = dict(G.nodes(data=True))[edge[1]]
537    n1_pmids = re.split(",", node1['pmid'])
538    n2_pmids = re.split(",", node2['pmid'])
539    intersection = list(set(n1_pmids) & set(n2_pmids))
540
541    # create html
542    result_html = f"""<p>common articles between <strong>{node1['mention']}</strong> and <strong>{node2['mention']}</strong></p>
543  
544  <ul>"""
545    for pmid in intersection:
546        result_html += f"<li>{pmid}</li>"
547    result_html += "</ul>"
548    return result_html
549
550
551def _html_graph_communities(G, communities, net, colors):
552    for i, community in enumerate(communities):
553        for node_id in community:
554            node = net.get_node(node_id)
555            node_type = dict(G.nodes(data=True))[node_id]['type']
556            node['color'] = colors[i]
557
558            if (node_type == 'gene'):
559                node['shape'] = 'dot'
560            elif (node_type == 'disease'):
561                node['shape'] = 'diamond'
562            elif (node_type == 'drug'):
563                node['shape'] = 'star'
564
565
566def _find_community(node_id, communities):
567    """
568        Find the index of the community containing the node.
569    """
570    for i, community in enumerate(communities):
571        if node_id in community:
572            return i
573
574
575def html_mark_subgraph(G, subG, name="nodes", hide_isolated_nodes=True):
576    """
577        Create a html file, named <name>, containing the graph G with the subgraph subG marked.
578
579        Args:
580            G (networkx.Graph): The graph to be displayed.
581            subG (networkx.Graph): The subgraph to be marked.
582            name (str): The name of the html file. Default is "nodes".
583            hide_isolated_nodes (bool): If True, the isolated nodes are not displayed. Default is True.
584    """
585
586    if hide_isolated_nodes:
587        G.remove_nodes_from(list(nx.isolates(G)))
588        subG.remove_nodes_from(list(nx.isolates(subG)))
589
590    net = Network('500px', '600px', notebook=True)
591
592    for n, d in G.nodes(data=True):
593        color = '#cccccc80'
594        if n not in subG.nodes():
595            net.add_node(n, d['mention'], color='#cccccc80',
596                title=add_title_node(G, n, d))
597
598    for n, d in subG.nodes(data=True):
599        color = '#cccccc80'
600        if d['type'] == 'gene':
601            color = '#0DA3E4'
602        elif d['type'] == 'disease':
603            color = '#bf4d2d'
604        elif d['type'] == 'drug':
605            color = '#5ef951'
606        net.add_node(n, d['mention'], color=color,
607            title=add_title_node(G, n, d))
608
609    # add edges not in subG
610    for edge in G.edges(data='True'):
611        weight = G.get_edge_data(edge[0], edge[1])['weight']
612        net.add_edge(edge[0], edge[1], value=weight,
613            color='#cccccc70', title=add_title_edge(G, edge))
614
615    # add edges in subG
616    for edge in subG.edges(data='True'):
617        weight = G.get_edge_data(edge[0], edge[1])['weight']
618        net.add_edge(edge[0], edge[1], value=weight,
619            color='#F0EB5A', title=add_title_edge(subG, edge))
620
621    net.show_buttons(filter_=['physics'])
622    net.show(name + ".html")
623    IPython.display.HTML(filename=Path(os.getcwd())/(name + ".html"))
624
625
626def html_graph(G, name="nodes", communities=None, hide_isolated_nodes=True):
627    """
628        Create an html file, named <name>, representing the graph G using vis.js.
629
630        Args:
631            G (NetworkX graph): The graph to be displayed.
632            name (str): The name of the file to be created. Default is "nodes".
633            communities (List): A list of communities to be displayed. If None, no communities are displayed. Default is None.
634            hide_isolated_nodes (bool): If True, hide isolated nodes. Default is True.
635
636        Returns:
637            A pyvis network object and writes the file to the current directory.
638    """
639
640    if hide_isolated_nodes:
641        G.remove_nodes_from(list(nx.isolates(G)))
642
643    net = Network('500px', '600px', notebook=True)
644
645    for n, d in G.nodes(data=True):
646        color = '#ccc'
647        if d['type'] == 'gene':
648            color = '#0DA3E4'
649        elif d['type'] == 'disease':
650            color = '#bf4d2d'
651        elif d['type'] == 'drug':
652            color = '#5ef951'
653        net.add_node(n, d['mention'], color=color,
654                     title=add_title_node(G, n, d))
655
656    # if communities is not None and len(communities) > 0
657    colors = []
658    if communities is not None and len(communities) > 0:
659        # generate random color for each community
660        colors = ["#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)]) + "90"
661                  for i in range(len(communities))]
662        _html_graph_communities(G, communities, net, colors)
663
664    # add edges
665    for edge in G.edges(data='True'):
666        weight = G.get_edge_data(edge[0], edge[1])['weight']
667        if communities is not None and len(communities) > 0:
668            # find the community of the two nodes
669            community_node1_index = _find_community(edge[0], communities)
670            if edge[1] in communities[community_node1_index]:
671                net.add_edge(edge[0], edge[1], value=weight,
672                             color=colors[community_node1_index], title=add_title_edge(G, edge))
673            else:
674                net.add_edge(edge[0], edge[1], value=weight,
675                             color='#F0EB5A', title=add_title_edge(G, edge))
676        else:
677            net.add_edge(edge[0], edge[1], value=weight,
678                         color='#F0EB5A', title=add_title_edge(G, edge))
679
680    net.show_buttons(filter_=['physics'])
681    net.show(name + ".html")
682    IPython.display.HTML(filename=Path(os.getcwd())/(name + ".html"))
683
684    return net
685
686
687def filter_by_centrality(G, percentage_threshold: float = 0.1):
688    '''
689    Filter the graph by mantaining the 'percentage_threshold' of the nodes with the highest centrality.
690    Args:
691        G: networkx graph
692        percentage_threshold: float between 0 and 1. The percentage of most central nodes to keep.
693    '''
694    mesh_id_to_degree = dict()
695    for node_id in G.nodes():
696        mesh_id_to_degree[node_id] = G.degree(node_id)
697
698    sorted_mesh_ids = sorted(
699        mesh_id_to_degree, key=mesh_id_to_degree.get, reverse=True)
700
701    # take only the nodes with the highest centrality
702    nodes_to_keep = sorted_mesh_ids[:ceil(
703        len(sorted_mesh_ids)*percentage_threshold)]
704
705    return G.subgraph(nodes_to_keep).copy()
706
707
708def filter_by_name(G, name: str):
709    """
710    Filter the graph by mantaining the nodes with the given name.
711    Args:
712        G: networkx graph
713        name: str. The name of the nodes to keep.
714    Returns:
715        A list of nodes that include the given name.
716    """
717    nodes = []
718    for n, d in G.nodes(data=True):
719        if name in d['mention']:
720            d['mesh_id'] = n
721            nodes.append(d)
722
723    return nodes
724
725
726def filter_by_category(G, category: str, max_number_of_nodes: int = None, sort_by: str = 'degree'):
727    '''
728    Returns the nodes with the type category
729    Args:
730        G (Networkx Graph): graph to filter
731        category (str): category to filter by: 'gene', 'disease', 'drug'
732        max_number_of_nodes (int): maximum number of nodes to return. If None, all nodes are returned
733        sort_by (str): sort by 'degree' or 'name'
734    Returns:
735        A list of nodes with the given category
736    '''
737    if category not in ['gene', 'disease', 'drug']:
738        raise ValueError("category must be one of 'gene', 'disease', 'drug'")
739
740    nodes = []
741    for n, d in G.nodes(data=True):
742        if d['type'] == category:
743            d['mesh_id'] = n
744            nodes.append(d)
745
746    if sort_by == 'degree':
747        nodes = sorted(nodes, key=lambda node: G.degree(
748            node['mesh_id']), reverse=True)
749    elif sort_by == 'name':
750        nodes = sorted(nodes, key=lambda node: node['mention'])
751
752    if max_number_of_nodes is not None:
753        nodes = nodes[:max_number_of_nodes]
754
755    return nodes
756
757
758def expand_from_node(G, node, max_distance):
759    ''' 
760    Returns the subgraph that starting from the node id link all the node with distance k from id
761    Args:
762        G (Networkx Graph): graph to filter
763        node (str): node id
764        max_distance (int): maximum distance to expand
765    Returns:
766        A new graph containing the nodes with distance k from node
767    '''
768    if type(node) == List:
769        node = node[0]
770        print("node is a list, taking the first element")
771        print("Otherwise use 'expand_from_nodes'")
772    node_list = [node['mesh_id']]
773
774    for node, succ in nx.bfs_successors(G, node['mesh_id'], depth_limit=max_distance):
775        node_list += list(succ)
776
777    return G.subgraph(node_list).copy()
778
779
780def expand_from_nodes(G, nodes, max_distance: int):
781    '''
782    Returns the subgraph that starting from the node id link all the node with distance k from id 
783    Args:
784        G (Networkx Graph): graph to filter
785        nodes (list): list of nodes mesh_id
786        max_distance (int): maximum distance to expand
787    Returns:
788        A new graph containing the nodes with distance k from node
789    '''
790    if type(nodes[0]) == dict:
791        node_id_list = [node['mesh_id'] for node in nodes]
792
793        node_list = [node['mesh_id'] for node in nodes]
794    elif type(nodes[0]) == str:
795        node_id_list = nodes
796        node_list = nodes
797    else:
798        raise ValueError("nodes must be a list of mesh_id or a list of node objects")
799
800    for node_id in node_id_list:
801        for _, succ in nx.bfs_successors(G, node_id, depth_limit=max_distance):
802            node_list += list(succ)
803
804    return G.subgraph(node_list).copy()
805
806
807def get_common_articles(node1, node2):
808    '''
809    Returns the articles that both nodes mention
810    Args:
811        node1 (dict): node object
812        node2 (dict): node object
813    Returns:
814        A list of articles
815    '''
816    articles_node1 = node1['pmid']
817    articles_node2 = node2['pmid']
818    return list(set(articles_node1.split(',')).intersection(articles_node2.split(',')))
819
820
821def search_path(G, from_node, to_node):
822    '''
823    Returns the path between two nodes in the graph
824    Args:
825        G (Networkx Graph): graph to filter
826        from_node (str): node id
827        to_node (str): node id
828    Returns:
829        A list of nodes that form the path
830    '''
831    try:
832        return nx.shortest_path(G, source=from_node['mesh_id'], target=to_node['mesh_id'])
833    except nx.NetworkXNoPath:
834        return []
835
836
837def search_paths_to_category(G, from_node, to_category: str):
838    '''
839    Returns the paths from 'from_node' to nodes with the type 'to_category'
840    Args:
841        G (Networkx Graph): graph to filter
842        from_node (str): node id
843        to_category (str): to_category to filter by: 'gene', 'disease', 'drug'
844    Returns:
845        A list of paths
846    '''
847    if to_category not in ['gene', 'disease', 'drug']:
848        raise ValueError("to_category must be one of 'gene', 'disease', 'drug'")
849
850    targets_nodes = filter_by_category(G, to_category)
851    paths = []
852    for target_node in targets_nodes:
853        paths.append(search_path(G, from_node, target_node))
854
855    # remove empty paths
856    paths = [path for path in paths if len(path) > 0]
857
858    return paths
859
860
861def get_graph_from_path(G, path):
862    '''
863    Returns the subgraph that contains the nodes in the path
864    Args:
865        G (Networkx Graph): graph to filter
866        path (list): list of nodes mesh_id
867    Returns:
868        A new graph containing the nodes in the path
869    '''
870    return G.subgraph(path).copy()
871
872
873def get_graph_from_paths(G, paths):
874    '''
875    Returns the subgraph that contains the nodes in the paths
876    Args:
877        G (Networkx Graph): graph to filter
878        paths (list): list of paths
879    Returns:
880        A new graph containing the nodes in the paths
881    '''
882    nodes = [node for path in paths for node in path]
883    return G.subgraph(nodes).copy()