diff --git a/.github/workflows/ci-publish.yml b/.github/workflows/ci-publish.yml index d9e0027..045ec21 100644 --- a/.github/workflows/ci-publish.yml +++ b/.github/workflows/ci-publish.yml @@ -13,11 +13,11 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set up JDK 25 (Azul Zulu) if: "!startsWith(github.ref, 'refs/tags/v')" - uses: actions/setup-java@v4 + uses: actions/setup-java@v5 with: distribution: zulu java-version: '25' @@ -25,7 +25,7 @@ jobs: - name: Set up JDK 25 (Azul Zulu) for Maven Central if: startsWith(github.ref, 'refs/tags/v') - uses: actions/setup-java@v4 + uses: actions/setup-java@v5 with: distribution: zulu java-version: '25' diff --git a/pom.xml b/pom.xml index 6658a27..0d1e34a 100644 --- a/pom.xml +++ b/pom.xml @@ -242,7 +242,7 @@ org.sonatype.central central-publishing-maven-plugin - 0.6.0 + 0.11.0 true central diff --git a/src/main/java/ccd/model/AbstractCCD.java b/src/main/java/ccd/model/AbstractCCD.java index 625433c..73a8bce 100644 --- a/src/main/java/ccd/model/AbstractCCD.java +++ b/src/main/java/ccd/model/AbstractCCD.java @@ -871,6 +871,64 @@ public Tree sampleTree(HeightSettingStrategy heightStrategy) { return getTreeBasedOnStrategy(SamplingStrategy.Sampling, heightStrategy); } + /** As {@link #sampleTrees(int, HeightSettingStrategy)} with {@link HeightSettingStrategy#None}. */ + public List sampleTrees(int count) { + return sampleTrees(count, HeightSettingStrategy.None); + } + + /** + * Draws {@code count} trees, in parallel. Individual draws are independent -- they only read + * the clade DAG and its cached probabilities -- so the one-at-a-time + * {@link #sampleTree(HeightSettingStrategy)} loop leaves all but one core idle. This warms + * every lazily-built structure first (so the draws themselves are pure reads) and then fans + * the draws out. + * + *

The result is reproducible: the per-draw seeds are taken from the shared {@link #random} + * serially, before the fan-out, so a given {@link #setRandom} state yields the same + * trees in the same order however the draws happen to be scheduled. + * + * @param count number of trees to draw + * @param heightStrategy how node heights are set (see {@link HeightSettingStrategy}) + * @return the sampled trees, in draw order + */ + public List sampleTrees(int count, HeightSettingStrategy heightStrategy) { + if (count < 0) { + throw new IllegalArgumentException("count must be non-negative, got " + count); + } + if (count == 0) { + return List.of(); + } + + // Warm everything lazily built, serially -- the parallel draws below must only read. + tidyUpCacheIfDirty(); + computeCladeProbabilitiesIfDirty(); + if (heightStrategy == HeightSettingStrategy.CommonAncestorHeights) { + setupCommonAncestorHeightsIfDirty(); + } + prepareForSampling(); + + long[] seeds = new long[count]; + for (int i = 0; i < count; i++) { + seeds[i] = random.nextLong(); + } + + return java.util.stream.IntStream.range(0, count).parallel().mapToObj(i -> { + threadRandom.set(new Random(seeds[i])); + try { + return getTreeBasedOnStrategy(SamplingStrategy.Sampling, heightStrategy); + } finally { + threadRandom.remove(); + } + }).toList(); + } + + /** + * Hook for subclasses to precompute any per-clade state that a draw would otherwise solve + * lazily, so that {@link #sampleTrees} fans out over pure reads. Does nothing by default. + */ + protected void prepareForSampling() { + } + @Override public Tree getMAPTree() { return this.getMAPTree(HeightSettingStrategy.One); @@ -881,8 +939,24 @@ public Tree getMAPTree(HeightSettingStrategy heightStrategy) { return getTreeBasedOnStrategy(SamplingStrategy.MAP, heightStrategy); } - /* Helper for methods to assign indices to inner vertices */ - private int runningInnerIndex; + /* Helper for methods to assign indices to inner vertices. Thread-local so that concurrent + * draws (see sampleTrees) each number their own tree's inner nodes from + * getSizeOfLeavesArray() upwards, without treading on one another. */ + private final ThreadLocal runningInnerIndex = ThreadLocal.withInitial(() -> new int[1]); + + /* Per-thread RNG installed by sampleTrees for the duration of one draw; null (meaning "use + * the shared `random`") on any thread that is not inside a parallel sampling run. */ + private final ThreadLocal threadRandom = new ThreadLocal<>(); + + /** + * The RNG the calling thread should sample with: its own stream inside a + * {@link #sampleTrees} draw, otherwise the shared {@link #random} (so single-threaded + * sampling and {@link #setRandom} behave exactly as before). + */ + protected Random random() { + Random threadLocal = threadRandom.get(); + return (threadLocal != null) ? threadLocal : random; + } /* Helper for methods to assign indices to leaves */ // private int runningLeafIndex; @@ -896,7 +970,7 @@ protected Tree getTreeBasedOnStrategy(SamplingStrategy samplingStrategy, HeightS setupCommonAncestorHeightsIfDirty(); } - runningInnerIndex = this.getSizeOfLeavesArray(); + runningInnerIndex.get()[0] = this.getSizeOfLeavesArray(); Node root = getVertexBasedOnStrategy(this.rootClade, samplingStrategy, heightStrategy); if (this instanceof FilteredCCD) { @@ -1019,7 +1093,7 @@ protected Node buildInternalVertex(Clade clade, CladePartition partition, /** Returns and advances the running inner-node index used when materialising sampled trees. */ protected int nextRunningInnerIndex() { - return runningInnerIndex++; + return runningInnerIndex.get()[0]++; } /* Helper method */ @@ -1070,7 +1144,7 @@ protected CladePartition getPartitionBasedOnStrategy(Clade clade, SamplingStrate // the sum of probabilities over all partitions of a clade // should be 1, so we can sample with a random value - double sampleWithMe = random.nextDouble(); + double sampleWithMe = random().nextDouble(); double probabilitySum = 0; for (CladePartition nextPartition : partitions) { diff --git a/src/main/java/ccd/model/KRegCCD.java b/src/main/java/ccd/model/KRegCCD.java index 8202eb6..0018769 100644 --- a/src/main/java/ccd/model/KRegCCD.java +++ b/src/main/java/ccd/model/KRegCCD.java @@ -190,11 +190,17 @@ public enum SamplingFidelity {SELF_CONSISTENT, FULL_SUPPORT} private SamplingFidelity samplingFidelity = SamplingFidelity.SELF_CONSISTENT; - /** Reservoir state for {@link #sampleBoundary}: count of valid boundaries seen so far. */ - private int boundarySeen; - private double boundaryWeightSeen; // weighted-reservoir accumulator for FLAT boundary sampling - /** Reservoir state for {@link #sampleBoundary}: the boundary currently selected. */ - private BitSet[] boundaryPick; + /** + * Reservoir state for one {@link #sampleBoundary} call: the number of valid boundaries seen so + * far (or their pathcount-weighted total, for {@link NovelMode#FLAT}) and the boundary + * currently selected. Held per call rather than per instance so that the concurrent draws in + * {@link #sampleTrees} each keep their own reservoir. + */ + private static final class BoundaryReservoir { + int seen; + double weightSeen; + BitSet[] pick; + } /** * Strictly-positive fallback increment for novel-node heights when the region top is not @@ -223,7 +229,7 @@ private record CladeReg(boolean reservable, double logEps, double tail, * Default per-clade escape probability used when a tool builds a KRegCCD without * specifying one (e.g. via {@link CCDType#KRegCCD}); the RSV2 operating point. */ - public static final double DEFAULT_MU = 0.01; + public static final double DEFAULT_MU = 0.005; /** * Default additive-smoothing pseudocount for the split-expanded backbone (shared with * {@link RegCCD}'s default). @@ -694,6 +700,14 @@ private void computeDiscountedMap() { if (mapLogProbMemo != null) { return; } + // The DP below visits every clade and needs each one's reserve discount, so solve all the + // reserves up front in parallel instead of lazily inside the serial size-ordered loop -- + // exactly as getEntropyRecursive does. Each computeReg is an O(m^2) disjoint-pair pass, so + // on large clade sets that lazy path, not the DP arithmetic, dominates getMAPTree, and it + // runs on one core. Same values either way (the reserve of a clade depends only on its own + // observed subclades), so this changes cost only, not the MAP tree. + precomputeReserves(); + Map logProb = new HashMap<>(); Map argmax = new HashMap<>(); List clades = new ArrayList<>(getClades()); @@ -889,13 +903,29 @@ public double[] getEntropyMonteCarlo(int samples) { * KRegCCD distribution */ public double getEntropyRecursive() { - Map freeMemo = new HashMap<>(); - Map redMemo = new HashMap<>(); - List clades = new ArrayList<>(getClades()); - clades.sort((a, b) -> Integer.compare(a.size(), b.size())); - for (Clade c : clades) { - entropyRedForced(c, freeMemo, redMemo); - entropyFree(c, freeMemo, redMemo); + Map freeMemo = new ConcurrentHashMap<>(); + Map redMemo = new ConcurrentHashMap<>(); + // Both entropyRedForced(c) and entropyFree(c) read only strictly-smaller clades (a clade's + // partition children and its blue-region boundary parts are all smaller than it), so once + // every smaller size level is memoised, all clades of a given size are independent and can + // be computed in parallel -- exactly as precomputeReserves parallelises computeReg. Each + // task writes only its own clade's freeMemo/redMemo entries (a concurrent map), and the + // per-clade arithmetic is unchanged, so the entropy is identical to the serial computation. + List> bySize = new ArrayList<>(leafArraySize + 1); + for (int s = 0; s <= leafArraySize; s++) { + bySize.add(new ArrayList<>()); + } + for (Clade c : getClades()) { + bySize.get(c.size()).add(c); + } + for (int s = 1; s <= leafArraySize; s++) { + List level = bySize.get(s); + if (!level.isEmpty()) { + level.parallelStream().forEach(c -> { + entropyRedForced(c, freeMemo, redMemo); + entropyFree(c, freeMemo, redMemo); + }); + } } return freeMemo.getOrDefault(getRootClade(), 0.0); } @@ -984,10 +1014,14 @@ private BlueTerms blueContribution(Clade c, CladeReg reg, double eps, Map freeMemo, Map redMemo) { List subs = observedSubclades(c); BitSet cBits = c.getCladeInBits(); + int lastBoundary = reg.lastBoundary(); double[] acc = new double[2]; // [0] childEntropy, [1] extraLogPc (SHARED) enumOps.get()[0] = 0; try { - for (int m = 3; m <= reg.lastBoundary(); m++) { + // Orders 3 and 4 (the default reserve depth) via the fast sum-arithmetic pair pass. + blueContributionFast(subs, cBits, eps, lastBoundary, acc, freeMemo, redMemo); + // Deeper orders (only reached when boundaries 3 and 4 are both empty) keep the old walk. + for (int m = 5; m <= lastBoundary; m++) { blueWalk(cBits, subs, m, 0, BitSet.newBitSet(leafArraySize), new ArrayList<>(m), eps, acc, freeMemo, redMemo); } @@ -997,6 +1031,159 @@ private BlueTerms blueContribution(Clade c, CladeReg reg, double eps, return new BlueTerms(acc[0], acc[1]); } + /** + * Accumulates the blue-region entropy terms for boundary orders 3 and 4 in one disjoint-pair + * pass, the entropy counterpart of {@link #countN1N2}: it enumerates the same boundaries by + * weighted-sum complement lookups, but at each boundary adds {@code boundaryMass * sum of the + * boundary parts' forced-red entropy} to {@code acc[0]} (and, in SHARED mode, the per-region + * {@code eps^(m-2) * log pathcount} self-information to {@code acc[1]}). Emits exactly the same + * boundaries as the old {@link #blueWalk} for these orders. + */ + private void blueContributionFast(List subs, BitSet cBits, double eps, int lastBoundary, + double[] acc, Map freeMemo, + Map redMemo) { + int m = subs.size(); + BitSet[] sb = new BitSet[m]; + long[] partSum = new long[m]; + int[] partCard = new int[m]; + double[] redH = new double[m]; // forced-red entropy of each boundary part (memoised lookup) + for (int i = 0; i < m; i++) { + sb[i] = subs.get(i).getCladeInBits(); + partSum[i] = weightedSum(sb[i]); + partCard[i] = sb[i].cardinality(); + redH[i] = entropyRedForced(subs.get(i), freeMemo, redMemo); + } + long sumC = weightedSum(cBits); + int cardC = cBits.cardinality(); + boolean flat = novelMode == NovelMode.FLAT; + double eps1 = eps; // eps^(3-2) + double eps2 = eps * eps; // eps^(4-2) + boolean do4 = lastBoundary >= 4; + final long[] ops = enumOps.get(); + + int stb = Integer.highestOneBit(Math.max(1, m)) << 1; + int smask = stb - 1; + int[] subHead = new int[stb]; + java.util.Arrays.fill(subHead, -1); + int[] subNxt = new int[m]; + for (int i = 0; i < m; i++) { + int b = mixHash(partSum[i]) & smask; + subNxt[i] = subHead[b]; + subHead[b] = i; + } + + int cap = 256, P = 0; + int[] ei = new int[cap], ej = new int[cap], eh = new int[cap]; + long[] sU = new long[cap]; + BitSet ue = BitSet.newBitSet(leafArraySize); + BitSet uf = BitSet.newBitSet(leafArraySize); + for (int i = 0; i < m; i++) { + BitSet a = sb[i]; + for (int j = i + 1; j < m; j++) { + if (++ops[0] > opsBudget) { + throw BUDGET_EXCEEDED; + } + if (a.intersects(sb[j])) { + continue; + } + long s = partSum[i] + partSum[j]; + long sR = sumC - s; + boolean ueBuilt = false; + for (int d = subHead[mixHash(sR) & smask]; d >= 0; d = subNxt[d]) { + if (d <= j || partSum[d] != sR) { + continue; + } + if (partCard[i] + partCard[j] + partCard[d] != cardC) { + continue; + } + if (!ueBuilt) { + ue.clear(); + ue.or(a); + ue.or(sb[j]); + ueBuilt = true; + } + if (ue.intersects(sb[d])) { + continue; + } + BitSet[] parts = {sb[i], sb[j], sb[d]}; + int pc = countAllNovelResolutions(cBits, parts); + if (pc > 0) { + double childH = redH[i] + redH[j] + redH[d]; + acc[0] += (flat ? pc * eps1 : eps1) * childH; + if (!flat) { + acc[1] += eps1 * Math.log(pc); + } + } + } + if (do4) { + if (P == cap) { + cap <<= 1; + ei = java.util.Arrays.copyOf(ei, cap); + ej = java.util.Arrays.copyOf(ej, cap); + eh = java.util.Arrays.copyOf(eh, cap); + sU = java.util.Arrays.copyOf(sU, cap); + } + ei[P] = i; + ej[P] = j; + sU[P] = s; + eh[P] = mixHash(s); + P++; + } + } + } + if (!do4 || P == 0) { + return; + } + int tb = Integer.highestOneBit(P) << 1; + int mask = tb - 1; + int[] head = new int[tb]; + java.util.Arrays.fill(head, -1); + int[] nxt = new int[P]; + for (int e = 0; e < P; e++) { + int b = eh[e] & mask; + nxt[e] = head[b]; + head[b] = e; + } + for (int e = 0; e < P; e++) { + int i = ei[e], j = ej[e]; + long sR = sumC - sU[e]; + boolean ueBuilt = false; + for (int f = head[mixHash(sR) & mask]; f >= 0; f = nxt[f]) { + if (++ops[0] > opsBudget) { + throw BUDGET_EXCEEDED; + } + int k = ei[f], l = ej[f]; + if (k <= j) { + continue; + } + if (partCard[i] + partCard[j] + partCard[k] + partCard[l] != cardC) { + continue; + } + if (!ueBuilt) { + ue.clear(); + ue.or(sb[i]); + ue.or(sb[j]); + ueBuilt = true; + } + uf.clear(); + uf.or(sb[k]); + uf.or(sb[l]); + if (ue.intersects(uf)) { + continue; + } + BitSet[] parts = {sb[i], sb[j], sb[k], sb[l]}; + int pc = countAllNovelResolutions(cBits, parts); + if (pc > 0) { + double childH = redH[i] + redH[j] + redH[k] + redH[l]; + acc[0] += (flat ? pc * eps2 : eps2) * childH; + if (!flat) { + acc[1] += eps2 * Math.log(pc); + } + } + } + } + } + private void blueWalk(BitSet cBits, List subs, int m, int startIdx, BitSet used, List chosen, double eps, double[] acc, Map freeMemo, Map redMemo) { @@ -1270,7 +1457,7 @@ protected Node getVertexBasedOnStrategy(Clade clade, SamplingStrategy samplingSt for (double w : orderWeight) { escapeMass += w; } - if (escapeMass > 0 && random.nextDouble() < escapeMass) { + if (escapeMass > 0 && random().nextDouble() < escapeMass) { Node region = sampleBlueRegion(clade, reg, orderWeight, escapeMass, heightStrategy); if (region != null) { return region; @@ -1347,7 +1534,7 @@ private double[] escapeOrderWeights(Clade c, CladeReg reg, double eps) { */ private Node sampleBlueRegion(Clade c, CladeReg reg, double[] orderWeight, double escapeMass, HeightSettingStrategy heightStrategy) { - double target = random.nextDouble() * escapeMass; + double target = random().nextDouble() * escapeMass; double acc = 0.0; int m = -1; for (int mm = 3; mm < orderWeight.length; mm++) { @@ -1411,17 +1598,15 @@ private Node sampleBlueRegion(Clade c, CladeReg reg, double[] orderWeight, doubl * enumeration {@link #countNj} counts. Returns the chosen parts, or {@code null} if none. */ private BitSet[] sampleBoundary(Clade c, List subs, int m) { - boundarySeen = 0; - boundaryWeightSeen = 0.0; - boundaryPick = null; + BoundaryReservoir reservoir = new BoundaryReservoir(); enumOps.get()[0] = 0; sampleBoundaryWalk(c.getCladeInBits(), subs, m, 0, - BitSet.newBitSet(leafArraySize), new ArrayList<>(m)); - return boundaryPick; + BitSet.newBitSet(leafArraySize), new ArrayList<>(m), reservoir); + return reservoir.pick; } private void sampleBoundaryWalk(BitSet cBits, List subs, int m, int startIdx, - BitSet used, List chosen) { + BitSet used, List chosen, BoundaryReservoir reservoir) { if (++enumOps.get()[0] > opsBudget) { throw BUDGET_EXCEEDED; } @@ -1446,14 +1631,14 @@ private void sampleBoundaryWalk(BitSet cBits, List subs, int m, int start if (novelMode == NovelMode.FLAT) { // FLAT samples a boundary in proportion to its pathcount (so that, with the // uniform resolution pick below, every distinct novel tree is equiprobable). - boundaryWeightSeen += pc; - if (random.nextDouble() * boundaryWeightSeen < pc) { // weighted reservoir - boundaryPick = parts; + reservoir.weightSeen += pc; + if (random().nextDouble() * reservoir.weightSeen < pc) { // weighted reservoir + reservoir.pick = parts; } } else { - boundarySeen++; - if (random.nextInt(boundarySeen) == 0) { // uniform reservoir: keep with prob 1/seen - boundaryPick = parts; + reservoir.seen++; + if (random().nextInt(reservoir.seen) == 0) { // uniform reservoir: keep with prob 1/seen + reservoir.pick = parts; } } return; @@ -1466,7 +1651,7 @@ private void sampleBoundaryWalk(BitSet cBits, List subs, int m, int start chosen.add(pb); BitSet newUsed = (BitSet) used.clone(); newUsed.or(pb); - sampleBoundaryWalk(cBits, subs, m, i + 1, newUsed, chosen); + sampleBoundaryWalk(cBits, subs, m, i + 1, newUsed, chosen, reservoir); chosen.remove(chosen.size() - 1); } } @@ -1558,7 +1743,7 @@ private Shape sampleShape(int mask, BitSet[] parts, BitSet[] unionOf, int[] f) { } int low = mask & (-mask); int rest = mask ^ low; - int targetCount = random.nextInt(f[mask]); // f[mask] > 0 at every visited mask + int targetCount = random().nextInt(f[mask]); // f[mask] > 0 at every visited mask int acc = 0; int chosenS1 = -1; int chosenS2 = -1; @@ -1744,6 +1929,16 @@ private int countAllNovelResolutions(BitSet cBits, BitSet[] parts) { * paying that enumeration lazily — and unevenly — during the first samples. Idempotent (cache hits * after the first call). Useful before a large sampling loop (e.g. variational ELBO estimation). */ + /** + * Solves every clade's reserve before {@link #sampleTrees} fans out, so the concurrent draws + * only read {@code regCache} instead of racing to fill it (and so the reserve cost is paid + * once, in parallel, rather than unevenly across the first draws). + */ + @Override + protected void prepareForSampling() { + precomputeReserves(); + } + public void precomputeReserves() { // Clades are independent: each computeReg only reads immutable clade structure + writes its own // entry in the (concurrent) regCache, and the op-budget counter is thread-local. With tailMode @@ -1796,18 +1991,53 @@ private CladeReg computeReg(Clade c) { int last = 2; boolean anyNonzero = false; enumOps.get()[0] = 0; - for (int j = 3; j <= c.size(); j++) { - if (j > reserveBoundary && anyNonzero) { - break; // reached reserve depth and have a usable (positive) reserve - } + if (reserveBoundary == 3 || reserveBoundary == 4) { + // Fast path for the practical reserve depths: k = 1 (N_1 only) or k = 2 (N_1 and N_2), + // sharing one weighted-sum disjoint-pair pass. N_2 is skipped entirely when k = 1. + boolean computeN2 = reserveBoundary == 4 && c.size() >= 4; try { - n[j] = countNj(c, subs, j, false); // accumulates enumOps across j + long[] n12 = countN1N2(c, subs, computeN2); + n[3] = (int) n12[0]; + last = 3; + if (n[3] > 0) { + anyNonzero = true; + } + if (computeN2) { + n[4] = (int) n12[1]; + last = 4; + if (n[4] > 0) { + anyNonzero = true; + } + } } catch (BudgetExceeded e) { - break; + // pathological pair blow-up: leave N_1/N_2 = 0 (same capped-out reserve as before) } - last = j; - if (n[j] > 0) { - anyNonzero = true; + // Only climb past the reserve depth if all computed orders were empty (keeps eps positive). + for (int j = last + 1; j <= c.size() && !anyNonzero; j++) { + try { + n[j] = countNj(c, subs, j, false); + } catch (BudgetExceeded e) { + break; + } + last = j; + if (n[j] > 0) { + anyNonzero = true; + } + } + } else { + for (int j = 3; j <= c.size(); j++) { + if (j > reserveBoundary && anyNonzero) { + break; // reached reserve depth and have a usable (positive) reserve + } + try { + n[j] = countNj(c, subs, j, false); // accumulates enumOps across j + } catch (BudgetExceeded e) { + break; + } + last = j; + if (n[j] > 0) { + anyNonzero = true; + } } } double logEps = anyNonzero ? solveLogEps(n, last, mu) : Double.NEGATIVE_INFINITY; @@ -1921,11 +2151,317 @@ private List observedSubclades(Clade c) { * resolution. earlyExit returns 1 as soon as one is found. */ private int countNj(Clade c, List subs, int m, boolean earlyExit) { BitSet cBits = c.getCladeInBits(); + if (m == 4) { + // Boundary-4 (N_2) dominates construction: the generic enumerate is O(m^3) (choose 3 + // parts, derive the 4th as the complement), so on the big clades it burns the whole + // op-budget. A 4-part boundary is exactly two disjoint observed PAIRS whose unions are + // complementary within C, so we can meet in the middle: enumerate the O(m^2) disjoint + // pairs once, bucket them by their union bitset, then match each union U with C\U. This + // visits only combinations that actually tile C (no rejected triples) -- O(m^2) build + // plus one visit per valid boundary -- and returns exactly the same count as enumerate + // for both FLAT (sum of pathcounts) and SHARED (admissible-boundary indicator). + return countN2ViaPairs(cBits, subs, earlyExit); + } BitSet used = BitSet.newBitSet(leafArraySize); List chosen = new ArrayList<>(m); return enumerate(cBits, subs, m, 0, used, chosen, earlyExit); } + /** + * Meet-in-the-middle count of boundary-4 (N_2) partitions of {@code cBits} into observed + * subclades (see {@link #countNj}). A 4-part boundary {@code P1 j}, so + * each boundary is emitted exactly once (indices {@code iComplements are found by weighted-sum arithmetic, not bitset materialisation. + * {@link Clade#getCladeInBits()}'s {@code hashCode} weights word {@code w} by {@code w+1}, so + * {@code weightedSum} is additive over disjoint sets: a pair's union sum is the sum of its two + * parts' precomputed sums, and the complement's sum is {@code weightedSum(C) - union sum}. Pairs + * are hashed by that {@code long} sum (no per-pair union bitset or {@code hashCode}); a chain hit + * is confirmed exactly — cheaply and immune to sum wraparound — by the cardinality + * identity {@code |Pi| + |Pj| + |Pk| + |Pl| == |C|} together with disjointness, which for parts + * that are subclades of C is equivalent to tiling C. This is the dominant cost of building the + * reserves and is allocation-free per pair. Shares the per-clade {@link #enumOps}/{@link + * #opsBudget} guard with {@link #enumerate}, so a pathological clade still degrades gracefully to + * the same capped-out {@code N_2 = 0}. + */ + private int countN2ViaPairs(BitSet cBits, List subs, boolean earlyExit) { + int m = subs.size(); + BitSet[] sb = new BitSet[m]; + long[] partSum = new long[m]; + int[] partCard = new int[m]; + for (int i = 0; i < m; i++) { + sb[i] = subs.get(i).getCladeInBits(); + partSum[i] = weightedSum(sb[i]); + partCard[i] = sb[i].cardinality(); + } + long sumC = weightedSum(cBits); + int cardC = cBits.cardinality(); + final long[] ops = enumOps.get(); + + // Enumerate disjoint observed pairs into flat arrays; entry e is pair (ei[e], ej[e]) with + // union weighted-sum sU[e]. No per-pair bitset is built (sums are added from the parts). + int cap = 256, P = 0; + int[] ei = new int[cap], ej = new int[cap], eh = new int[cap]; + long[] sU = new long[cap]; + for (int i = 0; i < m; i++) { + BitSet a = sb[i]; + for (int j = i + 1; j < m; j++) { + if (++ops[0] > opsBudget) { + throw BUDGET_EXCEEDED; + } + if (a.intersects(sb[j])) { + continue; + } + if (P == cap) { + cap <<= 1; + ei = java.util.Arrays.copyOf(ei, cap); + ej = java.util.Arrays.copyOf(ej, cap); + eh = java.util.Arrays.copyOf(eh, cap); + sU = java.util.Arrays.copyOf(sU, cap); + } + long s = partSum[i] + partSum[j]; + ei[P] = i; + ej[P] = j; + sU[P] = s; + eh[P] = mixHash(s); + P++; + } + } + if (P == 0) { + return 0; + } + + // Chained open hash over the union sums, so a pair can be found by its complement's sum. + int tb = Integer.highestOneBit(P) << 1; // power of two in [P, 2P) + int mask = tb - 1; + int[] head = new int[tb]; + java.util.Arrays.fill(head, -1); + int[] nxt = new int[P]; + for (int e = 0; e < P; e++) { + int b = eh[e] & mask; + nxt[e] = head[b]; + head[b] = e; + } + + // For each pair (i,j) match pairs (k,l) with union sum == weightedSum(C∖union(i,j)) and k>j; + // confirm exactly with the cardinality/disjointness tiling test. + int count = 0; + BitSet ue = BitSet.newBitSet(leafArraySize); + BitSet uf = BitSet.newBitSet(leafArraySize); + for (int e = 0; e < P; e++) { + int i = ei[e], j = ej[e]; + long sR = sumC - sU[e]; + boolean ueBuilt = false; + for (int f = head[mixHash(sR) & mask]; f >= 0; f = nxt[f]) { + if (++ops[0] > opsBudget) { + throw BUDGET_EXCEEDED; + } + int k = ei[f], l = ej[f]; + if (k <= j) { + continue; // canonical i the four parts tile C exactly + } + BitSet[] parts = {sb[i], sb[j], sb[k], sb[l]}; + int res = countAllNovelResolutions(cBits, parts); + if (res > 0) { + count += (novelMode == NovelMode.FLAT) ? res : 1; + if (earlyExit) { + return count; + } + } + } + } + return count; + } + + /** + * Computes the boundary-3 ({@code N_1}) and boundary-4 ({@code N_2}) counts of {@code c} + * together in a single disjoint-observed-pair pass (returns {@code {N_1, N_2}}). This is + * the default reserve depth (k = 2) and by far the hottest part of building the reserves, so the + * two orders share the one O(m^2) pair enumeration instead of scanning it twice. + * + *

Both boundaries are matched by weighted-sum arithmetic (see {@link #countN2ViaPairs}): for a + * disjoint pair {@code (i,j)} the boundary-3 complement is the single observed subclade whose + * {@code weightedSum} equals {@code weightedSum(C) - union sum} (looked up in a sum-keyed hash of + * the subclades), and the boundary-4 complement is another disjoint pair with that sum. Every hit + * is confirmed exactly with the cardinality/disjointness tiling test, so sum collisions and + * wraparound are harmless. Honours the same {@link #enumOps}/{@link #opsBudget} guard. + * + * @param computeN2 when {@code false} (reserve depth k = 1, {@code eps} from {@code N_1} alone) + * the boundary-4 stash and match are skipped and {@code N_2} is returned as 0. + */ + private long[] countN1N2(Clade c, List subs, boolean computeN2) { + int m = subs.size(); + BitSet[] sb = new BitSet[m]; + long[] partSum = new long[m]; + int[] partCard = new int[m]; + for (int i = 0; i < m; i++) { + sb[i] = subs.get(i).getCladeInBits(); + partSum[i] = weightedSum(sb[i]); + partCard[i] = sb[i].cardinality(); + } + BitSet cBits = c.getCladeInBits(); + long sumC = weightedSum(cBits); + int cardC = c.size(); + final long[] ops = enumOps.get(); + + // Sum-keyed chained hash of the subclades, for the boundary-3 complement (a single subclade). + int stb = Integer.highestOneBit(Math.max(1, m)) << 1; + int smask = stb - 1; + int[] subHead = new int[stb]; + java.util.Arrays.fill(subHead, -1); + int[] subNxt = new int[m]; + for (int i = 0; i < m; i++) { + int b = mixHash(partSum[i]) & smask; + subNxt[i] = subHead[b]; + subHead[b] = i; + } + + // One disjoint-pair pass: score boundary-3 inline, stash pairs for the boundary-4 match. + long n1 = 0; + int cap = 256, P = 0; + int[] ei = new int[cap], ej = new int[cap], eh = new int[cap]; + long[] sU = new long[cap]; + BitSet ue = BitSet.newBitSet(leafArraySize); + BitSet uf = BitSet.newBitSet(leafArraySize); + for (int i = 0; i < m; i++) { + BitSet a = sb[i]; + for (int j = i + 1; j < m; j++) { + if (++ops[0] > opsBudget) { + throw BUDGET_EXCEEDED; + } + if (a.intersects(sb[j])) { + continue; + } + long s = partSum[i] + partSum[j]; + long sR = sumC - s; + // boundary-3: complement is one observed subclade d, canonical d > j + boolean ueBuilt = false; + for (int d = subHead[mixHash(sR) & smask]; d >= 0; d = subNxt[d]) { + if (d <= j || partSum[d] != sR) { + continue; + } + if (partCard[i] + partCard[j] + partCard[d] != cardC) { + continue; + } + if (!ueBuilt) { + ue.clear(); + ue.or(a); + ue.or(sb[j]); + ueBuilt = true; + } + if (ue.intersects(sb[d])) { + continue; + } + BitSet[] parts = {sb[i], sb[j], sb[d]}; + int res = countAllNovelResolutions(cBits, parts); + if (res > 0) { + n1 += (novelMode == NovelMode.FLAT) ? res : 1; + } + } + if (computeN2) { + if (P == cap) { + cap <<= 1; + ei = java.util.Arrays.copyOf(ei, cap); + ej = java.util.Arrays.copyOf(ej, cap); + eh = java.util.Arrays.copyOf(eh, cap); + sU = java.util.Arrays.copyOf(sU, cap); + } + ei[P] = i; + ej[P] = j; + sU[P] = s; + eh[P] = mixHash(s); + P++; + } + } + } + + // boundary-4: match each pair with the pair whose union sum is the complement's, k > j. + long n2 = 0; + if (computeN2 && P > 0) { + int tb = Integer.highestOneBit(P) << 1; + int mask = tb - 1; + int[] head = new int[tb]; + java.util.Arrays.fill(head, -1); + int[] nxt = new int[P]; + for (int e = 0; e < P; e++) { + int b = eh[e] & mask; + nxt[e] = head[b]; + head[b] = e; + } + for (int e = 0; e < P; e++) { + int i = ei[e], j = ej[e]; + long sR = sumC - sU[e]; + boolean ueBuilt = false; + for (int f = head[mixHash(sR) & mask]; f >= 0; f = nxt[f]) { + if (++ops[0] > opsBudget) { + throw BUDGET_EXCEEDED; + } + int k = ei[f], l = ej[f]; + if (k <= j) { + continue; + } + if (partCard[i] + partCard[j] + partCard[k] + partCard[l] != cardC) { + continue; + } + if (!ueBuilt) { + ue.clear(); + ue.or(sb[i]); + ue.or(sb[j]); + ueBuilt = true; + } + uf.clear(); + uf.or(sb[k]); + uf.or(sb[l]); + if (ue.intersects(uf)) { + continue; + } + BitSet[] parts = {sb[i], sb[j], sb[k], sb[l]}; + int res = countAllNovelResolutions(cBits, parts); + if (res > 0) { + n2 += (novelMode == NovelMode.FLAT) ? res : 1; + } + } + } + } + return new long[]{n1, n2}; + } + + /** Weighted bit-sum matching {@link Clade#getCladeInBits()}'s {@code hashCode} word weighting + * (word {@code w} weighted by {@code w+1}); additive over disjoint sets, so it lets the + * boundary matcher find a pair's complement by {@code long} arithmetic. Wraparound mod 2^64 is + * harmless — the matcher confirms hits with an exact cardinality/disjointness test. */ + private static long weightedSum(BitSet b) { + long s = 0; + for (int bit = b.nextSetBit(0); bit >= 0; bit = b.nextSetBit(bit + 1)) { + s += ((long) ((bit >>> 6) + 1)) << (bit & 63); + } + return s; + } + + /** Scalar mixer spreading a {@code long} union sum to a hash-table bucket. */ + private static int mixHash(long x) { + long z = x * 0x9E3779B97F4A7C15L; + return (int) (z ^ (z >>> 32)); + } + /* Recursive canonical enumeration: pick (m-1) parts at strictly increasing * indices (= increasing canonical order), then derive the last part as the * complement and require it to be the canonical-largest. */ diff --git a/src/main/java/ccd/tools/SubtreeMarginal.java b/src/main/java/ccd/tools/SubtreeMarginal.java new file mode 100644 index 0000000..4a78aab --- /dev/null +++ b/src/main/java/ccd/tools/SubtreeMarginal.java @@ -0,0 +1,610 @@ +package ccd.tools; + +import beast.base.core.Description; +import beast.base.core.Input; +import beast.base.core.Log; +import beast.base.evolution.tree.Node; +import beast.base.evolution.tree.Tree; +import beast.base.inference.Runnable; +import beastfx.app.tools.Application; +import beastfx.app.treeannotator.TreeAnnotator; +import beastfx.app.util.OutFile; +import beastfx.app.util.TreeFile; + +import java.io.PrintStream; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.IdentityHashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.Set; + +import ccd.model.AbstractCCD; +import ccd.model.CCDType; +import ccd.model.Clade; +import ccd.model.CladePartition; +import ccd.model.ITreeDistribution; +import ccd.model.KRegCCD; +import ccd.model.bitsets.BitSet; + +/** + * Compares, at the resolution of small (3- or 4-taxon) induced subtrees, the + * distributions over topologies implied by the empirical posterior sample and by + * CCD0, CCD1 and (optimised) regCCD built from it. + *

+ * Purpose: to establish that CCD1 is an adequate basis to build on. CCD1 + * conditions each split on its parent clade and so carries the correlation structure + * of the posterior; regularisation then smooths that basis to give it full support. + * The three models are therefore not competitors here. CCD0, which factorises over + * clades and cannot represent those correlations, is the negative control that shows + * the comparison is not trivial; regCCD is CCD1 plus smoothing, and is expected to + * sit slightly further from the posterior than CCD1 does -- a small, uniform + * perturbation is what smoothing is, and it is the price of the coverage + * that {@code miss} measures. The question this tool answers is whether CCD1's + * deviation from the posterior is small relative to the posterior's own sampling + * noise, not which model wins on total variation. + *

+ * Rationale: the marginal posterior probability of a full topology is not + * Monte-Carlo-estimable for non-trivial analyses (almost every sampled tree is + * unique). The induced distribution over the topologies of a fixed small taxon + * subset is, however, low-dimensional and reliably estimable, and interpolates + * between single-clade marginals and the full topology. This tool estimates those + * induced distributions by sampling from each model and restricting every sampled + * tree to each subset. The empirical side additionally reports a split-half + * "noise floor" (posterior first half vs second half) so that a model's deviation + * can be judged against the irreducible sampling scatter. + *

+ * Output is a tab-separated table with one row per (subset, induced shape): + * {@code subset_id, k, taxa, shape, emp, emp_h1, emp_h2, ccd0, ccd1, regccd} + * where each of the last six columns is a probability (frequency) under that + * distribution for that subset. Shapes not seen under a distribution get 0. + */ +@Description("Compare induced 3-/4-taxon subtree topology distributions between the " + + "posterior sample and CCD0/CCD1/regCCD, by sampling and restriction.") +public class SubtreeMarginal extends Runnable { + + final public Input treeInput = new Input<>("trees", + "trees file to construct the CCDs from and analyse", Input.Validate.REQUIRED); + final public Input outputInput = new Input<>("out", + "tab-separated output file", Input.Validate.REQUIRED); + final public Input burnInPercentageInput = new Input<>("burnin", + "percentage of trees used as burn-in (ignored)", 10); + final public Input subsetSizesInput = new Input<>("k", + "comma-separated induced-subtree sizes to analyse", "3,4"); + final public Input subsetsPerSizeInput = new Input<>("subsets", + "number of random taxon subsets drawn per size k", 200); + final public Input sampleSizeInput = new Input<>("length", + "number of trees sampled from each CCD model", 100000); + final public Input foldsInput = new Input<>("folds", + "number of cross-validation folds for regCCD (KRegCCD) parameter selection", 5); + final public Input modelsInput = new Input<>("models", + "also build CCD0/CCD1/regCCD and compare; if false, only the empirical posterior is " + + "summarised (calibration mode)", true); + final public Input exactInput = new Input<>("exact", + "compute EXACT induced-subtree probabilities for CCD0 and CCD1 (sum-product DP over the " + + "clade DAG) instead of sampling; gives true coverage-miss (regCCD is 0 by the " + + "full-coverage theorem). Overrides -models.", false); + final public Input maxTreesInput = new Input<>("maxTrees", + "if > 0, subsample the posterior down to this many trees (uniformly) before analysis", 0); + final public Input seedInput = new Input<>("seed", + "random seed (subset choice and model sampling)"); + final public Input minSubsetsInput = new Input<>("minSubsets", + "minimum number of informative subsets (ones where the posterior itself spreads over " + + "more than one induced shape) required before the CCD models are built and " + + "sampled. Guards against spending hours on a posterior with too little " + + "topological uncertainty to measure; set to 0 to force the run", 1); + + // Distribution column indices in the per-shape count arrays. + private static final int EMP = 0, EMP_H1 = 1, EMP_H2 = 2, CCD0 = 3, CCD1 = 4, REGCCD = 5; + private static final int NDIST = 6; + private static final String[] COLS = {"emp", "emp_h1", "emp_h2", "ccd0", "ccd1", "regccd"}; + + @Override + public void initAndValidate() { + } + + @Override + public void run() throws Exception { + long seed = (seedInput.get() != null) ? seedInput.get() : System.currentTimeMillis(); + Random random = new Random(seed); + + Log.info.println("# Induced subtree-distribution comparison"); + Log.info.println(" trees file: " + treeInput.get().getPath()); + Log.info.println(" burnin: " + burnInPercentageInput.get()); + Log.info.println(" sizes k: " + subsetSizesInput.get()); + Log.info.println(" subsets/k: " + subsetsPerSizeInput.get()); + Log.info.println(" #samples: " + sampleSizeInput.get()); + Log.info.println(" seed: " + seed); + + // Load the posterior sample. + TreeAnnotator.MemoryFriendlyTreeSet treeSet = + CCDToolUtil.getTreeSet(treeInput, burnInPercentageInput.get()); + List posterior = CCDToolUtil.treesFromSet(treeSet); + if (posterior.isEmpty()) { + throw new IllegalArgumentException("No trees left after burn-in."); + } + Log.info.println(" posterior: " + posterior.size() + " trees"); + + // Optionally subsample the posterior (uniformly) for a faster calibration pass. + if (maxTreesInput.get() > 0 && posterior.size() > maxTreesInput.get()) { + Collections.shuffle(posterior, random); + posterior = new ArrayList<>(posterior.subList(0, maxTreesInput.get())); + Log.info.println(" subsampled: " + posterior.size() + " trees"); + } + + // Taxon names from the first tree. + List taxa = new ArrayList<>(); + for (Node leaf : posterior.get(0).getExternalNodes()) { + taxa.add(leaf.getID()); + } + Collections.sort(taxa); + Log.info.println(" taxa: " + taxa.size()); + + // Model construction is deferred until after the empirical tally, so that a posterior with + // no topological uncertainty can be detected and the (expensive) model work skipped -- see + // the pre-flight check below. Their RNG seeds are drawn HERE, at the point in the sequence + // where the models used to be built, so that deferring construction leaves subset selection + // and model sampling bit-identical to before. + AbstractCCD ccd0 = null, ccd1 = null, regccd = null; + final boolean wantModels = modelsInput.get() && !exactInput.get(); + long ccd0Seed = 0, ccd1Seed = 0, regccdSeed = 0; + if (wantModels) { + ccd0Seed = random.nextLong(); + ccd1Seed = random.nextLong(); + regccdSeed = random.nextLong(); + } else { + Log.info.println("> empirical-only calibration mode (no CCD models built)"); + } + + // Draw the taxon subsets for each size k. + List subsets = new ArrayList<>(); + for (String tok : subsetSizesInput.get().split(",")) { + int k = Integer.parseInt(tok.trim()); + if (k < 2 || k > taxa.size()) { + Log.warning("skipping k=" + k + " (out of range for " + taxa.size() + " taxa)"); + continue; + } + drawSubsets(taxa, k, subsetsPerSizeInput.get(), random, subsets); + } + Log.info.println(" subsets: " + subsets.size()); + + if (exactInput.get()) { + runExact(treeSet, posterior, subsets); + return; + } + + // Per-subset shape counts across the six distributions. + int n = subsets.size(); + List> counts = new ArrayList<>(n); + double[] totals = new double[NDIST]; + for (int i = 0; i < n; i++) { + counts.add(new LinkedHashMap<>()); + } + + // Empirical side: iterate every posterior tree (exact frequencies) and, by index + // parity, feed the two split-half distributions used for the noise floor. + Log.info.println("> tallying empirical posterior (" + posterior.size() + " trees)..."); + int empReport = Math.max(1, posterior.size() / 10); + for (int t = 0; t < posterior.size(); t++) { + Tree tree = posterior.get(t); + Node root = tree.getRoot(); + Map leaves = leafMap(tree); + int half = (t % 2 == 0) ? EMP_H1 : EMP_H2; + for (int i = 0; i < n; i++) { + String shape = inducedShape(root, leaves, subsets.get(i).taxaSet); + double[] c = counts.get(i).computeIfAbsent(shape, s -> new double[NDIST]); + c[EMP]++; + c[half]++; + } + if ((t + 1) % empReport == 0) { + Log.info.println(" empirical: " + (t + 1) + "/" + posterior.size()); + } + } + totals[EMP] = posterior.size(); + totals[EMP_H1] = (posterior.size() + 1) / 2; + totals[EMP_H2] = posterior.size() / 2; + + // Pre-flight: a subset carries signal only if the POSTERIOR itself spreads over more than + // one induced shape. If the posterior is (near-)certain at this k, every subset is a point + // mass, the downstream analysis discards all of them, and the model work -- regCCD + // parameter selection plus three 50k-tree sampling runs, hours on a large dataset -- would + // produce a table with nothing in it. Check before paying for it. (Encountered on + // tornabene-2016: 1801 posterior trees, all the SAME topology, so 0 of 1000 subsets were + // usable at k=4 and no larger k would have helped either.) + int nonDegenerate = 0; + for (Map c : counts) { + int shapesSeen = 0; + for (double[] v : c.values()) { + if (v[EMP] > 0) { + shapesSeen++; + } + } + if (shapesSeen >= 2) { + nonDegenerate++; + } + } + Log.info.println(" distinct posterior topologies: " + countDistinctTopologies(posterior) + + " of " + posterior.size() + " trees"); + Log.info.println(" informative subsets (posterior spreads over >1 shape): " + + nonDegenerate + " / " + n); + + boolean buildModels = wantModels; + if (wantModels && nonDegenerate < minSubsetsInput.get()) { + Log.warning("SKIPPING the model comparison: only " + nonDegenerate + " of " + n + + " subsets are informative (need >= " + minSubsetsInput.get() + ")."); + Log.warning("The posterior has too little topological uncertainty at this k for the " + + "induced-subtree comparison to measure anything. Writing the empirical " + + "columns only; model columns will be 0. Raise -minSubsets to force the run."); + buildModels = false; + } + + if (buildModels) { + Log.info.println("> building models..."); + ccd0 = CCDToolUtil.getCCDTypeByName(treeSet, CCDType.CCD0); + ccd1 = CCDToolUtil.getCCDTypeByName(treeSet, CCDType.CCD1); + Log.info.println(" selecting regCCD (KRegCCD) parameters by " + foldsInput.get() + + "-fold cross-validation..."); + regccd = KRegCCD.withOptimisedParameters(posterior, foldsInput.get()); + Log.info.println(" " + regccd); + ccd0.setRandom(new Random(ccd0Seed)); + ccd1.setRandom(new Random(ccd1Seed)); + regccd.setRandom(new Random(regccdSeed)); + } + + // Model side: sample from each CCD and restrict. + if (buildModels) { + tallyModel(ccd0, CCD0, sampleSizeInput.get(), subsets, counts, "CCD0"); + tallyModel(ccd1, CCD1, sampleSizeInput.get(), subsets, counts, "CCD1"); + tallyModel(regccd, REGCCD, sampleSizeInput.get(), subsets, counts, "regCCD"); + totals[CCD0] = totals[CCD1] = totals[REGCCD] = sampleSizeInput.get(); + } + + // Write the table. + Log.info.println("> writing " + outputInput.get().getPath()); + try (PrintStream out = new PrintStream(outputInput.get())) { + out.print("subset_id\tk\ttaxa\tshape"); + for (String col : COLS) { + out.print("\t" + col); + } + out.println(); + for (int i = 0; i < n; i++) { + Subset s = subsets.get(i); + String taxaStr = String.join(",", s.taxaList); + for (Map.Entry e : counts.get(i).entrySet()) { + double[] c = e.getValue(); + out.print(i + "\t" + s.taxaList.size() + "\t" + taxaStr + "\t" + e.getKey()); + for (int d = 0; d < NDIST; d++) { + double p = (totals[d] > 0) ? c[d] / totals[d] : 0.0; + out.print("\t" + p); + } + out.println(); + } + } + } + Log.info.println("... done."); + } + + /** + * Number of distinct topologies in the sample, ignoring branch lengths: a one-line diagnostic + * for why a posterior yielded few informative subsets. A value of 1 means the chain returned a + * single tree, in which case no choice of k can give this comparison anything to measure. + */ + private static int countDistinctTopologies(List posterior) { + Set seen = new HashSet<>(); + for (Tree t : posterior) { + List clades = new ArrayList<>(); + collectClades(t.getRoot(), clades); + Collections.sort(clades); + seen.add(String.join("|", clades)); + } + return seen.size(); + } + + /** Appends a canonical string for every internal clade of the subtree at {@code v}. */ + private static List collectClades(Node v, List out) { + List taxa = new ArrayList<>(); + if (v.isLeaf()) { + taxa.add(v.getNr()); + } else { + for (Node child : v.getChildren()) { + taxa.addAll(collectClades(child, out)); + } + Collections.sort(taxa); + out.add(taxa.toString()); + } + return taxa; + } + + /** Sample {@code length} trees from the distribution and tally induced shapes per subset. */ + private void tallyModel(ITreeDistribution dist, int col, int length, + List subsets, List> counts, String label) { + Log.info.println("> sampling " + length + " trees from " + label + "..."); + int report = Math.max(1, length / 10); + for (int t = 0; t < length; t++) { + Tree tree = dist.sampleTree(); + Node root = tree.getRoot(); + Map leaves = leafMap(tree); + for (int i = 0; i < subsets.size(); i++) { + String shape = inducedShape(root, leaves, subsets.get(i).taxaSet); + double[] c = counts.get(i).computeIfAbsent(shape, s -> new double[NDIST]); + c[col]++; + } + if ((t + 1) % report == 0) { + Log.info.println(" " + label + ": " + (t + 1) + "/" + length); + } + } + } + + /** Leaf-label to leaf-node map for one tree (built once, reused across all subsets). */ + private static Map leafMap(Tree tree) { + Map map = new HashMap<>(); + for (Node leaf : tree.getExternalNodes()) { + map.put(leaf.getID(), leaf); + } + return map; + } + + /** + * Canonical labelled rooted topology induced on {@code keep}, computed without copying or + * mutating the tree. Marks the kept-leaf-to-root paths (O(k * height)), then reads the induced + * topology off the marked nodes, suppressing any node that subtends only one kept leaf. + */ + private static String inducedShape(Node root, Map leaves, Set keep) { + Map below = new IdentityHashMap<>(); + for (String taxon : keep) { + Node x = leaves.get(taxon); + while (x != null) { + below.merge(x, 1, Integer::sum); + x = x.getParent(); + } + } + return inducedKey(root, below); + } + + /** Recurse through nodes subtending a kept leaf; a node with a single such child is suppressed. */ + private static String inducedKey(Node node, Map below) { + if (node.isLeaf()) { + return node.getID(); + } + List relevant = new ArrayList<>(); + for (Node child : node.getChildren()) { + if (below.getOrDefault(child, 0) > 0) { + relevant.add(child); + } + } + if (relevant.size() == 1) { + return inducedKey(relevant.get(0), below); + } + List parts = new ArrayList<>(relevant.size()); + for (Node child : relevant) { + parts.add(inducedKey(child, below)); + } + Collections.sort(parts); + return "(" + String.join(",", parts) + ")"; + } + + /** Draw {@code count} distinct random size-{@code k} subsets of {@code taxa}. */ + private static void drawSubsets(List taxa, int k, int count, Random random, List out) { + Set seen = new HashSet<>(); + int n = taxa.size(); + int attempts = 0, maxAttempts = count * 50 + 1000; + while (seen.size() < count && attempts++ < maxAttempts) { + List idx = new ArrayList<>(); + for (int i = 0; i < n; i++) { + idx.add(i); + } + Collections.shuffle(idx, random); + List chosen = new ArrayList<>(); + for (int i = 0; i < k; i++) { + chosen.add(taxa.get(idx.get(i))); + } + Collections.sort(chosen); + String key = String.join(",", chosen); + if (seen.add(key)) { + out.add(new Subset(chosen)); + } + } + } + + /** A taxon subset, both as an ordered label list and as a lookup set. */ + private static final class Subset { + final List taxaList; + final Set taxaSet; + + Subset(List taxaList) { + this.taxaList = taxaList; + this.taxaSet = new HashSet<>(taxaList); + } + } + + // ===== Exact induced-subtree probabilities for CCD0/CCD1 (sum-product DP over the clade DAG) ===== + + /** + * Exact mode. For each subset, tabulate the empirical induced-shape distribution and the EXACT + * induced probability under CCD0 and CCD1. regCCD is not evaluated: its induced probability is + * positive for every shape (full-coverage theorem), so its coverage-miss is 0 by construction. + * Output: {@code subset_id, k, taxa, shape, emp, ccd0_exact, ccd1_exact, ccd0_massObs, ccd1_massObs} + * where the mass columns are the total CCD probability on the subset's observed shapes, so the + * analysis can form exact TV via the tail term (1 - massObs). + */ + private void runExact(TreeAnnotator.MemoryFriendlyTreeSet treeSet, List posterior, + List subsets) throws Exception { + Log.info.println("> exact induced-probability mode (CCD0, CCD1)"); + AbstractCCD ccd0 = CCDToolUtil.getCCDTypeByName(treeSet, CCDType.CCD0); + AbstractCCD ccd1 = CCDToolUtil.getCCDTypeByName(treeSet, CCDType.CCD1); + String[] taxaNames = ccd1.getSomeBaseTree().getTaxaNames(); + Map nameToIdx = new HashMap<>(); + for (int i = 0; i < taxaNames.length; i++) { + nameToIdx.put(taxaNames[i], i); + } + + int report = Math.max(1, subsets.size() / 10); + try (java.io.PrintStream out = new java.io.PrintStream(outputInput.get())) { + out.println("subset_id\tk\ttaxa\tshape\temp\tccd0_exact\tccd1_exact\tccd0_massObs\tccd1_massObs"); + double nTot = posterior.size(); + for (int si = 0; si < subsets.size(); si++) { + Subset s = subsets.get(si); + BitSet sBits = BitSet.newBitSet(taxaNames.length); + for (String name : s.taxaList) { + sBits.set(nameToIdx.get(name)); + } + // empirical induced distribution + one representative clade-map per distinct shape + Map shapeCount = new LinkedHashMap<>(); + Map> shapeTarget = new HashMap<>(); + for (Tree tree : posterior) { + Map leaves = leafMap(tree); + Map below = markBelow(tree.getRoot(), leaves, s.taxaSet); + Map clades = new HashMap<>(); + targetClades(tree.getRoot(), below, nameToIdx, taxaNames.length, clades); + String key = shapeKeyFromClades(clades); + shapeCount.merge(key, 1, Integer::sum); + shapeTarget.putIfAbsent(key, clades); + } + // exact CCD probabilities per observed shape (restricted-bits cache is target-independent) + Map rb0 = new IdentityHashMap<>(); + Map rb1 = new IdentityHashMap<>(); + Map exact = new HashMap<>(); + double mass0 = 0, mass1 = 0; + for (Map.Entry> e : shapeTarget.entrySet()) { + double p0 = exactInduced(ccd0, sBits, e.getValue(), rb0); + double p1 = exactInduced(ccd1, sBits, e.getValue(), rb1); + exact.put(e.getKey(), new double[]{p0, p1}); + mass0 += p0; + mass1 += p1; + } + String taxaStr = String.join(",", s.taxaList); + for (Map.Entry e : shapeCount.entrySet()) { + double[] px = exact.get(e.getKey()); + out.println(si + "\t" + s.taxaList.size() + "\t" + taxaStr + "\t" + e.getKey() + + "\t" + (e.getValue() / nTot) + "\t" + px[0] + "\t" + px[1] + + "\t" + mass0 + "\t" + mass1); + } + if ((si + 1) % report == 0) { + Log.info.println(" exact: " + (si + 1) + "/" + subsets.size()); + } + } + } + Log.info.println("... done."); + } + + /** Mark, for every ancestor of a kept leaf, how many kept leaves lie below it. */ + private static Map markBelow(Node root, Map leaves, Set keep) { + Map below = new IdentityHashMap<>(); + for (String taxon : keep) { + Node x = leaves.get(taxon); + while (x != null) { + below.merge(x, 1, Integer::sum); + x = x.getParent(); + } + } + return below; + } + + /** + * Build the induced target topology as a map from each internal induced clade (as a BitSet over + * CCD taxon indices) to its two child clades. Returns the clade below {@code node}. + */ + private static BitSet targetClades(Node node, Map below, Map nameToIdx, + int nTaxa, Map cladeChildren) { + if (node.isLeaf()) { + BitSet b = BitSet.newBitSet(nTaxa); + b.set(nameToIdx.get(node.getID())); + return b; + } + List rel = new ArrayList<>(); + for (Node child : node.getChildren()) { + if (below.getOrDefault(child, 0) > 0) { + rel.add(child); + } + } + if (rel.size() == 1) { + return targetClades(rel.get(0), below, nameToIdx, nTaxa, cladeChildren); + } + BitSet mine = BitSet.newBitSet(nTaxa); + BitSet[] kids = new BitSet[rel.size()]; + for (int i = 0; i < rel.size(); i++) { + kids[i] = targetClades(rel.get(i), below, nameToIdx, nTaxa, cladeChildren); + mine.or(kids[i]); + } + cladeChildren.put(mine, new BitSet[]{kids[0], kids[1]}); + return mine; + } + + /** Canonical key of an induced topology: the sorted set of its internal induced clades. */ + private static String shapeKeyFromClades(Map cladeChildren) { + List keys = new ArrayList<>(cladeChildren.size()); + for (BitSet b : cladeChildren.keySet()) { + keys.add(b.toString()); + } + Collections.sort(keys); + return String.join("|", keys); + } + + /** Exact probability that the CCD induces {@code cladeChildren} on the subset {@code sBits}. */ + private static double exactInduced(AbstractCCD ccd, BitSet sBits, + Map cladeChildren, Map rbCache) { + return gExact(ccd.getRootClade(), sBits, cladeChildren, rbCache, new IdentityHashMap<>()); + } + + /** + * Sum-product recursion: probability that the CCD subtree below clade {@code c} induces, on + * {@code c}'s S-restriction, the corresponding subtree of the target. Memoised per target. + */ + private static double gExact(Clade c, BitSet sBits, Map cladeChildren, + Map rbCache, Map memo) { + Double cached = memo.get(c); + if (cached != null) { + return cached; + } + BitSet a = restricted(c, sBits, rbCache); + double result; + if (a.cardinality() <= 1) { + result = 1.0; // empty or singleton restriction: trivially induced + } else { + BitSet[] kids = cladeChildren.get(a); + if (kids == null) { + result = 0.0; // c's restriction is not a clade of the target + } else { + BitSet left = kids[0], right = kids[1]; + double total = 0.0; + for (CladePartition part : c.getPartitions()) { + Clade[] cc = part.getChildClades(); + BitSet a1 = restricted(cc[0], sBits, rbCache); + BitSet a2 = restricted(cc[1], sBits, rbCache); + double p = part.getCCP(); + if (a1.isEmpty()) { + total += p * gExact(cc[1], sBits, cladeChildren, rbCache, memo); + } else if (a2.isEmpty()) { + total += p * gExact(cc[0], sBits, cladeChildren, rbCache, memo); + } else if ((a1.equals(left) && a2.equals(right)) + || (a1.equals(right) && a2.equals(left))) { + total += p * gExact(cc[0], sBits, cladeChildren, rbCache, memo) + * gExact(cc[1], sBits, cladeChildren, rbCache, memo); + } + } + result = total; + } + } + memo.put(c, result); + return result; + } + + /** S-restriction of a clade's taxon set, cached per clade (target-independent). */ + private static BitSet restricted(Clade c, BitSet sBits, Map cache) { + BitSet r = cache.get(c); + if (r == null) { + r = (BitSet) c.getCladeInBitsTaxaOnly().clone(); + r.and(sBits); + cache.put(c, r); + } + return r; + } + + public static void main(String[] args) throws Exception { + new Application(new SubtreeMarginal(), "Subtree Marginal Comparison", args); + } +} diff --git a/version.xml b/version.xml index 7503b38..cacb6cc 100644 --- a/version.xml +++ b/version.xml @@ -1,5 +1,6 @@ +