PythiaLabelGenerator 1.0.0__py3-none-any.whl

labelgenerator/__init__.py

@@ -0,0 +1,15 @@
+import importlib.metadata
+
+try:
+    __version__ = importlib.metadata.distribution(__name__).version
+except Exception:
+    __version__ = "unknown"
+
+# Required if the package was installed via PyPi...
+if __version__ == "unknown":
+    try:
+        from importlib.metadata import version
+
+        __version__ = version("PythiaLabelGenerator")
+    except Exception:
+        pass

labelgenerator/iqtree.py

@@ -0,0 +1,142 @@
+import pathlib
+import subprocess
+from typing import Optional
+
+from pypythia.custom_types import DataType
+
+from labelgenerator.iqtree_parser import get_iqtree_results
+
+
+def _iqtree_results_exist_and_done(prefix: pathlib.Path) -> bool:
+    iqtree_file = pathlib.Path(f"{prefix}.iqtree")
+    logfile = pathlib.Path(f"{prefix}.log")
+
+    if not iqtree_file.exists() or not logfile.exists():
+        return False
+
+    return (
+        "Total wall-clock time used" in iqtree_file.read_text()
+        and "Date and Time:" in logfile.read_text()
+    )
+
+
+def get_iqtree_model(data_type: DataType) -> str:
+    """
+    Get the IQ-TREE model for the given data type.
+
+    Args:
+        data_type (DataType): The data type to get the model for.
+
+    Returns:
+        str: The IQ-TREE model for the given data type.
+            For DNA data: GTR+G4+FO
+            For AA data: LG+G4+FO
+            For morphological data: MK
+
+    """
+    return {
+        DataType.DNA: "GTR+G4+FO",
+        DataType.AA: "LG+G4+FO",
+        DataType.MORPH: "MK",
+    }[data_type]
+
+
+def run_statstests(
+    msa: pathlib.Path,
+    ml_trees: pathlib.Path,
+    best_tree: pathlib.Path,
+    iqtree: pathlib.Path,
+    model: str,
+    prefix: pathlib.Path,
+    seed: int = 0,
+    threads: Optional[int] = None,
+    is_morph: bool = False,
+    redo: bool = False,
+) -> None:
+    """
+    Run IQ-TREE statistical tests on the given set of ML trees. Will run all available tests (bp-RELL, KH (+weighted), SH (+weighted), ELW, AU) using 10,000 RELL bootstrap replicates.
+
+    Args:
+        msa (pathlib.Path): Path to the MSA file for which the ML trees were inferred.
+        ml_trees (pathlib.Path): Path to the file containing the ML trees.
+        best_tree (pathlib.Path): Path to the best ML tree.
+        iqtree (pathlib.Path): Path to the IQ-TREE executable.
+        model (str): The model to use for IQ-TREE.
+        prefix (pathlib.Path): Prefix for the output files.
+        seed (int): Seed for the random number generator. Defaults to 0.
+        threads (Optional[int]): Number of threads to use for IQ-TREE. Defaults to None. In this case, uses the IQ-TREE autoconfiguration.
+        is_morph (bool): Whether the data type is morphological. Defaults to False.
+        redo (bool): Whether to redo the analysis even if the results already exist. Defaults to False.
+    """
+    if not redo and _iqtree_results_exist_and_done(prefix):
+        return
+
+    cmd = [
+        iqtree,
+        "-s",
+        msa,
+        "-m",
+        model,
+        "-pre",
+        prefix,
+        "-z",
+        ml_trees,
+        "-treediff",
+        "-te",
+        best_tree,
+        "-n",
+        0,
+        "-zb",
+        10000,
+        "-zw",
+        "-au",
+        "-seed",
+        seed,
+    ]
+
+    if is_morph:
+        cmd.extend(["-st", "MORPH"])
+
+    if threads is not None:
+        cmd.extend(["-nt", threads])
+
+    if redo:
+        cmd.append("-redo")
+
+    try:
+        subprocess.check_output(list(map(str, cmd)), encoding="utf-8")
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"Running IQ-TREE command failed: {e.stdout}")
+    except Exception as e:
+        raise RuntimeError("Running IQ-TREE command failed.") from e
+
+
+def filter_plausible_trees(
+    ml_trees: pathlib.Path,
+    iqtree_results: pathlib.Path,
+    plausible_ml_trees: pathlib.Path,
+) -> None:
+    """
+    Filter the plausible ML trees based on the IQ-TREE results. A tree is plausible if all statistical tests (bp-RELL, KH (+weighted), SH (+weighted), ELW, AU) are significant.
+
+    Args:
+        ml_trees (pathlib.Path): Path to the file containing the ML trees the IQ-TREE were performed for.
+        iqtree_results (pathlib.Path): Path to the IQ-TREE results file.
+        plausible_ml_trees (pathlib.Path): Path to the file to write the plausible ML trees to.
+
+    Raises:
+        ValueError: If the number of IQ-TREE results does not match the number of ML trees.
+
+    """
+    iqtree_results = get_iqtree_results(iqtree_results)
+    newick_trees = [t.strip() for t in ml_trees.read_text().splitlines()]
+
+    if not len(iqtree_results) == len(newick_trees):
+        raise ValueError(
+            "Number of IQ-TREE results does not match the number of ML trees."
+        )
+
+    plausible_indices = [i for i, x in enumerate(iqtree_results) if x["plausible"]]
+    plausible_trees = [newick_trees[i] for i in plausible_indices]
+
+    plausible_ml_trees.write_text("\n".join(plausible_trees))

