add multi-fasta inintial probe set generation support

Author: dsmutin

README.md

@@ -25,24 +25,6 @@ bash setup/install.sh
 
 ## Usage
 
-### Preparation
-
-`pipeline.py` relies on pre-prepared BLASTn databases. To create the required `true_base`, `false_base`, and `contig_table`, you can use the following script:
-
-```bash
-bash scripts/generator/prep_db.sh \
-  -n {DATABASE_NAME} \
-  -c {CONTIG_NAME} \
-  -t {TMP_DIR} \
-  [FASTA]
-```
-
-#### Arguments:
-- `-n DATABASE_NAME`:  Name of the output BLAST database (required).  
-- `-c CONTIG_TABLE`:  Output file to store contig names and their corresponding sequence headers (required).  
-- `-t TMP_DIR`:  Temporary directory for intermediate files (optional, defaults to `./.tmp`).  
-- `FASTA`:  List of input FASTA files (gzipped or uncompressed). 
-
 ### Generation
 
 PROBEst can be run using the following command:
@@ -59,7 +41,7 @@ python pipeline.py \
 **Blastn databases** and **contig table** are results of the ```prep_db.sh```
 
 #### Key arguments:
-- `-i INPUT`: Input FASTA file for probe generation.
+- `-i INPUT`: Input FASTA file (or directory with fasta / fasta.gz file) for the initial probe setgeneration.
 - `-tb TRUE_BASE`: Input BLASTn database path for primer adjusting.
 - `-fb FALSE_BASE`: Input BLASTn database path for non-specific testing.
 - `-c CONTIG_TABLE`: .tsv table with BLAST database information.
