covariance whitening interactive

%23%20%2F%2F%2F%20script%0A%23%20requires-python%20%3D%20%22%3E%3D3.11%22%0A%23%20dependencies%20%3D%20%5B%22marimo%3E%3D0.23.9%22%2C%20%22matplotlib%3D%3D3.10.8%22%2C%20%22numpy%3D%3D2.4.3%22%5D%0A%23%20%0A%23%20%5Btool.marimo.export%5D%0A%23%20lock_kind%20%3D%20%22resolved%22%0A%23%20%2F%2F%2F%0A%0Aimport%20marimo%0A%0A__generated_with%20%3D%20%220.25.1%22%0Aapp%20%3D%20marimo.App()%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20numpy%20as%20np%0A%0A%20%20%20%20plt.style.use(%22default%22)%0A%20%20%20%20BLUE%2C%20ORANGE%2C%20RED%2C%20GREY%20%3D%20%22%231f77b4%22%2C%20%22%23ff7f0e%22%2C%20%22%23d62728%22%2C%20%220.55%22%0A%20%20%20%20return%20BLUE%2C%20GREY%2C%20ORANGE%2C%20RED%2C%20mo%2C%20np%2C%20plt%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%20The%20Covariance%20Matrix%20Is%20a%20Shape%0A%0A%20%20%20%20A%20%242%5Ctimes2%24%20covariance%20matrix%0A%0A%20%20%20%20%24%24%5Cmathbf%20C_%7B%5Cmathbf%7BXX%7D%7D%3D%5Cbegin%7Bpmatrix%7D%5Csigma_1%5E2%20%26%20%5Crho%5C%2C%5Csigma_1%5Csigma_2%5C%5C%20%5Crho%5C%2C%5Csigma_1%5Csigma_2%20%26%20%5Csigma_2%5E2%5Cend%7Bpmatrix%7D%24%24%0A%0A%20%20%20%20is%20an%20**ellipse**%3A%20its%20eigenvectors%20are%20the%20axes%2C%20the%20square%20roots%20of%20its%20eigenvalues%20the%20half-lengths.%0A%20%20%20%20Change%20the%20three%20numbers%20below%20and%20watch%20four%20things%20move%20together%3A%0A%0A%20%20%20%201.%20the%20cloud%20of%20samples%20and%20its%20covariance%20ellipses%2C%0A%20%20%20%202.%20the%20**Mahalanobis%20distance**%20of%20a%20test%20point%2C%20against%20its%20coordinate-wise%20%24z%24-scores%2C%0A%20%20%20%203.%20the%20variance%20%24%5Cmathbf%20a%5E%5Cmathsf%7BT%7D%5Cmathbf%20C%5Cmathbf%20a%24%20in%20every%20direction%2C%0A%20%20%20%204.%20the%20cloud%20after%20**whitening**%2C%20by%20PCA%2C%20ZCA%20or%20Cholesky%3A%20three%20different%20routes%20to%20the%20same%20sphere.%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%20s1%20%3D%20mo.ui.slider(0.3%2C%202.0%2C%20step%3D0.05%2C%20value%3D1.2%2C%20label%3D%22%CF%83%E2%82%81%22)%0A%20%20%20%20s2%20%3D%20mo.ui.slider(0.3%2C%202.0%2C%20step%3D0.05%2C%20value%3D0.8%2C%20label%3D%22%CF%83%E2%82%82%22)%0A%20%20%20%20rho%20%3D%20mo.ui.slider(-0.95%2C%200.95%2C%20step%3D0.05%2C%20value%3D0.75%2C%20label%3D%22correlation%20%CF%81%22)%0A%20%20%20%20px%20%3D%20mo.ui.slider(-3.0%2C%203.0%2C%20step%3D0.1%2C%20value%3D-1.0%2C%20label%3D%22test%20point%20x%E2%82%81%22)%0A%20%20%20%20py%20%3D%20mo.ui.slider(-3.0%2C%203.0%2C%20step%3D0.1%2C%20value%3D0.7%2C%20label%3D%22test%20point%20x%E2%82%82%22)%0A%20%20%20%20method%20%3D%20mo.ui.dropdown(%5B%22PCA%22%2C%20%22ZCA%22%2C%20%22Cholesky%22%5D%2C%20value%3D%22PCA%22%2C%20label%3D%22whitener%22)%0A%20%20%20%20mo.vstack(%5Bmo.hstack(%5Bs1%2C%20s2%2C%20rho%5D)%2C%20mo.hstack(%5Bpx%2C%20py%2C%20method%5D)%5D)%0A%20%20%20%20return%20method%2C%20px%2C%20py%2C%20rho%2C%20s1%2C%20s2%0A%0A%0A%40app.cell%0Adef%20_(np%2C%20rho%2C%20s1%2C%20s2)%3A%0A%20%20%20%20C%20%3D%20np.array(%5B%5Bs1.value**2%2C%20rho.value%20*%20s1.value%20*%20s2.value%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5Brho.value%20*%20s1.value%20*%20s2.value%2C%20s2.value**2%5D%5D)%0A%20%20%20%20lam%2C%20U%20%3D%20np.linalg.eigh(C)%0A%20%20%20%20order%20%3D%20np.argsort(lam)%5B%3A%3A-1%5D%0A%20%20%20%20lam%2C%20U%20%3D%20lam%5Border%5D%2C%20U%5B%3A%2C%20order%5D%0A%20%20%20%20for%20_i%20in%20range(2)%3A%0A%20%20%20%20%20%20%20%20if%20U%5Bnp.argmax(np.abs(U%5B%3A%2C%20_i%5D))%2C%20_i%5D%20%3C%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20U%5B%3A%2C%20_i%5D%20%3D%20-U%5B%3A%2C%20_i%5D%0A%20%20%20%20_rng%20%3D%20np.random.default_rng(3)%0A%20%20%20%20X%20%3D%20_rng.standard_normal((1500%2C%202))%20%40%20np.linalg.cholesky(C).T%0A%20%20%20%20th%20%3D%20np.linspace(0%2C%202%20*%20np.pi%2C%20300)%0A%20%20%20%20circ%20%3D%20np.vstack(%5Bnp.cos(th)%2C%20np.sin(th)%5D)%0A%20%20%20%20return%20C%2C%20U%2C%20X%2C%20circ%2C%20lam%0A%0A%0A%40app.cell%0Adef%20_(BLUE%2C%20C%2C%20GREY%2C%20ORANGE%2C%20RED%2C%20U%2C%20X%2C%20circ%2C%20lam%2C%20mo%2C%20np%2C%20plt%2C%20px%2C%20py)%3A%0A%20%20%20%20p%20%3D%20np.array(%5Bpx.value%2C%20py.value%5D)%0A%20%20%20%20dM%20%3D%20float(np.sqrt(p%20%40%20np.linalg.solve(C%2C%20p)))%0A%20%20%20%20z%20%3D%20np.abs(p)%20%2F%20np.sqrt(np.diag(C))%0A%0A%20%20%20%20fig1%2C%20(a1%2C%20a2)%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(12.5%2C%205.0))%0A%20%20%20%20a1.scatter(X%5B%3A%2C%200%5D%2C%20X%5B%3A%2C%201%5D%2C%20s%3D4%2C%20color%3DBLUE%2C%20alpha%3D0.25)%0A%20%20%20%20for%20_r%2C%20_lw%20in%20%5B(1%2C%202.0)%2C%20(2%2C%201.4)%2C%20(3%2C%200.9)%5D%3A%0A%20%20%20%20%20%20%20%20_E%20%3D%20_r%20*%20(U%20%40%20np.diag(np.sqrt(lam))%20%40%20circ)%0A%20%20%20%20%20%20%20%20a1.plot(_E%5B0%5D%2C%20_E%5B1%5D%2C%20color%3DORANGE%2C%20lw%3D_lw)%0A%20%20%20%20for%20_i%20in%20range(2)%3A%0A%20%20%20%20%20%20%20%20_v%20%3D%20U%5B%3A%2C%20_i%5D%20*%20np.sqrt(lam%5B_i%5D)%0A%20%20%20%20%20%20%20%20a1.annotate(%22%22%2C%20xy%3D_v%2C%20xytext%3D(0%2C%200)%2C%20arrowprops%3Ddict(arrowstyle%3D%22-%3E%22%2C%20color%3DRED%2C%20lw%3D2))%0A%20%20%20%20_sd%20%3D%20np.sqrt(np.diag(C))%0A%20%20%20%20a1.add_patch(plt.Rectangle((-_sd%5B0%5D%2C%20-_sd%5B1%5D)%2C%202%20*%20_sd%5B0%5D%2C%202%20*%20_sd%5B1%5D%2C%20fill%3DFalse%2C%20ec%3DGREY%2C%20ls%3D%22--%22%2C%20lw%3D1.4))%0A%20%20%20%20a1.plot(*p%2C%20%22o%22%2C%20color%3DRED%2C%20ms%3D10)%0A%20%20%20%20a1.set_title(f%22test%20point%3A%20Mahalanobis%20d%20%3D%20%7BdM%3A.2f%7D%20%20%20(z-scores%20%7Bz%5B0%5D%3A.2f%7D%2C%20%7Bz%5B1%5D%3A.2f%7D)%22%2C%20fontsize%3D12)%0A%20%20%20%20a1.set_xlim(-5%2C%205)%3B%20a1.set_ylim(-5%2C%205)%3B%20a1.set_aspect(%22equal%22)%0A%20%20%20%20a1.set_xlabel(%22%24x_1%24%22)%3B%20a1.set_ylabel(%22%24x_2%24%22)%0A%0A%20%20%20%20_ang%20%3D%20np.linspace(0%2C%20np.pi%2C%20400)%0A%20%20%20%20_a%20%3D%20np.vstack(%5Bnp.cos(_ang)%2C%20np.sin(_ang)%5D)%0A%20%20%20%20_var%20%3D%20np.einsum(%22in%2Cij%2Cjn-%3En%22%2C%20_a%2C%20C%2C%20_a)%0A%20%20%20%20a2.plot(np.degrees(_ang)%2C%20_var%2C%20color%3DBLUE%2C%20lw%3D2.4)%0A%20%20%20%20a2.axhline(lam%5B0%5D%2C%20color%3DORANGE%2C%20ls%3D%22--%22)%3B%20a2.axhline(lam%5B1%5D%2C%20color%3DORANGE%2C%20ls%3D%22--%22)%0A%20%20%20%20a2.text(2%2C%20lam%5B0%5D%20*%201.03%2C%20r%22%24%5Clambda_%7B%5Cmax%7D%24%22%2C%20color%3DORANGE%2C%20fontsize%3D13)%0A%20%20%20%20a2.text(2%2C%20lam%5B1%5D%20*%201.03%20%2B%200.02%2C%20r%22%24%5Clambda_%7B%5Cmin%7D%24%22%2C%20color%3DORANGE%2C%20fontsize%3D13)%0A%20%20%20%20a2.set_xlabel(r%22direction%20of%20%24%5Cmathbf%7Ba%7D%24%20(degrees)%22)%3B%20a2.set_ylabel(r%22%24%5Cmathbf%7Ba%7D%5E%5Cmathsf%7BT%7D%5Cmathbf%7BC%7D%5Cmathbf%7Ba%7D%24%22)%0A%20%20%20%20a2.set_title(f%22variance%20in%20every%20direction%3B%20eigenvalue%20spread%20%7Blam%5B0%5D%20%2F%20lam%5B1%5D%3A.1f%7D%22%2C%20fontsize%3D12)%0A%20%20%20%20fig1.tight_layout()%0A%0A%20%20%20%20mo.vstack(%5B%0A%20%20%20%20%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20%20%20%20%20%22The%20dashed%20box%20marks%20one%20standard%20deviation%20in%20each%20coordinate.%20A%20point%20inside%20the%20box%20can%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%22still%20be%20far%20from%20the%20data%20in%20Mahalanobis%20distance%2C%20if%20it%20disobeys%20the%20correlation.%22%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20fig1%2C%0A%20%20%20%20%5D)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20ORANGE%2C%20U%2C%20X%2C%20circ%2C%20lam%2C%20method%2C%20mo%2C%20np%2C%20plt)%3A%0A%20%20%20%20if%20method.value%20%3D%3D%20%22PCA%22%3A%0A%20%20%20%20%20%20%20%20A%20%3D%20np.diag(lam%20**%20-0.5)%20%40%20U.T%0A%20%20%20%20%20%20%20%20note%20%3D%20%22PCA%3A%20rotate%20onto%20the%20eigenvectors%2C%20then%20divide%20each%20axis%20by%20its%20standard%20deviation.%20The%20output%20coordinates%20are%20the%20principal%20components%2C%20ordered%20by%20variance.%22%0A%20%20%20%20elif%20method.value%20%3D%3D%20%22ZCA%22%3A%0A%20%20%20%20%20%20%20%20A%20%3D%20U%20%40%20np.diag(lam%20**%20-0.5)%20%40%20U.T%0A%20%20%20%20%20%20%20%20note%20%3D%20%22ZCA%3A%20the%20same%2C%20then%20rotate%20back.%20Of%20all%20whiteners%20it%20changes%20the%20data%20least%2C%20so%20each%20output%20coordinate%20still%20resembles%20the%20input%20coordinate%20of%20the%20same%20index.%22%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20_L%20%3D%20np.linalg.cholesky(C)%0A%20%20%20%20%20%20%20%20A%20%3D%20np.linalg.inv(_L)%0A%20%20%20%20%20%20%20%20note%20%3D%20%22Cholesky%3A%20a%20lower-triangular%20map%2C%20so%20output%201%20uses%20only%20input%201%20and%20output%202%20uses%20inputs%201%20and%202.%20Causal%20whitening%3A%20the%20innovations%20of%20the%20ordered%20coordinates.%22%0A%20%20%20%20Wd%20%3D%20X%20%40%20A.T%0A%20%20%20%20_col%20%3D%20(X%5B%3A%2C%200%5D%20-%20X%5B%3A%2C%200%5D.min())%20%2F%20np.ptp(X%5B%3A%2C%200%5D)%0A%0A%20%20%20%20fig2%2C%20(b1%2C%20b2)%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(12.5%2C%205.0))%0A%20%20%20%20b1.scatter(X%5B%3A%2C%200%5D%2C%20X%5B%3A%2C%201%5D%2C%20s%3D5%2C%20c%3D_col%2C%20cmap%3D%22viridis%22%2C%20alpha%3D0.6)%0A%20%20%20%20_E%20%3D%202%20*%20(U%20%40%20np.diag(np.sqrt(lam))%20%40%20circ)%0A%20%20%20%20b1.plot(_E%5B0%5D%2C%20_E%5B1%5D%2C%20color%3DORANGE%2C%20lw%3D2)%0A%20%20%20%20b1.set_title(%22before%3A%20covariance%20C%22%2C%20fontsize%3D12)%0A%20%20%20%20b2.scatter(Wd%5B%3A%2C%200%5D%2C%20Wd%5B%3A%2C%201%5D%2C%20s%3D5%2C%20c%3D_col%2C%20cmap%3D%22viridis%22%2C%20alpha%3D0.6)%0A%20%20%20%20b2.plot(2%20*%20circ%5B0%5D%2C%202%20*%20circ%5B1%5D%2C%20color%3DORANGE%2C%20lw%3D2)%0A%20%20%20%20b2.set_title(f%22after%20%7Bmethod.value%7D%3A%20sample%20covariance%20%7Bnp.round(np.cov(Wd%2C%20rowvar%3DFalse)%2C%202).tolist()%7D%22%2C%20fontsize%3D11)%0A%20%20%20%20for%20_ax%20in%20(b1%2C%20b2)%3A%0A%20%20%20%20%20%20%20%20_ax.set_xlim(-5%2C%205)%3B%20_ax.set_ylim(-5%2C%205)%3B%20_ax.set_aspect(%22equal%22)%0A%20%20%20%20fig2.tight_layout()%0A%20%20%20%20mo.vstack(%5Bmo.md(f%22**%7Bnote%7D**%20The%20colors%20label%20each%20point%20by%20its%20original%20%24x_1%24%2C%20so%20you%20can%20follow%20where%20each%20point%20goes.%22)%2C%20fig2%5D)%0A%20%20%20%20return%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%23%20What%20to%20try%0A%0A%20%20%20%20-%20Set%20%24%5Crho%3D0%24%20and%20%24%5Csigma_1%3D%5Csigma_2%24%3A%20the%20ellipse%20is%20a%20circle%2C%20all%20directions%20have%20the%20same%20variance%2C%20and%20every%20whitener%20is%20(almost)%20the%20identity.%0A%20%20%20%20-%20Push%20%24%5Crho%24%20toward%20%24%5Cpm0.95%24%3A%20the%20eigenvalue%20spread%20explodes.%20A%20point%20just%20off%20the%20long%20axis%20gets%20a%20huge%20Mahalanobis%20distance%20while%20its%20%24z%24-scores%20stay%20small.%0A%20%20%20%20-%20Switch%20the%20whitener%3A%20the%20whitened%20cloud%20is%20always%20a%20sphere%2C%20but%20the%20colors%20land%20in%20different%20places.%20Whitening%20is%20unique%20only%20up%20to%20a%20rotation.%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