QPI 与费米面 nesting 波矢的关系

基本物理图像

QPI，全称为 quasiparticle interference，即准粒子干涉。它的基本物理图像是：在金属或其他具有准粒子激发的体系中，杂质会把一个动量态上的准粒子散射到另一个动量态。入射波和散射波在实空间发生干涉，从而导致局域态密度出现空间振荡。没有杂质时，体系具有平移对称性，准粒子态可以用动量 $\mathbf{k}$ 标记，局域态密度在空间中是均匀的。加入一个点杂质后，准粒子可以发生散射

$\mathbf{k}\rightarrow \mathbf{k}+\mathbf{q}.$

入射态和散射态叠加后，会在实空间产生局域态密度调制

$\delta \rho(\mathbf{r},\omega).$

对这个空间调制做傅里叶变换，就得到 QPI 图案：

$\delta \rho(\mathbf{q},\omega).$

因此，QPI 的核心含义是：

$\boxed{ \text{实空间 LDOS 振荡的傅里叶分量，对应准粒子的散射波矢 } \mathbf{q}. }\nonumber$

换句话说，QPI 把杂质诱导的实空间驻波转换成动量空间中的散射信息。一个亮斑或亮线对应某类强散射过程，其波矢为 $\mathbf{q}$。

从 LDOS 修正得到 QPI 响应函数

局域态密度可以由格林函数表示为

$\rho(\mathbf{r},\omega) = -\frac{1}{\pi} \operatorname{Im}G(\mathbf{r},\mathbf{r};\omega).$

有杂质后，总格林函数可以写成

$G = G_0+\delta G.$

在 $T$-matrix 形式下，杂质导致的格林函数修正为

$\delta G = G_0TG_0.$

这里 $G_0$ 是无杂质格林函数，$T$ 是杂质散射矩阵。于是 LDOS 的变化为

$\delta \rho(\mathbf{r},\omega) = -\frac{1}{\pi} \operatorname{Im} \delta G(\mathbf{r},\mathbf{r};\omega).$

对 $\delta \rho(\mathbf{r},\omega)$ 做傅里叶变换，可以得到

$\delta \rho(\mathbf{q},\omega) \propto -\frac{1}{\pi} \operatorname{Im} \sum_{\mathbf{k}} G_0(\mathbf{k}+\mathbf{q},\omega) T(\omega) G_0(\mathbf{k},\omega).$

对于最简单的弱非磁性点杂质，在 Born 近似下有

$T(\omega)\approx V_0.$

因此

$\delta \rho(\mathbf{q},\omega) \propto -\frac{V_0}{\pi} \operatorname{Im} \sum_{\mathbf{k}} G_0(\mathbf{k}+\mathbf{q},\omega) G_0(\mathbf{k},\omega).$

这个公式已经显示出 QPI 的本质：它是一个 particle-hole 型响应函数，描述准粒子从 $\mathbf{k}$ 散射到 $\mathbf{k}+\mathbf{q}$ 的过程。

谱函数近似与 JDOS 图像

为了更直接地看出 QPI 和费米面 nesting 的关系，引入单粒子谱函数

$A(\mathbf{k},\omega) = -\frac{1}{\pi} \operatorname{Im} G_0(\mathbf{k},\omega).$

在很多直观分析中，QPI 强度可以近似写成 joint density of states，简称 JDOS：

$I_{\mathrm{QPI}}(\mathbf{q},\omega) \sim \sum_{\mathbf{k}} A(\mathbf{k},\omega) A(\mathbf{k}+\mathbf{q},\omega).$

如果准粒子寿命较长，谱函数在能带上非常尖锐，可以近似为

$A(\mathbf{k},\omega) \approx \delta(\omega-\epsilon_{\mathbf{k}}).$

于是

$I_{\mathrm{QPI}}(\mathbf{q},\omega) \sim \sum_{\mathbf{k}} \delta(\omega-\epsilon_{\mathbf{k}}) \delta(\omega-\epsilon_{\mathbf{k}+\mathbf{q}}).$

这就是 JDOS：

$J(\mathbf{q},\omega) = \sum_{\mathbf{k}} \delta(\omega-\epsilon_{\mathbf{k}}) \delta(\omega-\epsilon_{\mathbf{k}+\mathbf{q}}).$

