phytreon 0.1.1__py3-none-any.whl

phytreon/__init__.py

@@ -0,0 +1,49 @@
+"""phytreon -- phylogenetic trees and publication figures in Python.
+
+Quick start
+-----------
+>>> import phytreon as pt
+>>> tr = pt.datasets.primates()
+>>> (pt.TreeFigure(tr)
+...     .tip_labels()
+...     .support_labels()).save("tree.pdf")
+
+The package is layered:
+
+* ``phytreon.core``        -- ``Tree`` / ``Node`` data model and I/O
+* ``phytreon.layout``      -- topology -> display coordinates
+* ``phytreon.infer``       -- distance-based inference (NJ / UPGMA), ML, parsimony
+* ``phytreon.comparative`` -- ancestral state reconstruction
+* ``phytreon.plot``        -- the ``TreeFigure`` builder + matplotlib/plotly backends
+"""
+from __future__ import annotations
+
+from . import core, layout, infer, comparative, plot, datasets, treeops
+from .core import Tree, Node
+from .plot import TreeFigure
+from .infer import (
+    neighbor_joining, upgma, tree_from_alignment, distance_matrix,
+    Alignment, align, read_fasta, trim, bootstrap_support, build_tree, infer_ml,
+    ml_tree, log_likelihood, model_finder, parsimony_tree, parsimony_score,
+)
+from .comparative import ace_parsimony, ace_ml, ace_continuous, stochastic_map
+from .treeops import (
+    rotate, flip, swap_children, ladderize, collapse_low_support,
+    scale_clade, cut_tree, midpoint_root, group_clade, group_otu,
+    robinson_foulds,
+)
+
+__version__ = "0.1.1"
+
+__all__ = [
+    "core", "layout", "infer", "comparative", "plot", "datasets",
+    "Tree", "Node",
+    "TreeFigure",
+    "neighbor_joining", "upgma", "tree_from_alignment", "distance_matrix",
+    "Alignment", "align", "read_fasta", "trim", "bootstrap_support",
+    "build_tree", "infer_ml", "ml_tree", "log_likelihood", "model_finder",
+    "parsimony_tree", "parsimony_score", "robinson_foulds",
+    "ace_parsimony", "ace_ml", "ace_continuous", "stochastic_map",
+    "rotate", "flip", "swap_children", "ladderize", "collapse_low_support",
+    "scale_clade", "cut_tree", "midpoint_root", "group_clade", "group_otu",
+]

phytreon/comparative/__init__.py

@@ -0,0 +1,5 @@
+"""Comparative methods (ancestral state reconstruction, stochastic mapping)."""
+from .ace import ace_parsimony, ace_ml, ace_continuous
+from .stochastic_mapping import stochastic_map
+
+__all__ = ["ace_parsimony", "ace_ml", "ace_continuous", "stochastic_map"]

phytreon/comparative/ace.py

