From ae1e21781ea8312ebb443e841cba976f49d8edc6 Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Tue, 18 Aug 2026 12:58:59 +1200 Subject: [PATCH 1/8] 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/8] 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); + } + } +} From 210cefe46fe4e37991b9e0e1fc743992ea245c5c Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Thu, 20 Aug 2026 09:52:54 +1200 Subject: [PATCH 3/8] Speed up MRegCCD boundary counting and fix BitSet copy MRegCCD enumerated novel resolutions recursively, which blew up at boundary size 4 and silently truncated against its op budget: on RSV2 at depth 4 the score moved from -60.6219 to -60.6443 when the budget was raised 20x, so the default was returning a truncated answer without saying so. The new MRegCCD overrides only countsFor, indexing disjoint pairs of observed subclades by the bitset they cover and counting each boundary order against that index rather than by enumeration. Boundary 4 needs a canonical form -- a 4-part boundary decomposes into two pairs in three ways -- or every such boundary is counted three times; the agreement test caught exactly that. The previous implementation is kept as MRegCCDSlow so the two can be checked against each other, and MRegCCDAgreementTest holds them to identical counts on 2199 clades. Both now default to reserve depth 4, matching KRegCCD's default reserve so the two models look equally far past the CCD graph. Measured 4.3x faster at depth 4 and no longer truncating. BitSet.newBitSet used other.length() rather than other.size() to size the copy. length() is the index of the highest set bit plus one, so copying a set whose top words are empty returned an undersized BitSet, and the bitwise operations index the operand by this.words.length. Co-Authored-By: Claude Opus 5 (1M context) --- .gitignore | 2 +- src/main/java/ccd/model/MRegCCD.java | 688 +++--------------- src/main/java/ccd/model/MRegCCDSlow.java | 656 +++++++++++++++++ src/main/java/ccd/model/bitsets/BitSet.java | 5 +- .../java/ccd/model/MRegCCDAgreementTest.java | 133 ++++ src/test/java/ccd/model/MRegCCDTest.java | 28 +- .../ccd/model/bitsets/BitSetCopyTest.java | 54 ++ 7 files changed, 961 insertions(+), 605 deletions(-) create mode 100644 src/main/java/ccd/model/MRegCCDSlow.java create mode 100644 src/test/java/ccd/model/MRegCCDAgreementTest.java create mode 100644 src/test/java/ccd/model/bitsets/BitSetCopyTest.java diff --git a/.gitignore b/.gitignore index 255db57..fac52ce 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,4 @@ out/ .classpath .DS_Store /bin/ -.project \ No newline at end of file +.projecths_err_pid*.log diff --git a/src/main/java/ccd/model/MRegCCD.java b/src/main/java/ccd/model/MRegCCD.java index 2f11b7a..f553922 100644 --- a/src/main/java/ccd/model/MRegCCD.java +++ b/src/main/java/ccd/model/MRegCCD.java @@ -1,6 +1,5 @@ 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; @@ -9,462 +8,144 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; /** - * MRegCCD -- the one-parameter "per-new-split" regularised CCD. It unifies RegCCD's split-expansion - * {@code alpha} and KRegCCD's escape {@code mu} into a single per-clade escape rate, giving a - * full-support tree distribution with one hyperparameter {@code mu} (and no {@code alpha}). + * The one-parameter "per-new-split" regularised CCD, with the boundary counts computed the way + * {@link KRegCCD} computes its reserve rather than by direct recursive enumeration. * - *

The model is a plain {@link CCD1} backbone (raw conditional clade probabilities, no smoothing) - * extended with a per-clade escape reserve. The distribution is defined conditionally, clade by clade - * (chain rule over the observed-clade DAG; no global partition function): + *

The model is unchanged: this overrides only {@code countsFor}, so every probability, sample and + * point estimate is defined exactly as in {@link MRegCCDSlow}, which this overrides in a single + * method. The two must therefore agree wherever the reference implementation's op budget does not + * truncate its enumeration, which is what {@code MRegCCDAgreementTest} checks. + * + *

The reference implementation walks every ordered choice of boundary parts, which costs {@code O(m^(k-1))} in + * the number {@code m} of observed subclades of a clade and blows through a flat op budget on large + * analyses -- silently, since {@code countsFor} catches the overflow and leaves the remaining orders + * at zero, which also disables the tail correction that reads them. Here the boundaries are found by + * indexing instead: *

    - *
  • An observed split {@code C -> {L, R}} (seen in training) is priced - * {@code (1 - mu - tail(C)) * ccp(L|R)} when {@code C} can escape, else just {@code ccp(L|R)}.
    • - *
  • An escape at {@code C} resolves it through novel intermediate clades down to a - * boundary of {@code m} observed subclades (a maximal "blue" region). Such a region has - * {@code m - 1} new splits and is priced {@code eps(C)^(m-1)} -- one factor of {@code eps} per - * new split (so a recombination of two observed subclades, {@code m = 2}, costs one - * {@code eps}). This is the difference from KRegCCD, which makes recombinations representable in - * its {@code alpha}-expanded backbone and charges {@code eps} only per novel clade - * ({@code eps^(m-2)}).
    • + *
  • the disjoint pairs of observed subclades are enumerated once, {@code O(m^2)}, and grouped by + * the bitset they cover;
    • + *
  • a boundary of 2 is an observed subclade whose complement in {@code C} is observed;
    • + *
  • a boundary of 3 is an observed subclade whose complement is covered by a pair;
    • + *
  • a boundary of 4 is a pair whose complement is covered by another pair.
    • *
- * - *

The per-clade escape rate {@code eps(C)} is the root of {@code sum_{m>=2} M_m(C) eps^(m-1) = mu}, - * where {@code M_m(C)} counts the all-novel resolutions of {@code C} with an {@code m}-part boundary - * (the FLAT weighting: each distinct novel resolution counted once). Computing the full sum is - * #P-hard, so -- mirroring KRegCCD -- the orders {@code m = 2..reserveDepth} are enumerated exactly - * (bounded by an op-budget) and the omitted higher orders are a geometric tail correction added to - * {@code mu} (so observed splits are discounted by {@code 1 - mu - tail}, keeping the conditional - * properly normalised; truncating without the tail super-normalises). A clade with no escape route - * ({@code M_m = 0} for all computed {@code m}) is not reservable and keeps its raw CCP undiscounted. - * - *

