Mixture prior estimation for MASH

Mixture prior estimation for MASH#

Estimates the data-driven mixture of covariance matrices that a MASH model uses as its prior, from factor analyses and canonical patterns through residual-variance estimation to the fitted mixture.

Overview#

MASH borrows strength across conditions by putting a prior on effect-size vectors: a mixture of covariance matrices, each describing one pattern of sharing, such as an effect present in a single condition, shared equally everywhere, or shared with condition-specific magnitudes. This module estimates that mixture from the data before any MASH model is fitted, so one prior can serve several downstream analyses.

The work proceeds in four steps, in order.

Covariance components. flash and flash_nonneg derive candidate patterns by empirical Bayes factor analysis, the second constraining the factors to be non-negative; pca derives them from the leading principal components of the strong effects; canonical writes the fixed patterns that need no fitting, namely an identity matrix, one matrix per condition, equal effects, and a few simple heterogeneity models. All four are produced by default, which is what --mixture-components lists.
Residual variance. Correlated noise across conditions would otherwise be read as shared signal, so a residual covariance Vhat is estimated from the null effects. Five estimators are available: identity assumes no correlation, simple takes the empirical correlation of the null z-scores, simple_specific takes their covariance with a nearest-positive-definite correction, corshrink applies adaptive shrinkage to that correlation, and mle refines an estimate by EM on a random subset of the data.
Mixture prior. One of three engines fits the mixture weights and refines the data-driven components. ed_bovy is the extreme deconvolution exported by mashr; ud and ud_unconstrained are the ED and TED updates from udr, and the code marks both as opt-in, ud for known numerical issues and ud_unconstrained because it expects independent, identically distributed data on the z-score scale. The engine reads the component files from step 1 together with the Vhat estimate from step 2.
Plotting. plot_U draws the estimated covariance matrices as heatmaps ordered by their mixture weight.

Two settings tie the steps together. --effect-model decides whether effects are exchangeable on the effect-size scale (EE) or the z-score scale (EZ), and it appears in every file name from step 2 onwards. --vhat names which residual-variance estimate the engine should read, so it has to match the estimator that was actually run; the default is simple. mle is the exception in both directions: it needs a prior to exist already, and when --vhat mle is set the engines fall back to reading the simple estimate instead.

When to run it. After the per-condition summary statistics have been assembled into a single RDS by the effect-extraction step, and before fitting the MASH model, which consumes the prior produced here.

Input#

--data – the assembled summary statistics, an RDS holding effect sizes (.b), standard errors (.s) and z-scores (.z) for the strong, random and null subsets, plus XtX, each with one column per condition:

List of 10
 $ random.z: num [1:68, 1:8] -0.286 0.909 1.362 -0.807 -0.185 ...
 $ strong.z: num [1:17, 1:8] 8.77 1.16 9.16 -4.99 -1.02 ...
 $ null.z  : num [1:68, 1:8] -0.508 -1.149 -0.422 -0.206 0.351 ...
 $ random.b: num [1:68, 1:8] -0.00799 0.0248 0.03869 -0.02005 -0.00388 ...
 $ strong.b: num [1:17, 1:8] 0.2491 0.0927 0.3318 -0.1553 -0.1398 ...
 $ null.b  : num [1:68, 1:8] -0.03792 -0.02839 -0.03397 -0.00571 0.03438 ...
 $ null.s  : num [1:68, 1:8] 0.0746 0.0247 0.0805 0.0278 0.0979 ...
 $ random.s: num [1:68, 1:8] 0.028 0.0273 0.0284 0.0249 0.021 ...
 $ strong.s: num [1:17, 1:8] 0.0284 0.0798 0.0362 0.0311 0.1374 ...
 $ XtX     : num [1:8, 1:8] 31.06 6.09 1.02 27.12 9.89 ...

--output-prefix – prefix for every output file; derived from the --data basename when omitted.
--cwd – output directory, default mashr_workflow_output.
--effect-model – EE or EZ, default EE.
--vhat – which residual-variance estimate the mixture engines read: identity, simple, mle, vhat_corshrink_xcondition or vhat_simple_specific. Default simple.
--mixture-components – which component files the engines combine, default flash flash_nonneg pca canonical.
--output-suffix – suffix carried on output names, default all.

Step-specific:

--npc (pca) – number of principal components, default 2. The comment beside it in the code says it is set to 3 as in the mash paper, which no longer matches the value.
--n-subset and --max-iter (vhat_mle) – rows sampled for the EM estimate, default 6000, and the iteration cap, default 6.
--max-comp, --to-cor, --tol, --remove-label, --name (plot_U) – how many components to draw (-1 for all), whether to convert covariances to correlations, the tolerance, whether to drop labels, and a suffix added to the plot file name. Note that plot_U redefines --data to mean the prior RDS rather than the summary statistics.

Runtime:

--container, --job-size, --walltime, --mem, --numThreads – software environment and cluster resources.
--modular-script-dir – location of the R drivers, default code/script.

Output#

<prefix>.flash.rds, <prefix>.flash_nonneg.rds, <prefix>.pca.rds, <prefix>.canonical.rds – the covariance components from step 1, each a named list of condition-by-condition matrices. On the toy data pca yields PCA_1, PCA_2 and tPCA and canonical yields 13 matrices, while both FLASH files are empty (NULL), so FLASH contributes nothing to the toy prior.
<prefix>.<effect_model>.V_<method>.rds – the residual variance estimate from step 2, one condition-by-condition matrix. The method is part of the name, so estimates from different estimators coexist.

<prefix>.<effect_model>.prior.rds – the prior itself: the list of covariance matrices U, their mixture weights w, and the log-likelihood trace of the fit. Components that received zero weight are kept in the list:

List of 3
 $ U     :List of 17
 $ w     : Named num [1:17] 0.4714 0 0 0.0616 0.0949 ...
 $ loglik: num [1:1447] 0.705 3.455 3.719 4.142 5.068 ...
names(U): XtX, PCA_1, PCA_2, tPCA, identity, ALL, Ast, End, Exc, Inh, Mic, OPC, Oli, equal_effects, simple_het_1, simple_het_2, simple_het_3

<prefix>.<effect_model>.prior_c.log and <prefix>.<effect_model>.prior_loglike.log – engine diagnostics.
<plot_prefix>.pdf – heatmaps of the covariance matrices, ordered by mixture weight.

Each step also writes .stdout and .stderr beside its output. Example results for the toy data are under output/mixture_prior/.

Minimal Working Example#

Step 1. Covariance components#

All four run independently of each other, and all four are expected by the engines in step 3, since --mixture-components defaults to the full set.

FLASH factors#

Empirical Bayes factor analysis of the strong effects.

Timing: ~5-15 min (on toy dataset)

sos run pipeline/mixture_prior.ipynb flash \
    --output-prefix protocol_example \
    --data input/mash/protocol_example.mashr_input.rds \
    --cwd output/mixture_prior