forestpd 0.1.0__py3-none-any.whl

forestpd/__init__.py

@@ -0,0 +1,65 @@
+"""
+forestpd — maximizing phylogenetic diversity under phylogenetic and extinction uncertainty.
+
+Solves MaxMPD: given unrooted weighted trees on a shared taxon set, per-tree weights ``w_j``,
+integer taxon costs and a budget ``B``, choose ``S`` with ``c(S) <= B`` maximizing
+``sum_j w_j PD_{T_j}(S)`` — or, with per-taxon survival probabilities, its expected value over the
+taxa that survive. Given a maximum agreement forest (MAF) of size ``k``, it solves MaxMPD in
+``O(2^k * n * B^2)`` time, where ``n`` is the total size of the trees. ``k`` is taken per common
+cluster, so it measures how much the trees disagree locally.
+
+The public entry points are :class:`MaxMPDInstance` (build it once, then call
+:meth:`~MaxMPDInstance.optimize`, :meth:`~MaxMPDInstance.optimize_many` and
+:meth:`~MaxMPDInstance.evaluate_pd`) and the one-call functions :func:`optimize_maxmpd`,
+:func:`optimize_many_maxmpd` and :func:`evaluate_maxmpd`. The pipeline has three stages —
+:mod:`forestpd.cluster`, :mod:`forestpd.maf`, :mod:`forestpd.algo` — described in ``README.md``.
+"""
+
+from __future__ import annotations
+
+__version__ = "0.1.0"
+
+from .api import evaluate_maxmpd, optimize_many_maxmpd, optimize_maxmpd
+from .core import PhyloTree, TaxonRegistry
+from .exceptions import (
+    ForestPDError,
+    ForestPDImportError,
+    ForestPDInputError,
+    ForestPDValueError,
+    ForestPDWarning,
+)
+from .instance import MaxMPDInstance, MaxMPDResult, MaxMPDStructure
+from . import io, utils
+from .algo.convolve import HAVE_NUMBA as _HAVE_NUMBA
+
+if not _HAVE_NUMBA:  # pragma: no cover - depends on the environment
+    import warnings as _warnings
+
+    _warnings.warn(
+        "Numba could not be imported, so forestpd runs its NumPy/Python fallbacks: the same "
+        "results, but several times slower. Install numba (a dependency of forestpd) for full speed.",
+        ForestPDWarning,
+        stacklevel=2,
+    )
+
+__all__ = [
+    "__version__",
+    # Public API
+    "MaxMPDInstance",
+    "optimize_maxmpd",
+    "optimize_many_maxmpd",
+    "evaluate_maxmpd",
+    "MaxMPDResult",
+    "MaxMPDStructure",
+    "io",
+    "utils",
+    # Core types
+    "PhyloTree",
+    "TaxonRegistry",
+    # Exceptions
+    "ForestPDError",
+    "ForestPDValueError",
+    "ForestPDInputError",
+    "ForestPDImportError",
+    "ForestPDWarning",
+]

forestpd/_evaluate.py

