dayhoff-tools 1.1.4__tar.gz → 1.1.6__tar.gz

{dayhoff_tools-1.1.4 → dayhoff_tools-1.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dayhoff-tools
-Version: 1.1.4
+Version: 1.1.6
 Summary: Common tools for all the repos at Dayhoff Labs
 Author: Daniel Martin-Alarcon
 Author-email: dma@dayhofflabs.com

dayhoff_tools-1.1.6/dayhoff_tools/deployment/processors.py

@@ -0,0 +1,429 @@
+import logging
+import os
+import shlex
+import shutil
+import subprocess
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+class Processor(ABC):
+    """Processes data locally.  Abstract class for specific calculations.
+    Takes in a single file and produces a single file or folder of outputs."""
+
+    @abstractmethod
+    def run(self, input_file: str) -> str:
+        """Do the calculation, including reading from input_file
+        and writing to output_file"""
+        output_path = "output_file"
+
+        return output_path
+
+
+class BoltzPredictor(Processor):
+    """Processor for running Boltz docking predictions.
+
+    This class wraps the Boltz docking tool to predict protein structures
+    from sequence data.
+    """
+
+    def __init__(self, num_workers: int, boltz_options: str | None = None):
+        """Initialize the BoltzPredictor.
+
+        Args:
+            num_workers: Number of worker threads to use as a default.
+                         This can be overridden if --num_workers is present
+                         in boltz_options.
+            boltz_options: A string containing additional command-line options
+                           to pass to the Boltz predictor. Options should be
+                           space-separated (e.g., "--option1 value1 --option2").
+        """
+        self.num_workers = num_workers
+        self.boltz_options = boltz_options
+
+    def run(self, input_file: str) -> str:
+        """Run Boltz prediction on the input file.
+
+        Constructs the command using the input file, default number of workers,
+        and any additional options provided via `boltz_options`. If `--num_workers`
+        is specified in `boltz_options`, it overrides the default `num_workers`.
+
+        Args:
+            input_file: Path to the input file containing sequences
+
+        Returns:
+            Path to the output directory created by Boltz
+
+        Raises:
+            subprocess.CalledProcessError: If Boltz prediction fails
+        """
+        # Determine expected output directory name
+        input_base = os.path.splitext(os.path.basename(input_file))[0]
+        expected_output_dir = f"boltz_results_{input_base}"
+        logger.info(f"Expected output directory: {expected_output_dir}")
+
+        # Start building the command
+        cmd = ["boltz", "predict", input_file]
+
+        # Parse additional options if provided
+        additional_args = []
+        num_workers_in_opts = False
+        if self.boltz_options:
+            try:
+                parsed_opts = shlex.split(self.boltz_options)
+                additional_args.extend(parsed_opts)
+                if "--num_workers" in parsed_opts:
+                    num_workers_in_opts = True
+                    logger.info(
+                        f"Using --num_workers from BOLTZ_OPTIONS: {self.boltz_options}"
+                    )
+            except ValueError as e:
+                logger.error(f"Error parsing BOLTZ_OPTIONS '{self.boltz_options}': {e}")
+                # Decide if we should raise an error or proceed without options
+                # For now, proceed without the additional options
+                additional_args = []  # Clear potentially partially parsed args
+
+        # Add num_workers if not specified in options
+        if not num_workers_in_opts:
+            logger.info(f"Using default num_workers: {self.num_workers}")
+            cmd.extend(["--num_workers", str(self.num_workers)])
+
+        # Add the parsed additional arguments
+        cmd.extend(additional_args)
+
+        # Log the final command
+        # Use shlex.join for safer command logging, especially if paths/args have spaces
+        try:
+            safe_cmd_str = shlex.join(cmd)
+            logger.info(f"Running command: {safe_cmd_str}")
+        except AttributeError:  # shlex.join is Python 3.8+
+            logger.info(f"Running command: {' '.join(cmd)}")
+
+        # Stream output in real-time
+        process = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            bufsize=1,
+        )
+
+        stdout = process.stdout
+        if stdout:
+            for line in iter(stdout.readline, ""):
+                logger.info(f"BOLTZ: {line.rstrip()}")
+
+        # Wait for process to complete
+        return_code = process.wait()
+        if return_code != 0:
+            logger.error(f"Boltz prediction failed with exit code {return_code}")
+            raise subprocess.CalledProcessError(return_code, cmd)
+
+        logger.info(
+            f"Boltz prediction completed successfully. Output in {expected_output_dir}"
+        )
+        return expected_output_dir
+
+
+class MMSeqsProfileProcessor(Processor):
+    """Processor for running MMseqs2 profile searches.
+
+    This class wraps the MMseqs2 workflow to perform a profile-based search
+    against a target database using a query FASTA.
+    """
+
+    def __init__(
+        self,
+        query_fasta_path_in_image: str,
+        num_threads: int = 8,
+        mmseqs_args: dict | None = None,
+    ):
+        """Initialize the MMSeqsProfileProcessor.
+
+        Args:
+            query_fasta_path_in_image: Path to the query FASTA file. This path is expected
+                                       to be accessible within the execution environment (e.g.,
+                                       packaged in a Docker image).
+            num_threads: Number of threads to use for MMseqs2 commands.
+            mmseqs_args: A dictionary of additional MMseqs2 parameters.
+                         Expected keys: "memory_limit_gb", "evalue", "sensitivity",
+                         "max_seqs_search", "min_seq_id_cluster", "max_seqs_profile_msa".
+                         Defaults are used if not provided.
+        """
+        if not Path(query_fasta_path_in_image).is_file():
+            raise FileNotFoundError(
+                f"Query FASTA file not found at: {query_fasta_path_in_image}"
+            )
+        self.query_fasta_path = query_fasta_path_in_image
+        self.num_threads = str(num_threads)  # MMseqs2 expects string for threads
+
+        default_mmseqs_args = {
+            "memory_limit_gb": "30",
+            "evalue": "10",
+            "sensitivity": "7.5",
+            "max_seqs_search": "300",
+            "min_seq_id_cluster": "0.8",
+            "max_seqs_profile_msa": "1000",
+        }
+        if mmseqs_args:
+            self.mmseqs_args = {**default_mmseqs_args, **mmseqs_args}
+        else:
+            self.mmseqs_args = default_mmseqs_args
+
+        logger.info(
+            f"MMSeqsProfileProcessor initialized with query: {self.query_fasta_path}"
+        )
+        logger.info(f"MMSeqs args: {self.mmseqs_args}")
+        logger.info(f"Num threads: {self.num_threads}")
+
+    def _run_mmseqs_command(
+        self, command_parts: list[str], step_description: str, work_dir: Path
+    ):
+        """Runs an MMseqs2 command and logs its execution.
+
+        Args:
+            command_parts: A list of strings representing the command and its arguments.
+            step_description: A human-readable description of the MMseqs2 step.
+            work_dir: The working directory for the command.
+
+        Raises:
+            subprocess.CalledProcessError: If the MMseqs2 command returns a non-zero exit code.
+        """
+        full_command = " ".join(command_parts)
+        logger.info(f"Running MMseqs2 step in {work_dir}: {step_description}")
+        logger.info(f"Command: {full_command}")
+        try:
+            process = subprocess.run(
+                command_parts,
+                check=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                cwd=work_dir,  # Run command in the specified working directory
+            )
+            if process.stdout:
+                logger.info(f"MMseqs2 stdout: {process.stdout.strip()}")
+            if process.stderr:  # MMseqs often outputs informational messages to stderr
+                logger.info(f"MMseqs2 stderr: {process.stderr.strip()}")
+            logger.info(f"MMseqs2 step '{step_description}' completed successfully.")
+        except subprocess.CalledProcessError as e:
+            logger.error(f"MMseqs2 step '{step_description}' failed in {work_dir}.")
+            if e.stdout:
+                logger.error(f"MMseqs2 stdout: {e.stdout.strip()}")
+            if e.stderr:
+                logger.error(f"MMseqs2 stderr: {e.stderr.strip()}")
+            raise
+
+    def run(self, input_file: str) -> str:
+        """Run MMseqs2 profile search.
+
+        The input_file is the target FASTA. The query FASTA is provided
+        during initialization.
+
+        Args:
+            input_file: Path to the input target FASTA file.
+
+        Returns:
+            Path to the output directory containing results.m8 and results.fasta.
+
+        Raises:
+            subprocess.CalledProcessError: If any MMseqs2 command fails.
+            FileNotFoundError: If the input_file is not found.
+        """
+        if not Path(input_file).is_file():
+            raise FileNotFoundError(f"Input target FASTA file not found: {input_file}")
+
+        input_file_path = Path(input_file).resolve()  # Ensure absolute path
+        target_fasta_filename = input_file_path.name
+
+        # Create a unique base directory for this run's outputs and temp files
+        # This directory will be returned and subsequently uploaded by the Operator
+        run_base_dir_name = f"mmseqs_run_{Path(target_fasta_filename).stem}"
+        run_base_dir = Path(run_base_dir_name).resolve()
+        run_base_dir.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Created run base directory: {run_base_dir}")
+
+        # Define local paths within the run_base_dir
+        local_target_file = run_base_dir / target_fasta_filename
+        # Copy the target file into the run directory to keep inputs and outputs together
+        shutil.copy(input_file_path, local_target_file)
+        logger.info(f"Copied target file {input_file_path} to {local_target_file}")
+
+        # Query file is already specified by self.query_fasta_path (path in image)
+        local_query_file = Path(self.query_fasta_path).resolve()
+
+        # Temporary directory for MMseqs2 intermediate files, created inside run_base_dir
+        mmseqs_temp_dir = run_base_dir / "mmseqs_tmp"
+        mmseqs_temp_dir.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Created MMseqs2 temporary directory: {mmseqs_temp_dir}")
+
+        # Define output file paths directly within run_base_dir
+        local_results_m8_file = run_base_dir / "results.m8"
+        local_hits_fasta_file = run_base_dir / "results.fasta"
+
+        # --- MMseqs2 Workflow Paths (intermediate files in mmseqs_temp_dir) ---
+        query_db = mmseqs_temp_dir / "queryDB"
+        target_db = mmseqs_temp_dir / "targetDB"
+        # Ensure local_target_file is used for creating targetDB
+        target_db_input_file = local_target_file
+
+        query_db_cluster = mmseqs_temp_dir / "queryDB_cluster"
+        query_db_rep = mmseqs_temp_dir / "queryDB_rep"
+        aln_db = mmseqs_temp_dir / "alnDB"
+        profile_db = mmseqs_temp_dir / "profileDB"
+        result_db = mmseqs_temp_dir / "resultDB"
+        hits_db = mmseqs_temp_dir / "hitsDB"
+
+        try:
+            # 1. Create query database
+            self._run_mmseqs_command(
+                ["mmseqs", "createdb", str(local_query_file), str(query_db)],
+                "Create query DB",
+                run_base_dir,  # Working directory for the command
+            )
+
+            # 2. Create target database
+            self._run_mmseqs_command(
+                ["mmseqs", "createdb", str(target_db_input_file), str(target_db)],
+                "Create target DB",
+                run_base_dir,
+            )
+
+            # 3. Cluster query sequences
+            self._run_mmseqs_command(
+                [
+                    "mmseqs",
+                    "cluster",
+                    str(query_db),
+                    str(query_db_cluster),
+                    str(
+                        mmseqs_temp_dir / "tmp_cluster"
+                    ),  # MMseqs needs a temp dir for cluster
+                    "--min-seq-id",
+                    self.mmseqs_args["min_seq_id_cluster"],
+                    "--threads",
+                    self.num_threads,
+                ],
+                "Cluster query sequences",
+                run_base_dir,
+            )
+
+            # 4. Create representative set from query clusters
+            self._run_mmseqs_command(
+                [
+                    "mmseqs",
+                    "createsubdb",
+                    str(query_db_cluster),
+                    str(query_db),
+                    str(query_db_rep),
+                ],
+                "Create representative query set",
+                run_base_dir,
+            )
+
+            # 5. Create MSA for profile generation
+            self._run_mmseqs_command(
+                [
+                    "mmseqs",
+                    "search",
+                    str(query_db_rep),
+                    str(query_db),  # Search representative against full query DB
+                    str(aln_db),
+                    str(mmseqs_temp_dir / "tmp_search_msa"),  # Temp for this search
+                    "--max-seqs",
+                    self.mmseqs_args["max_seqs_profile_msa"],
+                    "--threads",
+                    self.num_threads,
+                ],
+                "Create MSA for profile",
+                run_base_dir,
+            )
+
+            # 6. Create profile database
+            self._run_mmseqs_command(
+                [
+                    "mmseqs",
+                    "result2profile",
+                    str(query_db_rep),  # Use query_db_rep as input for profile
+                    str(query_db),  # Full query DB as second arg
+                    str(aln_db),
+                    str(profile_db),
+                    "--threads",  # Added threads option
+                    self.num_threads,
+                ],
+                "Create profile DB",
+                run_base_dir,
+            )
+
+            # 7. Perform profile search
+            self._run_mmseqs_command(
+                [
+                    "mmseqs",
+                    "search",
+                    str(profile_db),
+                    str(target_db),
+                    str(result_db),
+                    str(mmseqs_temp_dir / "tmp_search_profile"),  # Temp for this search
+                    "--split-memory-limit",
+                    f"{self.mmseqs_args['memory_limit_gb']}G",
+                    "-e",
+                    self.mmseqs_args["evalue"],
+                    "--max-seqs",
+                    self.mmseqs_args["max_seqs_search"],
+                    "--threads",
+                    self.num_threads,
+                    "-s",
+                    self.mmseqs_args["sensitivity"],
+                ],
+                "Perform profile search",
+                run_base_dir,
+            )
+
+            # 8. Convert results to tabular format (M8)
+            self._run_mmseqs_command(
+                [
+                    "mmseqs",
+                    "convertalis",
+                    str(profile_db),  # Query DB used for search (profileDB)
+                    str(target_db),
+                    str(result_db),
+                    str(local_results_m8_file),  # Output M8 file
+                    "--threads",
+                    self.num_threads,
+                ],
+                "Convert results to M8",
+                run_base_dir,
+            )
+
+            # 9. Create subdatabase of hits from original target_db
+            self._run_mmseqs_command(
+                ["mmseqs", "createsubdb", str(result_db), str(target_db), str(hits_db)],
+                "Create hits subDB from target_db",
+                run_base_dir,
+            )
+
+            # 10. Convert hit sequences to FASTA
+            self._run_mmseqs_command(
+                ["mmseqs", "convert2fasta", str(hits_db), str(local_hits_fasta_file)],
+                "Convert hits to FASTA",
+                run_base_dir,
+            )
+
+            logger.info(
+                f"MMseqs2 workflow completed successfully. Outputs in {run_base_dir}"
+            )
+
+        finally:
+            # Clean up the MMseqs2 temporary directory
+            if mmseqs_temp_dir.exists():
+                shutil.rmtree(mmseqs_temp_dir)
+                logger.info(
+                    f"Cleaned up MMseqs2 temporary directory: {mmseqs_temp_dir}"
+                )
+            # The input_file (original target) is managed by the Operator
+            # The local_target_file (copy inside run_base_dir) will be cleaned up
+            # by the Operator when run_base_dir is deleted after upload.
+
+        return str(run_base_dir)  # Return the path to the directory containing results