labelgenerator/iqtree_parser.py

@@ -0,0 +1,183 @@
+import pathlib
+
+import regex
+
+# define some regex stuff
+blanks = r"\s+"  # matches >=1  subsequent whitespace characters
+sign = r"[-+]?"  # contains either a '-' or a '+' symbol or none of both
+# matches ints or floats of forms '1.105' or '1.105e-5' or '1.105e5' or '1.105e+5'
+float_re = r"\d+(?:\.\d+)?(?:[e][-+]?\d+)?"
+
+tree_id_re = r"\d+"  # tree ID is an int
+llh_re = rf"{sign}{float_re}"  # likelihood is a signed floating point
+deltaL_re = rf"{sign}{float_re}"  # deltaL is a signed floating point
+# test result entry is of form '0.123 +'
+test_result_re = rf"{float_re}{blanks}{sign}"
+
+stat_test_name = r"[a-zA-Z-]+"
+
+# table header is of form:
+# Tree      logL    deltaL  bp-RELL    p-KH     p-SH    p-WKH    p-WSH       c-ELW       p-AU
+table_header = rf"Tree{blanks}logL{blanks}deltaL{blanks}(?:({stat_test_name})\s*)*"
+table_header_re = regex.compile(table_header)
+
+# a table entry in the .iqtree file looks for example like this:
+# 5 -5708.931281 1.7785e-06  0.0051 -  0.498 +  0.987 +  0.498 +  0.987 +      0.05 +    0.453 +
+table_entry = rf"({tree_id_re}){blanks}({llh_re}){blanks}({deltaL_re}){blanks}(?:({test_result_re})\s*)*"
+table_entry_re = regex.compile(table_entry)
+
+# if there is only a single plausible tree
+# the line will look like this:
+# 1  -88.9544627       0
+table_entry_single_plausible_tree = (
+    rf"({tree_id_re}){blanks}({llh_re}){blanks}({deltaL_re})\s*"
+)
+table_entry_single_plausible_tree_re = regex.compile(table_entry_single_plausible_tree)
+
+START_STRING = "USER TREES"
+END_STRING = "TIME STAMP"
+
+TEST_NAMES = ["bp-RELL", "p-KH", "p-SH", "p-WKH", "p-WSH", "c-ELW", "p-AU"]
+
+
+def get_relevant_section(input_file: pathlib.Path) -> list[str]:
+    """
+    Returns the section between the START_STRING and END_STRING in the given file.
+
+    Args:
+        input_file (pathlib.Path): Path to the .iqtree file to extract the section from.
+
+    Returns:
+        A list of strings, each representing a line of the relevant section.
+
+    Raises:
+        ValueError: If the section between START_STRING and END_STRING is empty.
+    """
+    content = input_file.read_text().splitlines()
+
+    # now let's find the relevant lines
+    # the relevant lines are only between the start and end string
+    start = 0
+    end = 0
+
+    for i, line in enumerate(content):
+        if START_STRING in line:
+            start = i
+        if END_STRING in line:
+            end = i
+
+    if start == end:
+        raise ValueError(
+            f"The section between START_STRING {START_STRING} and END_STRING {END_STRING} is empty. Please check the input file {input_file}."
+        )
+
+    return content[start:end]
+
+
+def _get_default_entry() -> dict:
+    """
+    Returns a default entry for a single plausible tree.
+
+    Returns:
+        A dict containing the default entry for a single plausible tree.
+    """
+    return {
+        "plausible": 1,
+        "tests": {
+            "bp-RELL": {"score": 1, "significant": True},
+            "p-KH": {"score": 1, "significant": True},
+            "p-SH": {"score": 1, "significant": True},
+            "p-WKH": {"score": 1, "significant": True},
+            "p-WSH": {"score": 1, "significant": True},
+            "c-ELW": {"score": 1, "significant": True},
+            "p-AU": {"score": 1, "significant": True},
+        },
+    }
+
+
+def _regex_group_to_test_results(raw_results: list[str]) -> dict:
+    assert len(TEST_NAMES) == len(raw_results)
+
+    data = {"tests": {}}
+    num_passed = 0
+
+    for i, test in enumerate(TEST_NAMES):
+        test_result = raw_results[i]
+        score, significant = test_result.split(" ")
+        score = score.strip()
+        significant = significant.strip()
+        data["tests"][test] = {}
+        data["tests"][test]["score"] = float(score)
+        data["tests"][test]["significant"] = True if significant == "+" else False
+
+        if data["tests"][test]["significant"]:
+            num_passed += 1
+
+    data["plausible"] = num_passed == len(data["tests"].keys())
+    return data
+
+
+def get_cleaned_table_entries(table_section: list[str]) -> list[dict]:
+    """
+    Returns a list of dicts, each dict contains the iqtree test results for the respective tree.
+
+    Args:
+        table_section (list[str]): A list of strings, each representing a line of the relevant section.
+
+    Returns:
+        A list of dicts. Each dict contains the tree_id, llh, deltaL and all results of the performed
+            iqtree tests.
+
+    """
+    entries = []
+    for line in table_section:
+        line = line.strip()
+        # match the line against the regex defined above for a table entry
+        m = regex.match(table_entry_re, line)
+
+        # and match the line against the regex for a table entry in case of a single plausible tree
+        m_single_tree = regex.match(table_entry_single_plausible_tree_re, line)
+
+        if m:
+            # transform the raw results to a python dict
+            entry = _regex_group_to_test_results(m.captures(4))
+            entries.append(entry)
+        elif m_single_tree:
+            # if a match for a truncated table entry was found: we only have a single plausible tree
+            # => add the entry manually
+            entry = _get_default_entry()
+            entries.append(entry)
+        elif "= tree" in line:
+            # indicates that a tree is identical to one seen before
+            # => duplicate the results of this tree
+            _, id_of_identical_tree = line.rsplit(" ", 1)
+            id_of_identical_tree = int(id_of_identical_tree)
+
+            # IQ-Tree reports the results 1-indexed
+            # => to get the correct results we need to subtract one and access the entries
+            entry = entries[id_of_identical_tree - 1].copy()
+            entries.append(entry)
+
+    if not entries:
+        raise ValueError(
+            "No line in the given section matches the regex. Compare the regex and the given section. "
+            "Maybe the format has changed."
+        )
+
+    return entries
+
+
+def get_iqtree_results(iqtree_file: pathlib.Path) -> list[dict]:
+    """
+    Returns a list of dicts, each dict contains the iqtree test results for the respective tree.
+
+    Args:
+        iqtree_file (pathlib.Path): Path to the .iqtree file to extract the results from.
+
+    Returns:
+        A list of dicts. Each dict contains the tree_id, llh, deltaL and all results of the performed
+            iqtree tests.
+    """
+    section = get_relevant_section(iqtree_file)
+    entries = get_cleaned_table_entries(section)
+    return entries