Every tree on the taxon set has positive probability (full support), so {@link #containsTree} - * is always true and {@link #getLogProbabilityOfTree} is finite for all trees. - * - * @author Claude (CCD-Sophie) + * Each lookup is a hash probe rather than a search, so orders 2 to 4 cost {@code O(m^2)} in total. + * Orders beyond 4 fall back to the inherited enumeration, so this is a strict speed-up of the + * practical depths and never changes what is computed. */ -public class MRegCCD extends CCD1 { - - /** Default per-clade escape probability (the RSV2 operating point of the conditional model). */ - public static final double DEFAULT_MU = 0.0159; - - /** Default reserve depth: enumerate boundary sizes {@code m = 2..DEFAULT_RESERVE_DEPTH} exactly. */ - public static final int DEFAULT_RESERVE_DEPTH = 5; - - /** Per-clade enumeration-op budget (mirrors KRegCCD's; bounds the boundary enumeration). */ - private static final long OPS_BUDGET = Long.parseLong(System.getProperty("mreg.enumOps", "20000000")); - - private static final class BudgetExceeded extends RuntimeException { - BudgetExceeded() { - super(null, null, false, false); - } - } - - private static final BudgetExceeded BUDGET_EXCEEDED = new BudgetExceeded(); +public class MRegCCD extends MRegCCDSlow { - /** Per-clade escape probability (the single hyperparameter). */ - private final double mu; - - /** Max boundary size enumerated when solving eps; deeper orders are a geometric tail. */ - private final int reserveDepth; - - /** Whether the {@code (1 - mu - tail)} discount carries the geometric tail correction. */ - private final boolean useTail; - - /** Observed-clade bitsets (incl. leaves), sorted canonically; built lazily. */ - private List sortedCladeBits; - private final Map> subCache = new HashMap<>(); - private final Map countsCache = new HashMap<>(); - private long enumOps; + private final Map fastCounts = new ConcurrentHashMap<>(); public MRegCCD(List trees, double burnin, double mu) { - this(trees, burnin, mu, DEFAULT_RESERVE_DEPTH, true); + super(trees, burnin, mu); } public MRegCCD(List trees, double burnin, double mu, int reserveDepth, boolean useTail) { - super(trees, burnin); - validate(mu, reserveDepth); - this.mu = mu; - this.reserveDepth = reserveDepth; - this.useTail = useTail; + super(trees, burnin, mu, reserveDepth, useTail); } public MRegCCD(TreeSet treeSet, double mu) { - this(treeSet, mu, DEFAULT_RESERVE_DEPTH, true); + super(treeSet, mu); } public MRegCCD(TreeSet treeSet, double mu, int reserveDepth, boolean useTail) { - super(treeSet); - validate(mu, reserveDepth); - this.mu = mu; - this.reserveDepth = reserveDepth; - this.useTail = useTail; - } - - /** - * Builds an MRegCCD on {@code trees} with {@code mu} selected by maximising cross-validated - * held-out log-probability (see {@link ccd.algorithms.regularisation.MRegCCDParameterOptimiser}), - * rather than the fixed {@link #DEFAULT_MU}. The honest, no-peeking counterpart of - * {@code KRegCCD.withOptimisedParameters}. - */ - public static MRegCCD withOptimisedMu(List trees) { - double mu = ccd.algorithms.regularisation.MRegCCDParameterOptimiser.optimiseMu(trees).mu(); - return new MRegCCD(trees, 0.0, mu); - } - - private static void validate(double mu, int reserveDepth) { - if (mu <= 0 || mu >= 1) { - throw new IllegalArgumentException("mu must be in (0, 1), got " + mu); - } - if (reserveDepth < 2) { - throw new IllegalArgumentException("reserveDepth must be >= 2, got " + reserveDepth); - } - } - - /** The per-clade escape probability this model was built with. */ - public double getMu() { - return mu; - } - - public int getReserveDepth() { - return reserveDepth; - } - - /** - * Reserve counts {@code M_m(C)} by boundary size {@code m} (array index {@code m}, valid for - * {@code m = 2..min(|C|, reserveDepth)}); {@code M_m} is the number of all-novel resolutions of - * {@code C} with an {@code m}-part boundary. The first coefficient {@code M_2} (the {@code eps^1} - * term) is exactly the number of CCD0-expanded splits of {@code C} -- recombinations of two - * observed subclades whose split was never observed -- since those are the only escapes with no - * other novel (blue) clade. Exposed for inspection and cross-checks. - */ - public int[] reserveCounts(BitSet cladeInBits) { - return countsFor(cladeInBits).clone(); - } - - @Override - public String toString() { - return "MRegCCD [mu = " + mu + ", reserveDepth = " + reserveDepth + ", tail = " + useTail - + ", per-new-split, full support]"; - } - - /* ---------------------------------------------------------------------- - * Scoring - * ------------------------------------------------------------------- */ - - @Override - public double getLogProbabilityOfTree(Tree tree) { - return scoreTree(tree, mu); - } - - /** - * Full-support log-probability at an arbitrary escape probability {@code scoreMu}, reusing this - * model's ({@code mu}-independent) backbone and cached reserve counts. Lets a parameter search / - * cross-validation evaluate many {@code mu} on one trained model without rebuilding. For - * {@code scoreMu == mu} it equals {@link #getLogProbabilityOfTree(Tree)}. - */ - public double getLogProbabilityOfTree(Tree tree, double scoreMu) { - if (scoreMu <= 0 || scoreMu >= 1) { - throw new IllegalArgumentException("scoreMu must be in (0, 1), got " + scoreMu); - } - return scoreTree(tree, scoreMu); + super(treeSet, mu, reserveDepth, useTail); } @Override - public double getProbabilityOfTree(Tree tree) { - return Math.exp(getLogProbabilityOfTree(tree)); - } - - /** Always true: MRegCCD is full support, so every tree on this taxon set has positive probability. */ - @Override - public boolean containsTree(Tree tree) { - return true; - } - - private double scoreTree(Tree tree, double scoreMu) { - Map bits = new HashMap<>(); - computeBits(tree.getRoot(), bits); - double logp = 0.0; - for (Node v : tree.getNodesAsArray()) { - if (v.isLeaf()) { - continue; - } - BitSet vb = bits.get(v); - Clade c = getClade(vb); - if (c == null) { - continue; // novel clade: scored once at its maximal region's top - } - BitSet b1 = bits.get(v.getChildren().get(0)); - BitSet b2 = bits.get(v.getChildren().get(1)); - if (isSplitObserved(vb, b1, b2)) { - if (reservable(vb)) { // discount only clades that can actually escape - double resv = Math.min(scoreMu + (useTail ? tailFor(vb, scoreMu) : 0.0), 1 - 1e-12); - logp += Math.log(1.0 - resv); - } - logp += rawLogCCP(c, b1, b2); // raw CCD1 CCP - } else { - // region top: an observed clade resolved through a novel split. m-1 new splits. - int m = boundarySize(v, bits); - logp += (m - 1) * Math.log(epsFor(vb, scoreMu)); - } - } - return logp; - } - - /* ---------------------------------------------------------------------- - * Per-clade reserve (M_m counts -> eps, tail; mirrors KRegCCD.computeReg) - * ------------------------------------------------------------------- */ - - /** Whether clade {@code C} (given in bits) reserves any escape mass up to {@code reserveDepth}. */ - boolean reservable(BitSet C) { - for (int v : countsFor(C)) { - if (v > 0) { - return true; - } - } - return false; - } - - /** Escape root {@code eps} solving {@code sum_{m>=2} M_m eps^(m-1) = scoreMu} (monotone bisection). */ - double epsFor(BitSet C, double scoreMu) { - int[] n = countsFor(C); - if (!reservable(C)) { - return scoreMu; // crude fallback (no escape route within reserveDepth); should not be hit - } - return solveEps(n, scoreMu); - } - - /** Omitted-tail escape mass beyond the computed orders: geometric bound from the top two orders. */ - double tailFor(BitSet C, double scoreMu) { - int[] n = countsFor(C); - int last = n.length - 1; - if (last < 3) { - return 0.0; - } - int nLast = n[last], nPrev = n[last - 1]; - if (nLast <= 0 || nPrev <= 0) { - return 0.0; - } - double eps = epsFor(C, scoreMu); - double rho = ((double) nLast / nPrev) * eps; - if (rho <= 0 || rho >= 1) { - return 0.0; - } - return Math.min(nLast * Math.pow(eps, last - 1) * rho / (1 - rho), scoreMu); - } - - /** M_m counts (index m = boundary size, 2..min(|C|, reserveDepth)); cached, mu-independent. */ int[] countsFor(BitSet C) { - int[] cached = countsCache.get(C); + int[] cached = fastCounts.get(C); if (cached != null) { return cached; } int card = C.cardinality(); - int[] n = new int[Math.min(card, reserveDepth) + 1]; - if (card >= 2) { + int depth = Math.min(card, getReserveDepth()); + // The pair index only pays for itself once order 4 needs it; below that the inherited + // enumeration is cheaper, and identical by construction. + if (depth < 4) { + return super.countsFor(C); + } + int[] n = new int[depth + 1]; + if (card >= 2 && depth >= 2) { List subs = subclades(C); - enumOps = 0; - for (int m = 2; m < n.length; m++) { - try { - n[m] = countBoundaries(C, subs, m); - } catch (BudgetExceeded e) { - break; // deeper orders omitted (negligible, like the tail) - } - } - } - countsCache.put(C, n); - return n; - } - - private static double solveEps(int[] n, double mu) { - double lo = 0.0, hi = 1.0; - while (evalReserve(n, hi) < mu) { - hi *= 2.0; - } - for (int it = 0; it < 100; it++) { - double mid = 0.5 * (lo + hi); - if (evalReserve(n, mid) < mu) { - lo = mid; - } else { - hi = mid; - } - } - return 0.5 * (lo + hi); - } - - /** {@code sum_{m>=2} n[m] x^(m-1)}. */ - private static double evalReserve(int[] n, double x) { - double s = 0.0; - for (int m = 2; m < n.length; m++) { - if (n[m] > 0) { - s += n[m] * Math.pow(x, m - 1); - } - } - return s; - } - - /** Count m-part boundaries of C into observed subclades, weighted by their all-novel pathcount. */ - private int countBoundaries(BitSet C, List subs, int m) { - return enumerateBoundaries(C, subs, m, 0, BitSet.newBitSet(leafArraySize), new ArrayList<>(m)); - } - private int enumerateBoundaries(BitSet C, List subs, int m, int startIdx, - BitSet used, List chosen) { - if (++enumOps > OPS_BUDGET) { - throw BUDGET_EXCEEDED; - } - if (chosen.size() == m - 1) { - BitSet last = BitSet.newBitSet(C); - last.andNot(used); - if (last.isEmpty() || !isObs(last)) { - return 0; - } - if (compareBitSets(chosen.get(chosen.size() - 1), last) >= 0) { - return 0; // canonical: the derived last part must be the largest - } - BitSet[] parts = new BitSet[m]; - for (int i = 0; i < m - 1; i++) { - parts[i] = chosen.get(i); - } - parts[m - 1] = last; - return countAllNovelResolutions(C, parts); - } - int count = 0; - for (int i = startIdx; i < subs.size(); i++) { - BitSet pb = subs.get(i); - if (pb.intersects(used)) { - continue; + // Every disjoint pair of observed subclades, grouped by the bitset it covers. Only + // orders 3 and 4 consult this, so at depth 2 the index is not worth building: the O(m^2) + // pass would cost more than the order-2 scan it would serve. + Map> pairsByUnion = new HashMap<>(); + for (int i = 0; i < subs.size(); i++) { + BitSet a = subs.get(i); + for (int j = i + 1; j < subs.size(); j++) { + BitSet b = subs.get(j); + if (a.intersects(b)) { + continue; + } + BitSet union = BitSet.newBitSet(a); + union.or(b); + pairsByUnion.computeIfAbsent(union, k -> new ArrayList<>()) + .add(new BitSet[]{a, b}); + } } - chosen.add(pb); - BitSet newUsed = BitSet.newBitSet(used); - newUsed.or(pb); - count += enumerateBoundaries(C, subs, m, i + 1, newUsed, chosen); - chosen.remove(chosen.size() - 1); - } - return count; - } - /** - * Number of all-novel binary resolutions of C into the given observed parts (subset DP over the - * parts). A split is allowed iff: at the region root (full mask = C, an observed clade) the split - * is unobserved (a real escape); at an intermediate node the clade itself is novel (a maximal - * region stops at observed clades, matching {@link #boundarySize}). - */ - private int countAllNovelResolutions(BitSet C, BitSet[] parts) { - int k = parts.length; - if (k == 1) { - return 1; - } - int full = (1 << k) - 1; - BitSet[] unionOf = new BitSet[1 << k]; - unionOf[0] = BitSet.newBitSet(leafArraySize); - for (int mask = 1; mask <= full; mask++) { - int low = Integer.numberOfTrailingZeros(mask); - BitSet u = BitSet.newBitSet(unionOf[mask & (mask - 1)]); - u.or(parts[low]); - unionOf[mask] = u; - } - int[] f = new int[1 << k]; - for (int mask = 1; mask <= full; mask++) { - if (Integer.bitCount(mask) == 1) { - f[mask] = 1; - continue; - } - int low = mask & (-mask), rest = mask ^ low, count = 0; - for (int sub = rest; ; sub = (sub - 1) & rest) { - int s1 = sub | low, s2 = mask ^ s1; - if (s2 != 0 && splitAllowed(mask == full, unionOf[mask], unionOf[s1], unionOf[s2])) { - count += f[s1] * f[s2]; + for (BitSet d : subs) { + BitSet rest = BitSet.newBitSet(C); + rest.andNot(d); + if (rest.isEmpty()) { + continue; } - if (sub == 0) { - break; + // boundary 2: {d, rest}, counted once from its canonically smaller side + if (depth >= 2 && isObs(rest) && compareBitSets(d, rest) < 0) { + n[2] += countAllNovelResolutions(C, new BitSet[]{d, rest}); } - } - f[mask] = count; - } - return f[full]; - } - - private boolean splitAllowed(boolean top, BitSet union, BitSet a, BitSet b) { - return top ? !isSplitObserved(union, a, b) : !isObs(union); - } - - /* ---------------------------------------------------------------------- - * Observed-backbone queries (over the inherited CCD1 clade DAG) - * ------------------------------------------------------------------- */ - - private boolean isObs(BitSet x) { - return getClade(x) != null; // leaves are clades too - } - - private boolean isSplitObserved(BitSet parentBits, BitSet aBits, BitSet bBits) { - Clade parent = getClade(parentBits); - if (parent == null) { - return false; - } - Clade a = getClade(aBits); - Clade b = getClade(bBits); - if (a == null || b == null) { - return false; - } - return parent.getCladePartition(a, b) != null; - } - - private double rawLogCCP(Clade parent, BitSet aBits, BitSet bBits) { - CladePartition p = parent.getCladePartition(getClade(aBits), getClade(bBits)); - return p.getLogCCP(); - } - - /** Observed clades (incl. leaves) strictly contained in C, in canonical order; cached. */ - private List subclades(BitSet C) { - return subCache.computeIfAbsent(C, c -> { - List out = new ArrayList<>(); - int card = c.cardinality(); - for (BitSet x : sortedCladeBits()) { - if (x.cardinality() < card && subset(x, c)) { - out.add(x); + // boundary 3: {d} plus a pair covering the remainder + if (depth >= 3) { + for (BitSet[] p : pairsByUnion.getOrDefault(rest, List.of())) { + if (compareBitSets(d, p[0]) < 0) { // d must be the canonically first part + n[3] += countAllNovelResolutions(C, new BitSet[]{d, p[0], p[1]}); + } + } } } - return out; - }); - } - private List sortedCladeBits() { - if (sortedCladeBits == null) { - List all = new ArrayList<>(); - for (Clade c : getClades()) { - all.add(c.getCladeInBits()); + // boundary 4: a pair whose complement is covered by another pair + if (depth >= 4) { + for (Map.Entry> e : pairsByUnion.entrySet()) { + BitSet rest = BitSet.newBitSet(C); + rest.andNot(e.getKey()); + if (rest.isEmpty() || !subset(e.getKey(), C)) { + continue; + } + List others = pairsByUnion.get(rest); + if (others == null) { + continue; + } + for (BitSet[] p : e.getValue()) { + for (BitSet[] q : others) { + // A 4-part boundary splits into two pairs in three ways, so count only + // the one whose first pair holds the two canonically smallest parts: + // for parts w < x < y < z that is {w,x}|{y,z} and no other. + if (compareBitSets(p[1], q[0]) < 0) { + n[4] += countAllNovelResolutions(C, + new BitSet[]{p[0], p[1], q[0], q[1]}); + } + } + } + } } - all.sort(MRegCCD::compareBitSets); - sortedCladeBits = all; - } - return sortedCladeBits; - } - /** Boundary size of the maximal region rooted at v: count of maximal observed/leaf subclades below. */ - private int boundarySize(Node v, Map bits) { - int m = 0; - for (Node child : v.getChildren()) { - if (child.isLeaf() || getClade(bits.get(child)) != null) { - m++; - } else { - m += boundarySize(child, bits); + // orders beyond 4 are rare in practice; defer to the inherited enumeration + if (depth >= 5) { + int[] slow = super.countsFor(C); + for (int m = 5; m < n.length && m < slow.length; m++) { + n[m] = slow[m]; + } } } - return m; - } - - 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; + fastCounts.put(BitSet.newBitSet(C), n); + return n; } private static boolean subset(BitSet a, BitSet c) { @@ -472,175 +153,4 @@ private static boolean subset(BitSet a, BitSet c) { tmp.andNot(c); return tmp.isEmpty(); } - - /** Canonical total order on clade bitsets (lexicographic by set-bit indices). */ - private static int compareBitSets(BitSet a, BitSet b) { - int ia = a.nextSetBit(0), 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); - } - - /* ---------------------------------------------------------------------- - * Sampling (self-consistent) - * - * The PIT calibration test draws trees from the model and needs only each draw's log-probability - * (not the tree object), so we override sampleTreeLogProbability() with a direct simulation of the - * generative process and never materialise a Tree. At each reservable clade we escape with - * probability equal to its escape mass (= mu by the eps-solve, tail EXCLUDED) and otherwise take - * an observed (red) split ~ CCP; an escape draws a region order m proportional to M_m eps^(m-1), a - * boundary of m observed subclades proportional to its all-novel pathcount, and recurses into the - * boundary parts. The resolution shape within a region is not drawn -- every shape has the same - * weight eps^(m-1) and does not change the draw's log-probability -- so the simulation is cheap. - * - * The draw distribution exactly matches getLogProbabilityOfTree when the model is built with the - * tail OFF (then the red discount is 1 - mu, matching the escape mass), for trees whose regions are - * within reserveDepth; deeper regions (mass ~mu^reserveDepth) are never produced, the same - * self-consistent / full-support trade-off KRegCCD makes for its PIT. - * ------------------------------------------------------------------- */ - - @Override - public double sampleTreeLogProbability() { - return simulate(getRootClade()); - } - - private double simulate(Clade c) { - if (c.isLeaf()) { - return 0.0; - } - BitSet cb = c.getCladeInBits(); - if (reservable(cb)) { - double eps = epsFor(cb, mu); - int[] n = countsFor(cb); - double escapeMass = 0.0; - for (int m = 2; m < n.length; m++) { - if (n[m] > 0) { - escapeMass += n[m] * Math.pow(eps, m - 1); - } - } - if (random.nextDouble() < escapeMass) { - int m = sampleOrder(n, eps, escapeMass); - BitSet[] parts = sampleBoundaryParts(cb, subclades(cb), m); - double logp = (m - 1) * Math.log(eps); - if (parts != null) { - for (BitSet bp : parts) { - logp += simulate(getClade(bp)); - } - } - return logp; - } - CladePartition p = samplePartition(c); - double logp = Math.log(1.0 - escapeMass) + p.getLogCCP(); - return logp + simulate(p.getChildClades()[0]) + simulate(p.getChildClades()[1]); - } - CladePartition p = samplePartition(c); // non-reservable: observed split, no discount - return p.getLogCCP() + simulate(p.getChildClades()[0]) + simulate(p.getChildClades()[1]); - } - - /** Draws a region order m in {2..} with probability proportional to {@code M_m eps^(m-1)}. */ - private int sampleOrder(int[] n, double eps, double escapeMass) { - double target = random.nextDouble() * escapeMass, acc = 0.0; - for (int m = 2; m < n.length; m++) { - if (n[m] > 0) { - acc += n[m] * Math.pow(eps, m - 1); - if (target < acc) { - return m; - } - } - } - for (int m = n.length - 1; m >= 2; m--) { - if (n[m] > 0) { - return m; // numerical guard - } - } - return 2; - } - - /** Samples an observed (red) split of {@code c} with probability proportional to its CCP. */ - private CladePartition samplePartition(Clade c) { - List partitions = c.getPartitions(); - double target = random.nextDouble(), acc = 0.0; - for (CladePartition p : partitions) { - acc += p.getCCP(); - if (target < acc) { - return p; - } - } - return partitions.get(partitions.size() - 1); - } - - /** - * Weighted-reservoir samples one boundary of {@code c} into {@code m} observed subclades, - * proportional to its all-novel pathcount (so that, combined with order sampling, every distinct - * novel resolution is equiprobable at {@code eps^(m-1)}). Returns the parts, or {@code null} if - * none/op-budget. - */ - private BitSet[] sampleBoundaryParts(BitSet c, List subs, int m) { - boundaryPick = null; - boundaryWeightSeen = 0.0; - enumOps = 0; - try { - sampleBoundaryWalk(c, subs, m, 0, BitSet.newBitSet(leafArraySize), new ArrayList<>(m)); - } catch (BudgetExceeded e) { - return boundaryPick; // whatever was picked before the cap (may be null) - } - return boundaryPick; - } - - private BitSet[] boundaryPick; - private double boundaryWeightSeen; - - private void sampleBoundaryWalk(BitSet c, List subs, int m, int startIdx, - BitSet used, List chosen) { - if (++enumOps > OPS_BUDGET) { - throw BUDGET_EXCEEDED; - } - if (chosen.size() == m - 1) { - BitSet last = BitSet.newBitSet(c); - last.andNot(used); - if (last.isEmpty() || !isObs(last)) { - return; - } - if (compareBitSets(chosen.get(chosen.size() - 1), last) >= 0) { - return; - } - BitSet[] parts = new BitSet[m]; - for (int i = 0; i < m - 1; i++) { - parts[i] = chosen.get(i); - } - parts[m - 1] = last; - int pc = countAllNovelResolutions(c, parts); - if (pc <= 0) { - return; - } - boundaryWeightSeen += pc; - if (random.nextDouble() * boundaryWeightSeen < pc) { // weighted reservoir - boundaryPick = parts; - } - return; - } - for (int i = startIdx; i < subs.size(); i++) { - BitSet pb = subs.get(i); - if (pb.intersects(used)) { - continue; - } - chosen.add(pb); - BitSet newUsed = BitSet.newBitSet(used); - newUsed.or(pb); - sampleBoundaryWalk(c, subs, m, i + 1, newUsed, chosen); - chosen.remove(chosen.size() - 1); - } - } - - @Override - public Tree sampleTree(HeightSettingStrategy heightStrategy) { - throw new UnsupportedOperationException( - "MRegCCD materialised-tree sampling is not implemented; sampleTreeLogProbability() " - + "(used by the PIT) simulates draws without building trees."); - } } diff --git a/src/main/java/ccd/model/MRegCCDSlow.java b/src/main/java/ccd/model/MRegCCDSlow.java new file mode 100644 index 0000000..ace1207 --- /dev/null +++ b/src/main/java/ccd/model/MRegCCDSlow.java @@ -0,0 +1,656 @@ +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; + +/** + * The reference implementation of {@link MRegCCD}: identical model, boundary counts obtained by + * direct recursive enumeration rather than by indexing. Retained so that the faster implementation + * can be checked against it (see {@code MRegCCDAgreementTest}); prefer {@link MRegCCD} in use, which + * computes the same counts without the enumeration blow-up at boundary 4. + * + *

MRegCCDSlow -- the one-parameter "per-new-split" regularised CCD. It unifies RegCCD's split-expansion + * {@code alpha} and KRegCCD's escape {@code mu} into a single per-clade escape rate, giving a + * full-support tree distribution with one hyperparameter {@code mu} (and no {@code alpha}). + * + *

The model is a plain {@link CCD1} backbone (raw conditional clade probabilities, no smoothing) + * extended with a per-clade escape reserve. The distribution is defined conditionally, clade by clade + * (chain rule over the observed-clade DAG; no global partition function): + *

    + *
  • An observed split {@code C -> {L, R}} (seen in training) is priced + * {@code (1 - mu - tail(C)) * ccp(L|R)} when {@code C} can escape, else just {@code ccp(L|R)}.
    • + *
  • An escape at {@code C} resolves it through novel intermediate clades down to a + * boundary of {@code m} observed subclades (a maximal "blue" region). Such a region has + * {@code m - 1} new splits and is priced {@code eps(C)^(m-1)} -- one factor of {@code eps} per + * new split (so a recombination of two observed subclades, {@code m = 2}, costs one + * {@code eps}). This is the difference from KRegCCD, which makes recombinations representable in + * its {@code alpha}-expanded backbone and charges {@code eps} only per novel clade + * ({@code eps^(m-2)}).
    • + *
+ * + *

The per-clade escape rate {@code eps(C)} is the root of {@code sum_{m>=2} M_m(C) eps^(m-1) = mu}, + * where {@code M_m(C)} counts the all-novel resolutions of {@code C} with an {@code m}-part boundary + * (the FLAT weighting: each distinct novel resolution counted once). Computing the full sum is + * #P-hard, so -- mirroring KRegCCD -- the orders {@code m = 2..reserveDepth} are enumerated exactly + * (bounded by an op-budget) and the omitted higher orders are a geometric tail correction added to + * {@code mu} (so observed splits are discounted by {@code 1 - mu - tail}, keeping the conditional + * properly normalised; truncating without the tail super-normalises). A clade with no escape route + * ({@code M_m = 0} for all computed {@code m}) is not reservable and keeps its raw CCP undiscounted. + * + *

Every tree on the taxon set has positive probability (full support), so {@link #containsTree} + * is always true and {@link #getLogProbabilityOfTree} is finite for all trees. + * + * @author Claude (CCD-Sophie) + */ +public class MRegCCDSlow extends CCD1 { + + /** Default per-clade escape probability (the RSV2 operating point of the conditional model). */ + public static final double DEFAULT_MU = 0.0159; + + /** + * Default reserve depth: enumerate boundary sizes {@code m = 2..DEFAULT_RESERVE_DEPTH} exactly. + * Boundary 4 matches {@link KRegCCD}'s default reserve ({@code k = 2}, boundaries 3 and 4), so + * the two models look equally far past the CCD graph. It was 5, which went a boundary further + * than KRegCCD while the op budget below silently truncated the enumeration before reaching it. + */ + public static final int DEFAULT_RESERVE_DEPTH = 4; + + /** Per-clade enumeration-op budget (mirrors KRegCCD's; bounds the boundary enumeration). */ + private static final long OPS_BUDGET = Long.parseLong(System.getProperty("mreg.enumOps", "20000000")); + + private static final class BudgetExceeded extends RuntimeException { + BudgetExceeded() { + super(null, null, false, false); + } + } + + private static final BudgetExceeded BUDGET_EXCEEDED = new BudgetExceeded(); + + /** Per-clade escape probability (the single hyperparameter). */ + private final double mu; + + /** Max boundary size enumerated when solving eps; deeper orders are a geometric tail. */ + private final int reserveDepth; + + /** Whether the {@code (1 - mu - tail)} discount carries the geometric tail correction. */ + private final boolean useTail; + + /** Observed-clade bitsets (incl. leaves), sorted canonically; built lazily. */ + private List sortedCladeBits; + private final Map> subCache = new HashMap<>(); + private final Map countsCache = new HashMap<>(); + private long enumOps; + + public MRegCCDSlow(List trees, double burnin, double mu) { + this(trees, burnin, mu, DEFAULT_RESERVE_DEPTH, true); + } + + public MRegCCDSlow(List trees, double burnin, double mu, int reserveDepth, boolean useTail) { + super(trees, burnin); + validate(mu, reserveDepth); + this.mu = mu; + this.reserveDepth = reserveDepth; + this.useTail = useTail; + } + + public MRegCCDSlow(TreeSet treeSet, double mu) { + this(treeSet, mu, DEFAULT_RESERVE_DEPTH, true); + } + + public MRegCCDSlow(TreeSet treeSet, double mu, int reserveDepth, boolean useTail) { + super(treeSet); + validate(mu, reserveDepth); + this.mu = mu; + this.reserveDepth = reserveDepth; + this.useTail = useTail; + } + + /** + * Builds an MRegCCDSlow on {@code trees} with {@code mu} selected by maximising cross-validated + * held-out log-probability (see {@link ccd.algorithms.regularisation.MRegCCDParameterOptimiser}), + * rather than the fixed {@link #DEFAULT_MU}. The honest, no-peeking counterpart of + * {@code KRegCCD.withOptimisedParameters}. + */ + public static MRegCCDSlow withOptimisedMu(List trees) { + double mu = ccd.algorithms.regularisation.MRegCCDParameterOptimiser.optimiseMu(trees).mu(); + return new MRegCCDSlow(trees, 0.0, mu); + } + + private static void validate(double mu, int reserveDepth) { + if (mu <= 0 || mu >= 1) { + throw new IllegalArgumentException("mu must be in (0, 1), got " + mu); + } + if (reserveDepth < 2) { + throw new IllegalArgumentException("reserveDepth must be >= 2, got " + reserveDepth); + } + } + + /** The per-clade escape probability this model was built with. */ + public double getMu() { + return mu; + } + + public int getReserveDepth() { + return reserveDepth; + } + + /** + * Reserve counts {@code M_m(C)} by boundary size {@code m} (array index {@code m}, valid for + * {@code m = 2..min(|C|, reserveDepth)}); {@code M_m} is the number of all-novel resolutions of + * {@code C} with an {@code m}-part boundary. The first coefficient {@code M_2} (the {@code eps^1} + * term) is exactly the number of CCD0-expanded splits of {@code C} -- recombinations of two + * observed subclades whose split was never observed -- since those are the only escapes with no + * other novel (blue) clade. Exposed for inspection and cross-checks. + */ + public int[] reserveCounts(BitSet cladeInBits) { + return countsFor(cladeInBits).clone(); + } + + @Override + public String toString() { + return "MRegCCDSlow [mu = " + mu + ", reserveDepth = " + reserveDepth + ", tail = " + useTail + + ", per-new-split, full support]"; + } + + /* ---------------------------------------------------------------------- + * Scoring + * ------------------------------------------------------------------- */ + + @Override + public double getLogProbabilityOfTree(Tree tree) { + return scoreTree(tree, mu); + } + + /** + * Full-support log-probability at an arbitrary escape probability {@code scoreMu}, reusing this + * model's ({@code mu}-independent) backbone and cached reserve counts. Lets a parameter search / + * cross-validation evaluate many {@code mu} on one trained model without rebuilding. For + * {@code scoreMu == mu} it equals {@link #getLogProbabilityOfTree(Tree)}. + */ + public double getLogProbabilityOfTree(Tree tree, double scoreMu) { + if (scoreMu <= 0 || scoreMu >= 1) { + throw new IllegalArgumentException("scoreMu must be in (0, 1), got " + scoreMu); + } + return scoreTree(tree, scoreMu); + } + + @Override + public double getProbabilityOfTree(Tree tree) { + return Math.exp(getLogProbabilityOfTree(tree)); + } + + /** Always true: MRegCCDSlow is full support, so every tree on this taxon set has positive probability. */ + @Override + public boolean containsTree(Tree tree) { + return true; + } + + private double scoreTree(Tree tree, double scoreMu) { + Map bits = new HashMap<>(); + computeBits(tree.getRoot(), bits); + double logp = 0.0; + for (Node v : tree.getNodesAsArray()) { + if (v.isLeaf()) { + continue; + } + BitSet vb = bits.get(v); + Clade c = getClade(vb); + if (c == null) { + continue; // novel clade: scored once at its maximal region's top + } + BitSet b1 = bits.get(v.getChildren().get(0)); + BitSet b2 = bits.get(v.getChildren().get(1)); + if (isSplitObserved(vb, b1, b2)) { + if (reservable(vb)) { // discount only clades that can actually escape + double resv = Math.min(scoreMu + (useTail ? tailFor(vb, scoreMu) : 0.0), 1 - 1e-12); + logp += Math.log(1.0 - resv); + } + logp += rawLogCCP(c, b1, b2); // raw CCD1 CCP + } else { + // region top: an observed clade resolved through a novel split. m-1 new splits. + int m = boundarySize(v, bits); + logp += (m - 1) * Math.log(epsFor(vb, scoreMu)); + } + } + return logp; + } + + /* ---------------------------------------------------------------------- + * Per-clade reserve (M_m counts -> eps, tail; mirrors KRegCCD.computeReg) + * ------------------------------------------------------------------- */ + + /** Whether clade {@code C} (given in bits) reserves any escape mass up to {@code reserveDepth}. */ + boolean reservable(BitSet C) { + for (int v : countsFor(C)) { + if (v > 0) { + return true; + } + } + return false; + } + + /** Escape root {@code eps} solving {@code sum_{m>=2} M_m eps^(m-1) = scoreMu} (monotone bisection). */ + double epsFor(BitSet C, double scoreMu) { + int[] n = countsFor(C); + if (!reservable(C)) { + return scoreMu; // crude fallback (no escape route within reserveDepth); should not be hit + } + return solveEps(n, scoreMu); + } + + /** Omitted-tail escape mass beyond the computed orders: geometric bound from the top two orders. */ + double tailFor(BitSet C, double scoreMu) { + int[] n = countsFor(C); + int last = n.length - 1; + if (last < 3) { + return 0.0; + } + int nLast = n[last], nPrev = n[last - 1]; + if (nLast <= 0 || nPrev <= 0) { + return 0.0; + } + double eps = epsFor(C, scoreMu); + double rho = ((double) nLast / nPrev) * eps; + if (rho <= 0 || rho >= 1) { + return 0.0; + } + return Math.min(nLast * Math.pow(eps, last - 1) * rho / (1 - rho), scoreMu); + } + + /** M_m counts (index m = boundary size, 2..min(|C|, reserveDepth)); cached, mu-independent. */ + int[] countsFor(BitSet C) { + int[] cached = countsCache.get(C); + if (cached != null) { + return cached; + } + int card = C.cardinality(); + int[] n = new int[Math.min(card, reserveDepth) + 1]; + if (card >= 2) { + List subs = subclades(C); + enumOps = 0; + for (int m = 2; m < n.length; m++) { + try { + n[m] = countBoundaries(C, subs, m); + } catch (BudgetExceeded e) { + break; // deeper orders omitted (negligible, like the tail) + } + } + } + countsCache.put(C, n); + return n; + } + + private static double solveEps(int[] n, double mu) { + double lo = 0.0, hi = 1.0; + while (evalReserve(n, hi) < mu) { + hi *= 2.0; + } + for (int it = 0; it < 100; it++) { + double mid = 0.5 * (lo + hi); + if (evalReserve(n, mid) < mu) { + lo = mid; + } else { + hi = mid; + } + } + return 0.5 * (lo + hi); + } + + /** {@code sum_{m>=2} n[m] x^(m-1)}. */ + private static double evalReserve(int[] n, double x) { + double s = 0.0; + for (int m = 2; m < n.length; m++) { + if (n[m] > 0) { + s += n[m] * Math.pow(x, m - 1); + } + } + return s; + } + + /** Count m-part boundaries of C into observed subclades, weighted by their all-novel pathcount. */ + private int countBoundaries(BitSet C, List subs, int m) { + return enumerateBoundaries(C, subs, m, 0, BitSet.newBitSet(leafArraySize), new ArrayList<>(m)); + } + + private int enumerateBoundaries(BitSet C, List subs, int m, int startIdx, + BitSet used, List chosen) { + if (++enumOps > OPS_BUDGET) { + throw BUDGET_EXCEEDED; + } + if (chosen.size() == m - 1) { + BitSet last = BitSet.newBitSet(C); + last.andNot(used); + if (last.isEmpty() || !isObs(last)) { + return 0; + } + if (compareBitSets(chosen.get(chosen.size() - 1), last) >= 0) { + return 0; // canonical: the derived last part must be the largest + } + BitSet[] parts = new BitSet[m]; + for (int i = 0; i < m - 1; i++) { + parts[i] = chosen.get(i); + } + parts[m - 1] = last; + return countAllNovelResolutions(C, parts); + } + int count = 0; + for (int i = startIdx; i < subs.size(); i++) { + BitSet pb = subs.get(i); + if (pb.intersects(used)) { + continue; + } + chosen.add(pb); + BitSet newUsed = BitSet.newBitSet(used); + newUsed.or(pb); + count += enumerateBoundaries(C, subs, m, i + 1, newUsed, chosen); + chosen.remove(chosen.size() - 1); + } + return count; + } + + /** + * Number of all-novel binary resolutions of C into the given observed parts (subset DP over the + * parts). A split is allowed iff: at the region root (full mask = C, an observed clade) the split + * is unobserved (a real escape); at an intermediate node the clade itself is novel (a maximal + * region stops at observed clades, matching {@link #boundarySize}). + */ + int countAllNovelResolutions(BitSet C, BitSet[] parts) { + int k = parts.length; + if (k == 1) { + return 1; + } + int full = (1 << k) - 1; + BitSet[] unionOf = new BitSet[1 << k]; + unionOf[0] = BitSet.newBitSet(leafArraySize); + for (int mask = 1; mask <= full; mask++) { + int low = Integer.numberOfTrailingZeros(mask); + BitSet u = BitSet.newBitSet(unionOf[mask & (mask - 1)]); + u.or(parts[low]); + unionOf[mask] = u; + } + int[] f = new int[1 << k]; + for (int mask = 1; mask <= full; mask++) { + if (Integer.bitCount(mask) == 1) { + f[mask] = 1; + continue; + } + int low = mask & (-mask), rest = mask ^ low, count = 0; + for (int sub = rest; ; sub = (sub - 1) & rest) { + int s1 = sub | low, s2 = mask ^ s1; + if (s2 != 0 && splitAllowed(mask == full, unionOf[mask], unionOf[s1], unionOf[s2])) { + count += f[s1] * f[s2]; + } + if (sub == 0) { + break; + } + } + f[mask] = count; + } + return f[full]; + } + + private boolean splitAllowed(boolean top, BitSet union, BitSet a, BitSet b) { + return top ? !isSplitObserved(union, a, b) : !isObs(union); + } + + /* ---------------------------------------------------------------------- + * Observed-backbone queries (over the inherited CCD1 clade DAG) + * ------------------------------------------------------------------- */ + + boolean isObs(BitSet x) { + return getClade(x) != null; // leaves are clades too + } + + private boolean isSplitObserved(BitSet parentBits, BitSet aBits, BitSet bBits) { + Clade parent = getClade(parentBits); + if (parent == null) { + return false; + } + Clade a = getClade(aBits); + Clade b = getClade(bBits); + if (a == null || b == null) { + return false; + } + return parent.getCladePartition(a, b) != null; + } + + private double rawLogCCP(Clade parent, BitSet aBits, BitSet bBits) { + CladePartition p = parent.getCladePartition(getClade(aBits), getClade(bBits)); + return p.getLogCCP(); + } + + /** Observed clades (incl. leaves) strictly contained in C, in canonical order; cached. */ + List subclades(BitSet C) { + return subCache.computeIfAbsent(C, c -> { + List out = new ArrayList<>(); + int card = c.cardinality(); + for (BitSet x : sortedCladeBits()) { + if (x.cardinality() < card && subset(x, c)) { + out.add(x); + } + } + return out; + }); + } + + private List sortedCladeBits() { + if (sortedCladeBits == null) { + List all = new ArrayList<>(); + for (Clade c : getClades()) { + all.add(c.getCladeInBits()); + } + all.sort(MRegCCDSlow::compareBitSets); + sortedCladeBits = all; + } + return sortedCladeBits; + } + + /** Boundary size of the maximal region rooted at v: count of maximal observed/leaf subclades below. */ + private int boundarySize(Node v, Map bits) { + int m = 0; + for (Node child : v.getChildren()) { + if (child.isLeaf() || getClade(bits.get(child)) != null) { + m++; + } else { + m += boundarySize(child, bits); + } + } + return m; + } + + 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(); + } + + /** Canonical total order on clade bitsets (lexicographic by set-bit indices). */ + static int compareBitSets(BitSet a, BitSet b) { + int ia = a.nextSetBit(0), 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); + } + + /* ---------------------------------------------------------------------- + * Sampling (self-consistent) + * + * The PIT calibration test draws trees from the model and needs only each draw's log-probability + * (not the tree object), so we override sampleTreeLogProbability() with a direct simulation of the + * generative process and never materialise a Tree. At each reservable clade we escape with + * probability equal to its escape mass (= mu by the eps-solve, tail EXCLUDED) and otherwise take + * an observed (red) split ~ CCP; an escape draws a region order m proportional to M_m eps^(m-1), a + * boundary of m observed subclades proportional to its all-novel pathcount, and recurses into the + * boundary parts. The resolution shape within a region is not drawn -- every shape has the same + * weight eps^(m-1) and does not change the draw's log-probability -- so the simulation is cheap. + * + * The draw distribution exactly matches getLogProbabilityOfTree when the model is built with the + * tail OFF (then the red discount is 1 - mu, matching the escape mass), for trees whose regions are + * within reserveDepth; deeper regions (mass ~mu^reserveDepth) are never produced, the same + * self-consistent / full-support trade-off KRegCCD makes for its PIT. + * ------------------------------------------------------------------- */ + + @Override + public double sampleTreeLogProbability() { + return simulate(getRootClade()); + } + + private double simulate(Clade c) { + if (c.isLeaf()) { + return 0.0; + } + BitSet cb = c.getCladeInBits(); + if (reservable(cb)) { + double eps = epsFor(cb, mu); + int[] n = countsFor(cb); + double escapeMass = 0.0; + for (int m = 2; m < n.length; m++) { + if (n[m] > 0) { + escapeMass += n[m] * Math.pow(eps, m - 1); + } + } + if (random.nextDouble() < escapeMass) { + int m = sampleOrder(n, eps, escapeMass); + BitSet[] parts = sampleBoundaryParts(cb, subclades(cb), m); + double logp = (m - 1) * Math.log(eps); + if (parts != null) { + for (BitSet bp : parts) { + logp += simulate(getClade(bp)); + } + } + return logp; + } + CladePartition p = samplePartition(c); + double logp = Math.log(1.0 - escapeMass) + p.getLogCCP(); + return logp + simulate(p.getChildClades()[0]) + simulate(p.getChildClades()[1]); + } + CladePartition p = samplePartition(c); // non-reservable: observed split, no discount + return p.getLogCCP() + simulate(p.getChildClades()[0]) + simulate(p.getChildClades()[1]); + } + + /** Draws a region order m in {2..} with probability proportional to {@code M_m eps^(m-1)}. */ + private int sampleOrder(int[] n, double eps, double escapeMass) { + double target = random.nextDouble() * escapeMass, acc = 0.0; + for (int m = 2; m < n.length; m++) { + if (n[m] > 0) { + acc += n[m] * Math.pow(eps, m - 1); + if (target < acc) { + return m; + } + } + } + for (int m = n.length - 1; m >= 2; m--) { + if (n[m] > 0) { + return m; // numerical guard + } + } + return 2; + } + + /** Samples an observed (red) split of {@code c} with probability proportional to its CCP. */ + private CladePartition samplePartition(Clade c) { + List partitions = c.getPartitions(); + double target = random.nextDouble(), acc = 0.0; + for (CladePartition p : partitions) { + acc += p.getCCP(); + if (target < acc) { + return p; + } + } + return partitions.get(partitions.size() - 1); + } + + /** + * Weighted-reservoir samples one boundary of {@code c} into {@code m} observed subclades, + * proportional to its all-novel pathcount (so that, combined with order sampling, every distinct + * novel resolution is equiprobable at {@code eps^(m-1)}). Returns the parts, or {@code null} if + * none/op-budget. + */ + private BitSet[] sampleBoundaryParts(BitSet c, List subs, int m) { + boundaryPick = null; + boundaryWeightSeen = 0.0; + enumOps = 0; + try { + sampleBoundaryWalk(c, subs, m, 0, BitSet.newBitSet(leafArraySize), new ArrayList<>(m)); + } catch (BudgetExceeded e) { + return boundaryPick; // whatever was picked before the cap (may be null) + } + return boundaryPick; + } + + private BitSet[] boundaryPick; + private double boundaryWeightSeen; + + private void sampleBoundaryWalk(BitSet c, List subs, int m, int startIdx, + BitSet used, List chosen) { + if (++enumOps > OPS_BUDGET) { + throw BUDGET_EXCEEDED; + } + if (chosen.size() == m - 1) { + BitSet last = BitSet.newBitSet(c); + last.andNot(used); + if (last.isEmpty() || !isObs(last)) { + return; + } + if (compareBitSets(chosen.get(chosen.size() - 1), last) >= 0) { + return; + } + BitSet[] parts = new BitSet[m]; + for (int i = 0; i < m - 1; i++) { + parts[i] = chosen.get(i); + } + parts[m - 1] = last; + int pc = countAllNovelResolutions(c, parts); + if (pc <= 0) { + return; + } + boundaryWeightSeen += pc; + if (random.nextDouble() * boundaryWeightSeen < pc) { // weighted reservoir + boundaryPick = parts; + } + return; + } + for (int i = startIdx; i < subs.size(); i++) { + BitSet pb = subs.get(i); + if (pb.intersects(used)) { + continue; + } + chosen.add(pb); + BitSet newUsed = BitSet.newBitSet(used); + newUsed.or(pb); + sampleBoundaryWalk(c, subs, m, i + 1, newUsed, chosen); + chosen.remove(chosen.size() - 1); + } + } + + @Override + public Tree sampleTree(HeightSettingStrategy heightStrategy) { + throw new UnsupportedOperationException( + "MRegCCDSlow materialised-tree sampling is not implemented; sampleTreeLogProbability() " + + "(used by the PIT) simulates draws without building trees."); + } +} diff --git a/src/main/java/ccd/model/bitsets/BitSet.java b/src/main/java/ccd/model/bitsets/BitSet.java index a6b6c37..3c727fe 100644 --- a/src/main/java/ccd/model/bitsets/BitSet.java +++ b/src/main/java/ccd/model/bitsets/BitSet.java @@ -55,7 +55,10 @@ public static BitSet newBitSet(BitSet other) { if (other instanceof BitSet256 set) { return new BitSet256(set); } - BitSet b = new BitSet(other.length()); + // size(), not length(): length() is the index of the highest set bit plus one, so copying a + // set whose top words happen to be empty would return an undersized BitSet, and the bitwise + // operations below index the operand by this.words.length. + BitSet b = new BitSet(other.size()); b.or(other); return b; } diff --git a/src/test/java/ccd/model/MRegCCDAgreementTest.java b/src/test/java/ccd/model/MRegCCDAgreementTest.java new file mode 100644 index 0000000..a7d18f8 --- /dev/null +++ b/src/test/java/ccd/model/MRegCCDAgreementTest.java @@ -0,0 +1,133 @@ +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; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** MRegCCD must reproduce MRegCCDSlow exactly: same boundary counts, same tree probabilities. */ +public class MRegCCDAgreementTest { + + 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; + } + + @Test + public void boundaryCountsAndProbabilitiesAgree() { + int cladesChecked = 0; + for (int n : new int[]{6, 8, 10, 14}) { + for (int nTrees : new int[]{5, 25}) { + for (int depth : new int[]{2, 3, 4}) { + List tx = taxa(n); + MRegCCDSlow slow = new MRegCCDSlow(randomTrees(tx, nTrees, 21L), 0.0, 0.02, depth, true); + MRegCCD fast = new MRegCCD(randomTrees(tx, nTrees, 21L), 0.0, 0.02, depth, true); + + for (Clade c : slow.getClades()) { + BitSet cb = c.getCladeInBits(); + assertArrayEquals(slow.countsFor(cb), fast.countsFor(cb), + "boundary counts differ at " + n + " taxa, depth " + depth + + ", clade " + cb); + cladesChecked++; + } + + List probe = randomTrees(tx, 40, 99L); + for (Tree t : probe) { + assertEquals(slow.getLogProbabilityOfTree(t), fast.getLogProbabilityOfTree(t), + 1e-9, "tree probability differs at " + n + " taxa, depth " + depth); + } + } + } + } + System.out.printf("MRegCCD agrees with MRegCCDSlow on %d clades and every probed tree%n", + cladesChecked); + } + + /** + * The exact-normalisation guarantee is stated for the model, and MRegCCDTest checks it on the + * reference implementation. Since MRegCCD is the class callers get, and its fast path computes + * the counts that the reserve is solved from, check the property directly on it too. + */ + @Test + public void fastImplementationIsExactlyNormalisedAtFullDepth() { + for (int n : new int[]{5, 6}) { + for (double mu : new double[]{0.02, 0.1, 0.25}) { + List tx = taxa(n); + // full reserve depth: no omitted tail, so the model must normalise exactly + MRegCCD m = new MRegCCD(randomTrees(tx, 6, 5L), 0.0, mu, tx.size(), false); + double sum = 0.0; + for (Tree t : allRootedTopologies(tx)) { + sum += Math.exp(m.getLogProbabilityOfTree(t)); + } + System.out.printf("MRegCCD %d taxa mu=%.2f full-depth SUM = %.12f%n", n, mu, sum); + assertEquals(1.0, sum, 1e-9, + "MRegCCD at full reserve depth must be exactly normalised"); + } + } + } + + private static List allRootedTopologies(List taxa) { + List out = new ArrayList<>(); + for (String shape : shapes(taxa)) { + out.add(new TreeParser(taxa, shape + ";", 1, false)); + } + return out; + } + + 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/MRegCCDTest.java b/src/test/java/ccd/model/MRegCCDTest.java index ea8eb4a..4c5b8ad 100644 --- a/src/test/java/ccd/model/MRegCCDTest.java +++ b/src/test/java/ccd/model/MRegCCDTest.java @@ -13,7 +13,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue; /** - * Validates the one-parameter per-new-split {@link MRegCCD}: that with the full reserve depth it is an + * Validates the one-parameter per-new-split {@link MRegCCDSlow}: that with the full reserve depth it is an * exactly normalised distribution on enumerable taxon sets, and that truncating the reserve without a * tail correction super-normalises it (the artefact that made order-2 falsely appear to close the gap * to KRegCCD in the RSV2 experiment). @@ -64,7 +64,7 @@ private static List trees(List taxa, List shapes) { return out; } - private static double totalMass(MRegCCD m, List taxa) { + private static double totalMass(MRegCCDSlow m, List taxa) { double sum = 0.0; for (T t : allTopologies(taxa)) { Tree tree = new TreeParser(taxa, topo(t) + ";", 1, false); @@ -85,10 +85,10 @@ private void check5(double mu) { List taxa = Arrays.asList("A", "B", "C", "D", "E"); List train = trees(taxa, List.of(cat("A", "B", "C", "D", "E"), cat("D", "C", "B", "A", "E"))); - MRegCCD m = new MRegCCD(train, 0.0, mu, taxa.size(), false); // full depth -> no omitted tail + MRegCCDSlow m = new MRegCCDSlow(train, 0.0, mu, taxa.size(), false); // full depth -> no omitted tail double sum = totalMass(m, taxa); - System.out.printf("MRegCCD 5 taxa mu=%.2f full-depth SUM = %.12f%n", mu, sum); - assertEquals(1.0, sum, 1e-9, "MRegCCD at full reserve depth must be exactly normalised"); + System.out.printf("MRegCCDSlow 5 taxa mu=%.2f full-depth SUM = %.12f%n", mu, sum); + assertEquals(1.0, sum, 1e-9, "MRegCCDSlow at full reserve depth must be exactly normalised"); } private void check6(double mu) { @@ -96,10 +96,10 @@ private void check6(double mu) { List train = trees(taxa, List.of(new Node(cat("A", "B", "C", "D"), new Node(new Leaf("E"), new Leaf("F"))), new Node(cat("D", "C", "B", "A"), new Node(new Leaf("E"), new Leaf("F"))))); - MRegCCD m = new MRegCCD(train, 0.0, mu, taxa.size(), false); + MRegCCDSlow m = new MRegCCDSlow(train, 0.0, mu, taxa.size(), false); double sum = totalMass(m, taxa); - System.out.printf("MRegCCD 6 taxa mu=%.2f full-depth SUM = %.12f%n", mu, sum); - assertEquals(1.0, sum, 1e-9, "MRegCCD at full reserve depth must be exactly normalised"); + System.out.printf("MRegCCDSlow 6 taxa mu=%.2f full-depth SUM = %.12f%n", mu, sum); + assertEquals(1.0, sum, 1e-9, "MRegCCDSlow at full reserve depth must be exactly normalised"); } @Test @@ -112,7 +112,7 @@ public void m2EqualsCCD0ExpandedSplits() { for (int i = 0; i < all.size(); i += 47) picks.add(all.get(i)); // ~20 trees spread across the space List train = trees(taxa, picks); - MRegCCD mreg = new MRegCCD(train, 0.0, 0.05); + MRegCCDSlow mreg = new MRegCCDSlow(train, 0.0, 0.05); CCD0 ccd0 = new CCD0(train, 0); int checked = 0, withRecomb = 0; @@ -146,7 +146,7 @@ public void samplerMatchesScorer() { List picks = new ArrayList<>(); for (int i = 0; i < all.size(); i += 31) picks.add(all.get(i)); List train = trees(taxa, picks); - MRegCCD m = new MRegCCD(train, 0.0, 0.1, taxa.size(), false); // full depth, tail off + MRegCCDSlow m = new MRegCCDSlow(train, 0.0, 0.1, taxa.size(), false); // full depth, tail off // true entropy and normalisation by enumeration (scorer) double sum = 0.0, H = 0.0; @@ -169,7 +169,7 @@ public void samplerMatchesScorer() { } double hHat = s1 / N; double se = Math.sqrt(Math.max(0, s2 / N - hHat * hHat) / N); - System.out.printf("MRegCCD sampler: H_enum=%.5f H_MC=%.5f +/- %.5f (%.1f SE off)%n", + System.out.printf("MRegCCDSlow 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.01), "sampler entropy must match the scorer's enumerated entropy"); @@ -182,9 +182,9 @@ public void truncatedReserveSuperNormalises() { List.of(new Node(cat("A", "B", "C", "D"), new Node(new Leaf("E"), new Leaf("F"))), new Node(cat("D", "C", "B", "A"), new Node(new Leaf("E"), new Leaf("F"))))); double mu = 0.2; - double full = totalMass(new MRegCCD(train, 0.0, mu, taxa.size(), false), taxa); - double order2 = totalMass(new MRegCCD(train, 0.0, mu, 2, false), taxa); // M2 only, no tail - System.out.printf("MRegCCD 6 taxa mu=%.2f: full-depth SUM=%.9f order-2 SUM=%.9f%n", mu, full, order2); + double full = totalMass(new MRegCCDSlow(train, 0.0, mu, taxa.size(), false), taxa); + double order2 = totalMass(new MRegCCDSlow(train, 0.0, mu, 2, false), taxa); // M2 only, no tail + System.out.printf("MRegCCDSlow 6 taxa mu=%.2f: full-depth SUM=%.9f order-2 SUM=%.9f%n", mu, full, order2); assertEquals(1.0, full, 1e-9, "full depth normalised"); assertTrue(order2 > 1.0 + 1e-4, "order-2 reserve (no tail) must super-normalise (sum > 1), got " + order2); diff --git a/src/test/java/ccd/model/bitsets/BitSetCopyTest.java b/src/test/java/ccd/model/bitsets/BitSetCopyTest.java new file mode 100644 index 0000000..34944f8 --- /dev/null +++ b/src/test/java/ccd/model/bitsets/BitSetCopyTest.java @@ -0,0 +1,54 @@ +package ccd.model.bitsets; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Copies of a BitSet must keep the source's capacity, not shrink to its highest set bit. + * + *

The bitwise operations iterate over {@code this.words.length} and index the operand directly, + * so an undersized copy makes them throw. This bites only above 256 bits, where the generic BitSet + * is used instead of the fixed-size subclasses, and only when the top words are empty -- the common + * case for a clade that does not contain the highest-numbered taxa. + */ +public class BitSetCopyTest { + + @Test + public void copyKeepsCapacityAboveTheSpecialisedSizes() { + for (int nbits : new int[]{276, 320, 512, 1000}) { + BitSet full = BitSet.newBitSet(nbits); + full.set(nbits - 1); + + BitSet sparse = BitSet.newBitSet(nbits); + sparse.set(3); + BitSet copy = BitSet.newBitSet(sparse); + + assertEquals(sparse.size(), copy.size(), + "copy must keep the source capacity at " + nbits + " bits"); + + BitSet a = BitSet.newBitSet(full); + a.andNot(copy); + assertTrue(a.get(nbits - 1), "andNot must not clear unrelated high bits"); + + BitSet b = BitSet.newBitSet(copy); + b.andNot(full); + assertTrue(b.get(3), "andNot must keep the low bit"); + assertFalse(copy.intersects(full), "disjoint sets must not intersect"); + } + } + + @Test + public void copyOfAnEmptySetIsUsable() { + BitSet empty = BitSet.newBitSet(400); + BitSet copy = BitSet.newBitSet(empty); + assertEquals(empty.size(), copy.size(), "an empty copy must still have capacity"); + BitSet other = BitSet.newBitSet(400); + other.set(399); + copy.andNot(other); + copy.or(other); + assertTrue(copy.get(399)); + } +} From d0af71ccad40114975a8041d2c98809154c58a37 Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Thu, 20 Aug 2026 09:53:10 +1200 Subject: [PATCH 4/8] Add HeadToHeadBenchmark, a single-JVM driver for the model comparison The comparison was driven by a JUnit test forked once per dataset per model via Maven. Class loading of BEAST and CCD takes 30ms, so there was nothing to amortise by forking; project scan, dependency resolution, plugin loading and the surefire fork cost several seconds each, and a sweep spent over 99% of wall clock on startup -- 11m39s of which 2.4s was modelling. HeadToHeadBenchmark is a plain main that loops all datasets and models in one JVM. The full corpus of 101 datasets now takes about 26 minutes. It resumes: a sweep reads back the rows already in the CSV, skips those datasets and appends. Three separate runs were lost to a crash partway through (heap exhaustion, and a SIGBUS once swap was exhausted), each costing every dataset already scored. Failures are caught per model per dataset and recorded in the row, so a model that fails on one dataset costs that cell rather than the row or the run -- KRegCCD and MRegCCD both exhaust an 8GB heap on the 441-taxon stervander set while CRegCCD completes it. Hyperparameters are selected on a validation quarter disjoint from the fitted quarter. Drawing the validation trees from within the fitted set drove every parameter to its grid floor. The alpha grid is capped at 5: alpha is a per-split pseudocount, so beyond alpha = 1 the prior outweighs a real observation, and the only principled reason to exceed it is MCMC autocorrelation inflating counts by N/ESS -- the cap admits that down to ESS = 200 of 1000 training trees. Parameters are serialised with Double.toString. A "%.4f" format collapsed the alpha2 bands at 1.6e-5, 8e-5 and 4e-4 all to "0.0000", so the selected value could not be recovered from the CSV at all. Co-Authored-By: Claude Opus 5 (1M context) --- .../regularisation/NNIHeldOutComparison.java | 7 +- src/main/java/ccd/model/KRegCCD.java | 31 +- .../java/ccd/tools/HeadToHeadBenchmark.java | 267 +++++++++++ .../java/ccd/model/MRegDepthTimingTest.java | 58 +++ src/test/java/ccd/model/MunroParseCheck.java | 54 +++ .../ccd/model/RealDataHeadToHeadTest.java | 430 ++++++++++++++++-- 6 files changed, 789 insertions(+), 58 deletions(-) create mode 100644 src/main/java/ccd/tools/HeadToHeadBenchmark.java create mode 100644 src/test/java/ccd/model/MRegDepthTimingTest.java create mode 100644 src/test/java/ccd/model/MunroParseCheck.java diff --git a/src/main/java/ccd/algorithms/regularisation/NNIHeldOutComparison.java b/src/main/java/ccd/algorithms/regularisation/NNIHeldOutComparison.java index f7a71eb..760d08c 100644 --- a/src/main/java/ccd/algorithms/regularisation/NNIHeldOutComparison.java +++ b/src/main/java/ccd/algorithms/regularisation/NNIHeldOutComparison.java @@ -97,9 +97,10 @@ private static List modelSpecs(double alpha) { specs.add(new ModelSpec(String.format("NNIRegCCD[CO_OCCURRING,b=%.3f]", beta), tr -> new NNIRegCCD(tr, 0.0, PairingMode.CO_OCCURRING, alpha, beta))); } - // full-support KRegCCD (reserve depth 2); compare the three tail modes to - // show the normalisation effect on mean logP: NONE (super-normalised), - // BOUND (upper bound, sub-normalised), SAMPLED (Knuth, near-exact) + // full-support KRegCCD (reserve depth 2); compare the three tail modes to show the + // normalisation effect on mean logP. All three are sub-normalised by Theta(mu^2) (see + // KRegCCD.TailMode); they differ only in the O(mu^3) truncation term: NONE omits it, + // BOUND over-corrects it, SAMPLED estimates it (landing within noise of NONE). for (double mu : new double[]{0.01, 0.05, 0.1, 0.2}) { specs.add(new ModelSpec(String.format("KRegCCD[rd=2,none,mu=%.2f]", mu), tr -> new KRegCCD(tr, 0.0, mu, alpha, 2, KRegCCD.TailMode.NONE))); diff --git a/src/main/java/ccd/model/KRegCCD.java b/src/main/java/ccd/model/KRegCCD.java index 0018769..981231c 100644 --- a/src/main/java/ccd/model/KRegCCD.java +++ b/src/main/java/ccd/model/KRegCCD.java @@ -44,9 +44,11 @@ * {@code eps} is solved from only the cheap low orders {@code j = 1..k} (boundaries * {@code 3..k+2}; {@code N_1} is {@code O(m^2)}, {@code N_2} is {@code O(m^3)} in * the number {@code m} of observed subclades), while regions of any depth are still - * scored with that {@code eps}. The omitted deep orders contribute {@code O(mu^2/m)} - * to the reserve — a rounding error — so this barely affects probabilities while - * keeping the model tractable and full-support. Reserve depth {@code k = 2} (the + * scored with that {@code eps}. Since {@code eps ~ mu / N_1}, the omitted deep orders contribute + * {@code O(mu^(k+1))} to the reserve — {@code O(mu^3)} at the default {@code k = 2}, a rounding + * error — so this barely affects probabilities while keeping the model tractable and full-support. + * (Measured by exhaustive enumeration on 6- and 7-taxon sets: the truncation shifts total + * probability mass with a log-log slope in {@code mu} of 2.00 / 2.97 / 3.96 at {@code k = 1/2/3}.) Reserve depth {@code k = 2} (the * default) is recommended; an op-budget (scaled to the largest clade's N_1, overridable via * {@code -Dkreg.enumOps}) * bounds even the low-order enumeration on pathological clades, and a clade with @@ -112,13 +114,24 @@ private static final class BudgetExceeded extends RuntimeException { * How to correct for the omitted reserve tail (orders > reserve depth) when * discounting red splits by {@code 1 - mu - tail(C)}: *

    - *
  • {@link #NONE}: no correction ({@code tail = 0}); slightly - * super-normalised (inflates held-out scores by the tail).
    • - *
  • {@link #BOUND}: geometric upper bound on the tail; provably - * sub-normalised (never inflates).
    • - *
  • {@link #SAMPLED}: Knuth estimate of the actual tail; near-exactly - * normalised (up to Monte-Carlo noise).
    • + *
  • {@link #NONE}: no correction ({@code tail = 0}). The truncation alone inflates by + * {@code O(mu^(k+1))}, but that is dominated by the maximality deficit below, so the + * model is net sub-normalised.
    • + *
  • {@link #BOUND}: geometric upper bound on the tail; sub-normalised (never inflates).
    • + *
  • {@link #SAMPLED}: Knuth estimate of the actual tail. This removes the + * {@code O(mu^(k+1))} truncation, but NOT the maximality deficit, so it is + * not exactly normalised — it lands within Monte-Carlo noise of {@code NONE}.
    • *
+ * + *

All three modes are sub-normalised by {@code Theta(mu^2)}, and no tail correction + * removes it: a blue region is only scored at its maximal top, so each boundary part + * that is itself a reserving clade contributes {@code (1 - mu - tail)} rather than 1 — the + * {@code O(mu)}-per-reserving-boundary-part gap noted on {@link SamplingFidelity}. Summing + * {@link #getProbabilityOfTree} over every tree on 7 taxa (2 training trees, {@code k = 2}) + * gives total mass {@code -4.7e-3 / -6.3e-3 / -4.7e-3} from 1 for NONE/BOUND/SAMPLED at + * {@code mu = 0.05}, and {@code -2.6e-6 / -2.7e-6 / -2.7e-6} at {@code mu = 0.001}; + * the deficit scales as {@code mu^2} and persists at reserve depths where the truncation + * has fully converged. */ public enum TailMode {NONE, BOUND, SAMPLED} diff --git a/src/main/java/ccd/tools/HeadToHeadBenchmark.java b/src/main/java/ccd/tools/HeadToHeadBenchmark.java new file mode 100644 index 0000000..0b5200d --- /dev/null +++ b/src/main/java/ccd/tools/HeadToHeadBenchmark.java @@ -0,0 +1,267 @@ +package ccd.tools; + +import beast.base.evolution.tree.Tree; +import beastfx.app.treeannotator.TreeAnnotator; +import ccd.model.CCD1; +import ccd.model.CRegCCD; +import ccd.model.KRegCCD; +import ccd.model.MRegCCD; +import ccd.model.RegCCD; + +import java.io.File; +import java.io.FileWriter; +import java.io.PrintWriter; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.Set; +import java.util.List; +import java.util.Map; + +/** + * Held-out predictive comparison of the CCD variants across many prepared posteriors, in one JVM. + * + * 

+ *   java -cp ... ccd.tools.HeadToHeadBenchmark <prepared-dir> <out.csv> <pertree-dir> [models]
+ *
+ * + *

Each dataset must already have been prepared into {@code base.fit/val/train/test.trees}: the + * model is fitted on {@code train}, its hyperparameters chosen on the disjoint {@code fit}/{@code + * val} pair, and scored on {@code test}. Preparing separately means each source chain is parsed + * once rather than once per model, and running every dataset in one process means the work is + * modelling rather than process startup. + * + *

Per-tree log probabilities are written per dataset and model, so paired statistics can be + * computed afterwards between any two models without re-running anything. + */ +public class HeadToHeadBenchmark { + + private static final List ALL = + List.of("CCD1", "RegCCD", "KRegCCD", "MRegCCD", "CRegCCD"); + + public static void main(String[] args) throws Exception { + File prepared = new File(args[0]); + File out = new File(args[1]); + File perTree = new File(args[2]); + List models = args.length > 3 ? Arrays.asList(args[3].split(",")) : ALL; + perTree.mkdirs(); + + List bases = new ArrayList<>(); + for (File f : prepared.listFiles((d, n) -> n.endsWith(".train.trees"))) { + bases.add(f.getName().replaceAll("\\.train\\.trees$", "")); + } + bases.sort(String::compareTo); + System.out.printf("%d datasets, models %s%n", bases.size(), models); + + // Resume support. A sweep over this corpus takes about an hour, and a JVM that dies partway + // (heap exhaustion, SIGBUS from an exhausted swap, an external kill) used to cost every + // dataset already scored. So read back whatever rows the CSV already holds and skip those + // datasets, appending rather than truncating. Delete the CSV to force a clean run. + Set done = new LinkedHashSet<>(); + boolean header = false; + if (out.exists() && out.length() > 0) { + List lines = java.nio.file.Files.readAllLines(out.toPath()); + for (int i = 0; i < lines.size(); i++) { + String line = lines.get(i); + if (line.isBlank()) continue; + if (i == 0 && line.startsWith("dataset,")) { header = true; continue; } + done.add(line.substring(0, line.indexOf(','))); + } + System.out.printf("resuming: %d datasets already in %s, skipping them%n", + done.size(), out.getName()); + } + PrintWriter csv = new PrintWriter(new FileWriter(out, true)); + PrintWriter failed = new PrintWriter(new FileWriter(new File(out.getPath() + ".failed"), true)); + long t0 = System.nanoTime(); + + for (String base : bases) { + if (done.contains(base)) continue; + Map row = new LinkedHashMap<>(); + try { + List train = read(prepared, base, "train"); + List test = read(prepared, base, "test"); + row.put("dataset", base); + row.put("taxa", String.valueOf(train.get(0).getLeafNodeCount())); + row.put("nTrain", String.valueOf(train.size())); + row.put("nTest", String.valueOf(test.size())); + row.put("testProtocol", + new File(prepared, base + ".siblingtest").exists() ? "sibling-run" : "chain-half"); + row.put("CCD1_entropy", fmt(new CCD1(read(prepared, base, "train"), 0.0).getEntropy())); + + for (String m : models) { + runModel(m, prepared, base, test, row, perTree); + } + if (!header) { + csv.println(String.join(",", row.keySet())); + header = true; + } + csv.println(String.join(",", row.values())); + csv.flush(); + System.out.printf(" %-70s %s taxa%n", base.substring(0, Math.min(70, base.length())), + row.get("taxa")); + } catch (Throwable e) { + failed.printf("%s,%s: %s%n", base, e.getClass().getSimpleName(), e.getMessage()); + failed.flush(); + System.out.printf(" FAILED %s (%s)%n", base, e.getClass().getSimpleName()); + } + } + csv.close(); + failed.close(); + System.out.printf("done in %.1f s%n", (System.nanoTime() - t0) / 1e9); + } + + private static void runModel(String model, File dir, String base, List test, + Map row, File perTree) { + try { + List val = read(dir, base, "val"); + String params; + long build; + Scorer scorer; + + switch (model) { + case "CCD1" -> { + long t = System.nanoTime(); + CCD1 m = new CCD1(read(dir, base, "train"), 0.0); + build = ms(t); + params = "-"; + scorer = m::getLogProbabilityOfTree; + } + case "RegCCD" -> { + double best = pick(new double[]{0.01, 0.05, 0.1, 0.2, 0.4, 0.8, 1.0}, a -> { + RegCCD f = new RegCCD(read(dir, base, "fit"), 0.0, a); + return mean(f::getLogProbabilityOfTree, val); + }); + long t = System.nanoTime(); + RegCCD m = new RegCCD(read(dir, base, "train"), 0.0, best); + build = ms(t); + params = "alpha=" + fmt(best); + scorer = m::getLogProbabilityOfTree; + } + case "KRegCCD" -> { + KRegCCD f = new KRegCCD(read(dir, base, "fit"), 0.0, 0.005, 0.4); + double best = pick(new double[]{0.00002, 0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05}, + mu -> mean(t -> f.getLogProbabilityOfTree(t, mu), val)); + long t = System.nanoTime(); + KRegCCD m = new KRegCCD(read(dir, base, "train"), 0.0, best, 0.4); + m.precomputeReserves(); + build = ms(t); + params = "alpha=0.4;mu=" + fmt(best); + scorer = m::getLogProbabilityOfTree; + } + case "MRegCCD" -> { + MRegCCD f = new MRegCCD(read(dir, base, "fit"), 0.0, MRegCCD.DEFAULT_MU); + double best = pick(new double[]{0.0002, 0.001, 0.002, 0.008, 0.0159, 0.05, 0.1}, + mu -> mean(t -> f.getLogProbabilityOfTree(t, mu), val)); + long t = System.nanoTime(); + MRegCCD m = new MRegCCD(read(dir, base, "train"), 0.0, best); + build = ms(t); + params = "mu=" + fmt(best); + scorer = m::getLogProbabilityOfTree; + } + case "CRegCCD" -> { + CRegCCD f = new CRegCCD(read(dir, base, "fit"), 0.0); + // alpha is a per-split pseudocount, so alpha > 1 lets the prior outweigh a real + // observation on every split at once: at alpha = 12 a split seen once is only + // 1.08x as probable as one never seen, and the counts stop carrying information. + // The one reason to allow alpha > 1 is that MCMC samples are autocorrelated, so + // f(S) overstates the evidence by roughly N/ESS; capping at 5 admits that + // correction down to ESS = 200 out of the 1000 training trees and no further. + double[] gridA = {0.002, 0.01, 0.05, 0.2, 0.4, 1.0, 2.0, 5.0}; + // alpha1 and alpha2 are class totals rather than per-split pseudocounts, so the + // argument above does not bound them; the bands are set by what is ever chosen. + // No dataset selected alpha1 above 2.0 or alpha2 above 1.0, so the 12.0 band is + // dead weight for both; each keeps one band of headroom above its observed max + // so that the top selected value is interior rather than pinned to a ceiling. + double[] gridA1 = {0.002, 0.01, 0.05, 0.2, 0.4, 1.0, 2.0, 5.0}; + // alpha2 was selected at the old floor of 0.002 on 43% of datasets, which is a + // boundary rather than an optimum, so the floor drops four bands to locate it. + double[] gridA2 = {0.000016, 0.00008, 0.0004, 0.002, 0.01, 0.05, 0.2, 0.4, 1.0, 2.0}; + double bA = 0.4, b1 = 0.4, b2 = 0.05, bestScore = Double.NEGATIVE_INFINITY; + for (double a : gridA) { + for (double a1 : gridA1) { + for (double a2 : gridA2) { + double sc = mean(t -> f.getLogProbabilityOfTree(t, a, a1, a2), val); + if (sc > bestScore) { + bestScore = sc; bA = a; b1 = a1; b2 = a2; + } + } + } + } + long t = System.nanoTime(); + CRegCCD m = new CRegCCD(read(dir, base, "train"), 0.0, bA, b1, b2); + build = ms(t); + params = "alpha=" + fmt(bA) + ";alpha1=" + fmt(b1) + ";alpha2=" + fmt(b2); + scorer = m::getLogProbabilityOfTree; + } + default -> throw new IllegalArgumentException("unknown model " + model); + } + + long t = System.nanoTime(); + int covered = 0; + double sum = 0; + try (PrintWriter w = new PrintWriter(new File(perTree, base + "__" + model + ".txt"))) { + for (Tree x : test) { + double lp = scorer.logP(x); + w.println(lp); + if (Double.isFinite(lp)) { covered++; sum += lp; } + } + } + long score = ms(t); + row.put(model + "_params", params); + row.put(model + "_coverage", fmt(covered / (double) test.size())); + row.put(model + "_meanLogP", covered == test.size() ? fmt(sum / test.size()) : "-inf"); + row.put(model + "_constructMs", String.valueOf(build)); + row.put(model + "_scoreMs", String.valueOf(score)); + } catch (Throwable e) { + row.put(model + "_params", "FAILED:" + e.getClass().getSimpleName()); + row.put(model + "_coverage", ""); + row.put(model + "_meanLogP", ""); + row.put(model + "_constructMs", ""); + row.put(model + "_scoreMs", ""); + } + } + + interface Scorer { double logP(Tree t); } + interface Obj { double at(double x); } + + private static double pick(double[] grid, Obj f) { + double best = grid[0], bestScore = Double.NEGATIVE_INFINITY; + for (double x : grid) { + double s = f.at(x); + if (s > bestScore) { bestScore = s; best = x; } + } + return best; + } + + private static double mean(Scorer s, List trees) { + double sum = 0; int n = 0; + for (Tree t : trees) { + double lp = s.logP(t); + if (Double.isFinite(lp)) { sum += lp; n++; } + } + return n == 0 ? Double.NEGATIVE_INFINITY : sum / n; + } + + private static List read(File dir, String base, String part) { + try { + TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet( + new File(dir, base + "." + part + ".trees").getAbsolutePath(), 0); + ts.reset(); + List out = new ArrayList<>(); + while (ts.hasNext()) out.add(ts.next()); + return out; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static long ms(long t0) { return (System.nanoTime() - t0) / 1_000_000L; } + /** + * Formats a hyperparameter without losing it. A fixed "%.4f" collapsed every grid value below + * 5e-5 to "0.0000" -- the alpha2 bands at 1.6e-5, 8e-5 and 4e-4 and the KRegCCD mu band at + * 2e-5 all became the same unreadable string -- so the selected value could not be recovered + * from the CSV. Double.toString round-trips exactly and Python's float() parses it. + */ + private static String fmt(double d) { return Double.toString(d); } +} diff --git a/src/test/java/ccd/model/MRegDepthTimingTest.java b/src/test/java/ccd/model/MRegDepthTimingTest.java new file mode 100644 index 0000000..1d031f2 --- /dev/null +++ b/src/test/java/ccd/model/MRegDepthTimingTest.java @@ -0,0 +1,58 @@ +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; + +/** How much of MRegCCD's scoring cost is its reserve depth, versus the boundary enumeration itself? */ +public class MRegDepthTimingTest { + + private static final String PATH = System.getProperty("ccd.trees", ""); + + private static List read(int count, int skip) throws Exception { + TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(PATH, 10); + ts.reset(); + List all = new ArrayList<>(); + while (ts.hasNext()) { + all.add(ts.next()); + } + List pool = all.subList(skip * all.size() / 2, (skip + 1) * all.size() / 2); + 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(pool.get((int) (i * step))); + } + return out; + } + + @Test + public void depthVersusCost() throws Exception { + Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists()); + List test = read(200, 1); + System.out.printf("%n=== %s: MRegCCD cost by reserve depth ===%n", new File(PATH).getName()); + System.out.printf("%-7s %-9s %-12s %-12s %-14s%n", "depth", "impl", "construct", "score/200", "mean logP"); + for (int depth : new int[]{2, 3, 4}) { + for (String which : new String[]{"MRegCCDSlow", "MRegCCD"}) { + long t0 = System.nanoTime(); + MRegCCDSlow m = which.equals("MRegCCDSlow") + ? new MRegCCDSlow(read(500, 0), 0.0, MRegCCDSlow.DEFAULT_MU, depth, true) + : new MRegCCD(read(500, 0), 0.0, MRegCCDSlow.DEFAULT_MU, depth, true); + long build = (System.nanoTime() - t0) / 1_000_000L; + t0 = System.nanoTime(); + double sum = 0; + for (Tree t : test) { + sum += m.getLogProbabilityOfTree(t); + } + long score = (System.nanoTime() - t0) / 1_000_000L; + System.out.printf("%-7d %-9s %9dms %9dms %14.6f%n", + depth, which, build, score, sum / test.size()); + } + } + } +} diff --git a/src/test/java/ccd/model/MunroParseCheck.java b/src/test/java/ccd/model/MunroParseCheck.java new file mode 100644 index 0000000..91b918e --- /dev/null +++ b/src/test/java/ccd/model/MunroParseCheck.java @@ -0,0 +1,54 @@ +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.*; + +/** Checks the newick round-trip used by the comparison harness against direct tree objects. */ +public class MunroParseCheck { + private static final String PATH = System.getProperty("ccd.trees", ""); + + @Test + public void roundTripPreservesTopologies() throws Exception { + Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists()); + TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(PATH, 10); + ts.reset(); + List direct = new ArrayList<>(); + while (ts.hasNext() && direct.size() < 400) direct.add(ts.next()); + Tree t0 = direct.get(0); + List taxa = new ArrayList<>(List.of(t0.getTaxaNames())); + while (taxa.remove(null)) { } + System.out.println("taxa[0..3] = " + taxa.subList(0, Math.min(4, taxa.size()))); + String nwk = t0.getRoot().toShortNewick(false); + System.out.println("round-tripped newick head: " + nwk.substring(0, Math.min(150, nwk.length()))); + + // compare clade sets: direct vs reparsed + int mismatches = 0; + for (int i = 0; i < direct.size(); i++) { + String s = direct.get(i).getRoot().toShortNewick(false) + ";"; + Tree re = new beast.base.evolution.tree.TreeParser(taxa, s, 0, false); + Set a = clades(direct.get(i)), b = clades(re); + if (!a.equals(b)) mismatches++; + } + System.out.printf("clade-set mismatches after round trip: %d of %d trees%n", + mismatches, direct.size()); + } + + private static Set clades(Tree t) { + Set out = new TreeSet<>(); + collect(t.getRoot(), out); + return out; + } + + private static SortedSet collect(beast.base.evolution.tree.Node v, Set out) { + SortedSet s = new TreeSet<>(); + if (v.isLeaf()) { s.add(v.getID() != null ? v.getID() : String.valueOf(v.getNr())); } + else for (beast.base.evolution.tree.Node c : v.getChildren()) s.addAll(collect(c, out)); + out.add(String.join(",", s)); + return s; + } +} diff --git a/src/test/java/ccd/model/RealDataHeadToHeadTest.java b/src/test/java/ccd/model/RealDataHeadToHeadTest.java index df52b25..4438d4f 100644 --- a/src/test/java/ccd/model/RealDataHeadToHeadTest.java +++ b/src/test/java/ccd/model/RealDataHeadToHeadTest.java @@ -9,6 +9,7 @@ import java.io.File; import java.util.ArrayList; import java.util.List; +import java.util.Map; /** * Head-to-head held-out predictive comparison of the regularised CCD variants on a real posterior @@ -27,48 +28,135 @@ */ public class RealDataHeadToHeadTest { - private static final String PATH = System.getProperty("ccd.trees", ""); + private static 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; + /** Which models to run, comma separated; default all. Running one model at a time isolates + * cost and stops a model that fails on a dataset from losing the others' results for it. */ + private static final String MODELS = System.getProperty("ccd.models", ""); + /** Directory for per-tree log probabilities, one file per dataset and model. Paired statistics + * are computed from these afterwards, so models need not run together. */ + private static final String PERTREE = System.getProperty("ccd.pertree", ""); + /** + * A second chain to draw the test set from, for analyses that were run more than once. Splitting + * one chain in half conflates model quality with convergence: the halves differ partly because + * the sampler was still moving. An independent run removes that, so where replicate runs exist + * the test set is taken from a sibling run and the whole of the primary chain trains. + */ + private static final String TEST_TREES = System.getProperty("ccd.testtrees", ""); + + /** Directory of prepared subsets; when set, the source chain is never touched. */ + private static final String PREPARED = System.getProperty("ccd.prepared", ""); + + private static List readPrepared(String part) throws Exception { + String base = new File(PATH).getName().replaceAll("\\.trees$", ""); + File f = new File(PREPARED, base + "." + part + ".trees"); + TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(f.getAbsolutePath(), 0); + ts.reset(); + List out = new ArrayList<>(); + while (ts.hasNext()) { + out.add(ts.next()); + } + if (chainTrees == null) { + chainTrees = out.size(); + System.out.printf("prepared %s: %d trees, %d taxa%n", + base, out.size(), out.get(0).getLeafNodeCount()); + } + return out; + } + + private static boolean runs(String model) { + return MODELS.isEmpty() || ("," + MODELS + ",").contains("," + model + ","); + } + + private static void dumpPerTree(String model, Scorer s, List test) throws Exception { + if (PERTREE.isEmpty()) { + return; + } + new File(PERTREE).mkdirs(); + String name = new File(PATH).getName().replaceAll("[^A-Za-z0-9._-]", "_"); + try (java.io.PrintWriter w = new java.io.PrintWriter( + new File(PERTREE, name + "__" + model + ".txt"))) { + for (Tree t : test) { + w.println(s.logP(t)); + } + } + } + private interface Scorer { double logP(Tree t); } + private static Integer chainTrees; + + /** One cheap pass to size the stride before caching anything. */ + private static int countTrees() throws Exception { + TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(PATH, (int) BURNIN_PERCENT); + ts.reset(); + int n = 0; + while (ts.hasNext()) { + ts.next(); + n++; + } + return n; + } + private static List newickCache; private static List taxaCache; - /** Parses the tree file once, keeping the topologies as newick strings. */ + /** + * Parses the tree file once, keeping each topology as a newick string. + * + *

The topologies are written with node numbers rather than taxon labels, and the taxon list + * supplies the mapping. Labels cannot be round-tripped safely: bracketed genus names are standard + * for unvalidated nomenclature, so real files contain labels like + * {@code GluRS-B_AF_Bact_[Eubacterium]_eligens_...}, and a newick parser reads {@code [} as a + * comment opener. Stripping the brackets corrupts the name; keeping them fails to parse. + * MunroParseCheck verifies that the round trip preserves clade sets exactly. + */ + /** Upper bound on cached topologies; the experiment never needs more than a few thousand. */ + private static final int MAX_CACHED = 8000; + private static void load() throws Exception { if (newickCache != null) { return; } + int total = countTrees(); + int stride = Math.max(1, (total + MAX_CACHED - 1) / MAX_CACHED); TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(PATH, (int) BURNIN_PERCENT); ts.reset(); List nwk = new ArrayList<>(); List taxa = null; + // Keep at most MAX_CACHED topologies, evenly spaced through the chain. The experiment draws + // evenly spaced subsets anyway, so this changes nothing statistically, but it bounds memory: + // holding every tree of a 36k-tree chain on several hundred taxa exhausts the heap. + int chain = 0; while (ts.hasNext()) { Tree t = ts.next(); + chain++; 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 = new ArrayList<>(List.of(t.getTaxaNames())); + while (taxa.remove(null)) { + // getTaxaNames() is sized for all nodes on some inputs; drop the internal slots } - taxa.addAll(List.of(byNr)); } - nwk.add(t.getRoot().toNewick() + ";"); + if (stride > 1 && (chain - 1) % stride != 0) { + continue; + } + nwk.add(t.getRoot().toShortNewick(false) + ";"); } newickCache = nwk; taxaCache = taxa; - System.out.printf("loaded %d trees, %d taxa from %s%n", - nwk.size(), taxa.size(), new File(PATH).getName()); + chainTrees = chain; + System.out.printf("loaded %d of %d trees (stride %d), %d taxa from %s%n", + nwk.size(), chain, stride, 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. + * {@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 @@ -76,29 +164,56 @@ private static void load() throws Exception { */ 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 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 beast.base.evolution.tree.TreeParser(taxaCache, pool.get((int) (i * step)), 0, false)); + 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); + return PREPARED.isEmpty() ? read(count, 0.0, 0.5) : readPrepared("train"); } /** 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); + return PREPARED.isEmpty() ? read(count, 0.0, 0.25) : readPrepared("fit"); } private static List valSet(int count) throws Exception { - return read(count, 0.25, 0.5); + return PREPARED.isEmpty() ? read(count, 0.25, 0.5) : readPrepared("val"); + } + + /** Machine-readable output: one row per dataset, appended to -Dccd.csv if set. */ + private static final String CSV = System.getProperty("ccd.csv", ""); + private static final Map CSV_CELLS = new java.util.LinkedHashMap<>(); + + /** Milliseconds since a System.nanoTime() mark. */ + private static long ms(long t0) { + return (System.nanoTime() - t0) / 1_000_000L; + } + + private static void cell(String key, Object value) { + CSV_CELLS.put(key, String.valueOf(value)); + } + + private static void writeCsv() throws Exception { + if (CSV.isEmpty()) { + return; + } + File f = new File(CSV); + boolean header = !f.exists() || f.length() == 0; + try (java.io.PrintWriter w = new java.io.PrintWriter(new java.io.FileWriter(f, true))) { + if (header) { + w.println(String.join(",", CSV_CELLS.keySet())); + } + w.println(String.join(",", CSV_CELLS.values())); + } } private static double[] score(Scorer s, List test) { @@ -127,25 +242,183 @@ private static double meanAll(Scorer s, List test) { return sum / test.size(); } + /** + * Writes the four tree subsets this experiment uses -- fit, validation, train, test -- as small + * NEXUS files, so the source chain is parsed once rather than once per model. + * + *

Trees are written with a Translate block and numeric newick: taxon labels appear only in the + * Translate block, which keeps labels containing brackets (standard for unvalidated genus names) + * out of the newick, where a parser would read them as comments. + */ + @Test + public void prepareSubsets() throws Exception { + String outDir = System.getProperty("ccd.prepare", ""); + Assumptions.assumeTrue(!outDir.isEmpty() && !PATH.isEmpty() && new File(PATH).exists()); + String base = new File(PATH).getName().replaceAll("\\.trees$", ""); + File dir = new File(outDir); + dir.mkdirs(); + load(); + // With a sibling run the whole primary chain trains and the sibling supplies the test set; + // otherwise the chain is split in half, as the manuscript's RSV2 comparison does. + String[][] parts = TEST_TREES.isEmpty() + ? new String[][]{{"fit", "0.0", "0.25"}, {"val", "0.25", "0.5"}, + {"train", "0.0", "0.5"}, {"test", "0.5", "1.0"}} + : new String[][]{{"fit", "0.0", "0.5"}, {"val", "0.5", "1.0"}, + {"train", "0.0", "1.0"}}; + for (String[] part : parts) { + int count = part[0].equals("fit") || part[0].equals("val") ? N / 2 : N; + List sub = read(count, Double.parseDouble(part[1]), Double.parseDouble(part[2])); + File out = writeSubset(dir, base, part[0], sub); + System.out.printf("wrote %s (%d trees)%n", out.getName(), sub.size()); + } + if (!TEST_TREES.isEmpty()) { + writeSubset(dir, base, "test", readFrom(TEST_TREES, N)); + // marker so the run phase can record which protocol produced the test set + try (java.io.PrintWriter w = new java.io.PrintWriter( + new File(dir, base + ".siblingtest"))) { + w.println(new File(TEST_TREES).getName()); + } + System.out.printf("test set taken from sibling run %s%n", new File(TEST_TREES).getName()); + } + } + + /** Evenly spaced trees from the whole post-burn-in chain of another file. */ + private static List readFrom(String path, int count) throws Exception { + TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(path, (int) BURNIN_PERCENT); + ts.reset(); + List all = new ArrayList<>(); + while (ts.hasNext()) { + all.add(ts.next()); + } + List out = new ArrayList<>(); + double step = Math.max(1.0, all.size() / (double) count); + for (int i = 0; i < count && (int) (i * step) < all.size(); i++) { + out.add(all.get((int) (i * step))); + } + return out; + } + + /** Writes one subset as NEXUS with a Translate block and numeric newick. */ + private static File writeSubset(File dir, String base, String part, List sub) + throws Exception { + File out = new File(dir, base + "." + part + ".trees"); + try (java.io.PrintWriter w = new java.io.PrintWriter(out)) { + w.println("#NEXUS"); + w.println("Begin taxa;"); + w.println("\tDimensions ntax=" + taxaCache.size() + ";"); + w.println("\tTaxlabels"); + for (String t : taxaCache) { + w.println("\t\t'" + t + "'"); + } + w.println("\t\t;"); + w.println("End;"); + w.println("Begin trees;"); + w.println("\tTranslate"); + for (int i = 0; i < taxaCache.size(); i++) { + w.println("\t\t" + (i + 1) + " '" + taxaCache.get(i) + "'" + + (i + 1 < taxaCache.size() ? "," : "")); + } + w.println("\t\t;"); + for (int i = 0; i < sub.size(); i++) { + w.println("tree STATE_" + i + " = " + numericNewick(sub.get(i).getRoot()) + ";"); + } + w.println("End;"); + } + return out; + } + + /** Newick using 1-based taxon numbers, matching the Translate block written above. */ + private static String numericNewick(beast.base.evolution.tree.Node v) { + if (v.isLeaf()) { + return String.valueOf(v.getNr() + 1); + } + StringBuilder sb = new StringBuilder("("); + for (int i = 0; i < v.getChildCount(); i++) { + sb.append(i > 0 ? "," : "").append(numericNewick(v.getChild(i))); + } + return sb.append(")").toString(); + } + + /** + * Runs every prepared dataset in one JVM. A separate Maven invocation per dataset spends far + * more time starting up than modelling: sweeping CRegCCD over 90 datasets took 11 minutes of + * wall clock for 2.4 seconds of actual work. + */ + @Test + public void batchOverPreparedDatasets() throws Exception { + String batch = System.getProperty("ccd.batch", ""); + Assumptions.assumeTrue(!batch.isEmpty()); + File dir = new File(batch); + List bases = new ArrayList<>(); + for (File f : dir.listFiles((d, n) -> n.endsWith(".train.trees"))) { + bases.add(f.getName().replaceAll("\\.train\\.trees$", "")); + } + java.util.Collections.sort(bases); + System.out.printf("batch: %d prepared datasets%n", bases.size()); + int ok = 0; + for (String base : bases) { + PATH = new File(dir, base + ".trees").getAbsolutePath(); + chainTrees = null; + newickCache = null; + taxaCache = null; + CSV_CELLS.clear(); + try { + headToHeadOnRealData(); + ok++; + } catch (Throwable e) { + System.out.printf("FAILED %s (%s: %s)%n", base, + e.getClass().getSimpleName(), String.valueOf(e.getMessage())); + try (java.io.PrintWriter w = new java.io.PrintWriter( + new java.io.FileWriter(CSV + ".failed", true))) { + w.println(base + "," + e.getClass().getSimpleName()); + } + } + } + System.out.printf("batch complete: %d/%d%n", ok, bases.size()); + } + @Test public void headToHeadOnRealData() throws Exception { - Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(), + // in prepared mode the path names the dataset and locates its subsets; the original + // chain file need not be present + Assumptions.assumeTrue(!PATH.isEmpty() + && (new File(PATH).exists() || !PREPARED.isEmpty()), "set -Dccd.trees=/path/to/x.trees to run this comparison"); List probe = train(N); int nTaxa = probe.get(0).getLeafNodeCount(); + cell("dataset", new File(PATH).getName().replace(",", ";")); + // which protocol produced the test set: a sibling MCMC run, or the second half of this chain + cell("testProtocol", new File(PREPARED, new File(PATH).getName() + .replaceAll("\\.trees$", "") + ".siblingtest").exists() + ? "sibling-run" : "chain-half"); + cell("taxa", nTaxa); + cell("chainTrees", chainTrees); + cell("nTrain", probe.size()); + cell("nTest", N); + // CCD1 entropy of the training sample, recorded for every dataset regardless of which + // models run: it is the manuscript's measure of how much topological uncertainty a + // posterior actually holds, and datasets below a few nats cannot discriminate the models. + CCD1 entropyProbe = new CCD1(train(N), 0.0); + cell("CCD1_entropy", String.format("%.3f", entropyProbe.getEntropy())); 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 test = PREPARED.isEmpty() ? read(N, 0.5, 1.0) : readPrepared("test"); List val = valSet(N / 2); int nFit = N / 2; // ---- CCD1 ---- - CCD1 ccd1 = new CCD1(train(N), 0.0); + if (runs("CCD1")) { + List ccd1Trees = train(N); + long tCcd1 = System.nanoTime(); + CCD1 ccd1 = new CCD1(ccd1Trees, 0.0); + cell("CCD1_constructMs", ms(tCcd1)); report("CCD1", "-", ccd1::getLogProbabilityOfTree, test); + } // ---- RegCCD: alpha on validation ---- + if (runs("RegCCD")) { 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}) { @@ -156,10 +429,16 @@ public void headToHeadOnRealData() throws Exception { bestAlpha = alpha; } } - RegCCD reg = new RegCCD(train(N), 0.0, bestAlpha); + List regTrees = train(N); + long tReg = System.nanoTime(); + RegCCD reg = new RegCCD(regTrees, 0.0, bestAlpha); + cell("RegCCD_constructMs", ms(tReg)); report("RegCCD", String.format("alpha=%.2f", bestAlpha), reg::getLogProbabilityOfTree, test); + } // ---- KRegCCD: mu on validation at alpha = 0.4 (the manuscript's setting) ---- + KRegCCD kreg = null; + if (runs("KRegCCD")) { KRegCCD kFit = new KRegCCD(fitSet(nFit), 0.0, 0.005, 0.4); double bestMu = 0.005; double bestMuScore = Double.NEGATIVE_INFINITY; @@ -171,26 +450,46 @@ public void headToHeadOnRealData() throws Exception { bestMu = mu; } } - KRegCCD kreg = new KRegCCD(train(N), 0.0, bestMu, 0.4); + List kregTrees = train(N); + long tKreg = System.nanoTime(); + kreg = new KRegCCD(kregTrees, 0.0, bestMu, 0.4); + kreg.precomputeReserves(); // KRegCCD defers its reserve solve; charge it to construction + cell("KRegCCD_constructMs", ms(tKreg)); 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: mu on validation (isolated: a baseline crash must not lose the row) ---- + MRegCCD mreg = null; + try { + if (!runs("MRegCCD")) { + throw new IllegalStateException("skipped by ccd.models"); } + 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; + } + } + List mregTrees = train(N); + long tMreg = System.nanoTime(); + mreg = new MRegCCD(mregTrees, 0.0, bestMMu); + cell("MRegCCD_constructMs", ms(tMreg)); + report("MRegCCD", String.format("mu=%.5f", bestMMu), mreg::getLogProbabilityOfTree, test); + } catch (Throwable e) { + mreg = null; + if (runs("MRegCCD")) System.out.printf("MRegCCD SKIPPED (%s: %s)%n", + e.getClass().getSimpleName(), String.valueOf(e.getMessage())); } - 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: (alpha, alpha1, alpha2) on validation ---- + CRegCCD creg = null; + if (runs("CRegCCD")) { 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}; @@ -206,25 +505,48 @@ public void headToHeadOnRealData() throws Exception { } } } - CRegCCD creg = new CRegCCD(train(N), 0.0, bestC[1], bestC[2], bestC[3]); + List cregTrees = train(N); + long tCreg = System.nanoTime(); + creg = new CRegCCD(cregTrees, 0.0, bestC[1], bestC[2], bestC[3]); + cell("CRegCCD_constructMs", ms(tCreg)); 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 ---- + if (kreg == null) { + writeCsv(); + return; // paired statistics are computed from the per-tree dumps instead + } 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); + if (mreg != null) { + paired("MRegCCD", mreg::getLogProbabilityOfTree, kreg::getLogProbabilityOfTree, test); + } + if (creg != null) { + paired("CRegCCD", creg::getLogProbabilityOfTree, kreg::getLogProbabilityOfTree, test); + } + + writeCsv(); } - private static void report(String name, String params, Scorer s, List test) { + private static void report(String name, String params, Scorer s, List test) + throws Exception { + dumpPerTree(name, s, test); + long t0 = System.nanoTime(); double[] sc = score(s, test); + long scoreMs = ms(t0); // one pass over the test set double all = meanAll(s, test); + cell(name + "_scoreMs", scoreMs); + cell(name + "_params", params.replace(",", ";")); + cell(name + "_coverage", String.format("%.4f", sc[0] / test.size())); + cell(name + "_meanLogP", Double.isFinite(all) ? String.format("%.4f", all) : "-inf"); 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) { + // records name_vs_KRegCCD_{diff,se,wins,n} in addition to printing int wins = 0; double sumDiff = 0.0; double sumSq = 0.0; @@ -239,6 +561,10 @@ private static void paired(String name, Scorer a, Scorer baseline, List te int n = test.size(); double mean = sumDiff / n; double se = Math.sqrt(Math.max(0, sumSq / n - mean * mean) / n); + cell(name + "_vs_KRegCCD_diff", String.format("%.4f", mean)); + cell(name + "_vs_KRegCCD_se", String.format("%.4f", se)); + cell(name + "_vs_KRegCCD_wins", wins); + cell(name + "_vs_KRegCCD_n", n); System.out.printf(" %-9s mean log-ratio %+8.2f +/- %.2f nats/tree, better on %d/%d trees%n", name, mean, se, wins, n); } @@ -247,7 +573,10 @@ private static void paired(String name, Scorer a, Scorer baseline, List te * self-consistent with the scorer. */ @Test public void samplingOnRealData() throws Exception { - Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(), + // in prepared mode the path names the dataset and locates its subsets; the original + // chain file need not be present + Assumptions.assumeTrue(!PATH.isEmpty() + && (new File(PATH).exists() || !PREPARED.isEmpty()), "set -Dccd.trees=/path/to/x.trees to run this comparison"); CRegCCD ccd = new CRegCCD(train(N), 0.0); int nTaxa = ccd.getSizeOfLeavesArray(); @@ -273,7 +602,10 @@ public void samplingOnRealData() throws Exception { /** 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(), + // in prepared mode the path names the dataset and locates its subsets; the original + // chain file need not be present + Assumptions.assumeTrue(!PATH.isEmpty() + && (new File(PATH).exists() || !PREPARED.isEmpty()), "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); @@ -305,7 +637,10 @@ public void mapAndEntropyOnRealData() throws Exception { /** 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(), + // in prepared mode the path names the dataset and locates its subsets; the original + // chain file need not be present + Assumptions.assumeTrue(!PATH.isEmpty() + && (new File(PATH).exists() || !PREPARED.isEmpty()), "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()); @@ -328,7 +663,10 @@ public void entropyRecursionVersusMonteCarloOnRealData() throws Exception { /** 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(), + // in prepared mode the path names the dataset and locates its subsets; the original + // chain file need not be present + Assumptions.assumeTrue(!PATH.isEmpty() + && (new File(PATH).exists() || !PREPARED.isEmpty()), "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(); From 6638ce02fd0b6599c4df01e9c5230d829b23389b Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Thu, 20 Aug 2026 12:07:32 +1200 Subject: [PATCH 5/8] Add SplitClassCost: what each model charges for CRegCCD's split classes CRegCCD's four split classes -- observed, expanded, one novel clade, two novel clades -- are a property of the split and the training set rather than of any model, so all three full-support models can be asked what they charge for the same split. This tool classifies every split of every held-out tree directly from the training trees (no model internals) and regresses the per-tree log probability on the class counts crossed with the size of the clade being split. The regression has no intercept: the four class counts sum to the number of internal nodes, which is fixed for a given taxon set, so a design matrix with an intercept is singular and the split between intercept and coefficients is arbitrary. Without it each coefficient is directly the mean log probability contributed by one split of that class. The fitted mean is printed alongside the actual mean, since OLS forces them equal and the check catches a solver that is quietly wrong. On pela-zego this reproduces each model's benchmark mean exactly and explains R^2 = 0.996-0.998 of the per-tree variance, which is what licenses reading the coefficients as per-split costs even though KRegCCD and MRegCCD price maximal novel regions rather than individual splits. The result it was written to test: CRegCCD spreads the class total alpha2 uniformly over |A_2(C)| ~ 2^(m-1) splits, so its two-novel penalty falls linearly in the clade size m with slope -log 2, measured at -0.659 nats per taxon against the predicted -0.693. KRegCCD and MRegCCD raise an escape rate eps ~ mu/O(s^2) to a fixed power, so theirs grows only like log m. At m = 193 that is -129 nats per split against -14 and -6. Co-Authored-By: Claude Opus 5 (1M context) --- src/main/java/ccd/tools/SplitClassCost.java | 253 ++++++++++++++++++++ 1 file changed, 253 insertions(+) create mode 100644 src/main/java/ccd/tools/SplitClassCost.java diff --git a/src/main/java/ccd/tools/SplitClassCost.java b/src/main/java/ccd/tools/SplitClassCost.java new file mode 100644 index 0000000..ab19073 --- /dev/null +++ b/src/main/java/ccd/tools/SplitClassCost.java @@ -0,0 +1,253 @@ +package ccd.tools; + +import beast.base.evolution.tree.Node; +import beast.base.evolution.tree.Tree; +import beastfx.app.treeannotator.TreeAnnotator; +import ccd.model.*; + +import java.io.File; +import java.util.*; + +/** + * Measures what CRegCCD, KRegCCD and MRegCCD charge for each of CRegCCD's four split classes. + * + * 

+ *   java -cp ... ccd.tools.SplitClassCost <prepared-dir> <base> alpha alpha1 alpha2 muK muM
+ *
+ * + * + * The classes are a property of the split and the training set, not of any model, so they are + * computed here directly from the training trees rather than from any model's internals: + * 1 observed -- the split of C was seen in training + * 2 expanded -- C and both parts are observed clades, but never as this split of C + * 3 one novel -- exactly one part is an observed clade + * 4 two novel -- neither part is an observed clade + * + * Each held-out tree is then summarised by how many splits of each class it contains, and the + * per-tree log probability under each model is regressed on those counts. KRegCCD and MRegCCD + * price maximal novel *regions* rather than individual splits, so the fit is not exact by + * construction; the residual scatter is reported so that the extent of the approximation is + * visible rather than assumed away. + */ +public class SplitClassCost { + + static List read(File dir, String base, String part) throws Exception { + File f = new File(dir, base + "." + part + ".trees"); + TreeAnnotator.MemoryFriendlyTreeSet s = + new TreeAnnotator().new MemoryFriendlyTreeSet(f.getAbsolutePath(), 0); + s.reset(); + List out = new ArrayList<>(); + Tree t; + while ((t = s.next()) != null) out.add(t); + return out; + } + + /** Leaf-index bitset of the clade below a node, as a sorted string key. */ + static String clade(Node n, Map memo, int nTaxa) { + return bits(n, memo, nTaxa).toString(); + } + + static java.util.BitSet bits(Node n, Map memo, int nTaxa) { + java.util.BitSet b = memo.get(n); + if (b != null) return b; + b = new java.util.BitSet(nTaxa); + if (n.isLeaf()) { + b.set(n.getNr()); + } else { + for (Node c : n.getChildren()) b.or(bits(c, memo, nTaxa)); + } + memo.put(n, b); + return b; + } + + static void collect(Tree t, Set clades, Set splits, int nTaxa) { + Map memo = new HashMap<>(); + for (Node n : t.getNodesAsArray()) { + if (n.isLeaf()) continue; + clades.add(clade(n, memo, nTaxa)); + List ch = n.getChildren(); + if (ch.size() == 2) { + String l = clade(ch.get(0), memo, nTaxa), r = clade(ch.get(1), memo, nTaxa); + splits.add(clade(n, memo, nTaxa) + "|" + (l.compareTo(r) < 0 ? l + "," + r : r + "," + l)); + } + } + } + + /** Size bins for the parent clade: the analytic claim is about how the penalty scales with + * the number of taxa under the clade being split, so the bins are geometric. */ + static final int[] BIN_HI = {4, 8, 16, 32, 64, Integer.MAX_VALUE}; + static final String[] BIN_NAME = {"2-4", "5-8", "9-16", "17-32", "33-64", "65+"}; + /** mean parent-clade size actually seen in each (class, bin) cell, so the fitted penalty can be + * checked against an analytic prediction in m rather than against a bin label. */ + static double[] sumM = new double[4 * 6]; + + static int bin(int m) { + for (int i = 0; i < BIN_HI.length; i++) if (m <= BIN_HI[i]) return i; + return BIN_HI.length - 1; + } + + /** counts of the four classes crossed with parent-clade size bin, for one held-out tree. */ + static int[] classify(Tree t, Set clades, Set splits, int nTaxa) { + Map memo = new HashMap<>(); + int[] c = new int[4 * BIN_HI.length]; + for (Node n : t.getNodesAsArray()) { + if (n.isLeaf() || n.getChildren().size() != 2) continue; + String C = clade(n, memo, nTaxa); + String l = clade(n.getChildren().get(0), memo, nTaxa); + String r = clade(n.getChildren().get(1), memo, nTaxa); + boolean lo = clades.contains(l) || n.getChildren().get(0).isLeaf(); + boolean ro = clades.contains(r) || n.getChildren().get(1).isLeaf(); + String key = C + "|" + (l.compareTo(r) < 0 ? l + "," + r : r + "," + l); + int cls = splits.contains(key) ? 0 : (lo && ro) ? 1 : (lo || ro) ? 2 : 3; + int m = bits(n, memo, nTaxa).cardinality(); + int idx = cls * BIN_HI.length + bin(m); + c[idx]++; + sumM[idx] += m; + } + return c; + } + + /** + * Least squares WITHOUT an intercept; returns {b1..bk, R2}. + * + * The four class counts sum to the number of internal nodes, which is the same for every tree + * on a fixed taxon set, so a design matrix with an intercept is singular and the split between + * intercept and coefficients is arbitrary. Without the intercept each coefficient is directly + * the mean log probability contributed by one split of that class, which is the quantity of + * interest. R2 is taken about zero accordingly. + */ + static double[] ols(double[][] X, double[] y) { + int n = y.length, p = X[0].length; + double[][] A = new double[p][p]; + double[] b = new double[p]; + for (int i = 0; i < n; i++) { + double[] xi = X[i]; + for (int j = 0; j < p; j++) { + b[j] += xi[j] * y[i]; + for (int k = 0; k < p; k++) A[j][k] += xi[j] * xi[k]; + } + } + for (int i = 0; i < p; i++) A[i][i] += 1e-9; + double[] beta = solve(A, b); + double ssTot = 0, ssRes = 0; + for (int i = 0; i < n; i++) { + double pred = 0; + for (int j = 0; j < p; j++) pred += beta[j] * X[i][j]; + ssRes += (y[i] - pred) * (y[i] - pred); + ssTot += y[i] * y[i]; + } + double[] out = new double[p + 1]; + System.arraycopy(beta, 0, out, 0, p); + out[p] = 1 - ssRes / ssTot; + return out; + } + + static double[] solve(double[][] A, double[] b) { + int n = b.length; + double[][] M = new double[n][n + 1]; + for (int i = 0; i < n; i++) { + System.arraycopy(A[i], 0, M[i], 0, n); + M[i][n] = b[i]; + } + for (int col = 0; col < n; col++) { + int piv = col; + for (int r = col + 1; r < n; r++) if (Math.abs(M[r][col]) > Math.abs(M[piv][col])) piv = r; + double[] tmp = M[col]; M[col] = M[piv]; M[piv] = tmp; + for (int r = 0; r < n; r++) { + if (r == col || M[col][col] == 0) continue; + double f = M[r][col] / M[col][col]; + for (int c2 = col; c2 <= n; c2++) M[r][c2] -= f * M[col][c2]; + } + } + double[] x = new double[n]; + for (int i = 0; i < n; i++) x[i] = M[i][n] / M[i][i]; + return x; + } + + public static void main(String[] args) throws Exception { + File dir = new File(args[0]); + String base = args[1]; + double alpha = Double.parseDouble(args[2]), a1 = Double.parseDouble(args[3]); + double a2 = Double.parseDouble(args[4]); + double muK = Double.parseDouble(args[5]), muM = Double.parseDouble(args[6]); + + List test = read(dir, base, "test"); + int nTaxa = test.get(0).getLeafNodeCount(); + Set clades = new HashSet<>(), splits = new HashSet<>(); + for (Tree t : read(dir, base, "train")) collect(t, clades, splits, nTaxa); + + int n = test.size(); + int NB = BIN_HI.length, P = 4 * NB; + double[][] X = new double[n][P]; + int[] tot = new int[P]; + for (int i = 0; i < n; i++) { + int[] c = classify(test.get(i), clades, splits, nTaxa); + for (int j = 0; j < P; j++) { X[i][j] = c[j]; tot[j] += c[j]; } + } + System.out.printf("%s: %d taxa, %d held-out trees, %d internal splits each%n", + base, nTaxa, n, nTaxa - 1); + String[] CLS = {"observed", "expanded", "one-novel", "two-novel"}; + int tt = n * (nTaxa - 1); + System.out.print("split classes:"); + for (int k = 0; k < 4; k++) { + int sum = 0; + for (int bnd = 0; bnd < NB; bnd++) sum += tot[k * NB + bnd]; + System.out.printf(" %s %d (%.2f%%)", CLS[k], sum, 100.0 * sum / tt); + } + System.out.println(); + + Map ys = new LinkedHashMap<>(); + CRegCCD cre = new CRegCCD(read(dir, base, "train"), 0.0, alpha, a1, a2); + KRegCCD kre = new KRegCCD(read(dir, base, "train"), 0.0, muK, 0.4); + MRegCCD mre = new MRegCCD(read(dir, base, "train"), 0.0, muM); + double[] yc = new double[n], yk = new double[n], ym = new double[n]; + for (int i = 0; i < n; i++) { + yc[i] = cre.getLogProbabilityOfTree(test.get(i)); + yk[i] = kre.getLogProbabilityOfTree(test.get(i)); + ym[i] = mre.getLogProbabilityOfTree(test.get(i)); + } + ys.put("CRegCCD", yc); ys.put("KRegCCD", yk); ys.put("MRegCCD", ym); + + System.out.printf("%nmean log probability contributed per split, by class and parent clade size%n"); + for (Map.Entry e : ys.entrySet()) { + double[] r = ols(X, e.getValue()); + // OLS makes residuals orthogonal to every column, and the columns sum to a constant, + // so the fitted mean must equal the actual mean. Printing both is the check that the + // decomposition is real rather than a plausible-looking artefact of the solver. + double actual = 0, fitted = 0; + for (int i = 0; i < n; i++) { + actual += e.getValue()[i] / n; + for (int j = 0; j < P; j++) fitted += r[j] * X[i][j] / n; + } + System.out.printf("%n%s (mean logP %.2f, fitted %.2f, R2 %.4f)%n", + e.getKey(), actual, fitted, r[P]); + System.out.printf(" %-10s", "class\\size"); + for (String bn : BIN_NAME) System.out.printf("%9s", bn); + System.out.println(); + for (int k = 0; k < 4; k++) { + System.out.printf(" %-10s", CLS[k]); + for (int bnd = 0; bnd < NB; bnd++) { + if (tot[k * NB + bnd] < 20) System.out.printf("%9s", "."); + else System.out.printf("%9.2f", r[k * NB + bnd]); + } + System.out.println(); + } + if (e.getKey().equals("CRegCCD")) { + System.out.printf(" %-10s", "mean m"); + for (int bnd = 0; bnd < NB; bnd++) { + int cnt = tot[3 * NB + bnd]; + if (cnt < 20) System.out.printf("%9s", "."); + else System.out.printf("%9.1f", sumM[3 * NB + bnd] / cnt); + } + System.out.println(" <- for the two-novel row"); + } + } + System.out.println(); + System.out.println("Coefficients are nats per split of that class; more negative is a heavier"); + System.out.println("penalty. CRegCCD spreads the class total alpha2 uniformly over |A_4| ~ 2^(m-1)"); + System.out.println("splits, so its two-novel penalty should fall linearly in m with slope -log 2 ="); + System.out.println("-0.693 nats per taxon; KRegCCD and MRegCCD price the same split as a power of"); + System.out.println("an escape rate eps ~ mu / O(s^2), which is logarithmic in m. Compare the"); + System.out.println("two-novel row against the mean-m row to check that."); + } +} From 816ea050d228b7e1b2a60f1554bb36d3534a361f Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Thu, 20 Aug 2026 12:22:46 +1200 Subject: [PATCH 6/8] Count the factors each model scores, rather than regressing on them CRegCCD contributes one conditional split factor per internal node, including at novel clades. KRegCCD and MRegCCD cluster novel nodes into maximal regions and score each region once at its top, so nodes interior to a region never get a conditional distribution at all. On pela-zego that is 377.0 factors per held-out tree against 369.4. An attempt to measure this by regressing the score on class counts crossed with whether the clade was observed produced a coefficient of +43 nats for a single split, which is impossible: the interior-of-region counts are largely determined by the novel splits above them, so those columns are collinear and the coefficients are not identifiable. Counting the nodes is exact and needs no fit, so that is what this reports. Co-Authored-By: Claude Opus 5 (1M context) --- src/main/java/ccd/tools/SplitClassCost.java | 45 +++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/src/main/java/ccd/tools/SplitClassCost.java b/src/main/java/ccd/tools/SplitClassCost.java index ab19073..77d9adb 100644 --- a/src/main/java/ccd/tools/SplitClassCost.java +++ b/src/main/java/ccd/tools/SplitClassCost.java @@ -107,6 +107,15 @@ static int[] classify(Tree t, Set clades, Set splits, int nTaxa) return c; } + /** A node is blue for KRegCCD/MRegCCD purposes when its split introduces a novel clade. */ + static boolean isBlue(Node n, Set clades, Map memo, int nTaxa) { + if (n.isLeaf() || n.getChildren().size() != 2) return false; + for (Node ch : n.getChildren()) { + if (!ch.isLeaf() && !clades.contains(clade(ch, memo, nTaxa))) return true; + } + return false; + } + /** * Least squares WITHOUT an intercept; returns {b1..bk, R2}. * @@ -242,6 +251,42 @@ public static void main(String[] args) throws Exception { System.out.println(" <- for the two-novel row"); } } + // How many terms does each model actually score? CRegCCD is a chain-rule CCD over every + // clade in the tree, so it contributes one factor per internal node -- including the novel + // clades themselves. KRegCCD and MRegCCD cluster novel nodes into maximal regions and + // score each region once at its top, so the nodes interior to a region never receive a + // conditional split distribution at all. Counting the nodes is exact, unlike regressing + // the score on class counts: the interior count is largely determined by the novel splits + // above it, so those columns are collinear and their coefficients are not identifiable. + long internal = 0, novelClade = 0, regionTop = 0, blue = 0; + for (Tree t : test) { + Map memo = new HashMap<>(); + for (Node nd : t.getNodesAsArray()) { + if (nd.isLeaf() || nd.getChildren().size() != 2) continue; + internal++; + boolean isB = isBlue(nd, clades, memo, nTaxa); + if (isB) blue++; + if (!clades.contains(clade(nd, memo, nTaxa)) && !nd.isRoot()) novelClade++; + if (isB) { + Node par = nd.getParent(); + if (par == null || !isBlue(par, clades, memo, nTaxa)) regionTop++; + } + } + } + double perTree = 1.0 / test.size(); + System.out.printf("%nterms scored per held-out tree (%d internal nodes each)%n", nTaxa - 1); + System.out.printf(" internal nodes %8.2f%n", internal * perTree); + System.out.printf(" novel clades (no observed counterpart) %5.2f%n", novelClade * perTree); + System.out.printf(" blue nodes (split introduces a novel clade) %2.2f%n", blue * perTree); + System.out.printf(" maximal novel regions %8.2f%n", regionTop * perTree); + System.out.printf(" CRegCCD factors %8.2f (one per internal node)%n", + internal * perTree); + System.out.printf(" KRegCCD / MRegCCD factors %8.2f (red nodes + one per region)%n", + (internal - blue + regionTop) * perTree); + System.out.printf(" difference %8.2f factors per tree that CRegCCD%n", + (blue - regionTop) * perTree); + System.out.printf("%36s pays and the others do not%n", ""); + System.out.println(); System.out.println("Coefficients are nats per split of that class; more negative is a heavier"); System.out.println("penalty. CRegCCD spreads the class total alpha2 uniformly over |A_4| ~ 2^(m-1)"); From d8cde53bfcd7ed2d7fa4bfbe9765789346a5c6e9 Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Thu, 20 Aug 2026 13:27:36 +1200 Subject: [PATCH 7/8] Estimate the two-novel scaling directly rather than through bins The slope of the two-novel penalty against clade size was read off three binned means, which needed an arbitrary "too few splits to report" threshold and then fitted a line through the bins as though each carried equal weight. On pela-zego those bins hold 1, 4, 23, 859, 6 and 438 splits, and the top bin spans clade sizes 74 to 252, so a single coefficient for it averages a predicted penalty ranging from -51 to -174 nats. The size now enters as a covariate instead: per tree, the counts of the first three classes, the count of two-novel splits, and the sum over those splits of (m-1) in one fit and log m in the other. The coefficient on that column is the slope, estimated from every two-novel split at its proper weight, and the two fits can be compared on residuals to test which functional form holds. This sharpens the result rather than softening it. CRegCCD's slope is -0.7002 per taxon against the -log 2 = -0.6931 predicted by dividing a class total across |A_2(C)| ~ 2^(m-1), within 1% where the binned estimate was 5% off, and the linear form beats the logarithmic one on residuals 146316 to 250084. KRegCCD and MRegCCD are flat in m at -0.023 and -0.010 per taxon, and for them the two forms are indistinguishable, as expected when the dependence enters only through log m. Cell counts and clade-size ranges are printed alongside the binned table that remains, so a thin cell is visible rather than merely suppressed. Co-Authored-By: Claude Opus 5 (1M context) --- src/main/java/ccd/tools/SplitClassCost.java | 73 +++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/src/main/java/ccd/tools/SplitClassCost.java b/src/main/java/ccd/tools/SplitClassCost.java index 77d9adb..fc51a19 100644 --- a/src/main/java/ccd/tools/SplitClassCost.java +++ b/src/main/java/ccd/tools/SplitClassCost.java @@ -80,6 +80,9 @@ static void collect(Tree t, Set clades, Set splits, int nTaxa) { /** mean parent-clade size actually seen in each (class, bin) cell, so the fitted penalty can be * checked against an analytic prediction in m rather than against a bin label. */ static double[] sumM = new double[4 * 6]; + /** min and max parent-clade size per cell, so an open-ended top bin cannot hide its spread. */ + static int[] minM = new int[4 * 6]; + static int[] maxM = new int[4 * 6]; static int bin(int m) { for (int i = 0; i < BIN_HI.length; i++) if (m <= BIN_HI[i]) return i; @@ -103,10 +106,23 @@ static int[] classify(Tree t, Set clades, Set splits, int nTaxa) int idx = cls * BIN_HI.length + bin(m); c[idx]++; sumM[idx] += m; + if (minM[idx] == 0 || m < minM[idx]) minM[idx] = m; + if (m > maxM[idx]) maxM[idx] = m; } return c; } + /** Residual sum of squares of an intercept-free fit. */ + static double rss(double[][] X, double[] y, double[] beta) { + double t = 0; + for (int i = 0; i < y.length; i++) { + double pred = 0; + for (int j = 0; j < X[i].length; j++) pred += beta[j] * X[i][j]; + t += (y[i] - pred) * (y[i] - pred); + } + return t; + } + /** A node is blue for KRegCCD/MRegCCD purposes when its split introduces a novel clade. */ static boolean isBlue(Node n, Set clades, Map memo, int nTaxa) { if (n.isLeaf() || n.getChildren().size() != 2) return false; @@ -249,8 +265,65 @@ public static void main(String[] args) throws Exception { else System.out.printf("%9.1f", sumM[3 * NB + bnd] / cnt); } System.out.println(" <- for the two-novel row"); + System.out.printf(" %-10s", "range m"); + for (int bnd = 0; bnd < NB; bnd++) { + int cnt = tot[3 * NB + bnd]; + if (cnt < 20) System.out.printf("%9s", "."); + else System.out.printf("%9s", minM[3 * NB + bnd] + "-" + maxM[3 * NB + bnd]); + } + System.out.println(); + System.out.printf(" %-10s", "n splits"); + for (int bnd = 0; bnd < NB; bnd++) System.out.printf("%9d", tot[3 * NB + bnd]); + System.out.println(); } } + // Direct estimate of how the two-novel penalty scales with clade size, without binning. + // Binning forced an arbitrary "too few to report" threshold and then fitted a line through + // bin means as though each bin carried equal weight, when in this dataset the bins hold + // 23, 859 and 438 splits over clade sizes spanning 9 to 252. Here the size enters as a + // covariate instead: the design has, per tree, the counts of the first three classes, the + // count of two-novel splits, and the sum over those splits of (m-1) in one fit and of + // log m in the other. The coefficient on that last column is then the slope directly, + // estimated from every two-novel split with its proper weight. + // + // CRegCCD divides a class total by |A_2(C)| ~ 2^(m-1), so its slope against (m-1) should + // be -log 2 = -0.693 nats per taxon. KRegCCD and MRegCCD raise an escape rate to a power + // fixed by the region, with eps ~ mu/O(s^2), so theirs should be flat in m and instead + // linear in log m. + double[][] Xl = new double[n][5], Xg = new double[n][5]; + for (int i = 0; i < n; i++) { + Map memo = new HashMap<>(); + for (Node nd : test.get(i).getNodesAsArray()) { + if (nd.isLeaf() || nd.getChildren().size() != 2) continue; + String C = clade(nd, memo, nTaxa); + String l = clade(nd.getChildren().get(0), memo, nTaxa); + String r = clade(nd.getChildren().get(1), memo, nTaxa); + boolean lo = clades.contains(l) || nd.getChildren().get(0).isLeaf(); + boolean ro = clades.contains(r) || nd.getChildren().get(1).isLeaf(); + String key = C + "|" + (l.compareTo(r) < 0 ? l + "," + r : r + "," + l); + int cls = splits.contains(key) ? 0 : (lo && ro) ? 1 : (lo || ro) ? 2 : 3; + int m = bits(nd, memo, nTaxa).cardinality(); + if (cls < 3) { Xl[i][cls]++; Xg[i][cls]++; } + else { + Xl[i][3]++; Xg[i][3]++; + Xl[i][4] += m - 1; + Xg[i][4] += Math.log(m); + } + } + } + System.out.printf("%ntwo-novel penalty scaling, estimated from all %d two-novel splits%n", + (int) java.util.Arrays.stream(Xl).mapToDouble(v -> v[3]).sum()); + System.out.printf(" %-9s %14s %10s %14s %10s%n", "model", + "slope /(m-1)", "RSS", "slope /log m", "RSS"); + for (Map.Entry e : ys.entrySet()) { + double[] rl = ols(Xl, e.getValue()), rg = ols(Xg, e.getValue()); + System.out.printf(" %-9s %14.4f %10.0f %14.4f %10.0f%n", e.getKey(), + rl[4], rss(Xl, e.getValue(), rl), rg[4], rss(Xg, e.getValue(), rg)); + } + System.out.printf(" %-9s %14.4f %10s %14s %10s%n", "predicted", -Math.log(2), "", + "(flat)", ""); + System.out.println(" lower RSS is the better-supported functional form"); + // How many terms does each model actually score? CRegCCD is a chain-rule CCD over every // clade in the tree, so it contributes one factor per internal node -- including the novel // clades themselves. KRegCCD and MRegCCD cluster novel nodes into maximal regions and From b375839c56b5251a6f82f2ad00209f3a329027e4 Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Thu, 20 Aug 2026 17:03:42 +1200 Subject: [PATCH 8/8] Repair .gitignore: appending joined two patterns into one An earlier commit added hs_err_pid*.log by appending to a .gitignore whose last line had no trailing newline, producing ".projecths_err_pid*.log". That matched nothing, so both patterns were silently disabled: .project stopped being ignored and the crash logs it was meant to add never were. Also drops a __pycache__ entry that was appended here by a command run from the wrong working directory; it belongs to the benchmark scripts in the CCD_regularisation repository, not to this one. Co-Authored-By: Claude Opus 5 (1M context) --- .gitignore | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index fac52ce..e783d26 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,7 @@ out/ .classpath .DS_Store /bin/ -.projecths_err_pid*.log +.project + +# JVM crash logs from long benchmark runs +hs_err_pid*.log