wasserstein sampling

import%20marimo%0A%0A__generated_with%20%3D%20%220.20.4%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.gridspec%20as%20gridspec%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20scipy.stats%20as%20stats%0A%0A%20%20%20%20return%20gridspec%2C%20mo%2C%20np%2C%20plt%2C%20stats%0A%0A%0A%40app.cell%0Adef%20_(np)%3A%0A%20%20%20%20rng%20%3D%20np.random.default_rng(42)%0A%20%20%20%20return%20(rng%2C)%0A%0A%0A%40app.cell%0Adef%20_(rng)%3A%0A%20%20%20%20N%20%3D%2060%0A%20%20%20%20mu_true%20%3D%205.0%0A%20%20%20%20sigma_true%20%3D%201.0%0A%20%20%20%20historical%20%3D%20rng.normal(mu_true%2C%20sigma_true%2C%20size%3DN)%0A%20%20%20%20return%20N%2C%20historical%2C%20mu_true%2C%20sigma_true%0A%0A%0A%40app.cell%0Adef%20_(historical%2C%20np)%3A%0A%20%20%20%20empirical_sorted%20%3D%20np.sort(historical)%0A%20%20%20%20return%20(empirical_sorted%2C)%0A%0A%0A%40app.cell%0Adef%20_(np)%3A%0A%20%20%20%20def%20w1_sorted(x_sorted%2C%20y_sorted)%3A%0A%20%20%20%20%20%20%20%20%22%22%22W_1%20between%20two%20equal-weight%20N-point%20empirical%20distributions.%22%22%22%0A%20%20%20%20%20%20%20%20return%20np.mean(np.abs(x_sorted%20-%20y_sorted))%0A%0A%20%20%20%20return%20(w1_sorted%2C)%0A%0A%0A%40app.cell%0Adef%20_(np%2C%20stats%2C%20w1_sorted)%3A%0A%20%20%20%20x_test%20%3D%20np.array(%5B1.0%2C%202.0%2C%203.0%2C%204.0%5D)%0A%20%20%20%20y_test%20%3D%20np.array(%5B1.5%2C%202.0%2C%203.5%2C%204.5%5D)%0A%0A%20%20%20%20assert%20np.isclose(w1_sorted(x_test%2C%20y_test)%2C%20stats.wasserstein_distance(x_test%2C%20y_test))%2C%20(%0A%20%20%20%20%20%20%20%20%22Formula%20mismatch%20%E2%80%94%20check%20sorting!%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(np%2C%20rng%2C%20stats)%3A%0A%20%20%20%20def%20sample_wasserstein_ball_simple(eng%2C%20empirical%2C%20epsilon%2C%20n_samples)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20Sample%20from%20%0A%20%20%20%20%20%20%20%20the%20Wasserstein%20ball%20using%20normalize-then-scale.%0A%0A%20%20%20%20%20%20%20%20The%20logic%20in%20plain%20english%3A%0A%20%20%20%20%20%20%20%20%20%201.%20Draw%20a%20random%20direction%20in%20R%5EN%20(any%20symmetric%20distribution%20works).%0A%20%20%20%20%20%20%20%20%20%202.%20Normalize%20it%20so%20its%20L1%20radius%20equals%20exactly%20epsilon%20%E2%80%94%20this%20puts%0A%20%20%20%20%20%20%20%20%20%20%20%20%20it%20on%20the%20surface%20of%20the%20L1%20ball%20(the%20%22sphere%22).%0A%20%20%20%20%20%20%20%20%20%203.%20Scale%20by%20Uniform(0%2C%201)%20to%20land%20somewhere%20inside%20the%20ball%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20not%20just%20on%20its%20boundary.%0A%0A%20%20%20%20%20%20%20%20This%20guarantees%20W_1(P_hat%2C%20Q)%20%3C%3D%20epsilon%20for%20every%20sample%20by%20construction.%0A%20%20%20%20%20%20%20%20The%20tradeoff%20vs.%20Barthe%20et%20al.%20is%20that%20sampling%20is%20not%20perfectly%20uniform%0A%20%20%20%20%20%20%20%20over%20the%20ball%20%E2%80%94%20points%20cluster%20slightly%20toward%20the%20center%20%E2%80%94%20but%20the%0A%20%20%20%20%20%20%20%20distributions%20we%20generate%20are%20just%20as%20diverse%20and%20valid%20for%20intuition.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%0A%20%20%20%20%20%20%20%20N%20%3D%20len(empirical)%0A%0A%20%20%20%20%20%20%20%20raw_deltas%20%3D%20rng.normal(0%2C%201%2C%20size%3D(n_samples%2C%20N))%0A%20%20%20%20%20%20%20%20l1_radii%20%3D%20np.mean(np.abs(raw_deltas)%2C%20axis%3D1%2C%20keepdims%3DTrue)%0A%20%20%20%20%20%20%20%20sphere_deltas%20%3D%20raw_deltas%20*%20(epsilon%20%2F%20l1_radii)%0A%0A%20%20%20%20%20%20%20%20scale%20%3D%20rng.uniform(0%2C%201%2C%20size%3D(n_samples%2C%201))%0A%20%20%20%20%20%20%20%20scaled_deltas%20%3D%20sphere_deltas%20*%20scale%0A%0A%20%20%20%20%20%20%20%20distributions%2C%20w1_distances%20%3D%20%5B%5D%2C%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20delta%20in%20scaled_deltas%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20perturbed%20%3D%20np.sort(empirical%20%2B%20delta)%0A%20%20%20%20%20%20%20%20%20%20%20%20w1%20%3D%20stats.wasserstein_distance(empirical%2C%20perturbed)%0A%20%20%20%20%20%20%20%20%20%20%20%20distributions.append(perturbed)%0A%20%20%20%20%20%20%20%20%20%20%20%20w1_distances.append(w1)%0A%0A%20%20%20%20%20%20%20%20return%20distributions%2C%20np.array(w1_distances)%0A%0A%20%20%20%20return%20(sample_wasserstein_ball_simple%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20slider%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D0.01%2C%0A%20%20%20%20%20%20%20%20stop%3D1.5%2C%0A%20%20%20%20%20%20%20%20step%3D0.05%2C%0A%20%20%20%20%20%20%20%20label%3D%22Ball%20radius%20%CE%B5%22%2C%0A%20%20%20%20%20%20%20%20value%3D0.3%2C%0A%20%20%20%20%20%20%20%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20return%20(slider%2C)%0A%0A%0A%40app.cell%0Adef%20_(empirical_sorted%2C%20rng%2C%20sample_wasserstein_ball_simple%2C%20slider)%3A%0A%20%20%20%20ball_eps%20%3D%20slider.value%0A%20%20%20%20ball_dists%2C%20ball_w1s%20%3D%20sample_wasserstein_ball_simple(rng%2C%20empirical_sorted%2C%20ball_eps%2C%20300)%0A%20%20%20%20return%20ball_dists%2C%20ball_eps%2C%20ball_w1s%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20N%2C%0A%20%20%20%20ball_dists%2C%0A%20%20%20%20ball_eps%2C%0A%20%20%20%20ball_w1s%2C%0A%20%20%20%20empirical_sorted%2C%0A%20%20%20%20gridspec%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20mu_true%2C%0A%20%20%20%20np%2C%0A%20%20%20%20plt%2C%0A%20%20%20%20sigma_true%2C%0A%20%20%20%20slider%2C%0A%20%20%20%20stats%2C%0A)%3A%0A%20%20%20%20color%20%3D%20%22%232196F3%22%0A%20%20%20%20x_grid%20%3D%20np.linspace(mu_true%20-%204%20*%20sigma_true%2C%20mu_true%20%2B%204%20*%20sigma_true%2C%20500)%0A%0A%20%20%20%20fig%20%3D%20plt.figure(figsize%3D(14%2C%205))%0A%20%20%20%20fig.suptitle(%0A%20%20%20%20%20%20%20%20%22Wasserstein%20Ball%3A%20Sampling%20from%20the%20Ambiguity%20Set%20%20%7C%20%20%22%0A%20%20%20%20%20%20%20%20r%22Center%20%24%5Chat%7B%5Cmathbb%7BP%7D%7D%24%3A%20empirical%20distribution%2C%20%24N%3D60%24%2C%20%24%5Cmathcal%7BN%7D(5%2C1)%24%22%2C%0A%20%20%20%20%20%20%20%20fontsize%3D11%2C%0A%20%20%20%20)%0A%0A%20%20%20%20gs%20%3D%20gridspec.GridSpec(1%2C%202%2C%20figure%3Dfig%2C%20wspace%3D0.35)%0A%0A%20%20%20%20ax_cdf%20%3D%20fig.add_subplot(gs%5B0%2C%200%5D)%0A%0A%20%20%20%20for%20_q%20in%20ball_dists%5B%3A80%5D%3A%0A%20%20%20%20%20%20%20%20_cdf_y%20%3D%20np.arange(1%2C%20N%20%2B%201)%20%2F%20N%0A%20%20%20%20%20%20%20%20ax_cdf.step(_q%2C%20_cdf_y%2C%20color%3Dcolor%2C%20alpha%3D0.08%2C%20linewidth%3D0.8%2C%20where%3D%22post%22)%0A%0A%20%20%20%20_cdf_y%20%3D%20np.arange(1%2C%20N%20%2B%201)%20%2F%20N%0A%20%20%20%20ax_cdf.step(%0A%20%20%20%20%20%20%20%20empirical_sorted%2C%0A%20%20%20%20%20%20%20%20_cdf_y%2C%0A%20%20%20%20%20%20%20%20color%3D%22black%22%2C%0A%20%20%20%20%20%20%20%20linewidth%3D2.0%2C%0A%20%20%20%20%20%20%20%20where%3D%22post%22%2C%0A%20%20%20%20%20%20%20%20label%3Dr%22%24%5Chat%7B%5Cmathbb%7BP%7D%7D%24%20(empirical)%22%2C%0A%20%20%20%20%20%20%20%20zorder%3D5%2C%0A%20%20%20%20)%0A%20%20%20%20ax_cdf.plot(%0A%20%20%20%20%20%20%20%20x_grid%2C%0A%20%20%20%20%20%20%20%20stats.norm.cdf(x_grid%2C%20mu_true%2C%20sigma_true)%2C%0A%20%20%20%20%20%20%20%20color%3D%22gray%22%2C%0A%20%20%20%20%20%20%20%20linewidth%3D1.5%2C%0A%20%20%20%20%20%20%20%20linestyle%3D%22--%22%2C%0A%20%20%20%20%20%20%20%20label%3Dr%22True%20%24%5Cmathcal%7BN%7D(5%2C1)%24%22%2C%0A%20%20%20%20)%0A%20%20%20%20ax_cdf.set_title(f%22CDF%20ribbon%20%E2%80%94%20%CE%B5%20%3D%20%7Bball_eps%3A.2f%7D%22%2C%20fontsize%3D12%2C%20color%3Dcolor%2C%20fontweight%3D%22bold%22)%0A%20%20%20%20ax_cdf.set_xlabel(%22Value%22%2C%20fontsize%3D10)%0A%20%20%20%20ax_cdf.set_ylabel(%22CDF%22%2C%20fontsize%3D10)%0A%20%20%20%20ax_cdf.legend(fontsize%3D8%2C%20loc%3D%22upper%20left%22)%0A%20%20%20%20ax_cdf.set_xlim(x_grid%5B0%5D%2C%20x_grid%5B-1%5D)%0A%20%20%20%20ax_cdf.set_ylim(-0.02%2C%201.05)%0A%20%20%20%20ax_cdf.grid(True%2C%20alpha%3D0.3)%0A%20%20%20%20ax_cdf.text(%0A%20%20%20%20%20%20%20%200.98%2C%200.5%2C%0A%20%20%20%20%20%20%20%20f%22Ribbon%20width%5Cn%E2%88%9D%20%CE%B5%20%3D%20%7Bball_eps%3A.2f%7D%22%2C%0A%20%20%20%20%20%20%20%20transform%3Dax_cdf.transAxes%2C%0A%20%20%20%20%20%20%20%20fontsize%3D8%2C%20ha%3D%22right%22%2C%20va%3D%22center%22%2C%20color%3Dcolor%2C%0A%20%20%20%20%20%20%20%20bbox%3Ddict(boxstyle%3D%22round%2Cpad%3D0.3%22%2C%20facecolor%3D%22white%22%2C%20alpha%3D0.8)%2C%0A%20%20%20%20)%0A%0A%20%20%20%20ax_hist%20%3D%20fig.add_subplot(gs%5B0%2C%201%5D)%0A%20%20%20%20ax_hist.hist(ball_w1s%2C%20bins%3D30%2C%20color%3Dcolor%2C%20alpha%3D0.75%2C%20edgecolor%3D%22white%22%2C%20linewidth%3D0.5)%0A%20%20%20%20ax_hist.axvline(ball_eps%2C%20color%3D%22black%22%2C%20linewidth%3D2%2C%20linestyle%3D%22--%22%2C%20label%3Df%22%CE%B5%20%3D%20%7Bball_eps%3A.2f%7D%22)%0A%20%20%20%20ax_hist.axvline(%0A%20%20%20%20%20%20%20%20ball_w1s.mean()%2C%20color%3Dcolor%2C%20linewidth%3D1.5%2C%20linestyle%3D%22-%22%2C%0A%20%20%20%20%20%20%20%20label%3Df%22mean%20%3D%20%7Bball_w1s.mean()%3A.3f%7D%22%2C%0A%20%20%20%20)%0A%20%20%20%20ax_hist.set_xlabel(%22Actual%20W%E2%82%81%20distance%20to%20P%CC%82%22%2C%20fontsize%3D10)%0A%20%20%20%20ax_hist.set_ylabel(%22Count%22%2C%20fontsize%3D10)%0A%20%20%20%20ax_hist.set_title(f%22W%E2%82%81%20distances%20(%CE%B5%20%3D%20%7Bball_eps%3A.2f%7D)%22%2C%20fontsize%3D11)%0A%20%20%20%20ax_hist.legend(fontsize%3D9)%0A%20%20%20%20ax_hist.grid(True%2C%20alpha%3D0.3)%0A%20%20%20%20ax_hist.text(%0A%20%20%20%20%20%20%20%200.02%2C%200.95%2C%0A%20%20%20%20%20%20%20%20f%22Max%20%3D%20%7Bball_w1s.max()%3A.3f%7D%20%E2%89%A4%20%CE%B5%20%E2%9C%93%22%2C%0A%20%20%20%20%20%20%20%20transform%3Dax_hist.transAxes%2C%0A%20%20%20%20%20%20%20%20fontsize%3D9%2C%20va%3D%22top%22%2C%20color%3D%22green%22%2C%20fontweight%3D%22bold%22%2C%0A%20%20%20%20)%0A%0A%20%20%20%20mo.vstack(%5Bslider%2C%20mo.as_html(fig)%5D)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%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