From ae1e21781ea8312ebb443e841cba976f49d8edc6 Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Tue, 18 Aug 2026 12:58:59 +1200 Subject: [PATCH 1/2] Add CRegCCD: class-based regularised CCD with full support Smooths over every bipartition of a clade rather than only those in the CCD graph, with the pseudocount set by how many novel clades the bipartition introduces: a per-split alpha on the CCD0 split set (classes 0 and 1 pooled), and class totals alpha1, alpha2 for the one- and two-novel-clade classes. The exponentially large class is obtained by subtraction from 2^(m-1)-1, so it is never enumerated and there is no #P-hard counting. Per-class totals rather than per-split constants are essential: with a constant pseudocount the two-novel class swamps the data, leaving observed splits 6e-9 of the probability at a 40-taxon root clade. Properties, all by construction: exactly normalised by the chain rule (no partition function, no truncation); full support; regCCD nested exactly at alpha1 = alpha2 = 0; an observed split always outranks an expanded one at the same clade; and the mass held back shrinks as f(C) grows. Implements scoring, exact sampling, MAP and entropy. Sampling draws a class then a member, using rejection for the two-novel class, and scores each draw with the same routine the scorer uses, so the sampled and scored distributions coincide. MAP runs a DP over the observed-clade DAG with an optional wider search admitting one-novel-clade splits, plus a bound that can certify global optimality. Entropy is unbiased by Monte Carlo; a deterministic recursion is also provided, which approximates novel subclades as structureless and is therefore optimistic. Verified in CRegCCDTest and CRegCCDMapEntropyTest: total mass 1.000000000000 by enumeration on 4-7 taxa; agreement with RegCCD to 1.8e-15 at alpha1 = alpha2 = 0; class sizes partitioning all 2^(m-1)-1 bipartitions; sampled frequencies matching scored probabilities across all 105 five-taxon topologies; and the MAP search matching brute force in twelve configurations with the certificate firing in each. Co-Authored-By: Claude Opus 5 (1M context) --- src/main/java/ccd/model/CRegCCD.java | 1157 +++++++++++++++++ .../java/ccd/model/CRegCCDMapEntropyTest.java | 298 +++++ src/test/java/ccd/model/CRegCCDTest.java | 293 +++++ 3 files changed, 1748 insertions(+) create mode 100644 src/main/java/ccd/model/CRegCCD.java create mode 100644 src/test/java/ccd/model/CRegCCDMapEntropyTest.java create mode 100644 src/test/java/ccd/model/CRegCCDTest.java diff --git a/src/main/java/ccd/model/CRegCCD.java b/src/main/java/ccd/model/CRegCCD.java new file mode 100644 index 0000000..36e522a --- /dev/null +++ b/src/main/java/ccd/model/CRegCCD.java @@ -0,0 +1,1157 @@ +package ccd.model; + +import beast.base.evolution.tree.Node; +import beast.base.evolution.tree.Tree; +import beastfx.app.treeannotator.TreeAnnotator.TreeSet; +import ccd.model.bitsets.BitSet; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * CRegCCD -- the class-based regularised CCD (Jonathan's proposal): a full-support tree + * distribution obtained by additive smoothing over all bipartitions of every clade, with the + * pseudocount depending only on which of four classes a bipartition falls into. + * + *

At a clade {@code C} of {@code m} taxa, each of the {@code 2^(m-1) - 1} bipartitions + * {@code {A, B}} belongs to exactly one class: + *

    + *
  1. {@code A_1}: the split was observed in the sample;
    2. + *
  2. {@code A_2}: unobserved split, both {@code A} and {@code B} are observed clades + * (this is exactly the CCD0 split expansion);
    3. + *
  3. {@code A_3}: unobserved split, exactly one of {@code A}, {@code B} is an observed clade;
    4. + *
  4. {@code A_4}: neither {@code A} nor {@code B} is an observed clade.
    5. + *
+ * Only classes 1--3 are ever materialised; {@code |A_4|} follows in closed form as + * {@code (2^(m-1) - 1) - |A_1| - |A_2| - |A_3|}, so the whole of tree space is represented without + * enumerating it. There is no escape probability, no reserve equation and no region decomposition. + * + *

Per-class totals, not per-split constants. Each {@code a_j} is the total + * pseudocount mass of its class, so the per-split pseudocount is {@code a_j / |A_j|} and + * 

+ *   theta(S) = (f(S) + a_{j(S)} / |A_{j(S)}|) / (f(C) + sum over non-empty classes of a_j).
+ *
+ * This matters: {@code |A_4|} is essentially {@code 2^(m-1)}, so a constant per-split pseudocount + * would give class 4 all the mass and the data none (on 40 taxa the observed splits retain about + * {@code 6e-9} of the probability at a root clade; see {@code SplitClassSizeAnalysis}). With per-class + * totals the retained mass is independent of taxon count. Equivalently: draw a class, then draw + * uniformly within it. + * + *

Consequences, all by construction rather than by correction: + *

