services/plotting.py

"""
Plotly figure builders using the "standard" PDF-uncertainty style the
user prefers: a central curve + uncertainty band computed via LHAPDF's
own `PDFSet.uncertainty()`, and a Hessian-style correlation contour
between two flavors.

IMPORTANT: `PDFSet.uncertainty()` is only mathematically valid for
error types "hessian", "symmhessian", and "replicas". Every function
here assumes the caller already validated that with
`services.pdf_utils.load_validated_pdf_set` — these functions do not
re-check it themselves, to keep the validation in exactly one place.

Nothing in this module talks to Hugging Face or FastAPI directly; each
function takes an already-loaded `pdf_set` and returns a
`plotly.graph_objects.Figure`.
"""

import logging
import time
from typing import Optional

import numpy as np
import plotly.graph_objects as go

from services.colors import get_main_plot_label, get_parton_color, get_parton_name, get_parton_symbol

logger = logging.getLogger(__name__)


def build_standard_plot(
    pdf_set,
    pdfs,
    grid_name: str,
    q_scale: float,
    parton_id: int,
    plot_color: Optional[str] = None,
    x_vals=None,
    custom_title: Optional[str] = None,
) -> go.Figure:
    """
    Central value + uncertainty band for x*f(x, Q) of a single flavor.

    `pdfs` is `pdf_set.mkPDFs()`, passed in already-built (see
    services.pdf_utils.get_pdf_members) so repeat requests for the same
    PDF set don't reload every member's grid file from disk each time.
    """
    plot_color = plot_color or get_parton_color(parton_id)

    if x_vals is None:
        x_vals = np.logspace(-5, -0.01, 150)

    logger.info(
        "build_standard_plot: grid=%s parton_id=%s q_scale=%s n_points=%d members=%d",
        grid_name, parton_id, q_scale, len(x_vals), len(pdfs),
    )
    t0 = time.monotonic()
    central_values, upper_band, lower_band = [], [], []
    for x in x_vals:
        vals = [pdf.xfxQ(parton_id, x, q_scale) for pdf in pdfs]
        uncertainty = pdf_set.uncertainty(vals)

        central = uncertainty.central
        central_values.append(round(central, 6))
        upper_band.append(round(central + uncertainty.errplus, 6))
        lower_band.append(round(central - uncertainty.errminus, 6))
    logger.info("build_standard_plot: evaluated %d points in %.2fs", len(x_vals), time.monotonic() - t0)

    fig = go.Figure()

    x_vals_rounded = np.round(x_vals, 6)
    x_band = np.concatenate([x_vals_rounded, x_vals_rounded[::-1]])
    y_band = np.concatenate([upper_band, lower_band[::-1]])

    fig.add_trace(go.Scatter(
        x=x_band,
        y=y_band,
        fill="toself",
        fillcolor=plot_color,
        opacity=0.3,
        line=dict(color="rgba(255,255,255,0)"),
        name="Uncertainty",
        hoverinfo="skip",
    ))

    hover_custom_data = np.stack((upper_band, lower_band), axis=-1)
    fig.add_trace(go.Scatter(
        x=x_vals,
        y=central_values,
        mode="lines",
        line=dict(color=plot_color, width=2),
        name=f"{grid_name} Central <br>Q={q_scale} GeV",
        customdata=hover_custom_data,
        hovertemplate=(
            "<b>x:</b> %{x:.4e}<br>"
            "<b>Central:</b> %{y:.4f}<br>"
            "<b>Max (Upper):</b> %{customdata[0]:.4f}<br>"
            "<b>Min (Lower):</b> %{customdata[1]:.4f}"
            "<extra></extra>"
        ),
    ))

    default_title = " " #f"Parton Distribution: {get_parton_name(parton_id)} at Q = {q_scale} GeV"
    fig.update_layout(
        title=custom_title if custom_title else default_title,
        xaxis_title="$x$",
        yaxis_title=f"$x {get_parton_symbol(parton_id)}(x, Q)$",
        template="plotly_white",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        autosize=True,
        legend=dict(x=0.02, y=0.98, xanchor="left", yanchor="top", bgcolor="rgba(255,255,255,0.8)"),
        margin=dict(l=60, r=40, t=60, b=60),
    )
    fig.update_xaxes(
        type="log",
        exponentformat="power",
        showline=True,
        linewidth=1.5,
        linecolor="black",
        mirror="allticks",
        ticks="inside",
        tickwidth=1.5,
        tickcolor="black",
        showgrid=False,
    )

    y_upper_limit = max(upper_band) * 1.30
    y_lower_limit = min(lower_band) * (1 - 0.20)
    fig.update_yaxes(
        range=[y_lower_limit, y_upper_limit],
        showline=True,
        zeroline=False,
        linewidth=1.5,
        linecolor="black",
        mirror="allticks",
        ticks="inside",
        tickwidth=1.5,
        tickcolor="black",
        showgrid=False,
    )
    return fig


