zombi2.traits

Level 4: a trait evolving on the species tree, alongside the genome chain rather than inside it. Two functions, because a continuous trait and a discrete one take genuinely different arguments. A third, simulate_traits, is for several traits that depend on each other — the case with no order to grow them in.

zombi2.traits.simulate_continuous

simulate_continuous(tree, *, start=0.0, rate=1.0, reverts_to=None, pull=None, correlation=None, at_speciation=None, regimes=None, seed=None, progress=False) -> TraitsResult

Evolve a continuous trait down a tree and return a TraitsResult. One process, its variants selected by knobs (SPEC §4): Brownian motion (bare rate), Ornstein–Uhlenbeck (add reverts_to + pull), early burst (a changing_at schedule on rate), and variable-rates BM (a Drift law on rate).

Correlated traits ride together in one call (the joint rule inside a level): pass start and rate as dicts keyed by trait name and a correlation={(a, b): ρ} overlay (each ρ ∈ [−1, 1]). The traits then diffuse jointly — the branch increment is drawn from MVN(0, Σ·dt) with Σ = D R D (D = diag(σ_i), R the correlation matrix), so at a tip the correlation between two traits is exactly their ρ. Add reverts_to and pull — one value shared, or a dict of one per trait — and it is multivariate Ornstein–Uhlenbeck in its diagonal-drift restriction: each trait reverts to its own optimum at its own strength, the correlation stays in the diffusion, and the branch covariance is Σ_ij·(1 − e^{−(α_i+α_j)·dt})/(α_i + α_j). One trait's deviation pulling another — a full drift matrix — is a different model and is refused by name, not read as a diagonal. A correlated run takes bare per-trait rates. Its log is widened rather than absent: a value is a per-trait vector, so trait_events.tsv gets one from/to column pair per trait.

tree is the complete species tree (a Tree, or a SpeciesResult whose complete_tree is used). The trait evolves on every lineage, extant and extinct alike, so the ancestral states are exact and complete; the observed dataset is the extant tips, result.values.

start is the value at t = 0 (the origin, root.birth_time): the root lineage diffuses over its own branch [0, first split] like any other, so a trait and a genome evolve over the same branch set, and each node's stored value is the trait at that node's end_time (node_values[root] is the value at the first split, not start).

rate is the variance-rate σ² (a scope(base) with verbs chained onto it), per lineage: each lineage diffuses independently at σ², never pooled across the tree. A bare number is Brownian motion (Normal(0, σ²·dt) over a branch); changing_at({…}) makes σ² change through time — early burst / ACDC — with the per-branch variance the exact integral ∫ σ²(t) dt; varying_among('lineages', Drift(LogNormal(0.0, …))) makes σ² drift branch-to-branch — variable-rates BM ("ClaDS for traits") — each lineage inheriting its parent's σ² times a lognormal kick drawn at the split; scaled_by(TotalDiversity(cap=…)) makes σ² slow as the clade fills up — diversity-dependent / ecological-limits trait evolution — σ² scaled by (1 − standing_diversity/cap) as the tree's lineages-through-time grows (the tree is a fixed input the trait reads); scaled_by(driver, {…}) makes σ² read another level — the driver grown first on this same tree and handed over as its result object or its written trait_events.tsv, so a lineage diffuses faster while the driver is in one state than another. A discrete driver switches mid-branch, and the per-branch variance is the integral across those pieces, so a branch that spends half its length in the fast state accrues exactly half the fast variance.

reverts_to (the optimum θ) and pull (the strength α > 0) turn the diffusion into Ornstein–Uhlenbeck — the value is pulled toward θ while it diffuses, the exact per-branch transition being Normal(θ + (x−θ)·e^{−α·dt}, σ²/(2α)·(1−e^{−2α·dt})). Give both or neither. The optimum and the pull compose with the σ² modifiers: a trait that bursts early and reverts to an optimum is one rate with one modifier and two arguments. A σ² that moves along the branch leaves the mean untouched (it never read σ²) and makes the variance the exact pull-weighted integral ∫ e^{−2α(t₁−s)}·σ²(s) ds, stepping where the schedule, the standing diversity or the driver steps. That weight is the whole difference from Brownian motion's ∫ σ²(s) ds: under OU, variance accrued early has been pulled back toward θ by the time the branch ends, so the two integrals differ by an order of magnitude on a typical branch.

at_speciation adds an on-speciation jump — Normal(0, at_speciation) on each daughter at every speciation (the punctuational mode), layered on top of the along-branch anagenesis. Under correlation= it takes one variance per trait (or one shared) and the jump is drawn under the same overlay the diffusion uses. regimes gives multi-optimum OU: pass a discrete TraitsResult (a stochastic map painted by simulate_discrete() on this same tree) and a per-regime reverts_to={regime: θ}, and the value follows OU toward whichever regime's optimum a branch is in; it takes at_speciation too, one jump variance shared across regimes, and it takes a bare σ² — a modified variance-rate with regimes is not implemented yet. Deterministic given seed.

Source code in zombi2/traits/continuous.py

