diff --git a/.gitignore b/.gitignore index 255db57..e783d26 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,7 @@ out/ .classpath .DS_Store /bin/ -.project \ No newline at end of file +.project + +# JVM crash logs from long benchmark runs +hs_err_pid*.log 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/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/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/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/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/main/java/ccd/tools/SplitClassCost.java b/src/main/java/ccd/tools/SplitClassCost.java new file mode 100644 index 0000000..fc51a19 --- /dev/null +++ b/src/main/java/ccd/tools/SplitClassCost.java @@ -0,0 +1,371 @@ +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]; + /** 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; + 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; + 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; + 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}. + * + * 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.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 + // 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)"); + 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."); + } +} 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; + } +} 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/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/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 new file mode 100644 index 0000000..4438d4f --- /dev/null +++ b/src/test/java/ccd/model/RealDataHeadToHeadTest.java @@ -0,0 +1,691 @@ +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; +import java.util.Map; + +/** + * 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 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 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<>(List.of(t.getTaxaNames())); + while (taxa.remove(null)) { + // getTaxaNames() is sized for all nodes on some inputs; drop the internal slots + } + } + if (stride > 1 && (chain - 1) % stride != 0) { + continue; + } + nwk.add(t.getRoot().toShortNewick(false) + ";"); + } + newickCache = nwk; + taxaCache = taxa; + 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. + * + *

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(); + 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)); + } + 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 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 PREPARED.isEmpty() ? read(count, 0.0, 0.25) : readPrepared("fit"); + } + + private static List valSet(int count) throws Exception { + 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) { + 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(); + } + + /** + * 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 { + // 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 = PREPARED.isEmpty() ? read(N, 0.5, 1.0) : readPrepared("test"); + List val = valSet(N / 2); + int nFit = N / 2; + + // ---- CCD1 ---- + 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}) { + RegCCD m = new RegCCD(fitSet(nFit), 0.0, alpha); + double sc = score(m::getLogProbabilityOfTree, val)[1]; + if (sc > bestAlphaScore) { + bestAlphaScore = sc; + bestAlpha = alpha; + } + } + 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; + 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; + } + } + 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 (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())); + } + + // ---- 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}; + 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}; + } + } + } + } + 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()); + 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) + 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; + 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); + 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); + } + + /** 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 { + // 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(); + 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 { + // 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); + + 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 { + // 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()); + 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 { + // 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(); + 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); + } + } +} 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)); + } +}