diff --git a/src/main/java/ccd/model/MRegCCD.java b/src/main/java/ccd/model/MRegCCD.java
index 2f11b7a..f553922 100644
--- a/src/main/java/ccd/model/MRegCCD.java
+++ b/src/main/java/ccd/model/MRegCCD.java
@@ -1,6 +1,5 @@
package ccd.model;
-import beast.base.evolution.tree.Node;
import beast.base.evolution.tree.Tree;
import beastfx.app.treeannotator.TreeAnnotator.TreeSet;
import ccd.model.bitsets.BitSet;
@@ -9,462 +8,144 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
/**
- * MRegCCD -- the one-parameter "per-new-split" regularised CCD. It unifies RegCCD's split-expansion
- * {@code alpha} and KRegCCD's escape {@code mu} into a single per-clade escape rate, giving a
- * full-support tree distribution with one hyperparameter {@code mu} (and no {@code alpha}).
+ * The one-parameter "per-new-split" regularised CCD, with the boundary counts computed the way
+ * {@link KRegCCD} computes its reserve rather than by direct recursive enumeration.
*
- * The model is a plain {@link CCD1} backbone (raw conditional clade probabilities, no smoothing)
- * extended with a per-clade escape reserve. The distribution is defined conditionally, clade by clade
- * (chain rule over the observed-clade DAG; no global partition function):
+ *
The model is unchanged: this overrides only {@code countsFor}, so every probability, sample and
+ * point estimate is defined exactly as in {@link MRegCCDSlow}, which this overrides in a single
+ * method. The two must therefore agree wherever the reference implementation's op budget does not
+ * truncate its enumeration, which is what {@code MRegCCDAgreementTest} checks.
+ *
+ *
The reference implementation walks every ordered choice of boundary parts, which costs {@code O(m^(k-1))} in
+ * the number {@code m} of observed subclades of a clade and blows through a flat op budget on large
+ * analyses -- silently, since {@code countsFor} catches the overflow and leaves the remaining orders
+ * at zero, which also disables the tail correction that reads them. Here the boundaries are found by
+ * indexing instead:
*
- * - An observed split {@code C -> {L, R}} (seen in training) is priced
- * {@code (1 - mu - tail(C)) * ccp(L|R)} when {@code C} can escape, else just {@code ccp(L|R)}.
- * - An escape at {@code C} resolves it through novel intermediate clades down to a
- * boundary of {@code m} observed subclades (a maximal "blue" region). Such a region has
- * {@code m - 1} new splits and is priced {@code eps(C)^(m-1)} -- one factor of {@code eps} per
- * new split (so a recombination of two observed subclades, {@code m = 2}, costs one
- * {@code eps}). This is the difference from KRegCCD, which makes recombinations representable in
- * its {@code alpha}-expanded backbone and charges {@code eps} only per novel clade
- * ({@code eps^(m-2)}).
+ * - the disjoint pairs of observed subclades are enumerated once, {@code O(m^2)}, and grouped by
+ * the bitset they cover;
+ * - a boundary of 2 is an observed subclade whose complement in {@code C} is observed;
+ * - a boundary of 3 is an observed subclade whose complement is covered by a pair;
+ * - a boundary of 4 is a pair whose complement is covered by another pair.
*
- *
- * The per-clade escape rate {@code eps(C)} is the root of {@code sum_{m>=2} M_m(C) eps^(m-1) = mu},
- * where {@code M_m(C)} counts the all-novel resolutions of {@code C} with an {@code m}-part boundary
- * (the FLAT weighting: each distinct novel resolution counted once). Computing the full sum is
- * #P-hard, so -- mirroring KRegCCD -- the orders {@code m = 2..reserveDepth} are enumerated exactly
- * (bounded by an op-budget) and the omitted higher orders are a geometric tail correction added to
- * {@code mu} (so observed splits are discounted by {@code 1 - mu - tail}, keeping the conditional
- * properly normalised; truncating without the tail super-normalises). A clade with no escape route
- * ({@code M_m = 0} for all computed {@code m}) is not reservable and keeps its raw CCP undiscounted.
- *
- *
Every tree on the taxon set has positive probability (full support), so {@link #containsTree}
- * is always true and {@link #getLogProbabilityOfTree} is finite for all trees.
- *
- * @author Claude (CCD-Sophie)
+ * Each lookup is a hash probe rather than a search, so orders 2 to 4 cost {@code O(m^2)} in total.
+ * Orders beyond 4 fall back to the inherited enumeration, so this is a strict speed-up of the
+ * practical depths and never changes what is computed.
*/
-public class MRegCCD extends CCD1 {
-
- /** Default per-clade escape probability (the RSV2 operating point of the conditional model). */
- public static final double DEFAULT_MU = 0.0159;
-
- /** Default reserve depth: enumerate boundary sizes {@code m = 2..DEFAULT_RESERVE_DEPTH} exactly. */
- public static final int DEFAULT_RESERVE_DEPTH = 5;
-
- /** Per-clade enumeration-op budget (mirrors KRegCCD's; bounds the boundary enumeration). */
- private static final long OPS_BUDGET = Long.parseLong(System.getProperty("mreg.enumOps", "20000000"));
-
- private static final class BudgetExceeded extends RuntimeException {
- BudgetExceeded() {
- super(null, null, false, false);
- }
- }
-
- private static final BudgetExceeded BUDGET_EXCEEDED = new BudgetExceeded();
+public class MRegCCD extends MRegCCDSlow {
- /** Per-clade escape probability (the single hyperparameter). */
- private final double mu;
-
- /** Max boundary size enumerated when solving eps; deeper orders are a geometric tail. */
- private final int reserveDepth;
-
- /** Whether the {@code (1 - mu - tail)} discount carries the geometric tail correction. */
- private final boolean useTail;
-
- /** Observed-clade bitsets (incl. leaves), sorted canonically; built lazily. */
- private List sortedCladeBits;
- private final Map> subCache = new HashMap<>();
- private final Map countsCache = new HashMap<>();
- private long enumOps;
+ private final Map fastCounts = new ConcurrentHashMap<>();
public MRegCCD(List trees, double burnin, double mu) {
- this(trees, burnin, mu, DEFAULT_RESERVE_DEPTH, true);
+ super(trees, burnin, mu);
}
public MRegCCD(List trees, double burnin, double mu, int reserveDepth, boolean useTail) {
- super(trees, burnin);
- validate(mu, reserveDepth);
- this.mu = mu;
- this.reserveDepth = reserveDepth;
- this.useTail = useTail;
+ super(trees, burnin, mu, reserveDepth, useTail);
}
public MRegCCD(TreeSet treeSet, double mu) {
- this(treeSet, mu, DEFAULT_RESERVE_DEPTH, true);
+ super(treeSet, mu);
}
public MRegCCD(TreeSet treeSet, double mu, int reserveDepth, boolean useTail) {
- super(treeSet);
- validate(mu, reserveDepth);
- this.mu = mu;
- this.reserveDepth = reserveDepth;
- this.useTail = useTail;
- }
-
- /**
- * Builds an MRegCCD on {@code trees} with {@code mu} selected by maximising cross-validated
- * held-out log-probability (see {@link ccd.algorithms.regularisation.MRegCCDParameterOptimiser}),
- * rather than the fixed {@link #DEFAULT_MU}. The honest, no-peeking counterpart of
- * {@code KRegCCD.withOptimisedParameters}.
- */
- public static MRegCCD withOptimisedMu(List trees) {
- double mu = ccd.algorithms.regularisation.MRegCCDParameterOptimiser.optimiseMu(trees).mu();
- return new MRegCCD(trees, 0.0, mu);
- }
-
- private static void validate(double mu, int reserveDepth) {
- if (mu <= 0 || mu >= 1) {
- throw new IllegalArgumentException("mu must be in (0, 1), got " + mu);
- }
- if (reserveDepth < 2) {
- throw new IllegalArgumentException("reserveDepth must be >= 2, got " + reserveDepth);
- }
- }
-
- /** The per-clade escape probability this model was built with. */
- public double getMu() {
- return mu;
- }
-
- public int getReserveDepth() {
- return reserveDepth;
- }
-
- /**
- * Reserve counts {@code M_m(C)} by boundary size {@code m} (array index {@code m}, valid for
- * {@code m = 2..min(|C|, reserveDepth)}); {@code M_m} is the number of all-novel resolutions of
- * {@code C} with an {@code m}-part boundary. The first coefficient {@code M_2} (the {@code eps^1}
- * term) is exactly the number of CCD0-expanded splits of {@code C} -- recombinations of two
- * observed subclades whose split was never observed -- since those are the only escapes with no
- * other novel (blue) clade. Exposed for inspection and cross-checks.
- */
- public int[] reserveCounts(BitSet cladeInBits) {
- return countsFor(cladeInBits).clone();
- }
-
- @Override
- public String toString() {
- return "MRegCCD [mu = " + mu + ", reserveDepth = " + reserveDepth + ", tail = " + useTail
- + ", per-new-split, full support]";
- }
-
- /* ----------------------------------------------------------------------
- * Scoring
- * ------------------------------------------------------------------- */
-
- @Override
- public double getLogProbabilityOfTree(Tree tree) {
- return scoreTree(tree, mu);
- }
-
- /**
- * Full-support log-probability at an arbitrary escape probability {@code scoreMu}, reusing this
- * model's ({@code mu}-independent) backbone and cached reserve counts. Lets a parameter search /
- * cross-validation evaluate many {@code mu} on one trained model without rebuilding. For
- * {@code scoreMu == mu} it equals {@link #getLogProbabilityOfTree(Tree)}.
- */
- public double getLogProbabilityOfTree(Tree tree, double scoreMu) {
- if (scoreMu <= 0 || scoreMu >= 1) {
- throw new IllegalArgumentException("scoreMu must be in (0, 1), got " + scoreMu);
- }
- return scoreTree(tree, scoreMu);
+ super(treeSet, mu, reserveDepth, useTail);
}
@Override
- public double getProbabilityOfTree(Tree tree) {
- return Math.exp(getLogProbabilityOfTree(tree));
- }
-
- /** Always true: MRegCCD is full support, so every tree on this taxon set has positive probability. */
- @Override
- public boolean containsTree(Tree tree) {
- return true;
- }
-
- private double scoreTree(Tree tree, double scoreMu) {
- Map bits = new HashMap<>();
- computeBits(tree.getRoot(), bits);
- double logp = 0.0;
- for (Node v : tree.getNodesAsArray()) {
- if (v.isLeaf()) {
- continue;
- }
- BitSet vb = bits.get(v);
- Clade c = getClade(vb);
- if (c == null) {
- continue; // novel clade: scored once at its maximal region's top
- }
- BitSet b1 = bits.get(v.getChildren().get(0));
- BitSet b2 = bits.get(v.getChildren().get(1));
- if (isSplitObserved(vb, b1, b2)) {
- if (reservable(vb)) { // discount only clades that can actually escape
- double resv = Math.min(scoreMu + (useTail ? tailFor(vb, scoreMu) : 0.0), 1 - 1e-12);
- logp += Math.log(1.0 - resv);
- }
- logp += rawLogCCP(c, b1, b2); // raw CCD1 CCP
- } else {
- // region top: an observed clade resolved through a novel split. m-1 new splits.
- int m = boundarySize(v, bits);
- logp += (m - 1) * Math.log(epsFor(vb, scoreMu));
- }
- }
- return logp;
- }
-
- /* ----------------------------------------------------------------------
- * Per-clade reserve (M_m counts -> eps, tail; mirrors KRegCCD.computeReg)
- * ------------------------------------------------------------------- */
-
- /** Whether clade {@code C} (given in bits) reserves any escape mass up to {@code reserveDepth}. */
- boolean reservable(BitSet C) {
- for (int v : countsFor(C)) {
- if (v > 0) {
- return true;
- }
- }
- return false;
- }
-
- /** Escape root {@code eps} solving {@code sum_{m>=2} M_m eps^(m-1) = scoreMu} (monotone bisection). */
- double epsFor(BitSet C, double scoreMu) {
- int[] n = countsFor(C);
- if (!reservable(C)) {
- return scoreMu; // crude fallback (no escape route within reserveDepth); should not be hit
- }
- return solveEps(n, scoreMu);
- }
-
- /** Omitted-tail escape mass beyond the computed orders: geometric bound from the top two orders. */
- double tailFor(BitSet C, double scoreMu) {
- int[] n = countsFor(C);
- int last = n.length - 1;
- if (last < 3) {
- return 0.0;
- }
- int nLast = n[last], nPrev = n[last - 1];
- if (nLast <= 0 || nPrev <= 0) {
- return 0.0;
- }
- double eps = epsFor(C, scoreMu);
- double rho = ((double) nLast / nPrev) * eps;
- if (rho <= 0 || rho >= 1) {
- return 0.0;
- }
- return Math.min(nLast * Math.pow(eps, last - 1) * rho / (1 - rho), scoreMu);
- }
-
- /** M_m counts (index m = boundary size, 2..min(|C|, reserveDepth)); cached, mu-independent. */
int[] countsFor(BitSet C) {
- int[] cached = countsCache.get(C);
+ int[] cached = fastCounts.get(C);
if (cached != null) {
return cached;
}
int card = C.cardinality();
- int[] n = new int[Math.min(card, reserveDepth) + 1];
- if (card >= 2) {
+ int depth = Math.min(card, getReserveDepth());
+ // The pair index only pays for itself once order 4 needs it; below that the inherited
+ // enumeration is cheaper, and identical by construction.
+ if (depth < 4) {
+ return super.countsFor(C);
+ }
+ int[] n = new int[depth + 1];
+ if (card >= 2 && depth >= 2) {
List subs = subclades(C);
- enumOps = 0;
- for (int m = 2; m < n.length; m++) {
- try {
- n[m] = countBoundaries(C, subs, m);
- } catch (BudgetExceeded e) {
- break; // deeper orders omitted (negligible, like the tail)
- }
- }
- }
- countsCache.put(C, n);
- return n;
- }
-
- private static double solveEps(int[] n, double mu) {
- double lo = 0.0, hi = 1.0;
- while (evalReserve(n, hi) < mu) {
- hi *= 2.0;
- }
- for (int it = 0; it < 100; it++) {
- double mid = 0.5 * (lo + hi);
- if (evalReserve(n, mid) < mu) {
- lo = mid;
- } else {
- hi = mid;
- }
- }
- return 0.5 * (lo + hi);
- }
-
- /** {@code sum_{m>=2} n[m] x^(m-1)}. */
- private static double evalReserve(int[] n, double x) {
- double s = 0.0;
- for (int m = 2; m < n.length; m++) {
- if (n[m] > 0) {
- s += n[m] * Math.pow(x, m - 1);
- }
- }
- return s;
- }
-
- /** Count m-part boundaries of C into observed subclades, weighted by their all-novel pathcount. */
- private int countBoundaries(BitSet C, List subs, int m) {
- return enumerateBoundaries(C, subs, m, 0, BitSet.newBitSet(leafArraySize), new ArrayList<>(m));
- }
- private int enumerateBoundaries(BitSet C, List subs, int m, int startIdx,
- BitSet used, List chosen) {
- if (++enumOps > OPS_BUDGET) {
- throw BUDGET_EXCEEDED;
- }
- if (chosen.size() == m - 1) {
- BitSet last = BitSet.newBitSet(C);
- last.andNot(used);
- if (last.isEmpty() || !isObs(last)) {
- return 0;
- }
- if (compareBitSets(chosen.get(chosen.size() - 1), last) >= 0) {
- return 0; // canonical: the derived last part must be the largest
- }
- BitSet[] parts = new BitSet[m];
- for (int i = 0; i < m - 1; i++) {
- parts[i] = chosen.get(i);
- }
- parts[m - 1] = last;
- return countAllNovelResolutions(C, parts);
- }
- int count = 0;
- for (int i = startIdx; i < subs.size(); i++) {
- BitSet pb = subs.get(i);
- if (pb.intersects(used)) {
- continue;
+ // Every disjoint pair of observed subclades, grouped by the bitset it covers. Only
+ // orders 3 and 4 consult this, so at depth 2 the index is not worth building: the O(m^2)
+ // pass would cost more than the order-2 scan it would serve.
+ Map> pairsByUnion = new HashMap<>();
+ for (int i = 0; i < subs.size(); i++) {
+ BitSet a = subs.get(i);
+ for (int j = i + 1; j < subs.size(); j++) {
+ BitSet b = subs.get(j);
+ if (a.intersects(b)) {
+ continue;
+ }
+ BitSet union = BitSet.newBitSet(a);
+ union.or(b);
+ pairsByUnion.computeIfAbsent(union, k -> new ArrayList<>())
+ .add(new BitSet[]{a, b});
+ }
}
- chosen.add(pb);
- BitSet newUsed = BitSet.newBitSet(used);
- newUsed.or(pb);
- count += enumerateBoundaries(C, subs, m, i + 1, newUsed, chosen);
- chosen.remove(chosen.size() - 1);
- }
- return count;
- }
- /**
- * Number of all-novel binary resolutions of C into the given observed parts (subset DP over the
- * parts). A split is allowed iff: at the region root (full mask = C, an observed clade) the split
- * is unobserved (a real escape); at an intermediate node the clade itself is novel (a maximal
- * region stops at observed clades, matching {@link #boundarySize}).
- */
- private int countAllNovelResolutions(BitSet C, BitSet[] parts) {
- int k = parts.length;
- if (k == 1) {
- return 1;
- }
- int full = (1 << k) - 1;
- BitSet[] unionOf = new BitSet[1 << k];
- unionOf[0] = BitSet.newBitSet(leafArraySize);
- for (int mask = 1; mask <= full; mask++) {
- int low = Integer.numberOfTrailingZeros(mask);
- BitSet u = BitSet.newBitSet(unionOf[mask & (mask - 1)]);
- u.or(parts[low]);
- unionOf[mask] = u;
- }
- int[] f = new int[1 << k];
- for (int mask = 1; mask <= full; mask++) {
- if (Integer.bitCount(mask) == 1) {
- f[mask] = 1;
- continue;
- }
- int low = mask & (-mask), rest = mask ^ low, count = 0;
- for (int sub = rest; ; sub = (sub - 1) & rest) {
- int s1 = sub | low, s2 = mask ^ s1;
- if (s2 != 0 && splitAllowed(mask == full, unionOf[mask], unionOf[s1], unionOf[s2])) {
- count += f[s1] * f[s2];
+ for (BitSet d : subs) {
+ BitSet rest = BitSet.newBitSet(C);
+ rest.andNot(d);
+ if (rest.isEmpty()) {
+ continue;
}
- if (sub == 0) {
- break;
+ // boundary 2: {d, rest}, counted once from its canonically smaller side
+ if (depth >= 2 && isObs(rest) && compareBitSets(d, rest) < 0) {
+ n[2] += countAllNovelResolutions(C, new BitSet[]{d, rest});
}
- }
- f[mask] = count;
- }
- return f[full];
- }
-
- private boolean splitAllowed(boolean top, BitSet union, BitSet a, BitSet b) {
- return top ? !isSplitObserved(union, a, b) : !isObs(union);
- }
-
- /* ----------------------------------------------------------------------
- * Observed-backbone queries (over the inherited CCD1 clade DAG)
- * ------------------------------------------------------------------- */
-
- private boolean isObs(BitSet x) {
- return getClade(x) != null; // leaves are clades too
- }
-
- private boolean isSplitObserved(BitSet parentBits, BitSet aBits, BitSet bBits) {
- Clade parent = getClade(parentBits);
- if (parent == null) {
- return false;
- }
- Clade a = getClade(aBits);
- Clade b = getClade(bBits);
- if (a == null || b == null) {
- return false;
- }
- return parent.getCladePartition(a, b) != null;
- }
-
- private double rawLogCCP(Clade parent, BitSet aBits, BitSet bBits) {
- CladePartition p = parent.getCladePartition(getClade(aBits), getClade(bBits));
- return p.getLogCCP();
- }
-
- /** Observed clades (incl. leaves) strictly contained in C, in canonical order; cached. */
- private List subclades(BitSet C) {
- return subCache.computeIfAbsent(C, c -> {
- List out = new ArrayList<>();
- int card = c.cardinality();
- for (BitSet x : sortedCladeBits()) {
- if (x.cardinality() < card && subset(x, c)) {
- out.add(x);
+ // boundary 3: {d} plus a pair covering the remainder
+ if (depth >= 3) {
+ for (BitSet[] p : pairsByUnion.getOrDefault(rest, List.of())) {
+ if (compareBitSets(d, p[0]) < 0) { // d must be the canonically first part
+ n[3] += countAllNovelResolutions(C, new BitSet[]{d, p[0], p[1]});
+ }
+ }
}
}
- return out;
- });
- }
- private List sortedCladeBits() {
- if (sortedCladeBits == null) {
- List all = new ArrayList<>();
- for (Clade c : getClades()) {
- all.add(c.getCladeInBits());
+ // boundary 4: a pair whose complement is covered by another pair
+ if (depth >= 4) {
+ for (Map.Entry> e : pairsByUnion.entrySet()) {
+ BitSet rest = BitSet.newBitSet(C);
+ rest.andNot(e.getKey());
+ if (rest.isEmpty() || !subset(e.getKey(), C)) {
+ continue;
+ }
+ List others = pairsByUnion.get(rest);
+ if (others == null) {
+ continue;
+ }
+ for (BitSet[] p : e.getValue()) {
+ for (BitSet[] q : others) {
+ // A 4-part boundary splits into two pairs in three ways, so count only
+ // the one whose first pair holds the two canonically smallest parts:
+ // for parts w < x < y < z that is {w,x}|{y,z} and no other.
+ if (compareBitSets(p[1], q[0]) < 0) {
+ n[4] += countAllNovelResolutions(C,
+ new BitSet[]{p[0], p[1], q[0], q[1]});
+ }
+ }
+ }
+ }
}
- all.sort(MRegCCD::compareBitSets);
- sortedCladeBits = all;
- }
- return sortedCladeBits;
- }
- /** Boundary size of the maximal region rooted at v: count of maximal observed/leaf subclades below. */
- private int boundarySize(Node v, Map bits) {
- int m = 0;
- for (Node child : v.getChildren()) {
- if (child.isLeaf() || getClade(bits.get(child)) != null) {
- m++;
- } else {
- m += boundarySize(child, bits);
+ // orders beyond 4 are rare in practice; defer to the inherited enumeration
+ if (depth >= 5) {
+ int[] slow = super.countsFor(C);
+ for (int m = 5; m < n.length && m < slow.length; m++) {
+ n[m] = slow[m];
+ }
}
}
- return m;
- }
-
- private BitSet computeBits(Node v, Map bits) {
- BitSet b = BitSet.newBitSet(leafArraySize);
- if (v.isLeaf()) {
- b.set(v.getNr());
- } else {
- b.or(computeBits(v.getChildren().get(0), bits));
- b.or(computeBits(v.getChildren().get(1), bits));
- }
- bits.put(v, b);
- return b;
+ fastCounts.put(BitSet.newBitSet(C), n);
+ return n;
}
private static boolean subset(BitSet a, BitSet c) {
@@ -472,175 +153,4 @@ private static boolean subset(BitSet a, BitSet c) {
tmp.andNot(c);
return tmp.isEmpty();
}
-
- /** Canonical total order on clade bitsets (lexicographic by set-bit indices). */
- private static int compareBitSets(BitSet a, BitSet b) {
- int ia = a.nextSetBit(0), ib = b.nextSetBit(0);
- while (ia >= 0 && ib >= 0) {
- if (ia != ib) {
- return Integer.compare(ia, ib);
- }
- ia = a.nextSetBit(ia + 1);
- ib = b.nextSetBit(ib + 1);
- }
- return Integer.compare(ia, ib);
- }
-
- /* ----------------------------------------------------------------------
- * Sampling (self-consistent)
- *
- * The PIT calibration test draws trees from the model and needs only each draw's log-probability
- * (not the tree object), so we override sampleTreeLogProbability() with a direct simulation of the
- * generative process and never materialise a Tree. At each reservable clade we escape with
- * probability equal to its escape mass (= mu by the eps-solve, tail EXCLUDED) and otherwise take
- * an observed (red) split ~ CCP; an escape draws a region order m proportional to M_m eps^(m-1), a
- * boundary of m observed subclades proportional to its all-novel pathcount, and recurses into the
- * boundary parts. The resolution shape within a region is not drawn -- every shape has the same
- * weight eps^(m-1) and does not change the draw's log-probability -- so the simulation is cheap.
- *
- * The draw distribution exactly matches getLogProbabilityOfTree when the model is built with the
- * tail OFF (then the red discount is 1 - mu, matching the escape mass), for trees whose regions are
- * within reserveDepth; deeper regions (mass ~mu^reserveDepth) are never produced, the same
- * self-consistent / full-support trade-off KRegCCD makes for its PIT.
- * ------------------------------------------------------------------- */
-
- @Override
- public double sampleTreeLogProbability() {
- return simulate(getRootClade());
- }
-
- private double simulate(Clade c) {
- if (c.isLeaf()) {
- return 0.0;
- }
- BitSet cb = c.getCladeInBits();
- if (reservable(cb)) {
- double eps = epsFor(cb, mu);
- int[] n = countsFor(cb);
- double escapeMass = 0.0;
- for (int m = 2; m < n.length; m++) {
- if (n[m] > 0) {
- escapeMass += n[m] * Math.pow(eps, m - 1);
- }
- }
- if (random.nextDouble() < escapeMass) {
- int m = sampleOrder(n, eps, escapeMass);
- BitSet[] parts = sampleBoundaryParts(cb, subclades(cb), m);
- double logp = (m - 1) * Math.log(eps);
- if (parts != null) {
- for (BitSet bp : parts) {
- logp += simulate(getClade(bp));
- }
- }
- return logp;
- }
- CladePartition p = samplePartition(c);
- double logp = Math.log(1.0 - escapeMass) + p.getLogCCP();
- return logp + simulate(p.getChildClades()[0]) + simulate(p.getChildClades()[1]);
- }
- CladePartition p = samplePartition(c); // non-reservable: observed split, no discount
- return p.getLogCCP() + simulate(p.getChildClades()[0]) + simulate(p.getChildClades()[1]);
- }
-
- /** Draws a region order m in {2..} with probability proportional to {@code M_m eps^(m-1)}. */
- private int sampleOrder(int[] n, double eps, double escapeMass) {
- double target = random.nextDouble() * escapeMass, acc = 0.0;
- for (int m = 2; m < n.length; m++) {
- if (n[m] > 0) {
- acc += n[m] * Math.pow(eps, m - 1);
- if (target < acc) {
- return m;
- }
- }
- }
- for (int m = n.length - 1; m >= 2; m--) {
- if (n[m] > 0) {
- return m; // numerical guard
- }
- }
- return 2;
- }
-
- /** Samples an observed (red) split of {@code c} with probability proportional to its CCP. */
- private CladePartition samplePartition(Clade c) {
- List partitions = c.getPartitions();
- double target = random.nextDouble(), acc = 0.0;
- for (CladePartition p : partitions) {
- acc += p.getCCP();
- if (target < acc) {
- return p;
- }
- }
- return partitions.get(partitions.size() - 1);
- }
-
- /**
- * Weighted-reservoir samples one boundary of {@code c} into {@code m} observed subclades,
- * proportional to its all-novel pathcount (so that, combined with order sampling, every distinct
- * novel resolution is equiprobable at {@code eps^(m-1)}). Returns the parts, or {@code null} if
- * none/op-budget.
- */
- private BitSet[] sampleBoundaryParts(BitSet c, List subs, int m) {
- boundaryPick = null;
- boundaryWeightSeen = 0.0;
- enumOps = 0;
- try {
- sampleBoundaryWalk(c, subs, m, 0, BitSet.newBitSet(leafArraySize), new ArrayList<>(m));
- } catch (BudgetExceeded e) {
- return boundaryPick; // whatever was picked before the cap (may be null)
- }
- return boundaryPick;
- }
-
- private BitSet[] boundaryPick;
- private double boundaryWeightSeen;
-
- private void sampleBoundaryWalk(BitSet c, List subs, int m, int startIdx,
- BitSet used, List chosen) {
- if (++enumOps > OPS_BUDGET) {
- throw BUDGET_EXCEEDED;
- }
- if (chosen.size() == m - 1) {
- BitSet last = BitSet.newBitSet(c);
- last.andNot(used);
- if (last.isEmpty() || !isObs(last)) {
- return;
- }
- if (compareBitSets(chosen.get(chosen.size() - 1), last) >= 0) {
- return;
- }
- BitSet[] parts = new BitSet[m];
- for (int i = 0; i < m - 1; i++) {
- parts[i] = chosen.get(i);
- }
- parts[m - 1] = last;
- int pc = countAllNovelResolutions(c, parts);
- if (pc <= 0) {
- return;
- }
- boundaryWeightSeen += pc;
- if (random.nextDouble() * boundaryWeightSeen < pc) { // weighted reservoir
- boundaryPick = parts;
- }
- return;
- }
- for (int i = startIdx; i < subs.size(); i++) {
- BitSet pb = subs.get(i);
- if (pb.intersects(used)) {
- continue;
- }
- chosen.add(pb);
- BitSet newUsed = BitSet.newBitSet(used);
- newUsed.or(pb);
- sampleBoundaryWalk(c, subs, m, i + 1, newUsed, chosen);
- chosen.remove(chosen.size() - 1);
- }
- }
-
- @Override
- public Tree sampleTree(HeightSettingStrategy heightStrategy) {
- throw new UnsupportedOperationException(
- "MRegCCD materialised-tree sampling is not implemented; sampleTreeLogProbability() "
- + "(used by the PIT) simulates draws without building trees.");
- }
}
diff --git a/src/main/java/ccd/model/MRegCCDSlow.java b/src/main/java/ccd/model/MRegCCDSlow.java
new file mode 100644
index 0000000..ace1207
--- /dev/null
+++ b/src/main/java/ccd/model/MRegCCDSlow.java
@@ -0,0 +1,656 @@
+package ccd.model;
+
+import beast.base.evolution.tree.Node;
+import beast.base.evolution.tree.Tree;
+import beastfx.app.treeannotator.TreeAnnotator.TreeSet;
+import ccd.model.bitsets.BitSet;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * The reference implementation of {@link MRegCCD}: identical model, boundary counts obtained by
+ * direct recursive enumeration rather than by indexing. Retained so that the faster implementation
+ * can be checked against it (see {@code MRegCCDAgreementTest}); prefer {@link MRegCCD} in use, which
+ * computes the same counts without the enumeration blow-up at boundary 4.
+ *
+ * MRegCCDSlow -- the one-parameter "per-new-split" regularised CCD. It unifies RegCCD's split-expansion
+ * {@code alpha} and KRegCCD's escape {@code mu} into a single per-clade escape rate, giving a
+ * full-support tree distribution with one hyperparameter {@code mu} (and no {@code alpha}).
+ *
+ *
The model is a plain {@link CCD1} backbone (raw conditional clade probabilities, no smoothing)
+ * extended with a per-clade escape reserve. The distribution is defined conditionally, clade by clade
+ * (chain rule over the observed-clade DAG; no global partition function):
+ *
+ * - An observed split {@code C -> {L, R}} (seen in training) is priced
+ * {@code (1 - mu - tail(C)) * ccp(L|R)} when {@code C} can escape, else just {@code ccp(L|R)}.
+ * - An escape at {@code C} resolves it through novel intermediate clades down to a
+ * boundary of {@code m} observed subclades (a maximal "blue" region). Such a region has
+ * {@code m - 1} new splits and is priced {@code eps(C)^(m-1)} -- one factor of {@code eps} per
+ * new split (so a recombination of two observed subclades, {@code m = 2}, costs one
+ * {@code eps}). This is the difference from KRegCCD, which makes recombinations representable in
+ * its {@code alpha}-expanded backbone and charges {@code eps} only per novel clade
+ * ({@code eps^(m-2)}).
+ *
+ *
+ * The per-clade escape rate {@code eps(C)} is the root of {@code sum_{m>=2} M_m(C) eps^(m-1) = mu},
+ * where {@code M_m(C)} counts the all-novel resolutions of {@code C} with an {@code m}-part boundary
+ * (the FLAT weighting: each distinct novel resolution counted once). Computing the full sum is
+ * #P-hard, so -- mirroring KRegCCD -- the orders {@code m = 2..reserveDepth} are enumerated exactly
+ * (bounded by an op-budget) and the omitted higher orders are a geometric tail correction added to
+ * {@code mu} (so observed splits are discounted by {@code 1 - mu - tail}, keeping the conditional
+ * properly normalised; truncating without the tail super-normalises). A clade with no escape route
+ * ({@code M_m = 0} for all computed {@code m}) is not reservable and keeps its raw CCP undiscounted.
+ *
+ *
Every tree on the taxon set has positive probability (full support), so {@link #containsTree}
+ * is always true and {@link #getLogProbabilityOfTree} is finite for all trees.
+ *
+ * @author Claude (CCD-Sophie)
+ */
+public class MRegCCDSlow extends CCD1 {
+
+ /** Default per-clade escape probability (the RSV2 operating point of the conditional model). */
+ public static final double DEFAULT_MU = 0.0159;
+
+ /**
+ * Default reserve depth: enumerate boundary sizes {@code m = 2..DEFAULT_RESERVE_DEPTH} exactly.
+ * Boundary 4 matches {@link KRegCCD}'s default reserve ({@code k = 2}, boundaries 3 and 4), so
+ * the two models look equally far past the CCD graph. It was 5, which went a boundary further
+ * than KRegCCD while the op budget below silently truncated the enumeration before reaching it.
+ */
+ public static final int DEFAULT_RESERVE_DEPTH = 4;
+
+ /** Per-clade enumeration-op budget (mirrors KRegCCD's; bounds the boundary enumeration). */
+ private static final long OPS_BUDGET = Long.parseLong(System.getProperty("mreg.enumOps", "20000000"));
+
+ private static final class BudgetExceeded extends RuntimeException {
+ BudgetExceeded() {
+ super(null, null, false, false);
+ }
+ }
+
+ private static final BudgetExceeded BUDGET_EXCEEDED = new BudgetExceeded();
+
+ /** Per-clade escape probability (the single hyperparameter). */
+ private final double mu;
+
+ /** Max boundary size enumerated when solving eps; deeper orders are a geometric tail. */
+ private final int reserveDepth;
+
+ /** Whether the {@code (1 - mu - tail)} discount carries the geometric tail correction. */
+ private final boolean useTail;
+
+ /** Observed-clade bitsets (incl. leaves), sorted canonically; built lazily. */
+ private List sortedCladeBits;
+ private final Map> subCache = new HashMap<>();
+ private final Map countsCache = new HashMap<>();
+ private long enumOps;
+
+ public MRegCCDSlow(List trees, double burnin, double mu) {
+ this(trees, burnin, mu, DEFAULT_RESERVE_DEPTH, true);
+ }
+
+ public MRegCCDSlow(List trees, double burnin, double mu, int reserveDepth, boolean useTail) {
+ super(trees, burnin);
+ validate(mu, reserveDepth);
+ this.mu = mu;
+ this.reserveDepth = reserveDepth;
+ this.useTail = useTail;
+ }
+
+ public MRegCCDSlow(TreeSet treeSet, double mu) {
+ this(treeSet, mu, DEFAULT_RESERVE_DEPTH, true);
+ }
+
+ public MRegCCDSlow(TreeSet treeSet, double mu, int reserveDepth, boolean useTail) {
+ super(treeSet);
+ validate(mu, reserveDepth);
+ this.mu = mu;
+ this.reserveDepth = reserveDepth;
+ this.useTail = useTail;
+ }
+
+ /**
+ * Builds an MRegCCDSlow on {@code trees} with {@code mu} selected by maximising cross-validated
+ * held-out log-probability (see {@link ccd.algorithms.regularisation.MRegCCDParameterOptimiser}),
+ * rather than the fixed {@link #DEFAULT_MU}. The honest, no-peeking counterpart of
+ * {@code KRegCCD.withOptimisedParameters}.
+ */
+ public static MRegCCDSlow withOptimisedMu(List trees) {
+ double mu = ccd.algorithms.regularisation.MRegCCDParameterOptimiser.optimiseMu(trees).mu();
+ return new MRegCCDSlow(trees, 0.0, mu);
+ }
+
+ private static void validate(double mu, int reserveDepth) {
+ if (mu <= 0 || mu >= 1) {
+ throw new IllegalArgumentException("mu must be in (0, 1), got " + mu);
+ }
+ if (reserveDepth < 2) {
+ throw new IllegalArgumentException("reserveDepth must be >= 2, got " + reserveDepth);
+ }
+ }
+
+ /** The per-clade escape probability this model was built with. */
+ public double getMu() {
+ return mu;
+ }
+
+ public int getReserveDepth() {
+ return reserveDepth;
+ }
+
+ /**
+ * Reserve counts {@code M_m(C)} by boundary size {@code m} (array index {@code m}, valid for
+ * {@code m = 2..min(|C|, reserveDepth)}); {@code M_m} is the number of all-novel resolutions of
+ * {@code C} with an {@code m}-part boundary. The first coefficient {@code M_2} (the {@code eps^1}
+ * term) is exactly the number of CCD0-expanded splits of {@code C} -- recombinations of two
+ * observed subclades whose split was never observed -- since those are the only escapes with no
+ * other novel (blue) clade. Exposed for inspection and cross-checks.
+ */
+ public int[] reserveCounts(BitSet cladeInBits) {
+ return countsFor(cladeInBits).clone();
+ }
+
+ @Override
+ public String toString() {
+ return "MRegCCDSlow [mu = " + mu + ", reserveDepth = " + reserveDepth + ", tail = " + useTail
+ + ", per-new-split, full support]";
+ }
+
+ /* ----------------------------------------------------------------------
+ * Scoring
+ * ------------------------------------------------------------------- */
+
+ @Override
+ public double getLogProbabilityOfTree(Tree tree) {
+ return scoreTree(tree, mu);
+ }
+
+ /**
+ * Full-support log-probability at an arbitrary escape probability {@code scoreMu}, reusing this
+ * model's ({@code mu}-independent) backbone and cached reserve counts. Lets a parameter search /
+ * cross-validation evaluate many {@code mu} on one trained model without rebuilding. For
+ * {@code scoreMu == mu} it equals {@link #getLogProbabilityOfTree(Tree)}.
+ */
+ public double getLogProbabilityOfTree(Tree tree, double scoreMu) {
+ if (scoreMu <= 0 || scoreMu >= 1) {
+ throw new IllegalArgumentException("scoreMu must be in (0, 1), got " + scoreMu);
+ }
+ return scoreTree(tree, scoreMu);
+ }
+
+ @Override
+ public double getProbabilityOfTree(Tree tree) {
+ return Math.exp(getLogProbabilityOfTree(tree));
+ }
+
+ /** Always true: MRegCCDSlow is full support, so every tree on this taxon set has positive probability. */
+ @Override
+ public boolean containsTree(Tree tree) {
+ return true;
+ }
+
+ private double scoreTree(Tree tree, double scoreMu) {
+ Map bits = new HashMap<>();
+ computeBits(tree.getRoot(), bits);
+ double logp = 0.0;
+ for (Node v : tree.getNodesAsArray()) {
+ if (v.isLeaf()) {
+ continue;
+ }
+ BitSet vb = bits.get(v);
+ Clade c = getClade(vb);
+ if (c == null) {
+ continue; // novel clade: scored once at its maximal region's top
+ }
+ BitSet b1 = bits.get(v.getChildren().get(0));
+ BitSet b2 = bits.get(v.getChildren().get(1));
+ if (isSplitObserved(vb, b1, b2)) {
+ if (reservable(vb)) { // discount only clades that can actually escape
+ double resv = Math.min(scoreMu + (useTail ? tailFor(vb, scoreMu) : 0.0), 1 - 1e-12);
+ logp += Math.log(1.0 - resv);
+ }
+ logp += rawLogCCP(c, b1, b2); // raw CCD1 CCP
+ } else {
+ // region top: an observed clade resolved through a novel split. m-1 new splits.
+ int m = boundarySize(v, bits);
+ logp += (m - 1) * Math.log(epsFor(vb, scoreMu));
+ }
+ }
+ return logp;
+ }
+
+ /* ----------------------------------------------------------------------
+ * Per-clade reserve (M_m counts -> eps, tail; mirrors KRegCCD.computeReg)
+ * ------------------------------------------------------------------- */
+
+ /** Whether clade {@code C} (given in bits) reserves any escape mass up to {@code reserveDepth}. */
+ boolean reservable(BitSet C) {
+ for (int v : countsFor(C)) {
+ if (v > 0) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /** Escape root {@code eps} solving {@code sum_{m>=2} M_m eps^(m-1) = scoreMu} (monotone bisection). */
+ double epsFor(BitSet C, double scoreMu) {
+ int[] n = countsFor(C);
+ if (!reservable(C)) {
+ return scoreMu; // crude fallback (no escape route within reserveDepth); should not be hit
+ }
+ return solveEps(n, scoreMu);
+ }
+
+ /** Omitted-tail escape mass beyond the computed orders: geometric bound from the top two orders. */
+ double tailFor(BitSet C, double scoreMu) {
+ int[] n = countsFor(C);
+ int last = n.length - 1;
+ if (last < 3) {
+ return 0.0;
+ }
+ int nLast = n[last], nPrev = n[last - 1];
+ if (nLast <= 0 || nPrev <= 0) {
+ return 0.0;
+ }
+ double eps = epsFor(C, scoreMu);
+ double rho = ((double) nLast / nPrev) * eps;
+ if (rho <= 0 || rho >= 1) {
+ return 0.0;
+ }
+ return Math.min(nLast * Math.pow(eps, last - 1) * rho / (1 - rho), scoreMu);
+ }
+
+ /** M_m counts (index m = boundary size, 2..min(|C|, reserveDepth)); cached, mu-independent. */
+ int[] countsFor(BitSet C) {
+ int[] cached = countsCache.get(C);
+ if (cached != null) {
+ return cached;
+ }
+ int card = C.cardinality();
+ int[] n = new int[Math.min(card, reserveDepth) + 1];
+ if (card >= 2) {
+ List subs = subclades(C);
+ enumOps = 0;
+ for (int m = 2; m < n.length; m++) {
+ try {
+ n[m] = countBoundaries(C, subs, m);
+ } catch (BudgetExceeded e) {
+ break; // deeper orders omitted (negligible, like the tail)
+ }
+ }
+ }
+ countsCache.put(C, n);
+ return n;
+ }
+
+ private static double solveEps(int[] n, double mu) {
+ double lo = 0.0, hi = 1.0;
+ while (evalReserve(n, hi) < mu) {
+ hi *= 2.0;
+ }
+ for (int it = 0; it < 100; it++) {
+ double mid = 0.5 * (lo + hi);
+ if (evalReserve(n, mid) < mu) {
+ lo = mid;
+ } else {
+ hi = mid;
+ }
+ }
+ return 0.5 * (lo + hi);
+ }
+
+ /** {@code sum_{m>=2} n[m] x^(m-1)}. */
+ private static double evalReserve(int[] n, double x) {
+ double s = 0.0;
+ for (int m = 2; m < n.length; m++) {
+ if (n[m] > 0) {
+ s += n[m] * Math.pow(x, m - 1);
+ }
+ }
+ return s;
+ }
+
+ /** Count m-part boundaries of C into observed subclades, weighted by their all-novel pathcount. */
+ private int countBoundaries(BitSet C, List subs, int m) {
+ return enumerateBoundaries(C, subs, m, 0, BitSet.newBitSet(leafArraySize), new ArrayList<>(m));
+ }
+
+ private int enumerateBoundaries(BitSet C, List subs, int m, int startIdx,
+ BitSet used, List chosen) {
+ if (++enumOps > OPS_BUDGET) {
+ throw BUDGET_EXCEEDED;
+ }
+ if (chosen.size() == m - 1) {
+ BitSet last = BitSet.newBitSet(C);
+ last.andNot(used);
+ if (last.isEmpty() || !isObs(last)) {
+ return 0;
+ }
+ if (compareBitSets(chosen.get(chosen.size() - 1), last) >= 0) {
+ return 0; // canonical: the derived last part must be the largest
+ }
+ BitSet[] parts = new BitSet[m];
+ for (int i = 0; i < m - 1; i++) {
+ parts[i] = chosen.get(i);
+ }
+ parts[m - 1] = last;
+ return countAllNovelResolutions(C, parts);
+ }
+ int count = 0;
+ for (int i = startIdx; i < subs.size(); i++) {
+ BitSet pb = subs.get(i);
+ if (pb.intersects(used)) {
+ continue;
+ }
+ chosen.add(pb);
+ BitSet newUsed = BitSet.newBitSet(used);
+ newUsed.or(pb);
+ count += enumerateBoundaries(C, subs, m, i + 1, newUsed, chosen);
+ chosen.remove(chosen.size() - 1);
+ }
+ return count;
+ }
+
+ /**
+ * Number of all-novel binary resolutions of C into the given observed parts (subset DP over the
+ * parts). A split is allowed iff: at the region root (full mask = C, an observed clade) the split
+ * is unobserved (a real escape); at an intermediate node the clade itself is novel (a maximal
+ * region stops at observed clades, matching {@link #boundarySize}).
+ */
+ int countAllNovelResolutions(BitSet C, BitSet[] parts) {
+ int k = parts.length;
+ if (k == 1) {
+ return 1;
+ }
+ int full = (1 << k) - 1;
+ BitSet[] unionOf = new BitSet[1 << k];
+ unionOf[0] = BitSet.newBitSet(leafArraySize);
+ for (int mask = 1; mask <= full; mask++) {
+ int low = Integer.numberOfTrailingZeros(mask);
+ BitSet u = BitSet.newBitSet(unionOf[mask & (mask - 1)]);
+ u.or(parts[low]);
+ unionOf[mask] = u;
+ }
+ int[] f = new int[1 << k];
+ for (int mask = 1; mask <= full; mask++) {
+ if (Integer.bitCount(mask) == 1) {
+ f[mask] = 1;
+ continue;
+ }
+ int low = mask & (-mask), rest = mask ^ low, count = 0;
+ for (int sub = rest; ; sub = (sub - 1) & rest) {
+ int s1 = sub | low, s2 = mask ^ s1;
+ if (s2 != 0 && splitAllowed(mask == full, unionOf[mask], unionOf[s1], unionOf[s2])) {
+ count += f[s1] * f[s2];
+ }
+ if (sub == 0) {
+ break;
+ }
+ }
+ f[mask] = count;
+ }
+ return f[full];
+ }
+
+ private boolean splitAllowed(boolean top, BitSet union, BitSet a, BitSet b) {
+ return top ? !isSplitObserved(union, a, b) : !isObs(union);
+ }
+
+ /* ----------------------------------------------------------------------
+ * Observed-backbone queries (over the inherited CCD1 clade DAG)
+ * ------------------------------------------------------------------- */
+
+ boolean isObs(BitSet x) {
+ return getClade(x) != null; // leaves are clades too
+ }
+
+ private boolean isSplitObserved(BitSet parentBits, BitSet aBits, BitSet bBits) {
+ Clade parent = getClade(parentBits);
+ if (parent == null) {
+ return false;
+ }
+ Clade a = getClade(aBits);
+ Clade b = getClade(bBits);
+ if (a == null || b == null) {
+ return false;
+ }
+ return parent.getCladePartition(a, b) != null;
+ }
+
+ private double rawLogCCP(Clade parent, BitSet aBits, BitSet bBits) {
+ CladePartition p = parent.getCladePartition(getClade(aBits), getClade(bBits));
+ return p.getLogCCP();
+ }
+
+ /** Observed clades (incl. leaves) strictly contained in C, in canonical order; cached. */
+ List subclades(BitSet C) {
+ return subCache.computeIfAbsent(C, c -> {
+ List out = new ArrayList<>();
+ int card = c.cardinality();
+ for (BitSet x : sortedCladeBits()) {
+ if (x.cardinality() < card && subset(x, c)) {
+ out.add(x);
+ }
+ }
+ return out;
+ });
+ }
+
+ private List sortedCladeBits() {
+ if (sortedCladeBits == null) {
+ List all = new ArrayList<>();
+ for (Clade c : getClades()) {
+ all.add(c.getCladeInBits());
+ }
+ all.sort(MRegCCDSlow::compareBitSets);
+ sortedCladeBits = all;
+ }
+ return sortedCladeBits;
+ }
+
+ /** Boundary size of the maximal region rooted at v: count of maximal observed/leaf subclades below. */
+ private int boundarySize(Node v, Map bits) {
+ int m = 0;
+ for (Node child : v.getChildren()) {
+ if (child.isLeaf() || getClade(bits.get(child)) != null) {
+ m++;
+ } else {
+ m += boundarySize(child, bits);
+ }
+ }
+ return m;
+ }
+
+ private BitSet computeBits(Node v, Map bits) {
+ BitSet b = BitSet.newBitSet(leafArraySize);
+ if (v.isLeaf()) {
+ b.set(v.getNr());
+ } else {
+ b.or(computeBits(v.getChildren().get(0), bits));
+ b.or(computeBits(v.getChildren().get(1), bits));
+ }
+ bits.put(v, b);
+ return b;
+ }
+
+ private static boolean subset(BitSet a, BitSet c) {
+ BitSet tmp = BitSet.newBitSet(a);
+ tmp.andNot(c);
+ return tmp.isEmpty();
+ }
+
+ /** Canonical total order on clade bitsets (lexicographic by set-bit indices). */
+ static int compareBitSets(BitSet a, BitSet b) {
+ int ia = a.nextSetBit(0), ib = b.nextSetBit(0);
+ while (ia >= 0 && ib >= 0) {
+ if (ia != ib) {
+ return Integer.compare(ia, ib);
+ }
+ ia = a.nextSetBit(ia + 1);
+ ib = b.nextSetBit(ib + 1);
+ }
+ return Integer.compare(ia, ib);
+ }
+
+ /* ----------------------------------------------------------------------
+ * Sampling (self-consistent)
+ *
+ * The PIT calibration test draws trees from the model and needs only each draw's log-probability
+ * (not the tree object), so we override sampleTreeLogProbability() with a direct simulation of the
+ * generative process and never materialise a Tree. At each reservable clade we escape with
+ * probability equal to its escape mass (= mu by the eps-solve, tail EXCLUDED) and otherwise take
+ * an observed (red) split ~ CCP; an escape draws a region order m proportional to M_m eps^(m-1), a
+ * boundary of m observed subclades proportional to its all-novel pathcount, and recurses into the
+ * boundary parts. The resolution shape within a region is not drawn -- every shape has the same
+ * weight eps^(m-1) and does not change the draw's log-probability -- so the simulation is cheap.
+ *
+ * The draw distribution exactly matches getLogProbabilityOfTree when the model is built with the
+ * tail OFF (then the red discount is 1 - mu, matching the escape mass), for trees whose regions are
+ * within reserveDepth; deeper regions (mass ~mu^reserveDepth) are never produced, the same
+ * self-consistent / full-support trade-off KRegCCD makes for its PIT.
+ * ------------------------------------------------------------------- */
+
+ @Override
+ public double sampleTreeLogProbability() {
+ return simulate(getRootClade());
+ }
+
+ private double simulate(Clade c) {
+ if (c.isLeaf()) {
+ return 0.0;
+ }
+ BitSet cb = c.getCladeInBits();
+ if (reservable(cb)) {
+ double eps = epsFor(cb, mu);
+ int[] n = countsFor(cb);
+ double escapeMass = 0.0;
+ for (int m = 2; m < n.length; m++) {
+ if (n[m] > 0) {
+ escapeMass += n[m] * Math.pow(eps, m - 1);
+ }
+ }
+ if (random.nextDouble() < escapeMass) {
+ int m = sampleOrder(n, eps, escapeMass);
+ BitSet[] parts = sampleBoundaryParts(cb, subclades(cb), m);
+ double logp = (m - 1) * Math.log(eps);
+ if (parts != null) {
+ for (BitSet bp : parts) {
+ logp += simulate(getClade(bp));
+ }
+ }
+ return logp;
+ }
+ CladePartition p = samplePartition(c);
+ double logp = Math.log(1.0 - escapeMass) + p.getLogCCP();
+ return logp + simulate(p.getChildClades()[0]) + simulate(p.getChildClades()[1]);
+ }
+ CladePartition p = samplePartition(c); // non-reservable: observed split, no discount
+ return p.getLogCCP() + simulate(p.getChildClades()[0]) + simulate(p.getChildClades()[1]);
+ }
+
+ /** Draws a region order m in {2..} with probability proportional to {@code M_m eps^(m-1)}. */
+ private int sampleOrder(int[] n, double eps, double escapeMass) {
+ double target = random.nextDouble() * escapeMass, acc = 0.0;
+ for (int m = 2; m < n.length; m++) {
+ if (n[m] > 0) {
+ acc += n[m] * Math.pow(eps, m - 1);
+ if (target < acc) {
+ return m;
+ }
+ }
+ }
+ for (int m = n.length - 1; m >= 2; m--) {
+ if (n[m] > 0) {
+ return m; // numerical guard
+ }
+ }
+ return 2;
+ }
+
+ /** Samples an observed (red) split of {@code c} with probability proportional to its CCP. */
+ private CladePartition samplePartition(Clade c) {
+ List partitions = c.getPartitions();
+ double target = random.nextDouble(), acc = 0.0;
+ for (CladePartition p : partitions) {
+ acc += p.getCCP();
+ if (target < acc) {
+ return p;
+ }
+ }
+ return partitions.get(partitions.size() - 1);
+ }
+
+ /**
+ * Weighted-reservoir samples one boundary of {@code c} into {@code m} observed subclades,
+ * proportional to its all-novel pathcount (so that, combined with order sampling, every distinct
+ * novel resolution is equiprobable at {@code eps^(m-1)}). Returns the parts, or {@code null} if
+ * none/op-budget.
+ */
+ private BitSet[] sampleBoundaryParts(BitSet c, List subs, int m) {
+ boundaryPick = null;
+ boundaryWeightSeen = 0.0;
+ enumOps = 0;
+ try {
+ sampleBoundaryWalk(c, subs, m, 0, BitSet.newBitSet(leafArraySize), new ArrayList<>(m));
+ } catch (BudgetExceeded e) {
+ return boundaryPick; // whatever was picked before the cap (may be null)
+ }
+ return boundaryPick;
+ }
+
+ private BitSet[] boundaryPick;
+ private double boundaryWeightSeen;
+
+ private void sampleBoundaryWalk(BitSet c, List subs, int m, int startIdx,
+ BitSet used, List chosen) {
+ if (++enumOps > OPS_BUDGET) {
+ throw BUDGET_EXCEEDED;
+ }
+ if (chosen.size() == m - 1) {
+ BitSet last = BitSet.newBitSet(c);
+ last.andNot(used);
+ if (last.isEmpty() || !isObs(last)) {
+ return;
+ }
+ if (compareBitSets(chosen.get(chosen.size() - 1), last) >= 0) {
+ return;
+ }
+ BitSet[] parts = new BitSet[m];
+ for (int i = 0; i < m - 1; i++) {
+ parts[i] = chosen.get(i);
+ }
+ parts[m - 1] = last;
+ int pc = countAllNovelResolutions(c, parts);
+ if (pc <= 0) {
+ return;
+ }
+ boundaryWeightSeen += pc;
+ if (random.nextDouble() * boundaryWeightSeen < pc) { // weighted reservoir
+ boundaryPick = parts;
+ }
+ return;
+ }
+ for (int i = startIdx; i < subs.size(); i++) {
+ BitSet pb = subs.get(i);
+ if (pb.intersects(used)) {
+ continue;
+ }
+ chosen.add(pb);
+ BitSet newUsed = BitSet.newBitSet(used);
+ newUsed.or(pb);
+ sampleBoundaryWalk(c, subs, m, i + 1, newUsed, chosen);
+ chosen.remove(chosen.size() - 1);
+ }
+ }
+
+ @Override
+ public Tree sampleTree(HeightSettingStrategy heightStrategy) {
+ throw new UnsupportedOperationException(
+ "MRegCCDSlow materialised-tree sampling is not implemented; sampleTreeLogProbability() "
+ + "(used by the PIT) simulates draws without building trees.");
+ }
+}
diff --git a/src/main/java/ccd/model/bitsets/BitSet.java b/src/main/java/ccd/model/bitsets/BitSet.java
index a6b6c37..3c727fe 100644
--- a/src/main/java/ccd/model/bitsets/BitSet.java
+++ b/src/main/java/ccd/model/bitsets/BitSet.java
@@ -55,7 +55,10 @@ public static BitSet newBitSet(BitSet other) {
if (other instanceof BitSet256 set) {
return new BitSet256(set);
}
- BitSet b = new BitSet(other.length());
+ // size(), not length(): length() is the index of the highest set bit plus one, so copying a
+ // set whose top words happen to be empty would return an undersized BitSet, and the bitwise
+ // operations below index the operand by this.words.length.
+ BitSet b = new BitSet(other.size());
b.or(other);
return b;
}
diff --git a/src/test/java/ccd/model/MRegCCDAgreementTest.java b/src/test/java/ccd/model/MRegCCDAgreementTest.java
new file mode 100644
index 0000000..a7d18f8
--- /dev/null
+++ b/src/test/java/ccd/model/MRegCCDAgreementTest.java
@@ -0,0 +1,133 @@
+package ccd.model;
+
+import beast.base.evolution.tree.Tree;
+import beast.base.evolution.tree.TreeParser;
+import ccd.model.bitsets.BitSet;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/** MRegCCD must reproduce MRegCCDSlow exactly: same boundary counts, same tree probabilities. */
+public class MRegCCDAgreementTest {
+
+ private static List taxa(int n) {
+ List out = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ out.add("T" + i);
+ }
+ return out;
+ }
+
+ private static List randomTrees(List taxa, int nTrees, long seed) {
+ Random rng = new Random(seed);
+ List out = new ArrayList<>();
+ for (int t = 0; t < nTrees; t++) {
+ List pool = new ArrayList<>(taxa);
+ while (pool.size() > 1) {
+ String a = pool.remove(rng.nextInt(pool.size()));
+ String b = pool.remove(rng.nextInt(pool.size()));
+ pool.add("(" + a + "," + b + ")");
+ }
+ out.add(new TreeParser(taxa, pool.get(0) + ";", 1, false));
+ }
+ return out;
+ }
+
+ @Test
+ public void boundaryCountsAndProbabilitiesAgree() {
+ int cladesChecked = 0;
+ for (int n : new int[]{6, 8, 10, 14}) {
+ for (int nTrees : new int[]{5, 25}) {
+ for (int depth : new int[]{2, 3, 4}) {
+ List tx = taxa(n);
+ MRegCCDSlow slow = new MRegCCDSlow(randomTrees(tx, nTrees, 21L), 0.0, 0.02, depth, true);
+ MRegCCD fast = new MRegCCD(randomTrees(tx, nTrees, 21L), 0.0, 0.02, depth, true);
+
+ for (Clade c : slow.getClades()) {
+ BitSet cb = c.getCladeInBits();
+ assertArrayEquals(slow.countsFor(cb), fast.countsFor(cb),
+ "boundary counts differ at " + n + " taxa, depth " + depth
+ + ", clade " + cb);
+ cladesChecked++;
+ }
+
+ List probe = randomTrees(tx, 40, 99L);
+ for (Tree t : probe) {
+ assertEquals(slow.getLogProbabilityOfTree(t), fast.getLogProbabilityOfTree(t),
+ 1e-9, "tree probability differs at " + n + " taxa, depth " + depth);
+ }
+ }
+ }
+ }
+ System.out.printf("MRegCCD agrees with MRegCCDSlow on %d clades and every probed tree%n",
+ cladesChecked);
+ }
+
+ /**
+ * The exact-normalisation guarantee is stated for the model, and MRegCCDTest checks it on the
+ * reference implementation. Since MRegCCD is the class callers get, and its fast path computes
+ * the counts that the reserve is solved from, check the property directly on it too.
+ */
+ @Test
+ public void fastImplementationIsExactlyNormalisedAtFullDepth() {
+ for (int n : new int[]{5, 6}) {
+ for (double mu : new double[]{0.02, 0.1, 0.25}) {
+ List tx = taxa(n);
+ // full reserve depth: no omitted tail, so the model must normalise exactly
+ MRegCCD m = new MRegCCD(randomTrees(tx, 6, 5L), 0.0, mu, tx.size(), false);
+ double sum = 0.0;
+ for (Tree t : allRootedTopologies(tx)) {
+ sum += Math.exp(m.getLogProbabilityOfTree(t));
+ }
+ System.out.printf("MRegCCD %d taxa mu=%.2f full-depth SUM = %.12f%n", n, mu, sum);
+ assertEquals(1.0, sum, 1e-9,
+ "MRegCCD at full reserve depth must be exactly normalised");
+ }
+ }
+ }
+
+ private static List allRootedTopologies(List taxa) {
+ List out = new ArrayList<>();
+ for (String shape : shapes(taxa)) {
+ out.add(new TreeParser(taxa, shape + ";", 1, false));
+ }
+ return out;
+ }
+
+ private static List shapes(List taxa) {
+ List out = new ArrayList<>();
+ if (taxa.size() == 1) {
+ out.add(taxa.get(0) + ":1");
+ return out;
+ }
+ String first = taxa.get(0);
+ List rest = taxa.subList(1, taxa.size());
+ int n = rest.size();
+ for (int mask = 0; mask < (1 << n); mask++) {
+ List left = new ArrayList<>();
+ left.add(first);
+ List right = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ if ((mask & (1 << i)) != 0) {
+ left.add(rest.get(i));
+ } else {
+ right.add(rest.get(i));
+ }
+ }
+ if (right.isEmpty()) {
+ continue;
+ }
+ for (String l : shapes(left)) {
+ for (String r : shapes(right)) {
+ out.add("(" + l + "," + r + "):1");
+ }
+ }
+ }
+ return out;
+ }
+}
diff --git a/src/test/java/ccd/model/MRegCCDTest.java b/src/test/java/ccd/model/MRegCCDTest.java
index ea8eb4a..4c5b8ad 100644
--- a/src/test/java/ccd/model/MRegCCDTest.java
+++ b/src/test/java/ccd/model/MRegCCDTest.java
@@ -13,7 +13,7 @@
import static org.junit.jupiter.api.Assertions.assertTrue;
/**
- * Validates the one-parameter per-new-split {@link MRegCCD}: that with the full reserve depth it is an
+ * Validates the one-parameter per-new-split {@link MRegCCDSlow}: that with the full reserve depth it is an
* exactly normalised distribution on enumerable taxon sets, and that truncating the reserve without a
* tail correction super-normalises it (the artefact that made order-2 falsely appear to close the gap
* to KRegCCD in the RSV2 experiment).
@@ -64,7 +64,7 @@ private static List trees(List taxa, List shapes) {
return out;
}
- private static double totalMass(MRegCCD m, List taxa) {
+ private static double totalMass(MRegCCDSlow m, List taxa) {
double sum = 0.0;
for (T t : allTopologies(taxa)) {
Tree tree = new TreeParser(taxa, topo(t) + ";", 1, false);
@@ -85,10 +85,10 @@ private void check5(double mu) {
List taxa = Arrays.asList("A", "B", "C", "D", "E");
List train = trees(taxa,
List.of(cat("A", "B", "C", "D", "E"), cat("D", "C", "B", "A", "E")));
- MRegCCD m = new MRegCCD(train, 0.0, mu, taxa.size(), false); // full depth -> no omitted tail
+ MRegCCDSlow m = new MRegCCDSlow(train, 0.0, mu, taxa.size(), false); // full depth -> no omitted tail
double sum = totalMass(m, taxa);
- System.out.printf("MRegCCD 5 taxa mu=%.2f full-depth SUM = %.12f%n", mu, sum);
- assertEquals(1.0, sum, 1e-9, "MRegCCD at full reserve depth must be exactly normalised");
+ System.out.printf("MRegCCDSlow 5 taxa mu=%.2f full-depth SUM = %.12f%n", mu, sum);
+ assertEquals(1.0, sum, 1e-9, "MRegCCDSlow at full reserve depth must be exactly normalised");
}
private void check6(double mu) {
@@ -96,10 +96,10 @@ private void check6(double mu) {
List train = trees(taxa,
List.of(new Node(cat("A", "B", "C", "D"), new Node(new Leaf("E"), new Leaf("F"))),
new Node(cat("D", "C", "B", "A"), new Node(new Leaf("E"), new Leaf("F")))));
- MRegCCD m = new MRegCCD(train, 0.0, mu, taxa.size(), false);
+ MRegCCDSlow m = new MRegCCDSlow(train, 0.0, mu, taxa.size(), false);
double sum = totalMass(m, taxa);
- System.out.printf("MRegCCD 6 taxa mu=%.2f full-depth SUM = %.12f%n", mu, sum);
- assertEquals(1.0, sum, 1e-9, "MRegCCD at full reserve depth must be exactly normalised");
+ System.out.printf("MRegCCDSlow 6 taxa mu=%.2f full-depth SUM = %.12f%n", mu, sum);
+ assertEquals(1.0, sum, 1e-9, "MRegCCDSlow at full reserve depth must be exactly normalised");
}
@Test
@@ -112,7 +112,7 @@ public void m2EqualsCCD0ExpandedSplits() {
for (int i = 0; i < all.size(); i += 47) picks.add(all.get(i)); // ~20 trees spread across the space
List train = trees(taxa, picks);
- MRegCCD mreg = new MRegCCD(train, 0.0, 0.05);
+ MRegCCDSlow mreg = new MRegCCDSlow(train, 0.0, 0.05);
CCD0 ccd0 = new CCD0(train, 0);
int checked = 0, withRecomb = 0;
@@ -146,7 +146,7 @@ public void samplerMatchesScorer() {
List picks = new ArrayList<>();
for (int i = 0; i < all.size(); i += 31) picks.add(all.get(i));
List train = trees(taxa, picks);
- MRegCCD m = new MRegCCD(train, 0.0, 0.1, taxa.size(), false); // full depth, tail off
+ MRegCCDSlow m = new MRegCCDSlow(train, 0.0, 0.1, taxa.size(), false); // full depth, tail off
// true entropy and normalisation by enumeration (scorer)
double sum = 0.0, H = 0.0;
@@ -169,7 +169,7 @@ public void samplerMatchesScorer() {
}
double hHat = s1 / N;
double se = Math.sqrt(Math.max(0, s2 / N - hHat * hHat) / N);
- System.out.printf("MRegCCD sampler: H_enum=%.5f H_MC=%.5f +/- %.5f (%.1f SE off)%n",
+ System.out.printf("MRegCCDSlow sampler: H_enum=%.5f H_MC=%.5f +/- %.5f (%.1f SE off)%n",
H, hHat, se, Math.abs(hHat - H) / se);
assertEquals(H, hHat, Math.max(5 * se, 0.01),
"sampler entropy must match the scorer's enumerated entropy");
@@ -182,9 +182,9 @@ public void truncatedReserveSuperNormalises() {
List.of(new Node(cat("A", "B", "C", "D"), new Node(new Leaf("E"), new Leaf("F"))),
new Node(cat("D", "C", "B", "A"), new Node(new Leaf("E"), new Leaf("F")))));
double mu = 0.2;
- double full = totalMass(new MRegCCD(train, 0.0, mu, taxa.size(), false), taxa);
- double order2 = totalMass(new MRegCCD(train, 0.0, mu, 2, false), taxa); // M2 only, no tail
- System.out.printf("MRegCCD 6 taxa mu=%.2f: full-depth SUM=%.9f order-2 SUM=%.9f%n", mu, full, order2);
+ double full = totalMass(new MRegCCDSlow(train, 0.0, mu, taxa.size(), false), taxa);
+ double order2 = totalMass(new MRegCCDSlow(train, 0.0, mu, 2, false), taxa); // M2 only, no tail
+ System.out.printf("MRegCCDSlow 6 taxa mu=%.2f: full-depth SUM=%.9f order-2 SUM=%.9f%n", mu, full, order2);
assertEquals(1.0, full, 1e-9, "full depth normalised");
assertTrue(order2 > 1.0 + 1e-4,
"order-2 reserve (no tail) must super-normalise (sum > 1), got " + order2);
diff --git a/src/test/java/ccd/model/MRegDepthTimingTest.java b/src/test/java/ccd/model/MRegDepthTimingTest.java
new file mode 100644
index 0000000..1d031f2
--- /dev/null
+++ b/src/test/java/ccd/model/MRegDepthTimingTest.java
@@ -0,0 +1,58 @@
+package ccd.model;
+
+import beast.base.evolution.tree.Tree;
+import beastfx.app.treeannotator.TreeAnnotator;
+import ccd.tools.CCDToolUtil;
+import org.junit.jupiter.api.Assumptions;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+
+/** How much of MRegCCD's scoring cost is its reserve depth, versus the boundary enumeration itself? */
+public class MRegDepthTimingTest {
+
+ private static final String PATH = System.getProperty("ccd.trees", "");
+
+ private static List read(int count, int skip) throws Exception {
+ TreeAnnotator.TreeSet ts = CCDToolUtil.getTreeSet(PATH, 10);
+ ts.reset();
+ List all = new ArrayList<>();
+ while (ts.hasNext()) {
+ all.add(ts.next());
+ }
+ List pool = all.subList(skip * all.size() / 2, (skip + 1) * all.size() / 2);
+ List out = new ArrayList<>();
+ double step = Math.max(1.0, pool.size() / (double) count);
+ for (int i = 0; i < count && (int) (i * step) < pool.size(); i++) {
+ out.add(pool.get((int) (i * step)));
+ }
+ return out;
+ }
+
+ @Test
+ public void depthVersusCost() throws Exception {
+ Assumptions.assumeTrue(!PATH.isEmpty() && new File(PATH).exists());
+ List test = read(200, 1);
+ System.out.printf("%n=== %s: MRegCCD cost by reserve depth ===%n", new File(PATH).getName());
+ System.out.printf("%-7s %-9s %-12s %-12s %-14s%n", "depth", "impl", "construct", "score/200", "mean logP");
+ for (int depth : new int[]{2, 3, 4}) {
+ for (String which : new String[]{"MRegCCDSlow", "MRegCCD"}) {
+ long t0 = System.nanoTime();
+ MRegCCDSlow m = which.equals("MRegCCDSlow")
+ ? new MRegCCDSlow(read(500, 0), 0.0, MRegCCDSlow.DEFAULT_MU, depth, true)
+ : new MRegCCD(read(500, 0), 0.0, MRegCCDSlow.DEFAULT_MU, depth, true);
+ long build = (System.nanoTime() - t0) / 1_000_000L;
+ t0 = System.nanoTime();
+ double sum = 0;
+ for (Tree t : test) {
+ sum += m.getLogProbabilityOfTree(t);
+ }
+ long score = (System.nanoTime() - t0) / 1_000_000L;
+ System.out.printf("%-7d %-9s %9dms %9dms %14.6f%n",
+ depth, which, build, score, sum / test.size());
+ }
+ }
+ }
+}
diff --git a/src/test/java/ccd/model/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));
+ }
+}