def simulate_continuous(tree, *, start=0.0, rate=1.0, reverts_to=None, pull=None,
                        correlation=None, at_speciation=None, regimes=None, seed=None,
                        progress=False) -> TraitsResult:
    """Evolve a continuous trait down a tree and return a `TraitsResult`. One process, its
    variants selected by knobs (SPEC §4): **Brownian motion** (bare ``rate``), **Ornstein–Uhlenbeck**
    (add ``reverts_to`` + ``pull``), **early burst** (a ``changing_at`` schedule on ``rate``), and
    **variable-rates BM** (a ``Drift`` law on ``rate``).

    **Correlated traits** ride together in **one call** (the joint rule inside a level): pass
    ``start`` and ``rate`` as dicts keyed by trait name and a ``correlation={(a, b): ρ}`` overlay
    (each ρ ∈ [−1, 1]). The traits then diffuse jointly — the branch increment is drawn from
    ``MVN(0, Σ·dt)`` with ``Σ = D R D`` (``D = diag(σ_i)``, ``R`` the correlation matrix), so at a tip
    the correlation between two traits is exactly their ρ. Add ``reverts_to`` and ``pull`` — one
    value shared, or a dict of one per trait — and it is **multivariate Ornstein–Uhlenbeck in its
    diagonal-drift restriction**: each trait reverts to its own optimum at its own strength, the
    correlation stays in the diffusion, and the branch covariance is
    ``Σ_ij·(1 − e^{−(α_i+α_j)·dt})/(α_i + α_j)``. One trait's deviation pulling *another* — a full
    drift matrix — is a different model and is refused by name, not read as a diagonal. A correlated
    run takes bare per-trait rates. Its log is **widened** rather than absent: a value is a per-trait
    vector, so ``trait_events.tsv`` gets one ``from``/``to`` column pair per trait.

    ``tree`` is the **complete** species tree (a `Tree`, or a
    `SpeciesResult` whose ``complete_tree`` is used). The trait evolves on
    **every** lineage, extant and extinct alike, so the ancestral states are exact and complete; the
    observed dataset is the extant tips, ``result.values``.

    ``start`` is the value at ``t = 0`` (the origin, ``root.birth_time``): the root lineage
    diffuses over its own branch ``[0, first split]`` like any other, so a trait and a genome evolve
    over the **same** branch set, and each node's stored value is the trait at that node's
    ``end_time`` (``node_values[root]`` is the value at the first split, not ``start``).

    ``rate`` is the variance-rate σ² (a ``scope(base)`` with verbs chained onto it), *per lineage*:
    each lineage diffuses independently at σ², never pooled across the tree. A bare number is
    Brownian motion (``Normal(0, σ²·dt)`` over a branch); ``changing_at({…})`` makes σ² change
    through time — early burst / ACDC — with the per-branch variance the exact integral
    ``∫ σ²(t) dt``;
    ``varying_among('lineages', Drift(LogNormal(0.0, …)))`` makes σ² **drift branch-to-branch** — variable-rates BM ("ClaDS
    for traits") — each lineage inheriting its parent's σ² times a lognormal kick drawn at the split;
    ``scaled_by(TotalDiversity(cap=…))`` makes σ² **slow as the clade fills up** — diversity-dependent /
    ecological-limits trait evolution — σ² scaled by ``(1 − standing_diversity/cap)`` as the tree's
    lineages-through-time grows (the tree is a fixed input the trait reads);
    ``scaled_by(driver, {…})`` makes σ² **read another level** — the driver grown first on this
    same tree and handed over as its result object or its written ``trait_events.tsv``, so a lineage
    diffuses faster while the driver is in one state than another. A discrete driver switches
    *mid-branch*, and the per-branch variance is the integral across those pieces, so a branch that
    spends half its length in the fast state accrues exactly half the fast variance.

    ``reverts_to`` (the optimum θ) and ``pull`` (the strength α > 0) turn the diffusion into
    Ornstein–Uhlenbeck — the value is pulled toward θ while it diffuses, the exact per-branch
    transition being ``Normal(θ + (x−θ)·e^{−α·dt}, σ²/(2α)·(1−e^{−2α·dt}))``. Give **both** or
    neither. The optimum and the pull **compose with the σ² modifiers**: a trait that bursts early
    and reverts to an optimum is one rate with one modifier and two arguments. A σ² that moves along
    the branch leaves the mean untouched (it never read σ²) and makes the variance the exact
    pull-weighted integral ``∫ e^{−2α(t₁−s)}·σ²(s) ds``, stepping where the schedule, the standing
    diversity or the driver steps. That weight is the whole difference from Brownian motion's
    ``∫ σ²(s) ds``: under OU, variance accrued early has been pulled back toward θ by the time the
    branch ends, so the two integrals differ by an order of magnitude on a typical branch.

    ``at_speciation`` adds an **on-speciation** jump — ``Normal(0, at_speciation)`` on each daughter at
    every speciation (the punctuational mode), layered on top of the along-branch anagenesis. Under
    ``correlation=`` it takes one variance per trait (or one shared) and the jump is drawn under the
    same overlay the diffusion uses.
    ``regimes`` gives **multi-optimum OU**: pass a discrete `TraitsResult` (a stochastic map
    painted by `simulate_discrete()` on this same tree) and a per-regime ``reverts_to={regime: θ}``,
    and the value follows OU toward whichever regime's optimum a branch is in; it takes
    ``at_speciation`` too, one jump variance shared across regimes, and it takes a bare σ² — a
    modified variance-rate with ``regimes`` is not implemented yet. Deterministic given ``seed``.
    """
    tree = as_tree(tree, level="traits")
    if regimes is not None:
        if correlation is not None:
            # `regimes` dispatches before the correlated engine and threads no correlation, so a
            # correlation passed here would be read by nothing and the run would silently be the
            # uncorrelated model (SPEC §5: refuse, never ignore). The obstacle is the same one that
            # keeps a modified σ² out of `regimes`: the branch covariance under multi-optimum OU is
            # ∫ρ_ij·σ_i·σ_j weighted by each regime's own pull, which this engine does not integrate.
            raise ValueError(
                "correlation= with regimes= is not implemented yet: multi-optimum OU evolves one "
                "trait, so there is no second trait for a correlation to be with. Use correlation= "
                "on its own (correlated BM/OU), or regimes= on its own (multi-optimum OU).")
        return _simulate_regimes(tree, start, rate, reverts_to, pull, regimes, at_speciation, seed,
                                 progress)
    if isinstance(start, dict) or isinstance(rate, dict) or correlation is not None:
        return _simulate_correlated(tree, start, rate, reverts_to, pull, correlation,
                                    at_speciation, seed, progress)
    if isinstance(start, bool) or not isinstance(start, (int, float)) or not math.isfinite(start):
        raise ValueError(f"start must be a finite number, got {start!r}")
    r = as_rate(rate, default_scope=PerLineage)
    assert r.scope is not None              # `as_rate` filled the level's default just above
    if r.scope is not PerLineage:
        raise ValueError(
            f"rate has a {r.scope.__name__} scope, but a continuous trait's variance-rate is "
            f"per lineage — write PerLineage(...), or a bare number, which is per lineage here."
        )
    # A schedule (early burst), a Drift among lineages (variable-rates BM), TotalDiversity
    # (diversity-dependent) and a driver (σ² read off another level) are the σ² modifiers this engine
    # supports; anything else is rejected loudly — the genome engine's discipline.
    for m in r.modifiers:
        if m.reads == (DRAWN, "families"):
            # not a missing feature: there is nothing here for it to mean
            raise ValueError(
                "rate varies among families, but a trait has no gene families — "
                "varying_among('families', ...) belongs on a genomes rate. For per-lineage "
                "heterogeneity here use varying_among('lineages', Drift(LogNormal(0.0, 0.3))) "
                "(variable-rates BM)."
            )
        if not is_implemented(m, IMPLEMENTED_MODIFIERS, "traits.continuous"):
            raise ValueError(
                f"rate carries {describe(m)}, which the continuous trait engine does not "
                f"support. It takes changing_at (early burst), "
                f"varying_among('lineages', Drift(...)) (variable-rates BM), "
                f"scaled_by(TotalDiversity(cap=...)) (diversity-dependent), "
                f"scaled_by(driver, mapping) (driven by another level), and "
                f"set_by(driver, mapping) (the driver supplies σ² itself, in σ²'s own units)."
            )
        if isinstance(m, Driven):
            check_not_a_kernel(m.mapping, label="rate")
    # the per-lineage modifiers σ² carries (variable-rates BM), asked for the same way every level
    # asks — and every one of them is kept, so two compose rather than the second going quietly.
    drift = tuple(m for m, _ in r.carried_modifiers(unit="lineages"))
    check_one_memory(drift, label="rate", unit="lineages")
    r.check_one_base("rate")
    has_diversity = any(isinstance(m, OnTotalDiversity) for m in r.modifiers)  # σ² reads the standing LTT

    # OU: reverts_to (θ) + pull (α) turn the diffusion into mean-reversion — both or neither.
    is_ou = reverts_to is not None or pull is not None
    if is_ou:
        if reverts_to is None or pull is None:
            raise ValueError(
                "Ornstein–Uhlenbeck needs both reverts_to (the optimum) and pull (the strength); "
                "give both, or neither for Brownian motion."
            )
        if isinstance(reverts_to, bool) or not isinstance(reverts_to, (int, float)) \
                or not math.isfinite(reverts_to):
            raise ValueError(f"reverts_to must be a finite number, got {reverts_to!r}")
        if isinstance(pull, bool) or not isinstance(pull, (int, float)) \
                or not math.isfinite(pull) or pull <= 0:
            raise ValueError(
                f"pull must be a finite positive number (omit it for Brownian motion), got {pull!r}"
            )
        theta, alpha = float(reverts_to), float(pull)

    # conditioning: a σ² written with scaled_by reads another level, grown first on this same tree. Resolve
    # each driver once into a trajectory (value + next-switch, keyed by the shared node id), from a
    # written trait log or a grown result handed over in memory. Undriven ⇒ empty, and the walk below
    # is exactly the walk it was — no driver, no lookup, no change to the draw order.
    trajs = _resolve_drivers(_driven_mods(r), tree, "traits.continuous")

    jump_sd = _at_speciation_jump_sd(at_speciation)  # on-speciation jump width (0 if not requested)

    rng, seed = stream("traits", seed)      # own stream, and a drawn seed if none was given
    ltt = _LTT(tree) if has_diversity else None  # the standing-diversity curve, when σ² reads it
    node_values: dict[int, float] = {}
    root = tree.nodes[tree.root]
    # the initial value at t=0 — the origin the log reconstructs from (SPEC §2). A diffusion cannot be
    # rebuilt from events, but the row keeps the file's shape uniform across trait kinds.
    events: list[Change] = [Change(root.birth_time, "initial", tree.root, None, float(start))]
    inh: dict[int, tuple[float, ...]] = {}  # each lineage's σ² drift factors (variable-rates BM),
                                           # one per carried modifier, constant along its branch
    for i in _preorder(tree, progress):
        node = tree.nodes[i]
        # the root starts from `start` at t=0; every other node from its parent's end value (parent
        # < i, already set). One uniform rule: node_values[i] is the trait at node i's end_time.
        x = float(start) if node.parent is None else node_values[node.parent]
        if node.parent is not None and jump_sd > 0.0:
            jumped = x + float(rng.normal(0.0, jump_sd))  # on speciation: a jump at the split…
            events.append(Change(node.birth_time, "on_speciation", i, x, jumped))
            x = jumped                                    # …then anagenesis along the branch
        # thread the inherited factor: the root's is 1.0, each daughter's is its parent's times a
        # lognormal kick drawn at the split (so σ² is autocorrelated down the tree). None ⇒ 1.0, no draw.
        if node.parent is None:
            inh[i] = values_at_birth(drift, rng)
        else:
            inh[i] = values_at_split(drift, inh[node.parent], rng)
        t0, t1 = node.birth_time, node.end_time
        if is_ou:
            e = math.exp(-alpha * (t1 - t0))       # mean-reversion toward θ over the branch
            mean = theta + (x - theta) * e         # the mean does not read σ², so a modified σ² leaves it
            # …and the variance is the pull-weighted integral, which for a bare σ² is exactly the
            # closed form σ²/(2α)·(1−e^{−2α·dt}) this branch used before the weight existed.
            var = _accrued_variance(r, t0, t1, inherited=math.prod(inh[i]), ltt=ltt, trajs=trajs, node_id=i,
                                    pull=alpha)
        else:
            mean = x                                # pure diffusion (BM / early burst / variable-rates)
            var = _accrued_variance(r, t0, t1, inherited=math.prod(inh[i]), ltt=ltt, trajs=trajs, node_id=i)
        std = math.sqrt(var) if var > 0.0 else 0.0
        node_values[i] = mean + (float(rng.normal(0.0, std)) if std > 0.0 else 0.0)

    return TraitsResult(tree, cast("dict[int, object]", node_values), events, seed)

