plot_spectral.py

#!/usr/bin/env python3
"""
Reconstruct and plot the spectral function S(k, omega) from CheMPS moments.

Implements the Chebyshev reconstruction formula (Eq. 5 / 13 of arXiv:1101.5895):

    A^BC(omega) = (2 W'/W*) / (pi * sqrt(1 - omega'^2))
                  * [g_0 * mu_0 + 2 * sum_{n=1}^{N-1} g_n * mu_n * T_n(omega')]

where omega' is the rescaled frequency and g_n are damping factors.

Usage:
    python3 plot_spectral.py [moments_json] [meta_json]

    Default: chemps_moment.json and chemps_meta.json
"""

import json
import sys
import numpy as np
import matplotlib.pyplot as plt


def jackson_damping(n, N):
    """Jackson damping factors g_n^J, Eq. (15) of arXiv:1101.5895."""
    return (
        (N - n + 1) * np.cos(np.pi * n / (N + 1))
        + np.sin(np.pi * n / (N + 1)) / np.tan(np.pi / (N + 1))
    ) / (N + 1)


def lorentz_damping(n, N, lam=4.0):
    """Lorentz damping factors g_n^L, Eq. (16) of arXiv:1101.5895."""
    return np.sinh(lam * (1 - n / N)) / np.sinh(lam)


def chebyshev_reconstruct(omega_prime, mu_n, damping="jackson", lam=4.0):
    """
    Reconstruct f(x) = (1/pi) * 1/sqrt(1-x^2) * [g0*mu0 + 2*sum g_n*mu_n*T_n(x)]

    This gives the spectral function in rescaled frequency omega'.
    To convert to physical omega, multiply by 2*W'/(a*W*) = 1/a
    and shift: omega = a * (omega' + W').

    Parameters
    ----------
    omega_prime : array, rescaled frequencies in (-1, 1)
    mu_n : array, Chebyshev moments [mu_0, mu_1, ..., mu_{N-1}]
    damping : "jackson", "lorentz", or "none"
    lam : Lorentz damping parameter (default 4.0)

    Returns
    -------
    f : array, spectral function values at omega_prime
    """
    N = len(mu_n)
    x = omega_prime

    # Compute damping factors
    if damping == "jackson":
        g = np.array([jackson_damping(n, N) for n in range(N)])
    elif damping == "lorentz":
        g = np.array([lorentz_damping(n, N, lam) for n in range(N)])
    else:
        g = np.ones(N)

    # Evaluate Chebyshev polynomials T_n(x) using recurrence
    # T_0(x) = 1, T_1(x) = x, T_{n+1}(x) = 2x T_n(x) - T_{n-1}(x)
    f = np.zeros_like(x, dtype=float)

    T_prev2 = np.ones_like(x)   # T_0
    T_prev1 = x.copy()          # T_1

    # n=0 term
    f += g[0] * mu_n[0].real * T_prev2

    if N > 1:
        # n=1 term
        f += 2.0 * g[1] * mu_n[1].real * T_prev1

    for n in range(2, N):
        T_curr = 2.0 * x * T_prev1 - T_prev2
        f += 2.0 * g[n] * mu_n[n].real * T_curr
        T_prev2 = T_prev1
        T_prev1 = T_curr

    # Multiply by weight function
    # Avoid division by zero at boundaries
    mask = np.abs(x) < 1.0 - 1e-10
    result = np.zeros_like(f)
    result[mask] = f[mask] / (np.pi * np.sqrt(1.0 - x[mask] ** 2))

    return result