labelgenerator/label.py

@@ -0,0 +1,231 @@
+import pathlib
+from typing import Optional
+
+from pypythia.msa import parse_msa
+
+from labelgenerator.iqtree import (
+    filter_plausible_trees,
+    get_iqtree_model,
+    run_statstests,
+)
+from labelgenerator.logger import log_runtime_information
+from labelgenerator.raxmlng import infer_ml_trees, rf_distance
+
+
+def get_label(
+    n_all: int,
+    rf_all: float,
+    n_unique_all: int,
+    n_plausible: int,
+    rf_plausible: float,
+    n_unique_plausible: int,
+) -> float:
+    r"""
+    Compute the ground truth difficulty label for the given input values.
+
+    The ground-truth difficulty is computed according to our definition published in [Haag _et al._ (2022)](https://doi.org/10.1093/molbev/msac254):
+
+    Let $N_{\text{all}}$ be the number of inferred ML trees.
+    We first compute the average pairwise relative Robinson-Foulds (RF) distance between all trees ($RF_{\text{all}}$), as
+    well as the number of unique tree topologies among the inferred trees ($N^*_{\text{all}}$).
+    We filter the inferred trees using likelihood-based statistical tests to obtain the set of $N_{\text{pl}}$ _plausible
+    trees_.
+    We again compute the average pairwise RF distance between the plausible trees ($RF_{\text{pl}}$) and the number of
+    unique tree topologies among the plausible trees ($N^*_{\text{pl}}$).
+
+    The difficulty is then computed as follows:
+
+    $$
+    \text{difficulty} = \frac{1}{5} \cdot \bigg[ RF_{\text{all}} + RF_{\text{pl}}
+    + \frac{N^*_{\text{all}}}{N_{\text{all}}} + \frac{N^*_{\text{pl}}}{N_{\text{pl}}}
+    + \left( 1 - \frac{N_{\text{pl}}}{N_{\text{all}}} \right) \bigg]
+    $$
+
+    Args:
+        n_all (int): Number of inferred ML trees.
+        rf_all (float): Average pairwise RF distance between all trees.
+        n_unique_all (int): Number of unique tree topologies among the inferred trees.
+        n_plausible (int): Number of plausible trees.
+        rf_plausible (float): Average pairwise RF distance between the plausible trees.
+        n_unique_plausible (int): Number of unique tree topologies among the plausible trees.
+
+    Returns:
+        float: Ground truth difficulty label.
+
+    Raises:
+        ValueError: If the input values are not within the expected range.
+            - Number of unique trees is higher than the total number of trees.
+            - Number of unique plausible trees is higher than the number of plausible trees.
+            - Number of plausible trees is higher than the total number of trees.
+            - RF distance for all trees is not between 0 and 1.
+            - RF distance for plausible trees is not between 0 and 1.
+
+    """
+    if n_unique_all > n_all:
+        raise ValueError(
+            "Number of unique trees cannot be higher than the total number of trees."
+        )
+    if n_unique_plausible > n_plausible:
+        raise ValueError(
+            "Number of unique plausible trees cannot be higher than the number of plausible trees."
+        )
+    if n_plausible > n_all:
+        raise ValueError(
+            "Number of plausible trees cannot be higher than the total number of trees."
+        )
+    if not 0 <= rf_all <= 1:
+        raise ValueError("RF distance for all trees must be between 0 and 1.")
+    if not 0 <= rf_plausible <= 1:
+        raise ValueError("RF distance for plausible trees must be between 0 and 1.")
+
+    proportion_unique_all = n_unique_all / n_all
+    proportion_unique_plausible = n_unique_plausible / n_plausible
+    proportion_plausible = n_plausible / n_all
+
+    total = (
+        rf_all
+        + proportion_unique_all
+        + rf_plausible
+        + proportion_unique_plausible
+        + (1 - proportion_plausible)
+    )
+    label = total / 5
+
+    eps = 1e-9
+    assert -eps <= label <= 1 + eps, (
+        f"Label {label} is not between 0 and 1. Check the input values."
+    )
+
+    return label
+
+
+def compute_label(
+    msa_file: pathlib.Path,
+    raxmlng: pathlib.Path,
+    iqtree: pathlib.Path,
+    prefix: pathlib.Path,
+    model: Optional[str] = None,
+    n_trees: int = 100,
+    seed: int = 0,
+    threads: Optional[int] = None,
+    redo: bool = False,
+    log_info: bool = True,
+) -> float:
+    """
+    Compute the ground truth difficulty label for the given input MSA by inferring ML trees and running statistical tests.
+    See `labelgenerator.label.get_label` for the definition of the ground truth difficulty.
+
+
+    Args:
+        msa_file (pathlib.Path): Path to the MSA file to compute the label for. Can be either in FASTA or PHYLIP format.
+        raxmlng (pathlib.Path): Path to the RAxML-NG executable.
+        iqtree (pathlib.Path): Path to the IQ-TREE executable.
+        prefix (pathlib.Path): Prefix for the output files.
+        model (str, optional): Substitution model to use for the ML tree inference. Defaults to None. In this case, the
+            model is inferred from the MSA based on the data type (`GTR+G` for DNA, `LG+G` for AA, and `MULTIx_GTR` for morphological data).
+        n_trees (int, optional): Number of ML trees to infer. Defaults to 100. Please note that the computed label is only comparable to
+            Pythia predictions if 100 trees are inferred.
+        seed (int, optional): Seed for the random number generator. Defaults to 0.
+        threads (int, optional): Number of threads to use for the ML tree inference. Defaults to None. In this case, the RAxML-NG and IQ-TREE autoconfigs are used.
+        redo (bool, optional): If True, the computations are redone even if the output files already exist. Defaults to False.
+        log_info (bool, optional): If True, runtime information is logged. Defaults to True.
+
+    Returns:
+        float: The ground truth difficulty label for the given MSA.
+
+    """
+    msa_obj = parse_msa(msa_file)
+    model = model or msa_obj.get_raxmlng_model()
+
+    # 1. Infer 100 ML trees for the given MSA using RAxML-NG
+    if log_info:
+        log_runtime_information(f"Inferring {n_trees} ML trees using RAxML-NG.")
+
+    infer_ml_trees(
+        msa=msa_file,
+        raxmlng=raxmlng,
+        model=model,
+        prefix=prefix,
+        n_trees=n_trees,
+        seed=seed,
+        threads=threads,
+        redo=redo,
+    )
+
+    # 2. RF-Distance ML trees
+    ml_trees = pathlib.Path(f"{prefix}.raxml.mlTrees")
+    rfdistance_prefix = pathlib.Path(f"{prefix}.rfdist")
+
+    if log_info:
+        log_runtime_information("Computing RF-Distance between ML trees.")
+
+    n_unique_all, rf_all = rf_distance(
+        ml_trees=ml_trees, prefix=rfdistance_prefix, raxmlng=raxmlng, redo=redo
+    )
+
+    if log_info:
+        log_runtime_information(f"> RF-Distance ML trees: {round(rf_all, 2)}")
+        log_runtime_information(f"> Unique topologies ML trees: {n_unique_all}")
+
+    # 3. IQ-TREE statistical tests
+    best_tree = pathlib.Path(f"{prefix}.raxml.bestTree")
+    iqtree_prefix = pathlib.Path(f"{prefix}.iqtree")
+    if log_info:
+        log_runtime_information("Running IQ-TREE statistical tests.")
+    run_statstests(
+        msa=msa_file,
+        ml_trees=ml_trees,
+        best_tree=best_tree,
+        iqtree=iqtree,
+        model=get_iqtree_model(msa_obj.data_type),
+        prefix=iqtree_prefix,
+        seed=seed,
+        threads=threads,
+        redo=redo,
+    )
+
+    # 4. Filter the plausible trees
+    iqtree_results = pathlib.Path(f"{iqtree_prefix}.iqtree")
+    plausible_ml_trees = pathlib.Path(f"{prefix}.raxml.plausibleTrees")
+    if log_info:
+        log_runtime_information("Filtering plausible ML trees.")
+    filter_plausible_trees(
+        ml_trees=ml_trees,
+        iqtree_results=iqtree_results,
+        plausible_ml_trees=plausible_ml_trees,
+    )
+
+    n_plausible_trees = sum(1 for _ in plausible_ml_trees.open())
+
+    if log_info:
+        log_runtime_information(f"> Found {n_plausible_trees} plausible trees.")
+
+    # 5. RF-Distance plausible trees
+    rfdistance_plausible_prefix = pathlib.Path(f"{prefix}.rfdist.plausible")
+    if log_info:
+        log_runtime_information("Computing RF-Distance between plausible ML trees.")
+    n_unique_plausible, rf_plausible = rf_distance(
+        ml_trees=plausible_ml_trees,
+        prefix=rfdistance_plausible_prefix,
+        raxmlng=raxmlng,
+        redo=redo,
+    )
+
+    if log_info:
+        log_runtime_information(
+            f"> RF-Distance plausible trees: {round(rf_plausible, 2)}"
+        )
+        log_runtime_information(
+            f"> Unique topologies plausible trees: {n_unique_plausible}"
+        )
+
+    # 6. Compute the ground truth difficulty
+    difficulty = get_label(
+        n_all=n_trees,
+        rf_all=rf_all,
+        n_unique_all=n_unique_all,
+        n_plausible=n_plausible_trees,
+        rf_plausible=rf_plausible,
+        n_unique_plausible=n_unique_plausible,
+    )
+    return difficulty

labelgenerator/logger.py

@@ -0,0 +1,32 @@
+import sys
+import textwrap
+import time
+
+import loguru
+from labelgenerator import __version__
+
+SCRIPT_START = time.perf_counter()
+
+
+logger = loguru.logger
+logger.remove()
+logger.add(sys.stderr, format="{message}")
+
+
+def get_header():
+    return textwrap.dedent(
+        f"Difficulty LabelGenerator version {__version__} released by The Exelixis Lab\n"
+        f"Developed by: Julia Haag\n"
+        f"Latest version: https://github.com/tschuelia/LabelGenerator\n"
+        f"Questions/problems/suggestions? Please open an issue on GitHub.\n",
+    )
+
+
+def log_runtime_information(message, log_runtime=True):
+    if log_runtime:
+        seconds = time.perf_counter() - SCRIPT_START
+        fmt_time = time.strftime("%H:%M:%S", time.gmtime(seconds))
+        time_string = f"[{fmt_time}] "
+    else:
+        time_string = ""
+    logger.info(f"{time_string}{message}")