XspecT 0.2.7__py3-none-any.whl → 0.4.1__py3-none-any.whl

xspect/main.py

@@ -1,28 +1,27 @@
 """Project CLI"""
 
 from pathlib import Path
-import datetime
-import uuid
+from uuid import uuid4
 import click
 import uvicorn
 from xspect import fastapi
 from xspect.download_models import download_test_models
-from xspect.train import train_ncbi
-from xspect.models.result import (
-    StepType,
-)
+from xspect.file_io import filter_sequences
+from xspect.train import train_from_directory, train_from_ncbi
 from xspect.definitions import (
-    get_xspect_runs_path,
-    fasta_endings,
-    fastq_endings,
     get_xspect_model_path,
 )
-from xspect.pipeline import ModelExecution, Pipeline, PipelineStep
 from xspect.mlst_feature.mlst_helper import pick_scheme, pick_scheme_from_models_dir
 from xspect.mlst_feature.pub_mlst_handler import PubMLSTHandler
 from xspect.models.probabilistic_filter_mlst_model import (
     ProbabilisticFilterMlstSchemeModel,
 )
+from xspect.model_management import (
+    get_genus_model,
+    get_model_metadata,
+    get_models,
+    get_species_model,
+)
 
 
 @click.group()
@@ -32,103 +31,131 @@ def cli():
 
 
 @cli.command()
-def download_models():
+def web():
+    """Open the XspecT web application."""
+    uvicorn.run(fastapi.app, host="0.0.0.0", port=8000)
+
+
+# # # # # # # # # # # # # # #
+# Model management commands #
+# # # # # # # # # # # # # # #
+@cli.group()
+def models():
+    """Model management commands."""
+
+
+@models.command(
+    help="Download models from the internet.",
+)
+def download():
     """Download models."""
     click.echo("Downloading models, this may take a while...")
-    download_test_models("https://xspect2.s3.eu-central-1.amazonaws.com/models.zip")
+    download_test_models("http://assets.adrianromberg.com/xspect-models.zip")
 
 
-@cli.command()
-@click.argument("genus")
-@click.argument("path", type=click.Path(exists=True, dir_okay=True, file_okay=True))
+@models.command(
+    name="list",
+    help="List all models in the model directory.",
+)
+def list_models():
+    """List models."""
+    available_models = get_models()
+    if not available_models:
+        click.echo("No models found.")
+        return
+    # todo: make this machine readable
+    click.echo("Models found:")
+    click.echo("--------------")
+    for model_type, names in available_models.items():
+        if not names:
+            continue
+        click.echo(f"  {model_type}:")
+        for name in names:
+            click.echo(f"    - {name}")
+
+
+@models.group()
+def train():
+    """Train models."""
+
+
+@train.command(
+    name="ncbi",
+    help="Train a species and a genus model based on NCBI data.",
+)
+@click.option("-g", "--genus", "model_genus", prompt=True)
+@click.option("--svm_steps", type=int, default=1)
 @click.option(
-    "-m",
-    "--meta/--no-meta",
-    help="Metagenome classification.",
-    default=False,
+    "--author",
+    help="Author of the model.",
+    default=None,
 )
 @click.option(
-    "-s",
-    "--step",
-    help="Sparse sampling step size (e. g. only every 500th kmer for step=500).",
-    default=1,
+    "--author-email",
+    help="Email of the author.",
+    default=None,
 )