def build_standard_ratio_plot(
    pdf_set,
    pdfs,
    grid_name: str,
    q_scale: float,
    parton_id: int,
    plot_color: Optional[str] = None,
    x_vals=None,
    custom_title: Optional[str] = None,
) -> go.Figure:
    """
    Same as build_standard_plot, but normalized to the central value at
    each x (i.e. central line sits at 1.0, band shows relative uncertainty).

    `pdfs` is `pdf_set.mkPDFs()`, passed in already-built — see
    services.pdf_utils.get_pdf_members.
    """
    plot_color = plot_color or get_parton_color(parton_id)

    if x_vals is None:
        x_vals = np.logspace(-5, -0.01, 150)

    logger.info(
        "build_standard_ratio_plot: grid=%s parton_id=%s q_scale=%s n_points=%d members=%d",
        grid_name, parton_id, q_scale, len(x_vals), len(pdfs),
    )
    t0 = time.monotonic()
    central_values, upper_band, lower_band = [], [], []
    for x in x_vals:
        vals = [pdf.xfxQ(parton_id, x, q_scale) for pdf in pdfs]
        uncertainty = pdf_set.uncertainty(vals)
        central1 = uncertainty.central

        if abs(central1) < 1e-10:
            # Treat a (numerically) zero central value as a flat, uninformative ratio.
            central_values.append(1.0)
            upper_band.append(1.0)
            lower_band.append(1.0)
        else:
            central_values.append(round(central1 / central1, 6))
            upper_band.append(round(1.0 + uncertainty.errplus / central1, 6))
            lower_band.append(round(1.0 - uncertainty.errminus / central1, 6))
    logger.info("build_standard_ratio_plot: evaluated %d points in %.2fs", len(x_vals), time.monotonic() - t0)

    fig = go.Figure()

    x_vals_rounded = np.round(x_vals, 6)
    x_band = np.concatenate([x_vals_rounded, x_vals_rounded[::-1]])
    y_band = np.concatenate([upper_band, lower_band[::-1]])

    fig.add_trace(go.Scatter(
        x=x_band,
        y=y_band,
        fill="toself",
        fillcolor=plot_color,
        opacity=0.3,
        line=dict(color="rgba(255,255,255,0)"),
        name="Uncertainty",
        hoverinfo="skip",
    ))

    hover_custom_data = np.stack((upper_band, lower_band), axis=-1)
    fig.add_trace(go.Scatter(
        x=x_vals,
        y=central_values,
        mode="lines",
        line=dict(color=plot_color, width=2),
        name=f"{grid_name} Central <br>Q={q_scale} GeV",
        customdata=hover_custom_data,
        hovertemplate=(
            "<b>x:</b> %{x:.4e}<br>"
            "<b>Central:</b> %{y:.4f}<br>"
            "<b>Max (Upper):</b> %{customdata[0]:.4f}<br>"
            "<b>Min (Lower):</b> %{customdata[1]:.4f}"
            "<extra></extra>"
        ),
    ))

    default_title = " " #f"Parton Distribution: {get_parton_name(parton_id)} at Q = {q_scale} GeV"
    fig.update_layout(
        title=custom_title if custom_title else default_title,
        xaxis_title="$x$",
        yaxis_title=f"$x {get_parton_symbol(parton_id)}(x, Q)/central$",
        template="plotly_white",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        autosize=True,
        legend=dict(x=0.02, y=0.98, xanchor="left", yanchor="top", bgcolor="rgba(255,255,255,0.8)"),
        margin=dict(l=60, r=40, t=60, b=60),
    )
    fig.update_xaxes(
        type="log",
        exponentformat="power",
        showline=True,
        linewidth=1.5,
        linecolor="black",
        mirror="allticks",
        ticks="inside",
        tickwidth=1.5,
        tickcolor="black",
        showgrid=False,
    )

    y_upper_limit = max(upper_band) * 1.30
    ymin_data = min(lower_band)
    y_lower_limit = ymin_data * (1 - 0.30) if ymin_data >= 0 else ymin_data * (1 + 0.30)

    fig.update_yaxes(
        range=[y_lower_limit, y_upper_limit],
        showline=True,
        zeroline=False,
        linewidth=1.5,
        linecolor="black",
        mirror="allticks",
        ticks="inside",
        tickwidth=1.5,
        tickcolor="black",
        showgrid=False,
    )
    return fig