这个表达式的物理意义非常清楚：对于一个给定的散射波矢 $\mathbf{q}$，它统计有多少个动量点 $\mathbf{k}$ 同时满足

$\epsilon_{\mathbf{k}}=\omega,$

以及

$\epsilon_{\mathbf{k}+\mathbf{q}}=\omega.$

也就是说，$\mathbf{q}$ 必须连接同一能量上的两个态。

当 $\omega=E_F$ 时，上述条件变成

$\epsilon_{\mathbf{k}}=E_F,$ $\epsilon_{\mathbf{k}+\mathbf{q}}=E_F.$

这说明 $\mathbf{k}$ 和 $\mathbf{k}+\mathbf{q}$ 都在费米面上。因此：

$\boxed{ \text{QPI 在费米能附近探测的是连接费米面上两点的散射波矢。} }\nonumber$

nesting 波矢增强 QPI

费米面 nesting 指的是存在一个波矢 $\mathbf{q}_N$，可以把一段费米面平移到另一段费米面上，使两段费米面大面积重合或近似平行：

$C_2 \approx C_1+\mathbf{q}_N.$

如果 $\mathbf{q}$ 只是连接费米面上的少数两个点，那么 JDOS 有贡献，但不一定很强。相反，如果 $\mathbf{q}=\mathbf{q}_N$ 能连接两段近似平行的费米面，那么很多 $\mathbf{k}$ 都会同时满足

$\epsilon_{\mathbf{k}}=E_F,\qquad \epsilon_{\mathbf{k}+\mathbf{q}_N}=E_F.$

因此

$J(\mathbf{q}_N,E_F) = \sum_{\mathbf{k}} \delta(E_F-\epsilon_{\mathbf{k}}) \delta(E_F-\epsilon_{\mathbf{k}+\mathbf{q}_N})$

会显著增强。所以 QPI 和 nesting 的关系可以概括为：

$\boxed{ \text{费米面 nesting 波矢 } \mathbf{q}_N \text{ 会在 QPI 中表现为强峰、亮线或增强轮廓。} }\nonumber$

更严格的奇异性条件

从连续形式出发，

$J(\mathbf{q},\omega) = \int \frac{d^2 k}{(2\pi)^2} \delta(\omega-\epsilon_{\mathbf{k}}) \delta(\omega-\epsilon_{\mathbf{k}+\mathbf{q}}).$

两个 $\delta$ 函数要求 $\mathbf{k}$ 同时位于两条等能线上：

$C_1:\epsilon_{\mathbf{k}}=\omega,\qquad C_2:\epsilon_{\mathbf{k}+\mathbf{q}}=\omega.$

用变量变换可以得到

$J(\mathbf{q},\omega) \sim \sum_{\mathbf{k}_0} \frac{1}{ \left| \nabla_{\mathbf{k}}\epsilon_{\mathbf{k}_0} \times \nabla_{\mathbf{k}}\epsilon_{\mathbf{k}_0+\mathbf{q}} \right| }.$

其中 $\mathbf{k}_0$ 满足

$\epsilon_{\mathbf{k}_0}=\omega,\qquad \epsilon_{\mathbf{k}_0+\mathbf{q}}=\omega.$

当

$\nabla_{\mathbf{k}}\epsilon_{\mathbf{k}_0} \parallel \nabla_{\mathbf{k}}\epsilon_{\mathbf{k}_0+\mathbf{q}},$

也就是

$\nabla_{\mathbf{k}}\epsilon_{\mathbf{k}_0} \times \nabla_{\mathbf{k}}\epsilon_{\mathbf{k}_0+\mathbf{q}} = 0$

时，分母变小，QPI/JDOS 增强。因为

$\mathbf{v}_{\mathbf{k}} = \nabla_{\mathbf{k}}\epsilon_{\mathbf{k}},$

所以这个条件也可以写成

$\mathbf{v}_{\mathbf{k}_0} \parallel \mathbf{v}_{\mathbf{k}_0+\mathbf{q}}.$

这意味着散射波矢 $\mathbf{q}$ 连接两个费米速度平行或反平行的等能面点。对于二维等能线而言，这等价于两条等能线在相应位置具有平行切线。

因此，QPI 的奇异增强条件是：