@@ -81,6 +63,25 @@ python test_parameters.py \
   -p {JSON}
 ```
 
+
+### Manual preparation
+
+`pipeline.py` relies on pre-prepared BLASTn databases. To create the required `true_base`, `false_base`, and `contig_table`, you can use the following script:
+
+```bash
+bash scripts/generator/prep_db.sh \
+  -n {DATABASE_NAME} \
+  -c {CONTIG_NAME} \
+  -t {TMP_DIR} \
+  [FASTA]
+```
+
+#### Arguments:
+- `-n DATABASE_NAME`:  Name of the output BLAST database (required).  
+- `-c CONTIG_TABLE`:  Output file to store contig names and their corresponding sequence headers (required).  
+- `-t TMP_DIR`:  Temporary directory for intermediate files (optional, defaults to `./.tmp`).  
+- `FASTA`:  List of input FASTA files (gzipped or uncompressed). 
+
 ### Web Application
 
 PROBEst includes a user-friendly web interface for probe generation. The web app provides:

pipeline.py

@@ -25,6 +25,7 @@
 import numpy as np
 import pandas as pd
 import re
+import gzip
 
 # Add src directory to path to ensure local modules can be imported
 script_dir = os.path.dirname(os.path.abspath(__file__))
@@ -35,7 +36,7 @@
 # 0 Imports: absolute import ----
 from PROBESt.evolution import mutate_position
 from PROBESt.bash_wrappers import uniline_fasta, blastn_function, probe_check_function
-from PROBESt.misc import write_stats
+from PROBESt.misc import write_stats, collect_fasta_files, process_multiple_inputs
 from PROBESt.merge import merge
 from PROBESt.args import arguments_parse
 from PROBESt.modeling import run_modeling
@@ -82,25 +83,18 @@ def merge_iter(iter: int):
     script_path=args.script_path
 )
 
+# Collect all FASTA files from inputs
+input_fasta_files = collect_fasta_files(args.input)
+print(f"Found {len(input_fasta_files)} FASTA file(s) to process")
+
 # Create TMP
 os.makedirs(out_dir(0), exist_ok=True)
 
-# Make uniline fasta
-uniline_fasta(args, out_dir(0))
-print("Input fasta parsed")
-
-# Template generation
-# Check if initial_generator argument exists (for backward compatibility)
+# Process each input FASTA separately for initial set generation
 initial_generator = getattr(args, 'initial_generator', 'primer3')
+process_multiple_inputs(args, input_fasta_files, out_dir(0), initial_generator)
 
-if initial_generator == "primer3":
-    from PROBESt.primer3 import initial_set_generation
-    initial_set_generation(args, out_dir(0))
-elif initial_generator == "oligominer":
-    from PROBESt.oligominer import initial_set_generation
-    initial_set_generation(args, out_dir(0))
-else:
-    raise ValueError(f"Unknown initial generator: {initial_generator}")
+# Merge the final combined output
 merge_iter(0)
 
 # Check if initial set generation produced any probes
@@ -376,5 +370,18 @@ def merge_iter(iter: int):
     final_output_fa = args.output + "/output.fa"
 
 # 8. Modeling and visualization ----
+# Create combined input FASTA for modeling (from all input files)
+combined_input_fa = os.path.join(args.output, ".tmp", "combined_input.fa")
+os.makedirs(os.path.dirname(combined_input_fa), exist_ok=True)
+with open(combined_input_fa, 'w') as outfile:
+    for fasta_file in input_fasta_files:
+        # Handle gzipped files
+        if fasta_file.endswith('.gz'):
+            with gzip.open(fasta_file, 'rt') as infile:
+                outfile.write(infile.read())
+        else:
+            with open(fasta_file, 'r') as infile:
+                outfile.write(infile.read())
+
 modeling_output = os.path.join(args.output, "modeling_results.tsv")
-run_modeling(args, args.input, final_output_fa, modeling_output)
+run_modeling(args, combined_input_fa, final_output_fa, modeling_output)

setup/test_generator.sh

@@ -119,10 +119,24 @@ test2(){
         --oligominer_path "$OLIGOMINER_DIR"
 }
 
+# Test multifasta input
+test3 () {
+    gzip -d data/test/general/fasta_base/true_base/*.gz
+    python pipeline.py \
+        -i data/test/general/fasta_base/true_base \
+        -o data/test/general/output_multifasta \
+        -tb data/test/general/fasta_base/true_base \
+        -fb data/test/general/fasta_base/false_base_1 \
+        data/test/general/fasta_base/false_base_2 \
+        -a FISH \
+        --PRIMER_PICK_PRIMER 1 \
+        --PRIMER_NUM_RETURN 1
+}
 
 
 # Execute the preparation function
 clear_previous_test
 prepare_blast
 test1
-test2
+test2
+test3

src/PROBESt/args.py

@@ -32,7 +32,8 @@ def arguments_parse():
     # Main arguments
     parser.add_argument("-i", "--input",
                         required=True,
-                        help="Input FASTA file for probe generation. Probes are generated for different contigs separately. Only gene-coding regions are recommended (.fna).")
+                        nargs="*",
+                        help="Input FASTA file(s) or directory(ies) for probe generation. Can be a single file/directory or multiple files/directories. If a directory is provided, all *.fa, *.fna, *.fasta files (and their .gz versions) will be processed. Probes are generated for different contigs separately. Only gene-coding regions are recommended (.fna).")
 
     parser.add_argument("-tb", "--true_base",
                         required=True,

src/PROBESt/misc.py

@@ -1,6 +1,8 @@
 import os
 import re
 import subprocess
+import glob
+from typing import List
 
 
 def write_stats(stats: dict, output_dir: str):
@@ -83,3 +85,145 @@ def pairing(x):
 
     # Return a list with the base name appended by "RIGHT" and "LEFT"
     return [clear_primer + "RIGHT", clear_primer + "LEFT"]
+
+
+def collect_fasta_files(inputs: List[str]) -> List[str]:
+    """
+    Collect all FASTA files from input paths (files or directories).
+    
+    Args:
+        inputs: List of input paths. Each can be:
+            - A FASTA file (*.fa, *.fna, *.fasta or their .gz versions)
+            - A directory containing FASTA files
+    
+    Returns:
+        List of paths to all FASTA files found.
+        
+    Raises:
+        ValueError: If an input path doesn't exist or no FASTA files are found.
+    """
+    fasta_extensions = ['.fa', '.fna', '.fasta', '.fa.gz', '.fna.gz', '.fasta.gz']
+    fasta_files = []
+    
+    for input_path in inputs:
+        if not os.path.exists(input_path):
+            raise ValueError(f"Input path does not exist: {input_path}")
+        
+        if os.path.isfile(input_path):
+            # Check if it's a FASTA file
+            if any(input_path.endswith(ext) for ext in fasta_extensions):
+                fasta_files.append(input_path)
+            else:
+                raise ValueError(f"Input file is not a recognized FASTA file: {input_path}")
+        elif os.path.isdir(input_path):
+            # Collect all FASTA files from directory
+            dir_files = []
+            for ext in fasta_extensions:
+                pattern = os.path.join(input_path, f"*{ext}")
+                dir_files.extend(glob.glob(pattern))
+            
+            if not dir_files:
+                raise ValueError(f"No FASTA files found in directory: {input_path}")
+            
+            fasta_files.extend(sorted(dir_files))
+        else:
+            raise ValueError(f"Input path is neither a file nor a directory: {input_path}")
+    
+    if not fasta_files:
+        raise ValueError("No FASTA files found in any of the provided inputs")
+    
+    return sorted(fasta_files)
+
+
+def process_multiple_inputs(
+    args: 'Namespace',
+    input_fasta_files: List[str],
+    output_dir: str,
+    initial_generator: str = 'primer3'
+) -> str:
+    """
+    Process multiple input FASTA files separately for initial set generation and merge their outputs.
+    
+    This function processes each input FASTA file independently:
+    1. Creates a temporary directory for each input
+    2. Runs uniline_fasta and initial_set_generation for each
+    3. Merges individual outputs
+    4. Combines all outputs into a single output.fa file
+    
+    Args:
+        args: Arguments object containing all pipeline parameters
+        input_fasta_files: List of paths to FASTA files to process
+        output_dir: Output directory path (e.g., .tmp/0/)
+        initial_generator: Initial generator to use ('primer3' or 'oligominer')
+    
+    Returns:
+        Path to the final merged output.fa file
+    
+    Raises:
+        ValueError: If initial_generator is unknown
+        ImportError: If required modules cannot be imported
+    """
+    # Import here to avoid circular dependencies
+    from PROBESt.bash_wrappers import uniline_fasta
+    from PROBESt.merge import merge
+    
+    # Create output directory if it doesn't exist
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Process each input FASTA separately for initial set generation
+    temp_output_dirs = []
+    
+    for idx, fasta_file in enumerate(input_fasta_files):
+        print(f"\nProcessing input {idx + 1}/{len(input_fasta_files)}: {fasta_file}")
+        
+        # Create temporary directory for this input
+        temp_dir = os.path.join(output_dir, f"input_{idx}/")
+        os.makedirs(temp_dir, exist_ok=True)
+        temp_output_dirs.append(temp_dir)
+        
+        # Create a temporary args object with this specific input file
+        temp_args = type('Args', (), {})()
+        for attr in dir(args):
+            if not attr.startswith('_'):
+                setattr(temp_args, attr, getattr(args, attr))
+        temp_args.input = fasta_file
+        
+        # Make uniline fasta for this input
+        uniline_fasta(temp_args, temp_dir)
+        print(f"Input fasta parsed: {fasta_file}")
+        
+        # Template generation for this input
+        if initial_generator == "primer3":
+            from PROBESt.primer3 import initial_set_generation
+            initial_set_generation(temp_args, temp_dir)
+        elif initial_generator == "oligominer":
+            from PROBESt.oligominer import initial_set_generation
+            initial_set_generation(temp_args, temp_dir)
+        else:
+            raise ValueError(f"Unknown initial generator: {initial_generator}")
+        
+        # Merge this input's output
+        merge(algo=temp_args.algorithm,
+              input=os.path.join(temp_dir, "output.fa"),
+              output=os.path.join(temp_dir, "merged.fa"),
+              tmp=os.path.join(temp_dir, "fasta_table.tsv"),
+              NNN=10,
+              script_path=temp_args.script_path)
+    
+    # Merge all outputs to the first directory's output.fa
+    print("\nMerging all initial set outputs...")
+    final_output_fa = os.path.join(output_dir, "output.fa")
+    with open(final_output_fa, 'w') as outfile:
+        for temp_dir in temp_output_dirs:
+            merged_fa = os.path.join(temp_dir, "merged.fa")
+            if os.path.exists(merged_fa) and os.path.getsize(merged_fa) > 0:
+                with open(merged_fa, 'r') as infile:
+                    outfile.write(infile.read())
+            else:
+                # If merged.fa doesn't exist, try output.fa
+                output_fa = os.path.join(temp_dir, "output.fa")
+                if os.path.exists(output_fa) and os.path.getsize(output_fa) > 0:
+                    with open(output_fa, 'r') as infile:
+                        outfile.write(infile.read())
+    
+    return final_output_fa