Add mixed models to Milo differential abundance testing - #1069
Merged
Conversation
Closes #747 A random intercept in the design switches da_nhoods to a negative binomial mixed model, so repeated measurements of the same donor, batch or timepoint are no longer treated as independent samples. The (1 | variable) syntax and the automatic switch follow R Milo. The model is fitted by pseudo-likelihood: every iteration linearises the negative binomial likelihood into a working response, solves the weighted mixed model for the fixed effects and the BLUPs, and updates the variance components by Fisher scoring. Neighbourhoods in which a group has no cells have no finite estimate and are reported as NaN rather than as an arbitrarily large fold change.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1069 +/- ##
==========================================
+ Coverage 74.15% 79.33% +5.17%
==========================================
Files 52 53 +1
Lines 7291 7505 +214
==========================================
+ Hits 5407 5954 +547
+ Misses 1884 1551 -333
🚀 New features to boost your workflow:
|
Solve against the Cholesky factor of the variance matrix instead of building its pseudoinverse, and take the score and Fisher information traces in the space of random effect levels rather than of samples. The estimates are unchanged to machine precision and a fit is 18x faster, which matters because a neighbourhood is fitted one at a time. Testing a coefficient that is constant within every level of a random intercept against the number of samples counts repeated measurements as independent observations. Under the null that inflated the type I error to 6.8-8.0% at alpha=0.05, so the degrees of freedom now follow the between-within rule, which brings it to 3.8-5.8%.
The outer products of the random effect matrices are identical for every neighbourhood, so build them once instead of per fit, and start the refit that follows the dispersion update from the previous solution rather than from ordinary least squares. Also covers fitting two crossed random intercepts, which until now only the formula parsing was tested for. Averaged over replicates each variance component recovers its planted value.
…spend The dispersion solved the Pearson equation against n - p degrees of freedom, but the means it is computed from come from a fit that also spent degrees of freedom on the random effects, so the residuals are smaller than that and the dispersion came out 3-8% below the planted value. Too small a dispersion understates the variance of the counts and with it the standard error. Subtracting the effective degrees of freedom of the random effects takes the worst type I error under the null from 7.5% to 6.2% at alpha=0.05, and power is unchanged.
Running da_nhoods twice dropped every column of the new result if any one of them was already present, so a second run with a solver that reports a different set of columns raised a KeyError on a column the first run never wrote. Switching between a mixed model and a fixed effects fit hit this because they report different columns. Track what each run writes and clear exactly that, which also stops a mixed model result from leaving its variance components behind in a fixed effects result.
…trasts The differential expression models build their design with FormulaicContrasts, so the mixed model now does too instead of reaching for formulaic directly. That also drops the extra dependency this needed and with it a mypy override. Since the design now carries its column names, model_contrasts works for a mixed model as well: the R style contrast is matched against the coefficients and tested as a linear combination, rather than being rejected. The degrees of freedom rule reads the contrast too, so a contrast that is constant within a random effect level is still tested against the number of levels. The Returns section broke a bullet list without a blank line, which failed the docs build with warnings as errors.
subset_samples reduced the design frame in place and each solver then reduced it again with a mask of the original length, which raised a KeyError for a mixed model. The same double reduction sits in the pydeseq2 and edgeR paths, where it only shows once a sample has no counts at all. The frame is now reduced once, right after the samples to keep are known, and unused categories are dropped whatever caused a sample to go, not only an explicit subset.
A moment estimator has to be told how many degrees of freedom the fit spent, and getting that wrong leaves the dispersion biased by an amount that depends on the size of the design: it came out between 1% and 16% above the planted value depending on the number of donors and samples. The dispersion is reported to the user, so that drift matters on its own. Maximising the Cox-Reid adjusted profile likelihood instead, with the penalised information of the mixed model coefficients, holds the bias at -2% to -6% across every design tried. The dispersion and the fit are iterated to consistency, because the adjustment is evaluated at the fitted means and those depend on the dispersion in turn. Type I error is unchanged: over 1500 replicates the two estimators deviate from the nominal 5% by 0.48 and 0.52 points on average.
A random intercept says donors differ in how many cells fall in a neighbourhood. It cannot say that the condition itself lands differently from donor to donor, which is the other half of what a mixed model is usually wanted for, and R Milo has no syntax for it either. `(variable | group)` now adds that, whether the variable is numeric or categorical, in which case each level beyond the reference gets its own variance. The variance matrix was already linear in the variance components, so a slope only contributes another indicator scaled by the variable and needs nothing new from the fit itself. Intercept and slope carry separate variances rather than being allowed to covary. Over replicates a planted intercept variance of 0.3 and slope variance of 0.4 come back as 0.292 and 0.400 at 30 donors, and a slope variance of zero comes back as 0.016.
They arrived last, they have parameter recovery tests but have not been through the null calibration study or either real dataset, and unlike everything else here there is no reference implementation to check them against, since R Milo has no syntax for them. Better on their own once they have had the same treatment as the rest.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #747
A random intercept in the design switches
da_nhoodsto a negative binomial mixed model, so repeated measurements of the same donor, batch or timepoint are no longer treated as independent samples:The
(1 | variable)syntax and the automatic switch on seeing a random effect both follow R Milo. Results gainSE,tvalue, one<variable>_varianceper random intercept,Dispersion,LogliklihoodandConverged.Three bugs turned up along the way that are independent of this feature and are fixed here as well.
da_nhoodsdropped every column of a new result if any one of them was already present, so a second run with a solver reporting a different set of columns raised aKeyErroron a column the first run never wrote.subset_samplesreduced the design frame and then each solver reduced it again with a mask of the original length. And a random effect with a single level let a raw LAPACK error escape. Say the word if you would rather have those as their own pull request.Implementation
R Milo's GLMM is a custom pseudo-likelihood NB-GLMM (36 KB of R over 113 KB of C++), not lme4, so this is an independent implementation of the same method rather than a port: each iteration linearises the negative binomial likelihood into a working response, solves the weighted mixed model for the fixed effects and the BLUPs, and updates the variance components by Fisher scoring.
Neighbourhoods where a group has no cells are completely separated and have no finite estimate. R Milo screens these with
checkSeparation; here they are reported as NaN and flagged unconverged rather than as an arbitrarily large fold change, which is what they degenerate to otherwise.Agreement with miloR 2.6.0
Same simulated data fitted by both implementations (20 donors x 8 samples, planted sigma = 1.0, dispersion = 0.15):
Effect estimates agree with miloR essentially exactly, and so does the t-value ranking that decides which neighbourhoods are called significant. Standard errors come out about 2.5x smaller than miloR's, which traces to miloR's dispersion estimate landing ~5x above the planted value rather than to the standard error formula.
Calibration
Agreement with another implementation says nothing about whether the p-values are correct, so they were checked against the null: data simulated with no condition effect, counting how often a neighbourhood is called significant at alpha = 0.05. 1500 replicates per scenario, so the Monte Carlo standard error is about 0.56 points.
Testing a coefficient that is constant within every level of a random intercept against the number of samples counts repeated measurements as independent observations, which is the pseudo-replication a mixed model exists to prevent. It matters most where there are few donors carrying many samples each:
With the rule in place every scenario tried, whether the condition varies within donors or between them, sits between 4.3% and 6.0% against a nominal 5%.
Two approximations were implemented for this and then removed after measuring them. A Kenward-Roger covariance adjustment changes the standard error by a factor of 1.0000, because a within-cluster contrast is orthogonal to the random intercept and a between-cluster one is already handled by the degrees of freedom. Satterthwaite degrees of freedom, with the dispersion treated as a variance component since it enters the variance matrix through its diagonal, reproduces the between-within rule to within one degree of freedom across balanced, unbalanced, partially confounded and small sample designs (138.0 against 139, 28.8 against 29, 18.0 against 18, 5.9 against 6) for identical type I error at ten times the runtime. The rule is kept; both are validation of it rather than replacements for it.
On real data
The simulations above generate data from exactly the model the code assumes, which favours any correctly implemented estimator. So the same question was asked of a real graph:
pbmc68k_reduced, eight donors each with their own composition, three samples per donor, and a condition that is a label on the donor with no effect of its own. Every neighbourhood called significant is a false positive.Aggregated over 8 seeds. Ignoring the donor structure calls false positives above the nominal rate and two survive multiple testing correction; the random intercept more than halves them and none survive.
The milo tutorial was also executed end to end against this branch, which runs the mixed model on
stephenson_2021withSiteas a random intercept: 4307 neighbourhoods, 6 of them with a group that has no cells and reported as NaN, 21 of the remaining 4301 not converged, and a site variance between 0.02 and 0.66.Degrees of freedom
The between-within rule was checked against the Satterthwaite approximation rather than taken on faith. Satterthwaite was implemented with the dispersion as an extra variance component, which is exact here because the dispersion enters the variance matrix through its diagonal so its derivative is the identity. It agrees with the rule across balanced, unbalanced, partially confounded and small sample designs:
Type I error was identical in every case while the fit took 10x as long (8.5 to 81.5 ms), so the rule is kept and the Satterthwaite code is not.
Performance
A neighbourhood is fitted one at a time, so a real run does this thousands of times. Solving against the Cholesky factor of the variance matrix instead of building its pseudoinverse, and taking the score and Fisher information traces in the space of random effect levels rather than of samples, is worth most of it. Timed in separate processes on 150 samples with 25 donors:
The Cholesky rewrite matches the pseudoinverse version to 8e-15.
The dispersion is estimated by maximising the Cox-Reid adjusted profile likelihood rather than by solving a moment equation, iterated with the fit to consistency. A moment estimator needs to be told how many degrees of freedom the fit spent, and getting that wrong left the dispersion between 1% and 16% above the planted value depending on the size of the design; the adjusted profile likelihood holds it at -2% to -6% everywhere. Type I error is unchanged either way, but the dispersion is reported to the user, so its accuracy matters on its own.
Notes
formulaicfor the fixed effects model matrix, which already ships in thedeextra;solverandmodel_contrastsdo not apply and the latter raises.log(norm.factors), which omits library size.The milo tutorial gained a section on this in scverse/pertpy-tutorials#71, which is merged, and the submodule pointer here picks it up, so nothing needs to follow. The notebook job of this pull request executes that section against this branch.