-def classify_species(genus, path, meta, step):
-    """Classify sample(s) from file or directory PATH."""
-    click.echo("Classifying...")
-    click.echo(f"Step: {step}")
-
-    file_paths = []
-    if Path(path).is_dir():
-        file_paths = [
-            f
-            for f in Path(path).iterdir()
-            if f.is_file() and f.suffix[1:] in fasta_endings + fastq_endings
-        ]
-    else:
-        file_paths = [Path(path)]
-
-    # define pipeline
-    pipeline = Pipeline(genus + " classification", "Test Author", "test@example.com")
-    species_execution = ModelExecution(
-        genus.lower() + "-species", sparse_sampling_step=step
-    )
-    if meta:
-        species_filtering_step = PipelineStep(
-            StepType.FILTERING, genus, 0.7, species_execution
-        )
-        genus_execution = ModelExecution(
-            genus.lower() + "-genus", sparse_sampling_step=step
-        )
-        genus_execution.add_pipeline_step(species_filtering_step)
-        pipeline.add_pipeline_step(genus_execution)
-    else:
-        pipeline.add_pipeline_step(species_execution)
-
-    for idx, file_path in enumerate(file_paths):
-        run = pipeline.run(file_path)
-        time_str = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
-        save_path = get_xspect_runs_path() / f"run_{time_str}_{uuid.uuid4()}.json"
-        run.save(save_path)
-        print(
-            f"[{idx+1}/{len(file_paths)}] Run finished. Results saved to '{save_path}'."
-        )
+def train_ncbi(model_genus, svm_steps, author, author_email):
+    """Train a species and a genus model based on NCBI data."""
+    click.echo(f"Training {model_genus} species and genus metagenome model.")
+    try:
+        train_from_ncbi(model_genus, svm_steps, author, author_email)
+    except ValueError as e:
+        click.echo(f"Error: {e}")
+        return
+    click.echo(f"Training of {model_genus} model finished.")
 
 
-@cli.command()
-@click.argument("genus")
+@train.command(
+    name="directory",
+    help="Train a species (and possibly a genus) model based on local data.",
+)
+@click.option("-g", "--genus", "model_genus", prompt=True)
 @click.option(
-    "-bf-path",
-    "--bf-assembly-path",
-    help="Path to assembly directory for Bloom filter training.",
-    type=click.Path(exists=True, dir_okay=True, file_okay=False),
+    "-i",
+    "--input-path",
+    type=click.Path(exists=True, dir_okay=True, file_okay=True),
+    prompt=True,
 )
 @click.option(
-    "-svm-path",
-    "--svm-assembly-path",
-    help="Path to assembly directory for SVM training.",
-    type=click.Path(exists=True, dir_okay=True, file_okay=False),
+    "--meta",
+    is_flag=True,
+    help="Train a metagenome model for the genus.",
+    default=True,
 )
 @click.option(
-    "-s",
-    "--svm-step",
+    "--svm-steps",
+    type=int,
     help="SVM Sparse sampling step size (e. g. only every 500th kmer for step=500).",
     default=1,
 )
