Select covariates with reversible-jump MCMC

Java example to command line

Run the same eight-model statistical target without writing Java

The checked-in CLI example files reproduce the Java example's sixteen observations, three candidates, observation SD 0.75, intercept prior N(0,5), coefficient priors N(0,2), and independent inclusion probability 0.30. Eight ordinary Stan files describe the subset models. A model-space JSON describes their probabilities and transitions. Sampling controls and data paths remain on the command line.

1. Build and select the executable JAR

From the JDistlib repository root in PowerShell:

.\gradlew.bat :jar :jdistlib-all:jar
$jar = (Get-ChildItem distribution/all/build/libs/jdistlib-all-*-SNAPSHOT.jar | Select-Object -First 1).FullName

The all-in-one JAR needs no companion libraries. The core JAR in build/libs/ needs its generated lib/ directory alongside it. Build from the current source checkout; older published JARs do not contain this new mode. On Linux/macOS, use ./gradlew and pass the generated JAR path to java -jar.

2. Define models and paired transitions

selection.rj.json contains all eight subsets. The intercept-only subset has prior probability 0.7^3 = 0.343; dose-only has 0.3 * 0.7^2 = 0.147. These are probabilities of individual models, not model sizes. A complete smaller two-model document illustrates the syntax:

{
  "schema": "jdistlib.rj/1",
  "models": {
    "intercept": {"script": "intercept.stan", "prior_probability": 0.7},
    "dose": {"script": "dose.stan", "prior_probability": 0.3}
  },
  "transitions": [{
    "type": "add_drop",
    "from": "intercept",
    "to": "dose",
    "copy": {"alpha": "alpha"},
    "add": {"dose": {"proposal": "normal", "mean": 0, "sd": 2}},
    "space": "unconstrained"
  }]
}

Script paths resolve relative to this JSON file. The add move copies alpha and proposes dose; its paired drop move removes dose and evaluates the reverse proposal density. The adapter supplies the unit insertion/removal Jacobian, and the sampler includes forward/reverse move-selection probabilities. The Normal proposal above is distinct from the parameter prior defined inside dose.stan.

3. Load data and validate the complete model space

java -jar $jar --run --sampler rjmcmc --model-space examples/cli/rjmcmc/selection.rj.json --data examples/cli/rjmcmc/observations.json --validate

Expected: Validated RJ model space: 8 models, 24 directed moves. observations.json contains N, the row-major matrix X, and vector y. To load separate files, replace its --data argument with:

--set N=16 --data X=examples/cli/rjmcmc/design.csv --data-column y=examples/cli/rjmcmc/response.csv:response

Repeat --data file.json to merge different variables from several files. CLI data paths resolve from the working directory. Each model receives its declared inputs; variables unused by every model, duplicate names, missing values, invalid dimensions, unknown JSON fields, incomplete mappings and disconnected transition graphs are rejected.

4. Sample and save checkpoints

java -jar $jar --run --sampler rjmcmc --model-space examples/cli/rjmcmc/selection.rj.json --data examples/cli/rjmcmc/observations.json --chains 4 --threads 4 --warmup 1500 --samples 4000 --seed 20260829 --within-scale 0.20 --within-target-accept 0.30 --target-jump-accept 0.25 --output build/rj-cli/results.txt --checkpoint build/rj-cli/state

--samples counts retained draws per chain. --thin 2 doubles post-warmup iterations. Each chain owns its adaptive within-model random walk, and adaptation freezes after warmup. Birth proposals remain fixed. Directed edges start with equal weights; --adapt-move-weights false disables weight adaptation. Chains start at default parameter values in successive declared models. Fresh-run seeded results are independent of worker count.

The full CLI example has twelve paired add/drop edges. The Java example additionally proposes direct swaps; the schedules share a statistical target but may mix differently and do not produce identical seeded draws. Independent tests compare each CLI model's normalized log joint with the Java formula and compare sampled probabilities with exact Gaussian marginal likelihoods.

5. Inspect model probabilities and inclusion diagnostics