@@ -0,0 +1,113 @@
+"""
+Fast evaluation of the objective of a taxon set: the trees as one weighted split table.
+
+The objective of a protected set ``S`` is a sum over all edges of all trees of
+``w_j * len_j(e) * P(one side keeps a survivor) * P(the other side does)``, where a side holding a
+taxon of ``S`` keeps one surely and otherwise with probability ``1 - prod (1 - a(x))`` over its
+taxa. Both probabilities depend on the edge only through its split, so the edges are grouped by
+split once (:func:`forestpd.io.splits.split_system`, compiled) and each distinct split carries its
+summed weight ``W_s``. Each split is stored by its smaller side, as a sparse row of taxon ids.
+
+Scoring a set is then two sparse sums per split, independent of the number of trees: the protected
+taxa on the stored side (the other side has the rest), and the summed ``log(1 - a(x))`` of its
+unprotected taxa (the other side has the rest of the total). This regroups the per-edge sum of
+the per-edge definition exactly; without survival probabilities it is the split
+diversity of ``S``, i.e. its weighted PD.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+import numpy as np
+
+from .core.tree import PhyloTree
+from .io import _splits_kernel
+from .io import splits as _splits
+
+
+class SplitTable:
+    """
+    The weighted splits of a set of trees, ready to score taxon sets.
+
+    Parameters
+    ----------
+    trees : Sequence[PhyloTree]
+        The trees, all on leaves ``0..n_taxa-1``.
+    weights : numpy.ndarray
+        ``float64[t]``; the per-tree weights.
+    log_loss : numpy.ndarray
+        ``float64[n_taxa]``; ``log(1 - a(x))`` per taxon.
+
+    Examples
+    --------
+    >>> from forestpd.core.tree import PhyloTree
+    >>> t = PhyloTree.from_edges(6, [(0, 4), (1, 4), (4, 5), (2, 5), (3, 5)], [1.0] * 5,
+    ...                          {i: i for i in range(4)})
+    >>> table = SplitTable([t, t], np.array([1.0, 0.5]), np.zeros(4))
+    >>> table.value([0, 2])          # the path 0 - 4 - 5 - 2, in both trees: 3 + 1.5
+    4.5
+    """
+
+    def __init__(self, trees: Sequence[PhyloTree], weights: np.ndarray, log_loss: np.ndarray):
+        n_taxa = trees[0].n_leaves
+        members, offsets, split_weights = _splits._flat_split_system(trees, weights)
+        if _splits.USE_COMPILED and _splits_kernel.HAVE_NUMBA:
+            members, offsets = _splits_kernel.smaller_sides(members, offsets, n_taxa)
+        else:
+            members, offsets = _smaller_sides(members, offsets, n_taxa)
+        self.n_taxa = n_taxa
+        self.members = members
+        self.weights = split_weights
+        self.log_loss = np.asarray(log_loss, dtype=np.float64)
+        # Rows are nonempty (every side holds a taxon), so reduceat over these starts is exact.
+        self.starts = offsets[:-1]
+
+    def value(self, protected: Sequence[int]) -> float:
+        """
+        The objective of protecting the taxa ``protected`` (taxon ids).
+
+        Parameters
+        ----------
+        protected : Sequence[int]
+            Distinct taxon ids.
+
+        Returns
+        -------
+        float
+            ``sum_s W_s P(side A keeps a survivor) P(side B does)``.
+        """
+        if self.weights.size == 0:
+            return 0.0
+        is_protected = np.zeros(self.n_taxa, dtype=bool)
+        is_protected[np.asarray(list(protected), dtype=np.int64)] = True
+        n_protected = int(is_protected.sum())
+        unprotected_loss = np.where(is_protected, 0.0, self.log_loss)
+        total_loss = float(unprotected_loss.sum())
+
+        on_side = is_protected[self.members]
+        protected_a = np.add.reduceat(on_side.astype(np.int64), self.starts)
+        loss_a = np.add.reduceat(unprotected_loss[self.members], self.starts)
+        protected_b = n_protected - protected_a
+        loss_b = total_loss - loss_a
+        keeps_a = np.where(protected_a > 0, 1.0, 0.0 - np.expm1(loss_a))
+        keeps_b = np.where(protected_b > 0, 1.0, 0.0 - np.expm1(loss_b))
+        return float(np.dot(self.weights, keeps_a * keeps_b))
+
+
+def _smaller_sides(
+    members: np.ndarray, offsets: np.ndarray, n_taxa: int
+) -> tuple[np.ndarray, np.ndarray]:
+    """The NumPy version of :func:`forestpd.io._splits_kernel.smaller_sides`."""
+    sides = []
+    everyone = np.arange(n_taxa)
+    for r in range(offsets.size - 1):
+        side = members[offsets[r] : offsets[r + 1]]
+        sides.append(side if 2 * side.size <= n_taxa else np.setdiff1d(everyone, side))
+    out_offsets = np.zeros(len(sides) + 1, dtype=np.int64)
+    np.cumsum([side.size for side in sides], out=out_offsets[1:])
+    flat = np.concatenate(sides) if sides else np.empty(0, dtype=np.int64)
+    return flat.astype(np.int64), out_offsets
+
+
+__all__ = ["SplitTable"]

forestpd/_validate.py

@@ -0,0 +1,278 @@
+"""
+Validation of the instance data: per-tree weights, taxon costs, survival probabilities and budgets.
+
+Shared by :class:`~forestpd.instance.MaxMPDInstance`, :mod:`forestpd.utils` and
+:mod:`forestpd.io.splits`; every function raises :class:`~forestpd.exceptions.ForestPDInputError`
+on invalid input.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Iterable, Mapping, Sequence
+
+import numpy as np
+
+from .core.taxa import TaxonRegistry
+from .core.tree import PhyloTree, has_degree_two
+from .exceptions import ForestPDInputError
+from .io.newick import ROOT_LABEL
+
+
+def weights_to_array(
+    weights: Sequence[float] | Mapping[int, float] | None, n_trees: int
+) -> np.ndarray:
+    """
+    Validate per-tree weights ``w_j`` and return them as an array (no rescaling).
+
+    Parameters
+    ----------
+    weights : Sequence[float] | Mapping[int, float] | None
+        Per-tree weights, either aligned with the trees or keyed by tree index (missing indices
+        weigh ``1.0``). ``None`` means all ``1.0``.
+    n_trees : int
+        Number of trees ``t``.
+
+    Returns
+    -------
+    numpy.ndarray
+        ``float64[t]`` array of weights.
+
+    Raises
+    ------
+    ForestPDInputError
+        If a weight is negative or non-finite, a mapping key is not a tree index, or the sequence
+        length is wrong.
+    """
+    if weights is None:
+        return np.ones(n_trees, dtype=np.float64)
+    if isinstance(weights, Mapping):
+        for key in weights:
+            if (
+                isinstance(key, (bool, np.bool_))
+                or not isinstance(key, (int, np.integer))
+                or not 0 <= key < n_trees
+            ):
+                raise ForestPDInputError(
+                    f"Weight given for tree {key!r}, but trees are indexed 0..{n_trees - 1}"
+                )
+        values = [float(weights.get(j, 1.0)) for j in range(n_trees)]
+    else:
+        values = [float(w) for w in weights]
+        if len(values) != n_trees:
+            raise ForestPDInputError(f"Expected {n_trees} weights, got {len(values)}")
+    for j, w in enumerate(values):
+        if not math.isfinite(w) or w < 0:
+            raise ForestPDInputError(
+                f"Weight for tree {j} must be finite and non-negative, got {w}"
+            )
+    return np.asarray(values, dtype=np.float64)
+
+
+def costs_to_array(
+    registry: TaxonRegistry, costs: Mapping[str, int] | Sequence[int] | None
+) -> np.ndarray:
+    """
+    Validate taxon costs and return them as an array indexed by taxon id.
+
+    Parameters
+    ----------
+    registry : TaxonRegistry
+        The taxon registry fixing the id of each label.
+    costs : Mapping[str, int] | Sequence[int] | None
+        Integer cost per taxon label (missing labels cost 1), or a sequence of costs indexed by
+        taxon id. ``None`` means unit costs.
+
+    Returns
+    -------
+    numpy.ndarray
+        ``int64[n]`` array of costs, indexed by taxon id.
+
+    Raises
+    ------
+    ForestPDInputError
+        If a cost is given for an unknown label, is not a positive integer, or a sequence
+        has the wrong length.
+    """
+    result = np.ones(registry.n, dtype=np.int64)
+    if costs is None:
+        return result
+    if isinstance(costs, Mapping):
+        items = [
+            (label, registry.id_of(label) if _user_label(label, registry) else -1, cost)
+            for label, cost in costs.items()
+        ]
+    else:
+        values = list(costs)
+        if len(values) != registry.n:
+            raise ForestPDInputError(f"Expected {registry.n} costs, got {len(values)}")
+        items = [(registry.label_of(i), i, cost) for i, cost in enumerate(values)]
+    for label, taxon, cost in items:
+        if taxon < 0:
+            raise ForestPDInputError(f"Cost given for unknown taxon: {label!r}")
+        if isinstance(cost, (bool, np.bool_)) or not isinstance(cost, (int, np.integer)):
+            raise ForestPDInputError(
+                f"Cost for taxon {label!r} must be an integer, got {type(cost).__name__}"
+            )
+        if cost < 1:
+            raise ForestPDInputError(f"Cost for taxon {label!r} must be at least 1, got {cost}")
+        result[taxon] = cost
+    return result
+
+
+def survival_to_array(
+    registry: TaxonRegistry, survival: Mapping[str, float] | Sequence[float] | None
+) -> np.ndarray:
+    """
+    Validate survival probabilities and return them as an array indexed by taxon id.
+
+    Parameters
+    ----------
+    registry : TaxonRegistry
+        The taxon registry fixing the id of each label.
+    survival : Mapping[str, float] | Sequence[float] | None
+        Probability per taxon label (missing labels have ``0``), or a sequence indexed by taxon
+        id. ``None`` means ``0`` for every taxon.
+
+    Returns
+    -------
+    numpy.ndarray
+        ``float64[n]`` array of probabilities, indexed by taxon id.
+
+    Raises
+    ------
+    ForestPDInputError
+        If a probability is given for an unknown label, is not a number in ``[0, 1)``, or a
+        sequence has the wrong length.
+    """
+    result = np.zeros(registry.n, dtype=np.float64)
+    if survival is None:
+        return result
+    if isinstance(survival, Mapping):
+        items = [
+            (label, registry.id_of(label) if _user_label(label, registry) else -1, value)
+            for label, value in survival.items()
+        ]
+    else:
+        values = list(survival)
+        if len(values) != registry.n:
+            raise ForestPDInputError(
+                f"Expected {registry.n} survival probabilities, got {len(values)}"
+            )
+        items = [(registry.label_of(i), i, value) for i, value in enumerate(values)]
+    for label, taxon, value in items:
+        if taxon < 0:
+            raise ForestPDInputError(f"Survival probability given for unknown taxon: {label!r}")
+        if isinstance(value, (bool, np.bool_)) or not isinstance(
+            value, (int, float, np.integer, np.floating)
+        ):
+            raise ForestPDInputError(
+                f"Survival probability for taxon {label!r} must be a number, "
+                f"got {type(value).__name__}"
+            )
+        if not (0.0 <= float(value) < 1.0):
+            raise ForestPDInputError(
+                f"Survival probability for taxon {label!r} must lie in [0, 1), got {value}"
+            )
+        result[taxon] = float(value)
+    return result
+
+
+def validate_budget(budget: int) -> None:
+    """
+    Raise :class:`ForestPDInputError` unless ``budget`` is a non-negative integer.
+
+    Parameters
+    ----------
+    budget : int
+        The total budget ``B``.
+
+    Raises
+    ------
+    ForestPDInputError
+        If ``budget`` is not an integer (Python or NumPy; ``bool`` is rejected) or is negative.
+    """
+    if isinstance(budget, (bool, np.bool_)) or not isinstance(budget, (int, np.integer)):
+        raise ForestPDInputError(f"budget must be an integer, got {type(budget).__name__}")
+    if budget < 0:
+        raise ForestPDInputError(f"budget must be non-negative, got {budget}")
+
+
+def validate_uniform(trees: Sequence[PhyloTree], registry: TaxonRegistry) -> None:
+    """
+    Check that every tree has the same leaf set, namely all taxa in ``registry``, and no
+    internal vertex of degree 2 (trees read from Newick or PhyloZoo never have one).
+
+    Parameters
+    ----------
+    trees : Sequence[PhyloTree]
+        The input trees.
+    registry : TaxonRegistry
+        The registry of all taxa in the instance.
+
+    Raises
+    ------
+    ForestPDInputError
+        If the trees do not all span the same taxon set, or one has an internal vertex of
+        degree 2.
+    """
+    for j, tree in enumerate(trees):
+        if tree.n_leaves != registry.n:
+            raise ForestPDInputError(
+                f"Tree {j} has {tree.n_leaves} taxa but the instance has {registry.n}. "
+                "All trees must carry the same taxon set."
+            )
+        if has_degree_two(tree):
+            raise ForestPDInputError(
+                f"Tree {j} has an internal vertex of degree 2; suppress it first with "
+                "forestpd.core.tree.suppress_degree_two (this does not change any PD)."
+            )
+
+
+def taxon_ids(registry: TaxonRegistry, taxa: Iterable[str]) -> set[int]:
+    """
+    The ids of a set of taxon labels given by a user.
+
+    Parameters
+    ----------
+    registry : TaxonRegistry
+        The taxon registry.
+    taxa : Iterable[str]
+        Taxon labels; a single string is rejected rather than read as its characters.
+
+    Returns
+    -------
+    set[int]
+        Their ids.
+
+    Raises
+    ------
+    ForestPDInputError
+        If ``taxa`` is a string, or a label is unknown (the reserved root label included).
+    """
+    if isinstance(taxa, str):
+        raise ForestPDInputError(
+            f"Expected a collection of taxon labels, got the string {taxa!r}; "
+            f"wrap a single label as [{taxa!r}]."
+        )
+    ids = set()
+    for label in taxa:
+        if not _user_label(label, registry):
+            raise ForestPDInputError(f"Unknown taxon: {label!r}")
+        ids.add(registry.id_of(label))
+    return ids
+
+
+def _user_label(label: object, registry: TaxonRegistry) -> bool:
+    """Whether ``label`` names a taxon a user may refer to (not the reserved root leaf)."""
+    return label != ROOT_LABEL and label in registry
+
+
+__all__ = [
+    "taxon_ids",
+    "weights_to_array",
+    "costs_to_array",
+    "survival_to_array",
+    "validate_budget",
+    "validate_uniform",
+]