-def train_species(genus, bf_assembly_path, svm_assembly_path, svm_step):
-    """Train model."""
-
-    if bf_assembly_path or svm_assembly_path:
-        raise NotImplementedError(
-            "Training with specific assembly paths is not yet implemented."
-        )
-    try:
-        train_ncbi(genus, svm_step=svm_step)
-    except ValueError as e:
-        raise click.ClickException(str(e)) from e
+@click.option(
+    "--author",
+    help="Author of the model.",
+    default=None,
+)
+@click.option(
+    "--author-email",
+    help="Email of the author.",
+    default=None,
+)
+def train_directory(model_genus, input_path, svm_steps, meta, author, author_email):
+    """Train a model based on data from a directory for a given genus."""
+    click.echo(f"Training {model_genus} model with {svm_steps} SVM steps.")
+    train_from_directory(
+        model_genus,
+        Path(input_path),
+        svm_step=svm_steps,
+        meta=meta,
+        author=author,
+        author_email=author_email,
+    )
 
 
-@cli.command()
+@train.command(
+    name="mlst",
+    help="Train a MLST model based on PubMLST data.",
+)
 @click.option(
     "-c",
     "--choose_schemes",
@@ -154,27 +181,262 @@ def train_mlst(choose_schemes):
     click.echo(f"Saved at {model.cobs_path}")
 
 
-@cli.command()
+# # # # # # # # # # # # # # #
+# Classification commands   #
+# # # # # # # # # # # # # # #
+@cli.group(
+    name="classify",
+    help="Classify sequences using XspecT models.",
+)
+def classify_seqs():
+    """Classification commands."""
+
+
+@classify_seqs.command(
+    name="genus",
+    help="Classify samples using a genus model.",
+)
+@click.option(
+    "-g",
+    "--genus",
+    "model_genus",
+    help="Genus of the model to classify.",
+    type=click.Choice(get_models().get("Genus"), None),
+    prompt=True,
+)
+@click.option(
+    "-i",
+    "--input-path",
+    help="Path to FASTA or FASTQ file for classification.",
+    type=click.Path(exists=True, dir_okay=True, file_okay=True),
+    prompt=True,
+)
+@click.option(
+    "-o",
+    "--output-path",
+    help="Path to the output file.",
+    type=click.Path(dir_okay=True, file_okay=True),
+    default=Path(".") / f"result_{uuid4()}.json",
+)
+def classify_genus(model_genus, input_path, output_path):
+    """Classify samples using a genus model."""
+    click.echo("Classifying...")
+    genus_model = get_genus_model(model_genus)
+    result = genus_model.predict(Path(input_path))
+    result.save(output_path)
+    click.echo(f"Result saved as {output_path}.")
+
+
+@classify_seqs.command(
+    name="species",
+    help="Classify samples using a species model.",
+)
+@click.option(
+    "-g",
+    "--genus",
+    "model_genus",
+    help="Genus of the model to classify.",
+    type=click.Choice(get_models().get("Species"), None),
+    prompt=True,
+)
+@click.option(
+    "-i",
+    "--input-path",
+    help="Path to FASTA or FASTQ file for classification.",
+    type=click.Path(exists=True, dir_okay=True, file_okay=True),
+    prompt=True,
+)
+@click.option(
+    "-o",
+    "--output-path",
+    help="Path to the output file.",
+    type=click.Path(dir_okay=True, file_okay=True),
+    default=Path(".") / f"result_{uuid4()}.json",
+)
+@click.option(
+    "--sparse-sampling-step",
+    type=int,
+    help="Sparse sampling step (e. g. only every 500th kmer for '--sparse-sampling-step 500').",
+    default=1,
+)
+def classify_species(model_genus, input_path, output_path, sparse_sampling_step):
+    """Classify samples using a species model."""
+    click.echo("Classifying...")
+    species_model = get_species_model(model_genus)
+    result = species_model.predict(Path(input_path), step=sparse_sampling_step)
+    result.save(output_path)
+    click.echo(f"Result saved as {output_path}.")
+
+
+@classify_seqs.command(
+    name="mlst",
+    help="Classify samples using a MLST model.",
+)
 @click.option(
-    "-p",
-    "--path",
+    "-i",
+    "--input-path",
     help="Path to FASTA-file for mlst identification.",
     type=click.Path(exists=True, dir_okay=True, file_okay=True),
+    prompt=True,
+)
+@click.option(
+    "-o",
+    "--output-path",
+    help="Path to the output file.",
+    type=click.Path(dir_okay=True, file_okay=True),
+    default=Path(".") / f"result_{uuid4()}.json",
 )
-def classify_mlst(path):
+def classify_mlst(input_path, output_path):
     """MLST classify a sample."""
     click.echo("Classifying...")
-    path = Path(path)
+    input_path = Path(input_path)
     scheme_path = pick_scheme_from_models_dir()
     model = ProbabilisticFilterMlstSchemeModel.load(scheme_path)
-    model.predict(scheme_path, path).save(model.model_display_name, path)
-    click.echo(f"Run saved at {get_xspect_runs_path()}.")
+    result = model.predict(scheme_path, input_path)
+    result.save(output_path)
+    click.echo(f"Result saved as {output_path}.")
 
 
-@cli.command()
-def api():
-    """Open the XspecT FastAPI."""
-    uvicorn.run(fastapi.app, host="0.0.0.0", port=8000)
+# # # # # # # # # # # # # # #
+# Filtering commands        #
+# # # # # # # # # # # # # # #
+@cli.group(
+    name="filter",
+    help="Filter sequences using XspecT models.",
+)
+def filter_seqs():
+    """Filter commands."""
+
+
+@filter_seqs.command(
+    name="genus",
+    help="Filter sequences using a genus model.",
+)
+@click.option(
+    "-g",
+    "--genus",
+    "model_genus",
+    help="Genus of the model to use for filtering.",
+    type=click.Choice(get_models().get("Species"), None),
+    prompt=True,
+)
+@click.option(
+    "-i",
+    "--input-path",
+    help="Path to FASTA or FASTQ file for classification.",
+    type=click.Path(exists=True, dir_okay=True, file_okay=True),
+    prompt=True,
+)
+@click.option(
+    "-o",
+    "--output-path",
+    help="Path to the output file.",
+    type=click.Path(dir_okay=True, file_okay=True),
+    prompt=True,
+)
+@click.option(
+    "--threshold",
+    type=float,
+    help="Threshold for filtering (default: 0.7).",
+    default=0.7,
+    prompt=True,
+)
+def filter_genus(model_genus, input_path, output_path, threshold):
+    """Filter samples using a genus model."""
+    click.echo("Filtering...")
+    genus_model = get_genus_model(model_genus)
+    result = genus_model.predict(Path(input_path))
+    included_ids = result.get_filtered_subsequence_labels(model_genus, threshold)
+    if not included_ids:
+        click.echo("No sequences found for the given genus.")
+        return
+
+    filter_sequences(
+        Path(input_path),
+        Path(output_path),
+        included_ids=included_ids,
+    )
+    click.echo(f"Filtered sequences saved at {output_path}.")
+
+
+@filter_seqs.command(
+    name="species",
+    help="Filter sequences using a species model.",
+)
+@click.option(
+    "-g",
+    "--genus",
+    "model_genus",
+    help="Genus of the model to use for filtering.",
+    type=click.Choice(get_models().get("Species"), None),
+    prompt=True,
+)
+@click.option(
+    "-s",
+    "--species",
+    "model_species",
+    help="Species of the model to filter for.",
+)
+@click.option(
+    "-i",
+    "--input-path",
+    help="Path to FASTA or FASTQ file for classification.",
+    type=click.Path(exists=True, dir_okay=True, file_okay=True),
+    prompt=True,
+)
+@click.option(
+    "-o",
+    "--output-path",
+    help="Path to the output file.",
+    type=click.Path(dir_okay=True, file_okay=True),
+    prompt=True,
+)
+@click.option(
+    "--threshold",
+    type=float,
+    help="Threshold for filtering (default: 0.7). Use -1 to filter for the highest scoring species.",
+    default=0.7,
+    prompt=True,
+)
+def filter_species(model_genus, model_species, input_path, output_path, threshold):
+    """Filter a sample using the species model."""
+
+    available_species = get_model_metadata(f"{model_genus}-species")["display_names"]
+    available_species = {
+        id: name.replace(f"{model_genus} ", "")
+        for id, name in available_species.items()
+    }
+    if not model_species:
+        sorted_available_species = sorted(available_species.values())
+        model_species = click.prompt(
+            f"Please enter the species name: {model_genus}",
+            type=click.Choice(sorted_available_species, case_sensitive=False),
+        )
+    if model_species not in available_species.values():
+        raise click.BadParameter(
+            f"Species '{model_species}' not found in the {model_genus} species model."
+        )
+
+    # get the species ID from the name
+    model_species = [
+        id
+        for id, name in available_species.items()
+        if name.lower() == model_species.lower()
+    ][0]
+
+    click.echo("Filtering...")
+    species_model = get_species_model(model_genus)
+    result = species_model.predict(Path(input_path))
+    included_ids = result.get_filtered_subsequence_labels(model_species, threshold)
+    if not included_ids:
+        click.echo("No sequences found for the given species.")
+        return
+    filter_sequences(
+        Path(input_path),
+        Path(output_path),
+        included_ids=included_ids,
+    )
+    click.echo(f"Filtered sequences saved at {output_path}.")
 
 
 if __name__ == "__main__":

xspect/mlst_feature/mlst_helper.py

@@ -144,12 +144,10 @@ class MlstResult:
         }
         return result
 
