Etetoolkit

from collections import Counter
from pathlib import Path

from ete4 import Tree


def load_and_validate(path: Path, parser=1) -> Tree:
    with path.open(encoding="utf-8") as handle:
        tree = Tree(handle, parser=parser)

    names = list(tree.leaf_names())
    empty = [leaf.id for leaf in tree.leaves() if not leaf.name]
    duplicate_counts = {
        name: count
        for name, count in Counter(names).items()
        if name and count > 1
    }

    if empty:
        raise ValueError(f"{path}: unnamed leaves at positions {empty[:10]}")
    if duplicate_counts:
        raise ValueError(f"{path}: duplicate leaf names {duplicate_counts}")
    if len(names) < 2:
        raise ValueError(f"{path}: expected at least two leaves")

    return tree


tree = load_and_validate(Path("tree.nw"), parser=1)

supports = [
    node.support
    for node in tree.traverse()
    if not node.is_leaf and node.support is not None
]
distances = [
    node.dist
    for node in tree.traverse()
    if not node.is_root and node.dist is not None
]

if any(distance < 0 for distance in distances):
    raise ValueError("Negative branch length detected")

print("support range:", (min(supports), max(supports)) if supports else None)
print("distance range:", (min(distances), max(distances)) if distances else None)

from collections import Counter
from pathlib import Path

from ete4 import Tree

input_path = Path("inferred_tree.nw")
output_path = Path("processed_tree.nw")

with input_path.open(encoding="utf-8") as handle:
    tree = Tree(handle, parser="support")

# Keep a known sample set and preserve retained pairwise distances.
keep = ["sample_A", "sample_B", "sample_C", "outgroup"]
counts = Counter(tree.leaf_names())
invalid = {name: counts[name] for name in keep if counts[name] != 1}
if invalid:
    raise ValueError(f"Requested tips must exist exactly once: {invalid}")
tips = {leaf.name: leaf for leaf in tree.leaves()}
tree.prune([tips[name] for name in keep], preserve_branch_length=True)

# Prefer a biological outgroup when justified.
tree.set_outgroup(tips["outgroup"])

# Stable presentation order only; this does not alter clade membership.
tree.ladderize()

tree.write(
    outfile=str(output_path),
    parser="support",
    props=[],
)

from ete4 import Tree

tree = Tree("((A:1,B:1):1,(C:2,D:2):1);")
tree.set_midpoint_outgroup()
print(tree.write(props=[]))

tree = Tree("((A:1,B:1):1,(C:2,D:2):1);")
candidate = tree.get_midpoint_outgroup()
candidate_id = candidate.id
tree.set_midpoint_outgroup()
print("midpoint candidate:", candidate_id)

from collections import Counter
from pathlib import Path

from ete4 import Tree


def load(path: Path, parser=1) -> Tree:
    with path.open(encoding="utf-8") as handle:
        return Tree(handle, parser=parser)


def assert_unique_leaf_names(tree: Tree, label: str) -> None:
    counts = Counter(tree.leaf_names())
    if any(not name for name in counts):
        raise ValueError(f"{label} has unnamed leaves")
    duplicates = sorted(name for name, count in counts.items() if count > 1)
    if duplicates:
        raise ValueError(f"{label} has duplicate leaves: {duplicates}")


tree_a = load(Path("method_a.nw"))
tree_b = load(Path("method_b.nw"))
assert_unique_leaf_names(tree_a, "method_a")
assert_unique_leaf_names(tree_b, "method_b")

(
    rf,
    max_rf,
    common,
    edges_a,
    edges_b,
    discarded_a,
    discarded_b,
) = tree_a.robinson_foulds(
    tree_b,
    unrooted_trees=True,
)

print(
    {
        "rf": rf,
        "max_rf": max_rf,
        "normalized_rf": rf / max_rf if max_rf > 0 else None,
        "common_leaf_count": len(common),
        "discarded_edges_a": len(discarded_a),
        "discarded_edges_b": len(discarded_b),
    }
)

summary = tree_a.compare(tree_b, unrooted=True)
print(summary["rf"], summary["max_rf"], summary["norm_rf"])

import csv
from pathlib import Path

from ete4 import Tree

tree = Tree("((sample_1,sample_2),sample_3);")

metadata: dict[str, dict[str, str]] = {}
with Path("metadata.tsv").open(encoding="utf-8", newline="") as handle:
    reader = csv.DictReader(handle, delimiter="\t")
    required = {"sample", "host", "location"}
    if not reader.fieldnames or not required.issubset(reader.fieldnames):
        raise ValueError(f"metadata.tsv must contain columns {sorted(required)}")
    if len(reader.fieldnames) != len(set(reader.fieldnames)):
        raise ValueError("metadata.tsv has duplicate column names")
    for row in reader:
        if None in row or any(value is None for value in row.values()):
            raise ValueError(f"metadata.tsv has a malformed row at line {reader.line_num}")
        sample = row["sample"].strip()
        if not sample or sample in metadata:
            raise ValueError(f"Empty or duplicate metadata key: {sample!r}")
        metadata[sample] = {
            "host": row["host"].strip(),
            "location": row["location"].strip(),
        }