def main():
    # Parse arguments
    moments_file = "chemps_moment.json"
    meta_file = "chemps_meta.json"

    if len(sys.argv) > 1:
        moments_file = sys.argv[1]
    if len(sys.argv) > 2:
        meta_file = sys.argv[2]

    # Load moments
    with open(moments_file) as f:
        moments_data = json.load(f)

    # Sort by n
    moments_data.sort(key=lambda x: x["n"])
    mu_n = np.array([m["mu_n_real"] + 1j * m["mu_n_imag"] for m in moments_data])
    N = len(mu_n)
    print(f"Loaded {N} Chebyshev moments from {moments_file}")

    # Load metadata
    try:
        with open(meta_file) as f:
            meta = json.load(f)
        L = meta["L"]
        k_value = meta["k_value"]
        E0 = meta["E0"]
        W_star = meta["W_star"]
        epsilon_t = meta["epsilon_t"]
        W_prime = meta["W_prime"]
        a = meta["a"]
        Dmax = meta["Dmax"]
        k_index = meta["k_index"]
        print(f"Metadata: L={L}, k={k_value:.4f} ({k_value/np.pi:.3f}*pi), "
              f"Dmax={Dmax}, a={a:.4f}, W*={W_star:.4f}")
    except FileNotFoundError:
        print(f"Warning: {meta_file} not found, using defaults")
        L = 20
        W_star = 2 * np.pi
        epsilon_t = 0.025
        W_prime = 1.0 - 0.5 * epsilon_t
        a = W_star / (2.0 * W_prime)
        E0 = 0.0
        k_value = np.pi / 2
        k_index = None
        Dmax = 32

    # Reconstruct spectral function
    # omega' in [-W', W'] maps to omega in [0, W*]
    # omega = a * (omega' + W')
    # omega' = omega / a - W'
    n_pts = 2000
    omega_prime = np.linspace(-W_prime + 1e-6, W_prime - 1e-6, n_pts)
    # excitation energy: omega_exc = a * (omega' + W') >= 0
    omega_exc = a * (omega_prime + W_prime)

    # Reconstruct with different dampings
    S_jackson = chebyshev_reconstruct(omega_prime, mu_n, damping="jackson")
    S_lorentz = chebyshev_reconstruct(omega_prime, mu_n, damping="lorentz", lam=4.0)
    S_undamped = chebyshev_reconstruct(omega_prime, mu_n, damping="none")

    # The factor 2W'/(a*W*) from Eq. (5) cancels with the rescaling:
    # A^BC(omega) = (1/a) * f(omega') where f is what chebyshev_reconstruct returns
    # So S(k, omega) = (1/a) * f(omega')
    S_jackson /= a
    S_lorentz /= a
    S_undamped /= a

    # ==============================
    # Plot: Evolution with N (multiple truncation orders)
    # ==============================
    # List of N values to show
    N_values = [30, 50, 70, 100]
    N_values = [n for n in N_values if n <= N]  # only use available moments
    if N not in N_values:
        N_values.append(N)
    N_values.sort()

    omega_plot = omega_exc / np.pi
    colors = plt.cm.viridis(np.linspace(0.15, 0.85, len(N_values)))

    fig, axes = plt.subplots(2, 1, figsize=(10, 8))

    # --- Panel (a): S(k, omega) for multiple N (Jackson damping) ---
    ax = axes[0]
    for i, N_trunc in enumerate(N_values):
        mu_trunc = mu_n[:N_trunc]
        S_trunc = chebyshev_reconstruct(omega_prime, mu_trunc, damping="jackson")
        S_trunc /= a
        ax.plot(omega_plot, S_trunc, color=colors[i], linewidth=1.5,
                label=f"N={N_trunc}")

    ax.set_xlabel(r"$\omega/\pi$")
    if k_index is not None:
        ax.set_ylabel(rf"$S(k={k_value/np.pi:.2f}\pi, \omega)$")
        ax.set_title(
            rf"Heisenberg chain $L={L}$, $D={Dmax}$, "
            rf"$k={k_value/np.pi:.2f}\pi$ — Jackson damping"
        )
    else:
        ax.set_ylabel(r"$S(k, \omega)$")
        ax.set_title(f"Heisenberg chain L={L}, D={Dmax} — Jackson damping")

    ax.set_xlim(left=0)
    ax.set_ylim(bottom=0)

    # Mark spinon thresholds
    omega_1 = (np.pi / 2.0) * abs(np.sin(k_value))
    omega_2 = np.pi * abs(np.sin(k_value / 2.0))
    ax.axvline(x=omega_1 / np.pi, color="gray", linestyle=":", alpha=0.7,
               label=rf"$\omega_1/\pi$ = {omega_1/np.pi:.3f}")
    ax.axvline(x=omega_2 / np.pi, color="gray", linestyle="--", alpha=0.5,
               label=rf"$\omega_2/\pi$ = {omega_2/np.pi:.3f}")
    ax.legend()

    # --- Panel (b): Chebyshev moments ---
    ax2 = axes[1]
    ns = np.arange(N)
    ax2.semilogy(ns, np.abs(mu_n.real), "b.-", markersize=3, label=r"$|\mathrm{Re}\,\mu_n|$")
    if np.any(np.abs(mu_n.imag) > 1e-15):
        ax2.semilogy(ns, np.abs(mu_n.imag) + 1e-20, "r.-", markersize=3,
                     label=r"$|\mathrm{Im}\,\mu_n|$")
    # Mark the N values used in the top panel
    for N_trunc in N_values:
        ax2.axvline(x=N_trunc - 1, color="gray", linestyle=":", alpha=0.3)

    ax2.set_xlabel("n")
    ax2.set_ylabel(r"$|\mu_n|$")
    ax2.set_title("Chebyshev moments")
    ax2.legend()

    plt.tight_layout()

    output_file = "spectral_function.pdf"
    plt.savefig(output_file, dpi=150, bbox_inches="tight")
    print(f"Saved plot to {output_file}")
    plt.savefig("spectral_function.png", dpi=150, bbox_inches="tight")
    print("Saved plot to spectral_function.png")

    plt.show()


if __name__ == "__main__":
    main()