XspecT 0.5.1__py3-none-any.whl → 0.5.3__py3-none-any.whl

xspect/classify.py

@@ -4,64 +4,77 @@ import xspect.model_management as mm
 from xspect.models.probabilistic_filter_mlst_model import (
     ProbabilisticFilterMlstSchemeModel,
 )
-from xspect.definitions import fasta_endings, fastq_endings
+from xspect.file_io import prepare_input_output_paths
 
 
 def classify_genus(
     model_genus: str, input_path: Path, output_path: Path, step: int = 1
 ):
-    """Classify the input file using the genus model."""
-    model = mm.get_genus_model(model_genus)
+    """
+    Classify the genus of sequences.
 
-    input_paths = []
-    input_is_dir = input_path.is_dir()
-    ending_wildcards = [f"*.{ending}" for ending in fasta_endings + fastq_endings]
+    This function classifies input files using the genus model.
+    The input path can be a file or directory
 
-    if input_is_dir:
-        input_paths = [p for e in ending_wildcards for p in input_path.glob(e)]
-    elif input_path.is_file():
-        input_paths = [input_path]
+    Args:
+        model_genus (str): The genus model slug.
+        input_path (Path): The path to the input file/directory containing sequences.
+        output_path (Path): The path to the output file where results will be saved.
+        step (int): The amount of kmers to be skipped.
+    """
+    model = mm.get_genus_model(model_genus)
+    input_paths, get_output_path = prepare_input_output_paths(input_path)
 
     for idx, current_path in enumerate(input_paths):
         result = model.predict(current_path, step=step)
         result.input_source = current_path.name
-        output_name = (
-            f"{output_path.stem}_{idx+1}{output_path.suffix}"
-            if input_is_dir
-            else output_path.name
-        )
-        result.save(output_path.parent / output_name)
-        print(f"Saved result as {output_name}")
+        cls_path = get_output_path(idx, output_path)
+        result.save(cls_path)
+        print(f"Saved result as {cls_path.name}")
 
 
-def classify_species(model_genus, input_path, output_path, step=1):
-    """Classify the input file using the species model."""
-    model = mm.get_species_model(model_genus)
+def classify_species(
+    model_genus: str, input_path: Path, output_path: Path, step: int = 1
+):
+    """
+    Classify the species of sequences.
 
-    input_paths = []
-    input_is_dir = input_path.is_dir()
-    ending_wildcards = [f"*.{ending}" for ending in fasta_endings + fastq_endings]
+    This function classifies input files using the species model.
+    The input path can be a file or directory
 
-    if input_is_dir:
-        input_paths = [p for e in ending_wildcards for p in input_path.glob(e)]
-    elif input_path.is_file():
-        input_paths = [input_path]
+    Args:
+        model_genus (str): The genus model slug.
+        input_path (Path): The path to the input file/directory containing sequences.
+        output_path (Path): The path to the output file where results will be saved.
+        step (int): The amount of kmers to be skipped.
+    """
+    model = mm.get_species_model(model_genus)
+    input_paths, get_output_path = prepare_input_output_paths(input_path)
 
     for idx, current_path in enumerate(input_paths):
         result = model.predict(current_path, step=step)
         result.input_source = current_path.name
-        output_name = (
-            f"{output_path.stem}_{idx+1}{output_path.suffix}"
-            if input_is_dir
-            else output_path.name
-        )
-        result.save(output_path.parent / output_name)
-        print(f"Saved result as {output_name}")
+        cls_path = get_output_path(idx, output_path)
+        result.save(cls_path)
+        print(f"Saved result as {cls_path.name}")
+
+
+def classify_mlst(input_path: Path, output_path: Path, limit: bool):
+    """
+    Classify the strain type using the specific MLST model.
 
+    Args:
+        input_path (Path): The path to the input file/directory containing sequences.
+        output_path (Path): The path to the output file where results will be saved.
+        limit (bool): A limit for the highest allele_id results that are shown.
+    """
 
