hani2027/trustworthy-ml

1"""2plotting.py -- a single, consistent visual grammar for every figure.3 4All charts in the sandbox are drawn with Plotly and share one typeset5"academic" look: white plate background, hairline grid, Merriweather for prose6labels and IBM Plex Mono for numbers and axis titles, and a restrained7navy / crimson / teal palette that matches the page theme. Keeping this in one8place is what makes a dozen different demos look like figures from the same9monograph rather than a pile of dashboards.10 11Nothing here imports Streamlit, so the module stays light and the styling can be12reasoned about (and unit-checked) on its own.13"""14from __future__ import annotations15 16import plotly.graph_objects as go17 18# --- Palette (kept in sync with theme.py by hand; small + stable) ----------- #19INK = "#0E2A47"20DEEP = "#0A1F38"21SLATE = "#5A6B7B"22PARCHMENT = "#F7F5EF"23CARD = "#FFFFFF"24CRIMSON = "#8A1C2B"25TEAL = "#0F6E66"26HAIRLINE = "#D9D4C7"27GOLD = "#A8842C"28 29SERIF = "Merriweather, Georgia, 'Times New Roman', serif"30MONO = "IBM Plex Mono, 'SFMono-Regular', Menlo, monospace"31 32# Categorical order used across the app: ink, then the two accents, then golds.33COLORWAY = [INK, CRIMSON, TEAL, GOLD, SLATE, "#3C5A78"]34 35# Monochrome navy ramp for the "true" loss surface; warm teal ramp for the36# worst-case (effective) surface so the two 3-D plots never get confused.37NAVY_SCALE = [[0.0, "#0A1F38"], [0.45, "#3C5A78"], [0.8, "#9FB0C0"], [1.0, "#ECE7DA"]]38TEAL_SCALE = [[0.0, "#0C4F49"], [0.45, "#2C7E76"], [0.8, "#9AC3BD"], [1.0, "#EFE9DB"]]39 40# Passed straight to st.plotly_chart(config=...).41PLOTLY_CONFIG = {42    "displayModeBar": False,43    "scrollZoom": False,44    "doubleClick": "reset",45    "displaylogo": False,46}47 48 49def _axis(title: str | None = None):50    """Common 2-D axis styling: hairline grid, mono ticks, no zero-line clutter."""51    return dict(52        title=dict(text=title or "", font=dict(family=MONO, size=12, color=INK)),53        showgrid=True, gridcolor=HAIRLINE, gridwidth=1,54        zeroline=False, showline=True, linecolor="#C4BFB1", linewidth=1,55        ticks="outside", ticklen=4, tickcolor="#C4BFB1",56        tickfont=dict(family=MONO, size=11, color=SLATE),57    )58 59 60def style_2d(fig: go.Figure, *, x_title: str = "", y_title: str = "",61             height: int = 360, y_log: bool = False, legend: bool = True,62             margin_t: int = 28) -> go.Figure:63    """Apply the house style to any 2-D Plotly figure, in place."""64    fig.update_layout(65        template="simple_white",66        paper_bgcolor=CARD, plot_bgcolor=CARD,67        colorway=COLORWAY,68        font=dict(family=SERIF, size=13, color=INK),69        height=height,70        margin=dict(l=58, r=22, t=margin_t, b=48),71        xaxis=_axis(x_title),72        yaxis=_axis(y_title),73        hoverlabel=dict(font=dict(family=MONO, size=12), bgcolor="#FFFFFF",74                        bordercolor=HAIRLINE),75        showlegend=legend,76        legend=dict(77            font=dict(family=MONO, size=11, color=INK),78            bgcolor="rgba(255,255,255,0.65)", bordercolor=HAIRLINE, borderwidth=1,79            orientation="h", yanchor="bottom", y=1.0, xanchor="right", x=1.0,80        ),81        bargap=0.18,82    )83    if y_log:84        fig.update_yaxes(type="log")85    return fig86 87 88def style_3d(fig: go.Figure, *, height: int = 460,89             z_title: str = "Loss", eye=(1.5, 1.4, 0.9)) -> go.Figure:90    """House style for 3-D surface plots (the SAM landscapes)."""91    _tfont = dict(family=MONO, size=11, color=SLATE)92    axis = dict(93        showbackground=True, backgroundcolor="#FCFBF7",94        gridcolor=HAIRLINE, zeroline=False,95        tickfont=dict(family=MONO, size=9, color=SLATE),96    )97 98    def _ax(text):99        return dict(title=dict(text=text, font=_tfont), **axis)100 101    fig.update_layout(102        paper_bgcolor=CARD,103        font=dict(family=SERIF, size=12, color=INK),104        height=height,105        margin=dict(l=2, r=2, t=24, b=2),106        showlegend=True,107        legend=dict(font=dict(family=MONO, size=11), bgcolor="rgba(255,255,255,0.7)",108                    bordercolor=HAIRLINE, borderwidth=1, x=0.0, y=1.0),109        scene=dict(110            xaxis=_ax("w₁"),111            yaxis=_ax("w₂"),112            zaxis=_ax(z_title),113            camera=dict(eye=dict(x=eye[0], y=eye[1], z=eye[2])),114            aspectmode="cube",115        ),116    )117    return fig118 119 120def empty_note(message: str, height: int = 320) -> go.Figure:121    """A blank styled canvas carrying a single centered note (e.g. error states)."""122    fig = go.Figure()123    fig.add_annotation(text=message, x=0.5, y=0.5, xref="paper", yref="paper",124                       showarrow=False,125                       font=dict(family=MONO, size=13, color=CRIMSON))126    fig.update_layout(paper_bgcolor=CARD, plot_bgcolor=CARD, height=height,127                      xaxis=dict(visible=False), yaxis=dict(visible=False),128                      margin=dict(l=10, r=10, t=10, b=10))129    return fig130 131 132# --------------------------------------------------------------------------- #133# Reusable directed-acyclic-graph figure (Simpson confounder + SCM diagrams)   #134# --------------------------------------------------------------------------- #135def causal_dag_figure(nodes: dict, edges: list, *,136                      labels: dict | None = None,137                      node_colors: dict | None = None,138                      edge_colors: dict | None = None,139                      label_positions: dict | None = None,140                      height: int = 340, title: str | None = None,141                      footnotes: list | None = None) -> go.Figure:142    """143    Draw a small causal graph: nodes as ringed dots with outside labels and144    directed edges as trimmed arrow annotations. Reused by the Simpson and145    counterfactual demos so every diagram in the app is visually identical.146 147    nodes           : {key: (x, y)}148    edges           : [(src_key, dst_key), ...]149    labels          : {key: display_text}  (defaults to the key)150    node_colors     : {key: hex}            (ring + dot tint; defaults to INK)151    edge_colors     : {(src,dst): hex}      (defaults to SLATE)152    label_positions : {key: plotly textposition}  (defaults chosen by geometry)153    footnotes       : list of small mono captions placed under the graph154    """155    labels = labels or {}156    node_colors = node_colors or {}157    edge_colors = edge_colors or {}158    label_positions = label_positions or {}159 160    xs = [p[0] for p in nodes.values()]161    ys = [p[1] for p in nodes.values()]162    x_min, x_max = min(xs), max(xs)163    y_min, y_max = min(ys), max(ys)164    span_x = max(x_max - x_min, 1e-6)165    span_y = max(y_max - y_min, 1e-6)166    pad_x = 0.35 * span_x + 0.15167    pad_y = 0.45 * span_y + 0.20168 169    fig = go.Figure()170 171    # ---- edges as arrows, trimmed so the head rests just off the node dot ----172    for (s, d) in edges:173        x0, y0 = nodes[s]174        x1, y1 = nodes[d]175        dx, dy = x1 - x0, y1 - y0176        dist = (dx * dx + dy * dy) ** 0.5 or 1.0177        ux, uy = dx / dist, dy / dist178        gap0, gap1 = 0.12 * dist + 0.03, 0.16 * dist + 0.05179        color = edge_colors.get((s, d), SLATE)180        fig.add_annotation(181            x=x1 - ux * gap1, y=y1 - uy * gap1,182            ax=x0 + ux * gap0, ay=y0 + uy * gap0,183            xref="x", yref="y", axref="x", ayref="y",184            showarrow=True, arrowhead=3, arrowsize=1.2, arrowwidth=1.7,185            arrowcolor=color, standoff=0, startstandoff=0,186        )187 188    # ---- nodes as ringed dots with labels just outside ----189    def default_pos(key):190        x, y = nodes[key]191        if y >= (y_min + y_max) / 2 + 1e-9:      # upper node -> label above192            return "top center"193        return "bottom center"194 195    nx, ny, ntext, tpos, fills, lines = [], [], [], [], [], []196    for key, (x, y) in nodes.items():197        nx.append(x); ny.append(y)198        ntext.append(labels.get(key, key))199        tpos.append(label_positions.get(key, default_pos(key)))200        col = node_colors.get(key, INK)201        fills.append(_tint(col, 0.10))202        lines.append(col)203 204    fig.add_trace(go.Scatter(205        x=nx, y=ny, mode="markers+text",206        text=ntext, textposition=tpos,207        textfont=dict(family=MONO, size=12, color=INK),208        marker=dict(size=22, color=fills,209                    line=dict(color=lines, width=2.2)),210        hoverinfo="skip", showlegend=False,211    ))212 213    fig.update_layout(214        paper_bgcolor=CARD, plot_bgcolor=CARD,215        height=height, margin=dict(l=10, r=10, t=30 if title else 12, b=10),216        title=dict(text=title or "", font=dict(family=MONO, size=12, color=SLATE),217                   x=0.02, xanchor="left"),218        xaxis=dict(visible=False, range=[x_min - pad_x, x_max + pad_x]),219        yaxis=dict(visible=False, range=[y_min - pad_y, y_max + pad_y],220                   scaleanchor="x", scaleratio=1),221        showlegend=False,222    )223 224    if footnotes:225        for i, note in enumerate(footnotes):226            fig.add_annotation(227                x=0.0, y=-0.02 - 0.07 * i, xref="paper", yref="paper",228                xanchor="left", yanchor="top", showarrow=False, text=note,229                font=dict(family=MONO, size=10.5, color=SLATE),230            )231    return fig232 233 234def _tint(hex_color: str, alpha: float) -> str:235    """Return an rgba() string of ``hex_color`` over white at the given alpha."""236    h = hex_color.lstrip("#")237    r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)238    r = int(r * alpha + 255 * (1 - alpha))239    g = int(g * alpha + 255 * (1 - alpha))240    b = int(b * alpha + 255 * (1 - alpha))241    return f"rgb({r},{g},{b})"242