@@ -0,0 +1,223 @@
+"""Ancestral character estimation.
+
+Two discrete-trait reconstructions are provided:
+
+* :func:`ace_parsimony` -- Fitch parsimony (fast, model-free).
+* :func:`ace_ml` -- marginal maximum likelihood under an equal-rates
+  Mk model (Felsenstein pruning for the likelihood, an up/down pass for
+  the marginal posterior at every internal node, single rate estimated
+  by ML).
+
+Both write results onto ``node.data`` so the plotting layer can map them
+to aesthetics (e.g. pie charts / coloured node points at ancestors).
+Continuous-trait reconstruction (PIC / REML) is a documented extension
+point -- see :func:`ace_continuous`.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+from ..core.tree import Node, Tree
+
+
+# --------------------------------------------------------------------------
+# Fitch parsimony
+# --------------------------------------------------------------------------
+def ace_parsimony(tree: Tree, trait: Dict[str, str]) -> Dict[str, object]:
+    """Fitch parsimony reconstruction of a discrete trait.
+
+    ``trait`` maps tip name -> state.  Tips missing from the dict are
+    treated as fully ambiguous.  Each node gets ``data['ace_state']``;
+    the function returns ``{"score": int, "states": [...]}``.
+    """
+    states = sorted({v for v in trait.values()})
+    full = set(states)
+
+    # down-pass: state sets + score
+    score = 0
+    for node in tree.traverse("postorder"):
+        if node.is_leaf:
+            s = trait.get(node.name)
+            node.data["_pset"] = {s} if s is not None else set(full)
+        else:
+            sets = [c.data["_pset"] for c in node.children]
+            inter = set.intersection(*sets) if sets else set(full)
+            if inter:
+                node.data["_pset"] = inter
+            else:
+                node.data["_pset"] = set.union(*sets)
+                score += 1
+
+    # up-pass: resolve, preferring the parent's chosen state
+    for node in tree.traverse("preorder"):
+        pset = node.data.pop("_pset")
+        if node.is_root:
+            chosen = sorted(pset)[0]
+        else:
+            parent_state = node.parent.data["ace_state"]
+            chosen = parent_state if parent_state in pset else sorted(pset)[0]
+        node.data["ace_state"] = chosen
+
+    return {"score": score, "states": states}
+
+
+# --------------------------------------------------------------------------
+# Marginal ML under an equal-rates Mk model
+# --------------------------------------------------------------------------
+def mk_q_builder(k: int, model: str = "ER"):
+    """Return ``(build_Q, n_params, init)`` for an Mk rate matrix.
+
+    ``model``: ``"ER"`` (equal rates, 1 param), ``"SYM"`` (symmetric,
+    k(k-1)/2 params), ``"ARD"`` (all rates different, k(k-1) params).
+    """
+    import numpy as np
+    if model == "ER":
+        def build(p):
+            Q = np.full((k, k), p[0], dtype=float)
+            np.fill_diagonal(Q, 0.0)
+            for i in range(k):
+                Q[i, i] = -Q[i].sum()
+            return Q
+        return build, 1, [1.0]
+    if model == "SYM":
+        ut = [(i, j) for i in range(k) for j in range(i + 1, k)]
+
+        def build(p):
+            Q = np.zeros((k, k))
+            for v, (i, j) in zip(p, ut):
+                Q[i, j] = Q[j, i] = v
+            for i in range(k):
+                Q[i, i] = -Q[i].sum()
+            return Q
+        return build, len(ut), [1.0] * len(ut)
+    if model == "ARD":
+        pairs = [(i, j) for i in range(k) for j in range(k) if i != j]
+
+        def build(p):
+            Q = np.zeros((k, k))
+            for v, (i, j) in zip(p, pairs):
+                Q[i, j] = v
+            for i in range(k):
+                Q[i, i] = -Q[i].sum()
+            return Q
+        return build, len(pairs), [1.0] * len(pairs)
+    raise ValueError(f"unknown Mk model {model!r}; use ER/SYM/ARD")
+
+
+def ace_ml(tree: Tree, trait: Dict[str, str], model: str = "ER",
+           default_length: float = 1.0) -> Dict[Node, Dict[str, float]]:
+    """Marginal ML ancestral states under an Mk model (``ER``/``SYM``/``ARD``).
+
+    Returns ``{node: {state: posterior_prob}}`` for internal nodes and writes
+    ``data['ace_probs']`` / ``data['ace_state']`` onto every node.  Fitted
+    rates are stored on ``tree.root.data['_ace_model']``.
+    """
+    import numpy as np
+    from scipy.linalg import expm
+    from scipy.optimize import minimize, minimize_scalar
+
+    states = sorted({v for v in trait.values()})
+    k = len(states)
+    idx = {s: i for i, s in enumerate(states)}
+    pi = np.full(k, 1.0 / k)
+    nodes = tree.nodes("postorder")
+    build_Q, nparams, init = mk_q_builder(k, model)
+
+    def branch_p(params, t):
+        return expm(build_Q(params) * t)
+
+    def down_partials(params) -> Dict[Node, "np.ndarray"]:
+        D: Dict[Node, np.ndarray] = {}
+        for node in nodes:
+            if node.is_leaf:
+                vec = np.zeros(k)
+                s = trait.get(node.name)
+                if s is None:
+                    vec[:] = 1.0
+                else:
+                    vec[idx[s]] = 1.0
+                D[node] = vec
+            else:
+                vec = np.ones(k)
+                for c in node.children:
+                    vec = vec * (branch_p(params, c.length or default_length) @ D[c])
+                D[node] = vec
+        return D
+
+    def neg_log_lik(params) -> float:
+        if np.any(np.asarray(params) <= 0):
+            return 1e18
+        D = down_partials(params)
+        return -np.log(float(pi @ D[tree.root]) + 1e-300)
+
+    if nparams == 1:
+        r = float(minimize_scalar(lambda x: neg_log_lik([x]),
+                                  bounds=(1e-4, 100.0), method="bounded").x)
+        params = [r]
+    else:
+        params = list(minimize(neg_log_lik, init, method="Nelder-Mead",
+                               options={"xatol": 1e-3, "fatol": 1e-3,
+                                        "maxiter": 400}).x)
+
+    D = down_partials(params)
+    U: Dict[Node, np.ndarray] = {tree.root: pi.copy()}
+    for node in tree.nodes("preorder"):
+        if node.is_root:
+            continue
+        p = node.parent
+        msg = U[p].copy()
+        for sib in p.children:
+            if sib is node:
+                continue
+            msg = msg * (branch_p(params, sib.length or default_length) @ D[sib])
+        U[node] = branch_p(params, node.length or default_length).T @ msg
+
+    result: Dict[Node, Dict[str, float]] = {}
+    for node in tree.traverse():
+        post = D[node] * U[node]
+        total = post.sum()
+        post = post / total if total > 0 else np.full(k, 1.0 / k)
+        probs = {states[i]: float(post[i]) for i in range(k)}
+        node.data["ace_probs"] = probs
+        node.data["ace_state"] = max(probs, key=probs.get)
+        if not node.is_leaf:
+            result[node] = probs
+
+    tree.root.data["_ace_model"] = {
+        "model": model, "rates": [float(p) for p in params], "states": states,
+        "rate": float(params[0]) if nparams == 1 else None,
+    }
+    return result
+
+
+# --------------------------------------------------------------------------
+# Continuous traits -- extension point
+# --------------------------------------------------------------------------
+def ace_continuous(tree: Tree, trait: Dict[str, float]) -> Dict[Node, float]:
+    """Ancestral states for a continuous trait via Felsenstein's
+    independent-contrasts weighted averaging (Brownian motion).
+
+    Returns ``{node: estimated_value}`` and writes ``data['ace_value']``.
+    """
+    # postorder: weighted average of children, weights = 1 / (branch len + accumulated var)
+    var: Dict[Node, float] = {}
+    val: Dict[Node, float] = {}
+    for node in tree.traverse("postorder"):
+        if node.is_leaf:
+            if node.name not in trait:
+                raise KeyError(f"missing trait value for tip {node.name!r}")
+            val[node] = float(trait[node.name])
+            var[node] = 0.0
+        else:
+            ws, xs, vs = [], [], []
+            for c in node.children:
+                v = (c.length or 1.0) + var[c]
+                ws.append(1.0 / v)
+                xs.append(val[c])
+                vs.append(v)
+            wsum = sum(ws)
+            val[node] = sum(w * x for w, x in zip(ws, xs)) / wsum
+            var[node] = 1.0 / wsum
+    for node in tree.traverse():
+        node.data["ace_value"] = val[node]
+    return val

phytreon/comparative/stochastic_mapping.py

@@ -0,0 +1,134 @@
+"""Stochastic character mapping.
+
+Simulates discrete-trait histories along the branches under an Mk model,
+conditioned on the observed tip states (Nielsen 2002; Huelsenbeck 2003):
+
+1. fit the rate by ML and compute the down-pass conditional likelihoods,
+2. for each replicate, sample internal node states (preorder, conditioned)
+   then simulate the CTMC path along every branch (rejection sampling),
+3. summarise over replicates: per-node posterior state probabilities and
+   per-branch average dwell-time fractions.
+
+Results are written to ``node.data`` (``ace_probs``, ``ace_state``,
+``paint_segments``) so :meth:`TreeFigure.painted_branches` can paint the
+branches and :meth:`TreeFigure.node_pies` can chart node posteriors.
+"""
+from __future__ import annotations
+
+from typing import Dict, Optional
+
+from ..core.tree import Tree
+
+
+def stochastic_map(tree: Tree, trait: Dict[str, str], n: int = 200,
+                   model: str = "ER", rate: Optional[float] = None,
+                   seed: Optional[int] = None) -> Tree:
+    """Stochastic character map a discrete ``trait`` ({tip: state}) over ``tree``.
+
+    ``model`` is the Mk model (``ER``/``SYM``/``ARD``); ``n`` is the number of
+    simulated histories (default 200 -- raise for smoother posteriors).
+    """
+    import numpy as np
+    from scipy.linalg import expm
+    from scipy.optimize import minimize, minimize_scalar
+    from .ace import mk_q_builder
+
+    states = sorted({v for v in trait.values()})
+    k = len(states)
+    idx = {s: i for i, s in enumerate(states)}
+    pi = np.full(k, 1.0 / k)
+    post = tree.nodes("postorder")
+    build_Q, nparams, init = mk_q_builder(k, model)
+
+    def down(params):
+        Q = build_Q(params)
+        D = {}
+        for node in post:
+            if node.is_leaf:
+                v = np.zeros(k)
+                s = trait.get(node.name)
+                if s is None:
+                    v[:] = 1.0
+                else:
+                    v[idx[s]] = 1.0
+                D[id(node)] = v
+            else:
+                v = np.ones(k)
+                for c in node.children:
+                    v = v * (expm(Q * max(c.length or 0.0, 1e-6)) @ D[id(c)])
+                D[id(node)] = v
+        return D
+
+    if rate is not None:
+        params = [rate] if nparams == 1 else [rate] * nparams
+    else:
+        def negll(params):
+            if np.any(np.asarray(params) <= 0):
+                return 1e18
+            return -np.log(pi @ down(params)[id(tree.root)] + 1e-300)
+        if nparams == 1:
+            params = [float(minimize_scalar(lambda r: negll([r]),
+                                            bounds=(1e-4, 100.0),
+                                            method="bounded").x)]
+        else:
+            params = list(minimize(negll, init, method="Nelder-Mead",
+                                   options={"maxiter": 400}).x)
+
+    D = down(params)
+    Q = build_Q(params)
+    rate = float(params[0])
+    P = {id(n): expm(Q * max(n.length or 0.0, 1e-6))
+         for n in tree.traverse() if not n.is_root}
+    rng = np.random.default_rng(seed)
+
+    counts = {id(n): np.zeros(k) for n in tree.traverse()}
+    dwell = {id(n): np.zeros(k) for n in tree.traverse() if not n.is_root}
+
+    def sim_branch(a, b, t):
+        for _ in range(50):
+            segs, cur, rem = [], a, t
+            while True:
+                q = -Q[cur, cur]
+                dt = rng.exponential(1.0 / q) if q > 0 else rem
+                if dt >= rem:
+                    segs.append((cur, rem))
+                    break
+                segs.append((cur, dt))
+                rem -= dt
+                pr = Q[cur].copy()
+                pr[cur] = 0.0
+                cur = int(rng.choice(k, p=pr / pr.sum()))
+            if cur == b:
+                return segs
+        return [(a, t / 2), (b, t / 2)] if a != b else [(a, t)]
+
+    pre = tree.nodes("preorder")
+    for _ in range(n):
+        s = {}
+        pr = pi * D[id(tree.root)]
+        s[id(tree.root)] = int(rng.choice(k, p=pr / pr.sum()))
+        counts[id(tree.root)][s[id(tree.root)]] += 1
+        for node in pre:
+            if node.is_root:
+                continue
+            sp = s[id(node.parent)]
+            pr = P[id(node)][sp] * D[id(node)]
+            si = int(rng.choice(k, p=pr / pr.sum()))
+            s[id(node)] = si
+            counts[id(node)][si] += 1
+            for st, dt in sim_branch(sp, si, max(node.length or 0.0, 1e-6)):
+                dwell[id(node)][st] += dt
+
+    for node in tree.traverse():
+        c = counts[id(node)]
+        node.data["ace_probs"] = {states[i]: float(c[i] / n) for i in range(k)}
+        node.data["ace_state"] = states[int(c.argmax())]
+    for node in tree.traverse():
+        if node.is_root:
+            continue
+        d = dwell[id(node)]
+        tot = d.sum() or 1.0
+        node.data["paint_segments"] = [(states[i], float(d[i] / tot))
+                                       for i in range(k) if d[i] > 0]
+    tree.data["stochastic_map"] = {"rate": rate, "states": states, "n": n}
+    return tree

phytreon/core/__init__.py

@@ -0,0 +1,5 @@
+"""Core data model and I/O."""
+from .tree import Node, Tree
+from . import io
+
+__all__ = ["Node", "Tree", "io"]

phytreon/core/io.py

@@ -0,0 +1,177 @@
+"""Tree input/output.
+
+We reuse Biopython's robust parsers/writers for the standard formats
+(Newick, Nexus, PhyloXML, NeXML) and convert to/from our :class:`Tree`.
+A tiny self-contained Newick parser is also provided so that the package
+has a zero-config fast path for the common case (and so the data model
+can be used even if Biopython is unavailable).
+"""
+from __future__ import annotations
+
+import io as _io
+import re
+from typing import Optional
+
+from .tree import Node, Tree
+
+_BIO_FORMATS = {"newick", "nexus", "phyloxml", "nexml", "cdao"}
+
+
+# --------------------------------------------------------------------------
+# Biopython bridge
+# --------------------------------------------------------------------------
+def from_biopython(bp_tree) -> Tree:
+    """Convert a ``Bio.Phylo`` tree into a phytreon :class:`Tree`."""
+    def convert(clade) -> Node:
+        node = Node(
+            name=clade.name,
+            length=clade.branch_length,
+            support=getattr(clade, "confidence", None),
+            comment=getattr(clade, "comment", None),
+        )
+        for child in clade.clades:
+            node.add_child(convert(child))
+        return node
+
+    root = convert(bp_tree.root)
+    return Tree(root=root, name=getattr(bp_tree, "name", None))
+
+
+def to_biopython(tree: Tree):
+    from Bio.Phylo.BaseTree import Clade, Tree as BioTree
+
+    def convert(node: Node) -> Clade:
+        clade = Clade(
+            branch_length=node.length,
+            name=node.name,
+            confidence=node.support,
+        )
+        clade.clades = [convert(c) for c in node.children]
+        return clade
+
+    return BioTree(root=convert(tree.root), name=tree.name)
+
+
+def read(path: str, fmt: str = "newick") -> Tree:
+    if fmt not in _BIO_FORMATS:
+        raise ValueError(f"unsupported format {fmt!r}; choose from {_BIO_FORMATS}")
+    from Bio import Phylo
+    return from_biopython(Phylo.read(path, fmt))
+
+
+def write(tree: Tree, path: Optional[str] = None, fmt: str = "newick") -> Optional[str]:
+    """Write to ``path``; if ``path`` is ``None`` return the string."""
+    if fmt == "newick" and path is None:
+        return to_newick(tree)
+    from Bio import Phylo
+    bp = to_biopython(tree)
+    if path is None:
+        buf = _io.StringIO()
+        Phylo.write(bp, buf, fmt)
+        return buf.getvalue()
+    Phylo.write(bp, path, fmt)
+    return None
+
+
+# --------------------------------------------------------------------------
+# Lightweight self-contained Newick parser / writer
+# --------------------------------------------------------------------------
+_TOKEN = re.compile(r"\s*([(),;])\s*|\s*([^(),;]+)\s*")
+
+
+def parse_newick(newick: str) -> Tree:
+    """Parse a Newick string into a :class:`Tree` (no external deps).
+
+    Supports branch lengths (``name:0.1``), internal labels / support
+    values, and quoted names.  Square-bracket comments (e.g. NHX) are
+    captured verbatim onto ``node.comment``.
+    """
+    s = newick.strip()
+    if not s.endswith(";"):
+        s += ";"
+
+    pos = 0
+    n = len(s)
+
+    def parse_clade() -> Node:
+        nonlocal pos
+        node = Node()
+        if s[pos] == "(":
+            pos += 1  # consume '('
+            while True:
+                node.add_child(parse_clade())
+                if s[pos] == ",":
+                    pos += 1
+                    continue
+                if s[pos] == ")":
+                    pos += 1
+                    break
+        _parse_label(node)
+        return node
+
+    def _parse_label(node: Node) -> None:
+        nonlocal pos
+        # optional name/support
+        start = pos
+        while pos < n and s[pos] not in "():,;[":
+            pos += 1
+        label = s[start:pos].strip().strip("'\"")
+        if label:
+            if node.is_leaf:
+                node.name = label
+            else:
+                # internal label is usually a support value
+                try:
+                    node.support = float(label)
+                except ValueError:
+                    node.name = label
+        # optional comment [...]
+        if pos < n and s[pos] == "[":
+            depth, cstart = 0, pos + 1
+            while pos < n:
+                if s[pos] == "[":
+                    depth += 1
+                elif s[pos] == "]":
+                    depth -= 1
+                    if depth == 0:
+                        node.comment = s[cstart:pos]
+                        pos += 1
+                        break
+                pos += 1
+        # optional branch length
+        if pos < n and s[pos] == ":":
+            pos += 1
+            start = pos
+            while pos < n and s[pos] not in "():,;[":
+                pos += 1
+            try:
+                node.length = float(s[start:pos])
+            except ValueError:
+                pass
+
+    root = parse_clade()
+    return Tree(root=root)
+
+
+def to_newick(tree: Tree, with_support: bool = True) -> str:
+    """Serialise a :class:`Tree` back to a Newick string."""
+    def fmt(node: Node) -> str:
+        if node.is_leaf:
+            label = node.name or ""
+        else:
+            inner = ",".join(fmt(c) for c in node.children)
+            sup = ""
+            if with_support and node.support is not None:
+                sup = _num(node.support)
+            elif node.name:
+                sup = node.name
+            label = f"({inner}){sup}"
+        if node.length is not None:
+            label += f":{_num(node.length)}"
+        return label
+
+    return fmt(tree.root) + ";"
+
+
+def _num(x: float) -> str:
+    return f"{x:g}"