Search 17,000 array-measured genes with sparse RJMCMC

Name the analysis precisely

This is transcriptome-wide association, not genetically predicted-expression TWAS

The example associates directly measured Affymetrix expression with a clinical phenotype. That makes it an array-based transcriptome-wide biomarker association analysis. In contemporary human genetics, “TWAS” often means testing genetically predicted expression using an expression-reference panel and GWAS summary statistics; this example does not make that claim. Its purpose is to exercise real high-dimensional model selection, repeated observations, covariate adjustment, accelerators, and restart safety.

Public input · downloaded by the user

Prepare GSE93272 from GEO

GSE93272 contains 275 human Affymetrix Human Genome U133 Plus 2.0 samples and extensive rheumatoid-arthritis clinical characteristics. GEO exposes the processed Series Matrix publicly; raw CEL files are much larger and are not needed for this computational example.

# Install once in R:
BiocManager::install(c("GEOquery", "Biobase", "AnnotationDbi", "hgu133plus2.db"))

# From the JDistlib repository root:
Rscript examples/gse93272/prepare-gse93272.R build/example-data/gse93272

The checked-in script downloads the GPL570 Series Matrix, retains complete rheumatoid-arthritis records for standardized CRP.DAS28, age, sex, batch, and RIN, maps probes to gene symbols, chooses the highest-IQR probe per symbol, and retains the top 17,000 genes by outcome-blind IQR. It writes clinical.tsv, feature-major expression.tsv.gz, and a provenance manifest. The Java loader verifies the exact sample order.

Mixed-effect marginal likelihood

Integrate the subject intercept analytically

For sample i from subject s, the normalized target is:

y[i] = alpha + age[i] b_age + female[i] b_female
       + batch2[i] b_batch + rin[i] b_rin
       + sum(x[i,j] beta[j], j in M) + u[s] + error[i]

u[s]     ~ Normal(0, tau)
error[i] ~ Normal(0, sigma)
alpha, clinical coefficients ~ Normal(0, 2.5)
log(sigma) ~ Normal(-0.5, 1); log(tau) ~ Normal(-1, 1)
beta[j] ~ Normal(0, 1)

The example integrates each Gaussian u[s] out exactly, so a state stores seven common parameters plus only the active gene coefficients. This is a true random-intercept mixed model without adding one latent value per subject. The complete determinant, quadratic form, and normalizing constants remain in the target.

The model size has a Poisson(3) prior truncated to 0–20. Conditional on size k, every one of the choose(17000, k) subsets is equally likely. The -log choose(p,k) term is essential: omitting it would assign the wrong prior odds to model sizes.

Scalable state and proposals

Store twenty indices, not a 17,000-bit model identifier

SparseSubsetState holds sorted integer candidate indices and their coefficients. Candidate count is independent of the old 62-bit subset identifier; the maximum active count is checked separately. Add, drop, and swap moves include the exact candidate density, coefficient birth density, and boundary-dependent reverse move probability.

SparseSubsetTarget target = new SparseSubsetTarget(
    commonNames, geneNames, 20, mixedModel);

SparseSubsetRjSampler sampler = new SparseSubsetRjSampler(
    residualInformedCandidates,
    new GaussianSparseCoefficientProposal(0, 1),
    0.05, 0.30);

The locally informed proposal ranks inactive genes using |X' V^-1 r|. It mixes 10% uniform probability into every proposal, preserving support for all inactive candidates, and evaluates the same state-dependent density in the reverse calculation. This changes efficiency, not the posterior target.

Bounded work units

Run one or more checkpointed segments

# Compile the core, complete accelerator JAR, and examples
./gradlew :jdistlib-all:jar compileDocumentationExamples

# macOS/Linux: force CPU for the first smoke run
java -cp "distribution/all/build/libs/*:build/documentation-examples" \
  examples.WorkedSparseTranscriptomeRjmcmcExample \
  --data=build/example-data/gse93272 --compute=cpu \
  --chain=1 --segments=1 --segment-transitions=1000 --warmup=50000 --thin=10

# Windows PowerShell: use the available accelerator automatically
java -cp "distribution/all/build/libs/*;build/documentation-examples" `
  examples.WorkedSparseTranscriptomeRjmcmcExample `
  --data=build/example-data/gse93272 --compute=auto `
  --chain=1 --segments=10 --segment-transitions=10000 --warmup=50000 --thin=10

Run a short CPU smoke first. Then launch independent output directories with different --chain values. A segment is a durability and scheduling unit, not a new chain: rerunning the same command detects sparse-rj.checkpoint and continues. Changing data, warmup, thinning, maximum model size, concrete backend, or chain number causes the fingerprint check to reject an unsafe continuation.

Reboot-safe continuation

Checkpoint the entire stochastic process

The versioned checkpoint contains the sparse state, current normalized log joint, 64-bit transition and retained-draw counters, exact random-engine state, move weights and counts, per-model-size random-walk adaptation, and online model-size, inclusion, coefficient, and common-parameter moments. Warmup may be interrupted too: continuation restores adaptation rather than silently restarting or freezing it.

1. Write a deterministic draw segmentFlush and force the temporary tidy TSV, then atomically replace draws-FIRST-END.tsv.
2. Write the checkpointSerialize, checksum, flush, force, and atomically replace sparse-rj.checkpoint.
3. Recover by replayIf power fails between those commits, the old checkpoint deterministically replays and replaces the same segment. If it fails during either temporary write, the last committed pair remains usable.

Model and option fingerprints include SHA-256 hashes of both prepared inputs. The loader rejects corrupt checksums, dimensions that disagree with the target, and unexpected serialized classes. Keep the output on a local filesystem with reliable atomic rename semantics; network filesystems can weaken those guarantees.

Honest acceleration boundary

Keep the matrix resident; do the small state work on CPU

The expensive proposal statistic multiplies a roughly samples-by-17,000 expression matrix by one or more score vectors. CUDA and OpenCL implementations upload that matrix once and execute prepared X'v products on device. Compute.AUTO uses the accelerator above its conservative work threshold. The sparse state mutation, mixed-model log joint, accept/reject decision, RNG, and checkpoint remain on CPU; Vulkan currently uses the prepared CPU fallback for this primitive.

A GeForce RTX 2080 has ample memory for this matrix, but GPU acceleration does not guarantee an end-to-end speedup. With only about 100–275 rows, kernel launch and result transfer can dominate, and every proposal still requires CPU likelihood work. Measure transitions per second after warmup on the actual prepared dataset with --compute=cpu and --compute=cuda. Extrapolate only from sustained measurements: runtime = requested transitions / transitions per second. Independent chains can run in separate processes, but concurrent GPU chains compete for the same device.

Scientific and computational review

Fit multiple chains before interpreting markers

summary.json is an online occupancy summary and each retained segment is a tidy ragged table. Compare independent chains for model-size occupancy and top-gene PIPs; inspect add/drop/swap acceptance, active-set turnover, and conditional coefficient stability. A long chain that remains in one correlated gene neighborhood is not convincing merely because it contains many transitions.

• Require repeated visits and reasonable effective information for inclusion indicators, not only a ranked PIP list.
• Examine correlated probe/gene groups: posterior mass may be distributed across substitutes.
• Perform sensitivity analyses for the model-size prior, active cap, birth scale, informed/uniform mixture, and preprocessing.
• Hold back outcome data or use external replication before calling a selected gene a biomarker.
• Treat the 20-gene cap as a computational/scientific prior and verify that posterior model size does not pile up at 20.

Troubleshooting

Start small and preserve evidence