-def classify_mlst(input_path, output_path):
-    """Classify the input file using the MLST model."""
     scheme_path = pick_scheme_from_models_dir()
     model = ProbabilisticFilterMlstSchemeModel.load(scheme_path)
-    result = model.predict(scheme_path, input_path)
-    result.save(output_path)
+    input_paths, get_output_path = prepare_input_output_paths(input_path)
+    for idx, current_path in enumerate(input_paths):
+        result = model.predict(scheme_path, current_path, step=1, limit=limit)
+        result.input_source = current_path.name
+        cls_path = get_output_path(idx, output_path)
+        result.save(cls_path)
+        print(f"Saved result as {cls_path.name}")

xspect/definitions.py

@@ -7,8 +7,16 @@ fasta_endings = ["fasta", "fna", "fa", "ffn", "frn"]
 fastq_endings = ["fastq", "fq"]
 
 
-def get_xspect_root_path():
-    """Return the root path for XspecT data."""
+def get_xspect_root_path() -> Path:
+    """
+    Return the root path for XspecT data.
+
+    Returns the path to the XspecT data directory, which can be located either in the user's home directory or in the current working directory.
+    If neither exists, it creates the directory in the user's home directory.
+
+    Returns:
+        Path: The path to the XspecT data directory.
+    """
 
     home_based_dir = Path.home() / "xspect-data"
     if home_based_dir.exists():
@@ -22,29 +30,61 @@ def get_xspect_root_path():
     return home_based_dir
 
 
-def get_xspect_model_path():
-    """Return the path to the XspecT models."""
+def get_xspect_model_path() -> Path:
+    """
+    Return the path to the XspecT models.
+
+    Returns the path to the XspecT models directory, which is located within the XspecT data directory.
+    If the directory does not exist, it creates the directory.
+
+    Returns:
+        Path: The path to the XspecT models directory.
+    """
     model_path = get_xspect_root_path() / "models"
     model_path.mkdir(exist_ok=True, parents=True)
     return model_path
 
 
-def get_xspect_upload_path():
-    """Return the path to the XspecT upload directory."""
+def get_xspect_upload_path() -> Path:
+    """
+    Return the path to the XspecT upload directory.
+
+    Returns the path to the XspecT uploads directory, which is located within the XspecT data directory.
+    If the directory does not exist, it creates the directory.
+
+    Returns:
+        Path: The path to the XspecT uploads directory.
+    """
     upload_path = get_xspect_root_path() / "uploads"
     upload_path.mkdir(exist_ok=True, parents=True)
     return upload_path
 
 
-def get_xspect_runs_path():
-    """Return the path to the XspecT runs directory."""
+def get_xspect_runs_path() -> Path:
+    """
+    Return the path to the XspecT runs directory.
+
+    Returns the path to the XspecT runs directory, which is located within the XspecT data directory.
+    If the directory does not exist, it creates the directory.
+
+    Returns:
+        Path: The path to the XspecT runs directory.
+    """
     runs_path = get_xspect_root_path() / "runs"
     runs_path.mkdir(exist_ok=True, parents=True)
     return runs_path
 
 
-def get_xspect_mlst_path():
-    """Return the path to the XspecT runs directory."""
+def get_xspect_mlst_path() -> Path:
+    """
+    Return the path to the XspecT MLST directory.
+
+    Returns the path to the XspecT MLST directory, which is located within the XspecT data directory.
+    If the directory does not exist, it creates the directory.
+
+    Returns:
+        Path: The path to the XspecT MLST directory.
+    """
     mlst_path = get_xspect_root_path() / "mlst"
     mlst_path.mkdir(exist_ok=True, parents=True)
     return mlst_path

xspect/download_models.py

@@ -8,8 +8,16 @@ import requests
 from xspect.definitions import get_xspect_model_path
 
 
