hani2027/trustworthy-ml

1"""2modules/sam_optimizer.py -- SAM vs SGD on a sharp-vs-flat landscape.3 4The true loss L(w) has a narrow deep local minimum and a wide global one. SGD5descends L; SAM descends the worst-case surface E(w)=max_{‖ε‖≤ρ} L(w+ε). As ρ6grows the sharp well fills in on E, so SAM stops settling in the trap and moves7to the flat basin. Both 3-D surfaces, both descent paths, and the sharpness8readouts are recomputed from ρ.9"""10from __future__ import annotations11 12import numpy as np13import plotly.graph_objects as go14import streamlit as st15 16import components as C17import plotting as P18from core import sam_core as S19from i18n import t20 21DEFAULT_RHO = 0.5522 23 24@st.cache_data(show_spinner=False)25def _demo(rho: float):26    return S.compute_demo(rho=rho)27 28 29def render(lang: str) -> None:30    C.eyebrow(t(lang, "pr_sam_eyebrow"))31    C.section_title(t(lang, "pr_sam"))32    C.demo_intro(t(lang, "pr_sam_what"), t(lang, "pr_sam_why"),33                 t(lang, "pr_sam_expect"),34                 labels=(t(lang, "di_what"), t(lang, "di_why"),35                         t(lang, "di_expect")))36 37    with st.container(border=True):38        rho = st.slider(t(lang, "pr_sam_ctrl_rho"), 0.0, 1.2, DEFAULT_RHO, 0.05)39 40    d = _demo(float(rho))41    Xt, Yt, Zt = d["true_surface"]42    Xe, Ye, Ze = d["eff_surface"]43    sgd, sam = d["sgd_traj"], d["sam_traj"]44    sgd_z, sam_z = d["sgd_loss_curve"], d["sam_loss_curve"]45 46    # ---- Figures 1 & 2: the two 3-D surfaces ----47    g_left, g_right = st.columns(2)48    with g_left:49        with st.container(border=True):50            f1 = go.Figure()51            f1.add_trace(go.Surface(x=Xt, y=Yt, z=Zt, colorscale=P.NAVY_SCALE,52                                    showscale=False, opacity=0.92))53            f1.add_trace(go.Scatter3d(54                x=sgd[:, 0], y=sgd[:, 1], z=np.asarray(sgd_z) + 0.04,55                mode="lines", name="SGD",56                line=dict(color=P.CRIMSON, width=5)))57            f1.add_trace(go.Scatter3d(58                x=sam[:, 0], y=sam[:, 1], z=np.asarray(sam_z) + 0.04,59                mode="lines", name="SAM",60                line=dict(color=P.TEAL, width=5)))61            P.style_3d(f1, z_title="Loss", height=440)62            st.plotly_chart(f1, use_container_width=True, config=P.PLOTLY_CONFIG)63            C.figure_caption(t(lang, "pr_sam_cap_true"))64 65    with g_right:66        with st.container(border=True):67            f2 = go.Figure()68            f2.add_trace(go.Surface(x=Xe, y=Ye, z=Ze, colorscale=P.TEAL_SCALE,69                                    showscale=False, opacity=0.92))70            f2.add_trace(go.Scatter3d(71                x=[S.SHARP_C[0]], y=[S.SHARP_C[1]], z=[d["sharp_eff"] + 0.05],72                mode="markers", name="sharp",73                marker=dict(size=4, color=P.CRIMSON)))74            f2.add_trace(go.Scatter3d(75                x=[S.FLAT_C[0]], y=[S.FLAT_C[1]], z=[d["flat_eff"] + 0.05],76                mode="markers", name="flat",77                marker=dict(size=4, color=P.INK)))78            P.style_3d(f2, z_title="E(w)", height=440)79            st.plotly_chart(f2, use_container_width=True, config=P.PLOTLY_CONFIG)80            C.figure_caption(t(lang, "pr_sam_cap_eff"))81 82    # ---- Figure 3: descent paths on the true-loss contour ----83    with st.container(border=True):84        f3 = go.Figure()85        f3.add_trace(go.Contour(86            x=Xt[0], y=Yt[:, 0], z=Zt, showscale=False,87            colorscale=P.NAVY_SCALE, opacity=0.55,88            contours=dict(coloring="heatmap"), hoverinfo="skip"))89        f3.add_trace(go.Scatter(x=sgd[:, 0], y=sgd[:, 1], mode="lines",90                                name="SGD", line=dict(color=P.CRIMSON, width=3)))91        f3.add_trace(go.Scatter(x=sam[:, 0], y=sam[:, 1], mode="lines",92                                name="SAM", line=dict(color=P.TEAL, width=3)))93        f3.add_trace(go.Scatter(x=[sgd[0, 0]], y=[sgd[0, 1]], mode="markers",94                                name="start",95                                marker=dict(size=10, color=P.INK, symbol="circle")))96        f3.add_trace(go.Scatter(x=[sgd[-1, 0]], y=[sgd[-1, 1]], mode="markers",97                                name="SGD end",98                                marker=dict(size=11, color=P.CRIMSON, symbol="x")))99        f3.add_trace(go.Scatter(x=[sam[-1, 0]], y=[sam[-1, 1]], mode="markers",100                                name="SAM end",101                                marker=dict(size=11, color=P.TEAL, symbol="star")))102        P.style_2d(f3, x_title="w₁", y_title="w₂", height=380)103        st.plotly_chart(f3, use_container_width=True, config=P.PLOTLY_CONFIG)104        C.figure_caption(t(lang, "pr_sam_cap_traj"))105 106    # ---- readouts ----107    sam_basin = d["sam_basin"]108    C.readout_strip([109        {"k": "SHARP SHARPNESS", "v": f"{d['sharp_sharpness']:.2f}", "color": "crimson"},110        {"k": "FLAT SHARPNESS", "v": f"{d['flat_sharpness']:.2f}", "color": "teal"},111        {"k": "SGD BASIN", "v": d["sgd_basin"].upper(),112         "color": "teal" if d["sgd_basin"] == "flat" else "crimson"},113        {"k": "SAM BASIN", "v": sam_basin.upper(),114         "color": "teal" if sam_basin == "flat" else "crimson"},115    ])116 117    # ---- status ----118    if sam_basin == "flat":119        C.key_idea(t(lang, "pr_sam_flat"))120    else:121        C.warn(t(lang, "pr_sam_sharp"))122