[Feature Request] - Add wrapper for CheckM2

[jolespin]

The rationale behind this is that one might bin MAGs and cluster them then profile the metagenomics to try to capture signal from MAGs that were present but not at sufficient coverage to bin. In this case, it would be good to wrap CheckM2 in some way.

The pipeline I'm envisioning:

• Run Diamond against CheckM2 database
• Build Leviathan database from CheckM2 hits
• Wrap CheckM2 for hits

Requirements will need to change and add CheckM2 but only the python module. I don't want to force Diamond dependency.

[jolespin]

CheckM2 Python API

Use CheckM2's internal Python modules to compute genome completeness and
contamination from your own Diamond BLASTP results.

Quick Start

# protein_to_genome: dict mapping each protein ID to its genome name
# genome_to_proteinfasta: dict mapping each genome name to its .faa path

df_diamond = prepare_diamond("diamond_output.tsv", protein_to_genome)
df_metadata = calculate_metadata(genome_to_proteinfasta)
df_results = predict_quality(df_diamond, df_metadata, threads=8)

Pipeline Overview

Protein FASTA files
  ├─ Diamond BLASTP against CheckM2 DB
  │    └─ Parse hits → KEGG annotations per genome
  │         ├─ KO gene presence/absence     (19,999 features)
  │         ├─ KO pathway completeness        (416 features)
  │         ├─ KO module completeness          (757 features)
  │         └─ KO category completeness         (47 features)
  └─ Protein metadata
       ├─ Amino acid counts (A..Y)             (20 features)
       ├─ AALength (total residues)              (1 feature)
       └─ CDS count                              (1 feature)
                           ↓
              Feature vector (~21,241 columns)
                           ↓
         ┌─────────────────┴─────────────────┐
   General Model (LightGBM)         Specific Model (Keras NN)
   → Completeness + Contamination   → Completeness only
         └─────────────────┬─────────────────┘
                           ↓
              Cosine similarity vs reference genomes
              → Select best completeness prediction
                           ↓
              Final: Completeness, Contamination, Model Used

Diamond Parameters

If you run Diamond yourself, use these settings for consistent results:

Parameter	Value
--outfmt	6
--max-target-seqs	1
--query-cover	80
--subject-cover	80
--id	30
--evalue	1e-05

Subject hits in the CheckM2 database are formatted as
{Ref100_hit}~{KEGG_annotation}. Parse the part after ~ to get
the Kegg_annotation column.

Functions

import gzip
import shutil
import tempfile

import numpy as np
import pandas as pd
from checkm2.keggData import KeggCalculator
from checkm2.metadata import MetadataCalculator
from checkm2.modelProcessing import modelProcessor
from checkm2.modelPostprocessing import modelProcessor as postModelProcessor


def prepare_diamond(
    diamond_output: str,
    protein_to_genome: dict,
) -> pd.DataFrame:
    """Parse Diamond BLASTP output into the format expected by predict_quality.

    Parameters
    ----------
    diamond_output : str or Path
        Path to Diamond outfmt 6 TSV file.

    protein_to_genome : dict or pd.Series
        Maps protein IDs (Diamond query column) to genome names.

    Returns
    -------
    pd.DataFrame
        Columns: GenomeID, Kegg_annotation
    """
    df_raw = pd.read_csv(
        str(diamond_output), sep="\t", usecols=[0, 1],
        names=["ProteinID", "subject"],
    )
    df_raw["Kegg_annotation"] = df_raw["subject"].str.split("~", n=1).str[1]
    df_raw["GenomeID"] = df_raw["ProteinID"].map(protein_to_genome)
    return df_raw[["GenomeID", "Kegg_annotation"]].dropna()


def calculate_metadata(genome_to_proteinfasta: dict) -> pd.DataFrame:
    """Calculate protein metadata from FASTA files.

    Parameters
    ----------
    genome_to_proteinfasta : dict
        Maps genome name (str) to protein FASTA file path (str or Path).
        Files may be gzipped (.gz).

    Returns
    -------
    pd.DataFrame
        Columns: Name, CDS, AALength, A, C, D, E, F, G, H, I, K, L, M, N,
        P, Q, R, S, T, V, W, Y
    """
    rows = []
    for genome_name, faa_path in genome_to_proteinfasta.items():
        faa_path = str(faa_path)

        if faa_path.endswith(".gz"):
            with tempfile.NamedTemporaryFile(suffix=".faa", delete=False) as tmp:
                with gzip.open(faa_path, "rb") as gz:
                    shutil.copyfileobj(gz, tmp)
                tmp_path = tmp.name
            mc = MetadataCalculator(tmp_path)
        else:
            mc = MetadataCalculator(faa_path)

        _, cds = mc.calculate_CDS()
        _, aa_length = mc.calculate_amino_acid_length()
        _, aa_list, aa_counts = mc.calculate_amino_acid_counts()
        row = {"Name": genome_name, "CDS": cds, "AALength": aa_length}
        row.update(dict(zip(aa_list, aa_counts)))
        rows.append(row)
    return pd.DataFrame(rows)


