plotting.py

from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


def _robust_limits(values: np.ndarray, q_low: float = 2.0, q_high: float = 98.0) -> tuple[float, float]:
    finite_values = values[np.isfinite(values)]
    if finite_values.size == 0:
        return -1.0, 1.0
    vmin = float(np.percentile(finite_values, q_low))
    vmax = float(np.percentile(finite_values, q_high))
    if np.isclose(vmax, vmin):
        vmax = vmin + 1e-9
    return vmin, vmax


def _packed_2d_values(packed: dict[str, np.ndarray]) -> np.ndarray:
    return np.concatenate([packed["xz"].ravel(), packed["zy"].ravel(), packed["yx"].ravel()])


def _packed_3d_values(packed: dict[str, np.ndarray]) -> np.ndarray:
    return packed["cloud_values"]


def cartesian_to_spherical(X: np.ndarray, Y: np.ndarray, Z: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    r = np.sqrt(X**2 + Y**2 + Z**2)
    r_safe = np.where(r > 0.0, r, 1e-12)
    theta = np.arccos(Z / r_safe)
    phi = np.mod(np.arctan2(Y, X), 2.0 * np.pi)
    return r_safe, theta, phi


def build_grid(n: int = 80, extent: float = 12.0) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    x = np.linspace(-extent, extent, n)
    X, Y, Z = np.meshgrid(x, x, x, indexing="ij")
    return x, X, Y, Z


def pack_density(
    rho: np.ndarray,
    X: np.ndarray,
    Y: np.ndarray,
    Z: np.ndarray,
    q_cloud: float = 99.7,
    eps: float = 1e-30,
    positive_y_cloud: bool = False,
) -> dict[str, np.ndarray]:
    mid = rho.shape[0] // 2
    xz = np.log10(rho[:, mid, :] + eps)
    zy = np.log10(rho[mid, :, :] + eps)
    yx = np.log10(rho[:, :, mid] + eps)

    threshold = np.percentile(rho, q_cloud)
    mask = rho >= threshold
    # Always keep only the positive-y half-space for 3D cloud points.
    # The parameter is kept for backward compatibility with existing calls.
    mask = mask & (Y >= 0.0)

    points = np.column_stack([X[mask], Y[mask], Z[mask]])
    cloud_values = np.log10(rho[mask] + eps)
    return {
        "xz": xz,
        "zy": zy,
        "yx": yx,
        "points": points,
        "cloud_values": cloud_values,
    }


def _plot_block(
    fig: plt.Figure,
    ax_xz: plt.Axes,
    ax_zy: plt.Axes,
    ax_yx: plt.Axes,
    ax_3d: plt.Axes,
    packed: dict[str, np.ndarray],
    x: np.ndarray,
    title_prefix: str,
    cmap_2d: str = "magma",
    cmap_3d: str | None = None,
    log2d_limits: tuple[float, float] | None = None,
    log3d_limits: tuple[float, float] | None = None,
    xyz_limits: tuple[tuple[float, float], tuple[float, float], tuple[float, float]] | None = None,
) -> None:
    xz = packed["xz"]
    zy = packed["zy"]
    yx = packed["yx"]
    points = packed["points"]
    cloud_values = packed["cloud_values"]
    if cmap_3d is None:
        cmap_3d = cmap_2d

    if log2d_limits is None:
        vmin2d, vmax2d = _robust_limits(np.concatenate([xz.ravel(), zy.ravel(), yx.ravel()]))
    else:
        vmin2d, vmax2d = log2d_limits

    if cloud_values.size == 0:
        vmin3d, vmax3d = (-1.0, 1.0) if log3d_limits is None else log3d_limits
    elif log3d_limits is None:
        vmin3d, vmax3d = _robust_limits(cloud_values)
    else:
        vmin3d, vmax3d = log3d_limits
    if np.isclose(vmax3d, vmin3d):
        vmax3d = vmin3d + 1e-9

    im1 = ax_xz.imshow(
        xz.T,
        origin="lower",
        extent=[x.min(), x.max(), x.min(), x.max()],
        cmap=cmap_2d,
        vmin=vmin2d,
        vmax=vmax2d,
    )
    ax_xz.set_title(f"{title_prefix}: x-z (y=0)")
    ax_xz.set_xlabel("x (a.u.)")
    ax_xz.set_ylabel("z (a.u.)")
    ax_xz.set_aspect("equal")
    cbar1 = fig.colorbar(im1, ax=ax_xz, fraction=0.046, pad=0.04)
    cbar1.set_label("log10(probability density)")

    im2 = ax_zy.imshow(
        zy.T,
        origin="lower",
        extent=[x.min(), x.max(), x.min(), x.max()],
        cmap=cmap_2d,
        vmin=vmin2d,
        vmax=vmax2d,
    )
    ax_zy.set_title(f"{title_prefix}: z-y (x=0)")
    ax_zy.set_xlabel("z (a.u.)")
    ax_zy.set_ylabel("y (a.u.)")
    ax_zy.set_aspect("equal")
    cbar2 = fig.colorbar(im2, ax=ax_zy, fraction=0.046, pad=0.04)
    cbar2.set_label("log10(probability density)")

    im3 = ax_yx.imshow(
        yx.T,
        origin="lower",
        extent=[x.min(), x.max(), x.min(), x.max()],
        cmap=cmap_2d,
        vmin=vmin2d,
        vmax=vmax2d,
    )
    ax_yx.set_title(f"{title_prefix}: y-x (z=0)")
    ax_yx.set_xlabel("y (a.u.)")
    ax_yx.set_ylabel("x (a.u.)")
    ax_yx.set_aspect("equal")
    cbar3 = fig.colorbar(im3, ax=ax_yx, fraction=0.046, pad=0.04)
    cbar3.set_label("log10(probability density)")

    scatter = ax_3d.scatter(
        points[:, 0],
        points[:, 1],
        points[:, 2],
        c=cloud_values,
        s=2,
        cmap=cmap_3d,
        vmin=vmin3d,
        vmax=vmax3d,
        alpha=0.7,
    )
    ax_3d.set_title(f"{title_prefix}: 3D cloud")
    ax_3d.set_xlabel("x (a.u.)")
    ax_3d.set_ylabel("y (a.u.)")
    ax_3d.set_zlabel("z (a.u.)")
    if xyz_limits is not None:
        (x_min, x_max), (y_min, y_max), (z_min, z_max) = xyz_limits
    elif points.shape[0] > 0:
        x_min, x_max = float(points[:, 0].min()), float(points[:, 0].max())
        y_max = float(points[:, 1].max())
        z_min, z_max = float(points[:, 2].min()), float(points[:, 2].max())
    else:
        x_min = y_min = z_min = float(x.min())
        x_max = y_max = z_max = float(x.max())
    ax_3d.set_xlim(x_min, x_max)
    ax_3d.set_ylim(-y_max, y_max)
    ax_3d.set_zlim(z_min, z_max)
    ax_3d.set_box_aspect((1, 1, 1))
    cbar4 = fig.colorbar(scatter, ax=ax_3d, fraction=0.046, pad=0.08)
    cbar4.set_label("log10(probability density)")


def draw_two_state_comparison(
    fig: plt.Figure,
    packed_left: dict[str, np.ndarray],
    packed_right: dict[str, np.ndarray],
    x: np.ndarray,
    left_title: str,
    right_title: str,
    suptitle: str,
    cmap_2d: str = "magma",
    cmap_3d: str | None = None,
    lock_colorbars: bool = True,
    log2d_limits: tuple[float, float] | None = None,
    log3d_limits: tuple[float, float] | None = None,
) -> None:
    fig.clf()

    ax00 = fig.add_subplot(2, 4, 1)
    ax01 = fig.add_subplot(2, 4, 2)
    ax10 = fig.add_subplot(2, 4, 5)
    ax11 = fig.add_subplot(2, 4, 6, projection="3d")

    ax02 = fig.add_subplot(2, 4, 3)
    ax03 = fig.add_subplot(2, 4, 4)
    ax12 = fig.add_subplot(2, 4, 7)
    ax13 = fig.add_subplot(2, 4, 8, projection="3d")

    shared_2d = log2d_limits
    shared_3d = log3d_limits
    if lock_colorbars:
        if shared_2d is None:
            cached_2d = getattr(fig, "_helium_two_state_log2d_limits", None)
            if cached_2d is None:
                all_2d = np.concatenate([_packed_2d_values(packed_left), _packed_2d_values(packed_right)])
                cached_2d = _robust_limits(all_2d)
                setattr(fig, "_helium_two_state_log2d_limits", cached_2d)
            shared_2d = cached_2d
        if shared_3d is None:
            cached_3d = getattr(fig, "_helium_two_state_log3d_limits", None)
            if cached_3d is None:
                left_cloud = _packed_3d_values(packed_left)
                right_cloud = _packed_3d_values(packed_right)
                if left_cloud.size > 0 and right_cloud.size > 0:
                    cloud = np.concatenate([left_cloud, right_cloud])
                elif left_cloud.size > 0:
                    cloud = left_cloud
                elif right_cloud.size > 0:
                    cloud = right_cloud
                else:
                    cloud = np.array([-1.0, 1.0])
                cached_3d = _robust_limits(cloud)
                setattr(fig, "_helium_two_state_log3d_limits", cached_3d)
            shared_3d = cached_3d

    _plot_block(
        fig,
        ax00,
        ax01,
        ax10,
        ax11,
        packed_left,
        x,
        left_title,
        cmap_2d=cmap_2d,
        cmap_3d=cmap_3d,
        log2d_limits=shared_2d,
        log3d_limits=shared_3d,
    )
    _plot_block(
        fig,
        ax02,
        ax03,
        ax12,
        ax13,
        packed_right,
        x,
        right_title,
        cmap_2d=cmap_2d,
        cmap_3d=cmap_3d,
        log2d_limits=shared_2d,
        log3d_limits=shared_3d,
    )

    fig.suptitle(suptitle, fontsize=15)

def draw_single_state_2x2(
    fig: plt.Figure,
    packed: dict[str, np.ndarray],
    x: np.ndarray,
    title: str,
    suptitle: str,
    cmap_2d: str = "magma",
    cmap_3d: str | None = None,
    log2d_limits: tuple[float, float] | None = None,
    log3d_limits: tuple[float, float] | None = None,
    xyz_limits: tuple[tuple[float, float], tuple[float, float], tuple[float, float]] | None = None,
    lock_colorbars: bool = True,
) -> None:
    fig.clf()

    ax1 = fig.add_subplot(2, 2, 1)
    ax2 = fig.add_subplot(2, 2, 2)
    ax3 = fig.add_subplot(2, 2, 3)
    ax4 = fig.add_subplot(2, 2, 4, projection="3d")

    shared_2d = log2d_limits
    shared_3d = log3d_limits
    if lock_colorbars:
        if shared_2d is None:
            cached_2d = getattr(fig, "_helium_single_2x2_log2d_limits", None)
            if cached_2d is None:
                cached_2d = _robust_limits(_packed_2d_values(packed))
                setattr(fig, "_helium_single_2x2_log2d_limits", cached_2d)
            shared_2d = cached_2d
        if shared_3d is None:
            cached_3d = getattr(fig, "_helium_single_2x2_log3d_limits", None)
            if cached_3d is None:
                cloud = _packed_3d_values(packed)
                if cloud.size == 0:
                    cloud = np.array([-1.0, 1.0])
                cached_3d = _robust_limits(cloud)
                setattr(fig, "_helium_single_2x2_log3d_limits", cached_3d)
            shared_3d = cached_3d

    _plot_block(
        fig,
        ax1,
        ax2,
        ax3,
        ax4,
        packed,
        x,
        title,
        cmap_2d=cmap_2d,
        cmap_3d=cmap_3d,
        log2d_limits=shared_2d,
        log3d_limits=shared_3d,
        xyz_limits=xyz_limits,
    )
    fig.suptitle(suptitle, fontsize=15)


def save_figure(fig: plt.Figure, out_path: Path, dpi: int = 300) -> Path:
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, dpi=dpi, bbox_inches="tight")
    return out_path