-    def save(self, display: str, file_path: Path) -> None:
-        """Saves the result inside the "runs" directory"""
-        file_name = str(file_path).split("/")[-1]
-        json_path = get_xspect_runs_path() / "MLST" / f"{file_name}-{display}.json"
-        json_path.parent.mkdir(exist_ok=True, parents=True)
+    def save(self, output_path: Path) -> None:
+        """Saves the result as a JSON file."""
+        output_path.parent.mkdir(exist_ok=True, parents=True)
         json_object = json.dumps(self.to_dict(), indent=4)
 
-        with open(json_path, "w", encoding="utf-8") as file:
+        with open(output_path, "w", encoding="utf-8") as file:
             file.write(json_object)

xspect/model_management.py

@@ -2,7 +2,6 @@
 
 from json import loads, dumps
 from pathlib import Path
-from xspect.models.probabilistic_filter_model import ProbabilisticFilterModel
 from xspect.models.probabilistic_single_filter_model import (
     ProbabilisticSingleFilterModel,
 )
@@ -24,23 +23,10 @@ def get_species_model(genus):
     return species_filter_model
 
 
-def get_model_by_slug(model_slug: str):
-    """Get a model by its slug."""
-    model_path = get_xspect_model_path() / (model_slug + ".json")
-    model_metadata = get_model_metadata(model_path)
-    if model_metadata["model_class"] == "ProbabilisticSingleFilterModel":
-        return ProbabilisticSingleFilterModel.load(model_path)
-    if model_metadata["model_class"] == "ProbabilisticFilterSVMModel":
-        return ProbabilisticFilterSVMModel.load(model_path)
-    if model_metadata["model_class"] == "ProbabilisticFilterModel":
-        return ProbabilisticFilterModel.load(model_path)
-    raise ValueError(f"Model class {model_metadata['model_class']} not recognized.")
-
-
 def get_model_metadata(model: str | Path):
     """Get the metadata of a model."""
     if isinstance(model, str):