unmatched_tree = []
for leaf in tree.leaves():
    values = metadata.get(leaf.name)
    if values is None:
        unmatched_tree.append(leaf.name)
        continue
    leaf.add_props(**values)

unmatched_metadata = sorted(set(metadata) - set(tree.leaf_names()))
if unmatched_tree or unmatched_metadata:
    raise ValueError(
        f"Unmatched tree leaves={unmatched_tree}; "
        f"unmatched metadata rows={unmatched_metadata}"
    )

tree.write(
    outfile="annotated.nhx",
    props=["host", "location"],
)

from ete4 import PhyloTree

gene_tree = PhyloTree(
    "((Hsa|gene1,Ptr|gene1),(Hsa|gene2,Mmu|gene1));",
    sp_naming_function=lambda name: name.split("|", 1)[0],
)

events = gene_tree.get_descendant_evol_events(sos_thr=0.0)
for event in events:
    relationship = {
        "S": "orthology",
        "D": "paralogy",
    }[event.etype]
    print(
        relationship,
        sorted(event.in_seqs),
        sorted(event.out_seqs),
    )

for node in gene_tree.traverse():
    event_type = node.props.get("evoltype")
    if event_type:
        print(node.id, event_type)

from ete4 import PhyloTree

gene_newick = (
    "((Dme_001,Dme_002),"
    "(((Cfa_001,Mms_001),((Hsa_001,Ptr_001),Mmu_001)),"
    "(Ptr_002,(Hsa_002,Mmu_002))));"
)
species_newick = "((((Hsa,Ptr),Mmu),(Mms,Cfa)),Dme);"

species_from_gene = lambda name: name.split("_", 1)[0]

gene_tree = PhyloTree(
    gene_newick,
    sp_naming_function=species_from_gene,
)
species_tree = PhyloTree(
    species_newick,
    sp_naming_function=lambda name: name,
)

gene_species = {leaf.species for leaf in gene_tree.leaves()}
species_tree_tips = set(species_tree.leaf_names())
missing = sorted(gene_species - species_tree_tips)
if missing:
    raise ValueError(f"Species absent from species tree: {missing}")

reconciled_tree, events = gene_tree.reconcile(species_tree)

for event in events:
    if event.etype == "S":
        print("orthology", event.inparalogs, event.orthologs)
    elif event.etype == "D":
        print("paralogy", event.inparalogs, event.outparalogs)

reconciled_tree.write(
    outfile="reconciled.nhx",
    props=["evoltype"],
)

from pathlib import Path

from ete4 import PhyloTree

tree = PhyloTree(
    "((Human_1,Chimp_1),(Human_2,(Chimp_2,Mouse_1)));",
    sp_naming_function=lambda name: name.split("_", 1)[0],
)

output_dir = Path("subfamilies")
output_dir.mkdir(parents=True, exist_ok=True)

for index, subtree in enumerate(tree.split_by_dups(), start=1):
    subtree.write(
        outfile=str(output_dir / f"subfamily_{index:03d}.nw"),
        props=[],
    )

tree_count, duplication_count, newicks = tree.get_speciation_trees(
    newick_only=True,
)
for newick in newicks:
    process(newick)

from ete4 import NCBITaxa

ncbi = NCBITaxa()
taxids = [9606, 9598, 10090]
tree = ncbi.get_topology(taxids, intermediate_nodes=False, annotate=True)
print(tree.to_str(props=["sci_name", "rank", "taxid"]))

from ete4 import GTDBTaxa

gtdb = GTDBTaxa()
taxa = ["p__Huberarchaeota", "o__Peptococcales", "f__Korarchaeaceae"]
tree = gtdb.get_topology(taxa, intermediate_nodes=True, annotate=True)
print(tree.to_str(props=["sci_name", "rank"]))

from pathlib import Path

from ete4 import Tree

input_dir = Path("trees")
output_dir = Path("processed")
output_dir.mkdir(parents=True, exist_ok=True)

for input_path in sorted(input_dir.glob("*.nw")):
    with input_path.open(encoding="utf-8") as handle:
        tree = Tree(handle, parser="support")

    names = list(tree.leaf_names())
    if len(names) != len(set(names)):
        raise ValueError(f"{input_path}: duplicate leaf names")

    tree.set_midpoint_outgroup()
    tree.ladderize()

    output_path = output_dir / input_path.name
    tree.write(
        outfile=str(output_path),
        parser="support",
        props=[],
    )

for leaf in tree.leaves():
    process(leaf)

leaves = list(tree.leaves())

leaf_cache = tree.get_cached_content(prop="name")
for node in tree.traverse("postorder"):
    descendant_names = leaf_cache[node]
    summarize(node, descendant_names)

from ete4 import Tree
from ete4.treematcher import TreePattern

tree = Tree("((K,((A,B),C),D),(E,F));")
pattern = TreePattern("(,,)", safer=True)

for match in pattern.search(tree):
    print("three-child node:", match.id, match.name)