def build_correlation_plot(
    pdfs,
    q_scale: float,
    parton_id_1: int,
    parton_id_2: int,
    x_vals=None,
    color_correlated: Optional[str] = None,
    color_anti_correlated: str = "black",
) -> go.Figure:
    """
    Hessian-style correlation contour between two flavors across an
    x-grid, at a fixed Q scale.

    `pdfs` is `pdf_set.mkPDFs()`, passed in already-built — see
    services.pdf_utils.get_pdf_members.
    """
    color_correlated = color_correlated or get_parton_color(parton_id_1)

    num_members = len(pdfs)

    if x_vals is None:
        x_vals = np.logspace(-4, -0.01, 35)

    x_vals_rounded = np.round(x_vals, 6)
    num_x = len(x_vals_rounded)

    logger.info(
        "build_correlation_plot: parton_id_1=%s parton_id_2=%s q_scale=%s n_points=%d members=%d",
        parton_id_1, parton_id_2, q_scale, num_x, num_members,
    )
    t0 = time.monotonic()
    matrix_p1 = np.zeros((num_members, num_x))
    matrix_p2 = np.zeros((num_members, num_x))
    for m_idx, pdf in enumerate(pdfs):
        for x_idx, x in enumerate(x_vals_rounded):
            matrix_p1[m_idx, x_idx] = pdf.xfxQ(parton_id_1, x, q_scale)
            matrix_p2[m_idx, x_idx] = pdf.xfxQ(parton_id_2, x, q_scale)
    logger.info("build_correlation_plot: evaluated %d members x %d points in %.2fs", num_members, num_x, time.monotonic() - t0)

    correlation_matrix = np.zeros((num_x, num_x))
    p1_central = matrix_p1[0, :]
    p2_central = matrix_p2[0, :]
    p1_dev = matrix_p1[1:, :] - p1_central
    p2_dev = matrix_p2[1:, :] - p2_central

    for i in range(num_x):
        for j in range(num_x):
            delta_x = p1_dev[:, i]
            delta_y = p2_dev[:, j]

            numerator = np.sum(delta_x * delta_y)
            sum_x_sq = np.sum(delta_x ** 2)
            sum_y_sq = np.sum(delta_y ** 2)
            denominator = np.sqrt(sum_x_sq * sum_y_sq)

            if denominator > 1e-10:
                corr_value = np.clip(numerator / denominator, -1.0, 1.0)
                correlation_matrix[j, i] = round(corr_value, 4)
            else:
                correlation_matrix[j, i] = 0.0

    name_1 = get_parton_symbol(parton_id_1)
    name_2 = get_parton_symbol(parton_id_2)

    fig = go.Figure(data=go.Contour(
        z=correlation_matrix,
        x=x_vals_rounded,
        y=x_vals_rounded,
        colorscale=[
            [0.0, color_anti_correlated],
            [0.5, "rgb(245, 245, 245)"],
            [1.0, color_correlated],
        ],
        zmin=-1.0,
        zmax=1.0,
        line=dict(width=1, color="rgba(0, 0, 0, 0.2)"),
        contours=dict(coloring="heatmap", showlines=False),
        ncontours=15,
        colorbar=dict(
            title=dict(text="Corr. Coefficient", side="right"),
            thickness=15,
            len=1.0,
        ),
        hovertemplate=(
            f"<b>x&#8321; ({name_1}):</b> %{{x:.4e}}<br>"
            f"<b>x&#8322; ({name_2}):</b> %{{y:.4e}}<br>"
            "<b>Correlation:</b> %{z:.4f}<extra></extra>"
        ),
    ))

    fig.update_layout(
        title=" ",
        template="plotly_white",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        margin=dict(l=70, r=20, t=70, b=70),
        autosize=True,
    )
    fig.update_xaxes(
        title_text=f"$x_1 \\text{{ in }} {name_1}(x_1, Q)$",
        type="log",
        exponentformat="power",
        showline=True,
        linewidth=1.5,
        linecolor="black",
        mirror="allticks",
        ticks="inside",
        tickwidth=1.5,
        tickcolor="black",
        showgrid=False,
    )
    fig.update_yaxes(
        title_text=f"$x_2 \\text{{ in }} {name_2}(x_2, Q)$",
        type="log",
        exponentformat="power",
        showline=True,
        linewidth=1.5,
        linecolor="black",
        mirror="allticks",
        ticks="inside",
        tickwidth=1.5,
        tickcolor="black",
        showgrid=False,
    )
    return fig


