jurkat cells nbln

import%20marimo%0A%0A__generated_with%20%3D%20%220.23.2%22%0Aapp%20%3D%20marimo.App()%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20return%20(mo%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Principled%20gene%E2%80%93gene%20correlations%20with%20NBLN%20%2B%20SVI%20cascade%20%2B%20sparse%20loadings%20prior%0A%0A%20%20%20%20*This%20tutorial%20is%20a%20sequel%20to*%20%5BModeling%20Assumptions%20for%20Single-Cell%20RNA-seq%20with%20%60scribe%60%5D(jurkat_cells.md).%20*It%20assumes%20you%20have%20already%20worked%20through%20that%20one%20(or%20are%20comfortable%20with%20negative%20binomial%20counts%2C%20variable%20capture%2C%20parameterizations%2C%20and%20low-rank%20variational%20guides)%20and%20want%20to%20ask%20a%20sharper%20question%3A%20when%20the%20model%20says%20two%20genes%20covary%2C%20what%20exactly%20is%20the%20model%20saying%20%E2%80%94%20and%20how%20do%20we%20keep%20it%20honest%3F*%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20The%20question%20we%20want%20to%20answer%0A%0A%20%20%20%20In%20the%20previous%20tutorial%20we%20ended%20at%20a%20low-rank-Gaussian%20variational%20guide%20over%20%24(%5Cunderline%7B%5Cmu%7D%2C%20%5Cphi)%24%20and%20looked%20at%20the%20resulting%20correlation%20heatmap.%20That%20heatmap%20is%20*useful*%20%E2%80%94%20it%20shows%20which%20gene-level%20parameters%20are%20coupled%20in%20the%20variational%20posterior%20%E2%80%94%20but%20it%20is%20**not**%20a%20correlation%20between%20gene%20expressions.%20It%20is%20a%20correlation%20between%20*parameters%20of%20a%20model%20whose%20generative%20story%20emits%20independent%20counts%20gene%20by%20gene*%20given%20the%20per-cell%20capture%20and%20the%20shared%20%24%5Cphi%24.%20Two%20genes%20look%20correlated%20there%20because%20their%20NB%20parameters%20are%20pulled%20by%20the%20same%20data%2C%20not%20because%20the%20model%20has%20any%20built-in%20mechanism%20for%20one%20gene's%20expression%20to%20track%20another.%0A%0A%20%20%20%20For%20the%20question%20*%22which%20genes%20covary%20biologically%3F%22*%20we%20need%20a%20generative%20model%20that%20lets%20gene%20expression%20vary%20jointly%20across%20cells.%20That%20is%20the%20move%20we%20make%20in%20this%20tutorial%3A%20from%20NB-with-shared-%24p%24%20models%20(NBDM%20%2F%20NBVCP)%20to%20the%20**Negative-Binomial%20LogNormal**%20(NBLN)%2C%20where%20the%20per-cell%20rates%20live%20in%20a%20correlated%20log-normal%20space%3A%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cunderline%7Bx%7D_c%20%5C%3B%5Csim%5C%3B%20%5Cmathcal%7BN%7D%5C!%5Cleft(%5Cunderline%7B%5Cmu%7D%2C%5C%3B%20%5Cunderline%7B%5Cunderline%7BW%7D%7D%5C%2C%5Cunderline%7B%5Cunderline%7BW%7D%7D%5E%7B%5C!%5Ctop%7D%20%2B%20%5Cmathrm%7Bdiag%7D(%5Cunderline%7Bd%7D)%5Cright)%2C%0A%20%20%20%20%5Cqquad%0A%20%20%20%20u_%7Bc%2Cg%7D%20%5Cmid%20%5Cunderline%7Bx%7D_c%20%5C%3B%5Csim%5C%3B%20%5Cmathrm%7BNB%7D%5C!%5Cleft(r_g%2C%5C%3B%20%5Cmathrm%7Blogit%7D%5E%7B-1%7D%5C!(x_%7Bc%2Cg%7D%20-%20%5Ceta_c)%5Cright).%0A%20%20%20%20%24%24%0A%0A%20%20%20%20Now%20gene-gene%20correlation%20is%20a%20**property%20of%20the%20model%20itself**%20%E2%80%94%20encoded%20in%20the%20loadings%20matrix%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7BG%5Ctimes%20k%7D%24%20%E2%80%94%20not%20a%20byproduct%20of%20a%20richer%20variational%20guide.%20Asking%20%22do%20these%20two%20genes%20covary%3F%22%20becomes%20asking%20%22what%20does%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%5Cunderline%7B%5Cunderline%7BW%7D%7D%5E%7B%5C!%5Ctop%7D%24%20say%3F%22%2C%20which%20has%20a%20clean%20answer.%0A%0A%20%20%20%20But%20NBLN%20is%20not%20a%20free%20lunch.%20It%20introduces%20a%20per-cell%20*rigid-translation%20gauge*%20that%20creates%20a%20numerical%20and%20conceptual%20hazard%2C%20and%20the%20loadings%20matrix%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20has%20many%20more%20columns%20than%20the%20data%20can%20support%20%E2%80%94%20so%20without%20care%2C%20the%20%22principled%22%20correlations%20we%20extract%20will%20be%20polluted%20by%20gauge%20slop%20and%20noise.%20The%20rest%20of%20this%20tutorial%20walks%20through%20how%20%60scribe%60%20deals%20with%20both%20problems%20using%20three%20composable%20pieces%3A%0A%0A%20%20%20%201.%20**SVI%20cascade**%20%E2%80%94%20anchor%20the%20NBLN%20Laplace%20fit%20on%20a%20converged%20NBVCP-SVI%20fit.%0A%20%20%20%202.%20**Cascade%20freeze**%20%E2%80%94%20pin%20the%20gauge-vulnerable%20parameters%20at%20their%20SVI%20values.%0A%20%20%20%203.%20**Loadings%20shrinkage**%20%E2%80%94%20let%20a%20sparse%20prior%20on%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20pick%20the%20effective%20rank%20from%20data.%0A%0A%20%20%20%20The%20payoff%20is%20a%20correlation%20structure%20on%20genes%20that%20you%20can%20defend%20as%20biologically%20meaningful%20rather%20than%20as%20an%20artifact%20of%20how%20the%20optimizer%20landed.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Why%20NBLN%20needs%20a%20per-cell%20latent%20layer%20(and%20NBDM%20doesn't)%0A%0A%20%20%20%20A%20practical%20question%20often%20comes%20up%20at%20this%20point%3A%20*the%20previous%20tutorial%20worked%20entirely%20with%20NBDM%20%2F%20NBVCP%20models%2C%20where%20each%20cell's%20biology%20was%20just%20a%20vector%20of%20NB%20draws%20%E2%80%94%20no%20extra%20%22latents%20per%20cell%22%20to%20fit.%20Why%20does%20NBLN%20suddenly%20need%20one%3F*%20The%20answer%20is%20the%20most%20important%20conceptual%20difference%20between%20the%20two%20model%20families%2C%20and%20it%20explains%20a%20lot%20of%20NBLN's%20complexity.%20Let%20us%20take%20it%20slowly.%0A%0A%20%20%20%20%23%23%23%20NBDM%3A%20the%20per-cell%20composition%20integrates%20out%20over%20the%20simplex%20in%20closed%20form%0A%0A%20%20%20%20In%20NBDM%20%2F%20NBVCP%2C%20the%20per-cell%20generative%20process%20is%20grounded%20in%20the%20Poisson-Gamma-Dirichlet%20hierarchy.%20As%20shown%20in%20the%20paper%2C%20when%20all%20genes%20share%20the%20same%20success%20probability%20%24p%24%2C%20the%20joint%20distribution%20of%20independent%20negative%20binomial%20UMI%20counts%20admits%20a%20two-level%20decomposition%3A%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cpi(%5Cunderline%7Bu%7D_c%20%5Cmid%20%5Cunderline%7Br%7D%2C%20p)%20%3D%0A%20%20%20%20%5Cint_0%5E%7B%5Cinfty%7D%20d%5CLambda%20%5C%2C%0A%20%20%20%20%5Coverbrace%7B%0A%20%20%20%20%5Cpi(u_%7BT%2Cc%7D%20%5Cmid%20%5CLambda)%0A%20%20%20%20%7D%5E%7B%5Ctext%7BPoisson%7D%7D%0A%20%20%20%20%5Ctimes%0A%20%20%20%20%5Cunderbrace%7B%0A%20%20%20%20%5Cpi(%5CLambda%20%5Cmid%20r_T%2C%20%5Ctheta)%0A%20%20%20%20%7D_%7B%5Ctext%7BGamma%7D%7D%0A%20%20%20%20%5Cint%20d%5E%7BG-1%7D%20%5Cunderline%7B%5Crho%7D%20%5C%2C%0A%20%20%20%20%5Coverbrace%7B%0A%20%20%20%20%5Cpi(%5Cunderline%7B%5Crho%7D%20%5Cmid%20%5Cunderline%7Br%7D)%0A%20%20%20%20%7D%5E%7B%5Ctext%7BDirichlet%7D%7D%0A%20%20%20%20%5Ctimes%0A%20%20%20%20%5Cunderbrace%7B%0A%20%20%20%20%5Cpi(%5Cunderline%7Bu%7D_c%20%5Cmid%20u_%7BT%2Cc%7D%2C%20%5Cunderline%7B%5Crho%7D)%0A%20%20%20%20%7D_%7B%5Ctext%7BMultinomial%7D%7D%2C%0A%20%20%20%20%24%24%0A%0A%20%20%20%20where%20%24%5Cunderline%7B%5Crho%7D%20%5Cin%20%5CDelta%5E%7BG-1%7D%24%20is%20the%20per-cell%20gene-expression%20composition%20on%20the%20%24(G-1)%24-simplex%20(the%20space%20of%20%22fractions%20of%20the%20transcriptome%20occupied%20by%20each%20gene%22)%2C%20%24u_%7BT%2Cc%7D%20%3D%20%5Csum_g%20u_%7Bc%2Cg%7D%24%20is%20the%20total%20UMI%20count%2C%20%24r_T%20%3D%20%5Csum_g%20r_g%24%2C%20and%20%24%5Ctheta%20%3D%20p%2F(1-p)%24.%0A%0A%20%20%20%20In%20principle%2C%20each%20cell%20%22samples%22%20a%20value%20for%20the%20two%20continuous%20latent%20quantities%3A%20a%20total%20rate%20%24%5CLambda_c%24%20and%20a%20composition%20%24%5Cunderline%7B%5Crho%7D_c%24.%20However%2C%20part%20of%20the%20beauty%20of%20the%20NBDM%20problem%20is%20that%20**both%20integrals%20collapse%20to%20closed%20forms**%2C%20freeing%20us%20from%20having%20to%20determine%20on%20a%20per-cell%20basis%20the%20value%20of%20either%20of%20these%20variables.%20The%20inner%20simplex%20integral%20(Dirichlet%20%24%5Ctimes%24%20Multinomial)%20evaluates%20to%20the%20Dirichlet-Multinomial%20(%24%5Cmathrm%7BDM%7D%24)%20distribution%2C%20and%20the%20outer%20total%20integral%20(Gamma%20%24%5Ctimes%24%20Poisson)%20evaluates%20to%20the%20Negative%20Binomial%20(%24%5Cmathrm%7BNB%7D%24).%20The%20result%20is%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cpi(%5Cunderline%7Bu%7D_c%20%5Cmid%20%5Cunderline%7Br%7D%2C%20p)%20%3D%0A%20%20%20%20%5Cmathrm%7BNB%7D(u_%7BT%2Cc%7D%20%5Cmid%20r_T%2C%20p)%20%5Ctimes%20%5Cmathrm%7BDM%7D(%5Cunderline%7Bu%7D_c%20%5Cmid%20u_%7BT%2Cc%7D%2C%20%5Cunderline%7Br%7D).%0A%20%20%20%20%24%24%0A%0A%20%20%20%20Neither%20%24%5CLambda_c%24%20nor%20%24%5Cunderline%7B%5Crho%7D_c%24%20survives%20as%20a%20leftover%20latent%20to%20be%20fitted%20or%20approximated%20per%20cell%20%E2%80%94%20the%20model%20is%20fully%20tractable%2C%20with%20no%20per-cell%20numerical%20work%20required.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20NBLN%3A%20the%20same%20integral%20has%20no%20closed%20form%0A%0A%20%20%20%20NBLN%20replaces%20%22independent%20NB%20per%20gene%20with%20a%20shared%20%24p%24%22%20with%20a%20*correlated*%20log-rate%20vector%20per%20cell%3A%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cunderline%7Bx%7D_c%20%5C%3B%5Csim%5C%3B%20%5Cmathcal%7BN%7D(%5Cunderline%7B%5Cmu%7D%2C%5C%3B%20%5Cunderline%7B%5Cunderline%7BW%7D%7D%5Cunderline%7B%5Cunderline%7BW%7D%7D%5E%7B%5C!%5Ctop%7D%20%2B%20%5Cmathrm%7Bdiag%7D(%5Cunderline%7Bd%7D))%2C%0A%20%20%20%20%5Cqquad%0A%20%20%20%20u_%7Bc%2Cg%7D%20%5Cmid%20%5Cunderline%7Bx%7D_c%20%5C%3B%5Csim%5C%3B%20%5Cmathrm%7BNB%7D%5C!%5Cleft(r_g%2C%5C%3B%20%5Cmathrm%7Blogit%7D%5E%7B-1%7D%5C!(x_%7Bc%2Cg%7D%20-%20%5Ceta_c)%5Cright).%0A%20%20%20%20%24%24%0A%0A%20%20%20%20In%20other%20words%2C%20instead%20of%20having%20a%20globally-shared%20%24p%24%20parameter%2C%20the%20success%20probability%20becomes%20gene-%20*and*%20cell-specific%3A%0A%0A%20%20%20%20%24%24%0A%20%20%20%20p_%7Bc%2Cg%7D%20%3D%20%5Cmathrm%7Blogit%7D%5E%7B-1%7D(x_%7Bc%2Cg%7D%20-%20%5Ceta_c)%20%3D%20%5Cfrac%7B1%7D%7B1%20%2B%20e%5E%7B-(x_%7Bc%2Cg%7D%20-%20%5Ceta_c)%7D%7D%2C%0A%20%20%20%20%24%24%0A%0A%20%20%20%20where%20%24%5Ceta_c%24%20is%20a%20per-cell%20offset%20that%20absorbs%20total%20expression%20(another%20reparameterization%20of%20the%20UMI%20capture%20probability%20that%20models%20how%20efficiently%20each%20cell%20was%20sequenced)%2C%20and%20%24x_%7Bc%2Cg%7D%24%20is%20the%20coordinate%20of%20the%20shared%20Gaussian%20latent%20for%20gene%20%24g%24%20in%20cell%20%24c%24.%20%20Every%20gene%20in%20every%20cell%20lands%20at%20a%20different%20point%20in%20%24%5B0%2C1%5D%24%2C%20driven%20by%20where%20%24%5Cunderline%7Bx%7D_c%24%20happens%20to%20fall.%0A%0A%20%20%20%20This%20single%20change%20destroys%20the%20algebraic%20glue%20that%20made%20NBDM%20tractable.%20The%20closed-form%20Dirichlet-Multinomial%20structure%20we%20described%20above%20relied%20on%20*all%20genes%20sharing%20the%20same%20%24p%24*%20%E2%80%94%20that%20is%20precisely%20the%20condition%20under%20which%20the%20composition%20%24%5Cunderline%7B%5Crho%7D%24%20is%20Dirichlet%20and%20the%20simplex%20integral%20closes.%20Once%20%24p%24%20becomes%20gene-specific%20through%20%24x_%7Bc%2Cg%7D%24%2C%20the%20genes%20no%20longer%20share%20a%20common%20%22currency%22%2C%20the%20composition%20is%20no%20longer%20Dirichlet%2C%20and%20the%20per-cell%20marginal%20requires%20the%20%24G%24-dimensional%20integral%20that%20follows.%0A%0A%20%20%20%20To%20get%20the%20cell-level%20marginal%20we%20would%20need%0A%0A%20%20%20%20%24%24%0A%20%20%20%20p(%5Cunderline%7Bu%7D_c%20%5Cmid%20%5Cunderline%7B%5Cmu%7D%2C%20%5Cunderline%7B%5Cunderline%7BW%7D%7D%2C%20%5Cunderline%7Bd%7D%2C%20%5Cunderline%7Br%7D%2C%20%5Ceta_c)%0A%20%20%20%20%5C%3B%3D%5C%3B%0A%20%20%20%20%5Cint_%7B%5Cmathbb%7BR%7D%5EG%7D%0A%20%20%20%20%5Cmathrm%7Bd%7D%5Cunderline%7Bx%7D_c%5C%3B%0A%20%20%20%20%5Cunderbrace%7B%5Cmathcal%7BN%7D%5C!%5Cleft(%5Cunderline%7Bx%7D_c%5Cmid%20%5Cunderline%7B%5Cmu%7D%2C%5C%2C%20%5Cunderline%7B%5Cunderline%7BW%7D%7D%5Cunderline%7B%5Cunderline%7BW%7D%7D%5E%7B%5C!%5Ctop%7D%20%2B%20%5Cmathrm%7Bdiag%7D(%5Cunderline%7Bd%7D)%5Cright)%7D_%7B%5Ctext%7BGaussian%20prior%20on%20log-rates%7D%7D%5C%3B%0A%20%20%20%20%5Cunderbrace%7B%5Cprod_%7Bg%3D1%7D%5EG%20%5Cmathrm%7BNB%7D%5C!%5Cleft(u_%7Bc%2Cg%7D%20%5C%2C%3B%5C%2C%20r_g%2C%5C%2C%20%5Cmathrm%7Blogit%7D%5E%7B-1%7D(x_%7Bc%2Cg%7D%20-%20%5Ceta_c)%5Cright)%7D_%7B%5Ctext%7BNB%20likelihood%20per%20gene%7D%7D.%0A%20%20%20%20%24%24%0A%0A%20%20%20%20This%20integral%20has%20**no%20closed%20form**.%20The%20Gaussian%20wants%20to%20be%20conjugate%20to%20a%20Gaussian%20likelihood%3B%20the%20NB%20likelihood%2C%20parameterized%20through%20a%20logit%20transform%20of%20the%20latent%2C%20is%20not%20Gaussian.%20There%20is%20no%20magic%20cancellation%20that%20removes%20%24%5Cunderline%7Bx%7D_c%24.%20We%20are%20stuck%20with%20a%20%24G%24-dimensional%20integral%20that%20has%20to%20be%20approximated%20***for%20every%20cell***.%0A%0A%20%20%20%20%23%23%23%20What%20%22Laplace%20inference%22%20actually%20does%20about%20it%0A%0A%20%20%20%20%60scribe%60's%20NBLN-Laplace%20path%20handles%20each%20cell's%20integral%20by%20approximating%20the%20integrand%20around%20its%20mode%3A%0A%0A%20%20%20%201.%20Find%20the%20%5BMAP%5D(https%3A%2F%2Fen.wikipedia.org%2Fwiki%2FMaximum_a_posteriori_estimation)%20%24%5Cunderline%7Bx%7D_c%5E%5Cast%24%20%E2%80%94%20the%20value%20of%20%24%5Cunderline%7Bx%7D_c%24%20that%20maximizes%20the%20integrand%20via%20Newton%20iteration%20per%20cell.%0A%20%20%20%202.%20Compute%20the%20Hessian%20of%20the%20negative%20log-integrand%20at%20that%20mode%20(curvature%20of%20the%20local%20fit).%0A%20%20%20%203.%20Approximate%20the%20integrand%20as%20a%20Gaussian%20centered%20at%20%24%5Cunderline%7Bx%7D_c%5E%5Cast%24%20with%20covariance%20equal%20to%20the%20inverse%20Hessian.%0A%0A%20%20%20%20Under%20that%20local-Gaussian%20approximation%20the%20%24G%24-dimensional%20integral%20collapses%20into%20a%20%60log-det%60%20term%3A%20a%20tractable%20scalar%20that%20the%20optimizer%20can%20differentiate.%0A%0A%20%20%20%20%23%23%23%20The%20outer%20loop%3A%20it%20is%20an%20EM%20algorithm%0A%0A%20%20%20%20The%20per-cell%20Laplace%20step%20above%20is%20only%20the%20inner%20half%20of%20the%20picture.%20NBLN-Laplace%20alternates%20between%20**two**%20optimization%20problems%2C%20in%20the%20spirit%20of%20the%20classical%20%5BExpectation-Maximization%20(EM)%5D(https%3A%2F%2Fen.wikipedia.org%2Fwiki%2FExpectation%25E2%2580%2593maximization_algorithm)%20algorithm%20%E2%80%94%20fix%20one%20block%20of%20unknowns%2C%20solve%20the%20other%2C%20swap%2C%20repeat%3A%0A%0A%20%20%20%20-%20**E-step%20(%22expectation%22)%3A**%20with%20the%20global%20parameters%20%24%5Ctheta%20%3D%20(%5Cunderline%7B%5Cmu%7D%2C%20%5Cunderline%7B%5Cunderline%7BW%7D%7D%2C%20%5Cunderline%7Bd%7D%2C%20%5Cldots)%24%20held%20fixed%2C%20find%20each%20cell's%20per-cell%20latent%20MAP%20%24%5Cunderline%7Bx%7D_c%5E%5Cast%24%20(and%20the%20Hessian%20needed%20for%20the%20%60log-det%60%20term)%20via%20the%20Newton%20iteration%20described%20above.%20%60laplace_config%60%20(%60n_newton_steps%60%2C%20%60damping%60%2C%20%60newton_max_step%60%2C%20%E2%80%A6)%20controls%20this%20inner%20step.%0A%20%20%20%20-%20**M-step%20(%22maximization%22)%3A**%20with%20the%20per-cell%20MAPs%20held%20fixed%2C%20take%20optimizer%20steps%20on%20the%20global%20parameters%20%24%5Ctheta%24%20using%20the%20Laplace-approximated%20ELBO%20as%20the%20objective%20%E2%80%94%20the%20same%20kind%20of%20stochastic-gradient%20loop%20as%20a%20regular%20SVI%20fit.%20%60n_steps%60%2C%20%60optimizer_config%60%2C%20and%20%60early_stopping%60%20control%20this%20outer%20step.%0A%0A%20%20%20%20The%20two%20halves%20take%20turns%3A%20refresh%20the%20per-cell%20MAPs%2C%20run%20a%20chunk%20of%20M-step%20gradient%20updates%2C%20refresh%20the%20MAPs%20again%2C%20run%20more%20M-step%20updates%2C%20and%20so%20on%20until%20the%20loss%20envelope%20settles.%0A%0A%20%20%20%20The%20reason%20this%20split%20is%20necessary%20goes%20back%20to%20the%20identifiability%20story%20above.%20The%20per-cell%20latents%20%24%5Cunderline%7Bx%7D_c%24%20are%20%22annoying%22%20%E2%80%94%20there%20are%20%24C%20%5Ctimes%20G%24%20of%20them%2C%20they%20have%20no%20closed-form%20posterior%2C%20and%20they%20tangle%20non-linearly%20with%20the%20global%20parameters%20through%20the%20NB-logit.%20Holding%20them%20fixed%20at%20their%20current%20MAPs%20lets%20the%20M-step%20look%20like%20a%20familiar%20SVI%2FAdam%20loop%20on%20a%20few%20global%20parameters.%20Holding%20the%20global%20parameters%20fixed%20lets%20each%20cell's%20E-step%20look%20like%20an%20isolated%2C%20well-conditioned%20Newton%20problem.%20Doing%20both%20at%20once%20is%20intractable%3B%20alternating%20between%20them%20is%20what%20makes%20NBLN-Laplace%20tick.%0A%0A%20%20%20%20This%20is%20why%20NBLN-Laplace%20explicitly%20**owns%20a%20per-cell%20%24%5Cunderline%7Bx%7D_c%24%20MAP**%20as%20a%20stored%20field%20on%20the%20result%20(%60x_loc%60)%20and%20why%20the%20Newton%20iterations%20show%20up%20in%20the%20loss%20accounting.%20It%20is%20not%20a%20quirk%20of%20the%20implementation%20%E2%80%94%20it%20is%20the%20inevitable%20cost%20of%20moving%20from%20a%20model%20whose%20per-cell%20integral%20is%20closed-form%20(NBDM)%20to%20one%20whose%20per-cell%20integral%20is%20not%20(NBLN).%20The%20reward%20is%20the%20joint%20correlation%20structure%20encoded%20in%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%5Cunderline%7B%5Cunderline%7BW%7D%7D%5E%7B%5C!%5Ctop%7D%24%2C%20which%20is%20the%20only%20thing%20NBDM%20cannot%20give%20us%20no%20matter%20how%20clever%20the%20variational%20guide%20is.%0A%0A%20%20%20%20Keep%20this%20picture%20in%20mind%20throughout%20the%20rest%20of%20the%20tutorial%3A%20every%20time%20you%20see%20a%20%22per-cell%20Newton%22%20or%20a%20%22latent%20%24%5Cunderline%7Bx%7D_c%24%22%2C%20that%20is%20the%20price%20of%20replacing%20%22independent%20NB%20per%20gene%22%20with%20%22correlated%20log-rates%20across%20genes%22.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20%23%20Import%20basic%20packages%0A%20%20%20%20from%20pathlib%20import%20Path%0A%20%20%20%20import%20pickle%0A%0A%20%20%20%20%23%20Import%20functionality%20to%20clear%20GPU%20memory%0A%20%20%20%20import%20gc%0A%20%20%20%20from%20jax%20import%20clear_caches%0A%0A%20%20%20%20%23%20Import%20our%20main%20package%0A%20%20%20%20import%20scribe%0A%0A%20%20%20%20%23%20Import%20useful%20tools%0A%20%20%20%20import%20numpy%20as%20np%0A%0A%20%20%20%20%23%20Import%20plotting%20packages%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%0A%20%20%20%20%23%20Set%20our%20plotting%20style%20(totally%20optional)%0A%20%20%20%20scribe.viz.matplotlib_style()%0A%20%20%20%20return%20Path%2C%20clear_caches%2C%20gc%2C%20np%2C%20pickle%2C%20plt%2C%20scribe%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20We%20use%20the%20same%20%5BJurkat%2010x%20dataset%5D(https%3A%2F%2Fwww.10xgenomics.com%2Fdatasets%2Fjurkat-cells-1-standard-1-1-0)%20as%20the%20first%20tutorial%20%E2%80%94%20a%20monoculture%20of%20~3%2C200%20cells.%20Because%20there%20is%20only%20one%20cell%20type%2C%20any%20%22correlation%20structure%22%20the%20NBLN%20model%20reports%20has%20to%20come%20from%20something%20other%20than%20population%20substructure%3A%20it%20must%20be%20either%20real%20co-regulation%20within%20Jurkat%20biology%2C%20residual%20technical%2Flibrary-size%20effects%20the%20model%20failed%20to%20isolate%2C%20or%20noise.%20Discriminating%20among%20those%20three%20is%20exactly%20what%20the%20diagnostics%20in%20this%20notebook%20are%20for.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Path%2C%20scribe)%3A%0A%20%20%20%20%23%20Define%20data%20directory%0A%20%20%20%20data_dir%20%3D%20Path(%0A%20%20%20%20%20%20%20%20%22%2Fpath%2Fto%2Fjurkat_cells%2F%22%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20Load%20the%20data%0A%20%20%20%20adata%20%3D%20scribe.data_loader.load_and_preprocess_anndata(%0A%20%20%20%20%20%20%20%20data_dir%2C%20return_jax%3DFalse%0A%20%20%20%20)%0A%0A%20%20%20%20adata%0A%20%20%20%20return%20adata%2C%20data_dir%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Stage%201%20%E2%80%94%20the%20SVI%20cascade%20source%0A%0A%20%20%20%20The%20first%20ingredient%20is%20a%20converged%20NBVCP-SVI%20fit%20on%20the%20same%20data.%20Why%20bother%20fitting%20an%20%22old%22%20model%20first%20%E2%80%94%20and%20what%20makes%20its%20parameters%20trustworthy%20as%20a%20seed%20for%20NBLN%3F%0A%0A%20%20%20%20The%20real%20reason%20is%20**structural%20identifiability**.%20NBLN%20couples%20three%20quantities%20%E2%80%94%20the%20per-gene%20dispersion%20%24r_g%24%2C%20the%20per-cell%20capture%20%24%5Ceta_c%24%2C%20and%20the%20per-cell%20latent%20log-rate%20%24%5Cunderline%7Bx%7D_c%24%20%E2%80%94%20through%20a%20likelihood%20that%20cannot%20pin%20them%20all%20down%20at%20once.%20The%20three%20knobs%20trade%20off%20against%20each%20other%3A%20shift%20%24%5Ceta_c%24%20and%20you%20can%20compensate%20with%20%24%5Cunderline%7Bx%7D_c%24%3B%20rescale%20%24r_g%24%20and%20you%20change%20how%20the%20NB%20curvature%20reads%20the%20latent%3B%20and%20the%20gene-gene%20covariance%20%24%5Cunderline%7B%5Cunderline%7B%5CSigma%7D%7D%20%3D%20%5Cunderline%7B%5Cunderline%7BW%7D%7D%5C%2C%5Cunderline%7B%5Cunderline%7BW%7D%7D%5E%7B%5C!%5Ctop%7D%20%2B%20%5Cmathrm%7Bdiag%7D(%5Cunderline%7Bd%7D)%24%20leaks%20into%20all%20of%20it%20through%20the%20per-cell%20Gaussian%20prior%20on%20%24%5Cunderline%7Bx%7D_c%24.%20Cold-starting%20NBLN-Laplace%20on%20the%20full%20parameter%20set%20is%20asking%20the%20optimizer%20to%20disentangle%20several%20near-equivalent%20answers%20from%20one%20likelihood%20surface%3B%20what%20survives%20is%20gauge%20slop%20rather%20than%20biology.%0A%0A%20%20%20%20The%20cascade%20trick%20exploits%20a%20**structural%20fact%20about%20the%20model%20family**%3A%20NBVCP%20is%20exactly%20the%20limit%20of%20NBLN%20as%20%24%5Cunderline%7B%5Cunderline%7B%5CSigma%7D%7D%20%5Cto%200%24.%20When%20the%20gene-gene%20covariance%20vanishes%2C%20the%20per-cell%20log-rate%20prior%20collapses%20to%20a%20point%2C%20the%20latent%20%24%5Cunderline%7Bx%7D_c%24%20drops%20out%20as%20a%20free%20parameter%2C%20and%20the%20NBLN%20likelihood%20reduces%20gene-by-gene%20to%20the%20variable-capture%20NB%20you%20fit%20in%20the%20previous%20tutorial.%0A%0A%20%20%20%20The%20crucial%20observation%20is%20that%20**%24r_g%24%20plays%20the%20same%20role%20in%20both%20models**%20and%20does%20not%20depend%20on%20%24%5Cunderline%7B%5Cunderline%7B%5CSigma%7D%7D%24.%20It%20is%20the%20per-gene%20dispersion%20of%20the%20NB%20observation%20channel%20%E2%80%94%20a%20property%20of%20how%20counts%20are%20generated%20*given*%20a%20log-rate%20%E2%80%94%20not%20of%20how%20log-rates%20covary%20across%20genes.%20So%20the%20value%20of%20%24r_g%24%20that%20best%20explains%20a%20gene's%20marginal%20count%20distribution%20under%20NBVCP%20is%20also%20the%20value%20that%20best%20explains%20it%20under%20NBLN.%20In%20NBVCP%20the%20identifiability%20landscape%20is%20dramatically%20friendlier%20(no%20%24%5Cunderline%7B%5Cunderline%7B%5CSigma%7D%7D%24-mediated%20coupling%2C%20no%20per-cell%20latent%20layer%2C%20the%20Dirichlet-multinomial%20story%20works%20out%20cleanly)%2C%20so%20the%20data%20pins%20%24r_g%24%20down%20sharply.%20We%20fit%20there%2C%20and%20carry%20%24r_g%24%20forward.%0A%0A%20%20%20%20Per-cell%20capture%20%24%5Ceta_c%24%20is%20the%20harder%20case.%20Even%20in%20NBVCP%20it%20is%20only%20partly%20identified%20by%20the%20data%20%E2%80%94%20we%20need%20a%20biology-informed%20prior%20(%60priors%3D%7B%22capture_efficiency%22%3A%20...%7D%60)%20to%20break%20the%20residual%20ambiguity.%20And%20in%20NBLN%20there%20is%20an%20additional%2C%20exact%20*rigid-translation%20gauge*%3A%20shift%20every%20%24x_%7Bc%2Cg%7D%24%20up%20by%20some%20per-cell%20constant%20%24%5CDelta_c%24%20and%20shift%20%24%5Ceta_c%24%20down%20by%20the%20same%20%24%5CDelta_c%24%2C%20and%20the%20observed%20counts%20are%20unchanged%20because%20the%20likelihood%20only%20sees%20%24x_%7Bc%2Cg%7D%20-%20%5Ceta_c%24.%20So%20the%20%24%5Ceta_c%24%20we%20extract%20from%20NBVCP%20is%20not%20the%20%22true%22%20%24%5Ceta_c%24%20in%20any%20absolute%20sense%2C%20and%20we%20should%20not%20pretend%20it%20is.%0A%0A%20%20%20%20Here%20is%20the%20payoff%20that%20makes%20the%20cascade%20defensible%3A%20**downstream%20compositional%20inference%20is%20robust%20to%20misspecification%20of%20%24%5Ceta_c%24**.%20The%20compositional-robustness%20theorem%20(worked%20out%20in%20the%20paper)%20shows%20that%20the%20per-cell%20composition%20%24%5Cunderline%7B%5Crho%7D_c%20%3D%20%5Cmathrm%7Bsoftmax%7D(%5Cunderline%7Bx%7D_c)%24%20is%20invariant%20under%20the%20rigid%20translation.%20Any%20biologically%20interesting%20question%20that%20can%20be%20phrased%20compositionally%20%E2%80%94%20differential%20expression%20in%20CLR%2FILR%20coordinates%2C%20gene-gene%20correlations%20from%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5Cperp%20%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5Cperp%5E%7B%5C!%5Ctop%7D%24%2C%20pathway%20enrichment%20on%20ALR%20balances%20%E2%80%94%20is%20invariant%20to%20whatever%20per-cell%20constant%20the%20gauge%20picked%20up.%20So%20when%20we%20hand%20%24%5Ceta_c%24%20to%20NBLN%20as%20a%20fixed%20quantity%2C%20we%20are%20fixing%20it%20in%20a%20coordinate%20system%20where%20the%20questions%20we%20care%20about%20don't%20see%20the%20error.%0A%0A%20%20%20%20The%20cascade%20is%20therefore%20not%20%22use%20a%20richer%20model's%20posterior%20to%20seed%20NBLN%22.%20It%20is%20sharper%20than%20that%3A%0A%0A%20%20%20%20%3E%20*Fit%20the%20two%20parameters%20NBLN%20cannot%20pin%20down%20cleanly%20%E2%80%94%20%24r_g%24%20and%20%24%5Ceta_c%24%20%E2%80%94%20in%20the%20limit%20where%20the%20model%20is%20well-identified%20(%24%5Cunderline%7B%5Cunderline%7B%5CSigma%7D%7D%20%5Cto%200%24%2C%20which%20is%20NBVCP)%2C%20and%20hand%20them%20to%20NBLN%20as%20fixed%20quantities.%20The%20%24r_g%24%20values%20are%20exact%20in%20the%20limit%20(the%20same%20parameter%2C%20the%20same%20role%2C%20just%20easier%20to%20fit).%20The%20%24%5Ceta_c%24%20values%20are%20only%20fixed%20up%20to%20a%20gauge%20that%20compositional%20questions%20don't%20see%20anyway.*%0A%0A%20%20%20%20That%20is%20the%20cascade.%20The%20next%20two%20pieces%20%E2%80%94%20the%20freeze%20and%20the%20loadings%20shrinkage%20prior%20%E2%80%94%20build%20on%20this%20foundation%20by%20pinning%20the%20cascade%20values%20structurally%20(Piece%202)%20and%20by%20regularizing%20the%20remaining%20free%20parameter%2C%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%2C%20so%20that%20what%20it%20picks%20up%20is%20data-supported%20and%20not%20noise%20(Piece%203).%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20We%20use%3A%0A%0A%20%20%20%20-%20%60model%3D%22nbvcp%22%60%20with%20%60parameterization%3D%22mean_odds%22%60%20%E2%80%94%20the%20same%20combination%20that%20gave%20clean%20PPCs%20in%20the%20first%20tutorial.%0A%20%20%20%20-%20%60priors%3D%7B%22capture_efficiency%22%3A%20(np.log(100_000)%2C%200.1)%7D%60%20%E2%80%94%20the%20**biology-informed%20capture-efficiency%20prior**.%20The%20two%20numbers%20are%20the%20mean%20and%20standard%20deviation%20of%20%24%5Clog(M_c%5C%2C%5Cnu_c)%20%3D%20x_%7Bc%2Cg%7D%20-%20%5Ceta_c%24%20on%20a%20reasonable%20mRNA%20scale%3A%20%60log(100%2C000)%60%20says%20%22this%20cell%20sample's%20true%20mRNA%20count%20is%20around%20%2410%5E5%24%20molecules%22%20(a%20typical%20mammalian-cell%20ballpark)%20and%20%600.1%60%20says%20%22we%20are%20quite%20confident%20in%20that%20scale%22.%20The%20number%20is%20not%20a%20precise%20measurement%20%E2%80%94%20it%20just%20has%20to%20break%20the%20gauge%20degeneracy.%0A%20%20%20%20-%20%60gene_coverage%3D0.99%60%20%E2%80%94%20pool%20low-coverage%20genes%20into%20one%20%22other%22%20pseudo-gene.%20Reduces%20dimensionality%20without%20throwing%20data%20away.%0A%20%20%20%20-%20%60early_stopping%3D%7B...checkpoint_dir...%7D%60%20%E2%80%94%20Orbax%20checkpoints%20let%20us%20interrupt%20and%20resume%20the%20(long)%20SVI%20fit%20without%20losing%20progress.%0A%0A%20%20%20%20The%20block%20below%20either%20loads%20a%20previously-fit%20cached%20pickle%20or%20runs%20SVI%20from%20scratch.%20On%20first%20run%2C%20expect%20this%20stage%20to%20take%20some%20time.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20data_dir%2C%20np%2C%20pickle%2C%20scribe)%3A%0A%20%20%20%20%23%20Define%20output%20directory%0A%20%20%20%20out_dir%20%3D%20data_dir%20%2F%20%22scribe_results%22%0A%20%20%20%20out_dir.mkdir(parents%3DTrue%2C%20exist_ok%3DTrue)%0A%0A%20%20%20%20%23%20Cascade-stage%20knobs%20(named%20exactly%20as%20in%20the%20companion%20EDA%20script).%0A%20%20%20%20_gene_coverage%20%3D%200.99%0A%20%20%20%20_model_type_svi%20%3D%20%22nbvcp%22%0A%20%20%20%20_parameterization_svi%20%3D%20%22mean_odds%22%0A%20%20%20%20_capture_efficiency_prior%20%3D%20(np.log(100_000)%2C%200.1)%0A%0A%20%20%20%20_out_path_svi%20%3D%20out_dir%20%2F%20(%0A%20%20%20%20%20%20%20%20f%22scribe_results_%7B_model_type_svi%7D-%22%0A%20%20%20%20%20%20%20%20f%22%7B_parameterization_svi%7D-vcp-%22%0A%20%20%20%20%20%20%20%20f%22%7B_gene_coverage%7Dgene_coverage-%22%0A%20%20%20%20%20%20%20%20f%22anchor_svi.pkl%22%0A%20%20%20%20)%0A%20%20%20%20_svi_checkpoint_dir%20%3D%20out_dir%20%2F%20(%0A%20%20%20%20%20%20%20%20f%22checkpoints_%7B_model_type_svi%7D-%22%0A%20%20%20%20%20%20%20%20f%22%7B_parameterization_svi%7D-vcp-%22%0A%20%20%20%20%20%20%20%20f%22%7B_gene_coverage%7Dgene_coverage-%22%0A%20%20%20%20%20%20%20%20f%22anchor_svi%22%0A%20%20%20%20)%0A%0A%20%20%20%20if%20_out_path_svi.exists()%3A%0A%20%20%20%20%20%20%20%20with%20open(_out_path_svi%2C%20%22rb%22)%20as%20_f%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20svi_results%20%3D%20pickle.load(_f)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20_svi_checkpoint_dir.mkdir(parents%3DTrue%2C%20exist_ok%3DTrue)%0A%20%20%20%20%20%20%20%20svi_results%20%3D%20scribe.fit(%0A%20%20%20%20%20%20%20%20%20%20%20%20adata%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20model%3D_model_type_svi%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20parameterization%3D_parameterization_svi%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20inference_method%3D%22svi%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20gene_coverage%3D_gene_coverage%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20n_steps%3D500_000%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20unconstrained%3DTrue%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20priors%3D%7B%22capture_efficiency%22%3A%20_capture_efficiency_prior%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20early_stopping%3D%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22enabled%22%3A%20True%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22checkpoint_dir%22%3A%20str(_svi_checkpoint_dir)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20with%20open(_out_path_svi%2C%20%22wb%22)%20as%20_f%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20pickle.dump(svi_results%2C%20_f)%0A%0A%20%20%20%20svi_results%0A%20%20%20%20return%20out_dir%2C%20svi_results%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Sanity-checking%20the%20cascade%20source%0A%0A%20%20%20%20Before%20we%20use%20this%20fit%20as%20an%20anchor%20for%20NBLN%2C%20it%20must%20clear%20three%20sanity%20checks%3A%0A%0A%20%20%20%201.%20**Training%20converged.**%20The%20ELBO%20curve%20should%20drop%20fast%20and%20then%20flatten.%20If%20the%20SVI%20source%20has%20not%20converged%2C%20neither%20will%20the%20cascade.%0A%20%20%20%202.%20**Marginal%20PPCs%20look%20right.**%20Per-gene%20UMI%20histograms%20should%20overlap%20with%20the%20model's%20predictive%20band%20across%20the%20abundance%20range.%0A%20%20%20%203.%20**Per-cell%20capture%20behaves%20sensibly.**%20With%20a%20tight%20prior%20on%20%24%5Clog(M_c%5Cnu_c)%24%2C%20the%20inferred%20per-cell%20%24%5Cnu_c%24%20should%20be%20a%20monotone%20function%20of%20library%20size%20and%20saturate%20near%201%20only%20for%20the%20deepest%20cells.%0A%0A%20%20%20%20Skipping%20any%20of%20these%20is%20a%20recipe%20for%20cascade%20contamination%20%E2%80%94%20garbage%20in%2C%20garbage%20out.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(scribe%2C%20svi_results)%3A%0A%20%20%20%20scribe.viz.plot_loss(svi_results%2C%20figsize%3D(6%2C%203))%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20training%20loss%20is%20clean%20%E2%80%94%20fast%20initial%20drop%20and%20a%20flat%20tail.%20Next%2C%20marginal%20PPCs.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20scribe%2C%20svi_results)%3A%0A%20%20%20%20scribe.viz.plot_ppc(%0A%20%20%20%20%20%20%20%20svi_results%2C%0A%20%20%20%20%20%20%20%20adata%2C%0A%20%20%20%20%20%20%20%20n_genes%3D16%2C%0A%20%20%20%20%20%20%20%20n_rows%3D4%2C%0A%20%20%20%20%20%20%20%20figsize%3D(8%2C%208)%2C%0A%20%20%20%20%20%20%20%20n_samples%3D512%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20PPC%20panels%20track%20the%20observed%20histograms%20well%20across%20abundance%20%E2%80%94%20the%20cascade%20source%20has%20a%20coherent%20marginal%20story.%0A%0A%20%20%20%20Now%20we%20look%20at%20the%20per-cell%20capture%20diagnostic.%20With%20the%20biology-informed%20prior%2C%20%24%5Cnu_c%24%20should%20rise%20gently%20with%20library%20size%20and%20asymptote%20near%201%20only%20for%20cells%20with%20very%20deep%20coverage.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20scribe%2C%20svi_results)%3A%0A%20%20%20%20scribe.viz.plot_p_capture_scaling(svi_results%2C%20counts%3Dadata%2C%20figsize%3D(4%2C%204))%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Capture%20saturates%20near%20%24%5Csim%24%2030K%20UMIs%20%E2%80%94%20the%20deepest%20Jurkat%20cells%20in%20this%20dataset%20are%20interpreted%20as%20essentially%20fully-captured%2C%20while%20the%20shallowest%20cells%20are%20routed%20through%20a%20lower%20effective%20%24%5Cnu_c%24.%20This%20is%20the%20model's%20account%20of%20%22library%20size%20as%20a%20technical%20channel%22%20rather%20than%20%22library%20size%20as%20biology%22.%20This%20per-cell%20%24%5Ceta_c%24%20posterior%20is%20what%20we%20will%20hand%20to%20NBLN-Laplace%20as%20a%20frozen%20anchor.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20A%20first%20look%20at%20compositional%20PPCs%0A%0A%20%20%20%20Marginal%20PPCs%20ask%20%22does%20the%20model%20reproduce%20the%20histogram%20of%20UMI%20counts%20gene%20by%20gene%3F%22.%20That%20is%20necessary%20but%20not%20sufficient.%20The%20other%20thing%20we%20care%20about%20is%20*whether%20the%20joint%20distribution%20of%20expression%20across%20genes%20within%20a%20cell%20is%20right*%20%E2%80%94%20what%20fraction%20of%20a%20cell's%20transcriptome%20is%20gene%20A%20versus%20gene%20B%20versus%20gene%20C%3F%20This%20is%20the%20**compositional**%20question%2C%20and%20it%20is%20the%20natural%20object%20for%20downstream%20differential%20expression.%0A%0A%20%20%20%20%60scribe.viz.plot_compositional_ppc%60%20overlays%20the%20empirical%20distribution%20of%20per-cell%20*fractions*%20%24u_%7Bc%2Cg%7D%20%2F%20%5Csum_%7Bg'%7D%20u_%7Bc%2Cg'%7D%24%20on%20the%20model's%20compositional%20predictive%20distribution%2C%20for%20an%20automatically%20chosen%20set%20of%20well-expressed%20genes.%20If%20the%20model%20has%20the%20marginal%20counts%20right%20but%20the%20joint%20structure%20wrong%2C%20this%20is%20where%20you%20would%20see%20it.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20scribe%2C%20svi_results)%3A%0A%20%20%20%20scribe.viz.plot_compositional_ppc(%0A%20%20%20%20%20%20%20%20svi_results%2C%0A%20%20%20%20%20%20%20%20adata%2C%0A%20%20%20%20%20%20%20%20n_genes%3D16%2C%0A%20%20%20%20%20%20%20%20n_rows%3D4%2C%0A%20%20%20%20%20%20%20%20n_samples%3D2048%2C%0A%20%20%20%20%20%20%20%20min_mean_umi%3D5.0%2C%0A%20%20%20%20%20%20%20%20figsize%3D(10%2C%2010)%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20marginal%20compositional%20fractions%20look%20fine%20%E2%80%94%20but%20there%20is%20a%20striking%20visual%20feature%20in%20these%20panels%20that%20is%20worth%20pausing%20on%3A%20**the%20empirical%20(data)%20distribution%20is%20much%20wider%20than%20the%20model's%20predictive%20distribution.**%20It%20looks%20like%20the%20model%20is%20%22underconfident%22%20about%20how%20much%20spread%20there%20should%20be%20in%20each%20gene's%20fraction%20across%20cells.%20Is%20the%20model%20wrong%3F%0A%0A%20%20%20%20No%20%E2%80%94%20the%20*empirical%20fraction%20itself*%20is%20a%20noisy%20estimate%2C%20and%20that%20is%20exactly%20the%20noise%20the%20model%20is%20trying%20to%20look%20through.%0A%0A%20%20%20%20Consider%20what%20goes%20into%20the%20empirical%20fraction%20%24%5Chat%7B%5Crho%7D_%7Bc%2Cg%7D%20%3D%20u_%7Bc%2Cg%7D%20%2F%20%5Csum_%7Bg'%7D%20u_%7Bc%2Cg'%7D%24%20that%20we%20plot%20for%20the%20data.%20To%20get%20one%20of%20those%20numbers%20we%3A%0A%0A%20%20%20%201.%20Started%20with%20a%20true%20biological%20mRNA%20composition%20%24%5Cunderline%7B%5Crho%7D_c%24%20for%20cell%20%24c%24.%0A%20%20%20%202.%20Captured%20a%20fraction%20%24%5Cnu_c%24%20of%20those%20molecules%20(the%20*technical%20capture%20noise*).%0A%20%20%20%203.%20Sequenced%20and%20counted%20the%20survivors%20%E2%80%94%20a%20finite-sample%20multinomial%20draw%20from%20a%20finite%20library%20(the%20*multinomial%20sampling%20noise*).%0A%20%20%20%204.%20Divided%20counts%20by%20total%20counts%2C%20which%20propagates%20both%20noise%20sources%20into%20the%20resulting%20fraction.%0A%0A%20%20%20%20The%20empirical%20fraction%20%24%5Chat%7B%5Crho%7D_%7Bc%2Cg%7D%24%20is%20therefore%20%24%5Crho_%7Bc%2Cg%7D%24%20plus%20*all%20the%20noise%20that%20happened%20between%20cell%20and%20FASTQ*.%20For%20a%20gene%20at%20fraction%20%2410%5E%7B-3%7D%24%20in%20a%20cell%20that%20ended%20up%20with%2010%2C000%20UMIs%2C%20the%20empirical%20fraction%20has%20roughly%20%24%5Csqrt%7B10%2F10%7B%2C%7D000%5E2%7D%20%5Capprox%203%20%5Ctimes%2010%5E%7B-4%7D%24%20of%20sampling%20noise%20on%20top%20of%20whatever%20the%20true%20biological%20fraction%20is%20%E2%80%94%20a%2030%25%20relative%20noise%20floor%20that%20has%20nothing%20to%20do%20with%20biology.%0A%0A%20%20%20%20What%20the%20**model**%20plots%20as%20its%20predictive%20distribution%2C%20on%20the%20other%20hand%2C%20is%20the%20*posterior%20over%20the%20true%20biological%20%24%5Crho_%7Bc%2Cg%7D%24*%20%E2%80%94%20the%20thing%20the%20data%20is%20a%20noisy%20view%20of.%20Of%20course%20it%20is%20narrower.%20That%20gap%20between%20empirical%20and%20predicted%20is%20not%20the%20model%20being%20overconfident%3B%20it%20is%20the%20model%20correctly%20distinguishing%20biological%20signal%20from%20observation%20noise%2C%20and%20reporting%20the%20signal%20alone.%20**That%20distillation%20is%20exactly%20the%20value-add%20of%20having%20a%20probabilistic%20model%20in%20the%20first%20place.**%0A%0A%20%20%20%20Two%20consequences%20worth%20flagging%20for%20downstream%20analysis%3A%0A%0A%20%20%20%20-%20**Naive%20composition%20estimators**%20(%24%5Chat%7B%5Crho%7D_%7Bc%2Cg%7D%20%3D%20u_%7Bc%2Cg%7D%20%2F%20%5Ctext%7Btotal%7D_c%24)%20are%20not%20%22the%20truth%22%20%E2%80%94%20they%20are%20a%20particular%20estimator%20with%20known%20noise%20properties%2C%20mostly%20dominated%20by%20sampling%20at%20low%20UMI%20counts.%20Differential-expression%20methods%20that%20use%20%24%5Chat%7B%5Crho%7D%24%20directly%20inherit%20that%20noise%20floor.%0A%20%20%20%20-%20**Size-factor%20normalization**%20(divide%20by%20a%20per-cell%20scalar)%20does%20not%20fix%20this%20either.%20It%20treats%20the%20total%20UMI%20count%20as%20a%20single%20fixed%20number%20per%20cell%2C%20when%20in%20fact%20that%20number%20is%20itself%20a%20draw%20from%20a%20distribution%3B%20and%20once%20divided%2C%20the%20per-cell%20variance%20of%20the%20result%20is%20still%20dominated%20by%20multinomial%20sampling%20at%20low%20counts.%20A%20*model*%20that%20owns%20both%20the%20biology%20and%20the%20technical%20noise%20gets%20to%20integrate%20over%20both%20%E2%80%94%20and%20that%20is%20what%20produces%20the%20narrow%20predictive%20bands%20here.%0A%0A%20%20%20%20Now%2C%20compositional%20**marginals**%20can%20still%20hide%20a%20multitude%20of%20sins%20even%20when%20the%20noise%20gap%20looks%20reasonable.%20What%20if%20gene%20A%20and%20gene%20B%20individually%20look%20right%2C%20but%20when%20you%20look%20at%20the%20*joint*%20distribution%20of%20%24(%5Crho_A%2C%20%5Crho_B)%24%20across%20cells%2C%20the%20model%20produces%20a%20round%20blob%20while%20the%20data%20clearly%20shows%20a%20tilted%20ellipse%3F%20That%20is%20a%20gene-gene%20correlation%20the%20model%20is%20missing%2C%20and%20a%20marginal%20check%20will%20never%20reveal%20it.%0A%0A%20%20%20%20The%20right%20tool%20for%20this%20question%20is%20a%20compositional%20**corner%20plot**%20%E2%80%94%20a%20grid%20of%20pairwise%202D%20projections%20plus%20marginals%20%E2%80%94%20comparing%20data%20and%20predictive%20samples%20panel%20by%20panel.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20scribe%2C%20svi_results)%3A%0A%20%20%20%20scribe.viz.plot_compositional_corner_ppc(%0A%20%20%20%20%20%20%20%20svi_results%2C%0A%20%20%20%20%20%20%20%20adata%2C%0A%20%20%20%20%20%20%20%20n_genes%3D6%2C%0A%20%20%20%20%20%20%20%20n_samples%3D2048%2C%0A%20%20%20%20%20%20%20%20min_mean_umi%3D5.0%2C%0A%20%20%20%20%20%20%20%20density_method%3D%22kde%22%2C%0A%20%20%20%20%20%20%20%20figsize%3D(12%2C%2012)%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Diagonals%20(1D%20marginals)%20match.%20Off-diagonals%20(2D%20joint%20compositions)%20do%20not.%20The%20data%20shows%20tilted%20ridges%20in%20several%20panels%20%E2%80%94%20pairs%20of%20genes%20whose%20compositional%20fractions%20covary%20across%20cells%20%E2%80%94%20that%20the%20SVI%20model's%20predictive%20contours%20render%20as%20round%20blobs.%20**The%20model's%20generative%20story%20does%20not%20encode%20that%20joint%20structure.**%20It%20cannot%3A%20NBVCP%20emits%20counts%20independently%20across%20genes%20given%20the%20per-cell%20capture%2C%20so%20the%20only%20cross-gene%20%22correlation%22%20available%20is%20what%20flows%20through%20the%20shared%20%24%5Cphi%24%20(or%20the%20shared%20posterior%20on%20cell-level%20quantities).%20That%20is%20fundamentally%20a%20*marginal*%20coupling%2C%20not%20a%20joint%20one.%0A%0A%20%20%20%20This%20is%20the%20gap%20NBLN%20closes.%20Its%20log-rate%20prior%20%24%5Cmathcal%7BN%7D(%5Cunderline%7B%5Cmu%7D%2C%20%5Cunderline%7B%5Cunderline%7BW%7D%7D%5Cunderline%7B%5Cunderline%7BW%7D%7D%5E%7B%5C!%5Ctop%7D%20%2B%20%5Cmathrm%7Bdiag%7D(%5Cunderline%7Bd%7D))%24%20is%20a%20**joint**%20model%20for%20cell-level%20expression%2C%20so%20the%20same%20picture%20above%20is%20the%20criterion%20we%20will%20use%20to%20judge%20whether%20NBLN%20succeeds%3A%20do%20the%20compositional%20corner%20panels%20move%20from%20round%20blobs%20to%20tilted%20ellipses%20that%20match%20the%20data%3F%0A%0A%20%20%20%20Before%20fitting%20NBLN%2C%20one%20more%20useful%20preview%20from%20the%20SVI%20cascade%20source.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20scribe%2C%20svi_results)%3A%0A%20%20%20%20scribe.viz.plot_corner_ppc(%0A%20%20%20%20%20%20%20%20svi_results%2C%0A%20%20%20%20%20%20%20%20adata%2C%0A%20%20%20%20%20%20%20%20n_genes%3D5%2C%0A%20%20%20%20%20%20%20%20n_samples%3D512%2C%0A%20%20%20%20%20%20%20%20figsize%3D(8%2C%208)%2C%0A%20%20%20%20%20%20%20%20n_contour_levels%3D3%2C%0A%20%20%20%20%20%20%20%20density_method%3D%22kde%22%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%60plot_corner_ppc%60%20is%20the%20*count-space*%20analogue%20of%20the%20compositional%20corner%3A%20instead%20of%20comparing%20fractions%2C%20it%20compares%20raw%20UMI%20joints.%20You%20see%20the%20same%20pattern%20there%20%E2%80%94%20marginals%20OK%2C%20off-diagonals%20systematically%20rounder%20than%20the%20data.%20NBVCP%20has%20no%20mechanism%20to%20fix%20this.%20NBLN%20does.%0A%0A%20%20%20%20There%20is%20one%20important%20nuance%20to%20call%20out%20before%20we%20move%20on.%20The%20data%20clouds%20in%20the%20count-space%20corner%20show%20strong-looking%20diagonal%20alignments%3A%20as%20gene%20A's%20UMI%20count%20goes%20up%20across%20cells%2C%20so%20does%20gene%20B's.%20*Most%20of%20that%20diagonal%20alignment%20is%20not%20gene-gene%20biology%20%E2%80%94%20it%20is%20library-size.*%20Deeply-sequenced%20cells%20score%20higher%20on%20essentially%20every%20gene%2C%20and%20shallowly-sequenced%20cells%20score%20lower%20on%20essentially%20every%20gene%2C%20simply%20because%20they%20have%20more%20(or%20fewer)%20total%20reads.%20So%20a%20count-space%20pair%20plot%20pools%20%22two%20genes%20covary%20because%20the%20cell%20happens%20to%20be%20deep%22%20with%20%22two%20genes%20covary%20because%20they%20belong%20to%20the%20same%20regulatory%20module%22%2C%20and%20the%20former%20dominates%20whenever%20library%20size%20varies%20(which%20is%20always%2C%20in%20real%20data).%0A%0A%20%20%20%20This%20is%20exactly%20why%20**compositional**%20corner%20plots%20are%20the%20better%20diagnostic%20for%20the%20joint-structure%20question.%20By%20dividing%20through%20by%20the%20total%20per%20cell%2C%20the%20composition%20removes%20the%20library-size%20axis%20and%20leaves%20whatever%20residual%20coupling%20actually%20represents%20co-regulation.%20The%20model%20is%20being%20asked%20to%20reproduce%20that%20residual%2C%20not%20the%20library-size%20echo.%20Keep%20this%20in%20mind%20when%20you%20scan%20the%20count-space%20corner%20panels%3A%20the%20empirical%20clouds%20will%20look%20more%20correlated%20than%20the%20compositional%20clouds%20for%20purely%20mechanical%20reasons%2C%20and%20the%20*gap*%20between%20them%20is%20roughly%20the%20size%20of%20the%20library-size%20effect%20we%20are%20about%20to%20subtract%20away.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Stage%202%20%E2%80%94%20NBLN-Laplace%20with%20the%20full%20cascade%0A%0A%20%20%20%20Now%20we%20run%20the%20NBLN-Laplace%20fit%20that%20consumes%20the%20SVI%20source.%20There%20are%20three%20new%20things%20going%20on%20under%20the%20hood%3B%20here%20is%20what%20each%20is%20doing%20and%20*why*.%0A%0A%20%20%20%20%23%23%23%20Piece%201%3A%20%60informative_priors_from%3Dsvi_results%60%0A%0A%20%20%20%20This%20tells%20%60scribe.fit%60%20to%20derive%20empirical%20Gaussian%20priors%20on%20the%20NBLN%20parameters%20from%20the%20SVI%20source's%20posterior%20samples.%20Concretely%2C%20posterior%20draws%20from%20the%20SVI%20fit%20are%20moment-matched%20in%20NBLN's%20target%20coordinate%20system%20to%20produce%20per-gene%20priors%20on%20the%20dispersion%20%24r_g%24%2C%20on%20the%20per-gene%20log-rate%20offset%20%24%5Cmu_g%24%2C%20and%20per-cell%20priors%20on%20%24%5Ceta_c%24.%20The%20NBLN-Laplace%20loss%20then%20includes%20these%20as%20proper%20Gaussian%20log-prob%20terms%20(in%20the%20unconstrained%20coordinate).%20The%20effect%3A%20optimization%20starts%20in%20a%20basin%20that%20the%20SVI%20source%20already%20established%20as%20%22consistent%20with%20the%20data%22%2C%20and%20the%20new%20NBLN-specific%20parameters%20(%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%2C%20%24%5Cunderline%7Bd%7D%24%2C%20and%20the%20latent%20%24%5Cunderline%7Bx%7D_c%24%20per%20cell)%20have%20somewhere%20sensible%20to%20anchor%20to.%0A%0A%20%20%20%20Without%20this%20anchor%2C%20NBLN-Laplace%20cold-starts%20on%20a%20non-convex%20landscape%20with%20a%20strong%20gauge%20degeneracy.%20The%20cascade%20replaces%20%22where%20do%20we%20start%3F%22%20with%20%22we%20start%20from%20a%20fit%20that%20is%20already%20calibrated%22.%0A%0A%20%20%20%20%23%23%23%23%20A%20Bayesian%20footnote%3A%20this%20is%20empirical%20Bayes%2C%20on%20purpose%0A%0A%20%20%20%20A%20statistically-careful%20reader%20will%20notice%20we%20are%20doing%20something%20unusual%3A%20we%20are%20*learning%20prior%20hyperparameters%20from%20the%20same%20data%20we%20are%20about%20to%20use%20to%20fit%20the%20posterior*.%20That%20is%20the%20textbook%20definition%20of%20**empirical%20Bayes**%2C%20and%20there%20is%20a%20fair%20amount%20of%20(warranted)%20caution%20in%20the%20Bayesian%20literature%20about%20it%20%E2%80%94%20done%20carelessly%2C%20it%20can%20underestimate%20posterior%20uncertainty.%0A%0A%20%20%20%20Why%20are%20we%20comfortable%20doing%20it%20here%3F%0A%0A%20%20%20%201.%20**The%20parameters%20we%20are%20turning%20into%20priors%20are%20the%20ones%20the%20correlated%20model%20has%20the%20most%20trouble%20pinning%20down%20on%20its%20own.**%20Per-cell%20capture%20%24%5Ceta_c%24%20is%20gauge-degenerate%20in%20NBLN%3B%20per-gene%20dispersion%20%24r_g%24%20couples%20nonlinearly%20to%20the%20latent%20log-rate%20through%20the%20NB.%20Estimating%20these%20in%20the%20*uncorrelated*%20model%20(NBVCP)%2C%20where%20the%20per-cell%20log-rates%20collapse%20into%20a%20tidy%20Dirichlet-multinomial%20story%20with%20no%20extra%20latent%20layer%2C%20is%20dramatically%20easier%20and%20more%20stable.%20We%20are%20using%20the%20easy%20fit%20to%20inform%20the%20hard%20fit%2C%20not%20running%20two%20copies%20of%20the%20same%20fit%20on%20the%20same%20data.%0A%0A%20%20%20%202.%20**In%20the%20large-data%20limit%2C%20empirical-Bayes%20double-dipping%20has%20a%20vanishing%20cost.**%20With%20~3%2C200%20cells%20and%20tens%20of%20thousands%20of%20UMIs%20each%2C%20the%20prior%20hyperparameters%20we%20learn%20are%20pinned%20down%20to%20a%20fraction%20of%20their%20plausible%20range%20%E2%80%94%20far%20tighter%20than%20the%20residual%20uncertainty%20in%20the%20NBLN%20posterior.%20So%20when%20NBLN-Laplace%20re-uses%20these%20priors%2C%20the%20%22data%20was%20already%20used%20to%20inform%20the%20prior%22%20effect%20is%20statistically%20dominated%20by%20the%20much%20larger%20amount%20of%20data%20being%20explained%20by%20the%20new%20model%20parameters.%20This%20is%20a%20property%20of%20large-N%20empirical%20Bayes%20that%20is%20worked%20out%20in%20detail%20in%20the%20Bayesian%20statistics%20literature%3B%20a%20useful%20intuition%20is%20that%20the%20priors%20are%20informative%20on%20a%20subset%20of%20nuisance%20parameters%20whose%20posterior%20the%20new%20model's%20data%20would%20barely%20sharpen%20anyway.%0A%0A%20%20%20%203.%20**The%20alternative%20%E2%80%94%20uninformative%20priors%20on%20%24r_g%24%20and%20%24%5Ceta_c%24%20%E2%80%94%20is%20genuinely%20worse.**%20Without%20the%20cascade%2C%20NBLN-Laplace%20has%20to%20estimate%20the%20gauge-vulnerable%20parameters%20and%20the%20new%20correlation%20parameters%20at%20the%20same%20time%2C%20and%20the%20optimizer%20wanders%20along%20the%20gauge%20axis%20because%20the%20likelihood%20is%20flat%20there.%20The%20empirical-Bayes%20step%20gives%20the%20optimizer%20a%20non-flat%20objective%20to%20follow.%20The%20(small)%20statistical%20cost%20of%20double-dipping%20is%20paid%20for%20many%20times%20over%20by%20the%20(large)%20numerical-stability%20gain.%0A%0A%20%20%20%20The%20right%20mental%20model%20is%20%22we%20are%20doing%20maximum-likelihood%20estimation%20of%20certain%20nuisance%20parameters%20under%20a%20simpler%2C%20well-identified%20model%2C%20and%20then%20conditioning%20on%20those%20estimates%20as%20we%20move%20to%20a%20richer%2C%20less-identified%20model.%22%20This%20is%20a%20well-trodden%20Bayesian%20workflow%20%E2%80%94%20much%20like%20using%20a%20converged%20SVI%20fit%20to%20initialize%20MCMC%20for%20fields%20where%20MCMC%20alone%20would%20not%20mix.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Piece%202%3A%20%60informative_priors_freeze%3D(%22dispersion%22%2C%20%22capture_efficiency%22)%60%0A%0A%20%20%20%20This%20is%20where%20the%20gauge%20gets%20killed.%20The%20default%20freeze%20tuple%20%E2%80%94%20equivalent%20to%20using%20the%20parameter%20names%20%60(%22r%22%2C%20%22eta%22)%60%20%E2%80%94%20tells%20NBLN%20to%20**not%20refine**%20%24r_g%24%20and%20%24%5Ceta_c%24%20further.%20They%20are%20pinned%20at%20the%20SVI%20cascade%20source's%20MAP%20values%20and%20excluded%20from%20the%20optimizer's%20parameter%20dict%20entirely.%20Only%20%24%5Cunderline%7B%5Cmu%7D%24%2C%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%2C%20%24%5Cunderline%7Bd%7D%24%2C%20and%20the%20per-cell%20latents%20update%20during%20NBLN-Laplace.%0A%0A%20%20%20%20Pinning%20%24%5Ceta_c%24%20structurally%20fixes%20the%20rigid-translation%20gauge%3A%20now%20there%20is%20no%20%24%5CDelta_c%24%20direction%20along%20which%20the%20loss%20is%20flat%2C%20because%20every%20%24%5Ceta_c%24%20is%20a%20constant.%20Pinning%20%24r_g%24%20at%20SVI's%20values%20is%20largely%20a%20numerical-stability%20choice%20%E2%80%94%20NBLN%20cannot%20meaningfully%20improve%20on%20the%20SVI%20source's%20dispersion%20posterior%20here%2C%20and%20refining%20%24r_g%24%20alongside%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20adds%20noise%20to%20the%20M-step%20without%20buying%20anything.%0A%0A%20%20%20%20Two%20important%20details%3A%0A%0A%20%20%20%20-%20The%20descriptive%20aliases%20%60%22dispersion%22%60%20and%20%60%22capture_efficiency%22%60%20resolve%20to%20the%20canonical%20names%20%60%22r%22%60%20and%20%60%22eta%22%60%20automatically%20(see%20%60FREEZE_KEY_ALIASES%60).%20Either%20form%20works.%0A%20%20%20%20-%20Frozen%20parameters%20are%20still%20**honest%20probability%20distributions**%20when%20you%20ask%20the%20result%20for%20them%20%E2%80%94%20they%20route%20through%20the%20embedded%20SVI%20cascade%20source%2C%20so%20%60result.get_distributions()%60%20returns%20moment-matched%20Normals%20reflecting%20the%20SVI%20posterior%2C%20not%20point%20masses.%20The%20point%20masses%20are%20only%20used%20during%20the%20optimization.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Piece%203%3A%20%60priors%3D%7B%22loadings%22%3A%20%7B%22type%22%3A%20%22horseshoe_columnwise%22%2C%20%22tau_scale%22%3A%201.0%7D%7D%60%0A%0A%20%20%20%20The%20loadings%20matrix%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20used%20to%20build%20the%20gene-gene%20low-rank%20covariance%20matrix%20has%20shape%20%24(G%2C%20k)%24%20with%20%24k%20%3D%2032%24.%20We%20do%20not%20know%20in%20advance%20how%20many%20of%20those%2032%20factors%20the%20data%20actually%20supports.%20If%20we%20let%20the%20optimizer%20do%20whatever%20it%20wants%2C%20we%20will%20get%20a%20fit%20that%20uses%20all%2032%20%E2%80%94%20but%20most%20of%20them%20will%20be%20fitting%20noise.%20Worse%2C%20in%20a%20low-N%20regime%20the%20resulting%20%22correlations%22%20will%20be%20dominated%20by%20overfitting%20structure.%0A%0A%20%20%20%20A%20column-wise%20horseshoe%20prior%20says%3A%20*each%20column%20of%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20has%20its%20own%20scale%20%24%5Clambda_k%24%2C%20and%20that%20scale%20is%20drawn%20from%20a%20hierarchical%20prior%20that%20strongly%20prefers%20small%20values%20but%20has%20heavy%20tails%20for%20the%20few%20columns%20that%20genuinely%20need%20to%20be%20large.*%20In%20effect%2C%20the%20prior%20performs%20**adaptive%20rank%20selection**%20during%20fitting%20%E2%80%94%20columns%20that%20the%20data%20does%20not%20support%20shrink%20toward%20zero%3B%20columns%20it%20does%20support%20remain%20unshrunk.%20Mathematically%20this%20is%20a%20sparsity-inducing%20prior%20in%20the%20*column*%20dimension%2C%20analogous%20to%20(but%20more%20flexible%20than)%20hard-truncating%20%24k%24%20to%20a%20smaller%20number.%0A%0A%20%20%20%20The%20benefit%20over%20manually%20picking%20%24k%24%3A%20you%20can%20set%20%24k%24%20generously%20(here%2C%2032)%2C%20let%20the%20prior%20figure%20out%20the%20effective%20rank%2C%20and%20inspect%20the%20spectrum%20after%20the%20fact%20to%20see%20what%20the%20data%20actually%20supported.%20We%20will%20do%20exactly%20that%20with%20the%20%60plot_w_shrinkage_spectrum%60%20diagnostic.%0A%0A%20%20%20%20%23%23%23%20Other%20knobs%0A%0A%20%20%20%20-%20%60parameterization%3D%22canonical%22%60%20%E2%80%94%20NBLN's%20only%20supported%20parameterization.%0A%20%20%20%20-%20%60vae_latent_dim%3D32%60%20%E2%80%94%20the%20%24k%24%20in%20the%20low-rank%20covariance.%20Generous%3B%20the%20shrinkage%20prior%20decides%20the%20effective%20value.%0A%20%20%20%20-%20%60d_mode%3D%22learned%22%60%20%E2%80%94%20let%20the%20per-gene%20residual%20variance%20%24d_g%24%20be%20optimized%20rather%20than%20fixed.%0A%20%20%20%20-%20%60optimizer_config%3D%7B...clipped_adam...%7D%60%20and%20%60laplace_config%3D%7B...newton_max_step...%7D%60%20%E2%80%94%20numerical-stability%20knobs%20for%20the%20inner%20Newton%20solve%20and%20the%20outer%20M-step%20in%20the%20EM%20optimization.%20NBLN%20is%20more%20demanding%20than%20NBVCP%20in%20this%20respect.%0A%20%20%20%20-%20%60n_steps%3D20_000%60%20is%20much%20smaller%20than%20the%20SVI%20source's%20%60500_000%60%20because%20the%20M-step%20here%20is%20doing%20a%20lighter%20job%3A%20the%20heavy%20lifting%20is%20in%20the%20SVI%20cascade%20source%2C%20and%20the%20freeze%20removes%20two-thirds%20of%20the%20parameter%20blocks%20from%20active%20optimization.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20clear_caches%2C%20gc%2C%20np%2C%20out_dir%2C%20pickle%2C%20scribe%2C%20svi_results)%3A%0A%20%20%20%20clear_caches()%0A%20%20%20%20gc.collect()%0A%0A%20%20%20%20%23%20Match%20the%20EDA%20file's%20naming%20convention%20so%20the%20cached%20pickle%20from%0A%20%20%20%20%23%20the%20exploration%20is%20picked%20up%20automatically%20when%20present.%0A%20%20%20%20_gene_coverage_laplace%20%3D%200.99%0A%20%20%20%20_latent_dim%20%3D%2032%0A%20%20%20%20_capture_efficiency_prior%20%3D%20(np.log(100_000)%2C%200.1)%0A%0A%20%20%20%20_out_path_laplace%20%3D%20out_dir%20%2F%20(%0A%20%20%20%20%20%20%20%20f%22scribe_results_nbln-%22%0A%20%20%20%20%20%20%20%20f%22canonical-vcp-%22%0A%20%20%20%20%20%20%20%20f%22%7B_latent_dim%7Dlatent-%22%0A%20%20%20%20%20%20%20%20f%22%7B_gene_coverage_laplace%7Dgene_coverage-%22%0A%20%20%20%20%20%20%20%20f%22anchor-learneddmode_%22%0A%20%20%20%20%20%20%20%20f%22horseshoe_columnwisewprior_laplace.pkl%22%0A%20%20%20%20)%0A%20%20%20%20_laplace_checkpoint_dir%20%3D%20out_dir%20%2F%20(%0A%20%20%20%20%20%20%20%20f%22checkpoints_nbln-%22%0A%20%20%20%20%20%20%20%20f%22canonical-vcp-%22%0A%20%20%20%20%20%20%20%20f%22%7B_latent_dim%7Dlatent-%22%0A%20%20%20%20%20%20%20%20f%22%7B_gene_coverage_laplace%7Dgene_coverage-%22%0A%20%20%20%20%20%20%20%20f%22horseshoe_columnwisewprior_%22%0A%20%20%20%20%20%20%20%20f%22anchor-learneddmode_laplace%22%0A%20%20%20%20)%0A%0A%20%20%20%20if%20_out_path_laplace.exists()%3A%0A%20%20%20%20%20%20%20%20with%20open(_out_path_laplace%2C%20%22rb%22)%20as%20_f%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20laplace_results%20%3D%20pickle.load(_f)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20_laplace_checkpoint_dir.mkdir(parents%3DTrue%2C%20exist_ok%3DTrue)%0A%20%20%20%20%20%20%20%20laplace_results%20%3D%20scribe.fit(%0A%20%20%20%20%20%20%20%20%20%20%20%20adata%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20model%3D%22nbln%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20parameterization%3D%22canonical%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20inference_method%3D%22laplace%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20informative_priors_from%3Dsvi_results%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20informative_priors_freeze%3D(%22dispersion%22%2C%20%22capture_efficiency%22)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20informative_priors_tau%3D1.0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20informative_priors_n_samples%3D1000%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20gene_coverage%3D_gene_coverage_laplace%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20d_mode%3D%22learned%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20latent_dim%3D_latent_dim%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20n_steps%3D20_000%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20batch_size%3DNone%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20unconstrained%3DTrue%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20priors%3D%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22capture_efficiency%22%3A%20_capture_efficiency_prior%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22loadings%22%3A%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22type%22%3A%20%22horseshoe_columnwise%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22tau_scale%22%3A%201.0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20early_stopping%3D%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22enabled%22%3A%20False%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22checkpoint_dir%22%3A%20str(_laplace_checkpoint_dir)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20optimizer_config%3D%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22name%22%3A%20%22clipped_adam%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22step_size%22%3A%201e-3%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22grad_clip_norm%22%3A%2010.0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20laplace_config%3D%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22n_newton_steps%22%3A%205%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22damping%22%3A%201e-6%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22newton_tolerance%22%3A%201e-2%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22newton_max_step%22%3A%2010.0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20with%20open(_out_path_laplace%2C%20%22wb%22)%20as%20_f%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20pickle.dump(laplace_results%2C%20_f)%0A%0A%20%20%20%20laplace_results%0A%20%20%20%20return%20(laplace_results%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Did%20training%20converge%3F%20(or%3A%20why%20NBLN-Laplace%20loss%20curves%20look%20weird)%0A%0A%20%20%20%20NBLN-Laplace's%20loss%20is%20a%20Laplace%20approximation%20to%20the%20ELBO%3A%20it%20is%20the%20negative%20log-likelihood%20of%20the%20data%20evaluated%20at%20each%20cell's%20MAP%20latent%2C%20plus%20a%20half-log-det%20term%20that%20accounts%20for%20posterior%20curvature%2C%20plus%20the%20global%20prior%20log-probs%20(including%20the%20cascade%20priors%20and%20the%20loadings%20shrinkage%20prior).%20For%20models%20like%20SVI%20or%20the%20Gaussian-conjugate%20paths%20we%20have%20seen%20before%2C%20the%20convergence%20story%20is%20simple%3A%20the%20loss%20drops%20fast%20and%20then%20flattens.%0A%0A%20%20%20%20For%20NBLN-Laplace%20it%20is%20**not**%20that%20simple%2C%20and%20you%20should%20not%20panic%20when%20you%20see%20what%20comes%20next.%0A%0A%20%20%20%20The%20loss%20curve%20typically%20does%20drop%20fast%20%E2%80%94%20it%20usually%20settles%20into%20a%20low-loss%20region%20within%20the%20first%20~10K%20outer-loop%20iterations%20%E2%80%94%20but%20past%20that%20point%20it%20is%20common%20to%20see%20the%20curve%20**oscillate%20wildly**%2C%20occasionally%20spiking%20back%20up%20by%20orders%20of%20magnitude%20before%20recovering.%20This%20is%20not%20a%20bug%20in%20the%20model%20or%20in%20the%20fit%3B%20it%20is%20a%20structural%20property%20of%20the%20optimization%20problem%20we%20are%20solving.%0A%0A%20%20%20%20Why%20does%20it%20happen%3F%20Two%20things%20compose%20into%20instability%3A%0A%0A%20%20%20%201.%20**The%20inner%20Newton%20solve%20is%20approximate.**%20Each%20outer-loop%20step%20needs%20the%20per-cell%20%24%5Cunderline%7Bx%7D_c%5E%5Cast%24%20MAP%20that%20satisfies%20%24%5Cnabla%20%5Clog%20%5Cpi(%5Cunderline%7Bu%7D_c%20%5Cmid%20%5Cunderline%7Bx%7D_c%2C%20%5Ctheta)%20%3D%200%24.%20We%20get%20there%20by%20a%20small%20fixed%20number%20of%20Newton%20iterations%20per%20cell%20(%60n_newton_steps%3D5%60)%2C%20with%20a%20step-size%20cap%20(%60newton_max_step%60)%20and%20damping.%20Near%20a%20well-shaped%20basin%20those%20few%20iterations%20land%20essentially%20on%20the%20mode.%20Far%20from%20the%20basin%20%E2%80%94%20or%20when%20the%20outer-loop%20%24%5Ctheta%24%20updates%20push%20the%20basin%20sharply%20%E2%80%94%20the%20inner%20Newton%20can%20take%20a%20step%20that%20lands%20somewhere%20with%20much%20higher%20loss.%20That%20higher%20loss%20feeds%20back%20into%20the%20outer%20step%2C%20the%20outer%20step%20over-corrects%2C%20and%20the%20cycle%20repeats.%0A%20%20%20%202.%20**The%20Hessian's%20%60log-det%60%20term%20is%20non-monotone.**%20The%20Laplace%20correction%20includes%20%24%5Ctfrac%7B1%7D%7B2%7D%20%5Clog%20%5Cdet(-H)%24%2C%20which%20is%20unbounded%20above%20and%20can%20change%20quickly%20as%20the%20parameters%20move%20around%20the%20curvature%20landscape.%20Tiny%20outer-loop%20steps%20can%20change%20%60log-det%60%20by%20a%20large%20amount%20even%20when%20the%20data%20fit%20is%20essentially%20unchanged.%0A%0A%20%20%20%20%60scribe%60's%20implementation%20has%20two%20insurance%20policies%20for%20this.%20The%20first%20is%20**gradient%20clipping%20and%20damping**%20on%20the%20outer%20Adam%20step%20(%60grad_clip_norm%3D10.0%60%2C%20%60damping%3D1e-6%60)%20%E2%80%94%20these%20prevent%20the%20worst%20spikes%20from%20happening%20at%20all%20by%20capping%20each%20outer%20move.%20The%20second%2C%20and%20the%20one%20that%20actually%20matters%20for%20your%20final%20result%2C%20is%20**best-loss%20snapshotting**%3A%20every%20outer-loop%20iteration%2C%20the%20driver%20keeps%20a%20copy%20of%20the%20parameters%20whenever%20a%20new%20global%20minimum%20of%20the%20loss%20is%20achieved%2C%20and%20at%20the%20end%20of%20training%20it%20returns%20*those*%20parameters%2C%20not%20the%20last-iteration%20parameters.%20So%20even%20if%20the%20curve%20looks%20frantic%20in%20the%20second%20half%20of%20training%2C%20the%20returned%20%60laplace_results%60%20object%20reflects%20the%20best-loss%20snapshot%2C%20not%20whatever%20volatile%20state%20the%20optimizer%20ended%20on.%20The%20(well-tested)%20machinery%20for%20this%20lives%20in%20the%20Laplace-EM%20driver%20and%20is%20on%20by%20default%3B%20you%20do%20not%20need%20to%20configure%20it.%0A%0A%20%20%20%20Practically%2C%20when%20you%20look%20at%20the%20loss%20curve%20below%2C%20focus%20on%20whether%20the%20*envelope*%20of%20the%20curve%20settles%20low%20and%20stays%20low.%20Late-stage%20oscillations%20on%20top%20of%20a%20low%20envelope%20are%20fine%20%E2%80%94%20%60scribe%60%20will%20pick%20out%20the%20best%20moment.%20A%20loss%20curve%20that%20drifts%20steadily%20upward%20and%20never%20settles%20is%20the%20actually-bad%20pattern%20that%20does%20deserve%20worry%2C%20and%20in%20our%20cascade-frozen%20regime%20it%20almost%20never%20happens%20because%20the%20cascade%20priors%20anchor%20the%20optimizer%20to%20a%20good%20basin.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(laplace_results%2C%20scribe)%3A%0A%20%20%20%20scribe.viz.plot_loss(laplace_results%2C%20figsize%3D(6%2C%203))%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20curve%20above%20shows%20exactly%20the%20picture%20we%20described%3A%20a%20clean%20initial%20drop%20into%20a%20low-loss%20region%2C%20then%20oscillations%20on%20top%20of%20that%20envelope.%20Because%20%60scribe%60%20tracks%20the%20running%20best-loss%20snapshot%2C%20the%20parameters%20in%20%60laplace_results%60%20correspond%20to%20the%20*bottom*%20of%20that%20envelope%20%E2%80%94%20not%20whatever%20spike%20the%20curve%20happened%20to%20end%20on.%0A%0A%20%20%20%20%23%23%23%20Sanity%20check%20the%20cascade-frozen%20quantities%0A%0A%20%20%20%20The%20freeze%20should%20mean%20that%20%24r_g%24%20and%20%24%5Ceta_c%24%20from%20the%20NBLN%20result%20are%20*the%20same*%20as%20the%20SVI%20source's%20values%20(up%20to%20the%20moment-matching%20transform).%20The%20%60frozen_params%60%20field%20on%20the%20result%20tells%20you%20which%20parameters%20were%20frozen%2C%20and%20%60cascade_source%60%20embeds%20the%20SVI%20results%20object%20by%20reference%20so%20you%20can%20recover%20full-fidelity%20posterior%20samples%20for%20any%20frozen%20quantity.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(laplace_results)%3A%0A%20%20%20%20print(%22frozen_params%3A%20%20%20%20%20%20%20%20%22%2C%20laplace_results.frozen_params)%0A%20%20%20%20print(%22cascade_source%3A%20%20%20%20%20%20%20%22%2C%20type(laplace_results.cascade_source).__name__)%0A%20%20%20%20print(%22w_prior_diagnostics%20keys%3A%22%2C%20list(%0A%20%20%20%20%20%20%20%20laplace_results.w_prior_diagnostics.keys()%0A%20%20%20%20))%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%60frozen_params%3D%7B'r'%2C%20'eta'%7D%60%20confirms%20the%20freeze%20contract%3B%20the%20%60cascade_source%60%20is%20the%20embedded%20SVI%20result%3B%20and%20%60w_prior_diagnostics%60%20is%20a%20fresh%20dict%20with%20the%20loadings-shrinkage%20summary%20we%20will%20inspect%20next.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Stage%203%20%E2%80%94%20diagnose%20the%20loadings%20spectrum%0A%0A%20%20%20%20With%20the%20cascade%20and%20the%20freeze%20handling%20capture%2Fdispersion%2C%20the%20only%20%22discovery%22%20the%20NBLN-Laplace%20fit%20is%20doing%20is%20on%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20and%20%24%5Cunderline%7Bd%7D%24.%20The%20horseshoe%20prior%20is%20supposed%20to%20have%20killed%20the%20noise%20columns%20of%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20%E2%80%94%20but%20we%20should%20*check*%2C%20because%20if%20we%20set%20%24k%3D32%24%20and%20only%20%24%5Csim%205%24%20columns%20survive%2C%20knowing%20that%20explicitly%20changes%20how%20we%20interpret%20the%20correlation%20structure.%0A%0A%20%20%20%20%60scribe.viz.plot_w_shrinkage_spectrum%60%20plots%20two%20things%20on%20the%20same%20axis%3A%0A%0A%20%20%20%201.%20**Solid%20line%3A**%20the%20per-column%20compositional%20norm%20%24%5C%7C%5Cunderline%7B%5Cunderline%7BW%7D%7D_%7B%5Cperp%2Ck%7D%5C%7C%24%20%E2%80%94%20the%20gauge-invariant%20%22size%22%20of%20each%20factor.%20This%20is%20the%20headline%20diagnostic%3A%20factors%20with%20large%20compositional%20norm%20carry%20biological%20signal%3B%20factors%20near%20zero%20are%20noise%20the%20prior%20shrunk%20away.%0A%20%20%20%202.%20**Dashed%20line%3A**%20the%20per-column%20aux%20scale%20%24%5Csigma_k%24%20%E2%80%94%20the%20variational%20parameter%20that%20controls%20how%20tightly%20column%20%24k%24%20is%20shrunk.%20Heavy-tailed%20horseshoe%20scales%20are%20slightly%20less%20identifiable%20than%20the%20column%20norms%2C%20so%20the%20column-norm%20spectrum%20is%20the%20headline.%20The%20aux%20scales%20should%20track%20the%20column%20norms%20qualitatively.%0A%0A%20%20%20%20The%20horizontal%20dashed%20line%20marks%20the%205%25-of-maximum%20threshold%20that%20%60column_norm_effective_rank%60%20(a%20number%20on%20the%20result)%20uses%20to%20count%20%22active%22%20factors.%20**Look%20for%20an%20elbow%3A**%20if%20the%20spectrum%20drops%20sharply%20at%20some%20%24k_%5Ctext%7Beff%7D%24%20and%20then%20flattens%2C%20the%20data%20supports%20%24k_%5Ctext%7Beff%7D%24%20factors.%20If%20the%20spectrum%20is%20flat%20all%20the%20way%20out%20to%20%24k%3D32%24%2C%20the%20prior%20was%20not%20strong%20enough%20and%20you%20need%20to%20tighten%20%60tau_scale%60.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(laplace_results%2C%20scribe)%3A%0A%20%20%20%20scribe.viz.plot_w_shrinkage_spectrum(laplace_results%2C%20figsize%3D(5%2C%203.5))%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20spectrum%20shows%20a%20clear%20elbow.%20A%20handful%20of%20leading%20factors%20carry%20essentially%20all%20the%20compositional%20signal%3B%20the%20long%20tail%20beyond%20the%20elbow%20is%20at%20or%20below%20the%205%25-of-max%20threshold.%20The%20%60column_norm_effective_rank%60%20field%20summarizes%20this%20as%20a%20single%20number.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(laplace_results%2C%20np)%3A%0A%20%20%20%20_diag%20%3D%20laplace_results.w_prior_diagnostics%0A%20%20%20%20_sigma_k%20%3D%20np.asarray(_diag%5B%22sigma_k%22%5D)%0A%20%20%20%20print(%22strategy_type%3A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22%2C%20_diag%5B%22strategy_type%22%5D)%0A%20%20%20%20print(%22column_norm_effective_rank%3A%20%20%20%20%20%22%2C%20_diag%5B%22column_norm_effective_rank%22%5D)%0A%20%20%20%20print(%22effective_rank%20(alias)%3A%20%20%20%20%20%20%20%20%20%22%2C%20_diag%5B%22effective_rank%22%5D)%0A%20%20%20%20print(%0A%20%20%20%20%20%20%20%20%22tau%20(global%20horseshoe%20scale)%3A%20%20%20%20%22%0A%20%20%20%20%20%20%20%20f%22%7Bfloat(_diag%5B'tau'%5D)%3A.4f%7D%22%0A%20%20%20%20)%0A%20%20%20%20print(%0A%20%20%20%20%20%20%20%20%22sigma_k%20min%2Fmedian%2Fmax%3A%20%20%20%20%20%20%20%20%20%20%22%0A%20%20%20%20%20%20%20%20f%22%7Bfloat(_sigma_k.min())%3A.4f%7D%20%2F%20%22%0A%20%20%20%20%20%20%20%20f%22%7Bfloat(np.median(_sigma_k))%3A.4f%7D%20%2F%20%22%0A%20%20%20%20%20%20%20%20f%22%7Bfloat(_sigma_k.max())%3A.4f%7D%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20effective%20rank%20is%20the%20headline%20number%20to%20report%20alongside%20the%20model.%20It%20tells%20you%3A%20%22out%20of%2032%20latent%20factors%20I%20gave%20the%20model%2C%20the%20data%20supports%20roughly%20%24k_%5Ctext%7Beff%7D%24%20of%20them%20after%20shrinkage.%22%20For%20Jurkat%20%E2%80%94%20a%20monoculture%20with%20limited%20biological%20heterogeneity%20%E2%80%94%20a%20small%20%24k_%5Ctext%7Beff%7D%24%20is%20sensible.%0A%0A%20%20%20%20Without%20the%20loadings%20prior%2C%20this%20number%20would%20have%20been%2032%20(by%20construction)%20and%20we%20would%20not%20be%20able%20to%20distinguish%20%22data-supported%20factor%22%20from%20%22noise%20direction%20the%20model%20fit%20anyway%22.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Cross-check%3A%20the%20singular%20spectrum%20of%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5Cperp%24%0A%0A%20%20%20%20The%20shrinkage-spectrum%20plot%20tells%20you%20the%20*per-column*%20compositional%20norm%20%E2%80%94%20but%20column%20norms%20are%20not%20the%20only%20thing%20that%20matters.%20Two%20columns%20with%20the%20same%20norm%20can%20encode%20very%20different%20amounts%20of%20information%20depending%20on%20how%20they%20line%20up%20with%20each%20other%20and%20with%20the%20data's%20actual%20covariance%20structure.%0A%0A%20%20%20%20The%20right%20object%20to%20look%20at%20for%20*how%20much%20variance%20the%20model%20assigns%20to%20each%20underlying%20mode*%20is%20the%20**singular%20value%20spectrum%20of%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5Cperp%24**%20%E2%80%94%20the%20gauge-invariant%20compositional%20loadings.%20The%20squared%20singular%20values%20%24%5Csigma_i%5E2%24%20are%20the%20variance%20contributions%20of%20mode%20%24i%24%20to%20the%20implied%20compositional%20covariance.%20If%20the%20loadings%20shrinkage%20was%20doing%20what%20we%20hoped%2C%20we%20expect%20to%20see%20the%20same%20elbow%20shape%3A%20a%20few%20large%20singular%20values%2C%20then%20a%20fast%20decay.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(laplace_results%2C%20np%2C%20plt)%3A%0A%20%20%20%20%23%20W_perp%20%3D%20W%20-%20mean(W%2C%20axis%3D0)%20is%20the%20gauge-invariant%20compositional%20loadings.%0A%20%20%20%20_W_perp%20%3D%20np.asarray(laplace_results.get_W_compositional())%0A%20%20%20%20_singular_values%20%3D%20np.linalg.svd(_W_perp%2C%20compute_uv%3DFalse)%0A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(5%2C%203.5))%0A%20%20%20%20_k_idx%20%3D%20np.arange(1%2C%20_singular_values.shape%5B0%5D%20%2B%201)%0A%20%20%20%20_ax.plot(_k_idx%2C%20_singular_values%2C%20marker%3D%22o%22%2C%20ms%3D4%2C%20lw%3D1.0)%0A%20%20%20%20_ax.set_yscale(%22log%22)%0A%20%20%20%20_ax.set_xlabel(%22rank%20index%20%24i%24%22)%0A%20%20%20%20_ax.set_ylabel(r%22%24%5Csigma_i(W_%5Cperp)%24%20%20(singular%20value)%22)%0A%20%20%20%20_ax.set_title(%22Singular%20spectrum%20of%20compositional%20loadings%22)%0A%20%20%20%20plt.show()%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20same%20elbow%20appears%20in%20the%20singular%20spectrum%3A%20a%20few%20dominant%20singular%20values%20capture%20most%20of%20the%20variance%2C%20then%20a%20long%20tail%20at%20or%20near%20the%20shrinkage%20floor.%20This%20is%20structurally%20what%20the%20horseshoe%20prior%20was%20designed%20to%20produce%2C%20and%20it%20is%20what%20we%20want%20for%20downstream%20interpretation%3A%20the%20dominant%20left-singular%20vectors%20of%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5Cperp%24%20are%20the%20**gene%20programs**%20the%20data%20supports%20%E2%80%94%20each%20one%20a%20linear%20combination%20of%20genes%20that%20co-vary%20across%20cells.%0A%0A%20%20%20%20A%20useful%20aside%20on%20what%20the%20gauge%20projection%20is%20doing.%20The%20raw%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20has%20a%20rank-1%20component%20aligned%20with%20the%20all-ones%20direction%20in%20gene%20space%2C%20which%20corresponds%20to%20%22uniform%20up%2Fdown%20across%20all%20genes%22%20%E2%80%94%20i.e.%2C%20a%20pure%20library-size%20or%20capture%20mode.%20Because%20the%20cascade%20freeze%20has%20already%20pinned%20%24%5Ceta_c%24%2C%20that%20component%20cannot%20be%20moving%20in%20response%20to%20library%20size%20during%20the%20M-step.%20But%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20can%20still%20accidentally%20pick%20up%20a%20small%20piece%20of%20it.%20Projecting%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20onto%20the%20subspace%20orthogonal%20to%20the%20all-ones%20direction%20(which%20is%20what%20%60get_W_compositional()%60%20returns)%20removes%20that%20piece%20by%20construction.%20The%20%60get_gauge_diagnostics()%60%20call%20summarizes%20how%20much%20of%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24's%20mass%20lives%20in%20the%20gauge%20versus%20orthogonal%20to%20it.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(laplace_results)%3A%0A%20%20%20%20_gauge_diag%20%3D%20laplace_results.get_gauge_diagnostics()%0A%20%20%20%20print(%22W_compositional_norm%20(W_perp)%3A%20%20%20%20%20%20%20%22%2C%20_gauge_diag%5B%22W_compositional_norm%22%5D)%0A%20%20%20%20print(%22W_all_ones_component_norm%20(W_para)%3A%20%20%22%2C%20_gauge_diag%5B%22W_all_ones_component_norm%22%5D)%0A%20%20%20%20print(%22gauge_contamination_ratio%3A%20%20%20%20%20%20%20%20%20%20%20%22%2C%20_gauge_diag%5B%22gauge_contamination_ratio%22%5D)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Interpreting%20the%20gauge-contamination%20ratio%0A%0A%20%20%20%20The%20ratio%20%24%5C%7C%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5C%7C%5C%7C%20%5C%2C%2F%5C%2C%20%5C%7C%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5Cperp%5C%7C%24%20has%20two%20very%20different%20%22healthy%22%20regimes%3A%0A%0A%20%20%20%20-%20**Without%20loadings%20shrinkage**%20(or%20with%20a%20very%20weak%20prior)%3A%20a%20ratio%20above%20%24%5Csim%200.2%24%20is%20a%20red%20flag%20%E2%80%94%20it%20means%20a%20sizable%20fraction%20of%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20is%20being%20absorbed%20by%20the%20all-ones%20gauge%20mode%20instead%20of%20representing%20real%20cross-gene%20correlation.%20A%20ratio%20below%20%24%5Csim%200.05%24%20is%20clean.%0A%20%20%20%20-%20**With%20column-wise%20horseshoe%20shrinkage**%20(the%20regime%20we%20are%20in)%3A%20ratios%20in%20the%20%240.5%24%E2%80%93%240.8%24%20range%20are%20*routine%20and%20benign*.%20Why%3F%20Because%20the%20shrinkage%20prior%20actively%20pulls%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5Cperp%24%20toward%20zero%20on%20noise%20columns%2C%20deflating%20the%20denominator.%20The%20numerator%20(%24%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5C%7C%24)%20is%20small%20in%20absolute%20terms%20but%20no%20longer%20small%20*relative*%20to%20the%20deflated%20denominator.%20The%20right%20question%20to%20ask%20is%20whether%20the%20singular%20spectrum%20has%20a%20clean%20elbow%20(yes%20%E2%80%94%20we%20just%20saw%20it)%20and%20whether%20the%20cascade%20freeze%20is%20doing%20its%20job%20(yes%20%E2%80%94%20%60frozen_params%60%20confirmed%20it).%20The%20ratio%20is%20a%20secondary%20diagnostic.%0A%0A%20%20%20%20The%20full%20theoretical%20story%20for%20these%20regimes%20is%20in%20%60docs%2Ftheory%2Floadings-shrinkage.md%60.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Stage%204%20%E2%80%94%20does%20the%20joint%20structure%20actually%20match%20the%20data%3F%0A%0A%20%20%20%20We%20have%20a%20fit%20that%20converged%2C%20a%20sparse%20shrinkage%20spectrum%2C%20and%20a%20clean%20gauge.%20But%20that%20is%20all%20*internal*%20diagnostics.%20The%20real%20test%20is%20whether%20the%20model's%20joint%20compositional%20predictive%20distribution%20matches%20the%20empirical%20joints%20%E2%80%94%20the%20same%20diagnostic%20we%20used%20on%20the%20SVI%20source%2C%20but%20now%20with%20NBLN's%20covariance%20structure%20in%20play.%0A%0A%20%20%20%20Marginals%20first.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20laplace_results%2C%20scribe)%3A%0A%20%20%20%20scribe.viz.plot_ppc(%0A%20%20%20%20%20%20%20%20laplace_results%2C%0A%20%20%20%20%20%20%20%20adata%2C%0A%20%20%20%20%20%20%20%20n_genes%3D16%2C%0A%20%20%20%20%20%20%20%20n_rows%3D4%2C%0A%20%20%20%20%20%20%20%20figsize%3D(8%2C%208)%2C%0A%20%20%20%20%20%20%20%20n_samples%3D512%2C%0A%20%20%20%20%20%20%20%20ppc_level%3D%22marginal%22%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Marginal%20PPCs%20look%20essentially%20as%20clean%20as%20the%20SVI%20source%20%E2%80%94%20the%20cascade%20succeeded%20in%20preserving%20the%20marginal%20fit%20while%20moving%20to%20a%20model%20with%20real%20joint%20structure.%0A%0A%20%20%20%20%60plot_mean_calibration%60%20is%20the%20global%20single-number-per-gene%20version%3A%20predicted%20mean%20UMIs%20vs.%20empirical%20mean%20UMIs%2C%20on%20log%E2%80%93log%20axes.%20Points%20hugging%20the%20identity%20line%20means%20the%20model's%20MAP%20reproduces%20gene-level%20expression%20levels%20across%20the%20dynamic%20range.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20laplace_results%2C%20scribe)%3A%0A%20%20%20%20scribe.viz.plot_mean_calibration(laplace_results%2C%20counts%3Dadata%2C%20figsize%3D(4%2C%204))%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Now%20the%20*joint*%20test%20%E2%80%94%20the%20one%20NBVCP%20failed%20earlier.%20We%20render%20the%20same%20compositional%20corner%20grid%20for%20the%20NBLN-Laplace%20fit%20and%20look%20for%20those%20tilted%20ellipses%20that%20the%20SVI%20source%20could%20not%20reproduce.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20laplace_results%2C%20scribe)%3A%0A%20%20%20%20scribe.viz.plot_compositional_corner_ppc(%0A%20%20%20%20%20%20%20%20laplace_results%2C%0A%20%20%20%20%20%20%20%20adata%2C%0A%20%20%20%20%20%20%20%20n_genes%3D6%2C%0A%20%20%20%20%20%20%20%20n_samples%3D2048%2C%0A%20%20%20%20%20%20%20%20min_mean_umi%3D5.0%2C%0A%20%20%20%20%20%20%20%20density_method%3D%22kde%22%2C%0A%20%20%20%20%20%20%20%20figsize%3D(12%2C%2012)%2C%0A%20%20%20%20%20%20%20%20gene_selection%3D%22correlation_diverse%22%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Compare%20this%20against%20the%20NBVCP-SVI%20compositional%20corner%20from%20earlier.%20The%20off-diagonal%20panels%20are%20no%20longer%20round%20blobs%20%E2%80%94%20the%20model's%20predictive%20contours%20now%20follow%20the%20tilts%20and%20ridges%20in%20the%20empirical%20joint.%20**This%20is%20the%20principled%20gene%E2%80%93gene%20correlation%20story%3A**%20the%20contours%20are%20coming%20from%20NBLN's%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%5Cunderline%7B%5Cunderline%7BW%7D%7D%5E%7B%5C!%5Ctop%7D%24%20structure%2C%20not%20from%20a%20richer%20variational%20guide%20on%20a%20model%20that%20emits%20independent%20counts.%0A%0A%20%20%20%20The%20argument%20%60gene_selection%3D%22correlation_diverse%22%60%20asks%20the%20plot%20to%20pick%20a%20set%20of%20genes%20whose%20pairwise%20compositional%20correlations%20span%20a%20range%20of%20magnitudes%20(not%20all%20clustered%20around%20the%20same%20value).%20That%20makes%20the%20panel%20grid%20more%20informative%20%E2%80%94%20you%20see%20panels%20where%20the%20model%20says%20%22tight%20ridge%22%20alongside%20panels%20where%20it%20says%20%22essentially%20independent%22%2C%20and%20you%20can%20check%20both%20against%20the%20data.%0A%0A%20%20%20%20A%20pair-level%20cross-check%20with%20the%20count-space%20corner%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20laplace_results%2C%20scribe)%3A%0A%20%20%20%20scribe.viz.plot_corner_ppc(%0A%20%20%20%20%20%20%20%20laplace_results%2C%0A%20%20%20%20%20%20%20%20adata%2C%0A%20%20%20%20%20%20%20%20n_genes%3D5%2C%0A%20%20%20%20%20%20%20%20n_samples%3D512%2C%0A%20%20%20%20%20%20%20%20figsize%3D(8%2C%208)%2C%0A%20%20%20%20%20%20%20%20n_contour_levels%3D3%2C%0A%20%20%20%20%20%20%20%20density_method%3D%22kde%22%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Count-space%20joint%20structure%20also%20matches.%20The%20model%20has%20internalized%20the%20dependence%20rather%20than%20smearing%20it.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Stage%205%20%E2%80%94%20reading%20the%20cross-gene%20correlation%20structure%0A%0A%20%20%20%20With%20the%20joints%20validated%2C%20we%20can%20interpret%20the%20correlation%20matrix%20that%20NBLN's%20covariance%20implies.%20%60scribe.viz.plot_correlation_heatmap%60%20constructs%20the%20implied%20compositional%20correlation%20across%20genes%20from%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5Cperp%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5Cperp%5E%7B%5C!%5Ctop%7D%20%2B%20%5Cmathrm%7Bdiag%7D(%5Cunderline%7Bd%7D)%24%2C%20picks%20the%20most%20variable%20subset%20for%20readability%2C%20and%20hierarchically%20clusters%20them.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20laplace_results%2C%20scribe)%3A%0A%20%20%20%20scribe.viz.plot_correlation_heatmap(%0A%20%20%20%20%20%20%20%20laplace_results%2C%0A%20%20%20%20%20%20%20%20counts%3Dadata%2C%0A%20%20%20%20%20%20%20%20n_genes%3D1000%2C%0A%20%20%20%20%20%20%20%20figsize%3D(7%2C%207)%2C%0A%20%20%20%20%20%20%20%20n_clusters%3D4%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20clustered%20heatmap%20reveals%20block%20structure%3A%20groups%20of%20co-expressed%20genes.%20Because%20we%20routed%20everything%20through%20the%20gauge-killed%2C%20sparsity-regularized%20NBLN%20model%2C%20these%20blocks%20correspond%20to%20**gene%20programs%20the%20data%20actually%20supports**%20rather%20than%20to%20library-size%20or%20noise%20modes.%0A%0A%20%20%20%20%23%23%23%20One%20more%20refinement%3A%20subtract%20the%20library-size%20direction%20explicitly%0A%0A%20%20%20%20Even%20with%20the%20cascade%20freeze%2C%20there%20is%20one%20more%20interpretive%20layer%20worth%20knowing%20about.%20The%20compositional%20projection%20(%24%5Cunderline%7B%5Cunderline%7BW%7D%7D_%5Cperp%20%3D%20(I%20-%20%5Cfrac%7B1%7D%7BG%7D%5Cmathbf%7B1%7D%5Cmathbf%7B1%7D%5E%5Ctop)%5Cunderline%7B%5Cunderline%7BW%7D%7D%24)%20removes%20the%20*uniform*%20all-ones%20component%20%E2%80%94%20%22every%20gene%20goes%20up%20together%22.%20But%20cells%20in%20real%20data%20do%20not%20vary%20*uniformly*%20in%20library%20size%3B%20they%20vary%20along%20a%20direction%20in%20gene%20space%20that%20is%20correlated%20with%20library%20size%20in%20a%20*non-uniform*%20way%20(highly-expressed%20genes%20contribute%20more%20to%20library%20size%20than%20rare%20ones).%20That%20cell-to-cell%20library-size%20axis%20is%20its%20own%20one-dimensional%20component%20of%20variation%20that%20you%20may%20want%20to%20subtract%20before%20interpreting%20the%20rest.%0A%0A%20%20%20%20The%20%60subtract_direction%3D%22library_size%22%60%20argument%20fits%20a%20library-size%20axis%20from%20the%20data%20and%20projects%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%5Cunderline%7B%5Cunderline%7BW%7D%7D%5E%7B%5C!%5Ctop%7D%24%20onto%20its%20orthogonal%20complement%20before%20computing%20correlations.%20What%20remains%20is%20correlation%20structure%20perpendicular%20to%20library%20size%20%E2%80%94%20what%20is%20usually%20meant%20by%20%22real%20co-regulation%22.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(adata%2C%20laplace_results%2C%20scribe)%3A%0A%20%20%20%20scribe.viz.plot_correlation_heatmap(%0A%20%20%20%20%20%20%20%20laplace_results%2C%0A%20%20%20%20%20%20%20%20counts%3Dadata%2C%0A%20%20%20%20%20%20%20%20n_genes%3D1000%2C%0A%20%20%20%20%20%20%20%20n_clusters%3D4%2C%0A%20%20%20%20%20%20%20%20figsize%3D(7%2C%207)%2C%0A%20%20%20%20%20%20%20%20subtract_direction%3D%22library_size%22%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20The%20library-size-projected%20heatmap%20typically%20shows%20**fewer%2C%20sharper%20blocks**%20than%20the%20raw%20one%20%E2%80%94%20the%20dominant%20block%20of%20%22all%20the%20housekeeping%20%2F%20highly-expressed%20genes%20co-vary%20together%22%20is%20partly%20a%20library-size%20echo%2C%20and%20once%20you%20take%20it%20out%20the%20remaining%20structure%20is%20closer%20to%20what%20most%20analysts%20would%20call%20cell-state%20or%20gene-program%20correlation.%0A%0A%20%20%20%20Whether%20you%20should%20look%20at%20the%20raw%20heatmap%20or%20the%20library-size-projected%20one%20depends%20on%20the%20question.%20For%20the%20question%20%22which%20genes%20look%20like%20they%20belong%20to%20the%20same%20regulatory%20module%3F%22%20the%20projected%20version%20is%20usually%20the%20right%20object.%20For%20the%20question%20%22which%20genes%20are%20highly%20expressed%20alongside%20which%20other%20highly-expressed%20genes%3F%22%20the%20raw%20one%20is%20fine.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Why%20this%20is%20the%20principled%20answer%0A%0A%20%20%20%20It%20helps%20to%20compare%20what%20we%20just%20did%20to%20the%20obvious%20alternatives.%0A%0A%20%20%20%20%23%23%23%20Alternative%201%3A%20empirical%20correlation%20of%20normalized%20counts%0A%0A%20%20%20%20The%20Scanpy%2FSeurat%20path%20is%3A%20normalize%20to%20a%20common%20library%20size%20(or%20log-transform%20with%20%24%5Clog(1%2Bx)%24)%2C%20then%20compute%20pairwise%20Pearson%20correlation%20across%20cells.%20This%20is%20fast%20and%20easy.%20But%3A%0A%0A%20%20%20%20-%20The%20normalization%20is%20**not%20a%20model**%3B%20it%20does%20not%20propagate%20uncertainty.%20If%20two%20genes%20have%20correlated%20*counts*%20across%20cells%20because%20both%20have%20high%20library-size%20dependence%20and%20you%20happened%20to%20normalize%20imperfectly%2C%20you%20cannot%20tell%20that%20apart%20from%20a%20real%20biological%20covariance.%0A%20%20%20%20-%20Low-coverage%20cells%20contribute%20as%20much%20weight%20as%20high-coverage%20cells%2C%20even%20though%20they%20carry%20much%20less%20information.%0A%20%20%20%20-%20There%20is%20no%20built-in%20mechanism%20for%20separating%20%22two%20genes%20co-vary%22%20from%20%22two%20genes%20happen%20to%20both%20be%20high%22.%0A%0A%20%20%20%20%23%23%23%20Alternative%202%3A%20low-rank%20variational%20guide%20on%20NBVCP%20(the%20previous%20tutorial)%0A%0A%20%20%20%20With%20%60guide_rank%3D128%60%20and%20%60joint_params%3D%22biological%22%60%2C%20the%20variational%20guide%20carries%20gene-gene%20posterior%20correlations.%20But%3A%0A%0A%20%20%20%20-%20The%20**generative**%20model%20still%20emits%20independent%20gene%20counts.%20Any%20joint%20structure%20is%20a%20property%20of%20the%20posterior%20over%20parameters%2C%20not%20of%20the%20model%20itself.%20You%20can%20read%20the%20heatmap%2C%20but%20the%20model%20does%20not%20actually%20believe%20in%20cross-gene%20covariance%20%E2%80%94%20the%20heatmap%20is%20summarizing%20how%20the%20data%20couples%20parameter%20estimates.%0A%20%20%20%20-%20This%20is%20not%20nothing%20%E2%80%94%20when%20many%20cells%20observe%20gene%20A%20together%20with%20gene%20B%2C%20posterior%20over%20parameters%20does%20pick%20that%20up.%20But%20it%20confounds%20%22co-expression%22%20with%20%22co-information%22%20in%20a%20way%20that%20is%20hard%20to%20disentangle.%0A%0A%20%20%20%20%23%23%23%20What%20NBLN-Laplace%20%2B%20cascade%20%2B%20shrinkage%20gives%20you%0A%0A%20%20%20%20-%20The%20generative%20model%20**has**%20gene-gene%20covariance%20in%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%5Cunderline%7B%5Cunderline%7BW%7D%7D%5E%7B%5C!%5Ctop%7D%24.%20The%20heatmap%20is%20computed%20from%20that%20object%20directly.%20Sampling%20from%20the%20predictive%20distribution%20produces%20cells%20with%20the%20right%20joint%20structure%20(we%20verified%20this%20with%20compositional%20corner%20PPCs).%0A%20%20%20%20-%20The%20cascade%20gives%20us%20a%20calibrated%20starting%20point%20and%20pins%20capture%2C%20so%20the%20answer%20is%20not%20contaminated%20by%20gauge%20degeneracies.%0A%20%20%20%20-%20The%20shrinkage%20prior%20performs%20adaptive%20rank%20selection%2C%20so%20the%20heatmap%20is%20built%20from%20data-supported%20factors%20only%2C%20not%20from%20noise%20the%20model%20fit%20because%20we%20let%20it.%0A%20%20%20%20-%20Every%20step%20is%20checkable.%20We%20can%20compare%20PPCs%20at%20the%20marginal%20level%20(%60plot_ppc%60)%2C%20at%20the%20compositional%201D%20level%20(%60plot_compositional_ppc%60)%2C%20at%20the%20compositional%202D%20level%20(%60plot_compositional_corner_ppc%60)%2C%20and%20at%20the%20gene-pair%20count%20level%20(%60plot_corner_ppc%60).%20Each%20is%20a%20different%20lens%20on%20%22does%20this%20model%20match%20the%20data%3F%22%2C%20and%20we%20passed%20all%20of%20them.%0A%0A%20%20%20%20None%20of%20these%20other%20workflows%20let%20you%20ask%20the%20question%20this%20carefully.%20The%20cost%20is%20a%20more%20involved%20pipeline%3B%20the%20benefit%20is%20that%20the%20resulting%20gene-gene%20correlation%20matrix%20is%20genuinely%20a%20property%20of%20a%20calibrated%20probabilistic%20model%2C%20not%20a%20heuristic%20on%20transformed%20counts.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Recap%0A%0A%20%20%20%20The%20three%20composable%20pieces%20of%20the%20cascade%20workflow%3A%0A%0A%20%20%20%201.%20**SVI%20cascade%20source**%20(%60informative_priors_from%3Dsvi_results%60)%20%E2%80%94%20fit%20NBVCP%20under%20SVI%20with%20a%20tight%20biology-informed%20capture-efficiency%20prior%2C%20then%20carry%20that%20posterior%20forward%20as%20empirical%20Gaussian%20priors%20on%20NBLN's%20dispersion%2C%20mean%20log-rate%2C%20and%20per-cell%20capture.%0A%20%20%20%202.%20**Cascade%20freeze**%20(%60informative_priors_freeze%3D(%22dispersion%22%2C%20%22capture_efficiency%22)%60)%20%E2%80%94%20pin%20the%20cascade-derived%20%24r%24%20and%20%24%5Ceta%24%20as%20constants%20during%20the%20NBLN%20M-step.%20This%20structurally%20kills%20the%20per-cell%20rigid-translation%20gauge%20so%20that%20%24%5Cunderline%7B%5Cunderline%7BW%7D%7D%24%20can%20encode%20cross-gene%20correlations%20without%20being%20polluted%20by%20library-size%20slop.%0A%20%20%20%203.%20**Loadings%20shrinkage**%20(%60priors%3D%7B%22loadings%22%3A%20%7B%22type%22%3A%20%22horseshoe_columnwise%22%2C%20...%7D%7D%60)%20%E2%80%94%20a%20column-wise%20horseshoe%20prior%20on%20the%20loadings%20matrix%20performs%20adaptive%20rank%20selection.%20With%20%60latent_dim%3D32%60%20and%20the%20prior%20active%2C%20only%20data-supported%20factors%20survive%3B%20the%20rest%20are%20shrunk%20away.%0A%0A%20%20%20%20The%20diagnostics%20that%20confirm%20the%20workflow%20is%20doing%20what%20we%20want%3A%0A%0A%20%20%20%20-%20%60plot_loss%60%20%E2%80%94%20training%20converges%20cleanly.%0A%20%20%20%20-%20%60plot_ppc%60%20%E2%80%94%20marginal%20count%20distributions%20match%20across%20the%20abundance%20range.%0A%20%20%20%20-%20%60plot_compositional_ppc%60%20%E2%80%94%20marginal%20compositional%20fractions%20match.%0A%20%20%20%20-%20%60plot_compositional_corner_ppc%60%20%E2%80%94%20*joint*%20compositional%20structure%20matches%3B%20this%20is%20what%20NBVCP%20could%20not%20do.%0A%20%20%20%20-%20%60plot_w_shrinkage_spectrum%60%20%E2%80%94%20the%20spectrum%20of%20compositional%20column%20norms%20has%20a%20clean%20elbow%20at%20%24k_%5Ctext%7Beff%7D%20%5Cll%2032%24.%0A%20%20%20%20-%20%60singular%20spectrum%20of%20W_perp%60%20%E2%80%94%20independent%20verification%20that%20variance%20is%20concentrated%20in%20a%20few%20modes.%0A%20%20%20%20-%20%60get_gauge_diagnostics()%60%20%E2%80%94%20the%20gauge-contamination%20ratio%20is%20benign%20for%20shrinkage-active%20fits%20(a%20ratio%20in%20%24%5B0.5%2C%200.8%5D%24%20is%20fine%20here).%0A%20%20%20%20-%20%60plot_correlation_heatmap%60%20%E2%80%94%20the%20headline%20gene-gene%20correlation%20structure%2C%20optionally%20with%20%60subtract_direction%3D%22library_size%22%60%20for%20additional%20interpretability.%0A%0A%20%20%20%20The%20result%20is%20a%20correlation%20matrix%20on%20genes%20you%20can%20defend%20as%20a%20property%20of%20a%20calibrated%20joint%20generative%20model%2C%20with%20a%20record%20of%20every%20assumption%20that%20went%20into%20it%20and%20a%20posterior%20predictive%20that%20demonstrably%20matches%20the%20data's%20joint%20structure.%0A%0A%20%20%20%20%23%23%23%20Where%20to%20go%20from%20here%0A%0A%20%20%20%20-%20**Differential%20expression.**%20With%20a%20calibrated%20joint%20model%20in%20hand%2C%20posterior%20contrasts%20between%20gene%20pairs%20become%20honest%20probabilistic%20statements.%20%60scribe%60's%20DE%20framework%20(CLR%2FILR%2C%20lfsr%2C%20PEFP)%20is%20the%20next%20stop.%0A%20%20%20%20-%20**Multiple%20cell%20types.**%20Jurkat%20is%20a%20monoculture.%20Heterogeneous%20tissues%20require%20either%20mixture%20extensions%20or%20per-cluster%20fits%3B%20the%20cascade%20machinery%20generalizes%20naturally.%0A%20%20%20%20-%20**Larger%20latent%20dim.**%20For%20deeply%20heterogeneous%20data%2C%20push%20%60latent_dim%60%20higher%20and%20let%20the%20shrinkage%20prior%20pick%20a%20larger%20%24k_%5Ctext%7Beff%7D%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A