zombi2.traits.simulate_discrete

simulate_discrete(tree, *, states, switch=None, start=None, liability=None, threshold=None, correlation=None, at_speciation=None, seed=None, progress=False) -> TraitsResult

Evolve a discrete-state trait down a tree and return a TraitsResult. Two mechanisms:

• Mk (switch=) — a continuous-time Markov chain over the states, simulated exactly by Gillespie along every branch, so each node's (state, duration) segments are the realized history (.history) and .events reads off the transitions. switch is a symmetric rate (0.1, or PerLineage(0.1) written out), a {"marine->terrestrial": 0.1} dict, or a k×k matrix of numbers (see _q_matrix()). start is the root state (a label in states; None draws one uniformly). A switch rate may be driven by another level grown first on this same tree — write it as a rate expression, switch=PerLineage(0.4).scaled_by(habitat, {"aquatic": 3.0}) or per transition, switch={"a->b": PerLineage(0.2).scaled_by(habitat, {"aquatic": 3.0}), "b->a": 0.2}. The driver switches mid-branch, so the generator is rebuilt at each of its switches and the branch is simulated piece by piece — the exact CTMC with a time-varying generator, not one sample per branch.
• Threshold (liability= + threshold=) — the Wright–Felsenstein model: a discrete state read off an underlying continuous Brownian liability (variance-rate liability), cut into states by the threshold cut point(s) (k−1 increasing cuts for k states). start is the initial liability (a number, default 0.0). Give liability as a dict + a correlation={(a, b): ρ} overlay to evolve correlated discrete traits jointly — their liabilities diffuse together (Σ = D R D) and each is cut by the shared thresholds. A threshold trait has no Gillespie map, so .history is None and .events empty.