-def download_test_models(url):
-    """Download models."""
+def download_test_models(url: str) -> None:
+    """
+    Download models from the specified URL.
+
+    This function downloads a zip file from the given URL, extracts its contents,
+    and copies the extracted files to the XspecT model directory.
+
+    Args:
+        url (str): The URL from which to download the models.
+    """
     with TemporaryDirectory() as tmp_dir:
         tmp_dir = Path(tmp_dir)
         download_path = tmp_dir / "models.zip"

xspect/file_io.py

@@ -6,12 +6,20 @@ from json import loads
 import os
 from pathlib import Path
 import zipfile
+from typing import Callable, Iterator
 from Bio import SeqIO
 from xspect.definitions import fasta_endings, fastq_endings
 
 
-def delete_zip_files(dir_path):
-    """Delete all zip files in the given directory."""
+def delete_zip_files(dir_path) -> None:
+    """
+    Delete all zip files in the given directory.
+
+    This function checks each file in the specified directory and removes it if it is a zip file.
+
+    Args:
+        dir_path (Path): Path to the directory where zip files should be deleted.
+    """
     files = os.listdir(dir_path)
     for file in files:
         if zipfile.is_zipfile(file):
@@ -19,45 +27,39 @@ def delete_zip_files(dir_path):
             os.remove(file_path)
 
 
-def extract_zip(zip_path: Path, unzipped_path: Path):
-    """Extracts all files from a zip file."""
+def extract_zip(zip_path: Path, unzipped_path: Path) -> None:
+    """
+    Extracts all files from a zip file.
+
+    Extracts the contents of the specified zip file to the given directory.
+
+    Args:
+        zip_path (Path): Path to the zip file to be extracted.
+        unzipped_path (Path): Path to the directory where the contents will be extracted.
+    """
     unzipped_path.mkdir(parents=True, exist_ok=True)
 
     with zipfile.ZipFile(zip_path) as item:
         item.extractall(unzipped_path)
 
 
-def concatenate_meta(path: Path, genus: str):
-    """Concatenates all species files to one fasta file.
-
-    :param path: Path to the directory with the concatenated fasta files.
-    :type path: Path
-    :param genus: Genus name.
-    :type genus: str
+def get_record_iterator(file_path: Path) -> Iterator:
     """
-    files_path = path / "concatenate"
-    meta_path = path / (genus + ".fasta")
-    files = os.listdir(files_path)
+    Returns a record iterator for a fasta or fastq file.
 
-    with open(meta_path, "w", encoding="utf-8") as meta_file:
-        # Write the header.
-        meta_header = f">{genus} metagenome\n"
-        meta_file.write(meta_header)
-
-        # Open each concatenated species file and write the sequence in the meta file.
-        for file in files:
-            file_ending = str(file).rsplit(".", maxsplit=1)[-1]
-            if file_ending in fasta_endings:
-                with open(
-                    (files_path / str(file)), "r", encoding="utf-8"
-                ) as species_file:
-                    for line in species_file:
-                        if line[0] != ">":
-                            meta_file.write(line.replace("\n", ""))
-
-
-def get_record_iterator(file_path: Path):
-    """Returns a record iterator for a fasta or fastq file."""
+    This function checks the file extension to determine if the file is in fasta or fastq format
+    and returns an iterator over the records in the file using Biopython's SeqIO module.
+
+    Args:
+        file_path (Path): Path to the fasta or fastq file.
+
+    Returns:
+        Iterator: An iterator over the records in the file.
+
+    Raises:
+        ValueError: If the file path is not a Path object, does not exist, is not a file,
+                    or has an invalid file format.
+    """
     if not isinstance(file_path, Path):
         raise ValueError("Path must be a Path object")
 
@@ -76,17 +78,18 @@ def get_record_iterator(file_path: Path):
     raise ValueError("Invalid file format, must be a fasta or fastq file")
 
 
-def get_records_by_id(file: Path, ids: list[str]):
-    """Return records with the specified ids."""
-    records = get_record_iterator(file)
-    return [record for record in records if record.id in ids]
-
+def concatenate_species_fasta_files(
+    input_folders: list[Path], output_directory: Path
+) -> None:
+    """
+    Concatenate fasta files from different species into one file per species.
 