{dayhoff_tools-1.1.4 → dayhoff_tools-1.1.6}/dayhoff_tools/fasta.py

@@ -264,65 +264,168 @@ def split_fasta(
     target_folder: str,
     base_name: str,
     sequences_per_file: int = 1000,
-    max_files=None,
+    max_files: int | None = None,
+    show_progress: bool = True,
+    target_chunk_size_bytes: int | None = None,
 ) -> int:
-    """Split a FASTA file into multiple smaller files within a target folder.
+    """Split a FASTA file into multiple smaller files within a target folder,
+    with an overall progress bar. Files can be split based on a target number
+    of sequences or an approximate target file size in bytes.
 
     Args:
         fasta_file (str): Path to the input FASTA file.
         target_folder (str): Path to the folder where output files will be saved.
         base_name (str): Used to make output filenames: eg, basename_1.fasta.
         sequences_per_file (int): Number of sequences per output file.
-        max_files (int, optional): Maximum number of files to create. If None, all sequences are processed.
+            This is used if target_chunk_size_bytes is None.
+        max_files (int, optional): Maximum number of files to create.
+            If None, all sequences are processed.
+        show_progress (bool): If True, display a progress bar based on
+            file size processed. Defaults to True.
+        target_chunk_size_bytes (int, optional): Approximate target size for
+            each output file in bytes. If set, this takes precedence over
+            sequences_per_file. The actual file size may be slightly larger to
+            ensure full FASTA entries. Defaults to None.
+
+    Returns:
+        int: The number of output files created.
     """
     # Ensure the target folder exists
     os.makedirs(target_folder, exist_ok=True)
 
-    # Initialize counters
-    file_count = 1
-    sequence_count = 0
-
-    # Open the large FASTA file for reading
-    with open(fasta_file, "r", buffering=1024 * 1024) as fasta:
-        # Prepare the first output file
-        output_file_path = os.path.join(
-            target_folder, f"{base_name}_{file_count}.fasta"
+    # We create output files lazily (on first sequence) so we don't end up with
+    # spurious empty files.  `files_created` tracks the number of *real* files
+    # present on disk when we finish.
+    files_created = 0
+    current_output_file_sequence_count = 0
+    current_output_file_bytes_written = 0
+    pbar: tqdm | None = None
+    output_file = None  # Will be opened when we encounter the first header line
+    output_file_path = ""
+
+    if target_chunk_size_bytes is not None:
+        print(
+            f"Splitting by target chunk size: {target_chunk_size_bytes / (1024*1024):.2f} MB"
         )
-        output_file = open(output_file_path, "w", buffering=1024 * 1024)
-
-        for line in fasta:
-            # Check if we've reached the maximum number of files, if specified
-            if max_files is not None and file_count > max_files:
-                break
+    else:
+        print(f"Splitting by sequences per file: {sequences_per_file}")
 
-            # If line starts with ">", it's the beginning of a new sequence
-            if line.startswith(">"):
-                sequence_count += 1
+    try:
+        # Open the large FASTA file for reading
+        with open(fasta_file, "r", buffering=1024 * 1024) as fasta:
+            if show_progress:
+                total_size = os.path.getsize(fasta_file)
+                pbar = tqdm(
+                    total=total_size,
+                    unit="B",
+                    unit_scale=True,
+                    desc=f"Splitting {os.path.basename(fasta_file)}",
+                )
 
-                # If we reached the limit, start a new file
-                if sequence_count > sequences_per_file:
-                    # Close current file and open a new one
-                    output_file.close()
-                    print(f"File written: {output_file}")
-                    file_count += 1
-                    sequence_count = 1  # Reset sequence count for the new file
+            # We create output files on demand.  The very first file is not
+            # opened until we see the first sequence header.  This prevents
+            # an empty file from being created when the input FASTA is empty
+            # or when `max_files` is reached before any data are written.
+            def _open_new_output_file():
+                nonlocal output_file, output_file_path, files_created
 
-                    # Check again after incrementing file_count
-                    if max_files is not None and file_count > max_files:
-                        break
+                files_created += 1
+                output_file_path = os.path.join(
+                    target_folder, f"{base_name}_{files_created}.fasta"
+                )
+                output_file = open(output_file_path, "w", buffering=1024 * 1024)
 
-                    output_file_path = os.path.join(
-                        target_folder, f"{base_name}_{file_count}.fasta"
+            # Helper for logging and closing the current file
+            def _close_current_output_file():
+                nonlocal output_file, current_output_file_sequence_count, current_output_file_bytes_written
+                if output_file and not output_file.closed:
+                    output_file.close()
+                    print(
+                        f"File written: {output_file_path} "
+                        f"(Sequences: {current_output_file_sequence_count}, "
+                        f"Bytes: {current_output_file_bytes_written} / {(current_output_file_bytes_written / (1024*1024)):.2f} MB)"
                     )
-                    output_file = open(output_file_path, "w", buffering=1024 * 1024)
-
-            # Write the line to the current output file
-            output_file.write(line)
 
-        # Close the last output file
-        output_file.close()
-
-    return file_count
+            for line in fasta:
+                line_bytes = len(line.encode("utf-8"))
+                if pbar:
+                    pbar.update(line_bytes)
+
+                # Note: we don't enforce `max_files` here; we enforce it only when we
+                # are about to create *another* file (see logic further below). This
+                # ensures we finish writing the current file before stopping.
+
+                # If line starts with ">", it's the beginning of a new sequence
+                if line.startswith(">"):
+                    # Decide whether we need to roll over to a new output file.
+                    needs_new_file = False  # reset each time we encounter a header
+
+                    if (
+                        output_file is not None
+                        and current_output_file_sequence_count > 0
+                    ):
+                        if target_chunk_size_bytes is not None:
+                            # Size-based splitting takes precedence over sequence count.
+                            if (
+                                current_output_file_bytes_written
+                                >= target_chunk_size_bytes
+                            ):
+                                needs_new_file = True
+                        else:
+                            # Fallback to sequence-count based splitting.
+                            if current_output_file_sequence_count >= sequences_per_file:
+                                needs_new_file = True
+
+                    if needs_new_file:
+                        _close_current_output_file()
+
+                        # Respect `max_files`: do not create another file if limit reached
+                        if max_files is not None and files_created >= max_files:
+                            break
+
+                        _open_new_output_file()
+                        current_output_file_sequence_count = 0
+                        current_output_file_bytes_written = 0
+
+                    # Opening first file if not already open
+                    if output_file is None:
+                        _open_new_output_file()
+
+                    current_output_file_sequence_count += 1
+
+                # Write the line to the current output file (which should now exist)
+                if output_file is not None:
+                    output_file.write(line)
+                    current_output_file_bytes_written += line_bytes
+
+            # After loop, ensure the last file is handled
+            _close_current_output_file()
+
+    finally:
+        if pbar:
+            pbar.close()
+        # Ensure the file is closed in case of an exception before the natural end
+        if output_file and not output_file.closed:
+            output_file.close()
+            # It's hard to know the state to print a meaningful message here if an exception occurred mid-file.
+            # The primary 'File written' messages are handled within the loop and at the end of normal processing.
+
+    # If the last file was empty and removed, and it was the only file, file_count might be 1.
+    # Adjust file_count if the last output file was empty and removed.
+    if os.path.exists(output_file_path) and os.path.getsize(output_file_path) == 0:
+        # This can happen if max_files is hit exactly when a new file is due to be created,
+        # or if the input file itself is empty or contains no FASTA entries after the last split point.
+        # We should not count this empty file if it was removed.
+        # However, file_count is already incremented *before* a new file is opened.
+        # The logic for removing empty files is tricky to perfectly align with file_count
+        # without more complex state tracking. The current return reflects the number of
+        # *attempted* file creations that weren't immediately curtailed by max_files.
+        # For simplicity, we'll return the file_count as is, understanding it might
+        # include an empty file that was subsequently removed if it was the very last one.
+        # A more robust approach might decrement file_count if the last created file path was removed.
+        pass
+
+    return files_created
 
 
 def subtract_fasta_files(file1: str, file2: str, output_file: str):

dayhoff_tools-1.1.6/dayhoff_tools/intake/gtdb.py

@@ -0,0 +1,269 @@
+import collections.abc
+import csv
+import gzip
+import pathlib
+import re
+
+from tqdm import tqdm
+
+_ACCESSION_REGEX = re.compile(r"(GC[AF]_[0-9]+\.[0-9]+)")
+
+
+def _extract_accession_from_filename(filename_str: str) -> str:
+    """
+    Extracts the genome assembly accession (e.g., GCA_XXXXXXXXX.X or GCF_XXXXXXXXX.X)
+    from a filename.
+
+    Args:
+        filename_str (str): The filename string.
+
+    Returns:
+        str: The extracted accession or "UNKNOWN_ACCESSION" if not found.
+    """
+    match = _ACCESSION_REGEX.search(filename_str)
+    if match:
+        return match.group(1)
+    return "UNKNOWN_ACCESSION"
+
+
+def process_gtdb_files_to_fasta(
+    gtdb_top_folder: str,
+    output_fasta_path: str,
+    chunk_size: int = 10000,
+) -> None:
+    """
+    Processes a top-level GTDB folder containing gzipped FASTA files (.faa.gz)
+    and combines all protein sequences into a single FASTA file.
+
+    Output is written in chunks for efficiency with large datasets.
+    A progress bar is displayed during processing.
+
+    Args:
+        gtdb_top_folder (str): Path to the top-level GTDB directory.
+        output_fasta_path (str): Path to write the combined FASTA file.
+        chunk_size (int, optional): Number of sequences to process before
+            writing a chunk to the output file. Defaults to 10000.
+    """
+    gtdb_path = pathlib.Path(gtdb_top_folder)
+    faa_files = list(gtdb_path.rglob("*.faa.gz"))
+
+    if not faa_files:
+        print(f"No .faa.gz files found in {gtdb_top_folder}")
+        return
+
+    fasta_entries_chunk = []
+    sequences_in_current_chunk = 0
+
+    with open(output_fasta_path, "w") as fasta_out_file:
+        current_header_id = None
+        current_sequence_lines = []
+
+        for faa_file_path in tqdm(faa_files, desc="Processing GTDB files to FASTA"):
+            try:
+                with gzip.open(faa_file_path, "rt") as gz_file:
+                    for line_content in gz_file:
+                        line = line_content.strip()
+                        if not line:  # Skip empty lines
+                            continue
+                        if line.startswith(">"):
+                            if current_header_id and current_sequence_lines:
+                                sequence_string = "".join(current_sequence_lines)
+                                fasta_entries_chunk.append(
+                                    f">{current_header_id}\n{sequence_string}\n"
+                                )
+                                sequences_in_current_chunk += 1
+
+                            # Parse new header
+                            header_content = line[1:]
+                            parts = header_content.split(None, 1)
+                            current_header_id = parts[0]
+                            current_sequence_lines = []
+
+                            if sequences_in_current_chunk >= chunk_size:
+                                if fasta_entries_chunk:
+                                    fasta_out_file.write("".join(fasta_entries_chunk))
+                                fasta_entries_chunk = []
+                                sequences_in_current_chunk = 0
+                        else:
+                            if current_header_id:
+                                current_sequence_lines.append(line)
+
+                    if current_header_id and current_sequence_lines:
+                        sequence_string = "".join(current_sequence_lines)
+                        fasta_entries_chunk.append(
+                            f">{current_header_id}\n{sequence_string}\n"
+                        )
+                        sequences_in_current_chunk += 1
+
+                # Reset state for the next file to ensure clean parsing start for that file
+                current_header_id = None
+                current_sequence_lines = []
+
+            except gzip.BadGzipFile:
+                tqdm.write(
+                    f"Warning: Skipping corrupted or non-gzipped file: {faa_file_path}"
+                )
+                current_header_id = None
+                current_sequence_lines = []
+            except Exception as e:
+                tqdm.write(f"Warning: Error processing file {faa_file_path}: {e}")
+                current_header_id = None
+                current_sequence_lines = []
+
+        if fasta_entries_chunk:
+            fasta_out_file.write("".join(fasta_entries_chunk))
+
+    print(f"Processing complete. Output FASTA file created: {output_fasta_path}")
+
+
+def process_gtdb_files_to_csv(
+    gtdb_top_folder: str,
+    output_csv_path: str,
+    chunk_size: int = 10000,
+) -> None:
+    """
+    Processes a top-level GTDB folder containing gzipped FASTA files (.faa.gz)
+    and creates a CSV file with detailed information for each sequence entry.
+
+    The CSV includes the genome assembly accession, original FASTA header ID,
+    and header description for each entry. Output is written in chunks for
+    efficiency with large datasets. A progress bar is displayed during processing.
+
+    Args:
+        gtdb_top_folder (str): Path to the top-level GTDB directory.
+        output_csv_path (str): Path to write the CSV file.
+        chunk_size (int, optional): Number of sequences to process before
+            writing a chunk to the output file. Defaults to 10000.
+    """
+    gtdb_path = pathlib.Path(gtdb_top_folder)
+    faa_files = list(gtdb_path.rglob("*.faa.gz"))
+
+    if not faa_files:
+        print(f"No .faa.gz files found in {gtdb_top_folder}")
+        return
+
+    def _serial_iter(paths):
+        """Yield the same structure as the parallel branch but serially."""
+        for p in paths:
+            row_generator_for_file, file_warnings = _csv_rows_from_single_faa(str(p))
+            yield row_generator_for_file, file_warnings
+
+    # Open output CSV for streaming writes.
+    with open(output_csv_path, "w", newline="") as csv_out_file:
+        csv_writer = csv.writer(csv_out_file)
+        csv_writer.writerow(
+            [
+                "genome_assembly_accession",
+                "original_fasta_header_id",
+                "original_fasta_header_description",
+            ]
+        )
+
+        rows_buffer: list[list[str]] = []
+
+        # Choose the iterator depending on workers.
+        result_iter = _serial_iter(faa_files)
+        progress_iter = tqdm(
+            result_iter, total=len(faa_files), desc="Processing GTDB files to CSV"
+        )
+
+        # Consume iterator and stream rows to disk in chunks.
+        for row_generator_for_file, file_warnings in progress_iter:
+            # Add rows to buffer and flush in chunk-size batches.
+            # This will consume the generator, and in doing so, populate file_warnings if errors occur.
+            for r in row_generator_for_file:
+                rows_buffer.append(r)
+                if len(rows_buffer) >= chunk_size:
+                    csv_writer.writerows(rows_buffer)
+                    rows_buffer.clear()
+
+            # Now that the generator for the file has been processed (or attempted),
+            # emit any warnings that were collected for this specific file.
+            for w in file_warnings:
+                tqdm.write(w)
+
+        # Flush remaining rows.
+        if rows_buffer:
+            csv_writer.writerows(rows_buffer)
+
+    print(f"Processing complete. Output CSV file created: {output_csv_path}")
+
+
+# ---------------------------------------------------------------------------
+# Helper functions (private)
+# ---------------------------------------------------------------------------
+
+
+def _csv_rows_from_single_faa(
+    faa_file_path: str,
+) -> tuple[collections.abc.Iterable[list[str]], list[str]]:
+    """Parse a single gzipped FASTA (`.faa.gz`) file into CSV rows.
+
+    Parameters
+    ----------
+    faa_file_path
+        Path (as ``str``) to the ``.faa.gz`` file.
+
+    Returns
+    -------
+    tuple[collections.abc.Iterable[list[str]], list[str]]
+        * First element – an iterable (generator) of CSV rows ``[accession, header_id, description]``.
+        * Second element – list of warning strings produced while processing
+          the file.  The caller is responsible for emitting them.
+    """
+    warnings: list[str] = []  # Outer scope warnings list
+    faa_path = pathlib.Path(faa_file_path)
+    current_file_accession = _extract_accession_from_filename(faa_path.name)
+
+    def _generate_rows_iter_inner() -> (
+        collections.abc.Iterable[list[str]]
+    ):  # Renamed for clarity
+        # Local parsing state for the generator
+        current_header_id_gen = None
+        current_header_desc_gen = ""
+        has_sequence_lines_gen = False
+
+        try:
+            with gzip.open(faa_file_path, "rt") as gz_file:
+                for line_content in gz_file:
+                    line = line_content.strip()
+                    if not line:
+                        continue
+                    if line.startswith(">"):
+                        if current_header_id_gen and has_sequence_lines_gen:
+                            yield [
+                                current_file_accession,
+                                current_header_id_gen,
+                                current_header_desc_gen,
+                            ]
+
+                        header_content = line[1:]
+                        parts = header_content.split(None, 1)
+                        current_header_id_gen = parts[0]
+                        current_header_desc_gen = parts[1] if len(parts) > 1 else ""
+                        has_sequence_lines_gen = False
+                    else:
+                        if current_header_id_gen:
+                            has_sequence_lines_gen = True
+
+            # Add final entry if the file ended after sequence lines.
+            if current_header_id_gen and has_sequence_lines_gen:
+                yield [
+                    current_file_accession,
+                    current_header_id_gen,
+                    current_header_desc_gen,
+                ]
+        except gzip.BadGzipFile:
+            # Exception handled inside the generator.
+            # Append to the outer warnings list and terminate generator.
+            warnings.append(
+                f"Warning: Skipping corrupted or non-gzipped file: {faa_file_path}"
+            )
+            return  # Stop generation
+        except Exception as exc:
+            warnings.append(f"Warning: Error processing file {faa_file_path}: {exc}")
+            return  # Stop generation
+
+    # Directly return the generator instance and the warnings list.
+    # The warnings list will be populated by the generator if errors occur during its execution.
+    return _generate_rows_iter_inner(), warnings

{dayhoff_tools-1.1.4 → dayhoff_tools-1.1.6}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "dayhoff-tools"
-version = "1.1.4"
+version = "1.1.6"
 description = "Common tools for all the repos at Dayhoff Labs"
 authors = [
     {name = "Daniel Martin-Alarcon", email = "dma@dayhofflabs.com"}

dayhoff_tools-1.1.4/dayhoff_tools/deployment/processors.py

@@ -1,125 +0,0 @@
-import logging
-import os
-import subprocess
-import shlex
-from abc import ABC, abstractmethod
-
-logger = logging.getLogger(__name__)
-
-
-class Processor(ABC):
-    """Processes data locally.  Abstract class for specific calculations.
-    Takes in a single file and produces a single file or folder of outputs."""
-
-    @abstractmethod
-    def run(self, input_file: str) -> str:
-        """Do the calculation, including reading from input_file
-        and writing to output_file"""
-        output_path = "output_file"
-
-        return output_path
-
-
-class BoltzPredictor(Processor):
-    """Processor for running Boltz docking predictions.
-
-    This class wraps the Boltz docking tool to predict protein structures
-    from sequence data.
-    """
-
-    def __init__(self, num_workers: int, boltz_options: str | None = None):
-        """Initialize the BoltzPredictor.
-
-        Args:
-            num_workers: Number of worker threads to use as a default.
-                         This can be overridden if --num_workers is present
-                         in boltz_options.
-            boltz_options: A string containing additional command-line options
-                           to pass to the Boltz predictor. Options should be
-                           space-separated (e.g., "--option1 value1 --option2").
-        """
-        self.num_workers = num_workers
-        self.boltz_options = boltz_options
-
-    def run(self, input_file: str) -> str:
-        """Run Boltz prediction on the input file.
-
-        Constructs the command using the input file, default number of workers,
-        and any additional options provided via `boltz_options`. If `--num_workers`
-        is specified in `boltz_options`, it overrides the default `num_workers`.
-
-        Args:
-            input_file: Path to the input file containing sequences
-
-        Returns:
-            Path to the output directory created by Boltz
-
-        Raises:
-            subprocess.CalledProcessError: If Boltz prediction fails
-        """
-        # Determine expected output directory name
-        input_base = os.path.splitext(os.path.basename(input_file))[0]
-        expected_output_dir = f"boltz_results_{input_base}"
-        logger.info(f"Expected output directory: {expected_output_dir}")
-
-        # Start building the command
-        cmd = ["boltz", "predict", input_file]
-
-        # Parse additional options if provided
-        additional_args = []
-        num_workers_in_opts = False
-        if self.boltz_options:
-            try:
-                parsed_opts = shlex.split(self.boltz_options)
-                additional_args.extend(parsed_opts)
-                if "--num_workers" in parsed_opts:
-                    num_workers_in_opts = True
-                    logger.info(
-                        f"Using --num_workers from BOLTZ_OPTIONS: {self.boltz_options}"
-                    )
-            except ValueError as e:
-                logger.error(f"Error parsing BOLTZ_OPTIONS '{self.boltz_options}': {e}")
-                # Decide if we should raise an error or proceed without options
-                # For now, proceed without the additional options
-                additional_args = []  # Clear potentially partially parsed args
-
-        # Add num_workers if not specified in options
-        if not num_workers_in_opts:
-            logger.info(f"Using default num_workers: {self.num_workers}")
-            cmd.extend(["--num_workers", str(self.num_workers)])
-
-        # Add the parsed additional arguments
-        cmd.extend(additional_args)
-
-        # Log the final command
-        # Use shlex.join for safer command logging, especially if paths/args have spaces
-        try:
-            safe_cmd_str = shlex.join(cmd)
-            logger.info(f"Running command: {safe_cmd_str}")
-        except AttributeError:  # shlex.join is Python 3.8+
-            logger.info(f"Running command: {' '.join(cmd)}")
-
-        # Stream output in real-time
-        process = subprocess.Popen(
-            cmd,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            text=True,
-            bufsize=1,
-        )
-
-        stdout = process.stdout
-        if stdout:
-            for line in iter(stdout.readline, ""):
-                logger.info(f"BOLTZ: {line.rstrip()}")
-
-        # Wait for process to complete
-        return_code = process.wait()
-        if return_code != 0:
-            logger.error(f"Boltz prediction failed with exit code {return_code}")
-            raise subprocess.CalledProcessError(return_code, cmd)
-
-        logger.info(
-            f"Boltz prediction completed successfully. Output in {expected_output_dir}"
-        )
-        return expected_output_dir