$\boxed{ \epsilon_{\mathbf{k}}=\omega,\quad \epsilon_{\mathbf{k}+\mathbf{q}}=\omega,\quad \mathbf{v}_{\mathbf{k}}\parallel \mathbf{v}_{\mathbf{k}+\mathbf{q}}. }$

当 $\omega=E_F$ 时，这正是费米面 nesting 的几何条件。

线积分形式：QPI 是费米面与其平移副本的重叠

从 JDOS 公式出发：

$J(\mathbf{q},\omega) = \int \frac{d^2 k}{(2\pi)^2} \delta(\omega-\epsilon_{\mathbf{k}}) \delta(\omega-\epsilon_{\mathbf{k}+\mathbf{q}}).$

先利用第一个 $\delta$ 函数把二维积分限制到等能线 $C_\omega$ 上。利用恒等式

$\int d^2k\,\delta(\omega-\epsilon_{\mathbf{k}}) = \oint_{C_\omega} \frac{dl}{|\mathbf{v}_{\mathbf{k}}|},$

可以得到

$J(\mathbf{q},\omega) = \oint_{C_\omega} \frac{dl}{(2\pi)^2|\mathbf{v}_{\mathbf{k}}|} \delta(\omega-\epsilon_{\mathbf{k}+\mathbf{q}}).$

这个表达式的含义是：沿着等能线 $C_\omega$ 扫描每一个点 $\mathbf{k}$，检查平移 $\mathbf{q}$ 之后的点 $\mathbf{k}+\mathbf{q}$ 是否仍然位于同一个等能线上。如果 $\mathbf{q}$ 是 nesting 波矢，那么一整段等能线经过 $\mathbf{q}$ 平移后仍然接近另一段等能线。因此这个线积分会累积大量贡献。

所以：

$\boxed{ J(\mathbf{q},\omega) \text{ 测量的是等能线与其平移副本之间的重叠程度。} }\nonumber$

在费米能处：

$\boxed{ J(\mathbf{q},E_F) \text{ 测量的是费米面与其平移副本之间的重叠程度。} }\nonumber$

这就是 QPI 和费米面 nesting 关系最直观的数学表达。

例子：二维圆形费米面

考虑二维自由电子模型：

$\epsilon_{\mathbf{k}} = \frac{k^2}{2m}-E_F.$

费米面为圆：

$|\mathbf{k}|=k_F.$

QPI/JDOS 为

$J(\mathbf{q},E_F) = \int d^2k\, \delta(k^2-k_F^2) \delta(|\mathbf{k}+\mathbf{q}|^2-k_F^2).$

该积分在$|\mathbf{q}|\leq 2k_F$时有贡献，并且在$|\mathbf{q}|=2k_F$附近增强。这是因为 $|\mathbf{q}|=2k_F$ 对应 backscattering：

$\mathbf{k}\rightarrow -\mathbf{k}.$

因此普通二维金属的 QPI 常常出现 $2k_F$ 特征。这个 $2k_F$ 可以看作最简单的费米面自 nesting 波矢。

实际 QPI 不等于纯 JDOS

虽然 JDOS 给出了 QPI 和 nesting 的基本关系，但真实 QPI 还包含散射矩阵元。更完整地说，

$I_{\mathrm{QPI}}(\mathbf{q},\omega) \sim \sum_{\mathbf{k}} M(\mathbf{k},\mathbf{k}+\mathbf{q}) A(\mathbf{k},\omega) A(\mathbf{k}+\mathbf{q},\omega).$

其中 $M(\mathbf{k},\mathbf{k}+\mathbf{q})$ 是矩阵元因子，可能来自：

自旋重叠；
轨道重叠；
杂质类型；
超导 coherence factor；
STM tip 的选择性。

因此，费米面 nesting 决定哪些 $\mathbf{q}$ 在几何上可能增强，而矩阵元决定这些几何允许的散射是否真的在 QPI 图案中显著。

可以总结为：

$\boxed{ \text{nesting 决定几何相空间，矩阵元决定可见强度。} }\nonumber$

如果一个 nesting 波矢连接的两个费米面态自旋完全相反，而杂质又是非磁性的，那么该散射可能被自旋选择规则压制。此时即使几何上存在 nesting，QPI 中也未必出现强峰。