def build_main_plot(
    pdf_set,
    pdfs,
    grid_name: str,
    q_scale: float,
    parton_ids: list,
    x_vals=None,
    custom_title: Optional[str] = None,
) -> go.Figure:
    """
    The classic multi-flavor "overview" plot (CT18 Fig. 2 style, see
    arXiv:1912.10053): every flavor in parton_ids overlaid on one figure as
    central value + 90% C.L. uncertainty band, log-x / linear-y. The gluon
    (parton_id 21) is always scaled down by a fixed factor of 5 so it fits
    on the same y-scale as the quark flavors, following that same
    convention -- its legend entry is labelled "g/5" accordingly (see
    services.colors.get_main_plot_label).

    `pdfs` is `pdf_set.mkPDFs()`, passed in already-built (see
    services.pdf_utils.get_pdf_members).
    """
    if x_vals is None:
        x_vals = np.logspace(-6, np.log10(0.9), 500)

    logger.info(
        "build_main_plot: grid=%s parton_ids=%s q_scale=%s n_points=%d members=%d",
        grid_name, parton_ids, q_scale, len(x_vals), len(pdfs),
    )
    t0 = time.monotonic()

    fig = go.Figure()
    x_vals_rounded = np.round(x_vals, 6)
    x_band = np.concatenate([x_vals_rounded, x_vals_rounded[::-1]])

    all_upper_bands = []

    for parton_id in parton_ids:
        plot_color = get_parton_color(parton_id)
        label = get_main_plot_label(parton_id)
        scale = 1.0 / 5.0 if parton_id == 21 else 1.0

        central_values, upper_band, lower_band = [], [], []
        for x in x_vals:
            vals = [pdf.xfxQ(parton_id, x, q_scale) for pdf in pdfs]
            uncertainty = pdf_set.uncertainty(vals)

            central = uncertainty.central * scale
            central_values.append(round(central, 6))
            upper_band.append(round(central + uncertainty.errplus * scale, 6))
            lower_band.append(round(central - uncertainty.errminus * scale, 6))

        all_upper_bands.extend(upper_band)

        y_band = np.concatenate([upper_band, lower_band[::-1]])
        fig.add_trace(go.Scatter(
            x=x_band,
            y=y_band,
            fill="toself",
            fillcolor=plot_color,
            opacity=0.3,
            line=dict(color="rgba(255,255,255,0)"),
            name=label + " (unc.)",
            legendgroup=label,
            showlegend=False,
            hoverinfo="skip",
        ))

        hover_custom_data = np.stack((upper_band, lower_band), axis=-1)
        fig.add_trace(go.Scatter(
            x=x_vals,
            y=central_values,
            mode="lines",
            line=dict(color=plot_color, width=2),
            name=f"${label}$",
            legendgroup=label,
            customdata=hover_custom_data,
            hovertemplate=(
                f"<b>Flavor:</b> {label}<br>"
                "<b>x:</b> %{x:.4e}<br>"
                "<b>Central:</b> %{y:.4f}<br>"
                "<b>Max (Upper):</b> %{customdata[0]:.4f}<br>"
                "<b>Min (Lower):</b> %{customdata[1]:.4f}"
                "<extra></extra>"
            ),
        ))

    logger.info(
        "build_main_plot: evaluated %d flavors x %d points in %.2fs",
        len(parton_ids), len(x_vals), time.monotonic() - t0,
    )

    default_title = " "
    fig.update_layout(
        title=custom_title if custom_title else default_title,
        xaxis_title="$x$",
        yaxis_title="$x f(x, Q)$",
        template="plotly_white",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        autosize=True,
        legend=dict(x=0.98, y=0.98, xanchor="right", yanchor="top", bgcolor="rgba(255,255,255,0.8)"),
        margin=dict(l=60, r=40, t=60, b=60),
    )
    fig.update_xaxes(
        type="log",
        exponentformat="power",
        showline=True,
        linewidth=1.5,
        linecolor="black",
        mirror="allticks",
        ticks="inside",
        tickwidth=1.5,
        tickcolor="black",
        showgrid=False,
    )

    y_upper_limit = max(all_upper_bands) * 1.15 if all_upper_bands else 1.0
    fig.update_yaxes(
        range=[0, y_upper_limit],
        showline=True,
        zeroline=False,
        linewidth=1.5,
        linecolor="black",
        mirror="allticks",
        ticks="inside",
        tickwidth=1.5,
        tickcolor="black",
        showgrid=False,
    )
    return fig