strainfish 0.2.1

An Ensemble Machine Learning algorithm for classification of bacterial strains.

StrainFish

strainfish is a weighted ensemble machine learning algorithm with multiple DNA sequence encoders and logic, specifically designed for classification of marker sequences.

Conceived and built by Kranti Konganti, HFP

v0.2.0

• Multiple DNA sequence encoders for NVIDIA GPU-accelerated training.
• A weighted Ensemble machine-learning model generation with sensible defaults.
• NVIDIA GPU-accelerated Learning and Prediction only!
• Important Note: This software is under active development and as such some features are experimental. Results should be thoroughly validated and independently verified before use in critical applications or publications.

Table of Contents

1. Installation
2. Quick Start
3. Training Models
   • Basic Training Command
   • Advanced Configuration
   • Encoding Methods
4. Making Predictions
   • Basic Prediction Command
   • Model Management
5. Configuration Options
   • XGBoost Parameters
   • RandomForest Parameters
   • SentencePiece Parameters
   • Imbalance Handling Parameters
6. Test Data and Examples
7. Dependencies
8. License

Installation

StrainFish requires Python 3.11+ and can be installed via pip:

pip install strainfish

Quick Start

Training a Model

To train a model on your DNA sequences:

strainfish train run \
  -f path/to/sequences.fasta \
  -l path/to/labels.csv \
  -o /path/to/models_output_dir/model_prefix

Predicting using a Model

To perform prediction using a trained model:

strainfish predict run \
  -f path/to/predict_sequences.fasta \
  -m /path/to/models_output_dir/model_prefix \
  -o path/to/results_directory

Training Models

StrainFish uses an ensemble approach for both training and prediction (XGBoost, RandomForest and NaiveBayes), and includes multiple DNA sequence encodings, though only one encoding can be used at a time during training.

Basic Training Command

strainfish train run \
  -f training_sequences.fasta \                             # Input FASTA file
  -l labels.csv \                                           # Labels CSV (id,label)
  -o /path/to/models_output_dir/model_prefix                # Output directory for models

Advanced Configuration

Optional StrainFish configuration options during training:

strainfish train run \
  -f training_sequences.fasta \
  -l labels.csv \
  -o model_output_dir \
  --encode-method tf \              # Encoding method: sm, sp, or tf
  --kmer 7 \                        # K-mer size for hashing
  --num-hashes 100 \                # Number of hashes per sequence
  --factor 21 \                     # Sequence overlap factor
  --chunk-size 200 \                # Size of DNA chunks
  --pseknc-weight 0.1 \             # Weight for PseKNC encoding
  --xgb-n-estimators 300 \          # XGBoost parameters
  --rf-n-estimators 100 \           # RandomForest parameters

Encoding Methods

StrainFish supports three DNA sequence encoding methods:

• tf (TF-IDF): TF-IDF vectorization (Default)
• sp (SentencePiece): Subword tokenization using SentencePiece models (Experimental)
• sm (SOMH): MinHash based approach with PseKNC and sequencing composition weights (AT/GC ratio) (Experimental)

Making Predictions

Basic Prediction Command

strainfish predict run \
  -f prediction_sequences.fasta \                        # Input FASTA file(s)
  -m /path/to/models_output_dir/model_prefix \           # Path to trained model
  -o results_dir                                         # Output directory for predictions

Model Management

List available models:

strainfish predict list-models
# Or list models stored at a particular models directory:
strainfish predict list-models -md /path/to/models_dir

Configuration Options

Tunable parameters for StrainFish.

XGBoost Parameters

View all configurable XGBoost parameters:

strainfish train show-xgb-params

Key parameters:

• --xgb-n-estimators: Number of boosting rounds
• --xgb-max-depth: Maximum tree depth
• --xgb-learning-rate: Learning rate for boosting
• --xgb-subsample: Subsample ratio of the training instance

RandomForest Parameters

View all configurable RandomForest parameters:

strainfish train show-rf-params

Key parameters:

• --rf-n-estimators: Number of trees in the forest
• --rf-max-depth: Maximum depth of the tree
• --rf-random-state: Random seed for reproducibility
• --rf-min-samples-leaf: Minimum samples required at a leaf node

SentencePiece Parameters

View all configurable SentencePiece parameters:

strainfish train show-sp-params

Key parameters:

• --sp-vocab-size: Vocabulary size for tokenization
• --sp-max-sentence-length: Maximum sentence length
• --sp-char-cov: Character coverage ratio

Imbalance Handling Parameters

View all imbalance handling parameters:

strainfish train show-imb-params

Key parameters:

• --imb-smote-k-neighbors: Number of neighbors for SMOTE
• --imb-enn-n-neighbors: Number of neighbors for ENN cleaning

Test Data and Examples

This repository includes test data in the tests/test_input/ directory:

• test.train.fasta: Training sequences in FASTA format
• test.train.csv: Labels file with id,label columns
• predict.fasta: Sequences for prediction using trained models

You can use these to test StrainFish functionality:

# Train a model using test data
strainfish train run \
  -f tests/test_input/test.train.fasta \
  -l tests/test_input/test.train.csv \
  -o test_output/test_model

# Make predictions on the trained model
strainfish predict run \
  -f tests/test_input/predict.fasta \
  -m test_output/test_model \
  -o prediction_results

Dependencies

StrainFish has the following main dependencies:

• Core ML Libraries: numpy, pandas, scikit-learn, xgboost, cuml (GPU-accelerated)
• Sequence Processing: biopython, sourmash, sentencepiece
• CLI Interface: rich, rich-click
• Utilities: joblib, psutil, humanize, pynvml
• Testing: pytest, pytest-cov

For a complete list of dependencies, see pyproject.toml.

License

This project is licensed under the MIT License - see the LICENSE.md file for details.