QPI 与费米面 nesting 的最终关系

从公式链条看，QPI 与费米面 nesting 的关系可以总结为：

$\delta \rho(\mathbf{q},\omega) \sim \sum_{\mathbf{k}} G_0(\mathbf{k}+\mathbf{q},\omega) G_0(\mathbf{k},\omega).$

在谱函数近似下：

$I_{\mathrm{QPI}}(\mathbf{q},\omega) \sim \sum_{\mathbf{k}} A(\mathbf{k},\omega) A(\mathbf{k}+\mathbf{q},\omega).$

在长寿命极限：

$I_{\mathrm{QPI}}(\mathbf{q},\omega) \sim \sum_{\mathbf{k}} \delta(\omega-\epsilon_{\mathbf{k}}) \delta(\omega-\epsilon_{\mathbf{k}+\mathbf{q}}).$

在费米能处：

$I_{\mathrm{QPI}}(\mathbf{q},E_F) \sim \sum_{\mathbf{k}} \delta(E_F-\epsilon_{\mathbf{k}}) \delta(E_F-\epsilon_{\mathbf{k}+\mathbf{q}}).$

因此，QPI 的强峰出现在能够连接大量费米面态的波矢上：

$\boxed{ \mathbf{q}\approx \mathbf{q}_N, \qquad \epsilon_{\mathbf{k}}=E_F, \qquad \epsilon_{\mathbf{k}+\mathbf{q}_N}=E_F. }$

更严格的增强条件为：

$\boxed{ \nabla_{\mathbf{k}}\epsilon_{\mathbf{k}} \parallel \nabla_{\mathbf{k}}\epsilon_{\mathbf{k}+\mathbf{q}_N}. }$

这正是费米面 nesting 的几何条件。

Summary

QPI 是杂质散射导致的 LDOS 干涉图案。其傅里叶分量 $\delta\rho(\mathbf{q},\omega)$ 对应准粒子从 $\mathbf{k}$ 到 $\mathbf{k}+\mathbf{q}$ 的散射过程。在弱散射和长寿命近似下，QPI 强度近似为等能面自相关，即 JDOS。当 $\omega=E_F$ 时，这个自相关就是费米面与其平移副本的重叠。因此，费米面 nesting 波矢会在 QPI 中表现为增强的峰、亮线或特征轮廓。

代码演示

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
QPI and JDOS nesting demo.

This script separately calculates:

1. JDOS approximation:
       JDOS(q, omega) = sum_k A(k, omega) A(k+q, omega)

2. Born-QPI response:
       Lambda(q, omega) = sum_k G(k+q, omega) G(k, omega)
       QPI(q, omega)    = -Im Lambda(q, omega) / pi

Model:
       epsilon(k) = -2t (cos kx + cos ky) - mu

At mu = 0, the square-lattice Fermi surface is nested by
       Q_N = (pi, pi).

Outputs:
    data_qpi_jdos/
        fs_points_mu_*.dat
        jdos_mu_*.dat
        qpi_mu_*.dat
        linecut_jdos_mu_*.dat
        linecut_qpi_mu_*.dat

    fig_qpi_jdos/
        fs_mu_*.png
        spectral_mu_*.png
        jdos_mu_*.png
        qpi_mu_*.png
        linecut_jdos_mu_*.png
        linecut_qpi_mu_*.png
        comparison_jdos_linecut.png
        comparison_qpi_linecut.png
