feat: add DNA RNA local blast

Author: CHERRY-ui8

graphgen/configs/search_dna_config.yaml

@@ -12,4 +12,6 @@ pipeline:
       ncbi_params:
         email: [email protected] # NCBI requires an email address
         tool: GraphGen # tool name for NCBI API
+        use_local_blast: true # whether to use local blast for DNA search
+        local_blast_db: /your_path/refseq_241 # path to local BLAST database (without .nhr extension)
 

graphgen/configs/search_protein_config.yaml

@@ -11,4 +11,5 @@ pipeline:
       data_sources: [uniprot] # data source for searcher, support: wikipedia, google, uniprot
       uniprot_params:
         use_local_blast: true # whether to use local blast for uniprot search
-        local_blast_db: /your_path/uniprot_sprot
+        local_blast_db: /your_path/2024_01/uniprot_sprot # format: /path/to/${RELEASE}/uniprot_sprot
+        # options: uniprot_sprot (recommended, high quality), uniprot_trembl, or uniprot_${RELEASE} (merged database)

graphgen/configs/search_rna_config.yaml

@@ -10,5 +10,7 @@ pipeline:
     params:
       data_sources: [rnacentral] # data source for searcher, support: wikipedia, google, uniprot, ncbi, rnacentral
       rnacentral_params:
-        {} # RNAcentral doesn't require additional parameters currently
+        use_local_blast: true # whether to use local blast for RNA search
+        local_blast_db: /your_path/refseq_rna_241 # format: /path/to/refseq_rna_${RELEASE}
+        # can also use DNA database with RNA sequences (if already built)
 

graphgen/models/searcher/db/ncbi_searcher.py

@@ -1,6 +1,9 @@
 import asyncio
 import logging
+import os
 import re
+import subprocess
+import tempfile
 import time
 from concurrent.futures import ThreadPoolExecutor
 from functools import lru_cache
@@ -38,11 +41,22 @@ class NCBISearch(BaseSearcher):
     Note: NCBI has rate limits (max 3 requests per second), delays are required between requests.
     """
 
-    def __init__(self, email: str = "[email protected]", tool: str = "GraphGen"):
+    def __init__(
+        self,
+        email: str = "[email protected]",
+        tool: str = "GraphGen",
+        use_local_blast: bool = False,
+        local_blast_db: str = "nt_db",
+    ):
         super().__init__()
         Entrez.email = email
         Entrez.tool = tool
         Entrez.timeout = 60  # 60 seconds timeout
+        self.use_local_blast = use_local_blast
+        self.local_blast_db = local_blast_db
+        if self.use_local_blast and not os.path.isfile(f"{self.local_blast_db}.nhr"):
+            logger.error("Local BLAST database files not found. Please check the path.")
+            self.use_local_blast = False
 
     @staticmethod
     def _safe_get(obj, key, default=None):
@@ -518,10 +532,47 @@ def get_best_hit(self, keyword: str) -> Optional[dict]:
             logger.error("Keyword %s not found: %s", keyword, e)
         return None
 
+    def _local_blast(self, seq: str, threshold: float) -> Optional[str]:
+        """
+        Perform local BLAST search using local BLAST database.
+        :param seq: The DNA sequence.
+        :param threshold: E-value threshold for BLAST search.
+        :return: The accession number of the best hit or None if not found.
+        """
+        try:
+            with tempfile.NamedTemporaryFile(
+                mode="w+", suffix=".fa", delete=False
+            ) as tmp:
+                tmp.write(f">query\n{seq}\n")
+                tmp_name = tmp.name
+
+            cmd = [
+                "blastn",
+                "-db",
+                self.local_blast_db,
+                "-query",
+                tmp_name,
+                "-evalue",
+                str(threshold),
+                "-max_target_seqs",
+                "1",
+                "-outfmt",
+                "6 sacc",  # only return accession
+            ]
+            logger.debug("Running local blastn: %s", " ".join(cmd))
+            out = subprocess.check_output(cmd, text=True).strip()
+            os.remove(tmp_name)
+            if out:
+                return out.split("\n", maxsplit=1)[0]
+            return None
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.error("Local blastn failed: %s", exc)
+            return None
+
     def search_by_sequence(self, sequence: str, threshold: float = 0.01) -> Optional[dict]:
         """
         Search NCBI with a DNA sequence using BLAST.
-        Note: This is a simplified version. For production, consider using local BLAST.
+        Tries local BLAST first if enabled, falls back to network BLAST.
         :param sequence: DNA sequence (FASTA format or raw sequence).
         :param threshold: E-value threshold for BLAST search.
         :return: A dictionary containing the best hit information or None if not found.
@@ -542,7 +593,16 @@ def search_by_sequence(self, sequence: str, threshold: float = 0.01) -> Optional
                     logger.error("Invalid DNA sequence provided.")
                 return None
 
-            # Use BLAST search (Note: requires network connection, may be slow)
+            # Try local BLAST first if enabled
+            accession = None
+            if self.use_local_blast:
+                accession = self._local_blast(seq, threshold)
+                if accession:
+                    logger.debug("Local BLAST found accession: %s", accession)
+                    return self.get_by_accession(accession)
+
+            # Fall back to network BLAST
+            logger.debug("Falling back to NCBIWWW.qblast.")
             logger.debug("Performing BLAST search for DNA sequence...")
             time.sleep(0.35)
 

