add fasta input support to the databases in pipeline

Author: dsmutin

pipeline.py

@@ -40,6 +40,7 @@
 from PROBESt.args import arguments_parse
 from PROBESt.modeling import run_modeling
 from PROBESt.dedegeneration import run_dedegeneration
+from PROBESt.prepare_blast import prepare_bases_if_needed
 
 # Functions
 
@@ -68,6 +69,19 @@ def merge_iter(iter: int):
 print("\n---- PROBESt v.0.2.0 ----\n")
 print("Arguments passed")
 
+# Create output directory early (needed for BLAST database preparation)
+os.makedirs(args.output, exist_ok=True)
+
+# 0.1. Prepare BLAST databases from FASTA directories if needed ----
+args.true_base, args.false_base = prepare_bases_if_needed(
+    true_base=args.true_base,
+    false_bases=args.false_base,
+    output_dir=args.output,
+    contig_table_path=args.contig_table,
+    tmp_dir=args.prep_db_tmp,
+    script_path=args.script_path
+)
+
 # Create TMP
 os.makedirs(out_dir(0), exist_ok=True)
 

src/PROBESt/args.py

@@ -44,13 +44,20 @@ def arguments_parse():
                         help="Path(s) to the BLAST database(s) for non-specific testing. These databases are used to filter out non-specific probes. Wildcards are not accepted.")
 
     parser.add_argument("-c", "--contig_table",
-                        required=True,
-                        help="Path to a .tsv table containing BLAST database information.")
+                        required=False,
+                        default=None,
+                        help="Path to a .tsv table containing BLAST database information. If not provided and FASTA directories are used for bases, it will be auto-generated in the output directory.")
 
     parser.add_argument("-o", "--output",
                         required=True,
                         help="Output directory path for storing results.")
 