tree is the complete species tree (a Tree or SpeciesResult); the trait evolves on every lineage (convention B: the root diffuses over its own branch), and .values reads the extant tips. On an Mk trait, at_speciation (a probability in [0, 1]) adds an on-speciation shift — each daughter hops to a uniformly-chosen other state with that chance at every speciation. Deterministic given seed.

Source code in zombi2/traits/discrete.py

def simulate_discrete(tree, *, states, switch=None, start=None, liability=None, threshold=None,
                      correlation=None, at_speciation=None, seed=None,
                      progress=False) -> TraitsResult:
    """Evolve a discrete-state trait down a tree and return a `TraitsResult`. Two mechanisms:

    - **Mk** (``switch=``) — a continuous-time Markov chain over the ``states``, simulated **exactly**
      by Gillespie along every branch, so each node's ``(state, duration)`` segments *are* the realized
      history (``.history``) and ``.events`` reads off the transitions. ``switch`` is a symmetric rate
      (``0.1``, or ``PerLineage(0.1)`` written out), a ``{"marine->terrestrial": 0.1}`` dict, or a
      ``k×k`` matrix of numbers (see `_q_matrix()`).
      ``start`` is the root state (a label in ``states``; ``None`` draws one uniformly). A switch rate
      may be **driven by another level** grown first on this same tree — write it as a rate
      expression, ``switch=PerLineage(0.4).scaled_by(habitat, {"aquatic": 3.0})`` or per transition,
      ``switch={"a->b": PerLineage(0.2).scaled_by(habitat, {"aquatic": 3.0}), "b->a": 0.2}``. The driver
      switches mid-branch, so the generator is rebuilt at each of its switches and the branch is
      simulated piece by piece — the exact CTMC with a time-varying generator, not one sample per
      branch.
    - **Threshold** (``liability=`` + ``threshold=``) — the Wright–Felsenstein model: a discrete state
      read off an underlying continuous Brownian **liability** (variance-rate ``liability``), cut into
      ``states`` by the ``threshold`` cut point(s) (``k−1`` increasing cuts for ``k`` states). ``start``
      is the initial *liability* (a number, default 0.0). Give ``liability`` as a dict + a
      ``correlation={(a, b): ρ}`` overlay to evolve **correlated** discrete traits jointly — their
      liabilities diffuse together (``Σ = D R D``) and each is cut by the shared thresholds. A threshold
      trait has no Gillespie map, so ``.history`` is ``None`` and ``.events`` empty.

    ``tree`` is the **complete** species tree (a `Tree` or
    `SpeciesResult`); the trait evolves on every lineage (convention B: the root
    diffuses over its own branch), and ``.values`` reads the extant tips. On an Mk trait,
    ``at_speciation`` (a probability in ``[0, 1]``) adds an **on-speciation** shift — each daughter hops
    to a uniformly-chosen other state with that chance at every speciation. Deterministic given ``seed``.
    """
    tree = as_tree(tree, level="traits")
    states = list(states)
    if len(states) < 2:
        raise ValueError(f"a discrete trait needs at least 2 states, got {states!r}")
    if len(set(states)) != len(states):
        raise ValueError(f"states must be unique, got {states!r}")
    if liability is not None or threshold is not None:
        if switch is not None:
            raise ValueError("give switch= (an Mk trait) OR liability=/threshold= (a threshold trait), not both")
        if at_speciation is not None:
            raise ValueError("at_speciation is not implemented for threshold traits yet — it applies to Mk (switch=) traits")
        return _simulate_threshold(tree, states, liability, threshold, start, correlation, seed,
                                   progress)
    if correlation is not None:
        raise ValueError("correlation= on a discrete trait needs the threshold model — give liability= and threshold=")
    if switch is None:
        raise ValueError("give switch= — the transition rate(s) between the discrete states.")
    # conditioning: a switch rate written with scaled_by reads another level, grown first on this same
    # tree. The generator is then a function of the driver, so it is built per stretch rather than
    # once. No modifier at all ⇒ one constant Q and the walk below is exactly the walk it was.
    #
    # The two questions are separate, and conflating them was a bug: *any* modifier puts the run on
    # the rebuild-per-stretch path (that is what makes the generator a function of the context),
    # while only a `Driven` names a driver to resolve a trajectory from. A third-party
    # modifier used to pass the gate, land in the driver list, and crash the resolver looking for a
    # `.key` it does not have.
    sw_mods = _switch_modifiers(switch)
    entries = _driven_entries(states, switch) if sw_mods else None
    Q = None if sw_mods else _q_matrix(states, switch)
    trajs = _resolve_drivers([m for m in sw_mods if isinstance(m, Driven)], tree, "traits.discrete")
    if at_speciation is not None and (isinstance(at_speciation, bool)
            or not isinstance(at_speciation, (int, float)) or not 0.0 <= at_speciation <= 1.0):
        raise ValueError(f"at_speciation must be a probability in [0, 1] (the shift chance), got {at_speciation!r}")
    shift = 0.0 if at_speciation is None else float(at_speciation)

    rng, seed = stream("traits", seed)      # own stream, and a drawn seed if none was given
    idx = {s: i for i, s in enumerate(states)}
    if start is None:
        start_i = int(rng.integers(len(states)))
    elif start in idx:
        start_i = idx[start]
    else:
        raise ValueError(
            f"start must be one of states={states} (or None for a uniform draw), got {start!r}"
        )

    node_values: dict[int, object] = {}
    root = tree.nodes[tree.root]
    # the initial state at t=0 — the origin the log reconstructs from: tree + this + the switches give
    # the driver on every lineage, so the event log is the driver file (no separate driver).
    events: list[Change] = [Change(root.birth_time, "initial", tree.root, None, states[start_i])]
    for i in _preorder(tree, progress):
        node = tree.nodes[i]
        # the root starts from `start` at t=0 and evolves over its own branch; every other node from
        # its parent's end state (parent < i, already set) — the same convention-B walk as continuous.
        cur = start_i if node.parent is None else idx[node_values[node.parent]]
        if node.parent is not None and shift > 0.0 and float(rng.random()) < shift:
            j = int(rng.integers(len(states) - 1))  # on speciation: hop to a uniform *other* state
            new = j if j < cur else j + 1
            events.append(Change(node.birth_time, "on_speciation", i, states[cur], states[new]))
            cur = new
        if Q is not None:
            end_i, segs = _gillespie(cur, node.end_time - node.birth_time, Q, rng)
        else:  # a driven switch rate: no constant generator — cut the branch where the driver switches
            end_i, segs = _gillespie_driven(cur, node, i, entries, len(states), trajs, rng)
        t = node.birth_time  # the transitions between the Gillespie segments are the on-branch events
        for (s1, d1), (s2, _d) in zip(segs, segs[1:]):
            t += d1
            events.append(Change(t, "on_branch", i, states[s1], states[s2]))
        node_values[i] = states[end_i]
    events.sort(key=lambda c: c.time)
    return TraitsResult(tree, node_values, events, seed, kind="discrete")