forestpd/algo/__init__.py

@@ -0,0 +1,45 @@
+"""
+Stage 3 of the pipeline: the phylogenetic-diversity dynamic program.
+
+Modules, in the order the algorithm uses them:
+
+===========================  =============================================================
+:mod:`.local`                ``LocalInstance`` and the ``f_2`` tables that link clusters
+:mod:`.profile`              per-cluster precomputation, done once before the ``2^k`` loop
+:mod:`.survival`             survival probabilities: side constants and the per-cluster sums
+:mod:`.convolve`             max-plus convolution (compiled with Numba)
+:mod:`.component_dp`         ``PD_i^{(S)}``, Lemma "compute-PD_i"
+:mod:`.local_solver`         Algorithm 3: one cluster, all signatures
+:mod:`.solver`               Algorithm 1: bottom-up pass over the tree of clusters
+:mod:`.witness`              reconstruction of an optimal taxon set
+===========================  =============================================================
+"""
+
+from __future__ import annotations
+
+from .component_dp import component_dp_table, component_table
+from .convolve import HAVE_NUMBA, fold_max_plus, max_plus
+from .local import LocalInstance, NEG_INF, ordinary_f2, pseudo_leaf_f2
+from .local_solver import LocalSolution, signature_weights, solve_local
+from .profile import ClusterProfile, ComponentTree, build_profile
+from .solver import Solution, solve
+
+__all__ = [
+    "LocalInstance",
+    "NEG_INF",
+    "ordinary_f2",
+    "pseudo_leaf_f2",
+    "ClusterProfile",
+    "ComponentTree",
+    "build_profile",
+    "max_plus",
+    "fold_max_plus",
+    "HAVE_NUMBA",
+    "component_table",
+    "component_dp_table",
+    "LocalSolution",
+    "solve_local",
+    "signature_weights",
+    "solve",
+    "Solution",
+]