bits = new HashMap<>();
+ computeBits(tree.getRoot(), bits);
+ double logp = 0.0;
+ for (Node v : tree.getNodesAsArray()) {
+ if (v.isLeaf()) {
+ continue;
+ }
+ logp += logSplitProbability(bits.get(v),
+ bits.get(v.getChildren().get(0)),
+ bits.get(v.getChildren().get(1)),
+ b0, b1, b2);
+ }
+ return logp;
+ }
+
+ /**
+ * Log conditional probability of the bipartition {@code {aBits, bBits}} of clade {@code cBits},
+ * for a clade that need not be observed. This is the whole model: every internal node of a tree
+ * contributes exactly one such factor.
+ */
+ double logSplitProbability(BitSet cBits, BitSet aBits, BitSet bBits,
+ double b0, double b1, double b2) {
+ double[] wz = logWeightsAndZ(cBits, b0, b1, b2);
+ int cls = splitClass(cBits, aBits, bBits);
+ double fS = (cls == 0)
+ ? observedPartition(getClade(cBits), aBits, bBits).getNumberOfOccurrences() : 0.0;
+ double logNumerator = (fS > 0) ? Math.log(fS + Math.exp(wz[cls])) : wz[cls];
+ return logNumerator - Math.log(wz[4]);
+ }
+
+ /**
+ * Which of the four classes the bipartition {@code {aBits, bBits}} of {@code cBits} belongs to,
+ * as a 0-based index (0 = observed split, 3 = neither child observed).
+ */
+ int splitClass(BitSet cBits, BitSet aBits, BitSet bBits) {
+ if (observedPartition(getClade(cBits), aBits, bBits) != null) {
+ return 0;
+ }
+ boolean aObs = getClade(aBits) != null;
+ boolean bObs = getClade(bBits) != null;
+ return (aObs && bObs) ? 1 : ((aObs || bObs) ? 2 : 3);
+ }
+
+ private CladePartition observedPartition(Clade c, BitSet aBits, BitSet bBits) {
+ if (c == null) {
+ return null;
+ }
+ Clade ca = getClade(aBits);
+ Clade cb = getClade(bBits);
+ if (ca == null || cb == null) {
+ return null;
+ }
+ return c.getCladePartition(ca, cb);
+ }
+
+ /* ----------------------------------------------------------------------
+ * Class sizes
+ * ------------------------------------------------------------------- */
+
+ /**
+ * Sizes {@code {|A_1|, |A_2|, |A_3|, |A_4|}} of the four split classes of clade {@code cBits}.
+ *
+ * {@code |A_1|} is the number of observed splits; a single pass over the observed subclades of
+ * {@code C} yields {@code |A_3|} (observed subclade whose complement is not observed) and the
+ * number of observed-clade pairs, from which {@code |A_2|} follows; {@code |A_4|} is the
+ * remainder of {@code 2^(m-1) - 1}. Cached, and independent of the pseudocounts.
+ */
+ double[] classSizes(BitSet cBits) {
+ double[] cached = sizeCache.get(cBits);
+ if (cached != null) {
+ return cached;
+ }
+ int m = cBits.cardinality();
+ Clade c = getClade(cBits);
+
+ double n1 = 0.0;
+ if (c != null) {
+ for (CladePartition p : c.getPartitions()) {
+ if (p.getNumberOfOccurrences() > 0) {
+ n1++;
+ }
+ }
+ }
+
+ // one pass over observed clades strictly inside C
+ int bothObservedEnds = 0; // counts each both-observed bipartition twice (once per side)
+ double n3 = 0.0;
+ for (BitSet d : sortedCladeBits()) {
+ if (d.cardinality() >= m || !subset(d, cBits)) {
+ continue;
+ }
+ BitSet complement = BitSet.newBitSet(cBits);
+ complement.andNot(d);
+ if (getClade(complement) != null) {
+ bothObservedEnds++;
+ } else {
+ n3++;
+ }
+ }
+ double n2 = Math.max(0.0, bothObservedEnds / 2.0 - n1);
+
+ double total = Math.pow(2.0, m - 1) - 1.0;
+ double n4 = Math.max(0.0, total - n1 - n2 - n3);
+
+ double[] size = {n1, n2, n3, n4};
+ sizeCache.put(BitSet.newBitSet(cBits), size);
+ return size;
+ }
+
+ private synchronized List sortedCladeBits() {
+ if (sortedCladeBits == null) {
+ List all = new ArrayList<>();
+ for (Clade c : getClades()) {
+ all.add(c.getCladeInBits());
+ }
+ all.sort(CRegCCD::compareBitSets);
+ sortedCladeBits = all;
+ }
+ return sortedCladeBits;
+ }
+
+ private BitSet computeBits(Node v, Map bits) {
+ BitSet b = BitSet.newBitSet(leafArraySize);
+ if (v.isLeaf()) {
+ b.set(v.getNr());
+ } else {
+ b.or(computeBits(v.getChildren().get(0), bits));
+ b.or(computeBits(v.getChildren().get(1), bits));
+ }
+ bits.put(v, b);
+ return b;
+ }
+
+ private static boolean subset(BitSet a, BitSet c) {
+ BitSet tmp = BitSet.newBitSet(a);
+ tmp.andNot(c);
+ return tmp.isEmpty();
+ }
+
+ /* ----------------------------------------------------------------------
+ * MAP tree
+ *
+ * The maximum over all of tree space is a DP over the subset lattice, so instead we run the DP
+ * over the observed-clade DAG using only the both-children-observed splits (classes 1 and 2 --
+ * exactly CCD0's split set) and then *certify* that no off-backbone tree can beat it.
+ *
+ * The certificate is a pair of upper-bound DPs over the same DAG. U(C) bounds the best subtree
+ * log-probability over ALL trees on C, and V(C) bounds it over trees that use at least one
+ * off-backbone (class 3 or 4) split. Any subtree contributes at most 0, so a novel child is
+ * bounded by 0; every class-3 split at C shares one theta, as does every class-4 split, because
+ * the model is uniform within a class. If best(root) > V(root), no tree using an off-backbone
+ * split anywhere can beat the backbone optimum, so the backbone MAP is the global MAP.
+ * ------------------------------------------------------------------- */
+
+ private volatile Map mapBest;
+ private volatile Map mapArg;
+ private volatile double offBackboneBound = Double.NaN;
+
+ /** All both-children-observed bipartitions of {@code cBits} (classes 1 and 2), each once. */
+ private List backboneSplits(BitSet cBits) {
+ int m = cBits.cardinality();
+ List out = new ArrayList<>();
+ for (BitSet d : sortedCladeBits()) {
+ if (d.cardinality() >= m || !subset(d, cBits)) {
+ continue;
+ }
+ BitSet complement = BitSet.newBitSet(cBits);
+ complement.andNot(d);
+ if (getClade(complement) == null || compareBitSets(d, complement) >= 0) {
+ continue;
+ }
+ out.add(new BitSet[]{d, complement});
+ }
+ return out;
+ }
+
+ /** Log theta shared by every split of the given off-backbone class at {@code cBits}. */
+ private double logThetaOfClass(BitSet cBits, int cls) {
+ double[] size = classSizes(cBits);
+ if (size[cls] <= 0) {
+ return Double.NEGATIVE_INFINITY;
+ }
+ double[] wz = logWeightsAndZ(cBits, alpha, alpha1, alpha2);
+ return wz[cls] - Math.log(wz[4]);
+ }
+
+ private synchronized void computeMAP() {
+ if (mapBest != null) {
+ return;
+ }
+ List clades = new ArrayList<>(getClades());
+ clades.sort(java.util.Comparator.comparingInt(Clade::size));
+
+ Map best = new HashMap<>();
+ Map arg = new HashMap<>();
+ Map upper = new HashMap<>(); // U: best over all trees
+ Map upperOff = new HashMap<>(); // V: best over trees using an off-backbone split
+
+ for (Clade c : clades) {
+ BitSet cb = c.getCladeInBits();
+ if (c.size() == 1) {
+ best.put(cb, 0.0);
+ upper.put(cb, 0.0);
+ upperOff.put(cb, Double.NEGATIVE_INFINITY);
+ continue;
+ }
+ double bBest = Double.NEGATIVE_INFINITY;
+ BitSet[] bArg = null;
+ double bUpper = Double.NEGATIVE_INFINITY;
+ double bOff = Double.NEGATIVE_INFINITY;
+
+ for (BitSet[] s : backboneSplits(cb)) {
+ Double l = best.get(s[0]);
+ Double r = best.get(s[1]);
+ if (l == null || r == null) {
+ continue;
+ }
+ double theta = logSplitProbability(cb, s[0], s[1], alpha, alpha1, alpha2);
+ double v = theta + l + r;
+ if (v > bBest) {
+ bBest = v;
+ bArg = s;
+ }
+ double ul = upper.get(s[0]);
+ double ur = upper.get(s[1]);
+ bUpper = Math.max(bUpper, theta + ul + ur);
+ double vl = upperOff.get(s[0]);
+ double vr = upperOff.get(s[1]);
+ bOff = Math.max(bOff, theta + Math.max(vl + ur, ul + vr));
+ }
+
+ // class 3: one child observed (bounded above by U of that child, novel side by 0)
+ double t3 = logThetaOfClass(cb, 2);
+ if (t3 > Double.NEGATIVE_INFINITY) {
+ double bestObservedSide = Double.NEGATIVE_INFINITY;
+ int m = cb.cardinality();
+ for (BitSet d : sortedCladeBits()) {
+ if (d.cardinality() >= m || !subset(d, cb)) {
+ continue;
+ }
+ BitSet complement = BitSet.newBitSet(cb);
+ complement.andNot(d);
+ if (getClade(complement) == null) { // exactly one side observed
+ Double u = upper.get(d);
+ if (u != null) {
+ bestObservedSide = Math.max(bestObservedSide, u);
+ }
+ }
+ }
+ if (bestObservedSide > Double.NEGATIVE_INFINITY) {
+ bUpper = Math.max(bUpper, t3 + bestObservedSide);
+ bOff = Math.max(bOff, t3 + bestObservedSide);
+ }
+ }
+
+ // class 4: both children novel, each bounded by 0
+ double t4 = logThetaOfClass(cb, 3);
+ if (t4 > Double.NEGATIVE_INFINITY) {
+ bUpper = Math.max(bUpper, t4);
+ bOff = Math.max(bOff, t4);
+ }
+
+ best.put(cb, bBest);
+ arg.put(cb, bArg);
+ upper.put(cb, bUpper);
+ upperOff.put(cb, bOff);
+ }
+
+ this.offBackboneBound = upperOff.get(getRootClade().getCladeInBits());
+ this.mapArg = arg;
+ this.mapBest = best;
+ }
+
+ /**
+ * Exact MAP over all trees that use no two-novel-clade split, by memoised recursion over
+ * classes {@code A_0} and {@code A_1}.
+ *
+ * Restricting to {@code A_0} keeps the recursion on the observed-clade DAG. Admitting
+ * {@code A_1} as well -- peel off an observed clade, leave a novel remainder -- widens the state
+ * space to clades of the form {@code root} minus a union of disjoint observed clades. That set
+ * can in principle be large, so the search is capped by {@link #MAP_STATE_BUDGET} distinct
+ * clades; in practice it stays small because an {@code A_1} split is expensive and the recursion
+ * only ever descends.
+ *
+ *
Returns {@code {best, viaA2}} for the clade: the best log probability using only
+ * {@code A_0}/{@code A_1} splits, and an upper bound on any subtree that uses an {@code A_2}
+ * split somewhere. The second is the certificate: if {@code best > viaA2} at the root, no tree
+ * containing a two-novel-clade split can reach the optimum, so the answer is the global MAP.
+ */
+ private static final long MAP_STATE_BUDGET =
+ Long.getLong("creg.mapStates", 4_000_000L);
+
+ private static final class BudgetExhausted extends RuntimeException {
+ BudgetExhausted() {
+ super(null, null, false, false);
+ }
+ }
+
+ private double[] solveFull(BitSet cBits, int a1Budget, List> memos,
+ long[] states) {
+ Map memo = memos.get(a1Budget);
+ double[] cached = memo.get(cBits);
+ if (cached != null) {
+ return cached;
+ }
+ if (++states[0] > MAP_STATE_BUDGET) {
+ throw new BudgetExhausted();
+ }
+ int m = cBits.cardinality();
+ if (m == 1) {
+ double[] leaf = {0.0, Double.NEGATIVE_INFINITY};
+ memo.put(BitSet.newBitSet(cBits), leaf);
+ return leaf;
+ }
+ double best = Double.NEGATIVE_INFINITY;
+ double viaA2 = Double.NEGATIVE_INFINITY;
+
+ for (BitSet d : sortedCladeBits()) {
+ if (d.cardinality() >= m || !subset(d, cBits)) {
+ continue;
+ }
+ BitSet complement = BitSet.newBitSet(cBits);
+ complement.andNot(d);
+ boolean complementObserved = getClade(complement) != null;
+ if (complementObserved && compareBitSets(d, complement) >= 0) {
+ continue; // both-observed bipartitions are reached from their smaller side only
+ }
+ int childBudget = complementObserved ? a1Budget : a1Budget - 1;
+ if (childBudget < 0) {
+ continue; // no A_1 split allowance left on this path
+ }
+ double theta = logSplitProbability(cBits, d, complement, alpha, alpha1, alpha2);
+ double[] left = solveFull(d, childBudget, memos, states);
+ double[] right = solveFull(complement, childBudget, memos, states);
+ best = Math.max(best, theta + left[0] + right[0]);
+ double ul = Math.max(left[0], left[1]);
+ double ur = Math.max(right[0], right[1]);
+ viaA2 = Math.max(viaA2, theta + Math.max(left[1] + ur, ul + right[1]));
+ }
+
+ // an A_2 split taken here; both children are novel and bounded above by zero
+ double[] size = classSizes(cBits);
+ if (size[3] > 0) {
+ viaA2 = Math.max(viaA2, logThetaOfClass(cBits, 3));
+ }
+
+ double[] value = {best, viaA2};
+ memo.put(BitSet.newBitSet(cBits), value);
+ return value;
+ }
+
+ /**
+ * Result of the MAP search.
+ *
+ * {@code a2Excluded} says only that no {@code A_2} split can improve the optimum within
+ * the searched class of trees, i.e. among trees using at most {@code a1Depth} one-novel-clade
+ * splits per path. It upgrades to a genuine global certificate ({@code certifiedGlobal}) only
+ * when that depth was not binding, so that every {@code A_0}/{@code A_1} tree was considered.
+ */
+ public record MapResult(double maxLogProbability, double offBackboneBound,
+ boolean a2Excluded, int a1Depth, boolean exhaustiveInA1,
+ long statesExplored, boolean complete) {
+
+ /** True only when the search covered every A_0/A_1 tree and excluded A_2 as well. */
+ public boolean certifiedGlobal() {
+ return complete && exhaustiveInA1 && a2Excluded;
+ }
+ }
+
+ /**
+ * Runs the {@code A_0}/{@code A_1} search and reports whether the optimum it found is provably
+ * the global MAP. {@code complete} is false when the state budget was exhausted, in which case
+ * the backbone DP result should be used instead.
+ */
+ public MapResult solveMAP() {
+ return solveMAP(Integer.MAX_VALUE / 2);
+ }
+
+ /**
+ * As {@link #solveMAP()} but allowing at most {@code maxA1} one-novel-clade splits on any
+ * root-to-leaf path. {@code maxA1 = 0} is the backbone DP; raising it enlarges the search until
+ * either the optimum stops improving or the state budget is exhausted.
+ */
+ public MapResult solveMAP(int maxA1) {
+ int cap = Math.min(maxA1, getSizeOfLeavesArray());
+ List> memos = new ArrayList<>();
+ for (int i = 0; i <= cap; i++) {
+ memos.add(new HashMap<>());
+ }
+ long[] states = {0};
+ boolean exhaustive = cap >= getSizeOfLeavesArray() - 2;
+ try {
+ double[] root = solveFull(getRootClade().getCladeInBits(), cap, memos, states);
+ return new MapResult(root[0], root[1], root[0] > root[1], cap, exhaustive,
+ states[0], true);
+ } catch (BudgetExhausted e) {
+ return new MapResult(getMaxLogTreeProbability(), getOffBackboneBound(),
+ false, cap, exhaustive, states[0], false);
+ }
+ }
+
+ /** Log probability of the backbone MAP tree. */
+ @Override
+ public double getMaxLogTreeProbability() {
+ computeMAP();
+ return mapBest.get(getRootClade().getCladeInBits());
+ }
+
+ /**
+ * Whether the backbone MAP tree is provably the global MAP over all of tree space: true when no
+ * tree using a class-3 or class-4 split anywhere can reach the backbone optimum.
+ */
+ public boolean isMAPCertifiedGlobal() {
+ computeMAP();
+ return getMaxLogTreeProbability() > offBackboneBound;
+ }
+
+ /** The certificate's upper bound on any tree that uses an off-backbone split. */
+ public double getOffBackboneBound() {
+ computeMAP();
+ return offBackboneBound;
+ }
+
+ /* ----------------------------------------------------------------------
+ * Entropy
+ *
+ * The sampler draws from exactly the scored distribution, so E[-log q] is an unbiased estimate
+ * of H(q) with no truncation to correct for. A deterministic recursion would instead have to
+ * approximate the subtree entropy of novel clades, so the Monte-Carlo estimator is both simpler
+ * and more accurate here; only its standard error stands between it and the exact value.
+ * ------------------------------------------------------------------- */
+
+ /* ----------------------------------------------------------------------
+ * Deterministic entropy recursion
+ *
+ * H(C) = H_split(C) + sum_S theta(S) [H(A_S) + H(B_S)], with H(leaf) = 0.
+ *
+ * The local term is closed form: within class j >= 2 every member has the same
+ * theta_j = a_j / (|A_j| Z), so that class contributes -(a_j/Z) log theta_j as a single term --
+ * the exponentially large class 4 is never enumerated. The expectation term is exact for
+ * classes 1 and 2 (both children observed, so the recursion stays on the clade DAG) and needs a
+ * value for the novel child of a class-3 split and for both children of a class-4 split.
+ *
+ * APPROXIMATION: a novel clade is treated as *fresh*, i.e. as containing no observed clades
+ * other than its singletons, so its subtree entropy depends only on its size and is given by a
+ * universal g(k) computed once by an O(n^2) recursion. Real novel clades usually do contain
+ * observed clades, so g overestimates their structure-free entropy; the error enters only
+ * through the class-3 and class-4 branches, whose total weight at a clade is (a_3 + a_4)/Z.
+ * Class-4 splits are grouped by the sizes of the two sides, whose counts follow from the
+ * binomials minus the enumerable classes, keeping the whole pass O(K + m) per clade.
+ * ------------------------------------------------------------------- */
+
+ /** g(k): subtree entropy of a fresh (no observed subclades but singletons) clade of size k. */
+ private volatile double[] freshEntropy;
+
+ private synchronized double[] freshEntropy() {
+ if (freshEntropy != null) {
+ return freshEntropy;
+ }
+ int n = getSizeOfLeavesArray();
+ double[] g = new double[Math.max(3, n + 1)];
+ g[1] = 0.0;
+ if (g.length > 2) {
+ g[2] = 0.0; // the single split of a novel cherry has probability one
+ }
+ for (int k = 3; k <= n; k++) {
+ double total = Math.pow(2.0, k - 1) - 1.0;
+ double n3 = k; // {leaf, rest}, rest unobserved since k-1 >= 2
+ double n4 = total - n3; // both sides of size >= 2, so both unobserved
+ double z = alpha1 + (n4 > 0 ? alpha2 : 0.0);
+
+ double logT3 = Math.log(alpha1) - Math.log(n3) - Math.log(z);
+ double h = -(alpha1 / z) * logT3;
+ double e = (alpha1 / z) * g[k - 1]; // class-3 children are {1, k-1}
+ if (n4 > 0) {
+ double logT4 = Math.log(alpha2) - Math.log(n4) - Math.log(z);
+ h -= (alpha2 / z) * logT4;
+ double weighted = 0.0;
+ for (int j = 2; j <= k / 2; j++) {
+ double cnt = binomial(k, j);
+ if (j == k - j) {
+ cnt /= 2.0;
+ }
+ weighted += cnt * (g[j] + g[k - j]);
+ }
+ e += Math.exp(logT4) * weighted;
+ }
+ g[k] = h + e;
+ }
+ freshEntropy = g;
+ return g;
+ }
+
+ /** Local split entropy at {@code cBits}: {@code -sum_S theta(S) log theta(S)}, closed form. */
+ private double localSplitEntropy(BitSet cBits) {
+ double[] size = classSizes(cBits);
+ double[] wz = logWeightsAndZ(cBits, alpha, alpha1, alpha2);
+ double z = wz[4];
+ Clade c = getClade(cBits);
+ double h = 0.0;
+ if (size[0] > 0) { // class 1 is explicit: theta varies with the split count
+ double perSplit = Math.exp(wz[0]);
+ for (CladePartition p : c.getPartitions()) {
+ if (p.getNumberOfOccurrences() == 0) {
+ continue;
+ }
+ double theta = (p.getNumberOfOccurrences() + perSplit) / z;
+ h -= theta * Math.log(theta);
+ }
+ }
+ for (int j = 1; j < 4; j++) { // classes 2-4 are uniform within the class
+ if (size[j] > 0) {
+ double logTheta = wz[j] - Math.log(z);
+ h -= size[j] * Math.exp(logTheta) * logTheta;
+ }
+ }
+ return h;
+ }
+
+ /**
+ * Deterministic entropy in nats, using the fresh-clade approximation for novel subclades. Exact
+ * whenever no class-3 or class-4 split leads to a novel clade that contains an observed clade.
+ */
+ public double getEntropyRecursive() {
+ double[] g = freshEntropy();
+ List clades = new ArrayList<>(getClades());
+ clades.sort(java.util.Comparator.comparingInt(Clade::size));
+ Map entropy = new HashMap<>();
+
+ for (Clade c : clades) {
+ BitSet cb = c.getCladeInBits();
+ int m = c.size();
+ if (m == 1) {
+ entropy.put(cb, 0.0);
+ continue;
+ }
+ double[] size = classSizes(cb);
+ double[] wz = logWeightsAndZ(cb, alpha, alpha1, alpha2);
+ double z = wz[4];
+
+ double e = 0.0;
+
+ // class 1: observed splits, both children observed
+ if (size[0] > 0) {
+ double perSplit = Math.exp(wz[0]);
+ for (CladePartition p : c.getPartitions()) {
+ if (p.getNumberOfOccurrences() == 0) {
+ continue;
+ }
+ double theta = (p.getNumberOfOccurrences() + perSplit) / z;
+ e += theta * (entropy.get(p.getChildClades()[0].getCladeInBits())
+ + entropy.get(p.getChildClades()[1].getCladeInBits()));
+ }
+ }
+
+ // classes 2 and 3, plus the size profile of everything that is not class 4
+ double t2 = (size[1] > 0) ? Math.exp(wz[1]) / z : 0.0;
+ double t3 = (size[2] > 0) ? Math.exp(wz[2]) / z : 0.0;
+ double[] nonClass4 = new double[m / 2 + 1];
+ for (BitSet d : sortedCladeBits()) {
+ if (d.cardinality() >= m || !subset(d, cb)) {
+ continue;
+ }
+ BitSet complement = BitSet.newBitSet(cb);
+ complement.andNot(d);
+ int j = Math.min(d.cardinality(), complement.cardinality());
+ if (getClade(complement) != null) {
+ if (compareBitSets(d, complement) < 0) {
+ nonClass4[j]++;
+ if (observedPartition(c, d, complement) == null) { // class 2
+ e += t2 * (entropy.get(d) + entropy.get(complement));
+ }
+ }
+ } else { // class 3: d observed, complement novel
+ nonClass4[j]++;
+ e += t3 * (entropy.get(d) + g[complement.cardinality()]);
+ }
+ }
+
+ // class 4, grouped by the sizes of the two sides
+ if (size[3] > 0) {
+ double logT4 = wz[3] - Math.log(z);
+ double weighted = 0.0;
+ for (int j = 1; j <= m / 2; j++) {
+ double totalPairs = binomial(m, j);
+ if (j == m - j) {
+ totalPairs /= 2.0;
+ }
+ double count4 = totalPairs - nonClass4[j];
+ if (count4 > 0) {
+ weighted += count4 * (g[j] + g[m - j]);
+ }
+ }
+ e += Math.exp(logT4) * weighted;
+ }
+
+ entropy.put(cb, localSplitEntropy(cb) + e);
+ }
+ return entropy.get(getRootClade().getCladeInBits());
+ }
+
+ private static double binomial(int n, int k) {
+ double r = 1.0;
+ for (int i = 1; i <= k; i++) {
+ r = r * (n - k + i) / i;
+ }
+ return r;
+ }
+
+ /** Draws used by {@link #getEntropy()}. */
+ public static final int DEFAULT_ENTROPY_SAMPLES = 100_000;
+
+ /**
+ * Unbiased Monte-Carlo entropy in nats.
+ *
+ * @param samples number of draws
+ * @return {@code {estimate, standard error}}
+ */
+ public double[] getEntropyMonteCarlo(int samples) {
+ double s1 = 0.0;
+ double s2 = 0.0;
+ for (int i = 0; i < samples; i++) {
+ double logp = sampleTreeLogProbability();
+ s1 += -logp;
+ s2 += logp * logp;
+ }
+ double mean = s1 / samples;
+ double se = Math.sqrt(Math.max(0.0, s2 / samples - mean * mean) / samples);
+ return new double[]{mean, se};
+ }
+
+ /** Monte-Carlo entropy at {@link #DEFAULT_ENTROPY_SAMPLES} draws. */
+ @Override
+ public double getEntropy() {
+ return getEntropyMonteCarlo(DEFAULT_ENTROPY_SAMPLES)[0];
+ }
+
+ /** Not applicable: the Lewis recursion assumes the distribution is supported on the CCD graph. */
+ @Override
+ public double getEntropyLewis() {
+ throw new UnsupportedOperationException(
+ "CRegCCD has support outside the CCD graph; use getEntropyMonteCarlo(samples).");
+ }
+
+ /* ----------------------------------------------------------------------
+ * Sampling
+ *
+ * The generative process is the model read forwards: at each clade draw a class with
+ * probability proportional to (its observed count + a_j) over the non-empty classes, then a
+ * member uniformly within that class, then recurse into both children. Classes 1-3 are
+ * explicitly enumerable in one pass over the observed clades; class 4 is drawn by rejection
+ * from uniform bipartitions, which accepts with probability |A_4| / (2^(m-1) - 1) -- close to
+ * 1 for any clade large enough for that to matter.
+ *
+ * Every drawn split is scored with the same {@link #logSplitProbability} the scorer uses, so
+ * the sampling distribution equals exp(getLogProbabilityOfTree) by construction: there is no
+ * truncation and no separate sampling fidelity to choose.
+ * ------------------------------------------------------------------- */
+
+ /** Simulates one draw and returns its log probability, without materialising a tree. */
+ @Override
+ public double sampleTreeLogProbability() {
+ return simulate(getRootClade().getCladeInBits());
+ }
+
+ private double simulate(BitSet cBits) {
+ if (cBits.cardinality() == 1) {
+ return 0.0;
+ }
+ BitSet[] split = sampleSplit(cBits);
+ return logSplitProbability(cBits, split[0], split[1], alpha, alpha1, alpha2)
+ + simulate(split[0]) + simulate(split[1]);
+ }
+
+ /**
+ * The inherited sampler only ever picks an observed clade partition, so it would draw from the
+ * observed-splits-only distribution and could never produce a novel clade. Random sampling is
+ * therefore overridden, as is MAP, which uses this model's own backbone DP rather than the
+ * inherited CCD1 conditional clade probabilities.
+ */
+ @Override
+ protected Node getVertexBasedOnStrategy(Clade clade, SamplingStrategy samplingStrategy,
+ HeightSettingStrategy heightStrategy) {
+ if (clade.isLeaf()) {
+ return super.getVertexBasedOnStrategy(clade, samplingStrategy, heightStrategy);
+ }
+ if (samplingStrategy == SamplingStrategy.Sampling) {
+ return sampleVertex(clade.getCladeInBits(), heightStrategy);
+ }
+ if (samplingStrategy == SamplingStrategy.MAP) {
+ computeMAP();
+ return mapVertex(clade.getCladeInBits(), heightStrategy);
+ }
+ return super.getVertexBasedOnStrategy(clade, samplingStrategy, heightStrategy);
+ }
+
+ /** Traceback of the backbone MAP DP. Every clade it visits is observed, by construction. */
+ private Node mapVertex(BitSet cBits, HeightSettingStrategy heightStrategy) {
+ if (cBits.cardinality() == 1) {
+ return super.getVertexBasedOnStrategy(getClade(cBits),
+ SamplingStrategy.MAP, heightStrategy);
+ }
+ BitSet[] split = mapArg.get(cBits);
+ if (split == null) {
+ throw new AssertionError("no backbone split for clade " + cBits);
+ }
+ Node left = mapVertex(split[0], heightStrategy);
+ Node right = mapVertex(split[1], heightStrategy);
+ double logFactor = logSplitProbability(cBits, split[0], split[1], alpha, alpha1, alpha2);
+ return buildVertex(cBits, left, right, logFactor, heightStrategy);
+ }
+
+ private Node sampleVertex(BitSet cBits, HeightSettingStrategy heightStrategy) {
+ if (cBits.cardinality() == 1) {
+ // leaves are always observed clades, so the inherited leaf construction applies
+ return super.getVertexBasedOnStrategy(getClade(cBits),
+ SamplingStrategy.Sampling, heightStrategy);
+ }
+ BitSet[] split = sampleSplit(cBits);
+ Node left = sampleVertex(split[0], heightStrategy);
+ Node right = sampleVertex(split[1], heightStrategy);
+ double logFactor = logSplitProbability(cBits, split[0], split[1], alpha, alpha1, alpha2);
+ return buildVertex(cBits, left, right, logFactor, heightStrategy);
+ }
+
+ /** Assembles an internal node from two resolved children, stamping the subtree probability. */
+ private Node buildVertex(BitSet cBits, Node left, Node right, double logFactor,
+ HeightSettingStrategy heightStrategy) {
+ Node vertex = new Node();
+ vertex.setNr(nextRunningInnerIndex());
+ vertex.addChild(left);
+ vertex.addChild(right);
+
+ Clade observed = getClade(cBits);
+ double support = (observed != null) ? observed.getProbability() : 0.0;
+ vertex.setMetaData(CLADE_SUPPORT_KEY, support);
+ String posteriorSupport = CLADE_SUPPORT_KEY + "=" + support;
+ vertex.metaDataString = (vertex.metaDataString != null)
+ ? vertex.metaDataString + "," + posteriorSupport : posteriorSupport;
+
+ double logP = (Double) left.getMetaData(LOG_PROB_SUBTREE_KEY)
+ + (Double) right.getMetaData(LOG_PROB_SUBTREE_KEY) + logFactor;
+ vertex.setMetaData(LOG_PROB_SUBTREE_KEY, logP);
+ vertex.setMetaData(PROB_SUBTREE_KEY, Math.exp(logP));
+
+ setSampledHeight(vertex, left, right, observed, heightStrategy);
+ return vertex;
+ }
+
+ /** Heights: {@code One} stacks by one; the height strategies use the clade's recorded height
+ * when it is available and strictly above both children, else a minimal positive increment. */
+ private void setSampledHeight(Node vertex, Node left, Node right, Clade observed,
+ HeightSettingStrategy heightStrategy) {
+ if (heightStrategy == null || heightStrategy == HeightSettingStrategy.None) {
+ return;
+ }
+ double maxChild = Math.max(left.getHeight(), right.getHeight());
+ if (heightStrategy == HeightSettingStrategy.One) {
+ vertex.setHeight(maxChild + 1.0);
+ return;
+ }
+ double recorded = Double.NaN;
+ if (observed != null) {
+ recorded = (heightStrategy == HeightSettingStrategy.CommonAncestorHeights)
+ ? observed.getCommonAncestorHeight() : observed.getMeanOccurredHeight();
+ }
+ vertex.setHeight(recorded > maxChild ? recorded : maxChild + NOVEL_HEIGHT_EPS);
+ }
+
+ /** Draws one bipartition of {@code cBits} from this model's conditional distribution. */
+ private BitSet[] sampleSplit(BitSet cBits) {
+ double[] size = classSizes(cBits);
+ Clade c = getClade(cBits);
+ double[] wz = logWeightsAndZ(cBits, alpha, alpha1, alpha2);
+
+ // total mass of each class: counts plus its share of the prior
+ double[] weight = new double[4];
+ for (int j = 0; j < 4; j++) {
+ weight[j] = (size[j] > 0) ? Math.exp(wz[j]) * size[j] : 0.0;
+ }
+ if (size[0] > 0) {
+ weight[0] += c.getNumberOfOccurrences();
+ }
+
+ double total = weight[0] + weight[1] + weight[2] + weight[3];
+ double target = random().nextDouble() * total;
+ int cls = -1;
+ double acc = 0.0;
+ for (int j = 0; j < 4; j++) {
+ if (weight[j] <= 0) {
+ continue;
+ }
+ acc += weight[j];
+ if (target < acc) {
+ cls = j;
+ break;
+ }
+ }
+ if (cls < 0) { // numerical guard: fall back to the last non-empty class
+ for (int j = 3; j >= 0; j--) {
+ if (size[j] > 0) {
+ cls = j;
+ break;
+ }
+ }
+ }
+
+ return switch (cls) {
+ case 0 -> sampleObservedSplit(c, size[0]);
+ case 1, 2 -> sampleEnumerableSplit(cBits, cls);
+ default -> sampleNovelSplit(cBits, size[3]);
+ };
+ }
+
+ /** Class 1: an observed split, with weight {@code f(S) + a_1/|A_1|}. */
+ private BitSet[] sampleObservedSplit(Clade c, double n1) {
+ double perSplit = Math.exp(logWeightsAndZ(c.getCladeInBits(), alpha, alpha1, alpha2)[0]);
+ double totalWeight = c.getNumberOfOccurrences() + perSplit * n1;
+ double target = random().nextDouble() * totalWeight;
+ double acc = 0.0;
+ CladePartition last = null;
+ for (CladePartition p : c.getPartitions()) {
+ if (p.getNumberOfOccurrences() == 0) {
+ continue;
+ }
+ last = p;
+ acc += p.getNumberOfOccurrences() + perSplit;
+ if (target < acc) {
+ return childBits(p);
+ }
+ }
+ return childBits(last); // numerical guard
+ }
+
+ private static BitSet[] childBits(CladePartition p) {
+ return new BitSet[]{p.getChildClades()[0].getCladeInBits(),
+ p.getChildClades()[1].getCladeInBits()};
+ }
+
+ /**
+ * Classes 2 and 3, drawn uniformly by reservoir sampling over the same single pass across the
+ * observed clades inside {@code cBits} that produced the class sizes. A both-observed
+ * bipartition is reached from either side, so it is only considered from its canonically
+ * smaller side; a one-observed bipartition is reached exactly once, from its observed side.
+ */
+ private BitSet[] sampleEnumerableSplit(BitSet cBits, int cls) {
+ int m = cBits.cardinality();
+ int seen = 0;
+ BitSet[] pick = null;
+ for (BitSet d : sortedCladeBits()) {
+ if (d.cardinality() >= m || !subset(d, cBits)) {
+ continue;
+ }
+ BitSet complement = BitSet.newBitSet(cBits);
+ complement.andNot(d);
+ boolean complementObserved = getClade(complement) != null;
+ boolean candidate;
+ if (complementObserved) {
+ candidate = cls == 1
+ && compareBitSets(d, complement) < 0
+ && observedPartition(getClade(cBits), d, complement) == null;
+ } else {
+ candidate = cls == 2;
+ }
+ if (candidate) {
+ seen++;
+ if (random().nextInt(seen) == 0) {
+ pick = new BitSet[]{d, complement};
+ }
+ }
+ }
+ return pick;
+ }
+
+ /**
+ * Class 4, drawn uniformly among bipartitions with neither side an observed clade. Uniform
+ * bipartitions are generated by pinning the lowest taxon to one side and flipping a fair coin
+ * for the rest, and rejected unless both sides are novel. When acceptance would be poor the
+ * bipartition set is necessarily small, so it is enumerated instead.
+ */
+ private BitSet[] sampleNovelSplit(BitSet cBits, double n4) {
+ int m = cBits.cardinality();
+ int[] idx = new int[m];
+ int k = 0;
+ for (int b = cBits.nextSetBit(0); b >= 0; b = cBits.nextSetBit(b + 1)) {
+ idx[k++] = b;
+ }
+ double total = Math.pow(2.0, m - 1) - 1.0;
+
+ if (n4 / total >= MIN_REJECTION_ACCEPTANCE) {
+ while (true) {
+ BitSet left = BitSet.newBitSet(leafArraySize);
+ BitSet right = BitSet.newBitSet(leafArraySize);
+ left.set(idx[0]);
+ for (int i = 1; i < m; i++) {
+ if (random().nextBoolean()) {
+ left.set(idx[i]);
+ } else {
+ right.set(idx[i]);
+ }
+ }
+ if (right.isEmpty()) {
+ continue;
+ }
+ if (getClade(left) == null && getClade(right) == null) {
+ return new BitSet[]{left, right};
+ }
+ }
+ }
+
+ // low acceptance => 2^(m-1) is small; enumerate and reservoir-sample
+ int seen = 0;
+ BitSet[] pick = null;
+ for (int mask = 0; mask < (1 << (m - 1)); mask++) {
+ BitSet left = BitSet.newBitSet(leafArraySize);
+ BitSet right = BitSet.newBitSet(leafArraySize);
+ left.set(idx[0]);
+ for (int i = 1; i < m; i++) {
+ if ((mask & (1 << (i - 1))) != 0) {
+ left.set(idx[i]);
+ } else {
+ right.set(idx[i]);
+ }
+ }
+ if (right.isEmpty()) {
+ continue;
+ }
+ if (getClade(left) == null && getClade(right) == null) {
+ seen++;
+ if (random().nextInt(seen) == 0) {
+ pick = new BitSet[]{left, right};
+ }
+ }
+ }
+ return pick;
+ }
+
+ private static int compareBitSets(BitSet a, BitSet b) {
+ int ia = a.nextSetBit(0);
+ int ib = b.nextSetBit(0);
+ while (ia >= 0 && ib >= 0) {
+ if (ia != ib) {
+ return Integer.compare(ia, ib);
+ }
+ ia = a.nextSetBit(ia + 1);
+ ib = b.nextSetBit(ib + 1);
+ }
+ return Integer.compare(ia, ib);
+ }
+}
diff --git a/src/test/java/ccd/model/CRegCCDMapEntropyTest.java b/src/test/java/ccd/model/CRegCCDMapEntropyTest.java
new file mode 100644
index 0000000..6188288
--- /dev/null
+++ b/src/test/java/ccd/model/CRegCCDMapEntropyTest.java
@@ -0,0 +1,298 @@
+package ccd.model;
+
+import beast.base.evolution.tree.Tree;
+import beast.base.evolution.tree.TreeParser;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Random;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/** {@link CRegCCD}: MAP tree (with its global-optimality certificate) and entropy. */
+public class CRegCCDMapEntropyTest {
+
+ private static List taxa(int n) {
+ List out = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ out.add("T" + i);
+ }
+ return out;
+ }
+
+ private static List randomTrees(List taxa, int nTrees, long seed) {
+ Random rng = new Random(seed);
+ List out = new ArrayList<>();
+ for (int t = 0; t < nTrees; t++) {
+ List pool = new ArrayList<>(taxa);
+ while (pool.size() > 1) {
+ String a = pool.remove(rng.nextInt(pool.size()));
+ String b = pool.remove(rng.nextInt(pool.size()));
+ pool.add("(" + a + "," + b + ")");
+ }
+ out.add(new TreeParser(taxa, pool.get(0) + ";", 1, false));
+ }
+ return out;
+ }
+
+ /**
+ * The MAP tree must be the argmax over ALL topologies, not just those on the backbone. Checked
+ * by exhaustive enumeration, together with the certificate that claims it.
+ */
+ @Test
+ public void mapTreeIsGlobalOptimum() {
+ for (int n : new int[]{5, 6, 7}) {
+ for (int nTrees : new int[]{2, 10, 40}) {
+ List tx = taxa(n);
+ CRegCCD ccd = new CRegCCD(randomTrees(tx, nTrees, 13L), 0.0, 0.4, 0.4, 0.4);
+
+ double bruteForce = Double.NEGATIVE_INFINITY;
+ for (Tree t : CRegCCDTest.allRootedTopologies(tx)) {
+ bruteForce = Math.max(bruteForce, ccd.getLogProbabilityOfTree(t));
+ }
+ double reported = ccd.getMaxLogTreeProbability();
+ boolean certified = ccd.isMAPCertifiedGlobal();
+ System.out.printf("CRegCCD MAP %d taxa, %2d trees: brute=%.6f DP=%.6f "
+ + "certified=%-5s (off-backbone bound %.3f)%n",
+ n, nTrees, bruteForce, reported, certified, ccd.getOffBackboneBound());
+
+ assertEquals(bruteForce, reported, 1e-9,
+ "backbone DP must find the global maximum (" + n + " taxa)");
+
+ // the returned tree must actually attain it
+ Tree map = ccd.getMAPTree();
+ assertEquals(reported, ccd.getLogProbabilityOfTree(map), 1e-9,
+ "returned MAP tree must attain the reported maximum");
+ assertEquals(tx.size(), map.getLeafNodeCount());
+
+ // the certificate must never be wrong when it fires
+ if (certified) {
+ assertTrue(ccd.getOffBackboneBound() < bruteForce + 1e-12,
+ "certificate claimed optimality but the bound exceeds the optimum");
+ }
+ }
+ }
+ }
+
+ @Test
+ public void entropyMatchesEnumeration() {
+ List tx = taxa(6);
+ CRegCCD ccd = new CRegCCD(randomTrees(tx, 8, 17L), 0.0, 0.5, 0.3, 0.2);
+ ccd.setRandom(new Random(2024L));
+
+ double exact = 0.0;
+ for (Tree t : CRegCCDTest.allRootedTopologies(tx)) {
+ double logp = ccd.getLogProbabilityOfTree(t);
+ exact -= Math.exp(logp) * logp;
+ }
+ double[] mc = ccd.getEntropyMonteCarlo(500_000);
+ System.out.printf("CRegCCD entropy: exact=%.5f MC=%.5f +/- %.5f (%.1f SE off)%n",
+ exact, mc[0], mc[1], Math.abs(mc[0] - exact) / mc[1]);
+ assertEquals(exact, mc[0], Math.max(5 * mc[1], 0.002), "MC entropy must match enumeration");
+
+ assertTrue(Double.isFinite(ccd.getEntropy()), "default entropy must be finite");
+ }
+
+ /* ------------------------------------------------------------------ *
+ * The manuscript's four-taxon example, worked symbolically.
+ * Sample: (((A,B),C),D) and (((D,C),B),A), each once.
+ * ------------------------------------------------------------------ */
+
+ private static final List TAXA4 = Arrays.asList("A", "B", "C", "D");
+
+ static CRegCCD exampleModel(double alpha, double alpha1, double alpha2) {
+ List trees = new ArrayList<>();
+ trees.add(new TreeParser(TAXA4, "(((A,B),C),D);", 1, false));
+ trees.add(new TreeParser(TAXA4, "(((D,C),B),A);", 1, false));
+ return new CRegCCD(trees, 0.0, alpha, alpha1, alpha2);
+ }
+
+ /** The 15 four-taxon trees, grouped by the six probability categories. */
+ static final String[][] CATEGORIES = {
+ {"(((A,B),C),D);", "(((C,D),B),A);"}, // sampled
+ {"((A,B),(C,D));"}, // both children observed
+ {"(((A,C),B),D);", "(((B,C),A),D);", "(((B,D),C),A);", "(((B,C),D),A);"}, // novel inside an observed 3-clade
+ {"(((C,D),A),B);", "(((A,B),D),C);"}, // root one-observed, reconnects
+ {"(((A,C),D),B);", "(((A,D),C),B);", "(((A,D),B),C);", "(((B,D),A),C);"}, // root one-observed, novel cherry
+ {"((A,C),(B,D));", "((A,D),(B,C));"} // root neither observed
+ };
+
+ /** POOLED mode: per-split alpha over the three CCD0 splits, class totals a3 and a4. */
+ static double[] categoryProbabilities(double alpha, double a3, double a4) {
+ double z = 2 + 3 * alpha + a3 + a4;
+ return new double[]{
+ (1 + alpha) * (1 + alpha) / (z * (1 + alpha + a3)),
+ alpha / z,
+ (1 + alpha) * (a3 / 2) / (z * (1 + alpha + a3)),
+ (a3 / 2) * alpha / (z * (alpha + a3)),
+ (a3 / 2) * (a3 / 2) / (z * (alpha + a3)),
+ (a4 / 2) / z
+ };
+ }
+
+ /**
+ * How large is the fresh-clade approximation's error? Compares the deterministic recursion
+ * against exact enumeration across taxon counts, training-set sizes and pseudocounts.
+ */
+ @Test
+ public void deterministicRecursionErrorVersusEnumeration() {
+ System.out.printf("%n%-6s %-7s %-22s %-12s %-12s %-11s %-9s%n",
+ "taxa", "trees", "pseudocounts", "exact H", "recursion", "abs err", "rel err");
+ double worstRel = 0.0;
+ for (int n : new int[]{5, 6, 7, 8}) {
+ for (int nTrees : new int[]{2, 10, 50}) {
+ for (double[] p : new double[][]{{0.4, 0.4, 0.4}, {2.0, 0.4, 0.05}, {1.0, 1.0, 1.0}}) {
+ List tx = taxa(n);
+ CRegCCD ccd = new CRegCCD(randomTrees(tx, nTrees, 5L), 0.0,
+ p[0], p[1], p[2]);
+ double exact = 0.0;
+ for (Tree t : CRegCCDTest.allRootedTopologies(tx)) {
+ double logp = ccd.getLogProbabilityOfTree(t);
+ exact -= Math.exp(logp) * logp;
+ }
+ double rec = ccd.getEntropyRecursive();
+ double abs = Math.abs(rec - exact);
+ double rel = abs / exact;
+ worstRel = Math.max(worstRel, rel);
+ System.out.printf("%-6d %-7d a=(%.2f,%.2f,%.2f)%-6s %-12.6f %-12.6f %-11.2e %-8.3f%%%n",
+ n, nTrees, p[0], p[1], p[2], "", exact, rec, abs, 100 * rel);
+ }
+ }
+ }
+ System.out.printf("worst relative error = %.3f%%%n", 100 * worstRel);
+ assertTrue(worstRel < 0.25, "recursion should be within 25% of exact, worst was " + worstRel);
+ }
+
+ @Test
+ public void fourTaxonExampleMatchesClosedFormPooled() {
+ for (double[] p : new double[][]{{0.4, 0.4, 0.4}, {0.4, 0.4, 0.05}, {1.0, 0.5, 0.25}}) {
+ CRegCCD ccd = exampleModel(p[0], p[1], p[2]);
+ double[] expected = categoryProbabilities(p[0], p[1], p[2]);
+ double total = 0.0;
+ int count = 0;
+ for (int g = 0; g < CATEGORIES.length; g++) {
+ for (String nwk : CATEGORIES[g]) {
+ Tree t = new TreeParser(TAXA4, nwk, 1, false);
+ assertEquals(expected[g], ccd.getProbabilityOfTree(t), 1e-12,
+ "pooled category " + (g + 1) + " tree " + nwk);
+ total += ccd.getProbabilityOfTree(t);
+ count++;
+ }
+ }
+ assertEquals(15, count);
+ assertEquals(1.0, total, 1e-12, "the 15 probabilities must sum to one");
+ System.out.printf("pooled four-taxon example alpha=%.2f a3=%.2f a4=%.2f: "
+ + "closed form verified, sum = %.12f%n", p[0], p[1], p[2], total);
+ }
+ }
+
+ /**
+ * Under POOLED an observed split must always outrank an expanded one at the same clade, since
+ * they share a per-split pseudocount and the observed one adds f(S) >= 1. Under SEPARATE that
+ * can fail.
+ */
+ @Test
+ public void pooledGuaranteesObservedOutranksExpanded() {
+ for (int n : new int[]{6, 8}) {
+ for (int nTrees : new int[]{3, 20}) {
+ List tx = taxa(n);
+ List training = randomTrees(tx, nTrees, 77L);
+ for (double alpha : new double[]{0.05, 0.4, 2.0, 10.0}) {
+ CRegCCD pooled = new CRegCCD(training, 0.0, alpha, 0.4, 0.05);
+ for (Clade c : pooled.getClades()) {
+ if (c.size() < 2) {
+ continue;
+ }
+ double[] size = pooled.classSizes(c.getCladeInBits());
+ if (size[0] <= 0 || size[1] <= 0) {
+ continue;
+ }
+ // every expanded split has the same probability; take the largest observed one
+ double worstObserved = Double.POSITIVE_INFINITY;
+ for (CladePartition p : c.getPartitions()) {
+ if (p.getNumberOfOccurrences() > 0) {
+ worstObserved = Math.min(worstObserved, p.getNumberOfOccurrences());
+ }
+ }
+ // numerators: observed f + alpha, expanded alpha
+ assertTrue(worstObserved + alpha > alpha,
+ "observed must outrank expanded at " + c);
+ }
+ }
+ }
+ }
+ System.out.println("pooled: observed splits outrank expanded splits at every clade");
+ }
+
+ /**
+ * regCCD is nested exactly: with no prior mass on the novel classes, CRegCCD's alpha is
+ * regCCD's additive-alpha smoothing over the CCD0 split set, so the two must agree on every
+ * tree that regCCD supports.
+ */
+ @Test
+ public void regCCDIsNestedAtAlphaOneTwoZero() {
+ for (int n : new int[]{5, 6, 7}) {
+ for (int nTrees : new int[]{3, 15}) {
+ List tx = taxa(n);
+ for (double alpha : new double[]{0.1, 0.4, 1.0}) {
+ CRegCCD creg = new CRegCCD(randomTrees(tx, nTrees, 31L), 0.0, alpha, 0.0, 0.0);
+ RegCCD reg = new RegCCD(randomTrees(tx, nTrees, 31L), 0.0, alpha);
+ int compared = 0;
+ double worst = 0.0;
+ for (Tree t : CRegCCDTest.allRootedTopologies(tx)) {
+ double a = reg.getLogProbabilityOfTree(t);
+ if (!Double.isFinite(a)) {
+ continue; // outside regCCD's support
+ }
+ worst = Math.max(worst, Math.abs(a - creg.getLogProbabilityOfTree(t)));
+ compared++;
+ }
+ System.out.printf("regCCD nesting %d taxa, %2d trees, alpha=%.1f: "
+ + "%d trees compared, max |diff| = %.2e%n", n, nTrees, alpha, compared, worst);
+ assertTrue(compared > 0, "regCCD must support some trees");
+ assertEquals(0.0, worst, 1e-9,
+ "CRegCCD(alpha, 0, 0) must equal regCCD(alpha)");
+ }
+ }
+ }
+ }
+
+ /**
+ * The A_0/A_1 search must equal the true global optimum by enumeration, and its certificate
+ * must never claim optimality wrongly.
+ */
+ @Test
+ public void exactMapSearchMatchesEnumeration() {
+ int certified = 0, total = 0;
+ for (int n : new int[]{5, 6, 7, 8}) {
+ for (int nTrees : new int[]{2, 5, 20}) {
+ List tx = taxa(n);
+ CRegCCD ccd = new CRegCCD(randomTrees(tx, nTrees, 13L), 0.0, 0.4, 0.4, 0.05);
+ double brute = Double.NEGATIVE_INFINITY;
+ for (Tree t : CRegCCDTest.allRootedTopologies(tx)) {
+ brute = Math.max(brute, ccd.getLogProbabilityOfTree(t));
+ }
+ CRegCCD.MapResult r = ccd.solveMAP();
+ total++;
+ if (r.certifiedGlobal()) {
+ certified++;
+ }
+ System.out.printf("%d taxa, %2d trees: brute=%.6f A0/A1=%.6f bound=%8.3f "
+ + "certified=%-5s states=%d%n",
+ n, nTrees, brute, r.maxLogProbability(), r.offBackboneBound(),
+ r.certifiedGlobal(), r.statesExplored());
+ assertTrue(r.complete(), "search must complete within the state budget");
+ assertEquals(brute, r.maxLogProbability(), 1e-9,
+ "A_0/A_1 search must find the global optimum");
+ if (r.certifiedGlobal()) {
+ assertTrue(r.offBackboneBound() < brute + 1e-12,
+ "certificate must not claim optimality when the bound exceeds it");
+ }
+ }
+ }
+ System.out.printf("certificate fired in %d of %d configurations%n", certified, total);
+ }
+}
diff --git a/src/test/java/ccd/model/CRegCCDTest.java b/src/test/java/ccd/model/CRegCCDTest.java
new file mode 100644
index 0000000..08cfdb6
--- /dev/null
+++ b/src/test/java/ccd/model/CRegCCDTest.java
@@ -0,0 +1,293 @@
+package ccd.model;
+
+import beast.base.evolution.tree.Tree;
+import beast.base.evolution.tree.TreeParser;
+import ccd.model.bitsets.BitSet;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Random;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/** {@link CRegCCD}: exact normalisation, class-size arithmetic, and full support. */
+public class CRegCCDTest {
+
+ private static List taxa(int n) {
+ List out = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ out.add("T" + i);
+ }
+ return out;
+ }
+
+ private static List randomTrees(List taxa, int nTrees, long seed) {
+ Random rng = new Random(seed);
+ List out = new ArrayList<>();
+ for (int t = 0; t < nTrees; t++) {
+ List pool = new ArrayList<>(taxa);
+ while (pool.size() > 1) {
+ String a = pool.remove(rng.nextInt(pool.size()));
+ String b = pool.remove(rng.nextInt(pool.size()));
+ pool.add("(" + a + "," + b + ")");
+ }
+ out.add(new TreeParser(taxa, pool.get(0) + ";", 1, false));
+ }
+ return out;
+ }
+
+ private static double totalMass(CRegCCD ccd, List taxa) {
+ double sum = 0.0;
+ for (Tree t : allRootedTopologies(taxa)) {
+ sum += Math.exp(ccd.getLogProbabilityOfTree(t));
+ }
+ return sum;
+ }
+
+ @Test
+ public void exactlyNormalisedOverTreeSpace() {
+ for (int n : new int[]{4, 5, 6, 7}) {
+ List tx = taxa(n);
+ for (int nTrees : new int[]{1, 3, 20}) {
+ List training = randomTrees(tx, nTrees, 7L);
+ CRegCCD ccd = new CRegCCD(training, 0.0, 0.4, 0.4, 0.4);
+ double mass = totalMass(ccd, tx);
+ System.out.printf("CRegCCD %d taxa, %2d training trees: totalMass = %.12f%n",
+ n, nTrees, mass);
+ assertEquals(1.0, mass, 1e-9,
+ "CRegCCD must be exactly normalised (" + n + " taxa, " + nTrees + " trees)");
+ }
+ }
+ }
+
+ @Test
+ public void normalisedAcrossPseudocountChoices() {
+ List tx = taxa(6);
+ List training = randomTrees(tx, 10, 11L);
+ CRegCCD ccd = new CRegCCD(training, 0.0);
+ for (double[] p : new double[][]{
+ {0.0, 0.4, 0.4, 0.4},
+ {1.0, 1.0, 1.0, 1.0},
+ {0.0, 2.0, 0.5, 0.05},
+ {0.3, 0.01, 5.0, 0.2}}) {
+ double sum = 0.0;
+ for (Tree t : allRootedTopologies(tx)) {
+ sum += Math.exp(ccd.getLogProbabilityOfTree(t, p[1], p[2], p[3]));
+ }
+ System.out.printf("CRegCCD 6 taxa a=(%.2f,%.2f,%.2f,%.2f): totalMass = %.12f%n",
+ p[0], p[1], p[2], p[3], sum);
+ assertEquals(1.0, sum, 1e-9, "normalisation must hold for any pseudocounts");
+ }
+ }
+
+ @Test
+ public void classSizesPartitionAllBipartitions() {
+ List tx = taxa(8);
+ List training = randomTrees(tx, 50, 3L);
+ CRegCCD ccd = new CRegCCD(training, 0.0);
+ CCD0 ccd0 = new CCD0(training, 0);
+ int checked = 0;
+ for (Clade c : ccd.getClades()) {
+ if (c.size() < 2) {
+ continue;
+ }
+ double[] s = ccd.classSizes(c.getCladeInBits());
+ double total = Math.pow(2.0, c.size() - 1) - 1.0;
+ assertEquals(total, s[0] + s[1] + s[2] + s[3], 1e-6,
+ "class sizes must partition all bipartitions of " + c);
+
+ // |A_1| + |A_2| is exactly the CCD0 split count (observed + expanded)
+ Clade c0 = ccd0.getClade(c.getCladeInBits());
+ if (c0 != null) {
+ assertEquals(c0.getPartitions().size(), (int) Math.round(s[0] + s[1]),
+ "|A_1|+|A_2| must equal the CCD0 partition count for " + c);
+ }
+ checked++;
+ }
+ System.out.printf("class-size arithmetic verified on %d clades%n", checked);
+ assertTrue(checked > 0);
+ }
+
+ @Test
+ public void fullSupportOnHeldOutTrees() {
+ List tx = taxa(30);
+ List training = randomTrees(tx, 100, 5L);
+ List heldOut = randomTrees(tx, 100, 99L);
+ CRegCCD ccd = new CRegCCD(training, 0.0);
+ int covered = 0;
+ double sum = 0.0;
+ for (Tree t : heldOut) {
+ double lp = ccd.getLogProbabilityOfTree(t);
+ if (Double.isFinite(lp)) {
+ covered++;
+ sum += lp;
+ }
+ }
+ System.out.printf("CRegCCD 30 taxa: %d/%d held-out trees with finite logP, mean = %.2f%n",
+ covered, heldOut.size(), sum / covered);
+ assertEquals(heldOut.size(), covered, "every held-out tree must have positive probability");
+ }
+
+ /**
+ * If the simulator draws from q AND reports the correct log q, then the mean sampled -log q
+ * equals the entropy computed by enumeration with the scorer. This checks the sampler and the
+ * scorer are the same distribution without needing per-topology counts.
+ */
+ @Test
+ public void samplerEntropyMatchesScorer() {
+ List tx = taxa(6);
+ List training = randomTrees(tx, 8, 17L);
+ CRegCCD ccd = new CRegCCD(training, 0.0, 0.5, 0.3, 0.2);
+ ccd.setRandom(new Random(31337L));
+
+ double mass = 0.0;
+ double h = 0.0;
+ for (Tree t : allRootedTopologies(tx)) {
+ double logp = ccd.getLogProbabilityOfTree(t);
+ double p = Math.exp(logp);
+ mass += p;
+ h -= p * logp;
+ }
+ assertEquals(1.0, mass, 1e-9, "scored distribution must be normalised");
+
+ int n = 1_000_000;
+ double s1 = 0.0;
+ double s2 = 0.0;
+ for (int i = 0; i < n; i++) {
+ double logp = ccd.sampleTreeLogProbability();
+ s1 += -logp;
+ s2 += logp * logp;
+ }
+ double hHat = s1 / n;
+ double se = Math.sqrt(Math.max(0, s2 / n - hHat * hHat) / n);
+ System.out.printf("CRegCCD sampler: H_enum=%.5f H_MC=%.5f +/- %.5f (%.1f SE off)%n",
+ h, hHat, se, Math.abs(hHat - h) / se);
+ assertEquals(h, hHat, Math.max(5 * se, 0.005),
+ "sampler entropy must match the scorer's enumerated entropy");
+ }
+
+ /**
+ * The direct check: sampled topology frequencies must match the scored probabilities. Compares
+ * every topology whose expected count is large enough for a normal approximation.
+ */
+ @Test
+ public void sampledFrequenciesMatchScoredProbabilities() {
+ List tx = taxa(5);
+ List training = randomTrees(tx, 5, 23L);
+ CRegCCD ccd = new CRegCCD(training, 0.0, 0.5, 0.3, 0.2);
+ ccd.setRandom(new Random(4242L));
+
+ java.util.Map expected = new java.util.HashMap<>();
+ for (Tree t : allRootedTopologies(tx)) {
+ expected.put(canonical(t), Math.exp(ccd.getLogProbabilityOfTree(t)));
+ }
+
+ int n = 1_000_000;
+ java.util.Map counts = new java.util.HashMap<>();
+ for (int i = 0; i < n; i++) {
+ counts.merge(canonical(ccd.sampleTree()), 1, Integer::sum);
+ }
+
+ int checked = 0;
+ double worst = 0.0;
+ String worstKey = null;
+ for (var e : expected.entrySet()) {
+ double p = e.getValue();
+ if (n * p < 30) {
+ continue; // too rare for a normal approximation
+ }
+ int obs = counts.getOrDefault(e.getKey(), 0);
+ double z = Math.abs(obs - n * p) / Math.sqrt(n * p * (1 - p));
+ if (z > worst) {
+ worst = z;
+ worstKey = e.getKey();
+ }
+ checked++;
+ }
+ System.out.printf("CRegCCD frequencies: %d topologies checked, worst |z| = %.2f (%s)%n",
+ checked, worst, worstKey);
+ assertTrue(checked >= 10, "expected a reasonable number of comparable topologies");
+ assertTrue(worst < 4.5, "sampled frequencies must match scored probabilities, worst z = " + worst);
+
+ // no sampled topology may fall outside the enumerated support
+ for (String key : counts.keySet()) {
+ assertTrue(expected.containsKey(key), "sampler produced an unrecognised topology " + key);
+ }
+ }
+
+ @Test
+ public void sampledTreesAreValidAndSelfConsistent() {
+ List tx = taxa(20);
+ List training = randomTrees(tx, 50, 8L);
+ CRegCCD ccd = new CRegCCD(training, 0.0);
+ ccd.setRandom(new Random(5L));
+ for (int i = 0; i < 200; i++) {
+ Tree t = ccd.sampleTree();
+ assertEquals(tx.size(), t.getLeafNodeCount(), "sampled tree must have every taxon");
+ assertEquals(2 * tx.size() - 1, t.getNodeCount(), "sampled tree must be binary");
+ double stamped = (Double) t.getRoot().getMetaData(CCD1.LOG_PROB_SUBTREE_KEY);
+ assertEquals(ccd.getLogProbabilityOfTree(t), stamped, 1e-9,
+ "stamped log probability must equal the scorer's");
+ }
+ System.out.println("CRegCCD: 200 sampled 20-taxon trees valid and self-consistent");
+ }
+
+ /** Canonical topology key: nested sorted taxon-index sets. */
+ private static String canonical(Tree t) {
+ return canonical(t.getRoot());
+ }
+
+ private static String canonical(beast.base.evolution.tree.Node v) {
+ if (v.isLeaf()) {
+ return String.valueOf(v.getNr());
+ }
+ String a = canonical(v.getChild(0));
+ String b = canonical(v.getChild(1));
+ return (a.compareTo(b) <= 0) ? "(" + a + "," + b + ")" : "(" + b + "," + a + ")";
+ }
+
+ /* --------------------------------------------------------------------- */
+
+ static List allRootedTopologies(List taxa) {
+ List trees = new ArrayList<>();
+ for (String shape : shapes(taxa)) {
+ trees.add(new TreeParser(taxa, shape + ";", 1, false));
+ }
+ return trees;
+ }
+
+ private static List shapes(List taxa) {
+ List out = new ArrayList<>();
+ if (taxa.size() == 1) {
+ out.add(taxa.get(0) + ":1");
+ return out;
+ }
+ String first = taxa.get(0);
+ List rest = taxa.subList(1, taxa.size());
+ int n = rest.size();
+ for (int mask = 0; mask < (1 << n); mask++) {
+ List left = new ArrayList<>();
+ left.add(first);
+ List right = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ if ((mask & (1 << i)) != 0) {
+ left.add(rest.get(i));
+ } else {
+ right.add(rest.get(i));
+ }
+ }
+ if (right.isEmpty()) {
+ continue;
+ }
+ for (String l : shapes(left)) {
+ for (String r : shapes(right)) {
+ out.add("(" + l + "," + r + "):1");
+ }
+ }
+ }
+ return out;
+ }
+}
diff --git a/src/test/java/ccd/model/ClassUsageAnalysis.java b/src/test/java/ccd/model/ClassUsageAnalysis.java
new file mode 100644
index 0000000..42ca47c
--- /dev/null
+++ b/src/test/java/ccd/model/ClassUsageAnalysis.java
@@ -0,0 +1,227 @@
+package ccd.model;
+
+import beast.base.evolution.tree.Node;
+import beast.base.evolution.tree.Tree;
+import beast.base.evolution.tree.TreeParser;
+import beastfx.app.treeannotator.TreeAnnotator;
+import ccd.model.bitsets.BitSet;
+import ccd.tools.CCDToolUtil;
+import org.junit.jupiter.api.Assumptions;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Where does a real held-out tree's probability actually go under {@link CRegCCD}?
+ *
+ * Walks every internal node of every held-out tree, classifies its split into the four classes,
+ * and tallies how many nodes fall in each class and how much log probability each class contributes.
+ * This measures the concern that a class-4 split (neither child observed) reconnects to the CCD only
+ * by chance: if class 4 is both rare and responsible for a large share of the total loss, the
+ * uniform-within-class-4 prior is the binding weakness.
+ *
+ *
Also reports, for each class-4 node encountered, how many observed clades survive intact inside
+ * the two novel children -- i.e. how much backbone a class-4 split destroys.
+ */
+public class ClassUsageAnalysis {
+
+ private static final String PATH = System.getProperty("ccd.trees", "");
+ private static final int N = Integer.parseInt(System.getProperty("ccd.n", "1000"));
+
+ private static List newickCache;
+ private static List taxaCache;
+
+ private static void load() throws Exception {
+ if (newickCache != null) {
+ return;
+ }
+ TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(PATH, 10);
+ ts.reset();
+ List nwk = new ArrayList<>();
+ List taxa = null;
+ while (ts.hasNext()) {
+ Tree t = ts.next();
+ if (taxa == null) {
+ String[] byNr = new String[t.getLeafNodeCount()];
+ for (Node leaf : t.getExternalNodes()) {
+ byNr[leaf.getNr()] = leaf.getID();
+ }
+ taxa = new ArrayList<>(List.of(byNr));
+ }
+ nwk.add(t.getRoot().toNewick() + ";");
+ }
+ newickCache = nwk;
+ taxaCache = taxa;
+ }
+
+ private static List read(int count, double from, double to) throws Exception {
+ load();
+ List pool = newickCache.subList((int) (from * newickCache.size()),
+ (int) (to * newickCache.size()));
+ List out = new ArrayList<>();
+ double step = Math.max(1.0, pool.size() / (double) count);
+ for (int i = 0; i < count && (int) (i * step) < pool.size(); i++) {
+ out.add(new TreeParser(taxaCache, pool.get((int) (i * step)), 0, false));
+ }
+ return out;
+ }
+
+ @Test
+ public void classUsageOnHeldOutTrees() throws Exception {
+ Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(),
+ "set -Dccd.trees=/path/to/x.trees");
+
+ CRegCCD ccd = new CRegCCD(read(N, 0.0, 0.5), 0.0,
+ Double.parseDouble(System.getProperty("ccd.alpha", "0.4")),
+ Double.parseDouble(System.getProperty("ccd.alpha1", "0.4")),
+ Double.parseDouble(System.getProperty("ccd.alpha2", "0.05")));
+ List test = read(N, 0.5, 1.0);
+
+ long[] nodes = new long[4];
+ double[] logp = new double[4];
+ long class4Nodes = 0;
+ long survivingObserved = 0;
+ long shatteredObserved = 0;
+
+ for (Tree t : test) {
+ Map bits = new HashMap<>();
+ computeBits(t.getRoot(), bits, ccd.getSizeOfLeavesArray());
+ for (Node v : t.getNodesAsArray()) {
+ if (v.isLeaf()) {
+ continue;
+ }
+ BitSet cb = bits.get(v);
+ BitSet ab = bits.get(v.getChildren().get(0));
+ BitSet bb = bits.get(v.getChildren().get(1));
+ int cls = ccd.splitClass(cb, ab, bb);
+ nodes[cls]++;
+ logp[cls] += ccd.logSplitProbability(cb, ab, bb,
+ ccd.getAlpha(), ccd.getAlpha1(), ccd.getAlpha2());
+ if (cls == 3) {
+ class4Nodes++;
+ // how much observed structure did this split preserve vs destroy?
+ for (Clade obs : ccd.getClades()) {
+ if (obs.size() < 2 || obs.size() >= cb.cardinality()) {
+ continue;
+ }
+ BitSet o = obs.getCladeInBits();
+ BitSet tmp = BitSet.newBitSet(o);
+ tmp.andNot(cb);
+ if (!tmp.isEmpty()) {
+ continue; // not inside this clade at all
+ }
+ if (subset(o, ab) || subset(o, bb)) {
+ survivingObserved++;
+ } else {
+ shatteredObserved++;
+ }
+ }
+ }
+ }
+ }
+
+ long totalNodes = nodes[0] + nodes[1] + nodes[2] + nodes[3];
+ double totalLogp = logp[0] + logp[1] + logp[2] + logp[3];
+ System.out.printf("%n=== %s: class usage over %d held-out trees (%s) ===%n",
+ new File(PATH).getName(), test.size(), ccd);
+ System.out.printf("%-28s %10s %8s %14s %10s %12s%n",
+ "class", "nodes", "% nodes", "total logP", "% logP", "mean logP");
+ String[] names = {"1 observed split", "2 both children obs.",
+ "3 one child observed", "4 neither observed"};
+ for (int j = 0; j < 4; j++) {
+ System.out.printf("%-28s %10d %7.2f%% %14.1f %9.2f%% %12.3f%n",
+ names[j], nodes[j], 100.0 * nodes[j] / totalNodes, logp[j],
+ 100.0 * logp[j] / totalLogp, nodes[j] == 0 ? 0 : logp[j] / nodes[j]);
+ }
+ System.out.printf("total mean logP per tree = %.2f%n", totalLogp / test.size());
+ if (class4Nodes > 0) {
+ long tot = survivingObserved + shatteredObserved;
+ System.out.printf("class-4 splits: %d; observed clades below them: %d intact (%.1f%%), "
+ + "%d shattered (%.1f%%)%n",
+ class4Nodes, survivingObserved, 100.0 * survivingObserved / tot,
+ shatteredObserved, 100.0 * shatteredObserved / tot);
+ } else {
+ System.out.println("no class-4 splits occurred in any held-out tree");
+ }
+ }
+
+ private static boolean subset(BitSet a, BitSet c) {
+ BitSet tmp = BitSet.newBitSet(a);
+ tmp.andNot(c);
+ return tmp.isEmpty();
+ }
+
+ private static BitSet computeBits(Node v, Map bits, int leafArraySize) {
+ BitSet b = BitSet.newBitSet(leafArraySize);
+ if (v.isLeaf()) {
+ b.set(v.getNr());
+ } else {
+ b.or(computeBits(v.getChildren().get(0), bits, leafArraySize));
+ b.or(computeBits(v.getChildren().get(1), bits, leafArraySize));
+ }
+ bits.put(v, b);
+ return b;
+ }
+
+
+ /**
+ * Size profile of the splits that introduce two novel clades: parent size and the sizes of the
+ * two novel children. If the smaller side is consistently small, grading class 2 by the size of
+ * the smaller side would concentrate its mass where the real novelty is, instead of on the
+ * balanced splits that dominate a uniform draw.
+ */
+ @Test
+ public void novelSplitSizeProfile() throws Exception {
+ Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(),
+ "set -Dccd.trees=/path/to/x.trees");
+ CRegCCD ccd = new CRegCCD(read(N, 0.0, 0.5), 0.0,
+ Double.parseDouble(System.getProperty("ccd.alpha", "0.4")),
+ Double.parseDouble(System.getProperty("ccd.alpha1", "0.4")),
+ Double.parseDouble(System.getProperty("ccd.alpha2", "0.05")));
+ List test = read(N, 0.5, 1.0);
+
+ System.out.printf("%n=== %s: size profile of two-novel-clade splits ===%n",
+ new File(PATH).getName());
+ System.out.printf("%-8s %-10s %-10s %-14s %-16s%n",
+ "parent m", "small side", "large side", "small/parent", "uniform E[small]");
+ int count = 0;
+ double sumFrac = 0.0;
+ java.util.Map smallSizes = new java.util.TreeMap<>();
+ for (Tree t : test) {
+ Map bits = new HashMap<>();
+ computeBits(t.getRoot(), bits, ccd.getSizeOfLeavesArray());
+ for (Node v : t.getNodesAsArray()) {
+ if (v.isLeaf()) {
+ continue;
+ }
+ BitSet cb = bits.get(v);
+ BitSet ab = bits.get(v.getChildren().get(0));
+ BitSet bb = bits.get(v.getChildren().get(1));
+ if (ccd.splitClass(cb, ab, bb) != 3) {
+ continue;
+ }
+ int m = cb.cardinality();
+ int small = Math.min(ab.cardinality(), bb.cardinality());
+ int large = Math.max(ab.cardinality(), bb.cardinality());
+ count++;
+ sumFrac += small / (double) m;
+ smallSizes.merge(small, 1, Integer::sum);
+ if (count <= 40) {
+ System.out.printf("%-8d %-10d %-10d %-14.3f %-16.1f%n",
+ m, small, large, small / (double) m, m / 2.0);
+ }
+ }
+ }
+ if (count == 0) {
+ System.out.println("no two-novel-clade splits in the held-out set");
+ return;
+ }
+ System.out.printf("%d such splits; mean smaller-side fraction = %.3f "
+ + "(a uniform bipartition would give ~0.5)%n", count, sumFrac / count);
+ System.out.println("distribution of smaller-side size: " + smallSizes);
+ }
+}
diff --git a/src/test/java/ccd/model/KRegNormalisationTest.java b/src/test/java/ccd/model/KRegNormalisationTest.java
new file mode 100644
index 0000000..7043cc6
--- /dev/null
+++ b/src/test/java/ccd/model/KRegNormalisationTest.java
@@ -0,0 +1,172 @@
+package ccd.model;
+
+import beast.base.evolution.tree.Tree;
+import beast.base.evolution.tree.TreeParser;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Total probability mass of {@link KRegCCD}, by brute-force enumeration of every rooted topology.
+ *
+ * Documents the model's normalisation exactly as the manuscript states it:
+ *
+ * - on four taxa the model is exactly normalised (no blue-region boundary part is itself a
+ * reserving clade, so the region decomposition is tight);
+ * - from six taxa on it is sub-normalised -- total mass is below 1, never above --
+ * and the deficit scales as {@code mu^2}, which is the maximality deficit rather than the
+ * {@code O(mu^(k+1))} reserve truncation (it persists at full reserve depth).
+ *
+ */
+public class KRegNormalisationTest {
+
+ private static final List TAXA4 = Arrays.asList("A", "B", "C", "D");
+ private static final List TAXA5 = Arrays.asList("A", "B", "C", "D", "E");
+ private static final List TAXA6 = Arrays.asList("A", "B", "C", "D", "E", "F");
+ private static final List TAXA7 = Arrays.asList("A", "B", "C", "D", "E", "F", "G");
+
+ private static List trees(List taxa, String... newicks) {
+ List out = new ArrayList<>();
+ for (String nwk : newicks) {
+ out.add(new TreeParser(taxa, nwk, 1, false));
+ }
+ return out;
+ }
+
+ private static List training4() {
+ return trees(TAXA4, "(((A:1,B:1):1,C:1):1,D:1):0;", "((A:1,B:1):1,(C:1,D:1):1):0;");
+ }
+
+ private static List training6() {
+ return trees(TAXA6,
+ "(((((A:1,B:1):1,C:1):1,D:1):1,E:1):1,F:1):0;",
+ "((((D:1,C:1):1,B:1):1,A:1):1,(E:1,F:1):1):0;");
+ }
+
+ private static List training7() {
+ return trees(TAXA7,
+ "((((((A:1,B:1):1,C:1):1,D:1):1,E:1):1,F:1):1,G:1):0;",
+ "(((((D:1,C:1):1,B:1):1,A:1):1,(E:1,F:1):1):1,G:1):0;");
+ }
+
+ /** Total mass under the full-support score. */
+ private static double totalMass(KRegCCD ccd, List taxa) {
+ double mass = 0.0;
+ for (Tree t : allRootedTopologies(taxa)) {
+ mass += Math.exp(ccd.getLogProbabilityOfTree(t));
+ }
+ return mass;
+ }
+
+ @Test
+ public void exactlyNormalisedOnFourTaxa() {
+ for (double mu : new double[]{0.001, 0.005, 0.05}) {
+ KRegCCD ccd = new KRegCCD(training4(), 0.0, mu, 0.4, 2, KRegCCD.TailMode.NONE,
+ KRegCCD.NovelMode.FLAT);
+ double mass = totalMass(ccd, TAXA4);
+ System.out.printf("KRegCCD 4 taxa mu=%-6.3f totalMass = %.12f%n", mu, mass);
+ assertEquals(1.0, mass, 1e-9, "four-taxon model must normalise exactly");
+ }
+ }
+
+ /** Five taxa: exactness is not a taxon-count property but depends on whether any escape region
+ * has a reserving clade on its boundary, which the training set determines. */
+ @Test
+ public void fiveTaxaNormalisationDependsOnTrainingSet() {
+ List caterpillar = trees(TAXA5, "((((A:1,B:1):1,C:1):1,D:1):1,E:1):0;");
+ List mixed = trees(TAXA5,
+ "((((A:1,B:1):1,C:1):1,D:1):1,E:1):0;",
+ "(((A:1,B:1):1,(C:1,D:1):1):1,E:1):0;");
+ for (double mu : new double[]{0.005, 0.05}) {
+ for (String name : new String[]{"caterpillar", "mixed"}) {
+ List training = name.equals("caterpillar") ? caterpillar : mixed;
+ KRegCCD ccd = new KRegCCD(training, 0.0, mu, 0.4, 2, KRegCCD.TailMode.NONE,
+ KRegCCD.NovelMode.FLAT);
+ double mass = totalMass(ccd, TAXA5);
+ System.out.printf("KRegCCD 5 taxa mu=%-6.3f %-12s totalMass = %.9f (1-mass = %+.3e)%n",
+ mu, name, mass, 1.0 - mass);
+ assertTrue(mass <= 1.0 + 1e-12, "must never super-normalise, got " + mass);
+ }
+ }
+ }
+
+ @Test
+ public void subNormalisedFromSixTaxaWithMuSquaredDeficit() {
+ for (List taxa : List.of(TAXA6, TAXA7)) {
+ List training = (taxa == TAXA6) ? training6() : training7();
+ double prevDeficit = Double.NaN;
+ double prevMu = Double.NaN;
+ for (double mu : new double[]{0.05, 0.005, 0.001}) {
+ for (KRegCCD.TailMode tm : KRegCCD.TailMode.values()) {
+ KRegCCD ccd = new KRegCCD(training, 0.0, mu, 0.4, 2, tm,
+ KRegCCD.NovelMode.FLAT);
+ double mass = totalMass(ccd, taxa);
+ double deficit = 1.0 - mass;
+ System.out.printf("KRegCCD %d taxa mu=%-6.3f %-8s totalMass = %.9f (1-mass = %+.3e)%n",
+ taxa.size(), mu, tm, mass, deficit);
+ assertTrue(mass <= 1.0 + 1e-12,
+ "model must never super-normalise, got " + mass);
+ if (tm == KRegCCD.TailMode.NONE) {
+ if (!Double.isNaN(prevDeficit)) {
+ // a mu^2 deficit shrinks by the square of the mu ratio
+ double ratio = prevDeficit / deficit;
+ double expected = (prevMu / mu) * (prevMu / mu);
+ System.out.printf(" deficit shrank %.1fx as mu fell %.0fx "
+ + "(mu^2 predicts %.0fx)%n",
+ ratio, prevMu / mu, expected);
+ assertEquals(expected, ratio, 0.25 * expected,
+ "deficit must scale as mu^2");
+ }
+ prevDeficit = deficit;
+ prevMu = mu;
+ }
+ }
+ }
+ }
+ }
+
+ private static List allRootedTopologies(List taxa) {
+ List trees = new ArrayList<>();
+ for (String shape : shapes(taxa)) {
+ trees.add(new TreeParser(taxa, shape + ";", 1, false));
+ }
+ return trees;
+ }
+
+ private static List shapes(List taxa) {
+ List out = new ArrayList<>();
+ if (taxa.size() == 1) {
+ out.add(taxa.get(0) + ":1");
+ return out;
+ }
+ String first = taxa.get(0);
+ List rest = taxa.subList(1, taxa.size());
+ int n = rest.size();
+ for (int mask = 0; mask < (1 << n); mask++) {
+ List left = new ArrayList<>();
+ left.add(first);
+ List right = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ if ((mask & (1 << i)) != 0) {
+ left.add(rest.get(i));
+ } else {
+ right.add(rest.get(i));
+ }
+ }
+ if (right.isEmpty()) {
+ continue;
+ }
+ for (String l : shapes(left)) {
+ for (String r : shapes(right)) {
+ out.add("(" + l + "," + r + "):1");
+ }
+ }
+ }
+ return out;
+ }
+}
diff --git a/src/test/java/ccd/model/RealDataHeadToHeadTest.java b/src/test/java/ccd/model/RealDataHeadToHeadTest.java
new file mode 100644
index 0000000..df52b25
--- /dev/null
+++ b/src/test/java/ccd/model/RealDataHeadToHeadTest.java
@@ -0,0 +1,353 @@
+package ccd.model;
+
+import beast.base.evolution.tree.Tree;
+import beastfx.app.treeannotator.TreeAnnotator;
+import ccd.tools.CCDToolUtil;
+import org.junit.jupiter.api.Assumptions;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Head-to-head held-out predictive comparison of the regularised CCD variants on a real posterior
+ * tree sample.
+ *
+ * Protocol follows the manuscript's RSV2 comparison: the model is trained on trees drawn from the
+ * first half of the chain and scored on trees from the second half, so the test trees are genuinely
+ * out of training. Each model's hyperparameters are selected on an inner fit/validation split of the
+ * training half alone, then the model is rebuilt on the whole training half and scored on the test
+ * half. Reported per model: support coverage (fraction of test trees with positive probability), mean
+ * log probability over all test trees, and -- for the full-support models -- the paired per-tree
+ * comparison against KRegCCD.
+ *
+ *
Point the test at a tree file with {@code -Dccd.trees=/path/to/x.trees}; it is skipped when no
+ * file is given. {@code -Dccd.n=1000} sets the number of training and test trees.
+ */
+public class RealDataHeadToHeadTest {
+
+ private static final String PATH = System.getProperty("ccd.trees", "");
+ private static final int N = Integer.parseInt(System.getProperty("ccd.n", "1000"));
+ private static final double BURNIN_PERCENT = 10;
+
+ private interface Scorer {
+ double logP(Tree t);
+ }
+
+ private static List newickCache;
+ private static List taxaCache;
+
+ /** Parses the tree file once, keeping the topologies as newick strings. */
+ private static void load() throws Exception {
+ if (newickCache != null) {
+ return;
+ }
+ TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(PATH, (int) BURNIN_PERCENT);
+ ts.reset();
+ List nwk = new ArrayList<>();
+ List taxa = null;
+ while (ts.hasNext()) {
+ Tree t = ts.next();
+ if (taxa == null) {
+ taxa = new ArrayList<>();
+ String[] byNr = new String[t.getLeafNodeCount()];
+ for (beast.base.evolution.tree.Node leaf : t.getExternalNodes()) {
+ byNr[leaf.getNr()] = leaf.getID();
+ }
+ taxa.addAll(List.of(byNr));
+ }
+ nwk.add(t.getRoot().toNewick() + ";");
+ }
+ newickCache = nwk;
+ taxaCache = taxa;
+ System.out.printf("loaded %d trees, %d taxa from %s%n",
+ nwk.size(), taxa.size(), new File(PATH).getName());
+ }
+
+ /**
+ * {@code count} evenly spaced trees from the chain segment {@code [from, to)} (as fractions of the
+ * post-burn-in chain), freshly parsed on every call because the CCD constructors take ownership of
+ * the trees they are given.
+ *
+ * Segments must be disjoint for hyperparameter selection to be honest: scoring trees that are
+ * also in the fitted set drives every regularisation parameter to zero, because the backbone
+ * already fits its own training trees and any reserved mass is then pure loss.
+ */
+ private static List read(int count, double from, double to) throws Exception {
+ load();
+ int lo = (int) (from * newickCache.size());
+ int hi = (int) (to * newickCache.size());
+ List pool = newickCache.subList(lo, hi);
+ List out = new ArrayList<>();
+ double step = Math.max(1.0, pool.size() / (double) count);
+ for (int i = 0; i < count && (int) (i * step) < pool.size(); i++) {
+ out.add(new beast.base.evolution.tree.TreeParser(taxaCache, pool.get((int) (i * step)), 0, false));
+ }
+ return out;
+ }
+
+ /** Trees for the final models and for scoring: first half of the chain trains, second half tests. */
+ private static List train(int count) throws Exception {
+ return read(count, 0.0, 0.5);
+ }
+
+ /** Disjoint inner split of the training half: [0, 0.25) fits, [0.25, 0.5) validates. */
+ private static List fitSet(int count) throws Exception {
+ return read(count, 0.0, 0.25);
+ }
+
+ private static List valSet(int count) throws Exception {
+ return read(count, 0.25, 0.5);
+ }
+
+ private static double[] score(Scorer s, List test) {
+ int covered = 0;
+ double sum = 0.0;
+ for (Tree t : test) {
+ double lp = s.logP(t);
+ if (Double.isFinite(lp)) {
+ covered++;
+ sum += lp;
+ }
+ }
+ return new double[]{covered, covered == 0 ? Double.NEGATIVE_INFINITY : sum / covered};
+ }
+
+ /** Mean log probability over every test tree, treating unsupported trees as -infinity. */
+ private static double meanAll(Scorer s, List test) {
+ double sum = 0.0;
+ for (Tree t : test) {
+ double lp = s.logP(t);
+ if (!Double.isFinite(lp)) {
+ return Double.NEGATIVE_INFINITY;
+ }
+ sum += lp;
+ }
+ return sum / test.size();
+ }
+
+ @Test
+ public void headToHeadOnRealData() throws Exception {
+ Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(),
+ "set -Dccd.trees=/path/to/x.trees to run this comparison");
+
+ List probe = train(N);
+ int nTaxa = probe.get(0).getLeafNodeCount();
+ System.out.printf("%n=== %s: %d taxa, %d train / %d test trees ===%n",
+ new File(PATH).getName(), nTaxa, probe.size(), N);
+
+ List test = read(N, 0.5, 1.0);
+ List val = valSet(N / 2);
+ int nFit = N / 2;
+
+ // ---- CCD1 ----
+ CCD1 ccd1 = new CCD1(train(N), 0.0);
+ report("CCD1", "-", ccd1::getLogProbabilityOfTree, test);
+
+ // ---- RegCCD: alpha on validation ----
+ double bestAlpha = 0.4;
+ double bestAlphaScore = Double.NEGATIVE_INFINITY;
+ for (double alpha : new double[]{0.01, 0.05, 0.1, 0.2, 0.4, 0.8, 1.0}) {
+ RegCCD m = new RegCCD(fitSet(nFit), 0.0, alpha);
+ double sc = score(m::getLogProbabilityOfTree, val)[1];
+ if (sc > bestAlphaScore) {
+ bestAlphaScore = sc;
+ bestAlpha = alpha;
+ }
+ }
+ RegCCD reg = new RegCCD(train(N), 0.0, bestAlpha);
+ report("RegCCD", String.format("alpha=%.2f", bestAlpha), reg::getLogProbabilityOfTree, test);
+
+ // ---- KRegCCD: mu on validation at alpha = 0.4 (the manuscript's setting) ----
+ KRegCCD kFit = new KRegCCD(fitSet(nFit), 0.0, 0.005, 0.4);
+ double bestMu = 0.005;
+ double bestMuScore = Double.NEGATIVE_INFINITY;
+ for (double mu : new double[]{0.00002, 0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05}) {
+ final double m = mu;
+ double sc = score(t -> kFit.getLogProbabilityOfTree(t, m), val)[1];
+ if (sc > bestMuScore) {
+ bestMuScore = sc;
+ bestMu = mu;
+ }
+ }
+ KRegCCD kreg = new KRegCCD(train(N), 0.0, bestMu, 0.4);
+ report("KRegCCD", String.format("alpha=0.4, mu=%.5f", bestMu),
+ kreg::getLogProbabilityOfTree, test);
+
+ // ---- MRegCCD: mu on validation ----
+ MRegCCD mFit = new MRegCCD(fitSet(nFit), 0.0, MRegCCD.DEFAULT_MU);
+ double bestMMu = MRegCCD.DEFAULT_MU;
+ double bestMMuScore = Double.NEGATIVE_INFINITY;
+ for (double mu : new double[]{0.00005, 0.0002, 0.001, 0.002, 0.008, 0.0159, 0.05, 0.1}) {
+ final double m = mu;
+ double sc = score(t -> mFit.getLogProbabilityOfTree(t, m), val)[1];
+ if (sc > bestMMuScore) {
+ bestMMuScore = sc;
+ bestMMu = mu;
+ }
+ }
+ MRegCCD mreg = new MRegCCD(train(N), 0.0, bestMMu);
+ report("MRegCCD", String.format("mu=%.5f", bestMMu), mreg::getLogProbabilityOfTree, test);
+
+ // ---- CRegCCD: (a2, a3, a4) on validation, a1 = 0 ----
+ CRegCCD cFit = new CRegCCD(fitSet(nFit), 0.0);
+ double[] grid = {0.002, 0.01, 0.05, 0.2, 0.4, 1.0, 2.0, 5.0, 12.0};
+ double[] bestC = {0.0, 0.4, 0.4, 0.4};
+ double bestCScore = Double.NEGATIVE_INFINITY;
+ for (double b2 : grid) {
+ for (double b3 : grid) {
+ for (double b4 : grid) {
+ double sc = score(t -> cFit.getLogProbabilityOfTree(t, b2, b3, b4), val)[1];
+ if (sc > bestCScore) {
+ bestCScore = sc;
+ bestC = new double[]{0.0, b2, b3, b4};
+ }
+ }
+ }
+ }
+ CRegCCD creg = new CRegCCD(train(N), 0.0, bestC[1], bestC[2], bestC[3]);
+ report("CRegCCD", String.format("alpha=%.3f, alpha1=%.3f, alpha2=%.3f", bestC[1], bestC[2], bestC[3]),
+ creg::getLogProbabilityOfTree, test);
+
+ // ---- paired comparison of the full-support models against KRegCCD ----
+ System.out.printf("%npaired per-tree comparison against KRegCCD (n = %d test trees):%n", test.size());
+ paired("MRegCCD", mreg::getLogProbabilityOfTree, kreg::getLogProbabilityOfTree, test);
+ paired("CRegCCD", creg::getLogProbabilityOfTree, kreg::getLogProbabilityOfTree, test);
+ }
+
+ private static void report(String name, String params, Scorer s, List test) {
+ double[] sc = score(s, test);
+ double all = meanAll(s, test);
+ System.out.printf("%-9s %-27s coverage %6.1f%% mean logP(covered) %10.2f mean logP(all) %s%n",
+ name, params, 100.0 * sc[0] / test.size(), sc[1],
+ Double.isFinite(all) ? String.format("%10.2f", all) : " -inf");
+ }
+
+ private static void paired(String name, Scorer a, Scorer baseline, List test) {
+ int wins = 0;
+ double sumDiff = 0.0;
+ double sumSq = 0.0;
+ for (Tree t : test) {
+ double d = a.logP(t) - baseline.logP(t);
+ if (d > 0) {
+ wins++;
+ }
+ sumDiff += d;
+ sumSq += d * d;
+ }
+ int n = test.size();
+ double mean = sumDiff / n;
+ double se = Math.sqrt(Math.max(0, sumSq / n - mean * mean) / n);
+ System.out.printf(" %-9s mean log-ratio %+8.2f +/- %.2f nats/tree, better on %d/%d trees%n",
+ name, mean, se, wins, n);
+ }
+
+ /** Scale check: sampling from a real 129-taxon posterior must be fast, valid and
+ * self-consistent with the scorer. */
+ @Test
+ public void samplingOnRealData() throws Exception {
+ Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(),
+ "set -Dccd.trees=/path/to/x.trees to run this comparison");
+ CRegCCD ccd = new CRegCCD(train(N), 0.0);
+ int nTaxa = ccd.getSizeOfLeavesArray();
+ long t0 = System.nanoTime();
+ int draws = 200;
+ for (int i = 0; i < draws; i++) {
+ Tree t = ccd.sampleTree();
+ if (t.getLeafNodeCount() != nTaxa || t.getNodeCount() != 2 * nTaxa - 1) {
+ throw new AssertionError("invalid sampled tree");
+ }
+ double stamped = (Double) t.getRoot().getMetaData(CCD1.LOG_PROB_SUBTREE_KEY);
+ double scored = ccd.getLogProbabilityOfTree(t);
+ if (Math.abs(stamped - scored) > 1e-9) {
+ throw new AssertionError("stamped " + stamped + " != scored " + scored);
+ }
+ }
+ double secs = (System.nanoTime() - t0) / 1e9;
+ System.out.printf("CRegCCD sampling on %s: %d taxa, %d trees in %.2f s (%.1f ms/tree), "
+ + "all valid and self-consistent%n",
+ new File(PATH).getName(), nTaxa, draws, secs, 1000 * secs / draws);
+ }
+
+ /** MAP and entropy on a real posterior: correctness certificate and wall-clock cost. */
+ @Test
+ public void mapAndEntropyOnRealData() throws Exception {
+ Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(),
+ "set -Dccd.trees=/path/to/x.trees to run this comparison");
+ CRegCCD creg = new CRegCCD(train(N), 0.0, 2.0, 0.4, 0.05);
+
+ long t0 = System.nanoTime();
+ double maxLog = creg.getMaxLogTreeProbability();
+ boolean certified = creg.isMAPCertifiedGlobal();
+ double bound = creg.getOffBackboneBound();
+ double mapSecs = (System.nanoTime() - t0) / 1e9;
+
+ t0 = System.nanoTime();
+ Tree map = creg.getMAPTree();
+ double treeSecs = (System.nanoTime() - t0) / 1e9;
+ double scored = creg.getLogProbabilityOfTree(map);
+
+ t0 = System.nanoTime();
+ double[] h = creg.getEntropyMonteCarlo(20_000);
+ double entSecs = (System.nanoTime() - t0) / 1e9;
+
+ System.out.printf("%n=== %s: CRegCCD MAP and entropy ===%n", new File(PATH).getName());
+ System.out.printf("MAP DP + certificate : %.2f s, max logP = %.4f, "
+ + "off-backbone bound = %.4f, certified global = %s%n",
+ mapSecs, maxLog, bound, certified);
+ System.out.printf("MAP tree traceback : %.2f s, scored logP = %.4f (matches: %s)%n",
+ treeSecs, scored, Math.abs(scored - maxLog) < 1e-9);
+ System.out.printf("entropy (20k draws) : %.2f s, H = %.3f +/- %.3f nats%n",
+ entSecs, h[0], h[1]);
+ }
+
+ /** Deterministic entropy recursion vs the unbiased Monte-Carlo estimator on a real posterior. */
+ @Test
+ public void entropyRecursionVersusMonteCarloOnRealData() throws Exception {
+ Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(),
+ "set -Dccd.trees=/path/to/x.trees to run this comparison");
+ System.out.printf("%n=== %s: CRegCCD entropy, recursion vs Monte Carlo ===%n",
+ new File(PATH).getName());
+ System.out.printf("%-22s %-11s %-9s %-22s %-9s %-10s%n",
+ "pseudocounts", "recursion", "rec (s)", "Monte Carlo", "MC (s)", "difference");
+ for (double[] p : new double[][]{{2.0, 0.4, 0.05}, {5.0, 2.0, 0.4}, {0.4, 0.4, 0.4}}) {
+ CRegCCD ccd = new CRegCCD(train(N), 0.0, p[0], p[1], p[2]);
+ long t0 = System.nanoTime();
+ double rec = ccd.getEntropyRecursive();
+ double recSecs = (System.nanoTime() - t0) / 1e9;
+ t0 = System.nanoTime();
+ double[] mc = ccd.getEntropyMonteCarlo(200_000);
+ double mcSecs = (System.nanoTime() - t0) / 1e9;
+ double diff = rec - mc[0];
+ System.out.printf("a=(%.2f,%.2f,%.2f)%-6s %-11.4f %-9.2f %8.4f +/-%.4f %-9.2f %+.4f (%+.3f%%)%n",
+ p[0], p[1], p[2], "", rec, recSecs, mc[0], mc[1], mcSecs, diff, 100 * diff / mc[0]);
+ }
+ }
+
+ /** Does the exact A_0/A_1 MAP search stay tractable on a real posterior? */
+ @Test
+ public void exactMapOnRealData() throws Exception {
+ Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists(),
+ "set -Dccd.trees=/path/to/x.trees to run this comparison");
+ CRegCCD creg = new CRegCCD(train(N), 0.0, 0.4, 0.4, 0.05);
+ double backbone = creg.getMaxLogTreeProbability();
+ System.out.printf("%n=== %s: MAP search by allowed A_1 depth ===%n", new File(PATH).getName());
+ System.out.printf("backbone (A_0 only) max logP = %.4f%n", backbone);
+ System.out.printf("%-6s %-12s %-12s %-12s %-10s %-8s%n",
+ "maxA1", "max logP", "improvement", "bound", "certified", "states");
+ for (int k = 0; k <= 3; k++) {
+ long t0 = System.nanoTime();
+ CRegCCD.MapResult r = creg.solveMAP(k);
+ double secs = (System.nanoTime() - t0) / 1e9;
+ if (!r.complete()) {
+ System.out.printf("%-6d exceeded the state budget after %d states (%.1f s)%n",
+ k, r.statesExplored(), secs);
+ break;
+ }
+ System.out.printf("%-6d %-12.4f %-12.4f %-12.4f %-10s %-8d (%.1f s)%n",
+ k, r.maxLogProbability(), r.maxLogProbability() - backbone,
+ r.offBackboneBound(), r.a2Excluded(), r.statesExplored(), secs);
+ }
+ }
+}
diff --git a/src/test/java/ccd/model/SplitClassSizeAnalysis.java b/src/test/java/ccd/model/SplitClassSizeAnalysis.java
new file mode 100644
index 0000000..6ba4e2b
--- /dev/null
+++ b/src/test/java/ccd/model/SplitClassSizeAnalysis.java
@@ -0,0 +1,110 @@
+package ccd.model;
+
+import beast.base.evolution.tree.Tree;
+import beast.base.evolution.tree.TreeParser;
+import ccd.model.bitsets.BitSet;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
+
+/**
+ * Exploratory: sizes of the four split classes at a clade, for the class-based smoothing proposal.
+ *
+ * At a clade {@code C} of {@code m} taxa every one of the {@code 2^(m-1) - 1} bipartitions falls
+ * into exactly one of: (1) observed split; (2) unobserved, both children observed clades
+ * (the CCD0 expansion); (3) unobserved, exactly one child observed; (4) unobserved, neither child
+ * observed. Classes 1-3 are at most polynomial in the number of observed clades; class 4 is
+ * essentially all of {@code 2^(m-1)}. This prints the four sizes and the probability the smoothed
+ * model would put on class 1, under a per-split constant pseudocount versus a per-class total.
+ */
+public class SplitClassSizeAnalysis {
+
+ private static List randomTrees(int nTaxa, int nTrees, long seed) {
+ List taxa = new ArrayList<>();
+ for (int i = 0; i < nTaxa; i++) {
+ taxa.add("T" + i);
+ }
+ Random rng = new Random(seed);
+ List out = new ArrayList<>();
+ for (int t = 0; t < nTrees; t++) {
+ List pool = new ArrayList<>(taxa);
+ while (pool.size() > 1) {
+ int i = rng.nextInt(pool.size());
+ String a = pool.remove(i);
+ int j = rng.nextInt(pool.size());
+ String b = pool.remove(j);
+ pool.add("(" + a + "," + b + ")");
+ }
+ out.add(new TreeParser(taxa, pool.get(0) + ";", 1, false));
+ }
+ return out;
+ }
+
+ @Test
+ public void classSizesAtRoot() {
+ System.out.printf("%-6s %-7s %-8s %-8s %-8s %-14s %-14s %-14s%n",
+ "taxa", "trees", "|A1|", "|A2|", "|A3|", "|A4|", "P(A1) per-split", "P(A1) per-class");
+ for (int nTaxa : new int[]{12, 20, 30, 40}) {
+ int nTrees = 1000;
+ List trees = randomTrees(nTaxa, nTrees, 42L);
+ CCD0 ccd0 = new CCD0(trees, 0);
+ Clade root = null;
+ for (Clade c : ccd0.getClades()) {
+ if (c.size() == nTaxa) {
+ root = c;
+ }
+ }
+ if (root == null) {
+ continue;
+ }
+
+ int a1 = 0;
+ int a2 = 0;
+ for (CladePartition p : root.getPartitions()) {
+ if (p.getNumberOfOccurrences() > 0) {
+ a1++;
+ } else {
+ a2++;
+ }
+ }
+
+ // |A3|: observed proper subclades whose complement within root is NOT an observed clade
+ BitSet rootBits = root.getCladeInBits();
+ int a3 = 0;
+ for (Clade d : ccd0.getClades()) {
+ if (d.size() >= root.size()) {
+ continue;
+ }
+ BitSet db = d.getCladeInBits();
+ BitSet tmp = (BitSet) db.clone();
+ tmp.and(rootBits);
+ if (!tmp.equals(db)) {
+ continue; // not a subclade of root
+ }
+ BitSet comp = (BitSet) rootBits.clone();
+ comp.andNot(db);
+ if (ccd0.getClade(comp) == null) {
+ a3++;
+ }
+ }
+
+ double total = Math.pow(2, nTaxa - 1) - 1;
+ double a4 = total - a1 - a2 - a3;
+
+ // per-split constant pseudocount alpha on every class
+ double alpha = 0.4;
+ double fC = nTrees;
+ double denomSplit = fC + alpha * (a1 + a2 + a3 + a4);
+ double pA1Split = (fC + alpha * a1) / denomSplit;
+
+ // per-class total pseudocount alpha (spread within each class)
+ double denomClass = fC + alpha * 4;
+ double pA1Class = (fC + alpha) / denomClass;
+
+ System.out.printf("%-6d %-7d %-8d %-8d %-8d %-14.4g %-14.6g %-14.6g%n",
+ nTaxa, nTrees, a1, a2, a3, a4, pA1Split, pA1Class);
+ }
+ }
+}