-def concatenate_species_fasta_files(input_folders: list[Path], output_directory: Path):
-    """Concatenate fasta files from different species into one file per species.
+    This function iterates through each species folder within the given input folder,
+    collects all fasta files, and concatenates their contents into a single fasta file
+    named after the species.
 
     Args:
-        input_species_folders (list[Path]): List of paths to species folders.
+        input_folders (list[Path]): List of paths to species folders.
         output_directory (Path): Path to the output directory.
     """
     for species_folder in input_folders:
@@ -105,15 +108,22 @@ def concatenate_species_fasta_files(input_folders: list[Path], output_directory:
                     f.write(f_in.read())
 
 
-def concatenate_metagenome(fasta_dir: Path, meta_path: Path):
-    """Concatenate all fasta files in a directory into one file.
+def concatenate_metagenome(fasta_dir: Path, meta_path: Path) -> None:
+    """
+    Concatenate all fasta files in a directory into one file.
+
+    This function searches for all fasta files in the specified directory and writes their contents
+    into a single output file. The output file will contain the concatenated sequences from all fasta files.
 
     Args:
         fasta_dir (Path): Path to the directory with the fasta files.
         meta_path (Path): Path to the output file.
     """
+    fasta_files = [
+        file for ending in fasta_endings for file in fasta_dir.glob(f"*.{ending}")
+    ]
     with open(meta_path, "w", encoding="utf-8") as meta_file:
-        for fasta_file in fasta_dir.glob("*.fasta"):
+        for fasta_file in fasta_files:
             with open(fasta_file, "r", encoding="utf-8") as f_in:
                 meta_file.write(f_in.read())
 
@@ -121,13 +131,21 @@ def concatenate_metagenome(fasta_dir: Path, meta_path: Path):
 def get_ncbi_dataset_accession_paths(
     ncbi_dataset_path: Path,
 ) -> dict[str, Path]:
-    """Get the paths of the NCBI dataset accessions.
+    """
+    Get the paths of the NCBI dataset accessions.
+
+    This function reads the dataset catalog from the NCBI dataset directory and returns a dictionary
+    mapping each accession to its corresponding file path. The first item in the dataset catalog is
+    assumed to be a data report, and is skipped.
 
     Args:
         ncbi_dataset_path (Path): Path to the NCBI dataset directory.
 
     Returns:
         dict[str, Path]: Dictionary with the accession as key and the path as value.
+
+    Raises:
+        ValueError: If the dataset path does not exist or is invalid.
     """
     data_path = ncbi_dataset_path / "ncbi_dataset" / "data"
     if not data_path.exists():
@@ -147,13 +165,19 @@ def filter_sequences(
     input_file: Path,
     output_file: Path,
     included_ids: list[str],
-):
-    """Filter sequences by IDs from an input file and save them to an output file.
+) -> None:
+    """
+    Filter sequences by IDs from an input file and save them to an output file.
+
+    This function reads a fasta or fastq file, filters the sequences based on the provided IDs,
+    and writes the matching sequences to an output file. If no IDs are provided, no output file
+    is created.
 
     Args:
         input_file (Path): Path to the input file.
         output_file (Path): Path to the output file.
-        included_ids (list[str], optional): List of IDs to include. If None, no output file is created.
+        included_ids (list[str], optional): List of IDs to include. If None, no output file
+            is created.
     """
     if not included_ids:
         print("No IDs provided, no output file will be created.")
@@ -163,3 +187,46 @@ def filter_sequences(
         for record in get_record_iterator(input_file):
             if record.id in included_ids:
                 SeqIO.write(record, out_f, "fasta")
+
+
+def prepare_input_output_paths(
+    input_path: Path,
+) -> tuple[list[Path], Callable[[int, Path], Path]]:
+    """
+    Processes the input path into a list of input paths and a function generating output paths.
+
+    This function checks if the input path is a directory or a file. If it is a directory,
+    it collects all files with specified fasta and fastq endings. If it is a file, it uses that file
+    as the input path. It then returns a list of input file paths and a function that generates
+    output paths based on the index of the input file and a specified output path.
+
+    Args:
+        input_path (Path): Path to the directory or file.
+
+    Returns:
+        tuple[list[Path], Callable[[int, Path], Path]]: A tuple containing:
+            - A list of input file paths
+            - A function that takes an index and the output path,
+              and returns the processed output path.
+
+    Raises:
+        ValueError: If the input path is invalid.
+    """
+    input_is_dir = input_path.is_dir()
+    ending_wildcards = [f"*.{ending}" for ending in fasta_endings + fastq_endings]
+
+    if input_is_dir:
+        input_paths = [p for e in ending_wildcards for p in input_path.glob(e)]
+    elif input_path.is_file():
+        input_paths = [input_path]
+    else:
+        raise ValueError("Invalid input path")
+
+    def get_output_path(idx: int, output_path: Path) -> Path:
+        return (
+            output_path.parent / f"{output_path.stem}_{idx+1}{output_path.suffix}"
+            if input_is_dir
+            else output_path
+        )
+
+    return input_paths, get_output_path

xspect/filter_sequences.py

@@ -1,7 +1,6 @@
 from pathlib import Path
 from xspect.model_management import get_genus_model, get_species_model
-from xspect.file_io import filter_sequences
-from xspect.definitions import fasta_endings, fastq_endings
+from xspect.file_io import filter_sequences, prepare_input_output_paths
 
 
 def filter_species(
@@ -11,67 +10,51 @@ def filter_species(
     output_path: Path,
     threshold: float,
     classification_output_path: Path | None = None,
+    sparse_sampling_step: int = 1,
 ):
-    """Filter sequences by species.
+    """
+    Filter sequences by species.
+
     This function filters sequences from the input file based on the species model.
-    It uses the genus model to identify the genus of the sequences and then applies
-    the species model to filter the sequences.
+    It uses the species model to identify the species of individual sequences and then applies
+    a threshold filter the sequences.
 
     Args:
-        model_genus (str): The genus model slug.
-        model_species (str): The species model slug.
+        model_genus (str): The genus of the species model.
+        model_species (str): The species to filter by.
         input_path (Path): The path to the input file containing sequences.
         output_path (Path): The path to the output file where filtered sequences will be saved.
-            above this threshold will be included in the output file. A threshold of -1 will
-            include only sequences if the species score is the highest among the
-            available species scores.
         classification_output_path (Path): Optional path to save the classification results.
         threshold (float): The threshold for filtering sequences. Only sequences with a score
             above this threshold will be included in the output file. A threshold of -1 will
             include only sequences if the species score is the highest among the
             available species scores.
+        sparse_sampling_step (int): The step size for sparse sampling. Defaults to 1.
     """
     species_model = get_species_model(model_genus)
-
-    input_paths = []
-    input_is_dir = input_path.is_dir()
-    ending_wildcards = [f"*.{ending}" for ending in fasta_endings + fastq_endings]
-
-    if input_is_dir:
-        input_paths = [p for e in ending_wildcards for p in input_path.glob(e)]
-    elif input_path.is_file():
-        input_paths = [input_path]
+    input_paths, get_output_path = prepare_input_output_paths(input_path)
 
     for idx, current_path in enumerate(input_paths):
-        result = species_model.predict(current_path)
+        result = species_model.predict(current_path, step=sparse_sampling_step)
         result.input_source = current_path.name
 
         if classification_output_path:
-            classification_output_name = (
-                f"{classification_output_path.stem}_{idx+1}{classification_output_path.suffix}"
-                if input_is_dir
-                else classification_output_path.name
-            )
-            result.save(classification_output_path.parent / classification_output_name)
+            cls_out = get_output_path(idx, classification_output_path)
+            result.save(cls_out)
             print(
-                f"Saved classification results from {current_path.name} as {classification_output_name}"
+                f"Saved classification results from {current_path.name} as {cls_out.name}"
             )
 
         included_ids = result.get_filtered_subsequence_labels(model_species, threshold)
         if not included_ids:
             print(f"No sequences found for the given species in {current_path.name}.")
             continue
-        output_name = (
-            f"{output_path.stem}_{idx+1}{output_path.suffix}"
-            if input_is_dir
-            else output_path.name
-        )
-        filter_sequences(
-            current_path,
-            output_path.parent / output_name,
-            included_ids,
+
+        filter_output_path = get_output_path(idx, output_path)
+        filter_sequences(current_path, filter_output_path, included_ids)
+        print(
+            f"Saved filtered sequences from {current_path.name} as {filter_output_path.name}"
         )
-        print(f"Saved filtered sequences from {current_path.name} as {output_name}")
 
 
 def filter_genus(
@@ -80,8 +63,11 @@ def filter_genus(
     output_path: Path,
     threshold: float,
     classification_output_path: Path | None = None,
+    sparse_sampling_step: int = 1,
 ):
-    """Filter sequences by genus.
+    """
+    Filter sequences by genus.
+
     This function filters sequences from the input file based on the genus model.
     It uses the genus model to identify the genus of the sequences and then applies
     the filtering based on the provided threshold.
@@ -93,46 +79,30 @@ def filter_genus(
         threshold (float): The threshold for filtering sequences. Only sequences with a score
             above this threshold will be included in the output file.
         classification_output_path (Path): Optional path to save the classification results.
+        sparse_sampling_step (int): The step size for sparse sampling. Defaults to 1.
 
     """
-    genus_model = get_genus_model(model_genus)
-
-    input_paths = []
-    input_is_dir = input_path.is_dir()
-    ending_wildcards = [f"*.{ending}" for ending in fasta_endings + fastq_endings]
-
-    if input_is_dir:
-        input_paths = [p for e in ending_wildcards for p in input_path.glob(e)]
-    elif input_path.is_file():
-        input_paths = [input_path]
+    model = get_genus_model(model_genus)
+    input_paths, get_output_path = prepare_input_output_paths(input_path)
 
     for idx, current_path in enumerate(input_paths):
-        result = genus_model.predict(current_path)
+        result = model.predict(current_path, step=sparse_sampling_step)
         result.input_source = current_path.name
 
         if classification_output_path:
-            classification_output_name = (
-                f"{classification_output_path.stem}_{idx+1}{classification_output_path.suffix}"
-                if input_is_dir
-                else classification_output_path.name
-            )
-            result.save(classification_output_path.parent / classification_output_name)
+            cls_out = get_output_path(idx, classification_output_path)
+            result.save(cls_out)
             print(
-                f"Saved classification results from {current_path.name} as {classification_output_name}"
+                f"Saved classification results from {current_path.name} as {cls_out.name}"
             )
 
         included_ids = result.get_filtered_subsequence_labels(model_genus, threshold)
         if not included_ids:
             print(f"No sequences found for the given genus in {current_path.name}.")
             continue
-        output_name = (
-            f"{output_path.stem}_{idx+1}{output_path.suffix}"
-            if input_is_dir
-            else output_path.name
-        )
-        filter_sequences(
-            current_path,
-            output_path.parent / output_name,
-            included_ids,
+
+        filter_output_path = get_output_path(idx, output_path)
+        filter_sequences(current_path, filter_output_path, included_ids)
+        print(
+            f"Saved filtered sequences from {current_path.name} as {filter_output_path.name}"
         )
-        print(f"Saved filtered sequences from {current_path.name} as {output_name}")