def predict_quality(
    df_diamond: pd.DataFrame,
    df_metadata: pd.DataFrame,
    threads: int = 1,
) -> pd.DataFrame:
    """Predict genome completeness and contamination from Diamond BLASTP results.

    Parameters
    ----------
    df_diamond : pd.DataFrame
        Diamond output with columns:
          - "GenomeID": genome identifier (str)
          - "Kegg_annotation": KEGG ortholog ID, e.g. "K00001" (str)

    df_metadata : pd.DataFrame
        Per-genome protein statistics as returned by calculate_metadata().
        "Name" column must match the "GenomeID" values in df_diamond.

    threads : int
        Threads for LightGBM prediction.

    Returns
    -------
    pd.DataFrame
        Columns: Name, Completeness, Contamination, Completeness_Model_Used
    """
    KeggCalc = KeggCalculator()

    # -- 1. Build KEGG feature vectors from Diamond hits --
    defaultKOs = KeggCalc.return_default_values_from_category("KO_Genes")
    df_diamond = df_diamond[
        df_diamond["Kegg_annotation"].isin(defaultKOs.keys())
    ].copy()

    annot_dict = dict(
        zip(
            sorted(df_diamond["GenomeID"].unique()),
            [x for _, x in df_diamond.groupby("GenomeID")],
        )
    )

    genome_names = sorted(df_metadata["Name"].values)

    kegg_genome_list = []
    for genome in genome_names:
        sub_dict = defaultKOs.copy()
        sub_dict["Name"] = genome
        if genome in annot_dict:
            ko_counts = annot_dict[genome]["Kegg_annotation"].value_counts()
            sub_dict.update(ko_counts)
        kegg_genome_list.append(sub_dict)

    df_ko_genes = pd.DataFrame(kegg_genome_list)

    df_ko_pathways = KeggCalc.calculate_KO_group("KO_Pathways", df_ko_genes.copy())
    df_ko_categories = KeggCalc.calculate_KO_group("KO_Categories", df_ko_genes.copy())
    df_ko_modules = KeggCalc.calculate_module_completeness(df_ko_genes.copy())

    df_diamond_features = pd.concat(
        [df_ko_genes, df_ko_pathways, df_ko_modules, df_ko_categories], axis=1
    )

    # -- 2. Arrange metadata in the correct column order --
    metadata_order = KeggCalc.return_proper_order("Metadata")
    metadata_order_with_name = ["Name"] + metadata_order
    df_metadata = df_metadata[metadata_order_with_name].copy()
    df_metadata.sort_values(by="Name", inplace=True)
    df_metadata.reset_index(drop=True, inplace=True)

    df_diamond_features.sort_values(by="Name", inplace=True)
    df_diamond_features.reset_index(drop=True, inplace=True)

    # -- 3. Concatenate into final feature matrix --
    df_features = pd.concat(
        [df_metadata, df_diamond_features.drop("Name", axis=1)], axis=1
    )
    vector_array = df_features.iloc[:, 1:].values.astype(float)

    # -- 4. Run ML models --
    modelProc = modelProcessor(threads)
    general_completeness, general_contamination = modelProc.run_prediction_general(vector_array)

    ko_list_length = len(defaultKOs)
    specific_model_vector_len = ko_list_length + len(metadata_order)
    specific_completeness, scaled_features = modelProc.run_prediction_specific(
        vector_array, specific_model_vector_len
    )

    # -- 5. Model selection via cosine similarity --
    postProc = postModelProcessor(threads)
    final_completeness, final_contamination, model_chosen, csm_array = (
        postProc.calculate_general_specific_ratio(
            vector_array[:, 20],
            scaled_features,
            general_completeness,
            general_contamination,
            specific_completeness,
        )
    )

    return pd.DataFrame(
        {
            "Name": genome_names,
            "Completeness": np.round(final_completeness, 2),
            "Contamination": np.round(final_contamination, 2),
            "Completeness_Model_Used": model_chosen,
        }
    )

Internals

Details you shouldn't need for normal use, but useful for debugging or
extending the pipeline.

Feature Ordering

Column order is defined in checkm2/data/feature_ordering.json and must
match exactly for the ML models. The order is:

1. Metadata (22 columns): A, C, D, E, F, G, H, I, K, L, M, N, P, Q, R, S, T, V, W, Y, AALength, CDS
2. KO Genes (19,999 columns): presence/absence of KEGG orthologs
3. KO Pathways (416 columns): fraction of KOs present per pathway
4. KO Modules (757 columns): fraction of KOs present per module
5. KO Categories (47 columns): fraction of KOs present per category

Models

• General model (general_model_COMP.gbm, model_CONT.gbm): LightGBM. Predicts both completeness and contamination.
• Specific model (specific_model_COMP.keras): Keras NN with MinMaxScaler (scaler.sav). Predicts completeness only.
• Contamination always comes from the general model.

Model Selection Logic

Both models run on every genome. The final completeness is chosen based on:

• Cosine similarity of input features vs reference genomes (min_ref_rsdata_v1.npz)
• Novelty ratio: general_completeness / cosine_similarity^2
• AA ratio: AALength / mean_completeness — values < 1500 indicate reduced genomes
• Highly reduced genomes (AA ratio < 1500 and mean completeness < 55) always use the general model
• Otherwise, lower novelty favors the specific (NN) model