-        model_path = get_xspect_model_path() / (model + ".json")
+        model_path = get_xspect_model_path() / (model.lower() + ".json")
     elif isinstance(model, Path):
         model_path = model
     else:
@@ -85,3 +71,9 @@ def get_models():
             model_metadata["model_display_name"]
         )
     return model_dict
+
+
+def get_model_display_names(model_slug: str):
+    """Get the display names included in a model."""
+    model_metadata = get_model_metadata(model_slug)
+    return list(model_metadata["display_names"].values())

xspect/models/probabilistic_filter_model.py

@@ -26,6 +26,7 @@ class ProbabilisticFilterModel:
         base_path: Path,
         fpr: float = 0.01,
         num_hashes: int = 7,
+        training_accessions: dict[str, list[str]] = None,
     ) -> None:
         if k < 1:
             raise ValueError("Invalid k value, must be greater than 0")
@@ -46,6 +47,7 @@ class ProbabilisticFilterModel:
         self.fpr = fpr
         self.num_hashes = num_hashes
         self.index = None
+        self.training_accessions = training_accessions
 
     def get_cobs_index_path(self) -> Path:
         """Returns the path to the cobs index"""
@@ -63,13 +65,19 @@ class ProbabilisticFilterModel:
             "display_names": self.display_names,
             "fpr": self.fpr,
             "num_hashes": self.num_hashes,
+            "training_accessions": self.training_accessions,
         }
 
     def slug(self) -> str:
         """Returns a slug representation of the model"""
         return slugify(self.model_display_name + "-" + str(self.model_type))
 
-    def fit(self, dir_path: Path, display_names: dict = None) -> None:
+    def fit(
+        self,
+        dir_path: Path,
+        display_names: dict = None,
+        training_accessions: dict[str, list[str]] = None,
+    ) -> None:
         """Adds filters to the model"""
 
         if display_names is None:
@@ -84,16 +92,18 @@ class ProbabilisticFilterModel:
         if not dir_path.is_dir():
             raise ValueError("Directory path must be a directory")
 
+        self.training_accessions = training_accessions
+
         doclist = cobs.DocumentList()
         for file in dir_path.iterdir():
             if file.is_file() and file.suffix[1:] in fasta_endings + fastq_endings:
                 # cobs only uses the file name to the first "." as the document name
-                if file.name in display_names:
-                    self.display_names[file.name.split(".")[0]] = display_names[
-                        file.name
+                if file.stem in display_names:
+                    self.display_names[file.stem.split(".")[0]] = display_names[
+                        file.stem
                     ]
                 else:
-                    self.display_names[file.name.split(".")[0]] = file.stem
+                    self.display_names[file.stem.split(".")[0]] = file.stem
                 doclist.add(str(file))
 
         if len(doclist) == 0:
@@ -200,6 +210,7 @@ class ProbabilisticFilterModel:
                 path.parent,
                 model_json["fpr"],
                 model_json["num_hashes"],
+                model_json["training_accessions"],
             )
             model.display_names = model_json["display_names"]