Aarondard5/constrained_multiOT
0
1import os2import sys3import warnings4warnings.simplefilter('ignore')5 6import streamlit as st7 8st.set_page_config(page_title="MultiOT Network", layout="wide")9 10import numpy as np11import pandas as pd12import networkx as nx13import matplotlib14matplotlib.use('Agg')15import matplotlib.pyplot as plt16 17CODE_DIR = os.path.dirname(os.path.abspath(__file__))18sys.path.insert(0, CODE_DIR)19 20import dynamics_capacity as dyn_e_constr21import initialization as init22import generate_planar as gpl23import tools as tl24import analysis_tools as atl25 26# ── Fixed parameters (match notebook) ────────────────────────────────27N = 40028seedG = 1129seedF = 230weigth = 10.31N_real = 1 # single realization for speed32tol = 1e-233coupling = 'l2'34dt = 0.0935tot_time = 50036alpha = 237seed_dyn = seedG + seedF38 39OUT_DIR = os.path.join(CODE_DIR, '..', 'data', 'output', 'planar')40os.makedirs(OUT_DIR, exist_ok=True)41 42# ── Build graph once ─────────────────────────────────────────────────43@st.cache_resource44def build_graph():45 G = gpl.planar_graph(N, L_min=0.0, L_max=None, domain=(0, 0, 1, 1), seed=seedG)46 label = f'planar_{N}_0.0_1.5'47 f_edges = os.path.join(OUT_DIR, label + str(seedG) + '.csv')48 init.from_nx_graphs2df([G], outfile=f_edges)49 graph, nnode, nedge, nodes, nodes_inter, nodeName2Id, _ = init.file2graph(f_edges, sep=',')50 f_coord = os.path.join(OUT_DIR, label + str(seedG) + 'coord.csv')51 init.from_nx_graph2coord([G], nodes, outfile=f_coord,52 mapping=nodeName2Id, nodes_inter=nodes_inter)53 # mode '1' = import from file; input_path arg is unused in this mode54 init.coord_generation('1', OUT_DIR, f_coord, graph, nnode, sep=',')55 length0 = init.eucledian_bias('eucl', graph, length_inter0=0.1, haversine_on=False)56 nx.set_edge_attributes(graph,57 {e: length0[i] for i, e in enumerate(graph.edges())},58 'length0')59 return G, graph, nnode, nedge, nodes, nodes_inter, nodeName2Id, length060 61G, graph, nnode, nedge, nodes, nodes_inter, nodeName2Id, length0 = build_graph()62 63# ── Forcing (cached per p) ────────────────────────────────────────────64@st.cache_data65def get_forcing(p_val):66 label = f'planar_{N}_{p_val}_1.5'67 f_path = os.path.join(OUT_DIR, label + f'_forcing{seedF}.csv')68 raw = tl.forcing_generate(G, p=p_val, weigth=weigth,69 pos=None, pos_center_of_mass=None, seed=seedF)70 pd.DataFrame(raw, columns=['source', 'sink', 'weight']).to_csv(71 f_path, sep=' ', header=True, index=False)72 forcing, _, _ = init.forcing_importing_from_file(73 f_path, nodes, nodes_inter, nodeName2Id,74 sep=' ', header=0, source='source', sink='sink')75 return forcing76 77# ── Dynamics (cached per beta, p, ce_pct) ────────────────────────────78@st.cache_data79def run_dynamics(beta_val, p_val, ce_pct):80 forcing = get_forcing(p_val)81 82 pflux_map_b = {0: beta_val, 'inter': 1.5, 'inter-super': 1.5}83 etypes = nx.get_edge_attributes(graph, 'etype')84 pflux_b = np.array([pflux_map_b[etypes[e]] for e in graph.edges()])85 86 # unconstrained87 tdens_u, _, _, _, _, flux_u, _ = dyn_e_constr.dyn(88 graph, nodes, pflux_b, nedge, length0, forcing,89 tol, removed_e_id=0, plot_cost=False, N_real=N_real,90 verbose=0, seed=seed_dyn, tdens0=None, coupling=coupling,91 constraint=False, capacity=False,92 time_step=dt, alpha=alpha, tot_time=tot_time93 )94 95 # constrained: c_e = ce_pct-th percentile of unconstrained tdens96 c_e = float(np.percentile(tdens_u, max(1, ce_pct)))97 tdens_c, _, _, _, _, flux_c, _ = dyn_e_constr.dyn(98 graph, nodes, pflux_b, nedge, length0, forcing,99 tol, removed_e_id=0, plot_cost=False, N_real=N_real,100 verbose=0, seed=seed_dyn, tdens0=None, coupling=coupling,101 constraint=True, capacity=c_e,102 time_step=dt, alpha=alpha, tot_time=tot_time103 )104 105 return flux_u, flux_c, c_e106 107# ── G_plot (cached, never changes) ───────────────────────────────────108@st.cache_resource109def build_gplot():110 colors_map = {0: 'steelblue', 1: 'r', 2: 'g', 3: 'magenta',111 'inter': 'black', 'inter-super': 'black'}112 G_p = nx.Graph(graph)113 etypes = nx.get_edge_attributes(graph, 'etype')114 for e in graph.edges:115 G_p[e[0]][e[1]]['color'] = colors_map[etypes[e]]116 return G_p117 118G_plot = build_gplot()119pos_plot = nx.get_node_attributes(G_plot, 'pos')120colors_e = list(nx.get_edge_attributes(G_plot, 'color').values())121 122# ═══════════════════════════════════════════════════════════════════════123# UI124# ═══════════════════════════════════════════════════════════════════════125st.title("Multi-commodity Optimal Transport — Interactive Network")126 127with st.sidebar:128 st.header("Parameters")129 beta_val = st.slider("β (branching exponent)", 0.1, 1.9, 1.5, step=0.1, format="%.1f")130 p_val = st.slider("p (destination spread)", 0.0, 1.0, 0.0, step=0.1, format="%.1f")131 st.markdown("---")132 ce_pct = st.slider("Edge capacity, c_e", 1, 100, 75, step=1)133 st.markdown("---")134 st.caption("**Source** ▲ green | **Sink** ◆ magenta")135 run_btn = st.button("Run", type="primary", use_container_width=True)136 137# ── Only compute when Run is clicked; reuse session state otherwise ───138if run_btn:139 with st.spinner("Running dynamics — this may take a few minutes …"):140 flux_u, flux_c, c_e = run_dynamics(141 round(beta_val, 1), round(p_val, 1), ce_pct)142 st.session_state['flux_u'] = flux_u143 st.session_state['flux_c'] = flux_c144 st.session_state['c_e'] = c_e145 st.session_state['p_cached'] = round(p_val, 1)146 147if 'flux_u' not in st.session_state:148 st.info("Set parameters in the sidebar and click **Run** to compute.")149 st.stop()150 151flux_u = st.session_state['flux_u']152flux_c = st.session_state['flux_c']153c_e = st.session_state['c_e']154 155forcing = get_forcing(st.session_state['p_cached'])156M = forcing.shape[0]157flux_norm1_u = np.linalg.norm(flux_u, axis=1, ord=1)158flux_norm1_c = np.linalg.norm(flux_c, axis=1, ord=1)159 160# Commodity slider — default to OD pair with second-longest path161avgL_u = atl.calculate_path_length(flux_u, length0)162default_comm = int(np.argsort(-avgL_u)[1])163commodity_idx = st.slider(164 f"Commodity (OD pair) — {M} total", 0, M - 1, default_comm)165 166source = int(np.where(forcing[commodity_idx] > 0)[0][0])167sink = int(np.where(forcing[commodity_idx] < 0)[0][0])168 169# ── Draw helper ───────────────────────────────────────────────────────170def draw_net(ax, flux, flux_norm1, title, info=None):171 w1, w0 = 0.01, 1.0172 # background: overall flux widths173 widths = [w1 * flux_norm1[i] for i in range(nedge)]174 nx.draw_networkx(G_plot, pos_plot, node_size=2, with_labels=False,175 edge_color=colors_e, width=widths, ax=ax)176 # green overlay: selected OD pair177 widths_od = [0.1 + w0 * flux[i, commodity_idx] for i in range(nedge)]178 nx.draw_networkx(G_plot, pos_plot, node_size=0, with_labels=False,179 edge_color='green', width=widths_od, ax=ax)180 # source: green triangle181 nx.draw_networkx_nodes(G_plot, pos_plot, node_shape='^',182 nodelist=[source], node_size=300,183 node_color=['green'], edgecolors='black', ax=ax)184 # sink: magenta diamond185 nx.draw_networkx_nodes(G_plot, pos_plot, node_shape='D',186 nodelist=[sink], node_size=300,187 node_color=['magenta'], edgecolors='black', ax=ax)188 ax.set_title(title, fontsize=13, pad=8)189 if info:190 ax.text(0.02, 0.97, info, transform=ax.transAxes, fontsize=9,191 va='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.7))192 193fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 7))194draw_net(ax1, flux_u, flux_norm1_u,195 f"Unconstrained β={beta_val:.1f} p={p_val:.1f}")196draw_net(ax2, flux_c, flux_norm1_c,197 f"Constrained β={beta_val:.1f} p={p_val:.1f} c_e={ce_pct}th pct",198 info=f"c_e = {c_e:.5f}")199 200plt.tight_layout()201st.pyplot(fig)202plt.close(fig)203 204st.caption(205 f"Commodity **{commodity_idx}** │ "206 f"Source node **{source}** ▲ │ "207 f"Sink node **{sink}** ◆ │ "208 f"Total commodities: **{M}**"209)210 211