"""

import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib.colors import LogNorm, TwoSlopeNorm


# ============================================================
# 1. Global settings
# ============================================================

USE_LATEX = True
FIG_DPI = 400

DATA_DIR = "data_qpi_jdos"
FIG_DIR = "fig_qpi_jdos"

os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(FIG_DIR, exist_ok=True)

plt.rcParams["figure.facecolor"] = "white"
plt.rcParams["axes.facecolor"] = "white"
plt.rcParams["savefig.facecolor"] = "white"

plt.rcParams["font.family"] = "serif"
plt.rcParams["font.serif"] = ["Times New Roman", "Times", "DejaVu Serif"]
plt.rcParams["mathtext.fontset"] = "stix"

plt.rcParams["text.usetex"] = USE_LATEX
if USE_LATEX:
    plt.rcParams["text.latex.preamble"] = r"""
    \usepackage{amsmath}
    \usepackage{mathptmx}
    """

plt.rcParams["axes.linewidth"] = 1.0
plt.rcParams["xtick.direction"] = "in"
plt.rcParams["ytick.direction"] = "in"
plt.rcParams["xtick.top"] = True
plt.rcParams["ytick.right"] = True

plt.rcParams["font.size"] = 14
plt.rcParams["axes.labelsize"] = 16
plt.rcParams["legend.fontsize"] = 12
plt.rcParams["xtick.labelsize"] = 13
plt.rcParams["ytick.labelsize"] = 13


# ============================================================
# 2. Parameters
# ============================================================

N_K = 512
T_HOP = 1.0

OMEGA = 0.0
ETA = 0.035

FS_CUT = 0.03

# mu=0: nested Fermi surface
# mu=-0.8: away from perfect nesting
MU_LIST = [0.0, -0.8]

USE_LOG_JDOS = True


# ============================================================
# 3. Basic helpers
# ============================================================

def savefig(fig, filename):
    path = os.path.join(FIG_DIR, filename)
    fig.savefig(path, dpi=FIG_DPI, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved: {path}")


def mu_tag(mu):
    if abs(mu) < 1e-12:
        return "0p0"
    return f"{mu:+.2f}".replace("+", "p").replace("-", "m").replace(".", "p")


def pi_ticks(ax):
    ax.set_xlim(-np.pi, np.pi)
    ax.set_ylim(-np.pi, np.pi)
    ax.set_xticks([-np.pi, 0.0, np.pi])
    ax.set_yticks([-np.pi, 0.0, np.pi])
    ax.set_xticklabels([r"$-\pi$", r"$0$", r"$\pi$"])
    ax.set_yticklabels([r"$-\pi$", r"$0$", r"$\pi$"])


def max_three_ticks(ax):
    ax.xaxis.set_major_locator(MaxNLocator(nbins=3))
    ax.yaxis.set_major_locator(MaxNLocator(nbins=3))


def max_three_cbar_ticks(cbar):
    cbar.ax.yaxis.set_major_locator(MaxNLocator(nbins=3))
    cbar.update_ticks()


def nesting_corner_points():
    return [-np.pi, np.pi, -np.pi, np.pi], [-np.pi, -np.pi, np.pi, np.pi]


# ============================================================
# 4. Model
# ============================================================

def square_lattice_energy(kx, ky, mu, t=T_HOP):
    return -2.0 * t * (np.cos(kx) + np.cos(ky)) - mu


def green_function(eps, omega=OMEGA, eta=ETA):
    return 1.0 / (omega - eps + 1j * eta)


def spectral_function(eps, omega=OMEGA, eta=ETA):
    return (eta / np.pi) / ((omega - eps) ** 2 + eta ** 2)


# ============================================================
# 5. JDOS and QPI
# ============================================================

def compute_jdos_fft(A):
    """
    JDOS(q) = sum_k A(k) A(k+q)

    For real A, this periodic autocorrelation is obtained by
        ifft2(|fft2(A)|^2).
    """
    F = np.fft.fft2(A)
    jdos = np.fft.ifft2(np.abs(F) ** 2).real
    jdos /= A.size
    return np.fft.fftshift(jdos)


def reverse_fft_indices(arr):
    """
    Return arr[-p_y, -p_x] on a periodic FFT grid.
    """
    ny, nx = arr.shape
    iy = (-np.arange(ny)) % ny
    ix = (-np.arange(nx)) % nx
    return arr[np.ix_(iy, ix)]


def compute_born_qpi_fft(G):
    """
    Lambda(q) = sum_k G(k+q) G(k).

    This is a periodic correlation without complex conjugation.

    In Fourier space:
        Lambda_hat(p) = G_hat(p) G_hat(-p).

    QPI(q) = -Im Lambda(q) / pi.
    """
    F = np.fft.fft2(G)
    F_neg = reverse_fft_indices(F)

    lam = np.fft.ifft2(F * F_neg)
    lam /= G.size
    lam = np.fft.fftshift(lam)

    qpi = -np.imag(lam) / np.pi

    return lam, qpi


# ============================================================
# 6. Save data
# ============================================================

def save_fs_points(kx, ky, eps, mu):
    tag = mu_tag(mu)
    mask = np.abs(eps - OMEGA) < FS_CUT

    data = np.column_stack([
        kx[mask],
        ky[mask],
        eps[mask],
    ])

    path = os.path.join(DATA_DIR, f"fs_points_mu_{tag}.dat")
    np.savetxt(path, data, header="kx ky epsilon")
    print(f"Saved: {path}")


def save_jdos_data(qx, qy, jdos, mu):
    tag = mu_tag(mu)
    QX, QY = np.meshgrid(qx, qy, indexing="xy")

    data = np.column_stack([
        QX.ravel(),
        QY.ravel(),
        jdos.ravel(),
    ])

    path = os.path.join(DATA_DIR, f"jdos_mu_{tag}.dat")
    np.savetxt(path, data, header="qx qy JDOS")
    print(f"Saved: {path}")


def save_qpi_data(qx, qy, lam, qpi, mu):
    tag = mu_tag(mu)
    QX, QY = np.meshgrid(qx, qy, indexing="xy")

    data = np.column_stack([
        QX.ravel(),
        QY.ravel(),
        np.real(lam).ravel(),
        np.imag(lam).ravel(),
        qpi.ravel(),
        np.abs(qpi).ravel(),
    ])

    path = os.path.join(DATA_DIR, f"qpi_mu_{tag}.dat")
    header = "qx qy Re_Lambda Im_Lambda QPI_minus_ImLambda_over_pi abs_QPI"
    np.savetxt(path, data, header=header)
    print(f"Saved: {path}")


def save_linecuts(qvals, jdos_cut, qpi_cut, mu):
    tag = mu_tag(mu)

    data_jdos = np.column_stack([qvals, jdos_cut])
    path_jdos = os.path.join(DATA_DIR, f"linecut_jdos_mu_{tag}.dat")
    np.savetxt(path_jdos, data_jdos, header="qx JDOS(qx, qy=-pi)")
    print(f"Saved: {path_jdos}")

    data_qpi = np.column_stack([qvals, qpi_cut])
    path_qpi = os.path.join(DATA_DIR, f"linecut_qpi_mu_{tag}.dat")
    np.savetxt(path_qpi, data_qpi, header="qx QPI(qx, qy=-pi)")
    print(f"Saved: {path_qpi}")


# ============================================================
# 7. Plotting: FS and spectral function
# ============================================================

def plot_fermi_surface(kx, ky, eps, mu):
    tag = mu_tag(mu)

    fig, ax = plt.subplots(figsize=(5.2, 4.6))

    im = ax.imshow(
        eps,
        origin="lower",
        extent=(-np.pi, np.pi, -np.pi, np.pi),
        cmap="RdBu_r",
        aspect="equal",
    )

    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label(r"$\epsilon_{\mathbf{k}}$")
    max_three_cbar_ticks(cbar)

    ax.contour(kx, ky, eps, levels=[OMEGA], linewidths=1.8)

    ax.set_xlabel(r"$k_x$")
    ax.set_ylabel(r"$k_y$")
    ax.set_title(rf"Fermi surface, $\mu={mu:.2f}$")
    pi_ticks(ax)
    ax.set_aspect("equal", adjustable="box")

    savefig(fig, f"fs_mu_{tag}.png")


def plot_spectral_function(A, mu):
    tag = mu_tag(mu)

    fig, ax = plt.subplots(figsize=(5.2, 4.6))

    im = ax.imshow(
        A,
        origin="lower",
        extent=(-np.pi, np.pi, -np.pi, np.pi),
        cmap="viridis",
        aspect="equal",
    )

    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label(r"$A(\mathbf{k},\omega)$")
    max_three_cbar_ticks(cbar)

    ax.set_xlabel(r"$k_x$")
    ax.set_ylabel(r"$k_y$")
    ax.set_title(rf"Spectral function, $\mu={mu:.2f}$")
    pi_ticks(ax)
    ax.set_aspect("equal", adjustable="box")

    savefig(fig, f"spectral_mu_{tag}.png")


# ============================================================
# 8. Plotting: JDOS and QPI separately
# ============================================================

def plot_jdos(jdos, mu):
    tag = mu_tag(mu)

    fig, ax = plt.subplots(figsize=(5.2, 4.6))

    if USE_LOG_JDOS:
        positive = jdos[jdos > 0]
        vmin = np.percentile(positive, 5)
        vmax = np.percentile(positive, 99.5)

        im = ax.imshow(
            jdos,
            origin="lower",
            extent=(-np.pi, np.pi, -np.pi, np.pi),
            cmap="magma",
            norm=LogNorm(vmin=max(vmin, 1e-14), vmax=vmax),
            aspect="equal",
        )
    else:
        im = ax.imshow(
            jdos,
            origin="lower",
            extent=(-np.pi, np.pi, -np.pi, np.pi),
            cmap="magma",
            aspect="equal",
        )

    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label(r"$J_{\mathrm{JDOS}}(\mathbf{q},\omega)$")
    max_three_cbar_ticks(cbar)

    xq, yq = nesting_corner_points()
    ax.scatter(
        xq,
        yq,
        marker="x",
        s=60,
        linewidths=1.4,
        label=r"$\mathbf{Q}_N=(\pi,\pi)$",
    )

    ax.set_xlabel(r"$q_x$")
    ax.set_ylabel(r"$q_y$")
    ax.set_title(rf"JDOS, $\mu={mu:.2f}$")
    pi_ticks(ax)
    ax.set_aspect("equal", adjustable="box")
    ax.legend(frameon=True, loc="lower left")

    savefig(fig, f"jdos_mu_{tag}.png")


def plot_qpi(qpi, mu):
    tag = mu_tag(mu)

    fig, ax = plt.subplots(figsize=(5.2, 4.6))

    vmax = np.nanpercentile(np.abs(qpi), 99.5)
    if vmax <= 0 or not np.isfinite(vmax):
        vmax = 1.0

    norm = TwoSlopeNorm(vmin=-vmax, vcenter=0.0, vmax=vmax)

    im = ax.imshow(
        qpi,
        origin="lower",
        extent=(-np.pi, np.pi, -np.pi, np.pi),
        cmap="RdBu_r",
        norm=norm,
        aspect="equal",
    )

    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label(r"$\delta\rho(\mathbf{q},\omega)$")
    max_three_cbar_ticks(cbar)

    xq, yq = nesting_corner_points()
    ax.scatter(
        xq,
        yq,
        marker="x",
        s=60,
        linewidths=1.4,
        label=r"$\mathbf{Q}_N=(\pi,\pi)$",
    )

    ax.set_xlabel(r"$q_x$")
    ax.set_ylabel(r"$q_y$")
    ax.set_title(rf"Born-QPI, $\mu={mu:.2f}$")
    pi_ticks(ax)
    ax.set_aspect("equal", adjustable="box")
    ax.legend(frameon=True, loc="lower left")

    savefig(fig, f"qpi_mu_{tag}.png")


def plot_linecut_jdos(qvals, cut, mu):
    tag = mu_tag(mu)

    fig, ax = plt.subplots(figsize=(5.2, 3.8))

    ax.plot(qvals, cut, lw=1.8)

    ax.set_xlabel(r"$q_x$")
    ax.set_ylabel(r"$J_{\mathrm{JDOS}}(q_x,-\pi)$")
    ax.set_title(rf"JDOS line cut, $\mu={mu:.2f}$")

    ax.set_xlim(-np.pi, np.pi)
    ax.set_xticks([-np.pi, 0.0, np.pi])
    ax.set_xticklabels([r"$-\pi$", r"$0$", r"$\pi$"])
    ax.yaxis.set_major_locator(MaxNLocator(nbins=3))

    savefig(fig, f"linecut_jdos_mu_{tag}.png")


def plot_linecut_qpi(qvals, cut, mu):
    tag = mu_tag(mu)

    fig, ax = plt.subplots(figsize=(5.2, 3.8))

    ax.plot(qvals, cut, lw=1.8)

    ax.set_xlabel(r"$q_x$")
    ax.set_ylabel(r"$\delta\rho(q_x,-\pi)$")
    ax.set_title(rf"Born-QPI line cut, $\mu={mu:.2f}$")

    ax.set_xlim(-np.pi, np.pi)
    ax.set_xticks([-np.pi, 0.0, np.pi])
    ax.set_xticklabels([r"$-\pi$", r"$0$", r"$\pi$"])
    ax.yaxis.set_major_locator(MaxNLocator(nbins=3))

    savefig(fig, f"linecut_qpi_mu_{tag}.png")


# ============================================================
# 9. Comparison plots
# ============================================================

def plot_comparison_linecuts(all_results):
    fig, ax = plt.subplots(figsize=(5.4, 4.0))

    for mu, result in all_results.items():
        qvals = result["qvals"]
        cut = result["jdos_linecut"]
        cut_norm = cut / np.max(cut)

        ax.plot(
            qvals,
            cut_norm,
            lw=1.8,
            label=rf"$\mu={mu:.2f}$",
        )

    ax.set_xlabel(r"$q_x$")
    ax.set_ylabel(r"Normalized $J_{\mathrm{JDOS}}(q_x,-\pi)$")
    ax.set_title(r"JDOS nesting enhancement")

    ax.set_xlim(-np.pi, np.pi)
    ax.set_xticks([-np.pi, 0.0, np.pi])
    ax.set_xticklabels([r"$-\pi$", r"$0$", r"$\pi$"])
    ax.yaxis.set_major_locator(MaxNLocator(nbins=3))
    ax.legend(frameon=True)

    savefig(fig, "comparison_jdos_linecut.png")

    fig, ax = plt.subplots(figsize=(5.4, 4.0))

    for mu, result in all_results.items():
        qvals = result["qvals"]
        cut = result["qpi_linecut"]

        max_abs = np.max(np.abs(cut))
        if max_abs > 0:
            cut_norm = cut / max_abs
        else:
            cut_norm = cut

        ax.plot(
            qvals,
            cut_norm,
            lw=1.8,
            label=rf"$\mu={mu:.2f}$",
        )

    ax.set_xlabel(r"$q_x$")
    ax.set_ylabel(r"Normalized $\delta\rho(q_x,-\pi)$")
    ax.set_title(r"Born-QPI nesting feature")

    ax.set_xlim(-np.pi, np.pi)
    ax.set_xticks([-np.pi, 0.0, np.pi])
    ax.set_xticklabels([r"$-\pi$", r"$0$", r"$\pi$"])
    ax.yaxis.set_major_locator(MaxNLocator(nbins=3))
    ax.legend(frameon=True)

    savefig(fig, "comparison_qpi_linecut.png")


# ============================================================
# 10. Main calculation
# ============================================================

def run_case(mu):
    print("=" * 60)
    print(f"Running JDOS and Born-QPI for mu = {mu:.4f}")

    k = np.linspace(-np.pi, np.pi, N_K, endpoint=False)
    q = np.linspace(-np.pi, np.pi, N_K, endpoint=False)

    KX, KY = np.meshgrid(k, k, indexing="xy")

    eps = square_lattice_energy(KX, KY, mu)
    G = green_function(eps)
    A = spectral_function(eps)

    jdos = compute_jdos_fft(A)
    lam, qpi = compute_born_qpi_fft(G)

    # q_y = -pi line cut. Q=(pi,pi) appears at BZ corners.
    qy_index = 0
    jdos_linecut = jdos[qy_index, :]
    qpi_linecut = qpi[qy_index, :]

    save_fs_points(KX, KY, eps, mu)
    save_jdos_data(q, q, jdos, mu)
    save_qpi_data(q, q, lam, qpi, mu)
    save_linecuts(q, jdos_linecut, qpi_linecut, mu)

    plot_fermi_surface(KX, KY, eps, mu)
    plot_spectral_function(A, mu)

    plot_jdos(jdos, mu)
    plot_qpi(qpi, mu)

    plot_linecut_jdos(q, jdos_linecut, mu)
    plot_linecut_qpi(q, qpi_linecut, mu)

    return {
        "qvals": q,
        "jdos": jdos,
        "qpi": qpi,
        "jdos_linecut": jdos_linecut,
        "qpi_linecut": qpi_linecut,
    }


def main():
    all_results = {}

    for mu in MU_LIST:
        all_results[mu] = run_case(mu)

    plot_comparison_linecuts(all_results)

    print("=" * 60)
    print("Finished.")
    print(f"Figures saved in: {FIG_DIR}")
    print(f"Data saved in: {DATA_DIR}")


if __name__ == "__main__":
    main()