GPU Portfolio Optimizer

GPU Portfolio Optimizer — Edgeless Lab

import%20marimo%0A%0A__generated_with%20%3D%20%220.23.13%22%0Aapp%20%3D%20marimo.App(%0A%20%20%20%20width%3D%22full%22%2C%0A%20%20%20%20app_title%3D%22GPU%20Portfolio%20Optimizer%20%E2%80%94%20Edgeless%20Lab%22%2C%0A)%0A%0A%0A%40app.cell(hide_code%3DTrue)%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%20%20%20%20import%20time%0A%0A%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20import%20jax%0A%20%20%20%20%20%20%20%20import%20jax.numpy%20as%20jnp%0A%20%20%20%20%20%20%20%20_devs%20%3D%20jax.devices()%0A%20%20%20%20%20%20%20%20HAS_GPU%20%3D%20bool(_devs)%20and%20_devs%5B0%5D.platform%20%3D%3D%20%22gpu%22%0A%20%20%20%20%20%20%20%20DEVICE%20%3D%20str(_devs%5B0%5D)%20if%20_devs%20else%20%22CPU%22%0A%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20jax%20%3D%20None%0A%20%20%20%20%20%20%20%20jnp%20%3D%20np%0A%20%20%20%20%20%20%20%20HAS_GPU%20%3D%20False%0A%20%20%20%20%20%20%20%20DEVICE%20%3D%20%22CPU%20(NumPy%20fallback%20%E2%80%94%20JAX%20unavailable)%22%0A%0A%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20from%20wigglystuff%20import%20Excalidraw%0A%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20Excalidraw%20%3D%20None%0A%0A%20%20%20%20%23%20---%20Edgeless%20brand%20palette%20-------------------------------------------%0A%20%20%20%20EDG%20%3D%20%7B%0A%20%20%20%20%20%20%20%20%22bg%22%3A%20%22%23050505%22%2C%0A%20%20%20%20%20%20%20%20%22panel%22%3A%20%22%23101014%22%2C%0A%20%20%20%20%20%20%20%20%22text%22%3A%20%22%23e5e5e5%22%2C%0A%20%20%20%20%20%20%20%20%22muted%22%3A%20%22%238a8a8a%22%2C%0A%20%20%20%20%20%20%20%20%22accent%22%3A%20%22%23ff6b35%22%2C%20%20%23%20primary%20brand%20orange%0A%20%20%20%20%20%20%20%20%22cyan%22%3A%20%22%233dd6d0%22%2C%0A%20%20%20%20%20%20%20%20%22green%22%3A%20%22%234ade80%22%2C%0A%20%20%20%20%20%20%20%20%22red%22%3A%20%22%23f4515b%22%2C%0A%20%20%20%20%20%20%20%20%22violet%22%3A%20%22%23a78bfa%22%2C%0A%20%20%20%20%20%20%20%20%22grid%22%3A%20%22%2326262c%22%2C%0A%20%20%20%20%7D%0A%0A%20%20%20%20plt.rcParams.update(%7B%0A%20%20%20%20%20%20%20%20%22figure.facecolor%22%3A%20EDG%5B%22bg%22%5D%2C%0A%20%20%20%20%20%20%20%20%22axes.facecolor%22%3A%20EDG%5B%22bg%22%5D%2C%0A%20%20%20%20%20%20%20%20%22savefig.facecolor%22%3A%20EDG%5B%22bg%22%5D%2C%0A%20%20%20%20%20%20%20%20%22axes.edgecolor%22%3A%20EDG%5B%22grid%22%5D%2C%0A%20%20%20%20%20%20%20%20%22axes.labelcolor%22%3A%20EDG%5B%22text%22%5D%2C%0A%20%20%20%20%20%20%20%20%22axes.titlecolor%22%3A%20EDG%5B%22text%22%5D%2C%0A%20%20%20%20%20%20%20%20%22text.color%22%3A%20EDG%5B%22text%22%5D%2C%0A%20%20%20%20%20%20%20%20%22xtick.color%22%3A%20EDG%5B%22text%22%5D%2C%0A%20%20%20%20%20%20%20%20%22ytick.color%22%3A%20EDG%5B%22text%22%5D%2C%0A%20%20%20%20%20%20%20%20%22grid.color%22%3A%20EDG%5B%22grid%22%5D%2C%0A%20%20%20%20%20%20%20%20%22legend.labelcolor%22%3A%20EDG%5B%22text%22%5D%2C%0A%20%20%20%20%20%20%20%20%22figure.titlesize%22%3A%2015%2C%0A%20%20%20%20%20%20%20%20%22axes.titlesize%22%3A%2013%2C%0A%20%20%20%20%7D)%0A%0A%20%20%20%20def%20style_ax(_ax%2C%20_grid%3DTrue)%3A%0A%20%20%20%20%20%20%20%20_ax.set_facecolor(EDG%5B%22bg%22%5D)%0A%20%20%20%20%20%20%20%20_ax.tick_params(colors%3DEDG%5B%22text%22%5D)%0A%20%20%20%20%20%20%20%20for%20_sp%20in%20_ax.spines.values()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_sp.set_color(EDG%5B%22grid%22%5D)%0A%20%20%20%20%20%20%20%20if%20_grid%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_ax.grid(True%2C%20alpha%3D0.25%2C%20color%3DEDG%5B%22grid%22%5D)%0A%20%20%20%20%20%20%20%20return%20_ax%0A%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20DEVICE%2C%0A%20%20%20%20%20%20%20%20EDG%2C%0A%20%20%20%20%20%20%20%20Excalidraw%2C%0A%20%20%20%20%20%20%20%20HAS_GPU%2C%0A%20%20%20%20%20%20%20%20jax%2C%0A%20%20%20%20%20%20%20%20jnp%2C%0A%20%20%20%20%20%20%20%20mo%2C%0A%20%20%20%20%20%20%20%20np%2C%0A%20%20%20%20%20%20%20%20plt%2C%0A%20%20%20%20%20%20%20%20style_ax%2C%0A%20%20%20%20%20%20%20%20time%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(DEVICE%2C%20HAS_GPU%2C%20mo)%3A%0A%20%20%20%20_badge_color%20%3D%20%22%234ade80%22%20if%20HAS_GPU%20else%20%22%23ff6b35%22%0A%20%20%20%20_badge%20%3D%20f%22%F0%9F%9F%A2%20GPU%20%C2%B7%20%7BDEVICE%7D%22%20if%20HAS_GPU%20else%20f%22%F0%9F%9F%A0%20%7BDEVICE%7D%22%0A%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20f%22%22%22%0A%20%20%20%20%20%20%20%20%23%20GPU%20Portfolio%20Optimizer%20%E2%80%94%20Markowitz%20at%20Scale%0A%0A%20%20%20%20%20%20%20%20Monte%20Carlo%20portfolio%20optimization%20on%20the%20GPU%3A%20sample%20hundreds%20of%20thousands%0A%20%20%20%20%20%20%20%20of%20random%20asset%20allocations%20in%20a%20single%20batched%20JAX%20kernel%2C%20evaluate%20every%0A%20%20%20%20%20%20%20%20portfolio's%20return%2C%20volatility%20and%20Sharpe%20ratio%20at%20once%2C%20and%20trace%20the%0A%20%20%20%20%20%20%20%20**efficient%20frontier**.%20Every%20slider%20re-runs%20the%20full%20optimization.%0A%0A%20%20%20%20%20%20%20%20**Active%20device%3A**%20%3Cspan%20style%3D%22color%3A%7B_badge_color%7D%22%3E%7B_badge%7D%3C%2Fspan%3E%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20n_assets%20%3D%20mo.ui.slider(5%2C%20500%2C%20step%3D5%2C%20value%3D50%2C%20label%3D%22Number%20of%20Assets%22)%0A%20%20%20%20n_portfolios%20%3D%20mo.ui.slider(1_000%2C%20500_000%2C%20step%3D1_000%2C%20value%3D50_000%2C%20label%3D%22Random%20Portfolios%22)%0A%20%20%20%20risk_free%20%3D%20mo.ui.slider(0.0%2C%200.08%2C%20step%3D0.001%2C%20value%3D0.02%2C%20label%3D%22Risk-Free%20Rate%22)%0A%20%20%20%20target_return%20%3D%20mo.ui.slider(0.05%2C%200.30%2C%20step%3D0.005%2C%20value%3D0.12%2C%20label%3D%22Target%20Return%22)%0A%20%20%20%20return%20n_assets%2C%20n_portfolios%2C%20risk_free%2C%20target_return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo%2C%20n_assets%2C%20n_portfolios%2C%20risk_free%2C%20target_return)%3A%0A%20%20%20%20mo.hstack(%5B%0A%20%20%20%20%20%20%20%20mo.vstack(%5Bn_assets%2C%20risk_free%5D%2C%20align%3D%22start%22)%2C%0A%20%20%20%20%20%20%20%20mo.vstack(%5Bn_portfolios%2C%20target_return%5D%2C%20align%3D%22start%22)%2C%0A%20%20%20%20%5D)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(n_assets%2C%20np)%3A%0A%20%20%20%20%23%20---%20Synthetic%20market%3A%20expected%20returns%20%2B%20covariance%20------------------%0A%20%20%20%20np.random.seed(42)%0A%20%20%20%20_n%20%3D%20n_assets.value%0A%20%20%20%20true_returns%20%3D%20np.random.uniform(0.05%2C%200.25%2C%20_n)%0A%20%20%20%20vols%20%3D%20np.random.uniform(0.10%2C%200.60%2C%20_n)%0A%20%20%20%20_base_corr%20%3D%200.3%0A%20%20%20%20_corr%20%3D%20np.eye(_n)%20*%20(1%20-%20_base_corr)%20%2B%20_base_corr%0A%20%20%20%20cov%20%3D%20np.outer(vols%2C%20vols)%20*%20_corr%0A%20%20%20%20cov%20%3D%20(cov%20%2B%20cov.T)%20%2F%202%20%20%23%20enforce%20symmetry%0A%20%20%20%20np.fill_diagonal(cov%2C%20vols%20**%202)%0A%20%20%20%20return%20cov%2C%20true_returns%2C%20vols%0A%0A%0A%40app.cell%0Adef%20_(HAS_GPU%2C%20jax%2C%20jnp%2C%20np%2C%20time)%3A%0A%20%20%20%20%23%20---%20GPU%20portfolio%20kernel%20---------------------------------------------%0A%20%20%20%20%23%20Identical%20math%20on%20GPU%20(jax.numpy)%20or%20CPU%20(numpy).%20The%20bottleneck%20is%20the%0A%20%20%20%20%23%20quadratic%20form%20w%E1%B5%80%CE%A3w%20for%20every%20portfolio%2C%20mapped%20to%20one%20einsum%20kernel.%0A%20%20%20%20def%20make_portfolios(xp%2C%20returns_np%2C%20cov_np%2C%20n_port%2C%20rf)%3A%0A%20%20%20%20%20%20%20%20n%20%3D%20returns_np.shape%5B0%5D%0A%20%20%20%20%20%20%20%20if%20HAS_GPU%20and%20xp%20is%20jnp%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20w%20%3D%20jax.random.dirichlet(jax.random.PRNGKey(0)%2C%20jnp.ones(n)%2C%20(n_port%2C))%0A%20%20%20%20%20%20%20%20%20%20%20%20ret_x%2C%20cov_x%20%3D%20jnp.asarray(returns_np)%2C%20jnp.asarray(cov_np)%0A%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20w%20%3D%20np.random.dirichlet(np.ones(n)%2C%20n_port)%0A%20%20%20%20%20%20%20%20%20%20%20%20ret_x%2C%20cov_x%20%3D%20returns_np%2C%20cov_np%0A%20%20%20%20%20%20%20%20port_ret%20%3D%20xp.dot(w%2C%20ret_x)%0A%20%20%20%20%20%20%20%20port_vol%20%3D%20xp.sqrt(xp.maximum(xp.einsum(%22pi%2Cij%2Cpj-%3Ep%22%2C%20w%2C%20cov_x%2C%20w)%2C%200.0))%0A%20%20%20%20%20%20%20%20sharpe%20%3D%20(port_ret%20-%20rf)%20%2F%20port_vol%0A%20%20%20%20%20%20%20%20%23%20np.asarray%20forces%20a%20device%20sync%2C%20so%20GPU%20timings%20are%20honest%0A%20%20%20%20%20%20%20%20return%20np.asarray(w)%2C%20np.asarray(port_ret)%2C%20np.asarray(port_vol)%2C%20np.asarray(sharpe)%0A%0A%20%20%20%20def%20time_portfolios(xp%2C%20returns_np%2C%20cov_np%2C%20n_port%2C%20rf)%3A%0A%20%20%20%20%20%20%20%20make_portfolios(xp%2C%20returns_np%2C%20cov_np%2C%20max(100%2C%20n_port%20%2F%2F%2010)%2C%20rf)%20%20%23%20warmup%0A%20%20%20%20%20%20%20%20t0%20%3D%20time.time()%0A%20%20%20%20%20%20%20%20make_portfolios(xp%2C%20returns_np%2C%20cov_np%2C%20n_port%2C%20rf)%0A%20%20%20%20%20%20%20%20return%20time.time()%20-%20t0%0A%0A%20%20%20%20return%20make_portfolios%2C%20time_portfolios%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(cov%2C%20jnp%2C%20make_portfolios%2C%20n_portfolios%2C%20np%2C%20risk_free%2C%20true_returns)%3A%0A%20%20%20%20_xp%20%3D%20jnp%20if%20make_portfolios%20else%20np%0A%20%20%20%20weights%2C%20port_returns%2C%20port_vols%2C%20sharpe%20%3D%20make_portfolios(%0A%20%20%20%20%20%20%20%20_xp%2C%20true_returns%2C%20cov%2C%20n_portfolios.value%2C%20risk_free.value%0A%20%20%20%20)%0A%20%20%20%20max_sharpe_idx%20%3D%20int(np.nanargmax(sharpe))%0A%20%20%20%20min_vol_idx%20%3D%20int(np.nanargmin(port_vols))%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20max_sharpe_idx%2C%0A%20%20%20%20%20%20%20%20min_vol_idx%2C%0A%20%20%20%20%20%20%20%20port_returns%2C%0A%20%20%20%20%20%20%20%20port_vols%2C%0A%20%20%20%20%20%20%20%20sharpe%2C%0A%20%20%20%20%20%20%20%20weights%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(%0A%20%20%20%20max_sharpe_idx%2C%0A%20%20%20%20min_vol_idx%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20port_returns%2C%0A%20%20%20%20port_vols%2C%0A%20%20%20%20risk_free%2C%0A%20%20%20%20sharpe%2C%0A)%3A%0A%20%20%20%20_ms_ret%2C%20_ms_vol%2C%20_ms_sr%20%3D%20port_returns%5Bmax_sharpe_idx%5D%2C%20port_vols%5Bmax_sharpe_idx%5D%2C%20sharpe%5Bmax_sharpe_idx%5D%0A%20%20%20%20_mv_ret%2C%20_mv_vol%20%3D%20port_returns%5Bmin_vol_idx%5D%2C%20port_vols%5Bmin_vol_idx%5D%0A%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20f%22%22%22%0A%20%20%20%20%20%20%20%20**%7Blen(port_returns)%3A%2C%7D%20random%20portfolios**%20evaluated%20in%20one%20batched%20kernel.%0A%0A%20%20%20%20%20%20%20%20%7C%20Portfolio%20%7C%20Return%20%7C%20Volatility%20%7C%20Sharpe%20%7C%0A%20%20%20%20%20%20%20%20%7C-----------%7C--------%7C-----------%7C--------%7C%0A%20%20%20%20%20%20%20%20%7C%20**Max%20Sharpe**%20%7C%20%7B_ms_ret%3A.2%25%7D%20%7C%20%7B_ms_vol%3A.2%25%7D%20%7C%20%7B_ms_sr%3A.2f%7D%20%7C%0A%20%20%20%20%20%20%20%20%7C%20**Min%20Vol**%20%7C%20%7B_mv_ret%3A.2%25%7D%20%7C%20%7B_mv_vol%3A.2%25%7D%20%7C%20%7B(_mv_ret%20-%20risk_free.value)%20%2F%20_mv_vol%3A.2f%7D%20%7C%0A%0A%20%20%20%20%20%20%20%20Each%20dot%20below%20is%20one%20portfolio%2C%20colored%20by%20Sharpe%20ratio.%20The%20curved%20upper%0A%20%20%20%20%20%20%20%20edge%20of%20the%20cloud%20is%20the%20**efficient%20frontier**%20%E2%80%94%20no%20portfolio%20sits%20above%20it.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(%0A%20%20%20%20EDG%2C%0A%20%20%20%20max_sharpe_idx%2C%0A%20%20%20%20min_vol_idx%2C%0A%20%20%20%20n_assets%2C%0A%20%20%20%20np%2C%0A%20%20%20%20plt%2C%0A%20%20%20%20port_returns%2C%0A%20%20%20%20port_vols%2C%0A%20%20%20%20risk_free%2C%0A%20%20%20%20sharpe%2C%0A%20%20%20%20style_ax%2C%0A%20%20%20%20true_returns%2C%0A%20%20%20%20vols%2C%0A%20%20%20%20weights%2C%0A)%3A%0A%20%20%20%20_rf%20%3D%20risk_free.value%0A%0A%20%20%20%20_fig%20%3D%20plt.figure(figsize%3D(16%2C%2010))%0A%20%20%20%20_gs%20%3D%20_fig.add_gridspec(2%2C%203%2C%20height_ratios%3D%5B2%2C%201%5D%2C%20hspace%3D0.3%2C%20wspace%3D0.3)%0A%0A%20%20%20%20%23%20Main%3A%20efficient%20frontier%0A%20%20%20%20_ax_main%20%3D%20style_ax(_fig.add_subplot(_gs%5B0%2C%20%3A2%5D))%0A%20%20%20%20_sc%20%3D%20_ax_main.scatter(port_vols%2C%20port_returns%2C%20c%3Dsharpe%2C%20cmap%3D%22viridis%22%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%20%20alpha%3D0.3%2C%20s%3D1%2C%20rasterized%3DTrue)%0A%20%20%20%20_cbar%20%3D%20plt.colorbar(_sc%2C%20ax%3D_ax_main)%0A%20%20%20%20_cbar.set_label(%22Sharpe%20Ratio%22%2C%20color%3DEDG%5B%22text%22%5D)%0A%20%20%20%20_cbar.ax.yaxis.set_tick_params(color%3DEDG%5B%22text%22%5D)%0A%20%20%20%20for%20_lbl%20in%20_cbar.ax.get_yticklabels()%3A%0A%20%20%20%20%20%20%20%20_lbl.set_color(EDG%5B%22text%22%5D)%0A%0A%20%20%20%20_ax_main.scatter(port_vols%5Bmax_sharpe_idx%5D%2C%20port_returns%5Bmax_sharpe_idx%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20color%3DEDG%5B%22green%22%5D%2C%20s%3D200%2C%20zorder%3D5%2C%20edgecolors%3D%22white%22%2C%20linewidths%3D2%2C%20label%3D%22Max%20Sharpe%22)%0A%20%20%20%20_ax_main.scatter(port_vols%5Bmin_vol_idx%5D%2C%20port_returns%5Bmin_vol_idx%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20color%3DEDG%5B%22cyan%22%5D%2C%20s%3D200%2C%20zorder%3D5%2C%20edgecolors%3D%22white%22%2C%20linewidths%3D2%2C%20label%3D%22Min%20Vol%22)%0A%20%20%20%20_ax_main.scatter(vols%2C%20true_returns%2C%20color%3DEDG%5B%22accent%22%5D%2C%20s%3D60%2C%20zorder%3D4%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20edgecolors%3D%22white%22%2C%20linewidths%3D0.5%2C%20label%3Df%22%7Bn_assets.value%7D%20Assets%22)%0A%0A%20%20%20%20_ms_vol%2C%20_ms_ret%20%3D%20port_vols%5Bmax_sharpe_idx%5D%2C%20port_returns%5Bmax_sharpe_idx%5D%0A%20%20%20%20_cml_x%20%3D%20np.linspace(0%2C%20port_vols.max()%20*%200.8%2C%20100)%0A%20%20%20%20_cml_y%20%3D%20_rf%20%2B%20(_ms_ret%20-%20_rf)%20%2F%20_ms_vol%20*%20_cml_x%0A%20%20%20%20_ax_main.plot(_cml_x%2C%20_cml_y%2C%20%22--%22%2C%20color%3DEDG%5B%22muted%22%5D%2C%20alpha%3D0.7%2C%20linewidth%3D1.5%2C%20label%3D%22Capital%20Market%20Line%22)%0A%0A%20%20%20%20_ax_main.set_xlabel(%22Volatility%20(%CF%83)%22%2C%20fontsize%3D13)%0A%20%20%20%20_ax_main.set_ylabel(%22Expected%20Return%22%2C%20fontsize%3D13)%0A%20%20%20%20_ax_main.set_title(f%22Efficient%20Frontier%20%E2%80%94%20%7Bn_assets.value%7D%20Assets%2C%20%7Blen(port_returns)%3A%2C%7D%20Portfolios%22%2C%20fontsize%3D14)%0A%20%20%20%20_ax_main.legend(facecolor%3DEDG%5B%22panel%22%5D%2C%20edgecolor%3DEDG%5B%22grid%22%5D%2C%20fontsize%3D10)%0A%20%20%20%20_ax_main.grid(True%2C%20alpha%3D0.1%2C%20color%3DEDG%5B%22grid%22%5D)%0A%0A%20%20%20%20%23%20Top-right%3A%20max-Sharpe%20allocation%0A%20%20%20%20_ax_pie%20%3D%20_fig.add_subplot(_gs%5B0%2C%202%5D)%0A%20%20%20%20_top_n%20%3D%20min(8%2C%20n_assets.value)%0A%20%20%20%20_top_idx%20%3D%20np.argsort(weights%5Bmax_sharpe_idx%5D)%5B-_top_n%3A%5D%0A%20%20%20%20_top_w%20%3D%20weights%5Bmax_sharpe_idx%5D%5B_top_idx%5D%0A%20%20%20%20_other%20%3D%201.0%20-%20_top_w.sum()%0A%20%20%20%20_labels%20%3D%20%5Bf%22A%7B_i%7D%22%20for%20_i%20in%20_top_idx%5D%20%2B%20(%5B%22Other%22%5D%20if%20_other%20%3E%200.01%20else%20%5B%5D)%0A%20%20%20%20_sizes%20%3D%20list(_top_w)%20%2B%20(%5B_other%5D%20if%20_other%20%3E%200.01%20else%20%5B%5D)%0A%20%20%20%20_colors%20%3D%20plt.cm.viridis(np.linspace(0.2%2C%200.9%2C%20len(_sizes)))%0A%20%20%20%20_ax_pie.pie(_sizes%2C%20labels%3D_labels%2C%20colors%3D_colors%2C%20textprops%3D%7B%22color%22%3A%20EDG%5B%22text%22%5D%2C%20%22fontsize%22%3A%208%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20autopct%3D%22%251.0f%25%25%22%2C%20pctdistance%3D0.75)%0A%20%20%20%20_ax_pie.set_title(%22Max%20Sharpe%20Allocation%22%2C%20fontsize%3D12)%0A%0A%20%20%20%20%23%20Bottom%3A%20per-asset%20risk%2Freturn%0A%20%20%20%20_ax_bar%20%3D%20style_ax(_fig.add_subplot(_gs%5B1%2C%20%3A%5D)%2C%20_grid%3DFalse)%0A%20%20%20%20_sorted_idx%20%3D%20np.argsort(true_returns)%0A%20%20%20%20_ax_bar.bar(range(n_assets.value)%2C%20true_returns%5B_sorted_idx%5D%2C%20color%3DEDG%5B%22cyan%22%5D%2C%20alpha%3D0.6%2C%20label%3D%22Return%22)%0A%20%20%20%20_ax_twin%20%3D%20_ax_bar.twinx()%0A%20%20%20%20_ax_twin.plot(range(n_assets.value)%2C%20vols%5B_sorted_idx%5D%2C%20color%3DEDG%5B%22accent%22%5D%2C%20linewidth%3D2%2C%20marker%3D%22o%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20markersize%3D3%2C%20label%3D%22Volatility%22)%0A%20%20%20%20_ax_bar.set_xlabel(f%22%7Bn_assets.value%7D%20Assets%20(sorted%20by%20return)%22)%0A%20%20%20%20_ax_bar.set_ylabel(%22Expected%20Return%22%2C%20color%3DEDG%5B%22cyan%22%5D)%0A%20%20%20%20_ax_twin.set_ylabel(%22Volatility%22%2C%20color%3DEDG%5B%22accent%22%5D)%0A%20%20%20%20_ax_twin.tick_params(colors%3DEDG%5B%22text%22%5D)%0A%20%20%20%20_ax_bar.set_title(%22Individual%20Asset%20Risk-Return%20Profiles%22%2C%20fontsize%3D12)%0A%20%20%20%20_ax_bar.legend(facecolor%3DEDG%5B%22panel%22%5D%2C%20edgecolor%3DEDG%5B%22grid%22%5D%2C%20loc%3D%22upper%20left%22)%0A%20%20%20%20_ax_twin.legend(facecolor%3DEDG%5B%22panel%22%5D%2C%20edgecolor%3DEDG%5B%22grid%22%5D%2C%20loc%3D%22upper%20right%22)%0A%0A%20%20%20%20_fig%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(%0A%20%20%20%20HAS_GPU%2C%0A%20%20%20%20cov%2C%0A%20%20%20%20jnp%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20np%2C%0A%20%20%20%20plt%2C%0A%20%20%20%20risk_free%2C%0A%20%20%20%20style_ax%2C%0A%20%20%20%20time_portfolios%2C%0A%20%20%20%20true_returns%2C%0A)%3A%0A%20%20%20%20_bench_n%20%3D%20200_000%0A%0A%20%20%20%20_cpu_t%20%3D%20time_portfolios(np%2C%20true_returns%2C%20cov%2C%20_bench_n%2C%20risk_free.value)%0A%20%20%20%20_gpu_t%20%3D%20time_portfolios(jnp%2C%20true_returns%2C%20cov%2C%20_bench_n%2C%20risk_free.value)%20if%20HAS_GPU%20else%20None%0A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(11%2C%203.6))%0A%20%20%20%20style_ax(_ax%2C%20_grid%3DFalse)%0A%20%20%20%20_ax.set_title(f%22GPU%20vs%20CPU%20%E2%80%94%20%7B_bench_n%3A%2C%7D%20portfolios%20%C3%97%20%7Blen(true_returns)%7D%20assets%22%2C%20pad%3D12)%0A%0A%20%20%20%20_bars%20%3D%20%5B(%22CPU%20%C2%B7%20NumPy%22%2C%20_cpu_t%20*%201000%2C%20%22%238a8a8a%22)%5D%0A%20%20%20%20if%20_gpu_t%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20_bars.append((%22GPU%20%C2%B7%20JAX%22%2C%20_gpu_t%20*%201000%2C%20%22%23ff6b35%22))%0A%20%20%20%20_ypos%20%3D%20range(len(_bars))%0A%20%20%20%20_ax.barh(list(_ypos)%2C%20%5B_b%5B1%5D%20for%20_b%20in%20_bars%5D%2C%20color%3D%5B_b%5B2%5D%20for%20_b%20in%20_bars%5D%2C%20height%3D0.55%2C%20alpha%3D0.9)%0A%20%20%20%20_ax.set_yticks(list(_ypos))%0A%20%20%20%20_ax.set_yticklabels(%5B_b%5B0%5D%20for%20_b%20in%20_bars%5D)%0A%20%20%20%20_ax.set_xlabel(%22Wall-clock%20time%20(ms%2C%20lower%20is%20better)%22)%0A%20%20%20%20for%20_i%2C%20_b%20in%20enumerate(_bars)%3A%0A%20%20%20%20%20%20%20%20_ax.text(_b%5B1%5D%2C%20_i%2C%20f%22%20%20%7B_b%5B1%5D%3A.1f%7D%20ms%22%2C%20va%3D%22center%22%2C%20color%3D_b%5B2%5D%2C%20fontweight%3D%22bold%22)%0A%20%20%20%20plt.tight_layout()%0A%0A%20%20%20%20if%20_gpu_t%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20_verdict%20%3D%20mo.md(%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22**GPU%20speedup%3A%20%7B_cpu_t%20%2F%20_gpu_t%3A.1f%7D%C3%97%20faster**%20%E2%80%94%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22CPU%20%7B_cpu_t%20*%201000%3A.1f%7D%20ms%20vs%20GPU%20%7B_gpu_t%20*%201000%3A.1f%7D%20ms%20for%20%7B_bench_n%3A%2C%7D%20portfolios.%22%0A%20%20%20%20%20%20%20%20).callout(kind%3D%22success%22)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20_verdict%20%3D%20mo.md(%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22CPU%20baseline%3A%20**%7B_cpu_t%20*%201000%3A.1f%7D%20ms**%20for%20%7B_bench_n%3A%2C%7D%20portfolios.%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%22No%20GPU%20detected%20%E2%80%94%20open%20this%20notebook%20on%20**MoLab%20GPU**%20to%20populate%20the%20JAX%20bar%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%22(typically%2050%E2%80%93100%C3%97%20faster%20on%20an%20RTX%20Pro%206000).%22%0A%20%20%20%20%20%20%20%20).callout(kind%3D%22warn%22)%0A%0A%20%20%20%20mo.vstack(%5B_verdict%2C%20_fig%5D)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(EDG%2C%20make_portfolios%2C%20np%2C%20plt%2C%20risk_free)%3A%0A%20%20%20%20%23%20---%20Edge%20case%20%2F%20failure%20panel%20----------------------------------------%0A%20%20%20%20_rf%20%3D%20risk_free.value%0A%20%20%20%20_n%20%3D%205%0A%20%20%20%20_base_vol%20%3D%20np.array(%5B0.2%2C%200.3%2C%200.25%2C%200.4%2C%200.15%5D)%0A%20%20%20%20_base_ret%20%3D%20np.array(%5B0.10%2C%200.12%2C%200.08%2C%200.15%2C%200.06%5D)%0A%20%20%20%20_eye%20%3D%20np.eye(_n)%0A%0A%20%20%20%20def%20_cov_from(_vol%2C%20_corr_scalar%3D0.3)%3A%0A%20%20%20%20%20%20%20%20_c%20%3D%20np.eye(_n)%20*%20(1%20-%20_corr_scalar)%20%2B%20_corr_scalar%0A%20%20%20%20%20%20%20%20_m%20%3D%20np.outer(_vol%2C%20_vol)%20*%20_c%0A%20%20%20%20%20%20%20%20np.fill_diagonal(_m%2C%20_vol%20**%202)%0A%20%20%20%20%20%20%20%20return%20(_m%20%2B%20_m.T)%20%2F%202%0A%0A%20%20%20%20def%20_guard(_ret%2C%20_cov)%3A%0A%20%20%20%20%20%20%20%20_ret%20%3D%20np.nan_to_num(_ret%2C%20nan%3D0.0)%0A%20%20%20%20%20%20%20%20_cov%20%3D%20np.nan_to_num(_cov%2C%20nan%3D0.0)%0A%20%20%20%20%20%20%20%20_d%20%3D%20np.maximum(np.diag(_cov)%2C%201e-6)%0A%20%20%20%20%20%20%20%20np.fill_diagonal(_cov%2C%20_d)%0A%20%20%20%20%20%20%20%20return%20_ret%2C%20_cov%20%2B%20_eye%20*%201e-8%0A%0A%20%20%20%20def%20_bad_count(_ret%2C%20_cov)%3A%0A%20%20%20%20%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_%2C%20_pr%2C%20_pv%2C%20_sr%20%3D%20make_portfolios(np%2C%20_ret%2C%20_cov%2C%202_000%2C%20_rf)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20int((~np.isfinite(_pv)).sum()%20%2B%20(~np.isfinite(_sr)).sum())%0A%20%20%20%20%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20-1%0A%0A%20%20%20%20_zero_vol%20%3D%20_base_vol.copy()%3B%20_zero_vol%5B0%5D%20%3D%200.0%0A%20%20%20%20_nan_ret%20%3D%20_base_ret.copy()%3B%20_nan_ret%5B2%5D%20%3D%20np.nan%0A%20%20%20%20_scenarios%20%3D%20%5B%0A%20%20%20%20%20%20%20%20(%22Zero-volatility%20asset%22%2C%20_base_ret%2C%20_cov_from(_zero_vol))%2C%0A%20%20%20%20%20%20%20%20(%22NaN%20in%20expected%20returns%22%2C%20_nan_ret%2C%20_cov_from(_base_vol))%2C%0A%20%20%20%20%20%20%20%20(%22Singular%20covariance%20(%CF%81%3D1)%22%2C%20_base_ret%2C%20_cov_from(_base_vol%2C%201.0))%2C%0A%20%20%20%20%20%20%20%20(%22All-negative%20returns%22%2C%20-_base_ret%2C%20_cov_from(_base_vol))%2C%0A%20%20%20%20%5D%0A%0A%20%20%20%20_rows%20%3D%20%5B%5B%22Scenario%22%2C%20%22Raw%20bad%22%2C%20%22Guarded%20bad%22%2C%20%22Status%22%5D%5D%0A%20%20%20%20_row_colors%20%3D%20%5BNone%5D%0A%20%20%20%20for%20_name%2C%20_ret%2C%20_cov%20in%20_scenarios%3A%0A%20%20%20%20%20%20%20%20_raw%20%3D%20_bad_count(_ret%2C%20_cov)%0A%20%20%20%20%20%20%20%20_gr%2C%20_gc%20%3D%20_guard(_ret%2C%20_cov)%0A%20%20%20%20%20%20%20%20_safe%20%3D%20_bad_count(_gr%2C%20_gc)%0A%20%20%20%20%20%20%20%20_ok%20%3D%20_safe%20%3D%3D%200%0A%20%20%20%20%20%20%20%20_rows.append(%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20_name%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22error%22%20if%20_raw%20%3C%200%20else%20str(_raw)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22error%22%20if%20_safe%20%3C%200%20else%20str(_safe)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%E2%9C%93%20clean%22%20if%20_ok%20else%20%22%E2%9C%97%20NaN%2FInf%22%2C%0A%20%20%20%20%20%20%20%20%5D)%0A%20%20%20%20%20%20%20%20_row_colors.append(EDG%5B%22green%22%5D%20if%20_ok%20else%20EDG%5B%22red%22%5D)%0A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(12%2C%203.0))%0A%20%20%20%20_ax.axis(%22off%22)%0A%20%20%20%20_table%20%3D%20_ax.table(cellText%3D_rows%2C%20cellLoc%3D%22left%22%2C%20loc%3D%22center%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20colWidths%3D%5B0.42%2C%200.18%2C%200.20%2C%200.20%5D)%0A%20%20%20%20_table.auto_set_font_size(False)%0A%20%20%20%20_table.set_fontsize(10)%0A%20%20%20%20_table.scale(1%2C%201.7)%0A%20%20%20%20for%20_key%2C%20_cell%20in%20_table.get_celld().items()%3A%0A%20%20%20%20%20%20%20%20_cell.set_facecolor(EDG%5B%22panel%22%5D)%0A%20%20%20%20%20%20%20%20_cell.set_edgecolor(EDG%5B%22grid%22%5D)%0A%20%20%20%20%20%20%20%20_cell.set_text_props(color%3DEDG%5B%22text%22%5D)%0A%20%20%20%20%20%20%20%20if%20_key%5B0%5D%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_cell.set_text_props(color%3DEDG%5B%22cyan%22%5D%2C%20fontweight%3D%22bold%22)%0A%20%20%20%20%20%20%20%20elif%20_key%5B1%5D%20%3D%3D%203%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_cell.set_text_props(color%3D_row_colors%5B_key%5B0%5D%5D%2C%20fontweight%3D%22bold%22)%0A%20%20%20%20_ax.set_title(%22Edge%20Cases%20%26%20Guards%20%E2%80%94%20degenerate%20inputs%20vs.%20defended%20kernel%22%2C%20pad%3D14%2C%20loc%3D%22left%22)%0A%20%20%20%20_fig%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(Excalidraw%2C%20mo)%3A%0A%20%20%20%20if%20Excalidraw%20is%20None%3A%0A%20%20%20%20%20%20%20%20whiteboard%20%3D%20mo.md(%22*Excalidraw%20whiteboard%20requires%20%60wigglystuff%60%20(preinstalled%20on%20MoLab).*%22)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20whiteboard%20%3D%20mo.ui.anywidget(Excalidraw(height%3D500%2C%20theme%3D%22dark%22))%0A%20%20%20%20return%20(whiteboard%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20%23%23%20Annotate%20the%20frontier%0A%0A%20%20%20%20Use%20the%20Excalidraw%20canvas%20to%20mark%20up%20the%20efficient%20frontier%2C%20sketch%20capital%0A%20%20%20%20allocation%20lines%2C%20or%20derive%20the%20Markowitz%20optimization%20by%20hand.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(whiteboard)%3A%0A%20%20%20%20whiteboard%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20---%0A%20%20%20%20%23%23%23%20How%20it%20works%0A%0A%20%20%20%20%7C%20Stage%20%7C%20What%20it%20does%20%7C%0A%20%20%20%20%7C-------%7C--------------%7C%0A%20%20%20%20%7C%20Synthetic%20market%20%7C%20Expected%20returns%20%2B%20a%20regularized%20covariance%20matrix%20for%20*N*%20assets%20%7C%0A%20%20%20%20%7C%20GPU%20sampling%20%7C%20%60jax.random.dirichlet%60%20draws%20every%20weight%20vector%20in%20one%20batch%20%7C%0A%20%20%20%20%7C%20Vectorized%20math%20%7C%20%60jnp.einsum('pi%2Cij%2Cpj-%3Ep'%2C%20w%2C%20%CE%A3%2C%20w)%60%20evaluates%20all%20w%E1%B5%80%CE%A3w%20at%20once%20%7C%0A%20%20%20%20%7C%20Efficient%20frontier%20%7C%20Max-Sharpe%20and%20Min-Vol%20portfolios%20%2B%20the%20Capital%20Market%20Line%20%7C%0A%20%20%20%20%7C%20Benchmark%20%7C%20The%20same%20kernel%20timed%20on%20CPU%20(NumPy)%20and%20GPU%20(JAX)%20%7C%0A%20%20%20%20%7C%20Edge%20cases%20%7C%20Zero-vol%20assets%2C%20NaN%20returns%20and%20singular%20covariance%20vs.%20the%20guarded%20kernel%20%7C%0A%0A%20%20%20%20The%20simulation%20and%20benchmark%20cells%20share%20one%20%60make_portfolios(xp%2C%20%E2%80%A6)%60%20kernel%0A%20%20%20%20that%20takes%20the%20array%20module%20as%20an%20argument%2C%20so%20CPU%20and%20GPU%20run%20identical%20code.%0A%0A%20%20%20%20**Caveat%3A**%20this%20is%20mean-variance%20optimization%20on%20*synthetic*%20moments.%0A%20%20%20%20Real%20deployments%20need%20shrinkage%2Frobust%20covariance%20estimators%20and%20weight%0A%20%20%20%20constraints%20%E2%80%94%20random%20sampling%20only%20approximates%20the%20true%20frontier%20and%20is%0A%20%20%20%20optimistic%20in-sample.%0A%0A%20%20%20%20%23%23%23%20Run%20on%20MoLab%20GPU%0A%0A%20%20%20%201.%20Open%20%5Bmolab.marimo.io%5D(https%3A%2F%2Fmolab.marimo.io)%20and%20upload%20this%20file%20(**New%20%E2%86%92%20Upload**).%0A%20%20%20%202.%20Pick%20a%20**GPU%20runtime**%20(NVIDIA%20RTX%20Pro%206000%2C%2096%20GB%20VRAM).%0A%20%20%20%203.%20**Run%20all**%20%E2%80%94%20JAX%20auto-detects%20the%20GPU%2C%20the%20device%20badge%20turns%20green%2C%20and%0A%20%20%20%20%20%20%20the%20benchmark%20cell%20fills%20in%20real%20GPU%20timings.%0A%0A%20%20%20%20With%20no%20GPU%20the%20notebook%20transparently%20falls%20back%20to%20NumPy%20on%20CPU%2C%20so%20every%0A%20%20%20%20cell%20still%20runs%20locally%20%E2%80%94%20just%20slower.%0A%0A%20%20%20%20*Built%20with%20JAX%20%C2%B7%20wigglystuff%20%C2%B7%20marimo%20for%20the%20marimo%20%C3%97%20alphaXiv%20competition.*%0A%0A%20%20%20%20%60%23python%20%23quant%20%23portfolio%20%23markowitz%20%23gpu%20%23jax%60%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