zombi2.traits.simulate_traits

simulate_traits(tree, traits, *, joint=False, seed=None, progress=False)

Evolve several traits at once along a fixed tree, each one able to read the others.

One trait is simulate_discrete() / simulate_continuous(), and a trait grown first and then read is conditioning — two ordinary runs. This is for the case with no order: trait A's switch rate reads trait B while B's reads A, so neither can be finished before the other starts. That is the trait level joined to itself (SPEC §3), and being one level with one kind of result it stays here rather than going to zombi2.joint.simulate.

traits is a list of discrete() specs, each with a name, and a rate reads another by scaled_by("traits:<name>", …). joint=True says the run is what it is, and is checked both ways: asking for it when no trait reads another is an error, and reading another without it is an error too.

at_speciation works here as it does in a run of its own: each trait carrying one hops on its own at the split, and the pair lands wherever the two hops leave it.

Returns {name: TraitsResult} — one complete result per trait, exactly what the single-trait runners return, so every reader of one works on these unchanged. Deterministic given seed.

Source code in zombi2/traits/discrete.py

def simulate_traits(tree, traits, *, joint=False, seed=None, progress=False):
    """Evolve **several traits at once** along a fixed tree, each one able to read the others.

    One trait is `simulate_discrete()` / `simulate_continuous()`, and a trait grown first and then
    read is conditioning — two ordinary runs. This is for the case with no order: trait A's switch
    rate reads trait B while B's reads A, so neither can be finished before the other starts. That is
    the trait level joined to itself (SPEC §3), and being one level with one kind of result it stays
    here rather than going to `zombi2.joint.simulate`.

    ``traits`` is a list of `discrete()` specs, each with a ``name``, and a rate reads another by
    ``scaled_by("traits:<name>", …)``. ``joint=True`` says the run is what it is, and is checked both
    ways: asking for it when no trait reads another is an error, and reading another without it is an
    error too.

    ``at_speciation`` works here as it does in a run of its own: each trait carrying one hops on its
    own at the split, and the pair lands wherever the two hops leave it.

    Returns ``{name: TraitsResult}`` — one complete result per trait, exactly what the single-trait
    runners return, so every reader of one works on these unchanged. Deterministic given ``seed``.
    """
    from ..params.connection import Driven

    tree = as_tree(tree, level="traits")
    specs = list(traits)
    if len(specs) < 2:
        raise ValueError(
            f"simulate_traits evolves several traits at once, and got {len(specs)}. One trait is "
            f"simulate_discrete(tree, ...); two that read each other are what this is for.")
    if len(specs) > 2:
        raise NotImplementedError(
            f"two traits reading each other is what this evolves today, and got {len(specs)}. Three "
            f"is the same idea over a bigger product of states, and is not built.")
    for spec in specs:
        if not isinstance(spec, DiscreteTrait):
            raise TypeError(
                f"simulate_traits takes discrete trait specs — traits.discrete(name='size', "
                f"states=[...], switch=...) — and got {spec!r}. A continuous trait in a cycle needs "
                f"its diffusion held still over short stretches, which is not built.")
        if not spec.name:
            raise ValueError(
                "each trait needs a name here, because a rate reads the other one by it: "
                "traits.discrete(name='habitat', ...) and scaled_by('traits:habitat', ...).")
    if len({s.name for s in specs}) != len(specs):
        raise ValueError(f"trait names must be unique, got {[s.name for s in specs]}")

    # each trait's own alphabet and its rate specs, left unsettled: a switch rate reading the other
    # trait is not one number, which is exactly what `_q_matrix` would demand
    resolved = [(list(s.states), _driven_entries(list(s.states), s.switch)) for s in specs]
    names = {k for s in specs for k in _driver_keys(s)}
    reads = 0
    for spec, other in ((specs[0], specs[1]), (specs[1], specs[0])):
        for entry in _driven_entries(list(spec.states), spec.switch):
            for m in entry[2].modifiers:
                if isinstance(m, Driven):
                    if m.driver not in names:
                        raise ValueError(
                            f"trait {spec.name!r} reads {m.driver!r}, which is not a trait in this "
                            f"run. The traits here are {sorted(s.name for s in specs)}, read as "
                            f'scaled_by("traits:<name>", ...).')
                    reads += 1
    if reads and not joint:
        raise ValueError(
            "a trait here reads another trait in this same run, so the two are joint — neither can "
            "be finished before the other starts. Say so with joint=True. To read a trait grown "
            "EARLIER, pass that run's result rather than a name, which is conditioning.")
    if joint and not reads:
        raise ValueError(
            "joint=True says the traits drive each other, but none reads another. Give a switch "
            'rate a scaled_by("traits:<name>", ...), or evolve them as separate runs.')

    rng, seed = stream("traits", seed)
    Q = _product_generator(specs, resolved)
    (a_states, _a), (b_states, _b) = resolved
    kb = len(b_states)
    starts = []
    for spec, (states, _e) in zip(specs, resolved):
        idx = {s: i for i, s in enumerate(states)}
        if spec.start is None:
            starts.append(int(rng.integers(len(states))))
        elif spec.start in idx:
            starts.append(idx[spec.start])
        else:
            raise ValueError(f"start must be one of states={states} (or None for a uniform draw), "
                             f"got {spec.start!r}")
    start = starts[0] * kb + starts[1]

    root = tree.nodes[tree.root]
    node_pairs: dict[int, int] = {}
    per_trait: list[list] = [
        [Change(root.birth_time, "initial", tree.root, None, a_states[starts[0]])],
        [Change(root.birth_time, "initial", tree.root, None, b_states[starts[1]])]]
    shifts = [0.0 if s.at_speciation is None else float(s.at_speciation) for s in specs]
    for i in _preorder(tree, progress):
        node = tree.nodes[i]
        cur = start if node.parent is None else node_pairs[node.parent]
        if node.parent is not None and any(shifts):
            # each trait hops on its own at the split, exactly as it would in a run of its own; the
            # pair simply lands wherever the two hops leave it
            parts = [cur // kb, cur % kb]
            for k, (shift, (states, _e)) in enumerate(zip(shifts, resolved)):
                if shift > 0.0 and float(rng.random()) < shift:
                    j = int(rng.integers(len(states) - 1))   # to a uniform *other* state
                    new = j if j < parts[k] else j + 1
                    per_trait[k].append(Change(node.birth_time, "on_speciation", i,
                                               states[parts[k]], states[new]))
                    parts[k] = new
            cur = parts[0] * kb + parts[1]
        end, segs = _gillespie(cur, node.end_time - node.birth_time, Q, rng)
        # one product move is one trait switching, so unpacking the segments splits the pair's
        # history back into the two the reader asked for, with no state left ambiguous
        t = node.birth_time
        for (s1, d1), (s2, _d) in zip(segs, segs[1:]):
            t += d1
            for k, (states, changes) in enumerate(((a_states, per_trait[0]),
                                                   (b_states, per_trait[1]))):
                was, now = (s1 // kb, s2 // kb) if k == 0 else (s1 % kb, s2 % kb)
                if was != now:
                    changes.append(Change(t, "on_branch", i, states[was], states[now]))
        node_pairs[i] = end
    out = {}
    for k, (spec, (states, _e)) in enumerate(zip(specs, resolved)):
        values = {i: states[(p // kb) if k == 0 else (p % kb)] for i, p in node_pairs.items()}
        per_trait[k].sort(key=lambda c: c.time)
        out[spec.name] = TraitsResult(tree, values, per_trait[k], seed, kind="discrete")
    return out

zombi2.traits.TraitsResult dataclass

TraitsResult(complete_tree: Tree, node_values: dict[int, object], events: list[Change] = list(), seed: int | None = None, kind: str = 'continuous')

What simulate_continuous / simulate_discrete returns: the complete_tree it ran on, node_values at every node (the value at each node — extant, extinct, and internal alike; a float for a continuous trait, a state label for a discrete / threshold one, a per-trait dict for correlated traits), the events log, the seed, and the kind ("continuous" / "discrete" / "threshold"). The observed dataset is the extant tips, .values.

events is the timestamped event log — the same shape as the genome level's and the source of truth for a discrete (Mk) trait, from which history (the per-branch stochastic character map) is derived. A continuous trait has no along-branch events, so its log holds the initial row and the on-speciation jumps (just the initial row without at_speciation) while node_values carries the diffusion; a threshold trait's crossings are un-timed, so its log is empty and it has no map.

values property

values: dict[str, object]

The observed trait dataset — the value at each extant tip (the comparative-data vector), keyed by the tip name the tree writes: n5, or e5 for a lineage that died.

Keyed by name because the only thing anyone does with this is join it to the tree, and the tree names its tips. It used to be keyed by the bare node id (5), which the written trait_values.tsv and every Newick label do not use — so a comparative dataset built in Python shared no keys at all with the tree beside it, and nothing said so. The two files a write() produces always did match; it was the in-memory pair that did not.

values_by_id is the old view, for code that joins on node ids. Internal and extinct nodes keep their exact ancestral / lineage values in node_values, which stays id-keyed: it is the run's own record, not a dataset to export.

values_by_id property

values_by_id: dict[int, object]

values, keyed by node id rather than tip name — the shape values had before it was keyed to match the tree. For joining against node_values, complete_tree.nodes or anything else that works in ids.

history cached property

history: dict[int, list] | None

The per-branch stochastic character map — {node: [(state, duration), …]} whose durations sum to the branch length — derived from the event log (a discrete / Mk trait only). None for a continuous trait (a diffusion has no map) and for a threshold trait (its liability crossings are un-timed).

summary

summary() -> dict

What this run produced, as a plain dict — the payload of trait_summary.json.

A trait's shape depends on its kind, so the summary does too. A discrete trait — and a threshold one, which reads a discrete state off a continuous liability — is described by its switches: how many, and how the tips ended up distributed over the states — which is the thing you look at first, because a run whose tips are all in one state has told you nothing. A continuous one is described by where the values got to, since there are no along-branch events to count; its log holds the on-speciation jumps, and that count is here so an empty one is visibly empty rather than ambiguous.

Source code in zombi2/traits/result.py

def summary(self) -> dict:
    """What this run produced, as a plain dict — the payload of ``trait_summary.json``.

    A trait's shape depends on its kind, so the summary does too. A **discrete** trait — and a
    **threshold** one, which reads a discrete state off a continuous liability — is described
    by its switches: how many, and how the tips ended up distributed over the states — which is the
    thing you look at first, because a run whose tips are all in one state has told you nothing.
    A **continuous** one is described by where the values got to, since there are no along-branch
    events to count; its log holds the on-speciation jumps, and that count is here so an
    empty one is visibly empty rather than ambiguous."""
    values = list(self.values.values())
    switches = sum(1 for e in self.events if e.kind == "on_branch")
    jumps = sum(1 for e in self.events if e.kind == "on_speciation")
    out: dict[str, object] = {
        "level": "traits",
        "seed": self.seed,
        "kind": self.kind,
        "tips": len(values),
        "nodes": len(self.node_values),
        "events": {"on_branch": switches, "on_speciation": jumps},
    }
    # A CORRELATED run holds one value per trait at every node, so there is no single
    # distribution to describe and each trait gets its own entry. This is tested before `kind`
    # because both kinds can be correlated, and both got it wrong: a continuous one raised
    # `TypeError: float() argument must be … not 'dict'`, and a threshold one counted whole
    # dicts as if they were states, giving keys like "{'a': 'b', 'b': 'b'}". The per-trait shape
    # matches what `trait_values.tsv` writes (one column per trait), so the summary and the table
    # describe the same object.
    if values and isinstance(values[0], dict):
        per_trait = [cast("dict[str, object]", v) for v in values]
        names = list(per_trait[0])
        root = cast("dict[str, object]", self.node_values[self.complete_tree.root])
        out["traits"] = names
        if self.kind == "continuous":
            out["values"] = {n: _stats([float(cast(float, v[n])) for v in per_trait])
                             for n in names}
            out["value_at_root_node"] = {n: float(cast(float, root[n])) for n in names}
        else:
            out["states"] = {}
            for n in names:
                counts = collections.Counter(str(v[n]) for v in per_trait)
                cast(dict, out["states"])[n] = dict(sorted(counts.items()))
            out["state_at_root_node"] = {n: str(root[n]) for n in names}
        return out
    # continuous is the only NUMERIC kind. A threshold trait reads a discrete state off a
    # continuous liability, so its values are states like a discrete trait's — taking a mean of
    # them is what the first version of this tried to do.
    if self.kind != "continuous":
        counts = collections.Counter(str(v) for v in values)
        out["states"] = dict(sorted(counts.items()))
        out["states_at_tips"] = len(counts)
        # a run whose tips all share one state is degenerate, and the number that says so
        out["most_common_share"] = (round(counts.most_common(1)[0][1] / len(values), 6)
                                    if values else None)
    else:
        numeric = [float(cast(float, v)) for v in values]
        out["values"] = _stats(numeric)
        # the root NODE, i.e. after diffusing along the stem — not the value the run started
        # from, which is `start` and belongs to no node
        out["value_at_root_node"] = float(
            cast(float, self.node_values[self.complete_tree.root]))
    return out

write

write(directory, outputs=None) -> None

Write chosen outputs to directory (created if needed); the default is the set zombi2 traits writes for this kind, so the command and the API leave the same directory (_DEFAULT_OUTPUTS). "values" → trait_values.tsv (the node<TAB>kind<TAB>trait table over every node — tips, extinct lineages and internal nodes; kind is the tip's fate — extant / extinct (/ unsampled under incomplete sampling) — or ancestor for an internal node, so the extant tips filter out with kind == "extant"); "events" → trait_events.tsv, the event log (time · kind · lineage · from · to) — one initial row at t=0 giving the initial state, then every switch in time order; "summary" → trait_summary.json, what came out, as JSON (summary); "tree" → trait_tree.nwk, the complete tree as Newick with every node annotated [&trait=…] (a trait tree, carrying the exact ancestral values; opens in FigTree / iTOL).

trait_events.tsv is also the driver file: a genome / sequence run drives a rate with scaled_by("trait_events.tsv", …), replaying it against the shared tree. A discrete trait's log reconstructs its state on every lineage exactly (that is what the initial row and the switch times are for); a continuous trait's diffusion cannot be rebuilt from events, so it carries only the initial row and any on-speciation jumps.

Source code in zombi2/traits/result.py

def write(self, directory, outputs=None) -> None:
    """Write chosen ``outputs`` to ``directory`` (created if needed); the default is the set
    ``zombi2 traits`` writes for this kind, so the command and the API leave the same directory
    (`_DEFAULT_OUTPUTS`). ``"values"`` →
    ``trait_values.tsv`` (the ``node<TAB>kind<TAB>trait`` table over **every** node — tips, extinct
    lineages and internal nodes; ``kind`` is the tip's fate — ``extant`` / ``extinct`` (/ ``unsampled``
    under incomplete sampling) — or ``ancestor`` for an internal node, so the extant tips filter out
    with ``kind == "extant"``); ``"events"`` →
    ``trait_events.tsv``, the event log (``time · kind · lineage · from · to``) — one ``initial``
    row at t=0 giving the initial state, then every switch in time order; ``"summary"`` →
    ``trait_summary.json``, what came out, as JSON (`summary`); ``"tree"`` →
    ``trait_tree.nwk``, the complete tree as Newick with **every** node annotated ``[&trait=…]``
    (a *trait tree*, carrying the exact ancestral values; opens in FigTree / iTOL).

    ``trait_events.tsv`` is also the **driver file**: a genome / sequence run drives a rate
    with ``scaled_by("trait_events.tsv", …)``, replaying it against the shared tree. A
    **discrete** trait's log reconstructs its state on every lineage exactly (that is what the
    ``initial`` row and the switch times are for); a continuous trait's diffusion cannot be rebuilt
    from events, so it carries only the ``initial`` row and any on-speciation jumps."""
    if outputs is None:
        outputs = _DEFAULT_OUTPUTS.get(self.kind, _WRITE_OUTPUTS)
    unknown = [o for o in outputs if o not in _WRITE_OUTPUTS]
    if unknown:
        raise ValueError(f"unknown write outputs {unknown}; choose from {list(_WRITE_OUTPUTS)}")
    d = pathlib.Path(directory)
    d.mkdir(parents=True, exist_ok=True)
    names = self.complete_tree.labels()   # e<id> for a lineage that died; n<id> for the rest
    if "values" in outputs:
        # every node — extant tips, extinct lineages (e<id>) and internal nodes (n<id>) alike, each
        # with its exact value and a `kind` column (the tip's fate, or `ancestor`) so a comparative
        # vector of exactly the extant tips is one `kind == "extant"` filter
        kinds = {i: _node_kind(n) for i, n in self.complete_tree.nodes.items()}
        (d / "trait_values.tsv").write_text(_values_tsv(self.node_values, names, kinds),
                                            encoding="utf-8")
    if "events" in outputs:
        (d / "trait_events.tsv").write_text(_events_tsv(self.events, names), encoding="utf-8")
    if "summary" in outputs:
        write_summary(d / "trait_summary.json", self.summary())
    if "tree" in outputs:
        (d / "trait_tree.nwk").write_text(
            _trait_newick(self.complete_tree, self.node_values) + "\n", encoding="utf-8")

zombi2.traits.Change dataclass

Change(time: float, kind: str, lineage: int, from_state: object, to_state: object)

A realized trait change — one entry of the event log, the trait twin of the genome level's GeneEdge. On lineage lineage at time (origin-forward, the species-tree clock) the state went from from_state to to_state. kind is "on_branch" — a switch along a branch (an Mk transition) — "on_speciation" — a jump at a speciation node (from at_speciation; for a continuous trait from_state / to_state are the pre- and post-jump values) — or "initial", one synthetic entry at t=0 giving the initial state the run started in (from_state None, time the root's birth_time). That row is what lets the log stand on its own: the tree plus the initial state plus the switches determines the trait on every lineage at every instant, so no separate driver file is needed.

zombi2.traits.DiscreteTrait dataclass

DiscreteTrait(states: tuple, switch: object, start: object = None, at_speciation: object = None, name: 'str | None' = None)

A discrete (Mk) trait process — its parameters bundled but not yet run (SPEC §4). simulate_discrete(tree, ...) is the runner that grows this on a fixed tree; a joint model instead takes this spec and grows the trait with the tree it drives (joint.simulate(species.birth_death(...), traits.discrete(...))), so neither can be simulated first. Same parameters as simulate_discrete() (the Mk half): states, switch (the rate spec), start (the root state, None = uniform), at_speciation (the on-speciation shift probability).

zombi2.traits.ContinuousTrait dataclass

ContinuousTrait(start: float = 0.0, rate: object = 1.0, reverts_to: float | None = None, pull: float | None = None, at_speciation: object = None, name: str | None = None)

A continuous (diffusing) trait process — its parameters bundled but not yet run (SPEC §4).

simulate_continuous is the runner that grows this on a tree that already exists. This spec is for the case where the tree does not: hand it to joint.simulate(species.birth_death(...), traits.continuous(...)) and the trait grows with the tree whose speciation it drives — QuaSSE. Same arguments, same meanings, no tree.

A diffusing driver moves at every instant, so the birth rate it drives is never constant and a Gillespie step has nothing to hold still. The run therefore slices: the driver is held fixed across a step of step and released at each boundary, which the driven rate declares — scaled_by("trait", Curve(f), step=0.05). That is an approximation, and the only one in a joint run: everything else races exactly. Halve step, rerun the same seed, and see whether the answer moves.

The fields are simulate_continuous's: start (the value at t=0), rate (the variance-rate σ², per lineage, a bare number or a changing_at skyline), reverts_to + pull (Ornstein–Uhlenbeck; give both or neither), at_speciation (the variance of a jump at each split) and name (what a rate calls it, "traits:<name>"; a run holding one trait also answers to "trait").