+ * + * @author Claude + */ +public class CRegCCD extends CCD1 { + + /** + * Per-split pseudocount on the CCD0 split set (splits introducing no novel clade). This is + * regCCD's {@code alpha}; the fitted value was 0.4 on every real data set tested. + */ + public static final double DEFAULT_ALPHA = 0.4; + /** Total prior mass for splits introducing one novel clade. */ + public static final double DEFAULT_ALPHA1 = 0.4; + /** Total prior mass for splits introducing two novel clades. */ + public static final double DEFAULT_ALPHA2 = 0.05; + + private final double alpha; + private final double alpha1; + private final double alpha2; + + /** Class sizes are parameter-independent, so they are computed once and reused across a search. + * Concurrent because {@link #sampleTrees} fans draws out over threads. */ + private final Map sizeCache = new java.util.concurrent.ConcurrentHashMap<>(); + private volatile List sortedCladeBits; + + /** + * Strictly-positive height increment for a novel internal node whose clade has no recorded + * height, so that branch lengths stay positive. + */ + private static final double NOVEL_HEIGHT_EPS = 1e-8; + + /** + * Draw a class-4 split by rejection while at least this fraction of bipartitions are class 4, + * which bounds the expected number of attempts by its reciprocal; below it, enumerate instead. + * Since classes 1-3 are only polynomially large, a small acceptance rate implies a small + * {@code 2^(m-1)}, so the enumeration branch is always cheap. + */ + private static final double MIN_REJECTION_ACCEPTANCE = 0.02; + + public CRegCCD(List trees, double burnin) { + this(trees, burnin, DEFAULT_ALPHA, DEFAULT_ALPHA1, DEFAULT_ALPHA2); + } + + public CRegCCD(List trees, double burnin, double alpha, double alpha1, double alpha2) { + super(trees, burnin); + validate(alpha, alpha1, alpha2); + this.alpha = alpha; + this.alpha1 = alpha1; + this.alpha2 = alpha2; + } + + public CRegCCD(TreeSet treeSet) { + this(treeSet, DEFAULT_ALPHA, DEFAULT_ALPHA1, DEFAULT_ALPHA2); + } + + public CRegCCD(TreeSet treeSet, double alpha, double alpha1, double alpha2) { + super(treeSet, false); + validate(alpha, alpha1, alpha2); + this.alpha = alpha; + this.alpha1 = alpha1; + this.alpha2 = alpha2; + } + + private static void validate(double alpha, double alpha1, double alpha2) { + if (alpha <= 0) { + throw new IllegalArgumentException("alpha must be > 0, got " + alpha); + } + if (alpha1 < 0 || alpha2 < 0) { + throw new IllegalArgumentException( + "alpha1 and alpha2 must be >= 0, got " + alpha1 + ", " + alpha2); + } + } + + /** regCCD's per-split pseudocount on the CCD0 split set. */ + public double getAlpha() { + return alpha; + } + + /** Total prior mass for splits introducing one novel clade. */ + public double getAlpha1() { + return alpha1; + } + + /** Total prior mass for splits introducing two novel clades. */ + public double getAlpha2() { + return alpha2; + } + + @Override + public String toString() { + return "CRegCCD(alpha=" + alpha + ", alpha1=" + alpha1 + ", alpha2=" + alpha2 + ")"; + } + + /** + * Per-split pseudocount of each split class at {@code cBits} (as logs, so that + * {@code alpha2/|A_2|} with an exponentially large class cannot underflow), plus the normaliser + * {@code Z} in the last slot. Indices 0 and 1 are the two halves of the CCD0 split set and share + * the per-split {@code alpha}; indices 2 and 3 are the one- and two-novel-clade classes, whose + * class totals are spread over their members. + */ + private double[] logWeightsAndZ(BitSet cBits, double alpha, double alpha1, double alpha2) { + double[] size = classSizes(cBits); + Clade c = getClade(cBits); + double z = (c != null) ? c.getNumberOfOccurrences() : 0.0; + double[] w = {Double.NEGATIVE_INFINITY, Double.NEGATIVE_INFINITY, + Double.NEGATIVE_INFINITY, Double.NEGATIVE_INFINITY}; + double n0 = size[0] + size[1]; + if (n0 > 0) { + w[0] = Math.log(alpha); + w[1] = w[0]; + z += alpha * n0; + } + if (size[2] > 0) { + w[2] = Math.log(alpha1) - Math.log(size[2]); + z += alpha1; + } + if (size[3] > 0) { + w[3] = Math.log(alpha2) - Math.log(size[3]); + z += alpha2; + } + return new double[]{w[0], w[1], w[2], w[3], z}; + } + + /* ---------------------------------------------------------------------- + * Scoring + * ------------------------------------------------------------------- */ + + @Override + public double getLogProbabilityOfTree(Tree tree) { + return scoreTree(tree, alpha, alpha1, alpha2); + } + + /** + * Log probability at pseudocounts other than this model's own, reusing the cached (parameter-free) + * class sizes. Lets a cross-validation sweep evaluate many parameter vectors on one trained model. + */ + public double getLogProbabilityOfTree(Tree tree, double b0, double b1, double b2) { + validate(b0, b1, b2); + return scoreTree(tree, b0, b1, b2); + } + + @Override + public double getProbabilityOfTree(Tree tree) { + return Math.exp(getLogProbabilityOfTree(tree)); + } + + /** Always true: CRegCCD has full support over the trees on its taxon set. */ + @Override + public boolean containsTree(Tree tree) { + return true; + } + + private double scoreTree(Tree tree, double b0, double b1, double b2) { + Map bits = new HashMap<>(); + computeBits(tree.getRoot(), bits); + double logp = 0.0; + for (Node v : tree.getNodesAsArray()) { + if (v.isLeaf()) { + continue; + } + logp += logSplitProbability(bits.get(v), + bits.get(v.getChildren().get(0)), + bits.get(v.getChildren().get(1)), + b0, b1, b2); + } + return logp; + } + + /** + * Log conditional probability of the bipartition {@code {aBits, bBits}} of clade {@code cBits}, + * for a clade that need not be observed. This is the whole model: every internal node of a tree + * contributes exactly one such factor. + */ + double logSplitProbability(BitSet cBits, BitSet aBits, BitSet bBits, + double b0, double b1, double b2) { + double[] wz = logWeightsAndZ(cBits, b0, b1, b2); + int cls = splitClass(cBits, aBits, bBits); + double fS = (cls == 0) + ? observedPartition(getClade(cBits), aBits, bBits).getNumberOfOccurrences() : 0.0; + double logNumerator = (fS > 0) ? Math.log(fS + Math.exp(wz[cls])) : wz[cls]; + return logNumerator - Math.log(wz[4]); + } + + /** + * Which of the four classes the bipartition {@code {aBits, bBits}} of {@code cBits} belongs to, + * as a 0-based index (0 = observed split, 3 = neither child observed). + */ + int splitClass(BitSet cBits, BitSet aBits, BitSet bBits) { + if (observedPartition(getClade(cBits), aBits, bBits) != null) { + return 0; + } + boolean aObs = getClade(aBits) != null; + boolean bObs = getClade(bBits) != null; + return (aObs && bObs) ? 1 : ((aObs || bObs) ? 2 : 3); + } + + private CladePartition observedPartition(Clade c, BitSet aBits, BitSet bBits) { + if (c == null) { + return null; + } + Clade ca = getClade(aBits); + Clade cb = getClade(bBits); + if (ca == null || cb == null) { + return null; + } + return c.getCladePartition(ca, cb); + } + + /* ---------------------------------------------------------------------- + * Class sizes + * ------------------------------------------------------------------- */ + + /** + * Sizes {@code {|A_1|, |A_2|, |A_3|, |A_4|}} of the four split classes of clade {@code cBits}. + * + *

{@code |A_1|} is the number of observed splits; a single pass over the observed subclades of + * {@code C} yields {@code |A_3|} (observed subclade whose complement is not observed) and the + * number of observed-clade pairs, from which {@code |A_2|} follows; {@code |A_4|} is the + * remainder of {@code 2^(m-1) - 1}. Cached, and independent of the pseudocounts. + */ + double[] classSizes(BitSet cBits) { + double[] cached = sizeCache.get(cBits); + if (cached != null) { + return cached; + } + int m = cBits.cardinality(); + Clade c = getClade(cBits); + + double n1 = 0.0; + if (c != null) { + for (CladePartition p : c.getPartitions()) { + if (p.getNumberOfOccurrences() > 0) { + n1++; + } + } + } + + // one pass over observed clades strictly inside C + int bothObservedEnds = 0; // counts each both-observed bipartition twice (once per side) + double n3 = 0.0; + for (BitSet d : sortedCladeBits()) { + if (d.cardinality() >= m || !subset(d, cBits)) { + continue; + } + BitSet complement = BitSet.newBitSet(cBits); + complement.andNot(d); + if (getClade(complement) != null) { + bothObservedEnds++; + } else { + n3++; + } + } + double n2 = Math.max(0.0, bothObservedEnds / 2.0 - n1); + + double total = Math.pow(2.0, m - 1) - 1.0; + double n4 = Math.max(0.0, total - n1 - n2 - n3); + + double[] size = {n1, n2, n3, n4}; + sizeCache.put(BitSet.newBitSet(cBits), size); + return size; + } + + private synchronized List sortedCladeBits() { + if (sortedCladeBits == null) { + List all = new ArrayList<>(); + for (Clade c : getClades()) { + all.add(c.getCladeInBits()); + } + all.sort(CRegCCD::compareBitSets); + sortedCladeBits = all; + } + return sortedCladeBits; + } + + private BitSet computeBits(Node v, Map bits) { + BitSet b = BitSet.newBitSet(leafArraySize); + if (v.isLeaf()) { + b.set(v.getNr()); + } else { + b.or(computeBits(v.getChildren().get(0), bits)); + b.or(computeBits(v.getChildren().get(1), bits)); + } + bits.put(v, b); + return b; + } + + private static boolean subset(BitSet a, BitSet c) { + BitSet tmp = BitSet.newBitSet(a); + tmp.andNot(c); + return tmp.isEmpty(); + } + + /* ---------------------------------------------------------------------- + * MAP tree + * + * The maximum over all of tree space is a DP over the subset lattice, so instead we run the DP + * over the observed-clade DAG using only the both-children-observed splits (classes 1 and 2 -- + * exactly CCD0's split set) and then *certify* that no off-backbone tree can beat it. + * + * The certificate is a pair of upper-bound DPs over the same DAG. U(C) bounds the best subtree + * log-probability over ALL trees on C, and V(C) bounds it over trees that use at least one + * off-backbone (class 3 or 4) split. Any subtree contributes at most 0, so a novel child is + * bounded by 0; every class-3 split at C shares one theta, as does every class-4 split, because + * the model is uniform within a class. If best(root) > V(root), no tree using an off-backbone + * split anywhere can beat the backbone optimum, so the backbone MAP is the global MAP. + * ------------------------------------------------------------------- */ + + private volatile Map mapBest; + private volatile Map mapArg; + private volatile double offBackboneBound = Double.NaN; + + /** All both-children-observed bipartitions of {@code cBits} (classes 1 and 2), each once. */ + private List backboneSplits(BitSet cBits) { + int m = cBits.cardinality(); + List out = new ArrayList<>(); + for (BitSet d : sortedCladeBits()) { + if (d.cardinality() >= m || !subset(d, cBits)) { + continue; + } + BitSet complement = BitSet.newBitSet(cBits); + complement.andNot(d); + if (getClade(complement) == null || compareBitSets(d, complement) >= 0) { + continue; + } + out.add(new BitSet[]{d, complement}); + } + return out; + } + + /** Log theta shared by every split of the given off-backbone class at {@code cBits}. */ + private double logThetaOfClass(BitSet cBits, int cls) { + double[] size = classSizes(cBits); + if (size[cls] <= 0) { + return Double.NEGATIVE_INFINITY; + } + double[] wz = logWeightsAndZ(cBits, alpha, alpha1, alpha2); + return wz[cls] - Math.log(wz[4]); + } + + private synchronized void computeMAP() { + if (mapBest != null) { + return; + } + List clades = new ArrayList<>(getClades()); + clades.sort(java.util.Comparator.comparingInt(Clade::size)); + + Map best = new HashMap<>(); + Map arg = new HashMap<>(); + Map upper = new HashMap<>(); // U: best over all trees + Map upperOff = new HashMap<>(); // V: best over trees using an off-backbone split + + for (Clade c : clades) { + BitSet cb = c.getCladeInBits(); + if (c.size() == 1) { + best.put(cb, 0.0); + upper.put(cb, 0.0); + upperOff.put(cb, Double.NEGATIVE_INFINITY); + continue; + } + double bBest = Double.NEGATIVE_INFINITY; + BitSet[] bArg = null; + double bUpper = Double.NEGATIVE_INFINITY; + double bOff = Double.NEGATIVE_INFINITY; + + for (BitSet[] s : backboneSplits(cb)) { + Double l = best.get(s[0]); + Double r = best.get(s[1]); + if (l == null || r == null) { + continue; + } + double theta = logSplitProbability(cb, s[0], s[1], alpha, alpha1, alpha2); + double v = theta + l + r; + if (v > bBest) { + bBest = v; + bArg = s; + } + double ul = upper.get(s[0]); + double ur = upper.get(s[1]); + bUpper = Math.max(bUpper, theta + ul + ur); + double vl = upperOff.get(s[0]); + double vr = upperOff.get(s[1]); + bOff = Math.max(bOff, theta + Math.max(vl + ur, ul + vr)); + } + + // class 3: one child observed (bounded above by U of that child, novel side by 0) + double t3 = logThetaOfClass(cb, 2); + if (t3 > Double.NEGATIVE_INFINITY) { + double bestObservedSide = Double.NEGATIVE_INFINITY; + int m = cb.cardinality(); + for (BitSet d : sortedCladeBits()) { + if (d.cardinality() >= m || !subset(d, cb)) { + continue; + } + BitSet complement = BitSet.newBitSet(cb); + complement.andNot(d); + if (getClade(complement) == null) { // exactly one side observed + Double u = upper.get(d); + if (u != null) { + bestObservedSide = Math.max(bestObservedSide, u); + } + } + } + if (bestObservedSide > Double.NEGATIVE_INFINITY) { + bUpper = Math.max(bUpper, t3 + bestObservedSide); + bOff = Math.max(bOff, t3 + bestObservedSide); + } + } + + // class 4: both children novel, each bounded by 0 + double t4 = logThetaOfClass(cb, 3); + if (t4 > Double.NEGATIVE_INFINITY) { + bUpper = Math.max(bUpper, t4); + bOff = Math.max(bOff, t4); + } + + best.put(cb, bBest); + arg.put(cb, bArg); + upper.put(cb, bUpper); + upperOff.put(cb, bOff); + } + + this.offBackboneBound = upperOff.get(getRootClade().getCladeInBits()); + this.mapArg = arg; + this.mapBest = best; + } + + /** + * Exact MAP over all trees that use no two-novel-clade split, by memoised recursion over + * classes {@code A_0} and {@code A_1}. + * + *

Restricting to {@code A_0} keeps the recursion on the observed-clade DAG. Admitting + * {@code A_1} as well -- peel off an observed clade, leave a novel remainder -- widens the state + * space to clades of the form {@code root} minus a union of disjoint observed clades. That set + * can in principle be large, so the search is capped by {@link #MAP_STATE_BUDGET} distinct + * clades; in practice it stays small because an {@code A_1} split is expensive and the recursion + * only ever descends. + * + *

Returns {@code {best, viaA2}} for the clade: the best log probability using only + * {@code A_0}/{@code A_1} splits, and an upper bound on any subtree that uses an {@code A_2} + * split somewhere. The second is the certificate: if {@code best > viaA2} at the root, no tree + * containing a two-novel-clade split can reach the optimum, so the answer is the global MAP. + */ + private static final long MAP_STATE_BUDGET = + Long.getLong("creg.mapStates", 4_000_000L); + + private static final class BudgetExhausted extends RuntimeException { + BudgetExhausted() { + super(null, null, false, false); + } + } + + private double[] solveFull(BitSet cBits, int a1Budget, List> memos, + long[] states) { + Map memo = memos.get(a1Budget); + double[] cached = memo.get(cBits); + if (cached != null) { + return cached; + } + if (++states[0] > MAP_STATE_BUDGET) { + throw new BudgetExhausted(); + } + int m = cBits.cardinality(); + if (m == 1) { + double[] leaf = {0.0, Double.NEGATIVE_INFINITY}; + memo.put(BitSet.newBitSet(cBits), leaf); + return leaf; + } + double best = Double.NEGATIVE_INFINITY; + double viaA2 = Double.NEGATIVE_INFINITY; + + for (BitSet d : sortedCladeBits()) { + if (d.cardinality() >= m || !subset(d, cBits)) { + continue; + } + BitSet complement = BitSet.newBitSet(cBits); + complement.andNot(d); + boolean complementObserved = getClade(complement) != null; + if (complementObserved && compareBitSets(d, complement) >= 0) { + continue; // both-observed bipartitions are reached from their smaller side only + } + int childBudget = complementObserved ? a1Budget : a1Budget - 1; + if (childBudget < 0) { + continue; // no A_1 split allowance left on this path + } + double theta = logSplitProbability(cBits, d, complement, alpha, alpha1, alpha2); + double[] left = solveFull(d, childBudget, memos, states); + double[] right = solveFull(complement, childBudget, memos, states); + best = Math.max(best, theta + left[0] + right[0]); + double ul = Math.max(left[0], left[1]); + double ur = Math.max(right[0], right[1]); + viaA2 = Math.max(viaA2, theta + Math.max(left[1] + ur, ul + right[1])); + } + + // an A_2 split taken here; both children are novel and bounded above by zero + double[] size = classSizes(cBits); + if (size[3] > 0) { + viaA2 = Math.max(viaA2, logThetaOfClass(cBits, 3)); + } + + double[] value = {best, viaA2}; + memo.put(BitSet.newBitSet(cBits), value); + return value; + } + + /** + * Result of the MAP search. + * + *

{@code a2Excluded} says only that no {@code A_2} split can improve the optimum within + * the searched class of trees, i.e. among trees using at most {@code a1Depth} one-novel-clade + * splits per path. It upgrades to a genuine global certificate ({@code certifiedGlobal}) only + * when that depth was not binding, so that every {@code A_0}/{@code A_1} tree was considered. + */ + public record MapResult(double maxLogProbability, double offBackboneBound, + boolean a2Excluded, int a1Depth, boolean exhaustiveInA1, + long statesExplored, boolean complete) { + + /** True only when the search covered every A_0/A_1 tree and excluded A_2 as well. */ + public boolean certifiedGlobal() { + return complete && exhaustiveInA1 && a2Excluded; + } + } + + /** + * Runs the {@code A_0}/{@code A_1} search and reports whether the optimum it found is provably + * the global MAP. {@code complete} is false when the state budget was exhausted, in which case + * the backbone DP result should be used instead. + */ + public MapResult solveMAP() { + return solveMAP(Integer.MAX_VALUE / 2); + } + + /** + * As {@link #solveMAP()} but allowing at most {@code maxA1} one-novel-clade splits on any + * root-to-leaf path. {@code maxA1 = 0} is the backbone DP; raising it enlarges the search until + * either the optimum stops improving or the state budget is exhausted. + */ + public MapResult solveMAP(int maxA1) { + int cap = Math.min(maxA1, getSizeOfLeavesArray()); + List> memos = new ArrayList<>(); + for (int i = 0; i <= cap; i++) { + memos.add(new HashMap<>()); + } + long[] states = {0}; + boolean exhaustive = cap >= getSizeOfLeavesArray() - 2; + try { + double[] root = solveFull(getRootClade().getCladeInBits(), cap, memos, states); + return new MapResult(root[0], root[1], root[0] > root[1], cap, exhaustive, + states[0], true); + } catch (BudgetExhausted e) { + return new MapResult(getMaxLogTreeProbability(), getOffBackboneBound(), + false, cap, exhaustive, states[0], false); + } + } + + /** Log probability of the backbone MAP tree. */ + @Override + public double getMaxLogTreeProbability() { + computeMAP(); + return mapBest.get(getRootClade().getCladeInBits()); + } + + /** + * Whether the backbone MAP tree is provably the global MAP over all of tree space: true when no + * tree using a class-3 or class-4 split anywhere can reach the backbone optimum. + */ + public boolean isMAPCertifiedGlobal() { + computeMAP(); + return getMaxLogTreeProbability() > offBackboneBound; + } + + /** The certificate's upper bound on any tree that uses an off-backbone split. */ + public double getOffBackboneBound() { + computeMAP(); + return offBackboneBound; + } + + /* ---------------------------------------------------------------------- + * Entropy + * + * The sampler draws from exactly the scored distribution, so E[-log q] is an unbiased estimate + * of H(q) with no truncation to correct for. A deterministic recursion would instead have to + * approximate the subtree entropy of novel clades, so the Monte-Carlo estimator is both simpler + * and more accurate here; only its standard error stands between it and the exact value. + * ------------------------------------------------------------------- */ + + /* ---------------------------------------------------------------------- + * Deterministic entropy recursion + * + * H(C) = H_split(C) + sum_S theta(S) [H(A_S) + H(B_S)], with H(leaf) = 0. + * + * The local term is closed form: within class j >= 2 every member has the same + * theta_j = a_j / (|A_j| Z), so that class contributes -(a_j/Z) log theta_j as a single term -- + * the exponentially large class 4 is never enumerated. The expectation term is exact for + * classes 1 and 2 (both children observed, so the recursion stays on the clade DAG) and needs a + * value for the novel child of a class-3 split and for both children of a class-4 split. + * + * APPROXIMATION: a novel clade is treated as *fresh*, i.e. as containing no observed clades + * other than its singletons, so its subtree entropy depends only on its size and is given by a + * universal g(k) computed once by an O(n^2) recursion. Real novel clades usually do contain + * observed clades, so g overestimates their structure-free entropy; the error enters only + * through the class-3 and class-4 branches, whose total weight at a clade is (a_3 + a_4)/Z. + * Class-4 splits are grouped by the sizes of the two sides, whose counts follow from the + * binomials minus the enumerable classes, keeping the whole pass O(K + m) per clade. + * ------------------------------------------------------------------- */ + + /** g(k): subtree entropy of a fresh (no observed subclades but singletons) clade of size k. */ + private volatile double[] freshEntropy; + + private synchronized double[] freshEntropy() { + if (freshEntropy != null) { + return freshEntropy; + } + int n = getSizeOfLeavesArray(); + double[] g = new double[Math.max(3, n + 1)]; + g[1] = 0.0; + if (g.length > 2) { + g[2] = 0.0; // the single split of a novel cherry has probability one + } + for (int k = 3; k <= n; k++) { + double total = Math.pow(2.0, k - 1) - 1.0; + double n3 = k; // {leaf, rest}, rest unobserved since k-1 >= 2 + double n4 = total - n3; // both sides of size >= 2, so both unobserved + double z = alpha1 + (n4 > 0 ? alpha2 : 0.0); + + double logT3 = Math.log(alpha1) - Math.log(n3) - Math.log(z); + double h = -(alpha1 / z) * logT3; + double e = (alpha1 / z) * g[k - 1]; // class-3 children are {1, k-1} + if (n4 > 0) { + double logT4 = Math.log(alpha2) - Math.log(n4) - Math.log(z); + h -= (alpha2 / z) * logT4; + double weighted = 0.0; + for (int j = 2; j <= k / 2; j++) { + double cnt = binomial(k, j); + if (j == k - j) { + cnt /= 2.0; + } + weighted += cnt * (g[j] + g[k - j]); + } + e += Math.exp(logT4) * weighted; + } + g[k] = h + e; + } + freshEntropy = g; + return g; + } + + /** Local split entropy at {@code cBits}: {@code -sum_S theta(S) log theta(S)}, closed form. */ + private double localSplitEntropy(BitSet cBits) { + double[] size = classSizes(cBits); + double[] wz = logWeightsAndZ(cBits, alpha, alpha1, alpha2); + double z = wz[4]; + Clade c = getClade(cBits); + double h = 0.0; + if (size[0] > 0) { // class 1 is explicit: theta varies with the split count + double perSplit = Math.exp(wz[0]); + for (CladePartition p : c.getPartitions()) { + if (p.getNumberOfOccurrences() == 0) { + continue; + } + double theta = (p.getNumberOfOccurrences() + perSplit) / z; + h -= theta * Math.log(theta); + } + } + for (int j = 1; j < 4; j++) { // classes 2-4 are uniform within the class + if (size[j] > 0) { + double logTheta = wz[j] - Math.log(z); + h -= size[j] * Math.exp(logTheta) * logTheta; + } + } + return h; + } + + /** + * Deterministic entropy in nats, using the fresh-clade approximation for novel subclades. Exact + * whenever no class-3 or class-4 split leads to a novel clade that contains an observed clade. + */ + public double getEntropyRecursive() { + double[] g = freshEntropy(); + List clades = new ArrayList<>(getClades()); + clades.sort(java.util.Comparator.comparingInt(Clade::size)); + Map entropy = new HashMap<>(); + + for (Clade c : clades) { + BitSet cb = c.getCladeInBits(); + int m = c.size(); + if (m == 1) { + entropy.put(cb, 0.0); + continue; + } + double[] size = classSizes(cb); + double[] wz = logWeightsAndZ(cb, alpha, alpha1, alpha2); + double z = wz[4]; + + double e = 0.0; + + // class 1: observed splits, both children observed + if (size[0] > 0) { + double perSplit = Math.exp(wz[0]); + for (CladePartition p : c.getPartitions()) { + if (p.getNumberOfOccurrences() == 0) { + continue; + } + double theta = (p.getNumberOfOccurrences() + perSplit) / z; + e += theta * (entropy.get(p.getChildClades()[0].getCladeInBits()) + + entropy.get(p.getChildClades()[1].getCladeInBits())); + } + } + + // classes 2 and 3, plus the size profile of everything that is not class 4 + double t2 = (size[1] > 0) ? Math.exp(wz[1]) / z : 0.0; + double t3 = (size[2] > 0) ? Math.exp(wz[2]) / z : 0.0; + double[] nonClass4 = new double[m / 2 + 1]; + for (BitSet d : sortedCladeBits()) { + if (d.cardinality() >= m || !subset(d, cb)) { + continue; + } + BitSet complement = BitSet.newBitSet(cb); + complement.andNot(d); + int j = Math.min(d.cardinality(), complement.cardinality()); + if (getClade(complement) != null) { + if (compareBitSets(d, complement) < 0) { + nonClass4[j]++; + if (observedPartition(c, d, complement) == null) { // class 2 + e += t2 * (entropy.get(d) + entropy.get(complement)); + } + } + } else { // class 3: d observed, complement novel + nonClass4[j]++; + e += t3 * (entropy.get(d) + g[complement.cardinality()]); + } + } + + // class 4, grouped by the sizes of the two sides + if (size[3] > 0) { + double logT4 = wz[3] - Math.log(z); + double weighted = 0.0; + for (int j = 1; j <= m / 2; j++) { + double totalPairs = binomial(m, j); + if (j == m - j) { + totalPairs /= 2.0; + } + double count4 = totalPairs - nonClass4[j]; + if (count4 > 0) { + weighted += count4 * (g[j] + g[m - j]); + } + } + e += Math.exp(logT4) * weighted; + } + + entropy.put(cb, localSplitEntropy(cb) + e); + } + return entropy.get(getRootClade().getCladeInBits()); + } + + private static double binomial(int n, int k) { + double r = 1.0; + for (int i = 1; i <= k; i++) { + r = r * (n - k + i) / i; + } + return r; + } + + /** Draws used by {@link #getEntropy()}. */ + public static final int DEFAULT_ENTROPY_SAMPLES = 100_000; + + /** + * Unbiased Monte-Carlo entropy in nats. + * + * @param samples number of draws + * @return {@code {estimate, standard error}} + */ + public double[] getEntropyMonteCarlo(int samples) { + double s1 = 0.0; + double s2 = 0.0; + for (int i = 0; i < samples; i++) { + double logp = sampleTreeLogProbability(); + s1 += -logp; + s2 += logp * logp; + } + double mean = s1 / samples; + double se = Math.sqrt(Math.max(0.0, s2 / samples - mean * mean) / samples); + return new double[]{mean, se}; + } + + /** Monte-Carlo entropy at {@link #DEFAULT_ENTROPY_SAMPLES} draws. */ + @Override + public double getEntropy() { + return getEntropyMonteCarlo(DEFAULT_ENTROPY_SAMPLES)[0]; + } + + /** Not applicable: the Lewis recursion assumes the distribution is supported on the CCD graph. */ + @Override + public double getEntropyLewis() { + throw new UnsupportedOperationException( + "CRegCCD has support outside the CCD graph; use getEntropyMonteCarlo(samples)."); + } + + /* ---------------------------------------------------------------------- + * Sampling + * + * The generative process is the model read forwards: at each clade draw a class with + * probability proportional to (its observed count + a_j) over the non-empty classes, then a + * member uniformly within that class, then recurse into both children. Classes 1-3 are + * explicitly enumerable in one pass over the observed clades; class 4 is drawn by rejection + * from uniform bipartitions, which accepts with probability |A_4| / (2^(m-1) - 1) -- close to + * 1 for any clade large enough for that to matter. + * + * Every drawn split is scored with the same {@link #logSplitProbability} the scorer uses, so + * the sampling distribution equals exp(getLogProbabilityOfTree) by construction: there is no + * truncation and no separate sampling fidelity to choose. + * ------------------------------------------------------------------- */ + + /** Simulates one draw and returns its log probability, without materialising a tree. */ + @Override + public double sampleTreeLogProbability() { + return simulate(getRootClade().getCladeInBits()); + } + + private double simulate(BitSet cBits) { + if (cBits.cardinality() == 1) { + return 0.0; + } + BitSet[] split = sampleSplit(cBits); + return logSplitProbability(cBits, split[0], split[1], alpha, alpha1, alpha2) + + simulate(split[0]) + simulate(split[1]); + } + + /** + * The inherited sampler only ever picks an observed clade partition, so it would draw from the + * observed-splits-only distribution and could never produce a novel clade. Random sampling is + * therefore overridden, as is MAP, which uses this model's own backbone DP rather than the + * inherited CCD1 conditional clade probabilities. + */ + @Override + protected Node getVertexBasedOnStrategy(Clade clade, SamplingStrategy samplingStrategy, + HeightSettingStrategy heightStrategy) { + if (clade.isLeaf()) { + return super.getVertexBasedOnStrategy(clade, samplingStrategy, heightStrategy); + } + if (samplingStrategy == SamplingStrategy.Sampling) { + return sampleVertex(clade.getCladeInBits(), heightStrategy); + } + if (samplingStrategy == SamplingStrategy.MAP) { + computeMAP(); + return mapVertex(clade.getCladeInBits(), heightStrategy); + } + return super.getVertexBasedOnStrategy(clade, samplingStrategy, heightStrategy); + } + + /** Traceback of the backbone MAP DP. Every clade it visits is observed, by construction. */ + private Node mapVertex(BitSet cBits, HeightSettingStrategy heightStrategy) { + if (cBits.cardinality() == 1) { + return super.getVertexBasedOnStrategy(getClade(cBits), + SamplingStrategy.MAP, heightStrategy); + } + BitSet[] split = mapArg.get(cBits); + if (split == null) { + throw new AssertionError("no backbone split for clade " + cBits); + } + Node left = mapVertex(split[0], heightStrategy); + Node right = mapVertex(split[1], heightStrategy); + double logFactor = logSplitProbability(cBits, split[0], split[1], alpha, alpha1, alpha2); + return buildVertex(cBits, left, right, logFactor, heightStrategy); + } + + private Node sampleVertex(BitSet cBits, HeightSettingStrategy heightStrategy) { + if (cBits.cardinality() == 1) { + // leaves are always observed clades, so the inherited leaf construction applies + return super.getVertexBasedOnStrategy(getClade(cBits), + SamplingStrategy.Sampling, heightStrategy); + } + BitSet[] split = sampleSplit(cBits); + Node left = sampleVertex(split[0], heightStrategy); + Node right = sampleVertex(split[1], heightStrategy); + double logFactor = logSplitProbability(cBits, split[0], split[1], alpha, alpha1, alpha2); + return buildVertex(cBits, left, right, logFactor, heightStrategy); + } + + /** Assembles an internal node from two resolved children, stamping the subtree probability. */ + private Node buildVertex(BitSet cBits, Node left, Node right, double logFactor, + HeightSettingStrategy heightStrategy) { + Node vertex = new Node(); + vertex.setNr(nextRunningInnerIndex()); + vertex.addChild(left); + vertex.addChild(right); + + Clade observed = getClade(cBits); + double support = (observed != null) ? observed.getProbability() : 0.0; + vertex.setMetaData(CLADE_SUPPORT_KEY, support); + String posteriorSupport = CLADE_SUPPORT_KEY + "=" + support; + vertex.metaDataString = (vertex.metaDataString != null) + ? vertex.metaDataString + "," + posteriorSupport : posteriorSupport; + + double logP = (Double) left.getMetaData(LOG_PROB_SUBTREE_KEY) + + (Double) right.getMetaData(LOG_PROB_SUBTREE_KEY) + logFactor; + vertex.setMetaData(LOG_PROB_SUBTREE_KEY, logP); + vertex.setMetaData(PROB_SUBTREE_KEY, Math.exp(logP)); + + setSampledHeight(vertex, left, right, observed, heightStrategy); + return vertex; + } + + /** Heights: {@code One} stacks by one; the height strategies use the clade's recorded height + * when it is available and strictly above both children, else a minimal positive increment. */ + private void setSampledHeight(Node vertex, Node left, Node right, Clade observed, + HeightSettingStrategy heightStrategy) { + if (heightStrategy == null || heightStrategy == HeightSettingStrategy.None) { + return; + } + double maxChild = Math.max(left.getHeight(), right.getHeight()); + if (heightStrategy == HeightSettingStrategy.One) { + vertex.setHeight(maxChild + 1.0); + return; + } + double recorded = Double.NaN; + if (observed != null) { + recorded = (heightStrategy == HeightSettingStrategy.CommonAncestorHeights) + ? observed.getCommonAncestorHeight() : observed.getMeanOccurredHeight(); + } + vertex.setHeight(recorded > maxChild ? recorded : maxChild + NOVEL_HEIGHT_EPS); + } + + /** Draws one bipartition of {@code cBits} from this model's conditional distribution. */ + private BitSet[] sampleSplit(BitSet cBits) { + double[] size = classSizes(cBits); + Clade c = getClade(cBits); + double[] wz = logWeightsAndZ(cBits, alpha, alpha1, alpha2); + + // total mass of each class: counts plus its share of the prior + double[] weight = new double[4]; + for (int j = 0; j < 4; j++) { + weight[j] = (size[j] > 0) ? Math.exp(wz[j]) * size[j] : 0.0; + } + if (size[0] > 0) { + weight[0] += c.getNumberOfOccurrences(); + } + + double total = weight[0] + weight[1] + weight[2] + weight[3]; + double target = random().nextDouble() * total; + int cls = -1; + double acc = 0.0; + for (int j = 0; j < 4; j++) { + if (weight[j] <= 0) { + continue; + } + acc += weight[j]; + if (target < acc) { + cls = j; + break; + } + } + if (cls < 0) { // numerical guard: fall back to the last non-empty class + for (int j = 3; j >= 0; j--) { + if (size[j] > 0) { + cls = j; + break; + } + } + } + + return switch (cls) { + case 0 -> sampleObservedSplit(c, size[0]); + case 1, 2 -> sampleEnumerableSplit(cBits, cls); + default -> sampleNovelSplit(cBits, size[3]); + }; + } + + /** Class 1: an observed split, with weight {@code f(S) + a_1/|A_1|}. */ + private BitSet[] sampleObservedSplit(Clade c, double n1) { + double perSplit = Math.exp(logWeightsAndZ(c.getCladeInBits(), alpha, alpha1, alpha2)[0]); + double totalWeight = c.getNumberOfOccurrences() + perSplit * n1; + double target = random().nextDouble() * totalWeight; + double acc = 0.0; + CladePartition last = null; + for (CladePartition p : c.getPartitions()) { + if (p.getNumberOfOccurrences() == 0) { + continue; + } + last = p; + acc += p.getNumberOfOccurrences() + perSplit; + if (target < acc) { + return childBits(p); + } + } + return childBits(last); // numerical guard + } + + private static BitSet[] childBits(CladePartition p) { + return new BitSet[]{p.getChildClades()[0].getCladeInBits(), + p.getChildClades()[1].getCladeInBits()}; + } + + /** + * Classes 2 and 3, drawn uniformly by reservoir sampling over the same single pass across the + * observed clades inside {@code cBits} that produced the class sizes. A both-observed + * bipartition is reached from either side, so it is only considered from its canonically + * smaller side; a one-observed bipartition is reached exactly once, from its observed side. + */ + private BitSet[] sampleEnumerableSplit(BitSet cBits, int cls) { + int m = cBits.cardinality(); + int seen = 0; + BitSet[] pick = null; + for (BitSet d : sortedCladeBits()) { + if (d.cardinality() >= m || !subset(d, cBits)) { + continue; + } + BitSet complement = BitSet.newBitSet(cBits); + complement.andNot(d); + boolean complementObserved = getClade(complement) != null; + boolean candidate; + if (complementObserved) { + candidate = cls == 1 + && compareBitSets(d, complement) < 0 + && observedPartition(getClade(cBits), d, complement) == null; + } else { + candidate = cls == 2; + } + if (candidate) { + seen++; + if (random().nextInt(seen) == 0) { + pick = new BitSet[]{d, complement}; + } + } + } + return pick; + } + + /** + * Class 4, drawn uniformly among bipartitions with neither side an observed clade. Uniform + * bipartitions are generated by pinning the lowest taxon to one side and flipping a fair coin + * for the rest, and rejected unless both sides are novel. When acceptance would be poor the + * bipartition set is necessarily small, so it is enumerated instead. + */ + private BitSet[] sampleNovelSplit(BitSet cBits, double n4) { + int m = cBits.cardinality(); + int[] idx = new int[m]; + int k = 0; + for (int b = cBits.nextSetBit(0); b >= 0; b = cBits.nextSetBit(b + 1)) { + idx[k++] = b; + } + double total = Math.pow(2.0, m - 1) - 1.0; + + if (n4 / total >= MIN_REJECTION_ACCEPTANCE) { + while (true) { + BitSet left = BitSet.newBitSet(leafArraySize); + BitSet right = BitSet.newBitSet(leafArraySize); + left.set(idx[0]); + for (int i = 1; i < m; i++) { + if (random().nextBoolean()) { + left.set(idx[i]); + } else { + right.set(idx[i]); + } + } + if (right.isEmpty()) { + continue; + } + if (getClade(left) == null && getClade(right) == null) { + return new BitSet[]{left, right}; + } + } + } + + // low acceptance => 2^(m-1) is small; enumerate and reservoir-sample + int seen = 0; + BitSet[] pick = null; + for (int mask = 0; mask < (1 << (m - 1)); mask++) { + BitSet left = BitSet.newBitSet(leafArraySize); + BitSet right = BitSet.newBitSet(leafArraySize); + left.set(idx[0]); + for (int i = 1; i < m; i++) { + if ((mask & (1 << (i - 1))) != 0) { + left.set(idx[i]); + } else { + right.set(idx[i]); + } + } + if (right.isEmpty()) { + continue; + } + if (getClade(left) == null && getClade(right) == null) { + seen++; + if (random().nextInt(seen) == 0) { + pick = new BitSet[]{left, right}; + } + } + } + return pick; + } + + private static int compareBitSets(BitSet a, BitSet b) { + int ia = a.nextSetBit(0); + int ib = b.nextSetBit(0); + while (ia >= 0 && ib >= 0) { + if (ia != ib) { + return Integer.compare(ia, ib); + } + ia = a.nextSetBit(ia + 1); + ib = b.nextSetBit(ib + 1); + } + return Integer.compare(ia, ib); + } +} diff --git a/src/test/java/ccd/model/CRegCCDMapEntropyTest.java b/src/test/java/ccd/model/CRegCCDMapEntropyTest.java new file mode 100644 index 0000000..6188288 --- /dev/null +++ b/src/test/java/ccd/model/CRegCCDMapEntropyTest.java @@ -0,0 +1,298 @@ +package ccd.model; + +import beast.base.evolution.tree.Tree; +import beast.base.evolution.tree.TreeParser; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** {@link CRegCCD}: MAP tree (with its global-optimality certificate) and entropy. */ +public class CRegCCDMapEntropyTest { + + private static List taxa(int n) { + List out = new ArrayList<>(); + for (int i = 0; i < n; i++) { + out.add("T" + i); + } + return out; + } + + private static List randomTrees(List taxa, int nTrees, long seed) { + Random rng = new Random(seed); + List out = new ArrayList<>(); + for (int t = 0; t < nTrees; t++) { + List pool = new ArrayList<>(taxa); + while (pool.size() > 1) { + String a = pool.remove(rng.nextInt(pool.size())); + String b = pool.remove(rng.nextInt(pool.size())); + pool.add("(" + a + "," + b + ")"); + } + out.add(new TreeParser(taxa, pool.get(0) + ";", 1, false)); + } + return out; + } + + /** + * The MAP tree must be the argmax over ALL topologies, not just those on the backbone. Checked + * by exhaustive enumeration, together with the certificate that claims it. + */ + @Test + public void mapTreeIsGlobalOptimum() { + for (int n : new int[]{5, 6, 7}) { + for (int nTrees : new int[]{2, 10, 40}) { + List tx = taxa(n); + CRegCCD ccd = new CRegCCD(randomTrees(tx, nTrees, 13L), 0.0, 0.4, 0.4, 0.4); + + double bruteForce = Double.NEGATIVE_INFINITY; + for (Tree t : CRegCCDTest.allRootedTopologies(tx)) { + bruteForce = Math.max(bruteForce, ccd.getLogProbabilityOfTree(t)); + } + double reported = ccd.getMaxLogTreeProbability(); + boolean certified = ccd.isMAPCertifiedGlobal(); + System.out.printf("CRegCCD MAP %d taxa, %2d trees: brute=%.6f DP=%.6f " + + "certified=%-5s (off-backbone bound %.3f)%n", + n, nTrees, bruteForce, reported, certified, ccd.getOffBackboneBound()); + + assertEquals(bruteForce, reported, 1e-9, + "backbone DP must find the global maximum (" + n + " taxa)"); + + // the returned tree must actually attain it + Tree map = ccd.getMAPTree(); + assertEquals(reported, ccd.getLogProbabilityOfTree(map), 1e-9, + "returned MAP tree must attain the reported maximum"); + assertEquals(tx.size(), map.getLeafNodeCount()); + + // the certificate must never be wrong when it fires + if (certified) { + assertTrue(ccd.getOffBackboneBound() < bruteForce + 1e-12, + "certificate claimed optimality but the bound exceeds the optimum"); + } + } + } + } + + @Test + public void entropyMatchesEnumeration() { + List tx = taxa(6); + CRegCCD ccd = new CRegCCD(randomTrees(tx, 8, 17L), 0.0, 0.5, 0.3, 0.2); + ccd.setRandom(new Random(2024L)); + + double exact = 0.0; + for (Tree t : CRegCCDTest.allRootedTopologies(tx)) { + double logp = ccd.getLogProbabilityOfTree(t); + exact -= Math.exp(logp) * logp; + } + double[] mc = ccd.getEntropyMonteCarlo(500_000); + System.out.printf("CRegCCD entropy: exact=%.5f MC=%.5f +/- %.5f (%.1f SE off)%n", + exact, mc[0], mc[1], Math.abs(mc[0] - exact) / mc[1]); + assertEquals(exact, mc[0], Math.max(5 * mc[1], 0.002), "MC entropy must match enumeration"); + + assertTrue(Double.isFinite(ccd.getEntropy()), "default entropy must be finite"); + } + + /* ------------------------------------------------------------------ * + * The manuscript's four-taxon example, worked symbolically. + * Sample: (((A,B),C),D) and (((D,C),B),A), each once. + * ------------------------------------------------------------------ */ + + private static final List TAXA4 = Arrays.asList("A", "B", "C", "D"); + + static CRegCCD exampleModel(double alpha, double alpha1, double alpha2) { + List trees = new ArrayList<>(); + trees.add(new TreeParser(TAXA4, "(((A,B),C),D);", 1, false)); + trees.add(new TreeParser(TAXA4, "(((D,C),B),A);", 1, false)); + return new CRegCCD(trees, 0.0, alpha, alpha1, alpha2); + } + + /** The 15 four-taxon trees, grouped by the six probability categories. */ + static final String[][] CATEGORIES = { + {"(((A,B),C),D);", "(((C,D),B),A);"}, // sampled + {"((A,B),(C,D));"}, // both children observed + {"(((A,C),B),D);", "(((B,C),A),D);", "(((B,D),C),A);", "(((B,C),D),A);"}, // novel inside an observed 3-clade + {"(((C,D),A),B);", "(((A,B),D),C);"}, // root one-observed, reconnects + {"(((A,C),D),B);", "(((A,D),C),B);", "(((A,D),B),C);", "(((B,D),A),C);"}, // root one-observed, novel cherry + {"((A,C),(B,D));", "((A,D),(B,C));"} // root neither observed + }; + + /** POOLED mode: per-split alpha over the three CCD0 splits, class totals a3 and a4. */ + static double[] categoryProbabilities(double alpha, double a3, double a4) { + double z = 2 + 3 * alpha + a3 + a4; + return new double[]{ + (1 + alpha) * (1 + alpha) / (z * (1 + alpha + a3)), + alpha / z, + (1 + alpha) * (a3 / 2) / (z * (1 + alpha + a3)), + (a3 / 2) * alpha / (z * (alpha + a3)), + (a3 / 2) * (a3 / 2) / (z * (alpha + a3)), + (a4 / 2) / z + }; + } + + /** + * How large is the fresh-clade approximation's error? Compares the deterministic recursion + * against exact enumeration across taxon counts, training-set sizes and pseudocounts. + */ + @Test + public void deterministicRecursionErrorVersusEnumeration() { + System.out.printf("%n%-6s %-7s %-22s %-12s %-12s %-11s %-9s%n", + "taxa", "trees", "pseudocounts", "exact H", "recursion", "abs err", "rel err"); + double worstRel = 0.0; + for (int n : new int[]{5, 6, 7, 8}) { + for (int nTrees : new int[]{2, 10, 50}) { + for (double[] p : new double[][]{{0.4, 0.4, 0.4}, {2.0, 0.4, 0.05}, {1.0, 1.0, 1.0}}) { + List tx = taxa(n); + CRegCCD ccd = new CRegCCD(randomTrees(tx, nTrees, 5L), 0.0, + p[0], p[1], p[2]); + double exact = 0.0; + for (Tree t : CRegCCDTest.allRootedTopologies(tx)) { + double logp = ccd.getLogProbabilityOfTree(t); + exact -= Math.exp(logp) * logp; + } + double rec = ccd.getEntropyRecursive(); + double abs = Math.abs(rec - exact); + double rel = abs / exact; + worstRel = Math.max(worstRel, rel); + System.out.printf("%-6d %-7d a=(%.2f,%.2f,%.2f)%-6s %-12.6f %-12.6f %-11.2e %-8.3f%%%n", + n, nTrees, p[0], p[1], p[2], "", exact, rec, abs, 100 * rel); + } + } + } + System.out.printf("worst relative error = %.3f%%%n", 100 * worstRel); + assertTrue(worstRel < 0.25, "recursion should be within 25% of exact, worst was " + worstRel); + } + + @Test + public void fourTaxonExampleMatchesClosedFormPooled() { + for (double[] p : new double[][]{{0.4, 0.4, 0.4}, {0.4, 0.4, 0.05}, {1.0, 0.5, 0.25}}) { + CRegCCD ccd = exampleModel(p[0], p[1], p[2]); + double[] expected = categoryProbabilities(p[0], p[1], p[2]); + double total = 0.0; + int count = 0; + for (int g = 0; g < CATEGORIES.length; g++) { + for (String nwk : CATEGORIES[g]) { + Tree t = new TreeParser(TAXA4, nwk, 1, false); + assertEquals(expected[g], ccd.getProbabilityOfTree(t), 1e-12, + "pooled category " + (g + 1) + " tree " + nwk); + total += ccd.getProbabilityOfTree(t); + count++; + } + } + assertEquals(15, count); + assertEquals(1.0, total, 1e-12, "the 15 probabilities must sum to one"); + System.out.printf("pooled four-taxon example alpha=%.2f a3=%.2f a4=%.2f: " + + "closed form verified, sum = %.12f%n", p[0], p[1], p[2], total); + } + } + + /** + * Under POOLED an observed split must always outrank an expanded one at the same clade, since + * they share a per-split pseudocount and the observed one adds f(S) >= 1. Under SEPARATE that + * can fail. + */ + @Test + public void pooledGuaranteesObservedOutranksExpanded() { + for (int n : new int[]{6, 8}) { + for (int nTrees : new int[]{3, 20}) { + List tx = taxa(n); + List training = randomTrees(tx, nTrees, 77L); + for (double alpha : new double[]{0.05, 0.4, 2.0, 10.0}) { + CRegCCD pooled = new CRegCCD(training, 0.0, alpha, 0.4, 0.05); + for (Clade c : pooled.getClades()) { + if (c.size() < 2) { + continue; + } + double[] size = pooled.classSizes(c.getCladeInBits()); + if (size[0] <= 0 || size[1] <= 0) { + continue; + } + // every expanded split has the same probability; take the largest observed one + double worstObserved = Double.POSITIVE_INFINITY; + for (CladePartition p : c.getPartitions()) { + if (p.getNumberOfOccurrences() > 0) { + worstObserved = Math.min(worstObserved, p.getNumberOfOccurrences()); + } + } + // numerators: observed f + alpha, expanded alpha + assertTrue(worstObserved + alpha > alpha, + "observed must outrank expanded at " + c); + } + } + } + } + System.out.println("pooled: observed splits outrank expanded splits at every clade"); + } + + /** + * regCCD is nested exactly: with no prior mass on the novel classes, CRegCCD's alpha is + * regCCD's additive-alpha smoothing over the CCD0 split set, so the two must agree on every + * tree that regCCD supports. + */ + @Test + public void regCCDIsNestedAtAlphaOneTwoZero() { + for (int n : new int[]{5, 6, 7}) { + for (int nTrees : new int[]{3, 15}) { + List tx = taxa(n); + for (double alpha : new double[]{0.1, 0.4, 1.0}) { + CRegCCD creg = new CRegCCD(randomTrees(tx, nTrees, 31L), 0.0, alpha, 0.0, 0.0); + RegCCD reg = new RegCCD(randomTrees(tx, nTrees, 31L), 0.0, alpha); + int compared = 0; + double worst = 0.0; + for (Tree t : CRegCCDTest.allRootedTopologies(tx)) { + double a = reg.getLogProbabilityOfTree(t); + if (!Double.isFinite(a)) { + continue; // outside regCCD's support + } + worst = Math.max(worst, Math.abs(a - creg.getLogProbabilityOfTree(t))); + compared++; + } + System.out.printf("regCCD nesting %d taxa, %2d trees, alpha=%.1f: " + + "%d trees compared, max |diff| = %.2e%n", n, nTrees, alpha, compared, worst); + assertTrue(compared > 0, "regCCD must support some trees"); + assertEquals(0.0, worst, 1e-9, + "CRegCCD(alpha, 0, 0) must equal regCCD(alpha)"); + } + } + } + } + + /** + * The A_0/A_1 search must equal the true global optimum by enumeration, and its certificate + * must never claim optimality wrongly. + */ + @Test + public void exactMapSearchMatchesEnumeration() { + int certified = 0, total = 0; + for (int n : new int[]{5, 6, 7, 8}) { + for (int nTrees : new int[]{2, 5, 20}) { + List tx = taxa(n); + CRegCCD ccd = new CRegCCD(randomTrees(tx, nTrees, 13L), 0.0, 0.4, 0.4, 0.05); + double brute = Double.NEGATIVE_INFINITY; + for (Tree t : CRegCCDTest.allRootedTopologies(tx)) { + brute = Math.max(brute, ccd.getLogProbabilityOfTree(t)); + } + CRegCCD.MapResult r = ccd.solveMAP(); + total++; + if (r.certifiedGlobal()) { + certified++; + } + System.out.printf("%d taxa, %2d trees: brute=%.6f A0/A1=%.6f bound=%8.3f " + + "certified=%-5s states=%d%n", + n, nTrees, brute, r.maxLogProbability(), r.offBackboneBound(), + r.certifiedGlobal(), r.statesExplored()); + assertTrue(r.complete(), "search must complete within the state budget"); + assertEquals(brute, r.maxLogProbability(), 1e-9, + "A_0/A_1 search must find the global optimum"); + if (r.certifiedGlobal()) { + assertTrue(r.offBackboneBound() < brute + 1e-12, + "certificate must not claim optimality when the bound exceeds it"); + } + } + } + System.out.printf("certificate fired in %d of %d configurations%n", certified, total); + } +} diff --git a/src/test/java/ccd/model/CRegCCDTest.java b/src/test/java/ccd/model/CRegCCDTest.java new file mode 100644 index 0000000..08cfdb6 --- /dev/null +++ b/src/test/java/ccd/model/CRegCCDTest.java @@ -0,0 +1,293 @@ +package ccd.model; + +import beast.base.evolution.tree.Tree; +import beast.base.evolution.tree.TreeParser; +import ccd.model.bitsets.BitSet; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** {@link CRegCCD}: exact normalisation, class-size arithmetic, and full support. */ +public class CRegCCDTest { + + private static List taxa(int n) { + List out = new ArrayList<>(); + for (int i = 0; i < n; i++) { + out.add("T" + i); + } + return out; + } + + private static List randomTrees(List taxa, int nTrees, long seed) { + Random rng = new Random(seed); + List out = new ArrayList<>(); + for (int t = 0; t < nTrees; t++) { + List pool = new ArrayList<>(taxa); + while (pool.size() > 1) { + String a = pool.remove(rng.nextInt(pool.size())); + String b = pool.remove(rng.nextInt(pool.size())); + pool.add("(" + a + "," + b + ")"); + } + out.add(new TreeParser(taxa, pool.get(0) + ";", 1, false)); + } + return out; + } + + private static double totalMass(CRegCCD ccd, List taxa) { + double sum = 0.0; + for (Tree t : allRootedTopologies(taxa)) { + sum += Math.exp(ccd.getLogProbabilityOfTree(t)); + } + return sum; + } + + @Test + public void exactlyNormalisedOverTreeSpace() { + for (int n : new int[]{4, 5, 6, 7}) { + List tx = taxa(n); + for (int nTrees : new int[]{1, 3, 20}) { + List training = randomTrees(tx, nTrees, 7L); + CRegCCD ccd = new CRegCCD(training, 0.0, 0.4, 0.4, 0.4); + double mass = totalMass(ccd, tx); + System.out.printf("CRegCCD %d taxa, %2d training trees: totalMass = %.12f%n", + n, nTrees, mass); + assertEquals(1.0, mass, 1e-9, + "CRegCCD must be exactly normalised (" + n + " taxa, " + nTrees + " trees)"); + } + } + } + + @Test + public void normalisedAcrossPseudocountChoices() { + List tx = taxa(6); + List training = randomTrees(tx, 10, 11L); + CRegCCD ccd = new CRegCCD(training, 0.0); + for (double[] p : new double[][]{ + {0.0, 0.4, 0.4, 0.4}, + {1.0, 1.0, 1.0, 1.0}, + {0.0, 2.0, 0.5, 0.05}, + {0.3, 0.01, 5.0, 0.2}}) { + double sum = 0.0; + for (Tree t : allRootedTopologies(tx)) { + sum += Math.exp(ccd.getLogProbabilityOfTree(t, p[1], p[2], p[3])); + } + System.out.printf("CRegCCD 6 taxa a=(%.2f,%.2f,%.2f,%.2f): totalMass = %.12f%n", + p[0], p[1], p[2], p[3], sum); + assertEquals(1.0, sum, 1e-9, "normalisation must hold for any pseudocounts"); + } + } + + @Test + public void classSizesPartitionAllBipartitions() { + List tx = taxa(8); + List training = randomTrees(tx, 50, 3L); + CRegCCD ccd = new CRegCCD(training, 0.0); + CCD0 ccd0 = new CCD0(training, 0); + int checked = 0; + for (Clade c : ccd.getClades()) { + if (c.size() < 2) { + continue; + } + double[] s = ccd.classSizes(c.getCladeInBits()); + double total = Math.pow(2.0, c.size() - 1) - 1.0; + assertEquals(total, s[0] + s[1] + s[2] + s[3], 1e-6, + "class sizes must partition all bipartitions of " + c); + + // |A_1| + |A_2| is exactly the CCD0 split count (observed + expanded) + Clade c0 = ccd0.getClade(c.getCladeInBits()); + if (c0 != null) { + assertEquals(c0.getPartitions().size(), (int) Math.round(s[0] + s[1]), + "|A_1|+|A_2| must equal the CCD0 partition count for " + c); + } + checked++; + } + System.out.printf("class-size arithmetic verified on %d clades%n", checked); + assertTrue(checked > 0); + } + + @Test + public void fullSupportOnHeldOutTrees() { + List tx = taxa(30); + List training = randomTrees(tx, 100, 5L); + List heldOut = randomTrees(tx, 100, 99L); + CRegCCD ccd = new CRegCCD(training, 0.0); + int covered = 0; + double sum = 0.0; + for (Tree t : heldOut) { + double lp = ccd.getLogProbabilityOfTree(t); + if (Double.isFinite(lp)) { + covered++; + sum += lp; + } + } + System.out.printf("CRegCCD 30 taxa: %d/%d held-out trees with finite logP, mean = %.2f%n", + covered, heldOut.size(), sum / covered); + assertEquals(heldOut.size(), covered, "every held-out tree must have positive probability"); + } + + /** + * If the simulator draws from q AND reports the correct log q, then the mean sampled -log q + * equals the entropy computed by enumeration with the scorer. This checks the sampler and the + * scorer are the same distribution without needing per-topology counts. + */ + @Test + public void samplerEntropyMatchesScorer() { + List tx = taxa(6); + List training = randomTrees(tx, 8, 17L); + CRegCCD ccd = new CRegCCD(training, 0.0, 0.5, 0.3, 0.2); + ccd.setRandom(new Random(31337L)); + + double mass = 0.0; + double h = 0.0; + for (Tree t : allRootedTopologies(tx)) { + double logp = ccd.getLogProbabilityOfTree(t); + double p = Math.exp(logp); + mass += p; + h -= p * logp; + } + assertEquals(1.0, mass, 1e-9, "scored distribution must be normalised"); + + int n = 1_000_000; + double s1 = 0.0; + double s2 = 0.0; + for (int i = 0; i < n; i++) { + double logp = ccd.sampleTreeLogProbability(); + s1 += -logp; + s2 += logp * logp; + } + double hHat = s1 / n; + double se = Math.sqrt(Math.max(0, s2 / n - hHat * hHat) / n); + System.out.printf("CRegCCD sampler: H_enum=%.5f H_MC=%.5f +/- %.5f (%.1f SE off)%n", + h, hHat, se, Math.abs(hHat - h) / se); + assertEquals(h, hHat, Math.max(5 * se, 0.005), + "sampler entropy must match the scorer's enumerated entropy"); + } + + /** + * The direct check: sampled topology frequencies must match the scored probabilities. Compares + * every topology whose expected count is large enough for a normal approximation. + */ + @Test + public void sampledFrequenciesMatchScoredProbabilities() { + List tx = taxa(5); + List training = randomTrees(tx, 5, 23L); + CRegCCD ccd = new CRegCCD(training, 0.0, 0.5, 0.3, 0.2); + ccd.setRandom(new Random(4242L)); + + java.util.Map expected = new java.util.HashMap<>(); + for (Tree t : allRootedTopologies(tx)) { + expected.put(canonical(t), Math.exp(ccd.getLogProbabilityOfTree(t))); + } + + int n = 1_000_000; + java.util.Map counts = new java.util.HashMap<>(); + for (int i = 0; i < n; i++) { + counts.merge(canonical(ccd.sampleTree()), 1, Integer::sum); + } + + int checked = 0; + double worst = 0.0; + String worstKey = null; + for (var e : expected.entrySet()) { + double p = e.getValue(); + if (n * p < 30) { + continue; // too rare for a normal approximation + } + int obs = counts.getOrDefault(e.getKey(), 0); + double z = Math.abs(obs - n * p) / Math.sqrt(n * p * (1 - p)); + if (z > worst) { + worst = z; + worstKey = e.getKey(); + } + checked++; + } + System.out.printf("CRegCCD frequencies: %d topologies checked, worst |z| = %.2f (%s)%n", + checked, worst, worstKey); + assertTrue(checked >= 10, "expected a reasonable number of comparable topologies"); + assertTrue(worst < 4.5, "sampled frequencies must match scored probabilities, worst z = " + worst); + + // no sampled topology may fall outside the enumerated support + for (String key : counts.keySet()) { + assertTrue(expected.containsKey(key), "sampler produced an unrecognised topology " + key); + } + } + + @Test + public void sampledTreesAreValidAndSelfConsistent() { + List tx = taxa(20); + List training = randomTrees(tx, 50, 8L); + CRegCCD ccd = new CRegCCD(training, 0.0); + ccd.setRandom(new Random(5L)); + for (int i = 0; i < 200; i++) { + Tree t = ccd.sampleTree(); + assertEquals(tx.size(), t.getLeafNodeCount(), "sampled tree must have every taxon"); + assertEquals(2 * tx.size() - 1, t.getNodeCount(), "sampled tree must be binary"); + double stamped = (Double) t.getRoot().getMetaData(CCD1.LOG_PROB_SUBTREE_KEY); + assertEquals(ccd.getLogProbabilityOfTree(t), stamped, 1e-9, + "stamped log probability must equal the scorer's"); + } + System.out.println("CRegCCD: 200 sampled 20-taxon trees valid and self-consistent"); + } + + /** Canonical topology key: nested sorted taxon-index sets. */ + private static String canonical(Tree t) { + return canonical(t.getRoot()); + } + + private static String canonical(beast.base.evolution.tree.Node v) { + if (v.isLeaf()) { + return String.valueOf(v.getNr()); + } + String a = canonical(v.getChild(0)); + String b = canonical(v.getChild(1)); + return (a.compareTo(b) <= 0) ? "(" + a + "," + b + ")" : "(" + b + "," + a + ")"; + } + + /* --------------------------------------------------------------------- */ + + static List allRootedTopologies(List taxa) { + List trees = new ArrayList<>(); + for (String shape : shapes(taxa)) { + trees.add(new TreeParser(taxa, shape + ";", 1, false)); + } + return trees; + } + + private static List shapes(List taxa) { + List out = new ArrayList<>(); + if (taxa.size() == 1) { + out.add(taxa.get(0) + ":1"); + return out; + } + String first = taxa.get(0); + List rest = taxa.subList(1, taxa.size()); + int n = rest.size(); + for (int mask = 0; mask < (1 << n); mask++) { + List left = new ArrayList<>(); + left.add(first); + List right = new ArrayList<>(); + for (int i = 0; i < n; i++) { + if ((mask & (1 << i)) != 0) { + left.add(rest.get(i)); + } else { + right.add(rest.get(i)); + } + } + if (right.isEmpty()) { + continue; + } + for (String l : shapes(left)) { + for (String r : shapes(right)) { + out.add("(" + l + "," + r + "):1"); + } + } + } + return out; + } +} From fd365152a26c853241f9db333517811b56ecf6fa Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Tue, 18 Aug 2026 12:59:19 +1200 Subject: [PATCH 2/2] Add tests measuring KRegCCD normalisation and comparing the models on real data KRegNormalisationTest enumerates every rooted topology and asserts that KRegCCD is exactly normalised on four taxa but sub-normalised beyond it, with the deficit scaling as mu^2 and never exceeding one. This is the Theta(mu^2) maximality deficit, which is distinct from, and larger than, the O(mu^(k+1)) reserve truncation. SplitClassSizeAnalysis reports the four split-class sizes at a root clade, showing that a constant per-split pseudocount leaves the observed splits 6.3e-9 of the probability by 40 taxa. ClassUsageAnalysis attributes a held-out tree's log probability to the split classes and measures how much observed structure a two-novel-clade split destroys. RealDataHeadToHeadTest compares CCD1, regCCD, KRegCCD, MRegCCD and CRegCCD on a real posterior using the manuscript's RSV2 protocol, selecting each model's hyperparameters on a validation split disjoint from the fitted set. Co-Authored-By: Claude Opus 5 (1M context) --- .../java/ccd/model/ClassUsageAnalysis.java | 227 +++++++++++ .../java/ccd/model/KRegNormalisationTest.java | 172 +++++++++ .../ccd/model/RealDataHeadToHeadTest.java | 353 ++++++++++++++++++ .../ccd/model/SplitClassSizeAnalysis.java | 110 ++++++ 4 files changed, 862 insertions(+) create mode 100644 src/test/java/ccd/model/ClassUsageAnalysis.java create mode 100644 src/test/java/ccd/model/KRegNormalisationTest.java create mode 100644 src/test/java/ccd/model/RealDataHeadToHeadTest.java create mode 100644 src/test/java/ccd/model/SplitClassSizeAnalysis.java diff --git a/src/test/java/ccd/model/ClassUsageAnalysis.java b/src/test/java/ccd/model/ClassUsageAnalysis.java new file mode 100644 index 0000000..42ca47c --- /dev/null +++ b/src/test/java/ccd/model/ClassUsageAnalysis.java @@ -0,0 +1,227 @@ +package ccd.model; + +import beast.base.evolution.tree.Node; +import beast.base.evolution.tree.Tree; +import beast.base.evolution.tree.TreeParser; +import beastfx.app.treeannotator.TreeAnnotator; +import ccd.model.bitsets.BitSet; +import ccd.tools.CCDToolUtil; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.Test; + +import java.io.File; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Where does a real held-out tree's probability actually go under {@link CRegCCD}? + * + *

Walks every internal node of every held-out tree, classifies its split into the four classes, + * and tallies how many nodes fall in each class and how much log probability each class contributes. + * This measures the concern that a class-4 split (neither child observed) reconnects to the CCD only + * by chance: if class 4 is both rare and responsible for a large share of the total loss, the + * uniform-within-class-4 prior is the binding weakness. + * + *

Also reports, for each class-4 node encountered, how many observed clades survive intact inside + * the two novel children -- i.e. how much backbone a class-4 split destroys. + */ +public class ClassUsageAnalysis { + + private static final String PATH = System.getProperty("ccd.trees", ""); + private static final int N = Integer.parseInt(System.getProperty("ccd.n", "1000")); + + private static List newickCache; + private static List taxaCache; + + private static void load() throws Exception { + if (newickCache != null) { + return; + } + TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(PATH, 10); + ts.reset(); + List nwk = new ArrayList<>(); + List taxa = null; + while (ts.hasNext()) { + Tree t = ts.next(); + if (taxa == null) { + String[] byNr = new String[t.getLeafNodeCount()]; + for (Node leaf : t.getExternalNodes()) { + byNr[leaf.getNr()] = leaf.getID(); + } + taxa = new ArrayList<>(List.of(byNr)); + } + nwk.add(t.getRoot().toNewick() + ";"); + } + newickCache = nwk; + taxaCache = taxa; + } + + private static List read(int count, double from, double to) throws Exception { + load(); + List pool = newickCache.subList((int) (from * newickCache.size()), + (int) (to * newickCache.size())); + List out = new ArrayList<>(); + double step = Math.max(1.0, pool.size() / (double) count); + for (int i = 0; i < count && (int) (i * step) < pool.size(); i++) { + out.add(new TreeParser(taxaCache, pool.get((int) (i * step)), 0, false)); + } + return out; + } + + @Test + public void classUsageOnHeldOutTrees() throws Exception { + Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(), + "set -Dccd.trees=/path/to/x.trees"); + + CRegCCD ccd = new CRegCCD(read(N, 0.0, 0.5), 0.0, + Double.parseDouble(System.getProperty("ccd.alpha", "0.4")), + Double.parseDouble(System.getProperty("ccd.alpha1", "0.4")), + Double.parseDouble(System.getProperty("ccd.alpha2", "0.05"))); + List test = read(N, 0.5, 1.0); + + long[] nodes = new long[4]; + double[] logp = new double[4]; + long class4Nodes = 0; + long survivingObserved = 0; + long shatteredObserved = 0; + + for (Tree t : test) { + Map bits = new HashMap<>(); + computeBits(t.getRoot(), bits, ccd.getSizeOfLeavesArray()); + for (Node v : t.getNodesAsArray()) { + if (v.isLeaf()) { + continue; + } + BitSet cb = bits.get(v); + BitSet ab = bits.get(v.getChildren().get(0)); + BitSet bb = bits.get(v.getChildren().get(1)); + int cls = ccd.splitClass(cb, ab, bb); + nodes[cls]++; + logp[cls] += ccd.logSplitProbability(cb, ab, bb, + ccd.getAlpha(), ccd.getAlpha1(), ccd.getAlpha2()); + if (cls == 3) { + class4Nodes++; + // how much observed structure did this split preserve vs destroy? + for (Clade obs : ccd.getClades()) { + if (obs.size() < 2 || obs.size() >= cb.cardinality()) { + continue; + } + BitSet o = obs.getCladeInBits(); + BitSet tmp = BitSet.newBitSet(o); + tmp.andNot(cb); + if (!tmp.isEmpty()) { + continue; // not inside this clade at all + } + if (subset(o, ab) || subset(o, bb)) { + survivingObserved++; + } else { + shatteredObserved++; + } + } + } + } + } + + long totalNodes = nodes[0] + nodes[1] + nodes[2] + nodes[3]; + double totalLogp = logp[0] + logp[1] + logp[2] + logp[3]; + System.out.printf("%n=== %s: class usage over %d held-out trees (%s) ===%n", + new File(PATH).getName(), test.size(), ccd); + System.out.printf("%-28s %10s %8s %14s %10s %12s%n", + "class", "nodes", "% nodes", "total logP", "% logP", "mean logP"); + String[] names = {"1 observed split", "2 both children obs.", + "3 one child observed", "4 neither observed"}; + for (int j = 0; j < 4; j++) { + System.out.printf("%-28s %10d %7.2f%% %14.1f %9.2f%% %12.3f%n", + names[j], nodes[j], 100.0 * nodes[j] / totalNodes, logp[j], + 100.0 * logp[j] / totalLogp, nodes[j] == 0 ? 0 : logp[j] / nodes[j]); + } + System.out.printf("total mean logP per tree = %.2f%n", totalLogp / test.size()); + if (class4Nodes > 0) { + long tot = survivingObserved + shatteredObserved; + System.out.printf("class-4 splits: %d; observed clades below them: %d intact (%.1f%%), " + + "%d shattered (%.1f%%)%n", + class4Nodes, survivingObserved, 100.0 * survivingObserved / tot, + shatteredObserved, 100.0 * shatteredObserved / tot); + } else { + System.out.println("no class-4 splits occurred in any held-out tree"); + } + } + + private static boolean subset(BitSet a, BitSet c) { + BitSet tmp = BitSet.newBitSet(a); + tmp.andNot(c); + return tmp.isEmpty(); + } + + private static BitSet computeBits(Node v, Map bits, int leafArraySize) { + BitSet b = BitSet.newBitSet(leafArraySize); + if (v.isLeaf()) { + b.set(v.getNr()); + } else { + b.or(computeBits(v.getChildren().get(0), bits, leafArraySize)); + b.or(computeBits(v.getChildren().get(1), bits, leafArraySize)); + } + bits.put(v, b); + return b; + } + + + /** + * Size profile of the splits that introduce two novel clades: parent size and the sizes of the + * two novel children. If the smaller side is consistently small, grading class 2 by the size of + * the smaller side would concentrate its mass where the real novelty is, instead of on the + * balanced splits that dominate a uniform draw. + */ + @Test + public void novelSplitSizeProfile() throws Exception { + Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(), + "set -Dccd.trees=/path/to/x.trees"); + CRegCCD ccd = new CRegCCD(read(N, 0.0, 0.5), 0.0, + Double.parseDouble(System.getProperty("ccd.alpha", "0.4")), + Double.parseDouble(System.getProperty("ccd.alpha1", "0.4")), + Double.parseDouble(System.getProperty("ccd.alpha2", "0.05"))); + List test = read(N, 0.5, 1.0); + + System.out.printf("%n=== %s: size profile of two-novel-clade splits ===%n", + new File(PATH).getName()); + System.out.printf("%-8s %-10s %-10s %-14s %-16s%n", + "parent m", "small side", "large side", "small/parent", "uniform E[small]"); + int count = 0; + double sumFrac = 0.0; + java.util.Map smallSizes = new java.util.TreeMap<>(); + for (Tree t : test) { + Map bits = new HashMap<>(); + computeBits(t.getRoot(), bits, ccd.getSizeOfLeavesArray()); + for (Node v : t.getNodesAsArray()) { + if (v.isLeaf()) { + continue; + } + BitSet cb = bits.get(v); + BitSet ab = bits.get(v.getChildren().get(0)); + BitSet bb = bits.get(v.getChildren().get(1)); + if (ccd.splitClass(cb, ab, bb) != 3) { + continue; + } + int m = cb.cardinality(); + int small = Math.min(ab.cardinality(), bb.cardinality()); + int large = Math.max(ab.cardinality(), bb.cardinality()); + count++; + sumFrac += small / (double) m; + smallSizes.merge(small, 1, Integer::sum); + if (count <= 40) { + System.out.printf("%-8d %-10d %-10d %-14.3f %-16.1f%n", + m, small, large, small / (double) m, m / 2.0); + } + } + } + if (count == 0) { + System.out.println("no two-novel-clade splits in the held-out set"); + return; + } + System.out.printf("%d such splits; mean smaller-side fraction = %.3f " + + "(a uniform bipartition would give ~0.5)%n", count, sumFrac / count); + System.out.println("distribution of smaller-side size: " + smallSizes); + } +} diff --git a/src/test/java/ccd/model/KRegNormalisationTest.java b/src/test/java/ccd/model/KRegNormalisationTest.java new file mode 100644 index 0000000..7043cc6 --- /dev/null +++ b/src/test/java/ccd/model/KRegNormalisationTest.java @@ -0,0 +1,172 @@ +package ccd.model; + +import beast.base.evolution.tree.Tree; +import beast.base.evolution.tree.TreeParser; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Total probability mass of {@link KRegCCD}, by brute-force enumeration of every rooted topology. + * + *

Documents the model's normalisation exactly as the manuscript states it: + *

    + *
  • on four taxa the model is exactly normalised (no blue-region boundary part is itself a + * reserving clade, so the region decomposition is tight);
    • + *
  • from six taxa on it is sub-normalised -- total mass is below 1, never above -- + * and the deficit scales as {@code mu^2}, which is the maximality deficit rather than the + * {@code O(mu^(k+1))} reserve truncation (it persists at full reserve depth).
    • + *
+ */ +public class KRegNormalisationTest { + + private static final List TAXA4 = Arrays.asList("A", "B", "C", "D"); + private static final List TAXA5 = Arrays.asList("A", "B", "C", "D", "E"); + private static final List TAXA6 = Arrays.asList("A", "B", "C", "D", "E", "F"); + private static final List TAXA7 = Arrays.asList("A", "B", "C", "D", "E", "F", "G"); + + private static List trees(List taxa, String... newicks) { + List out = new ArrayList<>(); + for (String nwk : newicks) { + out.add(new TreeParser(taxa, nwk, 1, false)); + } + return out; + } + + private static List training4() { + return trees(TAXA4, "(((A:1,B:1):1,C:1):1,D:1):0;", "((A:1,B:1):1,(C:1,D:1):1):0;"); + } + + private static List training6() { + return trees(TAXA6, + "(((((A:1,B:1):1,C:1):1,D:1):1,E:1):1,F:1):0;", + "((((D:1,C:1):1,B:1):1,A:1):1,(E:1,F:1):1):0;"); + } + + private static List training7() { + return trees(TAXA7, + "((((((A:1,B:1):1,C:1):1,D:1):1,E:1):1,F:1):1,G:1):0;", + "(((((D:1,C:1):1,B:1):1,A:1):1,(E:1,F:1):1):1,G:1):0;"); + } + + /** Total mass under the full-support score. */ + private static double totalMass(KRegCCD ccd, List taxa) { + double mass = 0.0; + for (Tree t : allRootedTopologies(taxa)) { + mass += Math.exp(ccd.getLogProbabilityOfTree(t)); + } + return mass; + } + + @Test + public void exactlyNormalisedOnFourTaxa() { + for (double mu : new double[]{0.001, 0.005, 0.05}) { + KRegCCD ccd = new KRegCCD(training4(), 0.0, mu, 0.4, 2, KRegCCD.TailMode.NONE, + KRegCCD.NovelMode.FLAT); + double mass = totalMass(ccd, TAXA4); + System.out.printf("KRegCCD 4 taxa mu=%-6.3f totalMass = %.12f%n", mu, mass); + assertEquals(1.0, mass, 1e-9, "four-taxon model must normalise exactly"); + } + } + + /** Five taxa: exactness is not a taxon-count property but depends on whether any escape region + * has a reserving clade on its boundary, which the training set determines. */ + @Test + public void fiveTaxaNormalisationDependsOnTrainingSet() { + List caterpillar = trees(TAXA5, "((((A:1,B:1):1,C:1):1,D:1):1,E:1):0;"); + List mixed = trees(TAXA5, + "((((A:1,B:1):1,C:1):1,D:1):1,E:1):0;", + "(((A:1,B:1):1,(C:1,D:1):1):1,E:1):0;"); + for (double mu : new double[]{0.005, 0.05}) { + for (String name : new String[]{"caterpillar", "mixed"}) { + List training = name.equals("caterpillar") ? caterpillar : mixed; + KRegCCD ccd = new KRegCCD(training, 0.0, mu, 0.4, 2, KRegCCD.TailMode.NONE, + KRegCCD.NovelMode.FLAT); + double mass = totalMass(ccd, TAXA5); + System.out.printf("KRegCCD 5 taxa mu=%-6.3f %-12s totalMass = %.9f (1-mass = %+.3e)%n", + mu, name, mass, 1.0 - mass); + assertTrue(mass <= 1.0 + 1e-12, "must never super-normalise, got " + mass); + } + } + } + + @Test + public void subNormalisedFromSixTaxaWithMuSquaredDeficit() { + for (List taxa : List.of(TAXA6, TAXA7)) { + List training = (taxa == TAXA6) ? training6() : training7(); + double prevDeficit = Double.NaN; + double prevMu = Double.NaN; + for (double mu : new double[]{0.05, 0.005, 0.001}) { + for (KRegCCD.TailMode tm : KRegCCD.TailMode.values()) { + KRegCCD ccd = new KRegCCD(training, 0.0, mu, 0.4, 2, tm, + KRegCCD.NovelMode.FLAT); + double mass = totalMass(ccd, taxa); + double deficit = 1.0 - mass; + System.out.printf("KRegCCD %d taxa mu=%-6.3f %-8s totalMass = %.9f (1-mass = %+.3e)%n", + taxa.size(), mu, tm, mass, deficit); + assertTrue(mass <= 1.0 + 1e-12, + "model must never super-normalise, got " + mass); + if (tm == KRegCCD.TailMode.NONE) { + if (!Double.isNaN(prevDeficit)) { + // a mu^2 deficit shrinks by the square of the mu ratio + double ratio = prevDeficit / deficit; + double expected = (prevMu / mu) * (prevMu / mu); + System.out.printf(" deficit shrank %.1fx as mu fell %.0fx " + + "(mu^2 predicts %.0fx)%n", + ratio, prevMu / mu, expected); + assertEquals(expected, ratio, 0.25 * expected, + "deficit must scale as mu^2"); + } + prevDeficit = deficit; + prevMu = mu; + } + } + } + } + } + + private static List allRootedTopologies(List taxa) { + List trees = new ArrayList<>(); + for (String shape : shapes(taxa)) { + trees.add(new TreeParser(taxa, shape + ";", 1, false)); + } + return trees; + } + + private static List shapes(List taxa) { + List out = new ArrayList<>(); + if (taxa.size() == 1) { + out.add(taxa.get(0) + ":1"); + return out; + } + String first = taxa.get(0); + List rest = taxa.subList(1, taxa.size()); + int n = rest.size(); + for (int mask = 0; mask < (1 << n); mask++) { + List left = new ArrayList<>(); + left.add(first); + List right = new ArrayList<>(); + for (int i = 0; i < n; i++) { + if ((mask & (1 << i)) != 0) { + left.add(rest.get(i)); + } else { + right.add(rest.get(i)); + } + } + if (right.isEmpty()) { + continue; + } + for (String l : shapes(left)) { + for (String r : shapes(right)) { + out.add("(" + l + "," + r + "):1"); + } + } + } + return out; + } +} diff --git a/src/test/java/ccd/model/RealDataHeadToHeadTest.java b/src/test/java/ccd/model/RealDataHeadToHeadTest.java new file mode 100644 index 0000000..df52b25 --- /dev/null +++ b/src/test/java/ccd/model/RealDataHeadToHeadTest.java @@ -0,0 +1,353 @@ +package ccd.model; + +import beast.base.evolution.tree.Tree; +import beastfx.app.treeannotator.TreeAnnotator; +import ccd.tools.CCDToolUtil; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.Test; + +import java.io.File; +import java.util.ArrayList; +import java.util.List; + +/** + * Head-to-head held-out predictive comparison of the regularised CCD variants on a real posterior + * tree sample. + * + *

Protocol follows the manuscript's RSV2 comparison: the model is trained on trees drawn from the + * first half of the chain and scored on trees from the second half, so the test trees are genuinely + * out of training. Each model's hyperparameters are selected on an inner fit/validation split of the + * training half alone, then the model is rebuilt on the whole training half and scored on the test + * half. Reported per model: support coverage (fraction of test trees with positive probability), mean + * log probability over all test trees, and -- for the full-support models -- the paired per-tree + * comparison against KRegCCD. + * + *

Point the test at a tree file with {@code -Dccd.trees=/path/to/x.trees}; it is skipped when no + * file is given. {@code -Dccd.n=1000} sets the number of training and test trees. + */ +public class RealDataHeadToHeadTest { + + private static final String PATH = System.getProperty("ccd.trees", ""); + private static final int N = Integer.parseInt(System.getProperty("ccd.n", "1000")); + private static final double BURNIN_PERCENT = 10; + + private interface Scorer { + double logP(Tree t); + } + + private static List newickCache; + private static List taxaCache; + + /** Parses the tree file once, keeping the topologies as newick strings. */ + private static void load() throws Exception { + if (newickCache != null) { + return; + } + TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(PATH, (int) BURNIN_PERCENT); + ts.reset(); + List nwk = new ArrayList<>(); + List taxa = null; + while (ts.hasNext()) { + Tree t = ts.next(); + if (taxa == null) { + taxa = new ArrayList<>(); + String[] byNr = new String[t.getLeafNodeCount()]; + for (beast.base.evolution.tree.Node leaf : t.getExternalNodes()) { + byNr[leaf.getNr()] = leaf.getID(); + } + taxa.addAll(List.of(byNr)); + } + nwk.add(t.getRoot().toNewick() + ";"); + } + newickCache = nwk; + taxaCache = taxa; + System.out.printf("loaded %d trees, %d taxa from %s%n", + nwk.size(), taxa.size(), new File(PATH).getName()); + } + + /** + * {@code count} evenly spaced trees from the chain segment {@code [from, to)} (as fractions of the + * post-burn-in chain), freshly parsed on every call because the CCD constructors take ownership of + * the trees they are given. + * + *

Segments must be disjoint for hyperparameter selection to be honest: scoring trees that are + * also in the fitted set drives every regularisation parameter to zero, because the backbone + * already fits its own training trees and any reserved mass is then pure loss. + */ + private static List read(int count, double from, double to) throws Exception { + load(); + int lo = (int) (from * newickCache.size()); + int hi = (int) (to * newickCache.size()); + List pool = newickCache.subList(lo, hi); + List out = new ArrayList<>(); + double step = Math.max(1.0, pool.size() / (double) count); + for (int i = 0; i < count && (int) (i * step) < pool.size(); i++) { + out.add(new beast.base.evolution.tree.TreeParser(taxaCache, pool.get((int) (i * step)), 0, false)); + } + return out; + } + + /** Trees for the final models and for scoring: first half of the chain trains, second half tests. */ + private static List train(int count) throws Exception { + return read(count, 0.0, 0.5); + } + + /** Disjoint inner split of the training half: [0, 0.25) fits, [0.25, 0.5) validates. */ + private static List fitSet(int count) throws Exception { + return read(count, 0.0, 0.25); + } + + private static List valSet(int count) throws Exception { + return read(count, 0.25, 0.5); + } + + private static double[] score(Scorer s, List test) { + int covered = 0; + double sum = 0.0; + for (Tree t : test) { + double lp = s.logP(t); + if (Double.isFinite(lp)) { + covered++; + sum += lp; + } + } + return new double[]{covered, covered == 0 ? Double.NEGATIVE_INFINITY : sum / covered}; + } + + /** Mean log probability over every test tree, treating unsupported trees as -infinity. */ + private static double meanAll(Scorer s, List test) { + double sum = 0.0; + for (Tree t : test) { + double lp = s.logP(t); + if (!Double.isFinite(lp)) { + return Double.NEGATIVE_INFINITY; + } + sum += lp; + } + return sum / test.size(); + } + + @Test + public void headToHeadOnRealData() throws Exception { + Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(), + "set -Dccd.trees=/path/to/x.trees to run this comparison"); + + List probe = train(N); + int nTaxa = probe.get(0).getLeafNodeCount(); + System.out.printf("%n=== %s: %d taxa, %d train / %d test trees ===%n", + new File(PATH).getName(), nTaxa, probe.size(), N); + + List test = read(N, 0.5, 1.0); + List val = valSet(N / 2); + int nFit = N / 2; + + // ---- CCD1 ---- + CCD1 ccd1 = new CCD1(train(N), 0.0); + report("CCD1", "-", ccd1::getLogProbabilityOfTree, test); + + // ---- RegCCD: alpha on validation ---- + double bestAlpha = 0.4; + double bestAlphaScore = Double.NEGATIVE_INFINITY; + for (double alpha : new double[]{0.01, 0.05, 0.1, 0.2, 0.4, 0.8, 1.0}) { + RegCCD m = new RegCCD(fitSet(nFit), 0.0, alpha); + double sc = score(m::getLogProbabilityOfTree, val)[1]; + if (sc > bestAlphaScore) { + bestAlphaScore = sc; + bestAlpha = alpha; + } + } + RegCCD reg = new RegCCD(train(N), 0.0, bestAlpha); + report("RegCCD", String.format("alpha=%.2f", bestAlpha), reg::getLogProbabilityOfTree, test); + + // ---- KRegCCD: mu on validation at alpha = 0.4 (the manuscript's setting) ---- + KRegCCD kFit = new KRegCCD(fitSet(nFit), 0.0, 0.005, 0.4); + double bestMu = 0.005; + double bestMuScore = Double.NEGATIVE_INFINITY; + for (double mu : new double[]{0.00002, 0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05}) { + final double m = mu; + double sc = score(t -> kFit.getLogProbabilityOfTree(t, m), val)[1]; + if (sc > bestMuScore) { + bestMuScore = sc; + bestMu = mu; + } + } + KRegCCD kreg = new KRegCCD(train(N), 0.0, bestMu, 0.4); + report("KRegCCD", String.format("alpha=0.4, mu=%.5f", bestMu), + kreg::getLogProbabilityOfTree, test); + + // ---- MRegCCD: mu on validation ---- + MRegCCD mFit = new MRegCCD(fitSet(nFit), 0.0, MRegCCD.DEFAULT_MU); + double bestMMu = MRegCCD.DEFAULT_MU; + double bestMMuScore = Double.NEGATIVE_INFINITY; + for (double mu : new double[]{0.00005, 0.0002, 0.001, 0.002, 0.008, 0.0159, 0.05, 0.1}) { + final double m = mu; + double sc = score(t -> mFit.getLogProbabilityOfTree(t, m), val)[1]; + if (sc > bestMMuScore) { + bestMMuScore = sc; + bestMMu = mu; + } + } + MRegCCD mreg = new MRegCCD(train(N), 0.0, bestMMu); + report("MRegCCD", String.format("mu=%.5f", bestMMu), mreg::getLogProbabilityOfTree, test); + + // ---- CRegCCD: (a2, a3, a4) on validation, a1 = 0 ---- + CRegCCD cFit = new CRegCCD(fitSet(nFit), 0.0); + double[] grid = {0.002, 0.01, 0.05, 0.2, 0.4, 1.0, 2.0, 5.0, 12.0}; + double[] bestC = {0.0, 0.4, 0.4, 0.4}; + double bestCScore = Double.NEGATIVE_INFINITY; + for (double b2 : grid) { + for (double b3 : grid) { + for (double b4 : grid) { + double sc = score(t -> cFit.getLogProbabilityOfTree(t, b2, b3, b4), val)[1]; + if (sc > bestCScore) { + bestCScore = sc; + bestC = new double[]{0.0, b2, b3, b4}; + } + } + } + } + CRegCCD creg = new CRegCCD(train(N), 0.0, bestC[1], bestC[2], bestC[3]); + report("CRegCCD", String.format("alpha=%.3f, alpha1=%.3f, alpha2=%.3f", bestC[1], bestC[2], bestC[3]), + creg::getLogProbabilityOfTree, test); + + // ---- paired comparison of the full-support models against KRegCCD ---- + System.out.printf("%npaired per-tree comparison against KRegCCD (n = %d test trees):%n", test.size()); + paired("MRegCCD", mreg::getLogProbabilityOfTree, kreg::getLogProbabilityOfTree, test); + paired("CRegCCD", creg::getLogProbabilityOfTree, kreg::getLogProbabilityOfTree, test); + } + + private static void report(String name, String params, Scorer s, List test) { + double[] sc = score(s, test); + double all = meanAll(s, test); + System.out.printf("%-9s %-27s coverage %6.1f%% mean logP(covered) %10.2f mean logP(all) %s%n", + name, params, 100.0 * sc[0] / test.size(), sc[1], + Double.isFinite(all) ? String.format("%10.2f", all) : " -inf"); + } + + private static void paired(String name, Scorer a, Scorer baseline, List test) { + int wins = 0; + double sumDiff = 0.0; + double sumSq = 0.0; + for (Tree t : test) { + double d = a.logP(t) - baseline.logP(t); + if (d > 0) { + wins++; + } + sumDiff += d; + sumSq += d * d; + } + int n = test.size(); + double mean = sumDiff / n; + double se = Math.sqrt(Math.max(0, sumSq / n - mean * mean) / n); + System.out.printf(" %-9s mean log-ratio %+8.2f +/- %.2f nats/tree, better on %d/%d trees%n", + name, mean, se, wins, n); + } + + /** Scale check: sampling from a real 129-taxon posterior must be fast, valid and + * self-consistent with the scorer. */ + @Test + public void samplingOnRealData() throws Exception { + Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(), + "set -Dccd.trees=/path/to/x.trees to run this comparison"); + CRegCCD ccd = new CRegCCD(train(N), 0.0); + int nTaxa = ccd.getSizeOfLeavesArray(); + long t0 = System.nanoTime(); + int draws = 200; + for (int i = 0; i < draws; i++) { + Tree t = ccd.sampleTree(); + if (t.getLeafNodeCount() != nTaxa || t.getNodeCount() != 2 * nTaxa - 1) { + throw new AssertionError("invalid sampled tree"); + } + double stamped = (Double) t.getRoot().getMetaData(CCD1.LOG_PROB_SUBTREE_KEY); + double scored = ccd.getLogProbabilityOfTree(t); + if (Math.abs(stamped - scored) > 1e-9) { + throw new AssertionError("stamped " + stamped + " != scored " + scored); + } + } + double secs = (System.nanoTime() - t0) / 1e9; + System.out.printf("CRegCCD sampling on %s: %d taxa, %d trees in %.2f s (%.1f ms/tree), " + + "all valid and self-consistent%n", + new File(PATH).getName(), nTaxa, draws, secs, 1000 * secs / draws); + } + + /** MAP and entropy on a real posterior: correctness certificate and wall-clock cost. */ + @Test + public void mapAndEntropyOnRealData() throws Exception { + Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(), + "set -Dccd.trees=/path/to/x.trees to run this comparison"); + CRegCCD creg = new CRegCCD(train(N), 0.0, 2.0, 0.4, 0.05); + + long t0 = System.nanoTime(); + double maxLog = creg.getMaxLogTreeProbability(); + boolean certified = creg.isMAPCertifiedGlobal(); + double bound = creg.getOffBackboneBound(); + double mapSecs = (System.nanoTime() - t0) / 1e9; + + t0 = System.nanoTime(); + Tree map = creg.getMAPTree(); + double treeSecs = (System.nanoTime() - t0) / 1e9; + double scored = creg.getLogProbabilityOfTree(map); + + t0 = System.nanoTime(); + double[] h = creg.getEntropyMonteCarlo(20_000); + double entSecs = (System.nanoTime() - t0) / 1e9; + + System.out.printf("%n=== %s: CRegCCD MAP and entropy ===%n", new File(PATH).getName()); + System.out.printf("MAP DP + certificate : %.2f s, max logP = %.4f, " + + "off-backbone bound = %.4f, certified global = %s%n", + mapSecs, maxLog, bound, certified); + System.out.printf("MAP tree traceback : %.2f s, scored logP = %.4f (matches: %s)%n", + treeSecs, scored, Math.abs(scored - maxLog) < 1e-9); + System.out.printf("entropy (20k draws) : %.2f s, H = %.3f +/- %.3f nats%n", + entSecs, h[0], h[1]); + } + + /** Deterministic entropy recursion vs the unbiased Monte-Carlo estimator on a real posterior. */ + @Test + public void entropyRecursionVersusMonteCarloOnRealData() throws Exception { + Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(), + "set -Dccd.trees=/path/to/x.trees to run this comparison"); + System.out.printf("%n=== %s: CRegCCD entropy, recursion vs Monte Carlo ===%n", + new File(PATH).getName()); + System.out.printf("%-22s %-11s %-9s %-22s %-9s %-10s%n", + "pseudocounts", "recursion", "rec (s)", "Monte Carlo", "MC (s)", "difference"); + for (double[] p : new double[][]{{2.0, 0.4, 0.05}, {5.0, 2.0, 0.4}, {0.4, 0.4, 0.4}}) { + CRegCCD ccd = new CRegCCD(train(N), 0.0, p[0], p[1], p[2]); + long t0 = System.nanoTime(); + double rec = ccd.getEntropyRecursive(); + double recSecs = (System.nanoTime() - t0) / 1e9; + t0 = System.nanoTime(); + double[] mc = ccd.getEntropyMonteCarlo(200_000); + double mcSecs = (System.nanoTime() - t0) / 1e9; + double diff = rec - mc[0]; + System.out.printf("a=(%.2f,%.2f,%.2f)%-6s %-11.4f %-9.2f %8.4f +/-%.4f %-9.2f %+.4f (%+.3f%%)%n", + p[0], p[1], p[2], "", rec, recSecs, mc[0], mc[1], mcSecs, diff, 100 * diff / mc[0]); + } + } + + /** Does the exact A_0/A_1 MAP search stay tractable on a real posterior? */ + @Test + public void exactMapOnRealData() throws Exception { + Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(), + "set -Dccd.trees=/path/to/x.trees to run this comparison"); + CRegCCD creg = new CRegCCD(train(N), 0.0, 0.4, 0.4, 0.05); + double backbone = creg.getMaxLogTreeProbability(); + System.out.printf("%n=== %s: MAP search by allowed A_1 depth ===%n", new File(PATH).getName()); + System.out.printf("backbone (A_0 only) max logP = %.4f%n", backbone); + System.out.printf("%-6s %-12s %-12s %-12s %-10s %-8s%n", + "maxA1", "max logP", "improvement", "bound", "certified", "states"); + for (int k = 0; k <= 3; k++) { + long t0 = System.nanoTime(); + CRegCCD.MapResult r = creg.solveMAP(k); + double secs = (System.nanoTime() - t0) / 1e9; + if (!r.complete()) { + System.out.printf("%-6d exceeded the state budget after %d states (%.1f s)%n", + k, r.statesExplored(), secs); + break; + } + System.out.printf("%-6d %-12.4f %-12.4f %-12.4f %-10s %-8d (%.1f s)%n", + k, r.maxLogProbability(), r.maxLogProbability() - backbone, + r.offBackboneBound(), r.a2Excluded(), r.statesExplored(), secs); + } + } +} diff --git a/src/test/java/ccd/model/SplitClassSizeAnalysis.java b/src/test/java/ccd/model/SplitClassSizeAnalysis.java new file mode 100644 index 0000000..6ba4e2b --- /dev/null +++ b/src/test/java/ccd/model/SplitClassSizeAnalysis.java @@ -0,0 +1,110 @@ +package ccd.model; + +import beast.base.evolution.tree.Tree; +import beast.base.evolution.tree.TreeParser; +import ccd.model.bitsets.BitSet; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.Random; + +/** + * Exploratory: sizes of the four split classes at a clade, for the class-based smoothing proposal. + * + *

At a clade {@code C} of {@code m} taxa every one of the {@code 2^(m-1) - 1} bipartitions falls + * into exactly one of: (1) observed split; (2) unobserved, both children observed clades + * (the CCD0 expansion); (3) unobserved, exactly one child observed; (4) unobserved, neither child + * observed. Classes 1-3 are at most polynomial in the number of observed clades; class 4 is + * essentially all of {@code 2^(m-1)}. This prints the four sizes and the probability the smoothed + * model would put on class 1, under a per-split constant pseudocount versus a per-class total. + */ +public class SplitClassSizeAnalysis { + + private static List randomTrees(int nTaxa, int nTrees, long seed) { + List taxa = new ArrayList<>(); + for (int i = 0; i < nTaxa; i++) { + taxa.add("T" + i); + } + Random rng = new Random(seed); + List out = new ArrayList<>(); + for (int t = 0; t < nTrees; t++) { + List pool = new ArrayList<>(taxa); + while (pool.size() > 1) { + int i = rng.nextInt(pool.size()); + String a = pool.remove(i); + int j = rng.nextInt(pool.size()); + String b = pool.remove(j); + pool.add("(" + a + "," + b + ")"); + } + out.add(new TreeParser(taxa, pool.get(0) + ";", 1, false)); + } + return out; + } + + @Test + public void classSizesAtRoot() { + System.out.printf("%-6s %-7s %-8s %-8s %-8s %-14s %-14s %-14s%n", + "taxa", "trees", "|A1|", "|A2|", "|A3|", "|A4|", "P(A1) per-split", "P(A1) per-class"); + for (int nTaxa : new int[]{12, 20, 30, 40}) { + int nTrees = 1000; + List trees = randomTrees(nTaxa, nTrees, 42L); + CCD0 ccd0 = new CCD0(trees, 0); + Clade root = null; + for (Clade c : ccd0.getClades()) { + if (c.size() == nTaxa) { + root = c; + } + } + if (root == null) { + continue; + } + + int a1 = 0; + int a2 = 0; + for (CladePartition p : root.getPartitions()) { + if (p.getNumberOfOccurrences() > 0) { + a1++; + } else { + a2++; + } + } + + // |A3|: observed proper subclades whose complement within root is NOT an observed clade + BitSet rootBits = root.getCladeInBits(); + int a3 = 0; + for (Clade d : ccd0.getClades()) { + if (d.size() >= root.size()) { + continue; + } + BitSet db = d.getCladeInBits(); + BitSet tmp = (BitSet) db.clone(); + tmp.and(rootBits); + if (!tmp.equals(db)) { + continue; // not a subclade of root + } + BitSet comp = (BitSet) rootBits.clone(); + comp.andNot(db); + if (ccd0.getClade(comp) == null) { + a3++; + } + } + + double total = Math.pow(2, nTaxa - 1) - 1; + double a4 = total - a1 - a2 - a3; + + // per-split constant pseudocount alpha on every class + double alpha = 0.4; + double fC = nTrees; + double denomSplit = fC + alpha * (a1 + a2 + a3 + a4); + double pA1Split = (fC + alpha * a1) / denomSplit; + + // per-class total pseudocount alpha (spread within each class) + double denomClass = fC + alpha * 4; + double pA1Class = (fC + alpha) / denomClass; + + System.out.printf("%-6d %-7d %-8d %-8d %-8d %-14.4g %-14.6g %-14.6g%n", + nTaxa, nTrees, a1, a2, a3, a4, pA1Split, pA1Class); + } + } +}