graphgen/models/searcher/db/rnacentral_searcher.py

@@ -1,5 +1,8 @@
 import asyncio
+import os
 import re
+import subprocess
+import tempfile
 from typing import Dict, Optional, List, Any
 
 import aiohttp
@@ -23,10 +26,15 @@ class RNACentralSearch(BaseSearcher):
     API Documentation: https://rnacentral.org/api/v1
     """
 
-    def __init__(self):
+    def __init__(self, use_local_blast: bool = False, local_blast_db: str = "rna_db"):
         super().__init__()
         self.base_url = "https://rnacentral.org/api/v1"
         self.headers = {"Accept": "application/json"}
+        self.use_local_blast = use_local_blast
+        self.local_blast_db = local_blast_db
+        if self.use_local_blast and not os.path.isfile(f"{self.local_blast_db}.nhr"):
+            logger.error("Local BLAST database files not found. Please check the path.")
+            self.use_local_blast = False
 
     async def _fetch_all_xrefs(self, xrefs_url: str, session: aiohttp.ClientSession) -> List[Dict]:
         """
@@ -294,11 +302,50 @@ async def get_best_hit(self, keyword: str) -> Optional[dict]:
             logger.error("Keyword %s not found: %s", keyword, e)
             return None
 
-    async def search_by_sequence(self, sequence: str) -> Optional[dict]:
+    def _local_blast(self, seq: str, threshold: float) -> Optional[str]:
+        """
+        Perform local BLAST search using local BLAST database.
+        :param seq: The RNA sequence.
+        :param threshold: E-value threshold for BLAST search.
+        :return: The accession/ID of the best hit or None if not found.
+        """
+        try:
+            with tempfile.NamedTemporaryFile(
+                mode="w+", suffix=".fa", delete=False
+            ) as tmp:
+                tmp.write(f">query\n{seq}\n")
+                tmp_name = tmp.name
+
+            cmd = [
+                "blastn",
+                "-db",
+                self.local_blast_db,
+                "-query",
+                tmp_name,
+                "-evalue",
+                str(threshold),
+                "-max_target_seqs",
+                "1",
+                "-outfmt",
+                "6 sacc",  # only return accession
+            ]
+            logger.debug("Running local blastn for RNA: %s", " ".join(cmd))
+            out = subprocess.check_output(cmd, text=True).strip()
+            os.remove(tmp_name)
+            if out:
+                return out.split("\n", maxsplit=1)[0]
+            return None
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.error("Local blastn failed: %s", exc)
+            return None
+
+    async def search_by_sequence(self, sequence: str, threshold: float = 0.01) -> Optional[dict]:
         """
         Search RNAcentral with an RNA sequence.
+        Tries local BLAST first if enabled, falls back to RNAcentral API.
         Unified approach: Find RNA ID from sequence search, then call get_by_rna_id() for complete information.
         :param sequence: RNA sequence (FASTA format or raw sequence).
+        :param threshold: E-value threshold for BLAST search.
         :return: A dictionary containing complete RNA information or None if not found.
         """
         try:
@@ -318,7 +365,23 @@ async def search_by_sequence(self, sequence: str) -> Optional[dict]:
                 logger.error("Empty RNA sequence provided.")
                 return None
 
-            # RNAcentral API supports sequence search
+            # Try local BLAST first if enabled
+            if self.use_local_blast:
+                accession = self._local_blast(seq, threshold)
+                if accession:
+                    logger.debug("Local BLAST found accession: %s", accession)
+                    # Try to get RNA ID from accession (may need conversion)
+                    # For now, try using accession as RNA ID or search by it
+                    result = await self.get_by_rna_id(accession)
+                    if result:
+                        return result
+                    # If not found by ID, try keyword search
+                    result = await self.get_best_hit(accession)
+                    if result:
+                        return result
+
+            # Fall back to RNAcentral API
+            logger.debug("Falling back to RNAcentral API.")
             async with aiohttp.ClientSession() as session:
                 search_url = f"{self.base_url}/rna"
                 params = {"sequence": seq, "format": "json"}
@@ -373,7 +436,7 @@ async def search_by_sequence(self, sequence: str) -> Optional[dict]:
         reraise=True,
     )
     async def search(
-        self, query: str, threshold: float = 0.7, **kwargs
+        self, query: str, threshold: float = 0.1, **kwargs
     ) -> Optional[Dict]:
         """
         Search RNAcentral with either an RNAcentral ID, keyword, or RNA sequence.
@@ -395,7 +458,7 @@ async def search(
         if query.startswith(">") or (
             re.fullmatch(r"[AUCGN\s]+", query, re.I) and "U" in query.upper()
         ):
-            result = await self.search_by_sequence(query)
+            result = await self.search_by_sequence(query, threshold)
         # check if RNAcentral ID (typically starts with URS)
         elif re.fullmatch(r"URS\d+", query, re.I):
             result = await self.get_by_rna_id(query)