results.txt is tidy CSV with columns chain,draw,model_id,model_name,kind,parameter,value,log_joint. An absent coefficient has no row; it is not an observed zero. The # diagnostics= comment contains JSON with model occupancy, model/inclusion ESS, R-hat and MCSE, move acceptance and invalid proposals, model changes, and parameter summaries conditional on presence. Extract it in PowerShell:

$line = Get-Content build/rj-cli/results.txt | Where-Object { $_.StartsWith('# diagnostics=') }
$report = $line.Substring('# diagnostics='.Length) | ConvertFrom-Json
$report | ConvertTo-Json -Depth 20

Parameters sharing a name must have the same meaning and scale across models. A parameter present in only some models receives an inclusion indicator. The catalog lists all declared models; the diagnostic report lists visited models, with separate warnings for unvisited models. Inspect mixing warnings: zero observed occupancy is not proof of zero posterior probability, and a single chain cannot estimate R-hat.

6. Continue the saved chains

The command above writes state.chain-1.rjckpt through state.chain-4.rjckpt. Restore with the same model/data, chain count, thinning and kernel controls:

java -jar $jar --run --sampler rjmcmc --model-space examples/cli/rjmcmc/selection.rj.json --data examples/cli/rjmcmc/observations.json --chains 4 --samples 1000 --within-scale 0.20 --within-target-accept 0.30 --target-jump-accept 0.25 --resume build/rj-cli/state --output build/rj-cli/continued.txt --checkpoint build/rj-cli/continued-state

Omit --seed: saved random streams are restored, with zero new warmup. Resume currently executes chains sequentially and writes only new draws. Fingerprints reject changed model/data or kernel settings and swapped chain files. Tests verify exact continuation of sampled states on the same runtime. Generated quantities use a cloned post-run stream, so their random values can differ between split and uninterrupted exports. Use a new checkpoint prefix to preserve the previous run; output files are staged then replaced individually, not as a single multi-file transaction.

Current CLI scope

Schema 1 supports finite enumerated model families, at least one unbounded scalar real parameter per model, name-preserving copy maps, and independent Normal proposals for added coordinates. Bounds, vectors, offsets/multipliers, renaming, custom Jacobians, split/merge and direct swaps require other adapters. Large sparse candidate spaces still use the Java SparseSubsetRjSampler API. CLI within-model updates are CPU adaptive random walks; NUTS-specific flags fail in RJ mode. Existing --run --input NUTS and three-argument source generation remain available.

Run --run --sampler rjmcmc --help for controls, or read the complete CLI guide.

Scientific question

Which predictors belong in the regression?

The checked-in example has sixteen observations and three candidate predictors: dose, genotype, and a correlated proxy_marker. The response was generated from the first two, while the proxy creates a realistic competing-locus model that exercises swap moves. A model identifier is a bit mask, so model 3 contains dose and genotype while model 0 contains only the common intercept.

For subset M, the model is:

y[i] ~ Normal(alpha + sum(beta[j] * x[i,j], j in M), 0.75)
alpha ~ Normal(0, 5)
beta[j] ~ Normal(0, 2)
Pr(j in M) = 0.30 independently

Correct target

Keep every model-dependent constant

SubsetSelectionTarget takes a Java callback that receives common parameters, sorted active candidate indices, and active coefficients. Unlike an ordinary within-model target, it must return the complete normalized log joint. Constants that differ between models directly affect posterior model probabilities.

SubsetSelectionTarget target = new SubsetSelectionTarget(
    new String[] {"alpha"},
    new String[] {"dose", "genotype", "proxy_marker"},
    (common, active, coefficients) -> {
        double result = normalLogDensity(common[0], 0, 5);
        result += active.length * Math.log(0.30)
            + (3 - active.length) * Math.log(0.70);
        for (double beta : coefficients)
            result += normalLogDensity(beta, 0, 2);
        // Add every normalized observation density here.
        return result;
    });

Detailed balance

Add, drop, swap, then update within the model

