From 658360af009ac4f65bf7cc34708933c1d265d36d Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Wed, 19 Aug 2026 11:41:19 +1200 Subject: [PATCH 1/2] Fix undersized BitSet copies above 256 bits newBitSet(BitSet) sized the copy with other.length() -- the index of the highest set bit plus one -- where it needed other.size(), the capacity. A set whose top words happen to be empty therefore came back with a shorter word array. That is not only a crash. and/or/xor/andNot/intersects all iterate over this.words.length and index the operand directly, so an undersized operand throws ArrayIndexOutOfBoundsException while an undersized *receiver* silently drops the high words instead. It bites only above 256 bits, where the generic BitSet replaces the fixed-size subclasses, and only when the top words are empty -- the normal case for a clade that does not contain the highest-numbered taxa. It surfaced as a crash in MRegCCD on a 276-taxon analysis and cannot affect analyses of 256 taxa or fewer. BitSetCopyTest covers 276/320/512/1000 bits and the empty set, exercising both directions of andNot and intersects. Co-Authored-By: Claude Opus 5 (1M context) --- src/main/java/ccd/model/bitsets/BitSet.java | 5 +- .../ccd/model/bitsets/BitSetCopyTest.java | 54 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 src/test/java/ccd/model/bitsets/BitSetCopyTest.java diff --git a/src/main/java/ccd/model/bitsets/BitSet.java b/src/main/java/ccd/model/bitsets/BitSet.java index a6b6c37..3c727fe 100644 --- a/src/main/java/ccd/model/bitsets/BitSet.java +++ b/src/main/java/ccd/model/bitsets/BitSet.java @@ -55,7 +55,10 @@ public static BitSet newBitSet(BitSet other) { if (other instanceof BitSet256 set) { return new BitSet256(set); } - BitSet b = new BitSet(other.length()); + // size(), not length(): length() is the index of the highest set bit plus one, so copying a + // set whose top words happen to be empty would return an undersized BitSet, and the bitwise + // operations below index the operand by this.words.length. + BitSet b = new BitSet(other.size()); b.or(other); return b; } diff --git a/src/test/java/ccd/model/bitsets/BitSetCopyTest.java b/src/test/java/ccd/model/bitsets/BitSetCopyTest.java new file mode 100644 index 0000000..34944f8 --- /dev/null +++ b/src/test/java/ccd/model/bitsets/BitSetCopyTest.java @@ -0,0 +1,54 @@ +package ccd.model.bitsets; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Copies of a BitSet must keep the source's capacity, not shrink to its highest set bit. + * + *

The bitwise operations iterate over {@code this.words.length} and index the operand directly, + * so an undersized copy makes them throw. This bites only above 256 bits, where the generic BitSet + * is used instead of the fixed-size subclasses, and only when the top words are empty -- the common + * case for a clade that does not contain the highest-numbered taxa. + */ +public class BitSetCopyTest { + + @Test + public void copyKeepsCapacityAboveTheSpecialisedSizes() { + for (int nbits : new int[]{276, 320, 512, 1000}) { + BitSet full = BitSet.newBitSet(nbits); + full.set(nbits - 1); + + BitSet sparse = BitSet.newBitSet(nbits); + sparse.set(3); + BitSet copy = BitSet.newBitSet(sparse); + + assertEquals(sparse.size(), copy.size(), + "copy must keep the source capacity at " + nbits + " bits"); + + BitSet a = BitSet.newBitSet(full); + a.andNot(copy); + assertTrue(a.get(nbits - 1), "andNot must not clear unrelated high bits"); + + BitSet b = BitSet.newBitSet(copy); + b.andNot(full); + assertTrue(b.get(3), "andNot must keep the low bit"); + assertFalse(copy.intersects(full), "disjoint sets must not intersect"); + } + } + + @Test + public void copyOfAnEmptySetIsUsable() { + BitSet empty = BitSet.newBitSet(400); + BitSet copy = BitSet.newBitSet(empty); + assertEquals(empty.size(), copy.size(), "an empty copy must still have capacity"); + BitSet other = BitSet.newBitSet(400); + other.set(399); + copy.andNot(other); + copy.or(other); + assertTrue(copy.get(399)); + } +} From b8439636d459b682e4255974f6422e3f5a002ec2 Mon Sep 17 00:00:00 2001 From: Alexei Drummond Date: Wed, 19 Aug 2026 11:41:38 +1200 Subject: [PATCH 2/2] Compute MRegCCD boundary counts by indexing, and default to boundary 4 MRegCCD found its boundaries by walking every ordered choice of parts, costing O(m^(k-1)) in the number m of observed subclades of a clade, and guarded that with a flat 20M op budget. On large analyses the budget bit before the enumeration finished, and countsFor then caught the overflow and left the remaining orders at zero -- which also silently disabled the tail correction, since tailFor reads the top two orders. Nothing was reported to the caller. Measured on RSV2 (129 taxa) the shipped defaults returned -60.6219 where the untruncated answer is -60.6443. The counts are now found by indexing the disjoint pairs of observed subclades by the bitset they cover, once per clade in O(m^2): a boundary of 2 is a subclade whose complement is observed, of 3 a subclade whose complement is covered by a pair, of 4 a pair whose complement is covered by another pair. This is how KRegCCD computes its own reserve, which counts the same partitions indexed by novel-clade count rather than boundary size. Orders beyond 4, and depths below 4, defer to the direct enumeration, which is cheaper there. The former implementation is kept as MRegCCDSlow, and MRegCCD extends it overriding countsFor alone, so the model is defined in exactly one place and the two can be checked against each other. MRegCCDAgreementTest confirms identical boundary counts on 2199 clades and identical tree probabilities across 6/8/10/14 taxa and depths 2 to 4, and checks the exact-normalisation guarantee directly on the class callers get. Both now default to reserve depth 4, matching KRegCCD's default reserve (k = 2, boundaries 3 and 4), so the two models look equally far past the CCD graph. The previous default of 5 went a boundary further than KRegCCD while the op budget stopped it from ever getting there. On RSV2 at the new default this is 4.3x faster than the direct enumeration (1034ms against 4396ms to score 200 trees) and returns the untruncated value. MRegCCDParameterOptimiser needs no change: it names MRegCCD, which is now the indexed implementation. Co-Authored-By: Claude Opus 5 (1M context) --- src/main/java/ccd/model/MRegCCD.java | 688 +++--------------- src/main/java/ccd/model/MRegCCDSlow.java | 656 +++++++++++++++++ .../java/ccd/model/MRegCCDAgreementTest.java | 133 ++++ src/test/java/ccd/model/MRegCCDTest.java | 28 +- .../java/ccd/model/MRegDepthTimingTest.java | 58 ++ 5 files changed, 960 insertions(+), 603 deletions(-) create mode 100644 src/main/java/ccd/model/MRegCCDSlow.java create mode 100644 src/test/java/ccd/model/MRegCCDAgreementTest.java create mode 100644 src/test/java/ccd/model/MRegDepthTimingTest.java 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: *

- * - *

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): + *

+ * + *

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/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()); + } + } + } +}