%23%20%2F%2F%2F%20script%0A%23%20dependencies%20%3D%20%5B%0A%23%20%20%20%20%20%22marimo%3E%3D0.25.1%22%2C%0A%23%20%20%20%20%20%22matplotlib%3D%3D3.11.2%22%2C%0A%23%20%20%20%20%20%22numpy%3D%3D2.5.4%22%2C%0A%23%20%20%20%20%20%22scipy%3D%3D1.18.1%22%2C%0A%23%20%5D%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(width%3D%22medium%22%2C%20app_title%3D%22Computational%20Optimal%20Transport%22)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20itertools%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%20%20%20%20from%20scipy.linalg%20import%20sqrtm%0A%20%20%20%20from%20scipy.optimize%20import%20linear_sum_assignment%2C%20linprog%0A%20%20%20%20from%20scipy.stats%20import%20norm%0A%0A%20%20%20%20return%20itertools%2C%20linear_sum_assignment%2C%20linprog%2C%20mo%2C%20norm%2C%20np%2C%20plt%2C%20sqrtm%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%20Computational%20Optimal%20Transport%0A%20%20%20%20**Universidad%20Nacional%20de%20Colombia**%20%E2%80%94%20Daniel%20Melo%0A%0A%20%20%20%20%3Csmall%3EG.%20Peyr%C3%A9%20y%20M.%20Cuturi%2C%20*Computational%20Optimal%20Transport*%2C%20Foundations%20and%20Trends%20in%20Machine%20Learning%2C%202019.%3C%2Fsmall%3E%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(linprog%2C%20np)%3A%0A%20%20%20%20def%20marginal_matrix(n%2C%20m)%3A%0A%20%20%20%20%20%20%20%20return%20np.vstack(%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20np.kron(np.ones((1%2C%20m))%2C%20np.eye(n))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20np.kron(np.eye(m)%2C%20np.ones((1%2C%20n)))%2C%0A%20%20%20%20%20%20%20%20%5D)%0A%0A%20%20%20%20def%20kantorovich(a%2C%20b%2C%20C)%3A%0A%20%20%20%20%20%20%20%20n%2C%20m%20%3D%20C.shape%0A%20%20%20%20%20%20%20%20res%20%3D%20linprog(%0A%20%20%20%20%20%20%20%20%20%20%20%20C.flatten(order%3D%22F%22)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20A_eq%3Dmarginal_matrix(n%2C%20m)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20b_eq%3Dnp.concatenate(%5Ba%2C%20b%5D)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20bounds%3D(0%2C%20None)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20method%3D%22highs-ds%22%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20return%20res.x.reshape((n%2C%20m)%2C%20order%3D%22F%22)%2C%20res.fun%2C%20res%0A%0A%20%20%20%20def%20kantorovich_dual(a%2C%20b%2C%20C)%3A%0A%20%20%20%20%20%20%20%20n%2C%20m%20%3D%20C.shape%0A%20%20%20%20%20%20%20%20res%20%3D%20linprog(%0A%20%20%20%20%20%20%20%20%20%20%20%20-np.concatenate(%5Ba%2C%20b%5D)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20A_ub%3Dmarginal_matrix(n%2C%20m).T%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20b_ub%3DC.flatten(order%3D%22F%22)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20bounds%3D(None%2C%20None)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20method%3D%22highs%22%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20return%20res.x%5B%3An%5D%2C%20res.x%5Bn%3A%5D%2C%20-res.fun%0A%0A%20%20%20%20def%20pairwise_dist(X%2C%20Y)%3A%0A%20%20%20%20%20%20%20%20return%20np.linalg.norm(X%5B%3A%2C%20None%2C%20%3A%5D%20-%20Y%5BNone%2C%20%3A%2C%20%3A%5D%2C%20axis%3D-1)%0A%0A%20%20%20%20return%20kantorovich%2C%20kantorovich_dual%2C%20marginal_matrix%2C%20pairwise_dist%0A%0A%0A%40app.cell%0Adef%20_(np)%3A%0A%20%20%20%20def%20fmt_vec(v%2C%20fmt%3D%22%7B%3Ag%7D%22)%3A%0A%20%20%20%20%20%20%20%20return%20r%22%5Cleft(%22%20%2B%20r%22%2C%5C%20%22.join(fmt.format(t)%20for%20t%20in%20np.ravel(v))%20%2B%20r%22%5Cright)%22%0A%0A%20%20%20%20def%20sci(v%2C%20d%3D2)%3A%0A%20%20%20%20%20%20%20%20if%20v%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%220%22%0A%20%20%20%20%20%20%20%20mant%2C%20exp%20%3D%20f%22%7Bv%3A.%7Bd%7De%7D%22.split(%22e%22)%0A%20%20%20%20%20%20%20%20return%20rf%22%7Bmant%7D%5Ctimes%2010%5E%7B%7B%7Bint(exp)%7D%7D%7D%22%0A%0A%20%20%20%20def%20bmat(M%2C%20fmt%3D%22%7B%3A.4f%7D%22)%3A%0A%20%20%20%20%20%20%20%20M%20%3D%20np.atleast_2d(M)%0A%20%20%20%20%20%20%20%20tol%20%3D%201e-9%0A%0A%20%20%20%20%20%20%20%20def%20cell(x)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%220%22%20if%20abs(x)%20%3C%20tol%20else%20fmt.format(x)%0A%0A%20%20%20%20%20%20%20%20body%20%3D%20r%22%20%5C%5C%20%22.join(%22%20%26%20%22.join(cell(x)%20for%20x%20in%20row)%20for%20row%20in%20M)%0A%20%20%20%20%20%20%20%20return%20rf%22%5Cbegin%7B%7Bbmatrix%7D%7D%20%7Bbody%7D%20%5Cend%7B%7Bbmatrix%7D%7D%22%0A%0A%20%20%20%20def%20md_table(rows%2C%20header%3D(%22%22%2C%20%22Value%22))%3A%0A%20%20%20%20%20%20%20%20head%20%3D%20%22%7C%20%22%20%2B%20%22%20%7C%20%22.join(header)%20%2B%20%22%20%7C%22%0A%20%20%20%20%20%20%20%20sep%20%3D%20%22%7C%22%20%2B%20%22%7C%22.join(%5B%22%3A---%22%5D%20%2B%20%5B%22%3A---%3A%22%5D%20*%20(len(header)%20-%201))%20%2B%20%22%7C%22%0A%20%20%20%20%20%20%20%20body%20%3D%20%22%5Cn%22.join(%22%7C%20%22%20%2B%20%22%20%7C%20%22.join(r)%20%2B%20%22%20%7C%22%20for%20r%20in%20rows)%0A%20%20%20%20%20%20%20%20return%20%22%5Cn%22.join(%5Bhead%2C%20sep%2C%20body%5D)%0A%0A%20%20%20%20return%20bmat%2C%20fmt_vec%2C%20md_table%2C%20sci%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%201.%20Histograms%20and%20Measures%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%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Two%20ways%20to%20describe%20mass%0A%0A%20%20%20%20**Histograms.**%20Mass%20as%20a%20vector%20of%20weights%3A%0A%20%20%20%20%24%24%0A%20%20%20%20%5CSigma_n%20%5Ccoloneqq%20%5CBig%5C%7B%20%5Cmathbf%7Ba%7D%5Cin%5Cmathbb%7BR%7D%5En_%2B%20%5C%3B%3A%5C%3B%20%5Ctextstyle%5Csum_%7Bi%3D1%7D%5En%20a_i%20%3D%201%5CBig%5C%7D.%0A%20%20%20%20%24%24%0A%0A%20%20%20%20**Discrete%20measures.**%20Weights%20placed%20at%20positions%3A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Calpha%20%3D%20%5Csum_%7Bi%3D1%7D%5En%20a_i%5C%2C%5Cdelta_%7Bx_i%7D%2C%20%5Cqquad%20x_i%5Cin%5Cmathcal%7BX%7D.%0A%20%20%20%20%24%24%0A%0A%20%20%20%20**General%20measures.**%20%24%5Calpha%5Cin%5Cmathcal%7BM%7D_%2B%5E1(%5Cmathcal%7BX%7D)%24%2C%20a%20Radon%20measure%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cint_%7B%5Cmathcal%7BX%7D%7D%20f%5C%2C%5Cmathrm%7Bd%7D%5Calpha%20%5C%3B%5CBig(%3D%20%5Csum_%7Bi%7D%20a_i%20f(x_i)%20%5Ctext%7B%20if%20%24%5Calpha%24%20is%20discrete%7D%5CBig).%0A%20%20%20%20%24%24%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%20seed%20%3D%20mo.ui.slider(0%2C%20100%2C%20value%3D0%2C%20label%3D%22seed%22)%0A%20%20%20%20n_pts%20%3D%20mo.ui.slider(2%2C%2015%2C%20value%3D5%2C%20label%3D%22n%22)%0A%20%20%20%20m_pts%20%3D%20mo.ui.slider(2%2C%2015%2C%20value%3D6%2C%20label%3D%22m%22)%0A%20%20%20%20mo.center(mo.hstack(%5Bseed%2C%20n_pts%2C%20m_pts%5D))%0A%20%20%20%20return%20m_pts%2C%20n_pts%2C%20seed%0A%0A%0A%40app.cell%0Adef%20_(m_pts%2C%20n_pts%2C%20np%2C%20seed)%3A%0A%20%20%20%20_rng%20%3D%20np.random.default_rng(seed.value)%0A%20%20%20%20n%2C%20m%20%3D%20n_pts.value%2C%20m_pts.value%0A%20%20%20%20a%20%3D%20_rng.random(n)%0A%20%20%20%20a%20%2F%3D%20a.sum()%0A%20%20%20%20b%20%3D%20_rng.random(m)%0A%20%20%20%20b%20%2F%3D%20b.sum()%0A%20%20%20%20X%20%3D%20_rng.random((n%2C%202))%20*%20%5B0.45%2C%201.0%5D%0A%20%20%20%20Y%20%3D%20_rng.random((m%2C%202))%20*%20%5B0.45%2C%201.0%5D%20%2B%20%5B0.55%2C%200.0%5D%0A%20%20%20%20return%20X%2C%20Y%2C%20a%2C%20b%2C%20m%2C%20n%0A%0A%0A%40app.cell%0Adef%20_(X%2C%20Y%2C%20a%2C%20b%2C%20md_table%2C%20mo%2C%20np%2C%20plt)%3A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(10%2C%204))%0A%20%20%20%20_ax%5B0%5D.scatter(*X.T%2C%20s%3D3000%20*%20a%2C%20color%3D%22%231f4e79%22%2C%20alpha%3D0.8%2C%20label%3Dr%22%24%5Calpha%24%22)%0A%20%20%20%20_ax%5B0%5D.scatter(*Y.T%2C%20s%3D3000%20*%20b%2C%20color%3D%22%23d66e1e%22%2C%20alpha%3D0.85%2C%20label%3Dr%22%24%5Cbeta%24%22)%0A%20%20%20%20_ax%5B0%5D.set_title(r%22%24d%3D2%24%22)%0A%20%20%20%20_ax%5B0%5D.set_aspect(%22equal%22)%0A%20%20%20%20_ax%5B0%5D.legend()%0A%20%20%20%20_ax%5B1%5D.vlines(X%5B%3A%2C%200%5D%2C%200%2C%20a%2C%20color%3D%22%231f4e79%22%2C%20lw%3D2)%0A%20%20%20%20_ax%5B1%5D.plot(X%5B%3A%2C%200%5D%2C%20a%2C%20%22o%22%2C%20color%3D%22%231f4e79%22%2C%20label%3Dr%22%24%5Calpha%24%22)%0A%20%20%20%20_ax%5B1%5D.vlines(Y%5B%3A%2C%200%5D%2C%200%2C%20b%2C%20color%3D%22%23d66e1e%22%2C%20lw%3D2)%0A%20%20%20%20_ax%5B1%5D.plot(Y%5B%3A%2C%200%5D%2C%20b%2C%20%22o%22%2C%20color%3D%22%23d66e1e%22%2C%20label%3Dr%22%24%5Cbeta%24%22)%0A%20%20%20%20_ax%5B1%5D.set_title(r%22%24d%3D1%24%22)%0A%20%20%20%20_ax%5B1%5D.legend()%0A%20%20%20%20_fig.tight_layout()%0A%0A%20%20%20%20_f%20%3D%20lambda%20x%3A%20np.sum(x**2%2C%20axis%3D-1)%0A%20%20%20%20_exact%20%3D%20np.sum(a%20*%20_f(X))%0A%20%20%20%20_rng%20%3D%20np.random.default_rng(0)%0A%20%20%20%20_idx%20%3D%20_rng.choice(len(a)%2C%20size%3D100_000%2C%20p%3Da)%0A%20%20%20%20_mc%20%3D%20_f(X%5B_idx%5D).mean()%0A%0A%20%20%20%20tab_hist%20%3D%20md_table(%5B%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmin_i%20a_i%2C%5C%20%5Csum_i%20a_i%24%22%2C%20rf%22%24%7Ba.min()%3A.4f%7D%2C%5C%20%7Ba.sum()%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmin_j%20b_j%2C%5C%20%5Csum_j%20b_j%24%22%2C%20rf%22%24%7Bb.min()%3A.4f%7D%2C%5C%20%7Bb.sum()%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Csum_i%20a_i%20f(x_i)%24%2C%20with%20%24f(x)%3D%5ClVert%20x%5CrVert%5E2%24%22%2C%20rf%22%24%7B_exact%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cint%20f%5C%2C%5Cmathrm%7Bd%7D%5Calpha%24%20(Monte%20Carlo%2C%20%2410%5E5%24%20samples)%22%2C%20rf%22%24%7B_mc%3A.6f%7D%24%22)%2C%0A%20%20%20%20%5D%2C%20(%22Check%22%2C%20%22Value%22))%0A%0A%20%20%20%20mo.center(mo.as_html(_fig))%0A%20%20%20%20return%20(tab_hist%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20tab_hist)%3A%0A%20%20%20%20mo.center(mo.md(tab_hist))%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%202.%20Assignment%20and%20Monge%20Problem%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%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Assignment%0A%0A%20%20%20%20**Approach.**%20Given%20a%20cost%20matrix%20%24(%5Cmathbf%7BC%7D_%7Bi%2Cj%7D)_%7Bi%5Cin%5Cllbracket%20n%20%5Crrbracket%2C%5C%2C%20j%5Cin%5Cllbracket%20m%20%5Crrbracket%7D%24%20and%20assuming%20%24n%3Dm%24%2C%20solve%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmin_%7B%5Csigma%5Cin%5Cmathrm%7BPerm%7D(n)%7D%20%5Cfrac1n%5Csum_%7Bi%3D1%7D%5En%20%5Cmathbf%7BC%7D_%7Bi%2C%5Csigma(i)%7D.%0A%20%20%20%20%24%24%0A%20%20%20%20Combinatorial%20problem%3A%20%24n!%24%20candidates.%0A%0A%20%20%20%20*Is%20the%20solution%20unique%3F*%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20fmt_vec%2C%0A%20%20%20%20itertools%2C%0A%20%20%20%20linear_sum_assignment%2C%0A%20%20%20%20md_table%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20np%2C%0A%20%20%20%20pairwise_dist%2C%0A%20%20%20%20plt%2C%0A)%3A%0A%20%20%20%20_rng%20%3D%20np.random.default_rng(1)%0A%20%20%20%20n_assign%20%3D%206%0A%20%20%20%20_Xa%20%3D%20_rng.random((n_assign%2C%202))%0A%20%20%20%20_Ya%20%3D%20_rng.random((n_assign%2C%202))%20%2B%20%5B1.2%2C%200.0%5D%0A%20%20%20%20C_assign%20%3D%20pairwise_dist(_Xa%2C%20_Ya)%20**%202%0A%0A%20%20%20%20_perms%20%3D%20list(itertools.permutations(range(n_assign)))%0A%20%20%20%20_costs%20%3D%20np.array(%5BC_assign%5Bnp.arange(n_assign)%2C%20s%5D.mean()%20for%20s%20in%20_perms%5D)%0A%20%20%20%20_sigma_bf%20%3D%20np.array(_perms%5B_costs.argmin()%5D)%0A%20%20%20%20_%2C%20sigma_assign%20%3D%20linear_sum_assignment(C_assign)%0A%20%20%20%20cost_assign%20%3D%20C_assign%5Bnp.arange(n_assign)%2C%20sigma_assign%5D.mean()%0A%20%20%20%20_n_opt%20%3D%20int(np.sum(np.isclose(_costs%2C%20_costs.min())))%0A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(6%2C%203.5))%0A%20%20%20%20_ax.scatter(*_Xa.T%2C%20color%3D%22%231f4e79%22%2C%20s%3D80%2C%20zorder%3D3)%0A%20%20%20%20_ax.scatter(*_Ya.T%2C%20color%3D%22%23d66e1e%22%2C%20s%3D80%2C%20zorder%3D3)%0A%20%20%20%20for%20_i%2C%20_j%20in%20enumerate(sigma_assign)%3A%0A%20%20%20%20%20%20%20%20_ax.annotate(%22%22%2C%20xy%3D_Ya%5B_j%5D%2C%20xytext%3D_Xa%5B_i%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20arrowprops%3Ddict(arrowstyle%3D%22-%7C%3E%22%2C%20color%3D%22gray%22))%0A%20%20%20%20_ax.set_aspect(%22equal%22)%0A%20%20%20%20_ax.set_title(r%22Optimal%20%24%5Csigma%24%2C%20%24C_%7Bij%7D%3D%5C%7Cx_i-y_j%5C%7C%5E2%24%22)%0A%0A%20%20%20%20tab_assign%20%3D%20md_table(%5B%0A%20%20%20%20%20%20%20%20(r%22Candidates%20%24n!%24%22%2C%20rf%22%24%7Blen(_perms)%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22Brute%20force%3A%20%24%5Csigma%5E%5Cstar%24%2C%20cost%22%2C%20rf%22%24%7Bfmt_vec(_sigma_bf)%7D%2C%5C%20%7B_costs.min()%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22Hungarian%3A%20%24%5Csigma%5E%5Cstar%24%2C%20cost%22%2C%20rf%22%24%7Bfmt_vec(sigma_assign)%7D%2C%5C%20%7Bcost_assign%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22Number%20of%20optimal%20%24%5Csigma%24%22%2C%20rf%22%24%7B_n_opt%7D%24%22)%2C%0A%20%20%20%20%5D)%0A%0A%20%20%20%20mo.center(mo.as_html(_fig))%0A%20%20%20%20return%20C_assign%2C%20cost_assign%2C%20n_assign%2C%20sigma_assign%2C%20tab_assign%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20tab_assign)%3A%0A%20%20%20%20mo.center(mo.md(tab_assign))%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(fmt_vec%2C%20itertools%2C%20md_table%2C%20mo%2C%20np)%3A%0A%20%20%20%20_x%20%3D%20np.array(%5B0.0%2C%201.0%5D)%0A%20%20%20%20_y%20%3D%20np.array(%5B1.0%2C%202.0%5D)%0A%20%20%20%20_rows%20%3D%20%5B%5D%0A%20%20%20%20for%20_p%20in%20%5B1%2C%202%5D%3A%0A%20%20%20%20%20%20%20%20_C%20%3D%20np.abs(_x%5B%3A%2C%20None%5D%20-%20_y%5BNone%2C%20%3A%5D)%20**%20_p%0A%20%20%20%20%20%20%20%20for%20_s%20in%20itertools.permutations(range(2))%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_cost%20%3D%20_C%5Bnp.arange(2)%2C%20list(_s)%5D.mean()%0A%20%20%20%20%20%20%20%20%20%20%20%20_rows.append((rf%22%24%7B_p%7D%24%22%2C%20rf%22%24%7Bfmt_vec(_s)%7D%24%22%2C%20rf%22%24%7B_cost%3A.3f%7D%24%22))%0A%0A%20%20%20%20mo.center(mo.vstack(%5B%0A%20%20%20%20%20%20%20%20mo.md(rf%22%24%5Cmathbf%7B%7Bx%7D%7D%20%3D%20%7Bfmt_vec(_x)%7D%24%2C%20%24%5Cmathbf%7B%7By%7D%7D%20%3D%20%7Bfmt_vec(_y)%7D%24%2C%20%24%5Cmathbf%7B%7BC%7D%7D_%7B%7Bij%7D%7D%20%3D%20%5Clvert%20x_i%20-%20y_j%5Crvert%5Ep%24%22)%2C%0A%20%20%20%20%20%20%20%20mo.md(md_table(_rows%2C%20(%22%24p%24%22%2C%20r%22%24%5Csigma%24%22%2C%20r%22%24%5Cfrac1n%5Csum_i%20%5Cmathbf%7BC%7D_%7Bi%2C%5Csigma(i)%7D%24%22)))%2C%0A%20%20%20%20%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%23%20Monge%20Problem%0A%0A%20%20%20%20**Discrete%20Monge.**%20Find%20%24T%3A%5C%7Bx_i%5C%7D%5Cto%5C%7By_j%5C%7D%24%20that%20*pushes*%20%24%5Calpha%24%20onto%20%24%5Cbeta%24%3A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cforall%20j%2C%5C%3B%20b_j%20%3D%20%5Csum_%7Bi%3A%5C%2CT(x_i)%3Dy_j%7D%20a_i%20%5Cqquad%5CLongleftrightarrow%5Cqquad%20T_%7B%5C%23%7D%5Calpha%20%3D%20%5Cbeta.%0A%20%20%20%20%24%24%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmin_T%5CBig%5C%7B%20%5Csum_i%20c(x_i%2C%20T(x_i))%3A%20T_%7B%5C%23%7D%5Calpha%20%3D%20%5Cbeta%5CBig%5C%7D%0A%20%20%20%20%24%24%0A%0A%20%20%20%20**General%20Monge%20(1781).**%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmin_%7BT%7D%20%5CBig%5C%7B%20%5Cint_%7B%5Cmathcal%7BX%7D%7D%20c%5Cbig(x%2CT(x)%5Cbig)%5C%2C%5Cmathrm%7Bd%7D%5Calpha(x)%20%5C%3B%3A%5C%3B%20T_%7B%5C%23%7D%5Calpha%3D%5Cbeta%5CBig%5C%7D.%0A%20%20%20%20%24%24%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(fmt_vec%2C%20itertools%2C%20md_table%2C%20mo%2C%20np)%3A%0A%20%20%20%20_x%20%3D%20np.array(%5B0.0%2C%201.0%2C%202.0%2C%203.0%5D)%0A%20%20%20%20_a%20%3D%20np.array(%5B0.1%2C%200.4%2C%200.3%2C%200.2%5D)%0A%20%20%20%20_y%20%3D%20np.array(%5B0.5%2C%202.5%5D)%0A%20%20%20%20_b%20%3D%20np.array(%5B0.5%2C%200.5%5D)%0A%0A%20%20%20%20def%20_push(T%2C%20w%2C%20k)%3A%0A%20%20%20%20%20%20%20%20return%20np.bincount(T%2C%20weights%3Dw%2C%20minlength%3Dk)%0A%0A%20%20%20%20_rows%20%3D%20%5B%5D%0A%20%20%20%20_n_maps%20%3D%200%0A%20%20%20%20for%20_T%20in%20itertools.product(range(len(_y))%2C%20repeat%3Dlen(_x))%3A%0A%20%20%20%20%20%20%20%20_n_maps%20%2B%3D%201%0A%20%20%20%20%20%20%20%20_T%20%3D%20np.array(_T)%0A%20%20%20%20%20%20%20%20if%20np.allclose(_push(_T%2C%20_a%2C%20len(_y))%2C%20_b)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_cost%20%3D%20np.sum((_x%20-%20_y%5B_T%5D)%20**%202)%0A%20%20%20%20%20%20%20%20%20%20%20%20_rows.append((rf%22%24%7Bfmt_vec(_y%5B_T%5D)%7D%24%22%2C%20rf%22%24%7B_cost%3A.3f%7D%24%22))%0A%0A%20%20%20%20_a1%20%3D%20np.array(%5B1.0%5D)%0A%20%20%20%20_b1%20%3D%20np.array(%5B0.5%2C%200.5%5D)%0A%20%20%20%20_n_feasible_1%20%3D%20int(sum(%0A%20%20%20%20%20%20%20%20np.allclose(_push(np.array(_T1)%2C%20_a1%2C%202)%2C%20_b1)%0A%20%20%20%20%20%20%20%20for%20_T1%20in%20itertools.product(range(2)%2C%20repeat%3D1)%0A%20%20%20%20))%0A%0A%20%20%20%20mo.vstack(%5B%0A%20%20%20%20%20%20%20%20mo.md(rf%22%22%22%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cbegin%7B%7Baligned%7D%7D%0A%20%20%20%20%5Cmathbf%7B%7Bx%7D%7D%20%26%3D%20%7Bfmt_vec(_x)%7D%2C%20%26%20%5Cmathbf%7B%7Ba%7D%7D%20%26%3D%20%7Bfmt_vec(_a)%7D%2C%20%5C%5C%0A%20%20%20%20%5Cmathbf%7B%7By%7D%7D%20%26%3D%20%7Bfmt_vec(_y)%7D%2C%20%26%20%5Cmathbf%7B%7Bb%7D%7D%20%26%3D%20%7Bfmt_vec(_b)%7D%2C%20%26%20c(x%2Cy)%20%26%3D%20%7Cx-y%7C%5E2%0A%20%20%20%20%5Cend%7B%7Baligned%7D%7D%0A%20%20%20%20%24%24%0A%0A%20%20%20%20%24%7B_n_maps%7D%24%20maps%20%24T%3A%5C%7B%7Bx_i%5C%7D%7D%5Cto%5C%7B%7By_j%5C%7D%7D%24%20such%20that%20%24T_%7B%7B%5C%23%7D%7D%5Calpha%20%3D%20%5Cbeta%24%3A%0A%20%20%20%20%22%22%22)%2C%0A%20%20%20%20%20%20%20%20mo.center(mo.md(md_table(_rows%2C%20(r%22%24%5Cbig(T(x_1)%2C%5Cdots%2CT(x_4)%5Cbig)%24%22%2C%20r%22%24%5Csum_i%20c%5Cbig(x_i%2C%20T(x_i)%5Cbig)%24%22))))%2C%0A%20%20%20%20%20%20%20%20mo.callout(mo.md(%0A%20%20%20%20%20%20%20%20%20%20%20%20rf%22%24%5Calpha%20%3D%20%5Cdelta_0%24%2C%20%24%5Cbeta%20%3D%20%5Ctfrac12%5Cdelta_0%20%2B%20%5Ctfrac12%5Cdelta_1%24%3A%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20rf%22number%20of%20maps%20with%20%24T_%7B%7B%5C%23%7D%7D%5Calpha%20%3D%20%5Cbeta%24%20is%20%24%5Cmathbf%7B%7B%7B_n_feasible_1%7D%7D%7D%24.%22%0A%20%20%20%20%20%20%20%20)%2C%20kind%3D%22warn%22)%2C%0A%20%20%20%20%5D%2C%20align%3D%22center%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(md_table%2C%20mo%2C%20norm%2C%20np%2C%20plt)%3A%0A%20%20%20%20_ma%2C%20_sa%2C%20_mb%2C%20_sb%20%3D%200.0%2C%201.0%2C%203.0%2C%200.5%0A%20%20%20%20_T%20%3D%20lambda%20x%3A%20_mb%20%2B%20_sb%20%2F%20_sa%20*%20(x%20-%20_ma)%0A%20%20%20%20_rng%20%3D%20np.random.default_rng(2)%0A%20%20%20%20_xs%20%3D%20_rng.normal(_ma%2C%20_sa%2C%20200_000)%0A%20%20%20%20_tx%20%3D%20_T(_xs)%0A%20%20%20%20_grid%20%3D%20np.linspace(-4%2C%205.5%2C%20400)%0A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(7%2C%203))%0A%20%20%20%20_ax.hist(_xs%2C%20bins%3D150%2C%20density%3DTrue%2C%20color%3D%22%231f4e79%22%2C%20alpha%3D0.35)%0A%20%20%20%20_ax.hist(_tx%2C%20bins%3D150%2C%20density%3DTrue%2C%20color%3D%22%23d66e1e%22%2C%20alpha%3D0.35)%0A%20%20%20%20_ax.plot(_grid%2C%20norm.pdf(_grid%2C%20_ma%2C%20_sa)%2C%20color%3D%22%231f4e79%22%2C%20label%3Dr%22%24%5Calpha%24%22)%0A%20%20%20%20_ax.plot(_grid%2C%20norm.pdf(_grid%2C%20_mb%2C%20_sb)%2C%20color%3D%22%23d66e1e%22%2C%20label%3Dr%22%24%5Cbeta%24%22)%0A%20%20%20%20_ax.set_title(r%22%24T(x)%3D%5Cmu_%5Cbeta%2B%5Cfrac%7B%5Csigma_%5Cbeta%7D%7B%5Csigma_%5Calpha%7D(x-%5Cmu_%5Calpha)%24%2C%20histogram%20of%20%24T(X)%24%22)%0A%20%20%20%20_ax.legend()%0A%0A%20%20%20%20_mc%20%3D%20np.mean((_xs%20-%20_tx)%20**%202)%0A%20%20%20%20_exact%20%3D%20(_ma%20-%20_mb)%20**%202%20%2B%20(_sa%20-%20_sb)%20**%202%0A%20%20%20%20tab_monge%20%3D%20md_table(%5B%0A%20%20%20%20%20%20%20%20(r%22%24%5Coperatorname%7Bmean%7D%5Cbig(T(X)%5Cbig)%2C%5C%20%5Coperatorname%7Bstd%7D%5Cbig(T(X)%5Cbig)%24%22%2C%20rf%22%24%7B_tx.mean()%3A.4f%7D%2C%5C%20%7B_tx.std()%3A.4f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmu_%5Cbeta%2C%5C%20%5Csigma_%5Cbeta%24%22%2C%20rf%22%24%7B_mb%3Ag%7D%2C%5C%20%7B_sb%3Ag%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cint%20%5Cbig(x-T(x)%5Cbig)%5E2%5C%2C%5Cmathrm%7Bd%7D%5Calpha(x)%24%20(Monte%20Carlo)%22%2C%20rf%22%24%7B_mc%3A.4f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24(%5Cmu_%5Calpha-%5Cmu_%5Cbeta)%5E2%20%2B%20(%5Csigma_%5Calpha-%5Csigma_%5Cbeta)%5E2%24%22%2C%20rf%22%24%7B_exact%3A.4f%7D%24%22)%2C%0A%20%20%20%20%5D)%0A%0A%20%20%20%20mo.center(mo.as_html(_fig))%0A%20%20%20%20return%20(tab_monge%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20tab_monge)%3A%0A%20%20%20%20mo.center(mo.md(tab_monge))%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%203.%20Kantorovich%20Relaxation%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%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Coupling%20Approach%0A%0A%20%20%20%20We%20change%20the%20Transport%20Map%20%24T%24%20for%20a%20transport%20matrix%20%24%5Cmathbf%7BP%7D%5Cin%5Cmathbb%7BR%7D_%2B%5E%7Bn%5Ctimes%20m%7D%24%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathbf%7BU%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%20%5Ccoloneqq%20%5Cbig%5C%7B%5Cmathbf%7BP%7D%5Cin%5Cmathbb%7BR%7D_%2B%5E%7Bn%5Ctimes%20m%7D%20%3A%20%5Cmathbf%7BP%7D%5Cmathbf%7B1%7D_m%3D%5Cmathbf%7Ba%7D%2C%5C%3B%20%5Cmathbf%7BP%7D%5E%7B%5Ctop%7D%5Cmathbf%7B1%7D_n%3D%5Cmathbf%7Bb%7D%5Cbig%5C%7D.%0A%20%20%20%20%24%24%0A%0A%20%20%20%20-%20%24%5Cmathbf%7BU%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%24%20is%20a%20convex%20polytope.%0A%20%20%20%20-%20%24%5Cmathbf%7Ba%7D%5Cmathbf%7Bb%7D%5E%5Ctop%5Cin%5Cmathbf%7BU%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%24.%0A%20%20%20%20-%20%24%5Cmathbf%7BP%7D%5Cin%5Cmathbf%7BU%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%5Ciff%5Cmathbf%7BP%7D%5E%5Ctop%5Cin%5Cmathbf%7BU%7D(%5Cmathbf%7Bb%7D%2C%5Cmathbf%7Ba%7D)%24.%0A%0A%20%20%20%20**Kantorovich%20Problem%3A**%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BC%7D%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%20%5Ccoloneqq%20%5Cmin_%7B%5Cmathbf%7BP%7D%5Cin%5Cmathbf%7BU%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%7D%20%5Clangle%5Cmathbf%7BC%7D%2C%5C%2C%5Cmathbf%7BP%7D%5Crangle%20%3D%20%5Csum_%7Bi%2Cj%7D%5Cmathbf%7BC%7D_%7Bij%7D%5Cmathbf%7BP%7D_%7Bij%7D.%0A%20%20%20%20%24%24%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(X%2C%20Y%2C%20a%2C%20b%2C%20kantorovich%2C%20pairwise_dist)%3A%0A%20%20%20%20C%20%3D%20pairwise_dist(X%2C%20Y)%20**%202%0A%20%20%20%20P_opt%2C%20L_opt%2C%20_%20%3D%20kantorovich(a%2C%20b%2C%20C)%0A%20%20%20%20return%20C%2C%20L_opt%2C%20P_opt%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20C%2C%0A%20%20%20%20L_opt%2C%0A%20%20%20%20P_opt%2C%0A%20%20%20%20X%2C%0A%20%20%20%20Y%2C%0A%20%20%20%20a%2C%0A%20%20%20%20b%2C%0A%20%20%20%20kantorovich%2C%0A%20%20%20%20m%2C%0A%20%20%20%20md_table%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20n%2C%0A%20%20%20%20np%2C%0A%20%20%20%20plt%2C%0A%20%20%20%20sci%2C%0A)%3A%0A%20%20%20%20_P_ind%20%3D%20np.outer(a%2C%20b)%0A%20%20%20%20_%2C%20_L_T%2C%20_%20%3D%20kantorovich(b%2C%20a%2C%20C.T)%0A%20%20%20%20_nnz%20%3D%20int(np.sum(P_opt%20%3E%201e-10))%0A%0A%20%20%20%20_fig%20%3D%20plt.figure(figsize%3D(11%2C%204))%0A%20%20%20%20_gs%20%3D%20_fig.add_gridspec(2%2C%203%2C%20width_ratios%3D%5B1%2C%204%2C%206%5D%2C%20height_ratios%3D%5B1%2C%204%5D)%0A%20%20%20%20_axP%20%3D%20_fig.add_subplot(_gs%5B1%2C%201%5D)%0A%20%20%20%20_axP.imshow(P_opt%2C%20cmap%3D%22Blues%22%2C%20aspect%3D%22auto%22)%0A%20%20%20%20_axP.set_xlabel(r%22%24%5Cmathbf%7BP%7D%24%22)%0A%20%20%20%20_axA%20%3D%20_fig.add_subplot(_gs%5B1%2C%200%5D%2C%20sharey%3D_axP)%0A%20%20%20%20_axA.barh(np.arange(n)%2C%20a%2C%20color%3D%22%231f4e79%22)%0A%20%20%20%20_axA.invert_xaxis()%0A%20%20%20%20_axA.set_xlabel(r%22%24%5Cmathbf%7Ba%7D%24%22)%0A%20%20%20%20_axB%20%3D%20_fig.add_subplot(_gs%5B0%2C%201%5D%2C%20sharex%3D_axP)%0A%20%20%20%20_axB.bar(np.arange(m)%2C%20b%2C%20color%3D%22%23d66e1e%22)%0A%20%20%20%20_axB.set_title(r%22%24%5Cmathbf%7Bb%7D%24%22)%0A%20%20%20%20_ax2%20%3D%20_fig.add_subplot(_gs%5B%3A%2C%202%5D)%0A%20%20%20%20for%20_i%20in%20range(n)%3A%0A%20%20%20%20%20%20%20%20for%20_j%20in%20range(m)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20P_opt%5B_i%2C%20_j%5D%20%3E%201e-10%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_ax2.plot(%5BX%5B_i%2C%200%5D%2C%20Y%5B_j%2C%200%5D%5D%2C%20%5BX%5B_i%2C%201%5D%2C%20Y%5B_j%2C%201%5D%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20color%3D%22gray%22%2C%20lw%3D40%20*%20P_opt%5B_i%2C%20_j%5D%2C%20alpha%3D0.7)%0A%20%20%20%20_ax2.scatter(*X.T%2C%20s%3D3000%20*%20a%2C%20color%3D%22%231f4e79%22%2C%20alpha%3D0.8%2C%20zorder%3D3)%0A%20%20%20%20_ax2.scatter(*Y.T%2C%20s%3D3000%20*%20b%2C%20color%3D%22%23d66e1e%22%2C%20alpha%3D0.85%2C%20zorder%3D3)%0A%20%20%20%20_ax2.set_aspect(%22equal%22)%0A%20%20%20%20_fig.tight_layout()%0A%0A%20%20%20%20_e_a%20%3D%20sci(np.linalg.norm(P_opt.sum(1)%20-%20a))%0A%20%20%20%20_e_b%20%3D%20sci(np.linalg.norm(P_opt.sum(0)%20-%20b))%0A%20%20%20%20_e_ia%20%3D%20sci(np.linalg.norm(_P_ind.sum(1)%20-%20a))%0A%20%20%20%20_e_ib%20%3D%20sci(np.linalg.norm(_P_ind.sum(0)%20-%20b))%0A%0A%20%20%20%20tab_kant%20%3D%20md_table(%5B%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BC%7D%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%24%22%2C%20rf%22%24%7BL_opt%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5ClVert%20%5Cmathbf%7BP%7D%5Cmathbf%7B1%7D_m%20-%20%5Cmathbf%7Ba%7D%20%5CrVert%2C%5C%20%5ClVert%20%5Cmathbf%7BP%7D%5E%5Ctop%5Cmathbf%7B1%7D_n%20-%20%5Cmathbf%7Bb%7D%20%5CrVert%24%22%2C%20rf%22%24%7B_e_a%7D%2C%5C%20%7B_e_b%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathbf%7Ba%7D%5Cmathbf%7Bb%7D%5E%5Ctop%20%5Cin%20%5Cmathbf%7BU%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%24%3A%20%24%5ClVert%20%5Cmathbf%7Ba%7D%5Cmathbf%7Bb%7D%5E%5Ctop%5Cmathbf%7B1%7D_m%20-%20%5Cmathbf%7Ba%7D%20%5CrVert%2C%5C%20%5ClVert%20(%5Cmathbf%7Ba%7D%5Cmathbf%7Bb%7D%5E%5Ctop)%5E%5Ctop%5Cmathbf%7B1%7D_n%20-%20%5Cmathbf%7Bb%7D%20%5CrVert%24%22%2C%20rf%22%24%7B_e_ia%7D%2C%5C%20%7B_e_ib%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Clangle%20%5Cmathbf%7BC%7D%2C%20%5Cmathbf%7Ba%7D%5Cmathbf%7Bb%7D%5E%5Ctop%20%5Crangle%24%22%2C%20rf%22%24%7Bnp.sum(C%20*%20_P_ind)%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BC%7D%5E%5Ctop%7D(%5Cmathbf%7Bb%7D%2C%5Cmathbf%7Ba%7D)%24%22%2C%20rf%22%24%7B_L_T%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5C%23%5C%7B(i%2Cj)%20%3A%20%5Cmathbf%7BP%7D_%7Bij%7D%20%3E%200%5C%7D%2C%5C%20n%2Bm-1%24%22%2C%20rf%22%24%7B_nnz%7D%2C%5C%20%7Bn%20%2B%20m%20-%201%7D%24%22)%2C%0A%20%20%20%20%5D%2C%20(%22Check%22%2C%20%22Value%22))%0A%0A%20%20%20%20mo.center(mo.as_html(_fig))%0A%20%20%20%20return%20(tab_kant%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20tab_kant)%3A%0A%20%20%20%20mo.center(mo.md(tab_kant))%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%23%20The%20reformulation%20%E2%80%94%20Birkhoff%0A%0A%20%20%20%20Suppose%20%24n%3Dm%24%20and%20%24%5Cmathbf%7Ba%7D%3D%5Cmathbf%7Bb%7D%3D%5Cmathbf%7B1%7D_n%2Fn%24%2C%20the%20set%20of%20permutation%20matrices%20are%20included%20in%20%24%5Cmathbf%7BU%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%24%2C%20so%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BC%7D%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%20%5C%3B%5Cle%5C%3B%20%5Cmin_%7B%5Csigma%7D%5C%20%5Cfrac1n%5Csum_i%20%5Cmathbf%7BC%7D_%7Bi%2C%5Csigma(i)%7D.%0A%20%20%20%20%24%24%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%20mo.md(r%22%22%22%0A%20%20%20%20**Proposition%20(Kantorovich%20for%20matching).**%20If%20%24n%3Dm%24%20and%20%24%5Cmathbf%7Ba%7D%3D%5Cmathbf%7Bb%7D%3D%5Cmathbf%7B1%7D_n%2Fn%24%2C%20then%20there%20exists%20an%20optimal%20solution%20for%20Kantorovich%2C%20which%20is%20a%20permutation%20matrix%20associated%20to%20an%20optimal%20permutation%20for%20the%20assignment%20problem.%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BC%7D%7D(%5Cmathbf%7B1%7D_n%2Fn%2C%5Cmathbf%7B1%7D_n%2Fn)%20%3D%20%5Cmin_%7B%5Csigma%5Cin%20%5Cmathrm%7BPerm%7D(n)%7D%5Clangle%5Cmathbf%7BC%7D%2C%5C%2C%5Cmathbf%7BP%7D_%5Csigma%5Crangle%0A%20%20%20%20%24%24%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20C_assign%2C%0A%20%20%20%20bmat%2C%0A%20%20%20%20cost_assign%2C%0A%20%20%20%20fmt_vec%2C%0A%20%20%20%20kantorovich%2C%0A%20%20%20%20md_table%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20n_assign%2C%0A%20%20%20%20np%2C%0A%20%20%20%20sigma_assign%2C%0A)%3A%0A%20%20%20%20_u%20%3D%20np.ones(n_assign)%20%2F%20n_assign%0A%20%20%20%20_P%2C%20_L%2C%20_%20%3D%20kantorovich(_u%2C%20_u%2C%20C_assign)%0A%20%20%20%20_Pn%20%3D%20_P%20*%20n_assign%0A%20%20%20%20_is_perm%20%3D%20bool(np.allclose(_Pn%2C%20np.round(_Pn))%20and%20np.allclose(np.round(_Pn).sum(1)%2C%201))%0A%20%20%20%20_sigma_lp%20%3D%20np.argmax(_P%2C%20axis%3D1)%0A%20%20%20%20_Pn_tex%20%3D%20bmat(_Pn%2C%20%22%7B%3A.3g%7D%22)%0A%0A%20%20%20%20_rows%20%3D%20%5B%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BC%7D%7D(%5Cmathbf%7B1%7D_n%2Fn%2C%5C%2C%20%5Cmathbf%7B1%7D_n%2Fn)%24%22%2C%20rf%22%24%7B_L%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmin_%7B%5Csigma%7D%20%5Cfrac%7B1%7D%7Bn%7D%20%5Csum_i%20%5Cmathbf%7BC%7D_%7Bi%2C%5Csigma(i)%7D%24%22%2C%20rf%22%24%7Bcost_assign%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathbf%7BP%7D%24%20is%20a%20permutation%20matrix%22%2C%20rf%22%24%5Ctext%7B%7B%7B_is_perm%7D%7D%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Csigma%24%20from%20LP%2C%20%24%5Csigma%24%20from%20Hungarian%22%2C%20rf%22%24%7Bfmt_vec(_sigma_lp)%7D%2C%5C%20%7Bfmt_vec(sigma_assign)%7D%24%22)%2C%0A%20%20%20%20%5D%0A%20%20%20%20mo.vstack(%5B%0A%20%20%20%20%20%20%20%20mo.md(rf%22%24%24%20n%5C%2C%5Cmathbf%7B%7BP%7D%7D%20%3D%20%7B_Pn_tex%7D%20%24%24%22)%2C%0A%20%20%20%20%20%20%20%20mo.center(mo.md(md_table(_rows)))%2C%0A%20%20%20%20%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%23%20Marginals%20and%20projections%0A%0A%20%20%20%20%24%24%0A%20%20%20%20P_%7B%5Cmathcal%7BX%7D%5C%23%7D%5Cpi(A)%20%3D%20%5Cpi%5Cbig(P_%7B%5Cmathcal%7BX%7D%7D%5E%7B-1%7D(A)%5Cbig)%20%3D%20%5Cpi(A%5Ctimes%5Cmathcal%7BY%7D).%0A%20%20%20%20%24%24%0A%0A%20%20%20%20**Marginal%20constraints.**%20%24%5Cpi%5Cin%5Cmathbf%7BU%7D(%5Calpha%2C%5Cbeta)%24%20iff%20for%20all%20measurable%20%24A%5Csubset%5Cmathcal%7BX%7D%24%2C%20%24B%5Csubset%5Cmathcal%7BY%7D%24%3A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cpi(A%5Ctimes%5Cmathcal%7BY%7D)%3D%5Calpha(A)%2C%20%5Cqquad%20%5Cpi(%5Cmathcal%7BX%7D%5Ctimes%20B)%3D%5Cbeta(B).%0A%20%20%20%20%24%24%0A%0A%20%20%20%20**With%20test%20functions.**%20Equivalently%2C%20for%20all%20continuous%20%24f%24%2C%20%24g%24%3A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cint_%7B%5Cmathcal%7BX%7D%5Ctimes%5Cmathcal%7BY%7D%7D%20f(x)%5C%2C%5Cmathrm%7Bd%7D%5Cpi(x%2Cy)%20%3D%20%5Cint_%7B%5Cmathcal%7BX%7D%7D%20f%5C%2C%5Cmathrm%7Bd%7D%5Calpha%2C%20%5Cqquad%20%5Cint_%7B%5Cmathcal%7BX%7D%5Ctimes%5Cmathcal%7BY%7D%7D%20g(y)%5C%2C%5Cmathrm%7Bd%7D%5Cpi(x%2Cy)%20%3D%20%5Cint_%7B%5Cmathcal%7BY%7D%7D%20g%5C%2C%5Cmathrm%7Bd%7D%5Cbeta.%0A%20%20%20%20%24%24%0A%0A%20%20%20%20**Discrete%20case.**%20%24%5Cpi%3D%5Csum_%7Bi%2Cj%7D%5Cmathbf%7BP%7D_%7Bi%2Cj%7D%5C%2C%5Cdelta_%7B(x_i%2Cy_j)%7D%24%3A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathbf%7BP%7D%5Cmathbf%7B1%7D_m%3D%5Cmathbf%7Ba%7D%2C%20%5Cqquad%20%5Cmathbf%7BP%7D%5E%5Ctop%5Cmathbf%7B1%7D_n%3D%5Cmathbf%7Bb%7D%0A%20%20%20%20%24%24%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(md_table%2C%20mo%2C%20norm%2C%20np%2C%20plt)%3A%0A%20%20%20%20_mx%2C%20_sx%2C%20_my%2C%20_sy%2C%20_rho%20%3D%201.5%2C%20np.sqrt(0.175)%2C%202.2%2C%20np.sqrt(0.3)%2C%200.8%0A%20%20%20%20_cov%20%3D%20np.array(%5B%5B_sx**2%2C%20_rho%20*%20_sx%20*%20_sy%5D%2C%20%5B_rho%20*%20_sx%20*%20_sy%2C%20_sy**2%5D%5D)%0A%20%20%20%20_rng%20%3D%20np.random.default_rng(4)%0A%20%20%20%20_S%20%3D%20_rng.multivariate_normal(%5B_mx%2C%20_my%5D%2C%20_cov%2C%20200_000)%0A%20%20%20%20_A%20%3D%20(1.2%2C%201.8)%0A%20%20%20%20_B%20%3D%20(2.0%2C%202.6)%0A%20%20%20%20_piA%20%3D%20np.mean((_S%5B%3A%2C%200%5D%20%3E%3D%20_A%5B0%5D)%20%26%20(_S%5B%3A%2C%200%5D%20%3C%3D%20_A%5B1%5D))%0A%20%20%20%20_alphaA%20%3D%20norm.cdf(_A%5B1%5D%2C%20_mx%2C%20_sx)%20-%20norm.cdf(_A%5B0%5D%2C%20_mx%2C%20_sx)%0A%20%20%20%20_piB%20%3D%20np.mean((_S%5B%3A%2C%201%5D%20%3E%3D%20_B%5B0%5D)%20%26%20(_S%5B%3A%2C%201%5D%20%3C%3D%20_B%5B1%5D))%0A%20%20%20%20_betaB%20%3D%20norm.cdf(_B%5B1%5D%2C%20_my%2C%20_sy)%20-%20norm.cdf(_B%5B0%5D%2C%20_my%2C%20_sy)%0A%20%20%20%20_Ef%20%3D%20np.mean(np.cos(_S%5B%3A%2C%200%5D))%0A%20%20%20%20_f_exact%20%3D%20np.exp(-_sx**2%20%2F%202)%20*%20np.cos(_mx)%0A%20%20%20%20_Eg%20%3D%20np.mean(_S%5B%3A%2C%201%5D%20**%202)%0A%20%20%20%20_g_exact%20%3D%20_my**2%20%2B%20_sy**2%0A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(4.5%2C%204.5))%0A%20%20%20%20_ax.hexbin(_S%5B%3A5000%2C%200%5D%2C%20_S%5B%3A5000%2C%201%5D%2C%20gridsize%3D35%2C%20cmap%3D%22Greens%22%2C%20mincnt%3D1)%0A%20%20%20%20_ax.axvspan(*_A%2C%20color%3D%22gray%22%2C%20alpha%3D0.2)%0A%20%20%20%20_ax.set_xlabel(%22x%22)%0A%20%20%20%20_ax.set_ylabel(%22y%22)%0A%20%20%20%20_ax.set_title(r%22%24%5Cpi%24%20and%20%24A%5Ctimes%5Cmathcal%7BY%7D%24%22)%0A%0A%20%20%20%20tab_marg%20%3D%20md_table(%5B%0A%20%20%20%20%20%20%20%20(rf%22%24%5Cpi(A%5Ctimes%5Cmathcal%7B%7BY%7D%7D)%24%20vs%20%24%5Calpha(A)%24%2C%20%24A%3D%5B%7B_A%5B0%5D%3Ag%7D%2C%5C%2C%7B_A%5B1%5D%3Ag%7D%5D%24%22%2C%20rf%22%24%7B_piA%3A.4f%7D%24%22%2C%20rf%22%24%7B_alphaA%3A.4f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(rf%22%24%5Cpi(%5Cmathcal%7B%7BX%7D%7D%5Ctimes%20B)%24%20vs%20%24%5Cbeta(B)%24%2C%20%24B%3D%5B%7B_B%5B0%5D%3Ag%7D%2C%5C%2C%7B_B%5B1%5D%3Ag%7D%5D%24%22%2C%20rf%22%24%7B_piB%3A.4f%7D%24%22%2C%20rf%22%24%7B_betaB%3A.4f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cint%20%5Ccos%20x%5C%2C%5Cmathrm%7Bd%7D%5Cpi%24%20vs%20%24%5Cint%20%5Ccos%5C%2C%5Cmathrm%7Bd%7D%5Calpha%24%22%2C%20rf%22%24%7B_Ef%3A.4f%7D%24%22%2C%20rf%22%24%7B_f_exact%3A.4f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cint%20y%5E2%5C%2C%5Cmathrm%7Bd%7D%5Cpi%24%20vs%20%24%5Cint%20y%5E2%5C%2C%5Cmathrm%7Bd%7D%5Cbeta%24%22%2C%20rf%22%24%7B_Eg%3A.4f%7D%24%22%2C%20rf%22%24%7B_g_exact%3A.4f%7D%24%22)%2C%0A%20%20%20%20%5D%2C%20(%22%22%2C%20r%22%24%5Cpi%24%20(Monte%20Carlo)%22%2C%20%22Marginal%20(exact)%22))%0A%0A%20%20%20%20mo.center(mo.as_html(_fig))%0A%20%20%20%20return%20(tab_marg%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20tab_marg)%3A%0A%20%20%20%20mo.center(mo.md(tab_marg))%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%23%20Kantorovich%20problem%20between%20arbitrary%20measures%0A%0A%20%20%20%20With%20%24P_%7B%5Cmathcal%7BX%7D%7D(x%2Cy)%3Dx%24%2C%20%24P_%7B%5Cmathcal%7BY%7D%7D(x%2Cy)%3Dy%24%20as%20projections%3A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathbf%7BU%7D(%5Calpha%2C%5Cbeta)%20%5Ccoloneqq%20%5Cbig%5C%7B%5Cpi%5Cin%5Cmathcal%7BM%7D_%2B%5E1(%5Cmathcal%7BX%7D%5Ctimes%5Cmathcal%7BY%7D)%20%3A%20P_%7B%5Cmathcal%7BX%7D%5C%23%7D%5Cpi%3D%5Calpha%2C%5C%3B%20P_%7B%5Cmathcal%7BY%7D%5C%23%7D%5Cpi%3D%5Cbeta%5Cbig%5C%7D%2C%0A%20%20%20%20%24%24%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathrm%7BL%7D_c(%5Calpha%2C%5Cbeta)%20%5Ccoloneqq%20%5Cmin_%7B%5Cpi%5Cin%5Cmathbf%7BU%7D(%5Calpha%2C%5Cbeta)%7D%20%5Cint_%7B%5Cmathcal%7BX%7D%5Ctimes%5Cmathcal%7BY%7D%7D%20c(x%2Cy)%5C%2Cd%5Cpi(x%2Cy).%0A%20%20%20%20%24%24%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(kantorovich%2C%20md_table%2C%20mo%2C%20np%2C%20plt)%3A%0A%20%20%20%20_N%20%3D%2040%0A%20%20%20%20_g%20%3D%20np.linspace(0.2%2C%203.6%2C%20_N)%0A%20%20%20%20_a%20%3D%20np.exp(-((_g%20-%201.5)%20**%202)%20%2F%200.35)%0A%20%20%20%20_a%20%2F%3D%20_a.sum()%0A%20%20%20%20_b%20%3D%20np.exp(-((_g%20-%202.2)%20**%202)%20%2F%200.6)%0A%20%20%20%20_b%20%2F%3D%20_b.sum()%0A%20%20%20%20_C%20%3D%20(_g%5B%3A%2C%20None%5D%20-%20_g%5BNone%2C%20%3A%5D)%20**%202%0A%20%20%20%20_P%2C%20_L%2C%20_%20%3D%20kantorovich(_a%2C%20_b%2C%20_C)%0A%0A%20%20%20%20_Fa%2C%20_Fb%20%3D%20np.cumsum(_a)%2C%20np.cumsum(_b)%0A%20%20%20%20_T%20%3D%20np.interp(_Fa%20-%20_a%20%2F%202%2C%20_Fb%20-%20_b%20%2F%202%2C%20_g)%0A%20%20%20%20_L_quant%20%3D%20np.sum(_a%20*%20(_g%20-%20_T)%20**%202)%0A%20%20%20%20_nnz%20%3D%20int(np.sum(_P%20%3E%201e-10))%0A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(4.5%2C%204.5))%0A%20%20%20%20_ax.imshow(_P.T%2C%20origin%3D%22lower%22%2C%20extent%3D%5B_g%5B0%5D%2C%20_g%5B-1%5D%2C%20_g%5B0%5D%2C%20_g%5B-1%5D%5D%2C%20cmap%3D%22Greens%22)%0A%20%20%20%20_ax.plot(_g%2C%20_T%2C%20color%3D%22%232e7d5a%22%2C%20lw%3D1%2C%20ls%3D%22--%22%2C%20label%3Dr%22%24T%3DF_%5Cbeta%5E%7B-1%7D%5Ccirc%20F_%5Calpha%24%22)%0A%20%20%20%20_ax.set_xlabel(%22x%22)%0A%20%20%20%20_ax.set_ylabel(%22y%22)%0A%20%20%20%20_ax.legend(loc%3D%22upper%20left%22)%0A%20%20%20%20_ax.set_title(r%22%24%5Cpi%3D(%5Cmathrm%7BId%7D%2CT)_%5C%23%5Calpha%24%22)%0A%0A%20%20%20%20tab_kant_gen%20%3D%20md_table(%5B%0A%20%20%20%20%20%20%20%20(%22Grid%20size%22%2C%20rf%22%24%7B_N%7D%5Ctimes%7B_N%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathrm%7BL%7D_c(%5Calpha%2C%5Cbeta)%24%20(LP)%22%2C%20rf%22%24%7B_L%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Csum_i%20a_i%20%5Cbig(x_i%20-%20T(x_i)%5Cbig)%5E2%24%22%2C%20rf%22%24%7B_L_quant%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5C%23%5C%7B(i%2Cj)%20%3A%20%5Cmathbf%7BP%7D_%7Bij%7D%20%3E%200%5C%7D%24%22%2C%20rf%22%24%7B_nnz%7D%24%22)%2C%0A%20%20%20%20%5D)%0A%0A%20%20%20%20mo.center(mo.as_html(_fig))%0A%20%20%20%20return%20(tab_kant_gen%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20tab_kant_gen)%3A%0A%20%20%20%20mo.center(mo.md(tab_kant_gen))%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%23%20Monge%E2%80%93Kantorovich%20equivalence%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%20mo.md(r%22%22%22%0A%20%20%20%20**Theorem%20(Brenier).**%20Let%20%24%5Cmathcal%7BX%7D%3D%5Cmathcal%7BY%7D%3D%5Cmathbb%7BR%7D%5Ed%24%2C%20%24c(x%2Cy)%3D%5C%7Cx-y%5C%7C%5E2%24%2C%20and%20let%20%24%5Calpha%24%20have%20a%20density.%20Then%20there%20exists%20a%20unique%20optimal%20coupling%2C%20supported%20on%20the%20graph%20of%20a%20map%3A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cpi%5E%5Cstar%20%3D%20(%5Cmathrm%7BId%7D%2CT)_%7B%5C%23%7D%5Calpha%2C%20%5Cqquad%20T%20%3D%20%5Cnabla%5Cvarphi%2C%5Cquad%20%5Cvarphi%3A%5Cmathbb%7BR%7D%5Ed%5Cto%5Cmathbb%7BR%7D%20%5Ctext%7B%20convex%7D%2C%0A%20%20%20%20%24%24%0A%20%20%20%20and%20%24T%24%20is%20the%20unique%20gradient%20of%20a%20convex%20function%20with%20%24(%5Cnabla%5Cvarphi)_%7B%5C%23%7D%5Calpha%3D%5Cbeta%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20bmat%2C%0A%20%20%20%20linear_sum_assignment%2C%0A%20%20%20%20md_table%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20np%2C%0A%20%20%20%20pairwise_dist%2C%0A%20%20%20%20plt%2C%0A%20%20%20%20sqrtm%2C%0A)%3A%0A%20%20%20%20_ma%20%3D%20np.array(%5B0.0%2C%200.0%5D)%0A%20%20%20%20_Sa%20%3D%20np.array(%5B%5B1.0%2C%200.3%5D%2C%20%5B0.3%2C%200.5%5D%5D)%0A%20%20%20%20_mb%20%3D%20np.array(%5B3.0%2C%201.0%5D)%0A%20%20%20%20_Sb%20%3D%20np.array(%5B%5B0.4%2C%20-0.2%5D%2C%20%5B-0.2%2C%201.2%5D%5D)%0A%20%20%20%20_Sa_h%20%3D%20np.real(sqrtm(_Sa))%0A%20%20%20%20_Sa_ih%20%3D%20np.linalg.inv(_Sa_h)%0A%20%20%20%20_M%20%3D%20np.real(sqrtm(_Sa_h%20%40%20_Sb%20%40%20_Sa_h))%0A%20%20%20%20_A%20%3D%20_Sa_ih%20%40%20_M%20%40%20_Sa_ih%0A%20%20%20%20_T%20%3D%20lambda%20x%3A%20_mb%20%2B%20(x%20-%20_ma)%20%40%20_A.T%0A%20%20%20%20_W2sq%20%3D%20np.sum((_ma%20-%20_mb)%20**%202)%20%2B%20np.trace(_Sa%20%2B%20_Sb%20-%202%20*%20_M)%0A%0A%20%20%20%20_rng%20%3D%20np.random.default_rng(5)%0A%20%20%20%20_xs%20%3D%20_rng.multivariate_normal(_ma%2C%20_Sa%2C%20100_000)%0A%20%20%20%20_tx%20%3D%20_T(_xs)%0A%0A%20%20%20%20_k%20%3D%20300%0A%20%20%20%20_xa%20%3D%20_rng.multivariate_normal(_ma%2C%20_Sa%2C%20_k)%0A%20%20%20%20_yb%20%3D%20_rng.multivariate_normal(_mb%2C%20_Sb%2C%20_k)%0A%20%20%20%20_%2C%20_s%20%3D%20linear_sum_assignment(pairwise_dist(_xa%2C%20_yb)%20**%202)%0A%20%20%20%20_W2sq_emp%20%3D%20np.mean(np.sum((_xa%20-%20_yb%5B_s%5D)%20**%202%2C%20axis%3D1))%0A%0A%20%20%20%20_show%20%3D%20_rng.choice(_k%2C%20size%3D25%2C%20replace%3DFalse)%0A%20%20%20%20_fig%2C%20_axs%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(11%2C%204)%2C%20sharex%3DTrue%2C%20sharey%3DTrue)%0A%20%20%20%20for%20_ax%20in%20_axs%3A%0A%20%20%20%20%20%20%20%20_ax.scatter(*_xa.T%2C%20s%3D8%2C%20color%3D%22%231f4e79%22%2C%20alpha%3D0.2%2C%20lw%3D0)%0A%20%20%20%20%20%20%20%20_ax.scatter(*_yb.T%2C%20s%3D8%2C%20color%3D%22%23d66e1e%22%2C%20alpha%3D0.2%2C%20lw%3D0)%0A%20%20%20%20%20%20%20%20_ax.scatter(*_xa%5B_show%5D.T%2C%20s%3D18%2C%20color%3D%22%231f4e79%22%2C%20zorder%3D3)%0A%20%20%20%20%20%20%20%20_ax.set_aspect(%22equal%22)%0A%20%20%20%20for%20_i%20in%20_show%3A%0A%20%20%20%20%20%20%20%20_axs%5B0%5D.annotate(%22%22%2C%20xy%3D_yb%5B_s%5B_i%5D%5D%2C%20xytext%3D_xa%5B_i%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20arrowprops%3Ddict(arrowstyle%3D%22-%7C%3E%22%2C%20color%3D%22gray%22%2C%20lw%3D0.8))%0A%20%20%20%20%20%20%20%20_axs%5B1%5D.annotate(%22%22%2C%20xy%3D_T(_xa%5B_i%5D)%2C%20xytext%3D_xa%5B_i%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20arrowprops%3Ddict(arrowstyle%3D%22-%7C%3E%22%2C%20color%3D%22%232e7d5a%22%2C%20lw%3D0.8))%0A%20%20%20%20_axs%5B0%5D.scatter(*_yb%5B_s%5B_show%5D%5D.T%2C%20s%3D18%2C%20color%3D%22%23d66e1e%22%2C%20zorder%3D3)%0A%20%20%20%20_axs%5B1%5D.scatter(*_T(_xa%5B_show%5D).T%2C%20s%3D18%2C%20color%3D%22%232e7d5a%22%2C%20zorder%3D3)%0A%20%20%20%20_axs%5B0%5D.set_title(rf%22Optimal%20assignment%20between%20%24%7B_k%7D%24%20samples%22)%0A%20%20%20%20_axs%5B1%5D.set_title(r%22Brenier%20map%20%24T(x)%3Dm_%5Cbeta%2BA(x-m_%5Calpha)%24%22)%0A%20%20%20%20_fig.tight_layout()%0A%0A%20%20%20%20_eig%20%3D%20np.linalg.eigvalsh(_A)%0A%20%20%20%20_cov_T%20%3D%20np.cov(_tx.T)%0A%20%20%20%20_mc%20%3D%20np.mean(np.sum((_xs%20-%20_tx)%20**%202%2C%20axis%3D1))%0A%20%20%20%20_sym%20%3D%20bool(np.allclose(_A%2C%20_A.T))%0A%0A%20%20%20%20tab_brenier%20%3D%20md_table(%5B%0A%20%20%20%20%20%20%20%20(r%22%24A%24%2C%20with%20%24T(x)%20%3D%20m_%5Cbeta%20%2B%20A(x%20-%20m_%5Calpha)%24%22%2C%20rf%22%24%7Bbmat(_A)%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24A%20%3D%20A%5E%5Ctop%24%2C%20eigenvalues%20of%20%24A%24%22%2C%20rf%22%24%5Ctext%7B%7B%7B_sym%7D%7D%7D%2C%5C%20%7Bbmat(_eig)%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Coperatorname%7BCov%7D%5Cbig(T(X)%5Cbig)%24%20vs%20%24%5CSigma_%5Cbeta%24%22%2C%20rf%22%24%7Bbmat(_cov_T%2C%20'%7B%3A.3f%7D')%7D%24%20vs%20%24%7Bbmat(_Sb%2C%20'%7B%3A.3f%7D')%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathcal%7BW%7D_2%5E2%24%20(closed%20form)%22%2C%20rf%22%24%7B_W2sq%3A.4f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathbb%7BE%7D%5C%2C%5ClVert%20X%20-%20T(X)%20%5CrVert%5E2%24%20(Monte%20Carlo)%22%2C%20rf%22%24%7B_mc%3A.4f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(rf%22%24%5Cmathcal%7B%7BW%7D%7D_2%5E2%24%20(assignment%20on%20%24%7B_k%7D%24%20samples)%22%2C%20rf%22%24%7B_W2sq_emp%3A.4f%7D%24%22)%2C%0A%20%20%20%20%5D)%0A%0A%20%20%20%20mo.center(mo.as_html(_fig))%0A%20%20%20%20return%20(tab_brenier%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20tab_brenier)%3A%0A%20%20%20%20mo.center(mo.md(tab_brenier))%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%204.%20Wasserstein%20distance%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%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Metric%20Properties%20of%20Optimal%20Transport%0A%20%20%20%20**Proposition.**%20Let%20%24n%3Dm%24%2C%20%24p%5Cge%201%24%2C%20and%20%24%5Cmathbf%7BC%7D%20%3D%20%5Cmathbf%7BD%7D%5Ep%20%3D%20(%5Cmathbf%7BD%7D%5Ep_%7Bi%2Cj%7D)_%7Bi%2Cj%7D%5Cin%5Cmathbb%7BR%7D%5E%7Bn%5Ctimes%20n%7D%24%20(entrywise%20power)%2C%20where%20%24%5Cmathbf%7BD%7D%5Cin%5Cmathbb%7BR%7D%5E%7Bn%5Ctimes%20n%7D_%2B%24%20is%20a%20distance%20on%20%24%5Cllbracket%20n%20%5Crrbracket%24%2C%20i.e.%0A%0A%20%20%20%201.%20%24%5Cmathbf%7BD%7D%24%20is%20symmetric%3B%0A%20%20%20%202.%20%24%5Cmathbf%7BD%7D_%7Bi%2Cj%7D%3D0%24%20if%20and%20only%20if%20%24i%3Dj%24%3B%0A%20%20%20%203.%20%24%5Cforall%5C%2C(i%2Cj%2Ck)%5Cin%5Cllbracket%20n%5Crrbracket%5E3%2C%5C%3B%20%5Cmathbf%7BD%7D_%7Bi%2Ck%7D%5Cle%20%5Cmathbf%7BD%7D_%7Bi%2Cj%7D%2B%5Cmathbf%7BD%7D_%7Bj%2Ck%7D%24.%0A%0A%20%20%20%20Then%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathcal%7BW%7D_p(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%20%5Ccoloneqq%20%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BD%7D%5Ep%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%5E%7B1%2Fp%7D%0A%20%20%20%20%24%24%0A%20%20%20%20defines%20the%20*%24p%24-Wasserstein%20distance*%20on%20%24%5CSigma_n%24%3A%20%24%5Cmathcal%7BW%7D_p%24%20is%20symmetric%2C%20positive%2C%20%24%5Cmathcal%7BW%7D_p(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%3D0%20%5Ciff%20%5Cmathbf%7Ba%7D%3D%5Cmathbf%7Bb%7D%24%2C%20and%20satisfies%20the%20triangle%20inequality%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cforall%5C%2C%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D%2C%5Cmathbf%7Bc%7D%5Cin%5CSigma_n%2C%5Cqquad%20%5Cmathcal%7BW%7D_p(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bc%7D)%20%5Cle%20%5Cmathcal%7BW%7D_p(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%20%2B%20%5Cmathcal%7BW%7D_p(%5Cmathbf%7Bb%7D%2C%5Cmathbf%7Bc%7D).%0A%20%20%20%20%24%24%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%20p_sel%20%3D%20mo.ui.dropdown(%5B%221%22%2C%20%222%22%2C%20%223%22%5D%2C%20value%3D%222%22%2C%20label%3D%22p%22)%0A%20%20%20%20n_triples%20%3D%20mo.ui.slider(10%2C%20200%2C%20value%3D100%2C%20step%3D10%2C%20label%3D%22random%20triples%22)%0A%20%20%20%20mo.center(mo.hstack(%5Bp_sel%2C%20n_triples%5D))%0A%20%20%20%20return%20n_triples%2C%20p_sel%0A%0A%0A%40app.cell%0Adef%20_(kantorovich%2C%20md_table%2C%20mo%2C%20n_triples%2C%20np%2C%20p_sel%2C%20pairwise_dist%2C%20sci)%3A%0A%20%20%20%20_p%20%3D%20int(p_sel.value)%0A%20%20%20%20_rng%20%3D%20np.random.default_rng(6)%0A%20%20%20%20_n%20%3D%206%0A%20%20%20%20_Z%20%3D%20_rng.random((_n%2C%202))%0A%20%20%20%20_D%20%3D%20pairwise_dist(_Z%2C%20_Z)%0A%0A%20%20%20%20def%20_W(u%2C%20v)%3A%0A%20%20%20%20%20%20%20%20return%20max(kantorovich(u%2C%20v%2C%20_D**_p)%5B1%5D%2C%200.0)%20**%20(1%20%2F%20_p)%0A%0A%20%20%20%20def%20_hist()%3A%0A%20%20%20%20%20%20%20%20_h%20%3D%20_rng.random(_n)%20**%203%0A%20%20%20%20%20%20%20%20return%20_h%20%2F%20_h.sum()%0A%0A%20%20%20%20_tri%2C%20_sym%2C%20_self%2C%20_pos%20%3D%20%5B%5D%2C%20%5B%5D%2C%20%5B%5D%2C%20%5B%5D%0A%20%20%20%20for%20_%20in%20range(n_triples.value)%3A%0A%20%20%20%20%20%20%20%20_a%2C%20_b%2C%20_c%20%3D%20_hist()%2C%20_hist()%2C%20_hist()%0A%20%20%20%20%20%20%20%20_wab%2C%20_wba%2C%20_wbc%2C%20_wac%20%3D%20_W(_a%2C%20_b)%2C%20_W(_b%2C%20_a)%2C%20_W(_b%2C%20_c)%2C%20_W(_a%2C%20_c)%0A%20%20%20%20%20%20%20%20_tri.append(_wac%20-%20_wab%20-%20_wbc)%0A%20%20%20%20%20%20%20%20_sym.append(abs(_wab%20-%20_wba))%0A%20%20%20%20%20%20%20%20_self.append(_W(_a%2C%20_a))%0A%20%20%20%20%20%20%20%20_pos.append(_wab)%0A%0A%20%20%20%20_x%20%3D%20np.array(%5B%5B0.0%5D%2C%20%5B1.0%5D%2C%20%5B2.0%5D%5D)%0A%20%20%20%20_D1%20%3D%20pairwise_dist(_x%2C%20_x)%0A%20%20%20%20_e%20%3D%20np.eye(3)%0A%20%20%20%20_W2%20%3D%20lambda%20u%2C%20v%3A%20kantorovich(u%2C%20v%2C%20_D1**2)%5B1%5D%20**%200.5%0A%0A%20%20%20%20_rows1%20%3D%20%5B%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmax%5C%2C%5Cbig%5B%5Cmathcal%7BW%7D_p(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bc%7D)%20-%20%5Cmathcal%7BW%7D_p(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%20-%20%5Cmathcal%7BW%7D_p(%5Cmathbf%7Bb%7D%2C%5Cmathbf%7Bc%7D)%5Cbig%5D%24%22%2C%20rf%22%24%7Bsci(max(_tri))%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmax%5C%2C%5Cbig%5Clvert%20%5Cmathcal%7BW%7D_p(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%20-%20%5Cmathcal%7BW%7D_p(%5Cmathbf%7Bb%7D%2C%5Cmathbf%7Ba%7D)%20%5Cbig%5Crvert%24%22%2C%20rf%22%24%7Bsci(max(_sym))%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmax%5C%2C%5Cmathcal%7BW%7D_p(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Ba%7D)%24%22%2C%20rf%22%24%7Bsci(max(_self))%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmin%5C%2C%5Cmathcal%7BW%7D_p(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%24%2C%20%24%5Cmathbf%7Ba%7D%5Cneq%5Cmathbf%7Bb%7D%24%22%2C%20rf%22%24%7Bmin(_pos)%3A.4f%7D%24%22)%2C%0A%20%20%20%20%5D%0A%20%20%20%20_rows2%20%3D%20%5B%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathbf%7BD%7D%5E2_%7B02%7D%24%20vs%20%24%5Cmathbf%7BD%7D%5E2_%7B01%7D%20%2B%20%5Cmathbf%7BD%7D%5E2_%7B12%7D%24%22%2C%20rf%22%24%7B_D1%5B0%2C%202%5D%20**%202%3A.1f%7D%24%20vs%20%24%7B_D1%5B0%2C%201%5D%20**%202%20%2B%20_D1%5B1%2C%202%5D%20**%202%3A.1f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathcal%7BW%7D_2(%5Cdelta_0%2C%5Cdelta_2)%24%20vs%20%24%5Cmathcal%7BW%7D_2(%5Cdelta_0%2C%5Cdelta_1)%20%2B%20%5Cmathcal%7BW%7D_2(%5Cdelta_1%2C%5Cdelta_2)%24%22%2C%20rf%22%24%7B_W2(_e%5B0%5D%2C%20_e%5B2%5D)%3A.4f%7D%24%20vs%20%24%7B_W2(_e%5B0%5D%2C%20_e%5B1%5D)%20%2B%20_W2(_e%5B1%5D%2C%20_e%5B2%5D)%3A.4f%7D%24%22)%2C%0A%20%20%20%20%5D%0A%20%20%20%20mo.vstack(%5B%0A%20%20%20%20%20%20%20%20mo.center(mo.md(md_table(_rows1%2C%20(rf%22Check%20(%24p%3D%7B_p%7D%24%2C%20%24n%3D%7B_n%7D%24)%22%2C%20%22Value%22))))%2C%0A%20%20%20%20%20%20%20%20mo.md(r%22Points%20%240%2C%201%2C%202%24%20on%20%24%5Cmathbb%7BR%7D%24%2C%20%24%5Cmathbf%7BD%7D_%7Bij%7D%20%3D%20%7Cx_i%20-%20x_j%7C%24%3A%22)%2C%0A%20%20%20%20%20%20%20%20mo.center(mo.md(md_table(_rows2)))%2C%0A%20%20%20%20%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%23%20Measure%20Metric%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%20mo.md(r%22%22%22%0A%20%20%20%20**Divergences**%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathrm%7BKL%7D(%5Calpha%5C%2C%5C%7C%5C%2C%5Cbeta)%20%3D%20%5Cint_%7B%5Cmathcal%7BX%7D%7D%20%5Clog%5CBig(%5Cfrac%7B%5Cmathrm%7Bd%7D%5Calpha%7D%7B%5Cmathrm%7Bd%7D%5Cbeta%7D%5CBig)%5C%2C%5Cmathrm%7Bd%7D%5Calpha%0A%20%20%20%20%24%24%0A%20%20%20%20%24%24%0A%20%20%20%20%5C%7C%5Calpha-%5Cbeta%5C%7C_%7BTV%7D%20%3D%20%5Cint_%7B%5Cmathcal%7BX%7D%7D%20%7C%5Cmathrm%7Bd%7D%5Calpha-%5Cmathrm%7Bd%7D%5Cbeta%7C%0A%20%20%20%20%24%24%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathcal%7BW%7D_2(%5Calpha%2C%5Cbeta)%20%3D%20%5CBig(%5Cmin_%7B%5Cpi%5Cin%5Cmathcal%20U(%5Calpha%2C%5Cbeta)%7D%20%5Cint_%7B%5Cmathcal%7BX%7D%5Ctimes%5Cmathcal%7BX%7D%7D%5C%7Cx-y%5C%7C%5E2%5C%2C%5Cmathrm%7Bd%7D%5Cpi(x%2Cy)%5CBig)%5E%7B1%2F2%7D%0A%20%20%20%20%24%24%0A%20%20%20%20%24%5Calpha%3D%5Cmathcal%20N(%5Cmu%2C%5Csigma%5E2)%24%2C%20%24%5Cbeta%3D%5Cmathcal%20N(%5Cmu%2B%5Ctau%2C%5Csigma%5E2)%24%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%20tau%20%3D%20mo.ui.slider(0.0%2C%204.0%2C%20step%3D0.1%2C%20value%3D2.2%2C%20label%3D%22%CF%84%22)%0A%20%20%20%20t_interp%20%3D%20mo.ui.slider(0.0%2C%201.0%2C%20step%3D0.05%2C%20value%3D0.5%2C%20label%3D%22t%22)%0A%20%20%20%20mo.center(mo.hstack(%5Btau%2C%20t_interp%5D))%0A%20%20%20%20return%20t_interp%2C%20tau%0A%0A%0A%40app.cell%0Adef%20_(md_table%2C%20mo%2C%20norm%2C%20np%2C%20plt%2C%20t_interp%2C%20tau)%3A%0A%20%20%20%20_mu%2C%20_s%20%3D%201.1%2C%200.25%0A%20%20%20%20_tau%2C%20_t%20%3D%20tau.value%2C%20t_interp.value%0A%20%20%20%20_x%20%3D%20np.linspace(-1.0%2C%206.5%2C%205000)%0A%20%20%20%20_dx%20%3D%20_x%5B1%5D%20-%20_x%5B0%5D%0A%20%20%20%20_pa%20%3D%20norm.pdf(_x%2C%20_mu%2C%20_s)%0A%20%20%20%20_pb%20%3D%20norm.pdf(_x%2C%20_mu%20%2B%20_tau%2C%20_s)%0A%0A%20%20%20%20_kl%20%3D%20np.sum(_pa%20*%20(norm.logpdf(_x%2C%20_mu%2C%20_s)%20-%20norm.logpdf(_x%2C%20_mu%20%2B%20_tau%2C%20_s)))%20*%20_dx%0A%20%20%20%20_tv%20%3D%20np.sum(np.abs(_pa%20-%20_pb))%20*%20_dx%0A%20%20%20%20_q%20%3D%20np.linspace(1e-5%2C%201%20-%201e-5%2C%2020000)%0A%20%20%20%20_w2%20%3D%20np.sqrt(np.mean((norm.ppf(_q%2C%20_mu%2C%20_s)%20-%20norm.ppf(_q%2C%20_mu%20%2B%20_tau%2C%20_s))%20**%202))%0A%0A%20%20%20%20_kl_ex%20%3D%20_tau**2%20%2F%20(2%20*%20_s**2)%0A%20%20%20%20_tv_ex%20%3D%202%20*%20(2%20*%20norm.cdf(_tau%20%2F%20(2%20*%20_s))%20-%201)%0A%0A%20%20%20%20_taus%20%3D%20np.linspace(0%2C%204%2C%20200)%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(1%2C%203%2C%20figsize%3D(13%2C%203.2))%0A%20%20%20%20_ax%5B0%5D.plot(_x%2C%20_pa%2C%20color%3D%22%231f4e79%22)%0A%20%20%20%20_ax%5B0%5D.plot(_x%2C%20_pb%2C%20color%3D%22%23d66e1e%22)%0A%20%20%20%20_ax%5B0%5D.plot(_x%2C%20(1%20-%20_t)%20*%20_pa%20%2B%20_t%20*%20_pb%2C%20color%3D%22%232e7d5a%22%2C%20ls%3D%22--%22)%0A%20%20%20%20_ax%5B0%5D.set_title(r%22%24(1-t)%5Calpha%2Bt%5Cbeta%24%22)%0A%20%20%20%20_ax%5B1%5D.plot(_x%2C%20_pa%2C%20color%3D%22%231f4e79%22)%0A%20%20%20%20_ax%5B1%5D.plot(_x%2C%20_pb%2C%20color%3D%22%23d66e1e%22)%0A%20%20%20%20_ax%5B1%5D.plot(_x%2C%20norm.pdf(_x%2C%20_mu%20%2B%20_t%20*%20_tau%2C%20_s)%2C%20color%3D%22%232e7d5a%22%2C%20ls%3D%22--%22)%0A%20%20%20%20_ax%5B1%5D.set_title(r%22%24%5Cmathcal%7BW%7D_2%24%20geodesic%3A%20%24%5Cmathcal%7BN%7D(%5Cmu%2Bt%5Ctau%2C%5Csigma%5E2)%24%22)%0A%20%20%20%20_ax%5B2%5D.plot(_taus%2C%20_taus**2%20%2F%20(2%20*%20_s**2)%2C%20label%3D%22KL%22)%0A%20%20%20%20_ax%5B2%5D.plot(_taus%2C%202%20*%20(2%20*%20norm.cdf(_taus%20%2F%20(2%20*%20_s))%20-%201)%2C%20label%3D%22TV%22)%0A%20%20%20%20_ax%5B2%5D.plot(_taus%2C%20_taus%2C%20label%3Dr%22%24%5Cmathcal%7BW%7D_2%24%22)%0A%20%20%20%20_ax%5B2%5D.axvline(_tau%2C%20color%3D%22gray%22%2C%20ls%3D%22%3A%22)%0A%20%20%20%20_ax%5B2%5D.set_ylim(0%2C%205)%0A%20%20%20%20_ax%5B2%5D.set_xlabel(r%22%24%5Ctau%24%22)%0A%20%20%20%20_ax%5B2%5D.legend()%0A%20%20%20%20_fig.tight_layout()%0A%0A%20%20%20%20tab_div%20%3D%20md_table(%5B%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathrm%7BKL%7D(%5Calpha%5C%2C%5CVert%5C%2C%5Cbeta)%24%22%2C%20rf%22%24%7B_kl%3A.4f%7D%24%22%2C%20rf%22%24%5Cfrac%7B%7B%5Ctau%5E2%7D%7D%7B%7B2%5Csigma%5E2%7D%7D%20%3D%20%7B_kl_ex%3A.4f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5ClVert%20%5Calpha%20-%20%5Cbeta%20%5CrVert_%7BTV%7D%24%22%2C%20rf%22%24%7B_tv%3A.4f%7D%24%22%2C%20rf%22%242%5Cbig(2%5CPhi(%5Ctfrac%7B%7B%5Ctau%7D%7D%7B%7B2%5Csigma%7D%7D)-1%5Cbig)%20%3D%20%7B_tv_ex%3A.4f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathcal%7BW%7D_2(%5Calpha%2C%5Cbeta)%24%22%2C%20rf%22%24%7B_w2%3A.4f%7D%24%22%2C%20rf%22%24%5Ctau%20%3D%20%7B_tau%3A.4f%7D%24%22)%2C%0A%20%20%20%20%5D%2C%20(rf%22%24%5Csigma%3D%7B_s%3Ag%7D%2C%5C%20%5Ctau%3D%7B_tau%3Ag%7D%24%22%2C%20%22Numerical%22%2C%20%22Closed%20form%22))%0A%0A%20%20%20%20mo.center(mo.as_html(_fig))%0A%20%20%20%20return%20(tab_div%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20tab_div)%3A%0A%20%20%20%20mo.center(mo.md(tab_div))%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%205.%20Dual%20Problem%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%20mo.vstack(%5B%0A%20%20%20%20%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Kantorovich%20Duality%20(discrete%20case)%0A%20%20%20%20**Proposition.**%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BC%7D%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%20%5C%3B%3D%5C%3B%20%5Cmax_%7B(%5Cmathbf%7Bf%7D%2C%5Cmathbf%7Bg%7D)%5Cin%5Cmathbf%7BR%7D(%5Cmathbf%7BC%7D)%7D%20%5Clangle%5Cmathbf%7Bf%7D%2C%5C%2C%5Cmathbf%7Ba%7D%5Crangle%20%2B%20%5Clangle%5Cmathbf%7Bg%7D%2C%5C%2C%5Cmathbf%7Bb%7D%5Crangle%2C%20%5Cqquad%20%5Cmathbf%7BR%7D(%5Cmathbf%7BC%7D)%20%5Ccoloneqq%20%5Cbig%5C%7B(%5Cmathbf%7Bf%7D%2C%5Cmathbf%7Bg%7D)%5Cin%5Cmathbb%7BR%7D%5En%5Ctimes%5Cmathbb%7BR%7D%5Em%20%3A%20%5Cmathbf%7Bf%7D_i%2B%5Cmathbf%7Bg%7D_j%5Cle%5Cmathbf%7BC%7D_%7Bij%7D%5Cbig%5C%7D.%0A%20%20%20%20%24%24%0A%20%20%20%20%22%22%22)%2C%0A%20%20%20%20%20%20%20%20mo.callout(mo.md(r%22%22%22%0A%20%20%20%20**Example.**%20A%20shipper%20charges%20%24%5Cmathbf%7Bf%7D_i%24%20to%20pick%20up%20at%20%24x_i%24%20and%20%24%5Cmathbf%7Bg%7D_j%24%20to%20deliver%20at%20%24y_j%24.%20To%20be%20competitive%2C%20the%20prices%20must%20satisfy%20%24%5Cmathbf%7Bf%7D_i%20%2B%20%5Cmathbf%7Bg%7D_j%20%5Cle%20%5Cmathbf%7BC%7D_%7Bij%7D%24%3B%20maximizing%20their%20profit%2C%20the%20shipper%20exactly%20matches%20the%20optimal%20cost.%0A%20%20%20%20%22%22%22)%2C%20kind%3D%22success%22)%2C%0A%20%20%20%20%5D)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20C%2C%0A%20%20%20%20L_opt%2C%0A%20%20%20%20P_opt%2C%0A%20%20%20%20a%2C%0A%20%20%20%20b%2C%0A%20%20%20%20fmt_vec%2C%0A%20%20%20%20kantorovich%2C%0A%20%20%20%20kantorovich_dual%2C%0A%20%20%20%20md_table%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20n%2C%0A%20%20%20%20np%2C%0A%20%20%20%20sci%2C%0A)%3A%0A%20%20%20%20_f%2C%20_g%2C%20_L_dual%20%3D%20kantorovich_dual(a%2C%20b%2C%20C)%0A%20%20%20%20_S%20%3D%20C%20-%20_f%5B%3A%2C%20None%5D%20-%20_g%5BNone%2C%20%3A%5D%0A%20%20%20%20_supp%20%3D%20P_opt%20%3E%201e-10%0A%20%20%20%20_%2C%20_%2C%20_res%20%3D%20kantorovich(a%2C%20b%2C%20C)%0A%20%20%20%20_f_lp%2C%20_g_lp%20%3D%20_res.eqlin.marginals%5B%3An%5D%2C%20_res.eqlin.marginals%5Bn%3A%5D%0A%20%20%20%20_c%20%3D%201.7%0A%20%20%20%20_L_shift%20%3D%20np.dot(_f%20%2B%20_c%2C%20a)%20%2B%20np.dot(_g%20-%20_c%2C%20b)%0A%20%20%20%20_L_mult%20%3D%20np.dot(_f_lp%2C%20a)%20%2B%20np.dot(_g_lp%2C%20b)%0A%0A%20%20%20%20_rows%20%3D%20%5B%0A%20%20%20%20%20%20%20%20(r%22Primal%20%24%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BC%7D%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%24%22%2C%20rf%22%24%7BL_opt%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22Dual%20%24%5Clangle%5Cmathbf%7Bf%7D%2C%5Cmathbf%7Ba%7D%5Crangle%20%2B%20%5Clangle%5Cmathbf%7Bg%7D%2C%5Cmathbf%7Bb%7D%5Crangle%24%22%2C%20rf%22%24%7B_L_dual%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22Dual%20from%20LP%20multipliers%22%2C%20rf%22%24%7B_L_mult%3A.6f%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmin_%7Bij%7D%5C%2C(%5Cmathbf%7BC%7D_%7Bij%7D%20-%20%5Cmathbf%7Bf%7D_i%20-%20%5Cmathbf%7Bg%7D_j)%20%5Cge%200%24%22%2C%20rf%22%24%7Bsci(_S.min())%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmax_%7B%5Cmathbf%7BP%7D_%7Bij%7D%3E0%7D%5C%2C%5Cbig%5Clvert%20%5Cmathbf%7BC%7D_%7Bij%7D%20-%20%5Cmathbf%7Bf%7D_i%20-%20%5Cmathbf%7Bg%7D_j%20%5Cbig%5Crvert%24%22%2C%20rf%22%24%7Bsci(np.abs(_S%5B_supp%5D).max())%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(rf%22%24%5Clangle%5Cmathbf%7B%7Bf%7D%7D%2Bc%2C%5Cmathbf%7B%7Ba%7D%7D%5Crangle%20%2B%20%5Clangle%5Cmathbf%7B%7Bg%7D%7D-c%2C%5Cmathbf%7B%7Bb%7D%7D%5Crangle%24%2C%20%24c%3D%7B_c%3Ag%7D%24%22%2C%20rf%22%24%7B_L_shift%3A.6f%7D%24%22)%2C%0A%20%20%20%20%5D%0A%20%20%20%20_f_tex%20%3D%20fmt_vec(_f%2C%20%22%7B%3A.4f%7D%22)%0A%20%20%20%20_g_tex%20%3D%20fmt_vec(_g%2C%20%22%7B%3A.4f%7D%22)%0A%0A%20%20%20%20mo.vstack(%5B%0A%20%20%20%20%20%20%20%20mo.md(rf%22%22%22%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cbegin%7B%7Baligned%7D%7D%0A%20%20%20%20%5Cmathbf%7B%7Bf%7D%7D%20%26%3D%20%7B_f_tex%7D%20%26%26%20%5Ctext%7B%7B(pick-up%20prices)%7D%7D%20%5C%5C%0A%20%20%20%20%5Cmathbf%7B%7Bg%7D%7D%20%26%3D%20%7B_g_tex%7D%20%26%26%20%5Ctext%7B%7B(delivery%20prices)%7D%7D%0A%20%20%20%20%5Cend%7B%7Baligned%7D%7D%0A%20%20%20%20%24%24%0A%20%20%20%20%22%22%22)%2C%0A%20%20%20%20%20%20%20%20mo.center(mo.md(md_table(_rows%2C%20(%22Check%22%2C%20%22Value%22))))%2C%0A%20%20%20%20%5D%2C%20align%3D%22center%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%206.%20The%20Kantorovich%20Linear%20Programs%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%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Primal%20in%20standard%20form%0A%0A%20%20%20%20We%20vectorize%20column-wise%3A%20%24%5Cmathbf%7Bp%7D%3D%5Coperatorname%7Bvec%7D(%5Cmathbf%7BP%7D)%5Cin%5Cmathbb%7BR%7D%5E%7Bnm%7D%24%2C%20%24%5Cmathbf%7Bc%7D%3D%5Coperatorname%7Bvec%7D(%5Cmathbf%7BC%7D)%24.%20Both%20marginal%20constraints%20are%20linear%20in%20%24%5Cmathbf%7Bp%7D%24%3A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathbf%7BP%7D%5Cmathbf%7B1%7D_m%20%3D%20(%5Cmathbf%7B1%7D_m%5E%5Ctop%5Cotimes%5Cmathbf%7BI%7D_n)%5C%2C%5Cmathbf%7Bp%7D%2C%20%5Cqquad%20%5Cmathbf%7BP%7D%5E%5Ctop%5Cmathbf%7B1%7D_n%20%3D%20(%5Cmathbf%7BI%7D_m%5Cotimes%5Cmathbf%7B1%7D_n%5E%5Ctop)%5C%2C%5Cmathbf%7Bp%7D.%0A%20%20%20%20%24%24%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%20mo.md(r%22%22%22%0A%20%20%20%20**Primal%20in%20standard%20form**%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BC%7D%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%20%3D%20%5Cmin_%7B%5Cmathbf%7Bp%7D%5Cin%5Cmathbb%7BR%7D%5E%7Bnm%7D_%2B%7D%5C%20%5Cmathbf%7Bc%7D%5E%5Ctop%5Cmathbf%7Bp%7D%20%5Cquad%5Ctext%7Bs.t.%7D%5Cquad%20%5Cmathbf%7BA%7D%5Cmathbf%7Bp%7D%20%3D%20%5Cbegin%7Bbmatrix%7D%5Cmathbf%7Ba%7D%5C%5C%20%5Cmathbf%7Bb%7D%5Cend%7Bbmatrix%7D%2C%20%5Cqquad%20%5Cmathbf%7BA%7D%20%3D%20%5Cbegin%7Bbmatrix%7D%5Cmathbf%7B1%7D_m%5E%5Ctop%5Cotimes%5Cmathbf%7BI%7D_n%5C%5C%20%5Cmathbf%7BI%7D_m%5Cotimes%5Cmathbf%7B1%7D_n%5E%5Ctop%5Cend%7Bbmatrix%7D%20%5Cin%5Cmathbb%7BR%7D%5E%7B(n%2Bm)%5Ctimes%20nm%7D.%0A%20%20%20%20%24%24%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20C%2C%0A%20%20%20%20L_opt%2C%0A%20%20%20%20P_opt%2C%0A%20%20%20%20a%2C%0A%20%20%20%20b%2C%0A%20%20%20%20bmat%2C%0A%20%20%20%20m%2C%0A%20%20%20%20marginal_matrix%2C%0A%20%20%20%20md_table%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20n%2C%0A%20%20%20%20np%2C%0A%20%20%20%20sci%2C%0A)%3A%0A%20%20%20%20_A2%20%3D%20marginal_matrix(2%2C%203)%0A%20%20%20%20_A%20%3D%20marginal_matrix(n%2C%20m)%0A%20%20%20%20_p%20%3D%20P_opt.flatten(order%3D%22F%22)%0A%20%20%20%20_c%20%3D%20C.flatten(order%3D%22F%22)%0A%20%20%20%20_lhs1%20%3D%20np.kron(np.ones((1%2C%20m))%2C%20np.eye(n))%20%40%20_p%0A%20%20%20%20_lhs2%20%3D%20np.kron(np.eye(m)%2C%20np.ones((1%2C%20n)))%20%40%20_p%0A%20%20%20%20_A2_tex%20%3D%20bmat(_A2%2C%20%22%7B%3A.0f%7D%22)%0A%0A%20%20%20%20_rows%20%3D%20%5B%0A%20%20%20%20%20%20%20%20(r%22Shape%20of%20%24%5Cmathbf%7BA%7D%24%22%2C%20rf%22%24%7B_A.shape%5B0%5D%7D%5Ctimes%7B_A.shape%5B1%5D%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Coperatorname%7Brank%7D%5Cmathbf%7BA%7D%2C%5C%20n%2Bm-1%24%22%2C%20rf%22%24%7Bnp.linalg.matrix_rank(_A)%7D%2C%5C%20%7Bn%20%2B%20m%20-%201%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5ClVert%20(%5Cmathbf%7B1%7D_m%5E%5Ctop%5Cotimes%5Cmathbf%7BI%7D_n)%5C%2C%5Cmathbf%7Bp%7D%20-%20%5Cmathbf%7BP%7D%5Cmathbf%7B1%7D_m%20%5CrVert%24%22%2C%20rf%22%24%7Bsci(np.linalg.norm(_lhs1%20-%20P_opt.sum(1)))%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5ClVert%20(%5Cmathbf%7BI%7D_m%5Cotimes%5Cmathbf%7B1%7D_n%5E%5Ctop)%5C%2C%5Cmathbf%7Bp%7D%20-%20%5Cmathbf%7BP%7D%5E%5Ctop%5Cmathbf%7B1%7D_n%20%5CrVert%24%22%2C%20rf%22%24%7Bsci(np.linalg.norm(_lhs2%20-%20P_opt.sum(0)))%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5ClVert%20%5Cmathbf%7BA%7D%5Cmathbf%7Bp%7D%20-%20%5B%5Cmathbf%7Ba%7D%3B%5Cmathbf%7Bb%7D%5D%20%5CrVert%24%22%2C%20rf%22%24%7Bsci(np.linalg.norm(_A%20%40%20_p%20-%20np.concatenate(%5Ba%2C%20b%5D)))%7D%24%22)%2C%0A%20%20%20%20%20%20%20%20(r%22%24%5Cmathbf%7Bc%7D%5E%5Ctop%5Cmathbf%7Bp%7D%24%20vs%20%24%5Cmathrm%7BL%7D_%7B%5Cmathbf%7BC%7D%7D(%5Cmathbf%7Ba%7D%2C%5Cmathbf%7Bb%7D)%24%22%2C%20rf%22%24%7B_c%20%40%20_p%3A.6f%7D%24%20vs%20%24%7BL_opt%3A.6f%7D%24%22)%2C%0A%20%20%20%20%5D%0A%20%20%20%20mo.vstack(%5B%0A%20%20%20%20%20%20%20%20mo.md(rf%22%24%24%20n%3D2%2C%5C%20m%3D3%3A%20%5Cqquad%20%5Cmathbf%7B%7BA%7D%7D%20%3D%20%7B_A2_tex%7D%20%24%24%22)%2C%0A%20%20%20%20%20%20%20%20mo.center(mo.md(md_table(_rows%2C%20(rf%22Check%20(%24n%3D%7Bn%7D%24%2C%20%24m%3D%7Bm%7D%24)%22%2C%20%22Value%22))))%2C%0A%20%20%20%20%5D%2C%20align%3D%22center%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
5afa4947c43944cc43d162883b7c4bf6