SubsetSelectionRj schedules three reversible structure moves. Add chooses uniformly among inactive predictors and samples a new coefficient from its declared birth density. Drop chooses uniformly among active predictors and evaluates the matching reverse birth density. Swap chooses one of each and preserves dimension. The sampler also includes the probability of selecting each move among those available at the current and proposed boundaries.

GaussianRjBirthProposal birth =
    new GaussianRjBirthProposal(0, 2);
ReversibleJumpSampler sampler =
    SubsetSelectionRj.sampler(birth, 0.20, 0.30);

The general engine evaluates:

log alpha = logJoint(proposed) - logJoint(current)
          + log q(reverse) - log q(forward)
          + log p(select reverse) - log p(select forward)
          + logAbsJacobian

CoordinateInsertionTransformation supplies the unit-Jacobian birth/death map and DimensionMatchingValidator checks total dimensions, inverse round trips, and reciprocal Jacobians. Custom ReversibleJumpMove implementations can supply nonzero Jacobians and arbitrary Java model structures.

Four independent chains

Warm up proposals, freeze them, and retain ragged draws

ReversibleJumpSamplingOptions options =
    ReversibleJumpSamplingOptions.builder()
        .warmupIterations(1500)
        .sampleIterations(4000)
        .targetJumpAcceptance(0.25)
        .adaptMoveWeights(true)
        .build();

ReversibleJumpResult chain = sampler.sample(
    target,
    target.state(0, new double[] {0}, new double[0]),
    options,
    new MersenneTwister(20260829));

Move weights, per-model random-walk scales, and adaptive birth proposals may change only during discarded warmup. They are frozen before retained sampling. Use a separate sampler instance and explicit random engine for each independent chain.

Compile and run the complete file

The classpath wildcard stays valid across JDistlib releases.

./gradlew compileDocumentationExamples
java -cp "build/libs/*:build/documentation-examples" \
  examples.WorkedReversibleJumpSelectionExample

# Windows
gradlew.bat compileDocumentationExamples
java -cp "build\libs\*;build\documentation-examples" examples.WorkedReversibleJumpSelectionExample

An optional first argument selects the output directory. The default is build/example-output/rjmcmc.

Review the trans-dimensional chain

Model occupancy alone is not enough

ReversibleJumpDiagnostics.analyze combines chains and reports posterior model and candidate inclusion probabilities with ESS/MCSE/R-hat, retained transition counts, round trips between the two most visited models, add/drop/swap acceptance and invalid proposals, and parameter summaries conditional on presence.

ReversibleJumpDiagnosticReport report =
    ReversibleJumpDiagnostics.analyze(target, chains);

double[] inclusion = report.inclusionProbabilities();
double[] inclusionEss = report.inclusionEffectiveSampleSizes();
double[] inclusionRhat = report.inclusionRHats();

Investigate a candidate with low inclusion ESS or R-hat above 1.01, a move direction with near-zero acceptance, no round trips, models seen in only one chain, or active coefficients with very few conditional draws. The example writes rjmcmc-diagnostics.json using schema jdistlib.rjmcmc-diagnostics/1 and rjmcmc-chain-1.csv in tidy ragged form.

Reproducibility

Resume the exact frozen process

The portable checkpoint includes the model identifier, variable-length parameter vector, log joint, completed iteration, cloned random stream, frozen move weights, and every component's adaptation state. Model and option fingerprints reject incompatible restores.

ReversibleJumpCheckpointIO.write(path, chain.checkpoint(),
    "worked-rj-linear-v1", "warmup-1500-samples-4000-v1");

PortableReversibleJumpCheckpoint restored =
    ReversibleJumpCheckpointIO.read(path,
        "worked-rj-linear-v1", "warmup-1500-samples-4000-v1");

ReversibleJumpResult continued = newSampler.resume(
    target, restored.checkpoint(),
    ReversibleJumpSamplingOptions.builder()
        .warmupIterations(0).sampleIterations(250).build());

Scientific sources

Methods used in this worked example