+    # BLAST database preparation arguments (for FASTA directory inputs)
+    parser.add_argument("--prep_db_tmp",
+                        required=False,
+                        default=None,
+                        help="Temporary directory for BLAST database preparation when using FASTA directories as input. If not provided, a temporary directory will be created in the output directory.")
+
     parser.add_argument("-t", "--threads",
                         required=False,
                         default="1",
@@ -286,4 +293,8 @@ def arguments_parse():
     if args.output_tmp == "":
         args.output_tmp = args.output + "/.tmp/"
 
+    # Set default contig_table if not provided
+    if args.contig_table is None:
+        args.contig_table = args.output + "/contigs.tsv"
+    
     return args

src/PROBESt/prepare_blast.py

@@ -0,0 +1,202 @@
+"""
+Module for preparing BLAST databases from FASTA directories.
+
+This module provides functions to detect whether input paths are directories
+containing FASTA files (instead of pre-built BLAST databases) and to convert
+them to BLAST databases using the prep_db.sh script.
+"""
+
+import os
+import subprocess
+import glob
+from typing import List, Tuple, Optional
+
+
+# Supported FASTA file extensions
+FASTA_EXTENSIONS = {'.fa', '.fasta', '.fna', '.fa.gz', '.fasta.gz', '.fna.gz'}
+
+
+def is_fasta_directory(path: str) -> bool:
+    """
+    Check if the given path is a directory containing FASTA files.
+    
+    Args:
+        path: Path to check.
+        
+    Returns:
+        True if path is a directory containing FASTA files, False otherwise.
+    """
+    if not os.path.isdir(path):
+        return False
+    
+    # Check for FASTA files in the directory
+    for ext in FASTA_EXTENSIONS:
+        pattern = os.path.join(path, f"*{ext}")
+        if glob.glob(pattern):
+            return True
+    
+    return False
+
+
+def get_fasta_files(directory: str) -> List[str]:
+    """
+    Get all FASTA files from a directory.
+    
+    Args:
+        directory: Path to the directory containing FASTA files.
+        
+    Returns:
+        List of paths to FASTA files.
+    """
+    fasta_files = []
+    for ext in FASTA_EXTENSIONS:
+        pattern = os.path.join(directory, f"*{ext}")
+        fasta_files.extend(glob.glob(pattern))
+    return sorted(fasta_files)
+
+
+def prepare_blast_database(
+    fasta_dir: str,
+    output_db_path: str,
+    contig_table_path: str,
+    tmp_dir: Optional[str] = None,
+    script_path: Optional[str] = None
+) -> str:
+    """
+    Prepare a BLAST database from a directory containing FASTA files.
+    
+    Uses the prep_db.sh script to merge FASTA files and create a BLAST database.
+    
+    Args:
+        fasta_dir: Path to directory containing FASTA files.
+        output_db_path: Path for the output BLAST database.
+        contig_table_path: Path for the contig names output file.
+        tmp_dir: Optional temporary directory for intermediate files.
+        script_path: Path to the directory containing prep_db.sh script.
+        
+    Returns:
+        Path to the created BLAST database.
+        
+    Raises:
+        ValueError: If the directory contains no FASTA files.
+        RuntimeError: If the prep_db.sh script fails.
+    """
+    fasta_files = get_fasta_files(fasta_dir)
+    
+    if not fasta_files:
+        raise ValueError(f"No FASTA files found in directory: {fasta_dir}")
+    
+    # Determine script path
+    if script_path is None:
+        script_path = os.path.join(
+            os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
+            "scripts", "generator"
+        )
+    
+    prep_db_script = os.path.join(script_path, "prep_db.sh")
+    
+    if not os.path.exists(prep_db_script):
+        raise FileNotFoundError(f"prep_db.sh script not found at: {prep_db_script}")
+    
+    # Ensure output directories exist
+    os.makedirs(os.path.dirname(output_db_path) or '.', exist_ok=True)
+    os.makedirs(os.path.dirname(contig_table_path) or '.', exist_ok=True)
+    
+    # Build command
+    cmd = [
+        "bash", prep_db_script,
+        "-n", output_db_path,
+        "-c", contig_table_path,
+    ]
+    
+    if tmp_dir:
+        cmd.extend(["-t", tmp_dir])
+    
+    cmd.extend(fasta_files)
+    
+    # Run prep_db.sh
+    result = subprocess.run(
+        cmd,
+        capture_output=True,
+        text=True
+    )
+    
+    if result.returncode != 0:
+        raise RuntimeError(
+            f"prep_db.sh failed with return code {result.returncode}.\n"
+            f"stdout: {result.stdout}\n"
+            f"stderr: {result.stderr}"
+        )
+    
+    return output_db_path
+
+
+def prepare_bases_if_needed(
+    true_base: str,
+    false_bases: List[str],
+    output_dir: str,
+    contig_table_path: str,
+    tmp_dir: Optional[str] = None,
+    script_path: Optional[str] = None
+) -> Tuple[str, List[str]]:
+    """
+    Prepare BLAST databases from FASTA directories if needed.
+    
+    Checks if true_base and false_bases are directories containing FASTA files.
+    If they are, creates BLAST databases from them. Otherwise, returns the
+    original paths (assumed to be pre-built BLAST databases).
+    
+    Args:
+        true_base: Path to true base (BLAST database or FASTA directory).
+        false_bases: List of paths to false bases (BLAST databases or FASTA directories).
+        output_dir: Output directory for created BLAST databases.
+        contig_table_path: Path for the contig names output file.
+        tmp_dir: Optional temporary directory for intermediate files.
+        script_path: Path to the directory containing prep_db.sh script.
+        
+    Returns:
+        Tuple of (processed_true_base, processed_false_bases) paths to BLAST databases.
+    """
+    blast_db_dir = os.path.join(output_dir, ".blast_db")
+    os.makedirs(blast_db_dir, exist_ok=True)
+    
+    # Process true base
+    if is_fasta_directory(true_base):
+        dir_name = os.path.basename(true_base.rstrip('/'))
+        processed_true_base = os.path.join(blast_db_dir, f"true_{dir_name}")
+        true_contig_table = contig_table_path if os.path.dirname(contig_table_path) else os.path.join(blast_db_dir, contig_table_path)
+        
+        print(f"Preparing BLAST database from FASTA directory: {true_base}")
+        prepare_blast_database(
+            fasta_dir=true_base,
+            output_db_path=processed_true_base,
+            contig_table_path=true_contig_table,
+            tmp_dir=tmp_dir,
+            script_path=script_path
+        )
+        print(f"Created BLAST database: {processed_true_base}")
+    else:
+        processed_true_base = true_base
+    
+    # Process false bases
+    processed_false_bases = []
+    for i, false_base in enumerate(false_bases):
+        if is_fasta_directory(false_base):
+            dir_name = os.path.basename(false_base.rstrip('/'))
+            processed_false_base = os.path.join(blast_db_dir, f"false_{i}_{dir_name}")
+            false_contig_table = os.path.join(blast_db_dir, f"contigs_false_{i}.tsv")
+            
+            print(f"Preparing BLAST database from FASTA directory: {false_base}")
+            prepare_blast_database(
+                fasta_dir=false_base,
+                output_db_path=processed_false_base,
+                contig_table_path=false_contig_table,
+                tmp_dir=tmp_dir,
+                script_path=script_path
+            )
+            print(f"Created BLAST database: {processed_false_base}")
+            processed_false_bases.append(processed_false_base)
+        else:
+            processed_false_bases.append(false_base)
+    
+    return processed_true_base, processed_false_bases