latent viz

%23%20%2F%2F%2F%20script%0A%23%20requires-python%20%3D%20%22%3E%3D3.13%22%0A%23%20dependencies%20%3D%20%5B%22marimo%22%2C%20%22numpy%3D%3D2.4.3%22%2C%20%22matplotlib%3D%3D3.10.8%22%5D%0A%23%20%0A%23%20%5Btool.marimo.export%5D%0A%23%20lock_kind%20%3D%20%22observed%22%0A%23%20%2F%2F%2F%0A%0Aimport%20marimo%0A%0A__generated_with%20%3D%20%220.25.2%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Latent%20variable%20models%3A%20simple%20conditionals%20make%20a%20funky%20marginal%0A%0A%20%20%20%20We%20posit%20%24p(x%2C%20z)%20%3D%20p(x%20%5Cmid%20z)%5C%2Cp(z)%24%20with%20latent%20%24z%24.%0A%0A%20%20%20%20-%20%24p(z)%20%3D%20%5Cmathcal%7BN%7D(0%2C%201)%24%0A%20%20%20%20-%20%24p(x%20%5Cmid%20z)%20%3D%20%5Cmathcal%7BN%7D%5Cbig(%5Cmu(z)%2C%5C%2C%20%5Csigma%5E2%5Cbig)%24%20with%20%24%5Cmu(z)%20%3D%20A%20%5Csin(%5Comega%20z)%24%20(a%20simple%20Gaussian%20per%20%24z%24)%0A%20%20%20%20-%20marginal%20%24p(x)%20%3D%20%5Cint%20p(x%20%5Cmid%20z)%5C%2Cp(z)%5C%2Cdz%24%20%E2%80%94%20multi-modal%20even%20though%20every%20conditional%20is%20unimodal.%0A%0A%20%20%20%20Below%20are%20two%20different%20latent-variable%20stories%20with%20the%20same%20observed%20marginal%20%24p(x)%24.%0A%0A%20%20%20%20**Model%201%3A**%20latent%20%24z%20%5Csim%20%5Cmathcal%7BN%7D(0%2C1)%24%2C%20then%20%24x%20%5Cmid%20z%20%5Csim%20%5Cmathcal%7BN%7D(A%5Csin(%5Comega%20z)%2C%20%5Csigma%5E2)%24.%0A%0A%20%20%20%20**Model%202%3A**%20latent%20%24m%20%3D%20A%5Csin(%5Comega%20z)%24%20directly%2C%20with%20the%20induced%20non-Gaussian%20prior%20%24p(m)%24%2C%20then%20%24x%20%5Cmid%20m%20%5Csim%20%5Cmathcal%7BN%7D(m%2C%20%5Csigma%5E2)%24.%0A%0A%20%20%20%20The%20ridgelines%20show%20that%20these%20are%20structurally%20different%20latent%20spaces%2C%20yet%20they%20integrate%20to%20the%20same%20distribution%20over%20observations.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20sigma%20%3D%20mo.ui.slider(0.05%2C%201.5%2C%20value%3D0.35%2C%20step%3D0.05%2C%20label%3D%22conditional%20width%20%20%CF%83%22)%0A%20%20%20%20amp%20%3D%20mo.ui.slider(0.5%2C%204.0%2C%20value%3D2.5%2C%20step%3D0.1%2C%20label%3D%22amplitude%20%20A%22)%0A%20%20%20%20freq%20%3D%20mo.ui.slider(0.3%2C%203.0%2C%20value%3D1.4%2C%20step%3D0.1%2C%20label%3D%22frequency%20%20%CF%89%22)%0A%20%20%20%20n_lines%20%3D%20mo.ui.slider(5%2C%20200%2C%20value%3D60%2C%20step%3D5%2C%20label%3D%22number%20of%20p(x%7Cz)%20curves%22)%0A%20%20%20%20mo.vstack(%5Bsigma%2C%20amp%2C%20freq%2C%20n_lines%5D)%0A%20%20%20%20return%20amp%2C%20freq%2C%20n_lines%2C%20sigma%0A%0A%0A%40app.cell%0Adef%20_(amp%2C%20freq%2C%20n_lines%2C%20np%2C%20plt%2C%20sigma)%3A%0A%20%20%20%20def%20normal_pdf(x%2C%20mu%2C%20s)%3A%0A%20%20%20%20%20%20%20%20return%20np.exp(-0.5%20*%20((x%20-%20mu)%20%2F%20s)%20**%202)%20%2F%20(s%20*%20np.sqrt(2%20*%20np.pi))%0A%0A%20%20%20%20x%20%3D%20np.linspace(-6%2C%206%2C%20800)%0A%0A%20%20%20%20%23%20fine%20grid%20over%20z%20for%20an%20accurate%20marginal%20p(x)%20%3D%20%E2%88%AB%20p(x%7Cz)%20p(z)%20dz%0A%20%20%20%20z_fine%20%3D%20np.linspace(-4%2C%204%2C%202000)%0A%20%20%20%20pz_fine%20%3D%20normal_pdf(z_fine%2C%200.0%2C%201.0)%0A%20%20%20%20mu_fine%20%3D%20amp.value%20*%20np.sin(freq.value%20*%20z_fine)%0A%20%20%20%20cond_fine%20%3D%20normal_pdf(x%5B%3A%2C%20None%5D%2C%20mu_fine%5BNone%2C%20%3A%5D%2C%20sigma.value)%20%20%23%20(x%2C%20z)%0A%20%20%20%20marginal%20%3D%20cond_fine%20%40%20pz_fine%20*%20(z_fine%5B1%5D%20-%20z_fine%5B0%5D)%0A%0A%20%20%20%20%23%20Model%201%20ridgelines%3A%20a%20subset%20of%20z%20values%2C%20weighted%20by%20p(z)%0A%20%20%20%20z_lines%20%3D%20np.linspace(-3%2C%203%2C%20n_lines.value)%0A%20%20%20%20pz_lines%20%3D%20normal_pdf(z_lines%2C%200.0%2C%201.0)%0A%20%20%20%20alpha_z%20%3D%20pz_lines%20%2F%20pz_lines.max()%0A%0A%20%20%20%20%23%20Model%202%20prior%20over%20m%20is%20the%20pushforward%20of%20z%20-%3E%20A%20sin(%CF%89z).%20%20Estimate%20it%0A%20%20%20%20%23%20only%20for%20visualization%3B%20the%20equality%20of%20the%20marginal%20follows%20by%20the%0A%20%20%20%20%23%20change%20of%20latent%20variable%20m%20%3D%20%CE%BC(z).%0A%20%20%20%20m_lines%20%3D%20np.linspace(-amp.value%2C%20amp.value%2C%20n_lines.value)%0A%20%20%20%20dz%20%3D%20z_fine%5B1%5D%20-%20z_fine%5B0%5D%0A%20%20%20%20kde_bw%20%3D%20max(0.05%2C%200.04%20*%20amp.value)%0A%20%20%20%20pm_lines%20%3D%20normal_pdf(m_lines%5B%3A%2C%20None%5D%2C%20mu_fine%5BNone%2C%20%3A%5D%2C%20kde_bw)%20%40%20pz_fine%20*%20dz%0A%20%20%20%20alpha_m%20%3D%20pm_lines%20%2F%20pm_lines.max()%0A%0A%20%20%20%20fig%20%3D%20plt.figure(figsize%3D(9.5%2C%207.2)%2C%20constrained_layout%3DTrue)%0A%20%20%20%20gs%20%3D%20fig.add_gridspec(2%2C%202%2C%20height_ratios%3D%5B1%2C%202.4%5D)%0A%20%20%20%20ax_top%20%3D%20fig.add_subplot(gs%5B0%2C%20%3A%5D)%0A%20%20%20%20ax_z%20%3D%20fig.add_subplot(gs%5B1%2C%200%5D%2C%20sharex%3Dax_top)%0A%20%20%20%20ax_m%20%3D%20fig.add_subplot(gs%5B1%2C%201%5D%2C%20sharex%3Dax_top)%0A%0A%20%20%20%20ax_top.fill_between(x%2C%20marginal%2C%20color%3D%22%233b4cc0%22%2C%20alpha%3D0.85)%0A%20%20%20%20ax_top.plot(x%2C%20marginal%2C%20color%3D%22%231a2580%22%2C%20lw%3D1.5%2C%20label%3D%22both%20models%22)%0A%20%20%20%20ax_top.set_title(%22same%20observed%20marginal%20%20p(x)%22)%0A%20%20%20%20ax_top.set_ylabel(%22density%22)%0A%20%20%20%20ax_top.set_ylim(bottom%3D0)%0A%20%20%20%20ax_top.legend(frameon%3DFalse)%0A%0A%20%20%20%20ridge_scale_z%20%3D%200.55%20%2F%20normal_pdf(0%2C%200%2C%20sigma.value)%0A%20%20%20%20for%20zi%2C%20ai%20in%20zip(z_lines%2C%20alpha_z)%3A%0A%20%20%20%20%20%20%20%20mu%20%3D%20amp.value%20*%20np.sin(freq.value%20*%20zi)%0A%20%20%20%20%20%20%20%20y%20%3D%20zi%20%2B%20ridge_scale_z%20*%20normal_pdf(x%2C%20mu%2C%20sigma.value)%0A%20%20%20%20%20%20%20%20ax_z.plot(x%2C%20y%2C%20color%3D%22%23b40426%22%2C%20alpha%3D0.12%20%2B%200.88%20*%20float(ai)%2C%20lw%3D1.1)%0A%20%20%20%20%20%20%20%20ax_z.fill_between(x%2C%20zi%2C%20y%2C%20color%3D%22%23b40426%22%2C%20alpha%3D0.04%20%2B%200.18%20*%20float(ai)%2C%20lw%3D0)%0A%20%20%20%20ax_z.plot(amp.value%20*%20np.sin(freq.value%20*%20z_fine)%2C%20z_fine%2C%20color%3D%22%23222222%22%2C%20lw%3D1.0%2C%20alpha%3D0.65)%0A%20%20%20%20ax_z.set_title(%22Model%201%3A%20latent%20z%20~%20N(0%2C1)%22)%0A%20%20%20%20ax_z.set_xlabel(%22x%22)%0A%20%20%20%20ax_z.set_ylabel(%22z%22)%0A%20%20%20%20ax_z.set_ylim(z_lines.min()%20-%200.2%2C%20z_lines.max()%20%2B%200.75)%0A%0A%20%20%20%20ridge_scale_m%20%3D%200.16%20*%20amp.value%20%2F%20normal_pdf(0%2C%200%2C%20sigma.value)%0A%20%20%20%20for%20mi%2C%20ai%20in%20zip(m_lines%2C%20alpha_m)%3A%0A%20%20%20%20%20%20%20%20y%20%3D%20mi%20%2B%20ridge_scale_m%20*%20normal_pdf(x%2C%20mi%2C%20sigma.value)%0A%20%20%20%20%20%20%20%20ax_m.plot(x%2C%20y%2C%20color%3D%22%23f46d43%22%2C%20alpha%3D0.12%20%2B%200.88%20*%20float(ai)%2C%20lw%3D1.1)%0A%20%20%20%20%20%20%20%20ax_m.fill_between(x%2C%20mi%2C%20y%2C%20color%3D%22%23f46d43%22%2C%20alpha%3D0.04%20%2B%200.18%20*%20float(ai)%2C%20lw%3D0)%0A%20%20%20%20ax_m.plot(m_lines%2C%20m_lines%2C%20color%3D%22%23222222%22%2C%20lw%3D1.0%2C%20alpha%3D0.65)%0A%20%20%20%20ax_m.set_title(%22Model%202%3A%20latent%20m%20%3D%20A%20sin(%CF%89z)%22)%0A%20%20%20%20ax_m.set_xlabel(%22x%22)%0A%20%20%20%20ax_m.set_ylabel(%22m%22)%0A%20%20%20%20ax_m.set_ylim(-amp.value%20-%200.08%20*%20amp.value%2C%20amp.value%20%2B%200.25%20*%20amp.value)%0A%0A%20%20%20%20fig%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%0A%20%20%20%20return%20mo%2C%20np%2C%20plt%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A