changing gene info source

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 
+from concurrency import execute_multithread
 from get_index import get_engine
-from protein_viz import get_gene_names, get_protein_name, render_html
+from protein_viz import get_gene_name, get_protein_name, render_html
 
 index_repo = "ronig/siamese_protein_index"
 model_repo = "ronig/protein_search_engine"
@@ -23,12 +24,7 @@ def limit_n_results(n):
 def update_dropdown_menu(search_res):
     choices = []
     for row in search_res:
-        gene = row["gene_ids"][0]
-        if gene != "Unknown":
-            choice_parts = [row["pdb_name"], row["chain_id"], gene]
-        else:
-            choice_parts = [row["pdb_name"], row["chain_id"]]
-        choice = ",".join(choice_parts)
+        choice = ",".join([row["pdb_name"], row["chain_id"]])
         choices.append(choice)
 
     return gr.Dropdown.update(
@@ -38,18 +34,26 @@ def update_dropdown_menu(search_res):
 
 def format_search_results(raw_search_results):
     formatted_search_results = {}
-    for res in raw_search_results:
-        prot = res["pdb_name"]
-        chain = res["chain_id"]
-        value = res["score"]
-        genes = ",".join(res["gene_ids"])
-        key = f"PDB: {prot} | Chain: {chain}"
-        if genes != "Unknown":
-            key += f" | Gene: {genes}"
+    for key, value in execute_multithread(
+        func=format_search_result,
+        inputs=({"raw_result": res} for res in raw_search_results),
+        n_workers=len(raw_search_results),
+    ):
         formatted_search_results[key] = value
     return formatted_search_results
 
 
+def format_search_result(raw_result):
+    prot = raw_result["pdb_name"]
+    chain = raw_result["chain_id"]
+    value = raw_result["score"]
+    gene_name = get_gene_name(pdb_id=prot, chain_id=chain)
+    key = f"PDB: {prot} | Chain: {chain}"
+    if gene_name != "Unknown":
+        key += f" | Gene: {gene_name}"
+    return key, value
+
+
 def switch_viz(new_choice):
     choice_parts = new_choice.split(",")
     pdb_id, chain = choice_parts[0], choice_parts[1]
@@ -57,10 +61,6 @@ def switch_viz(new_choice):
     protein_name = get_protein_name(pdb_id)
 
     new_value = f"""**PDB Title**: {protein_name}"""
-    if len(choice_parts) > 2:
-        gene = choice_parts[2]
-        gene_name = get_gene_names([gene])[0]
-        new_value += f"""\n\n**Gene Name**: {gene_name.title()}"""
 
     description_update = gr.Markdown.update(value=new_value, visible=True)
     return render_html(pdb_id=pdb_id, chain=chain), title_update, description_update

concurrency.py

@@ -0,0 +1,22 @@
+import concurrent.futures
+import itertools
+from typing import Callable, Iterable
+
+
+def execute_multithread(func: Callable, inputs: Iterable, n_workers):
+    with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as executor:
+        futures = {
+            executor.submit(func, **task)
+            for task in itertools.islice(inputs, n_workers)
+        }
+
+        while futures:
+            done, futures = concurrent.futures.wait(
+                futures, return_when=concurrent.futures.FIRST_COMPLETED
+            )
+
+            for future in done:
+                yield future.result()
+
+            for task in itertools.islice(inputs, len(done)):
+                futures.add(executor.submit(func, **task))

protein_viz.py

@@ -1,6 +1,3 @@
-from typing import List
-
-import mygene
 import requests
 
 
@@ -33,11 +30,42 @@ def render_html(pdb_id, chain):
     return iframe
 
 
-def get_gene_names(genes: List[str]):
-    mg = mygene.MyGeneInfo()
-    ginfo = mg.querymany(genes, scopes="ensembl.gene", verbose=False, fields="name")
-    gene_names = [gene["name"] for gene in ginfo]
-    return gene_names
+def get_gene_name(pdb_id, chain_id):
+    entity_id = get_polymer_entity_id(chain_id, pdb_id)
+    gene_name = get_gene_name_from_polymer_entity(pdb_id=pdb_id, entity_id=entity_id)
+    return gene_name
+
+
+def get_polymer_entity_id(chain_id, pdb_id):
+    url = (
+        f"https://data.rcsb.org/rest/v1/core/"
+        f"polymer_entity_instance/{pdb_id}/{chain_id}"
+    )
+    response = requests.get(url, timeout=1)
+    if response.ok:
+        res_data = response.json()
+        entity_id = int(
+            res_data["rcsb_polymer_entity_instance_container_identifiers"]["entity_id"]
+        )
+    else:
+        entity_id = None
+    return entity_id
+
+
+def get_gene_name_from_polymer_entity(pdb_id, entity_id):
+    if entity_id:
+        url = f"https://data.rcsb.org/rest/v1/core/polymer_entity/{pdb_id}/{entity_id}"
+        response = requests.get(url, timeout=1)
+        if response.ok:
+            res_data = response.json()
+            try:
+                gene_name = res_data["rcsb_entity_source_organism"][0][
+                    "rcsb_gene_name"
+                ][0]["value"]
+                return gene_name
+            except KeyError:
+                pass
+    return "Unknown"
 
 
 def get_protein_name(pdb_id: str):
@@ -48,4 +76,5 @@ def get_protein_name(pdb_id: str):
         protein_name = data["struct"]["title"]
     else:
         protein_name = "Unknown"
+
     return protein_name