AQC + 트로터 동역학을 활용한 중성자 산란 시뮬레이션 서버리스 워크플로우

from functools import partial

import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize

import quimb.tensor as qtn
from qiskit import QuantumCircuit
from qiskit.quantum_info import SparsePauliOp
from qiskit_addon_aqc_tensor.simulation import tensornetwork_from_circuit
from qiskit_addon_aqc_tensor.simulation.quimb import QuimbSimulator
from qiskit_ibm_catalog import QiskitServerless

#  Dynamical structure factor via discrete Fourier transform


def get_spectrum(n, Gjjc, dt, time_steps, q_steps, w_steps):
    """Compute the dynamical structure factor from the retarded Green's function.

    Uses the center-site approximation and a discrete Fourier transform.
    """
    green = Gjjc / 4  # sigma -> S=1/2
    omega_max = np.pi / dt
    qpoints = np.arange(0, 2 * np.pi, 2 * np.pi / q_steps)
    omegas = np.arange(0, omega_max, omega_max / w_steps)
    green_map = np.zeros((omegas.shape[0], qpoints.shape[0]))
    center = n // 2 - 1
    for iw, w in enumerate(omegas):
        exponent = np.exp(1j * w * dt * np.arange(1, time_steps + 1))
        S_w = np.dot(green.T, exponent) * dt
        for iq, q in enumerate(qpoints):
            q_matrix = np.exp(-1j * q * np.arange(-center, center + 2, 1))
            green_map[iw, iq] = np.imag(np.dot(S_w, q_matrix))
    return green_map


#  Plotting helpers


def plot_spectrum(
    dsf,
    dt,
    q_steps,
    w_steps,
    lower_bound=False,
    upper_bound=False,
    title=None,
):
    """Heat-map of the dynamical structure factor."""
    omega_max = np.pi / dt
    qpoints = np.arange(0, 2 * np.pi, 2 * np.pi / q_steps)
    omegas = np.arange(0, omega_max, omega_max / w_steps)
    x, y = np.meshgrid(qpoints, omegas)
    fig, ax = plt.subplots(figsize=(8, 5))
    c = ax.pcolormesh(x, y, dsf / np.max(dsf), cmap="viridis", shading="auto")
    fig.colorbar(c, ax=ax, label="Normalized intensity")
    if lower_bound:
        ax.plot(
            qpoints,
            np.pi * np.abs(np.sin(qpoints)) / 2,
            "--",
            color="white",
            lw=1.5,
            label="Lower bound",
        )
    if upper_bound:
        ax.plot(
            qpoints,
            np.pi * np.abs(np.sin(qpoints / 2)),
            "--",
            color="red",
            lw=1.5,
            label="Upper bound",
        )
    ax.set_ylim(0, 3.6)
    ax.set_xlim(0, 2 * np.pi - 2 * np.pi / q_steps)
    ax.set_xlabel(r"$q$", fontsize=16)
    ax.set_ylabel(r"$\tilde{\omega} = \omega / J$", fontsize=16)
    ax.set_xticks([0, np.pi / 2, np.pi, 3 * np.pi / 2, 2 * np.pi])
    ax.set_xticklabels(["0", r"$\pi/2$", r"$\pi$", r"$3\pi/2$", r"$2\pi$"])
    if lower_bound or upper_bound:
        ax.legend(loc="upper right", fontsize=11)
    if title:
        ax.set_title(title, fontsize=14)
    plt.tight_layout()
    plt.show()


def plot_green(n, Gjjc, time_steps, dt, title=None):
    """Heat-map of the retarded Green's function in real space and time."""
    fig, ax = plt.subplots(figsize=(8, 6))
    t_axis = np.arange(1, time_steps + 1) * dt
    site_axis = np.arange(n)
    x, y = np.meshgrid(t_axis, site_axis)
    c = ax.pcolormesh(
        x,
        y,
        np.real(Gjjc).T,
        cmap="RdBu",
        vmax=0.5,
        vmin=-0.5,
        shading="auto",
    )
    fig.colorbar(c, ax=ax, label=r"Re $G^R(j, j_c, t)$")
    ax.set_xlabel(r"Time  ($t / J^{-1}$)", fontsize=16)
    ax.set_ylabel("Site index $j$", fontsize=16)
    if title:
        ax.set_title(title, fontsize=14)
    plt.tight_layout()
    plt.show()


#  Variational ground-state ansatz (HVA)


def _apply_xxz_pair_gate(qc, q0, q1, theta):
    """Apply the parameterized XXZ-type two-qubit gate used in the HVA."""
    qc.cx(q0, q1)
    qc.rz(theta, q1)
    qc.h(q0)
    qc.rz(theta + np.pi / 2, q0)
    qc.cx(q0, q1)
    qc.rz(-theta, q1)
    qc.h(q1)
    qc.cx(q1, q0)
    qc.rz(np.pi / 2, q1)
    qc.rz(-np.pi / 2, q0)
    qc.h(q1)
    qc.h(q0)


def build_gs_ansatz(n, params, layers):
    """Build the Hamiltonian variational ansatz (HVA) circuit for
    ground-state preparation of the 1D Heisenberg model.

    Starts from a product of singlet pairs and applies alternating
    odd/even layers of parameterized XXZ gates. For layer r,
    params[2 * r] is the odd-layer (inter-pair) angle and
    params[2 * r + 1] is the even-layer (intra-pair) angle.
    """
    qc = QuantumCircuit(n)
    # Initial singlet product state
    for i in range(n // 2):
        qc.x(2 * i)
        qc.x(2 * i + 1)
        qc.h(2 * i + 1)
        qc.cx(2 * i + 1, 2 * i)
    # Variational layers
    for r in range(layers):
        for i in range(1, (n + 1) // 2):  # odd layer
            _apply_xxz_pair_gate(qc, 2 * i - 1, 2 * i, params[2 * r])
        for i in range(n // 2):  # even layer
            _apply_xxz_pair_gate(qc, 2 * i, 2 * i + 1, params[2 * r + 1])
    return qc


def prepare_ground_state(n, gs_layers=5, max_bond=128, cutoff=1e-8):
    """Prepare the KCuF3 (isotropic Heisenberg) ground state as a QuantumCircuit.

    Runs DMRG (quimb MPO + DMRG2) to get the chain's ground state, then optimizes
    the HVA angles to maximize the MPS overlap |<psi_ansatz|psi_DMRG>|^2. No exact
    diagonalization, so it scales to larger n.
    """
    J = Jz = 1.0
    builder = qtn.SpinHam1D(S=1 / 2)
    builder += J * 0.5, "+", "-"
    builder += J * 0.5, "-", "+"
    builder += Jz, "Z", "Z"
    H_mpo = builder.build_mpo(L=n)
    dmrg = qtn.DMRG2(H_mpo)
    dmrg.solve(tol=1e-8, verbosity=0)

    gs_sim = QuimbSimulator(
        quimb_circuit_factory=partial(
            qtn.CircuitMPS, gate_opts=dict(cutoff=cutoff, max_bond=max_bond)
        ),
        autodiff_backend="jax",
    )

    def gs_infidelity(params):
        psi = tensornetwork_from_circuit(
            build_gs_ansatz(n, params, gs_layers), gs_sim
        ).psi
        return 1 - abs(psi.H @ dmrg.state) ** 2

    # Seed and optimizer match the original tutorial. Each layer starts at
    # [0, pi/2]: an odd-layer angle of 0 makes the inter-pair gate the identity,
    # and an even-layer angle of pi/2 makes the intra-pair gate a SWAP (since
    # 0.5 * (XX + YY + ZZ) = SWAP - I/2). That puts the seed at the singlet-pair
    # product limit, which is already a decent approximation to the Heisenberg
    # ground state, so the optimizer only has to refine it. The small jitter
    # (fixed RNG seed, so runs are reproducible) breaks the exact symmetry
    # between layers; COBYQA then runs for up to 100 iterations.
    rng = np.random.default_rng(12345)
    x0 = np.tile([0.0, np.pi / 2], gs_layers) + rng.normal(
        scale=0.1, size=2 * gs_layers
    )
    result_gs = scipy.optimize.minimize(
        gs_infidelity, x0, method="COBYQA", options={"maxiter": 100}
    )
    print(f"DMRG ground-state energy: {dmrg.energy:.6f}")
    print(f"GS fidelity: {1 - result_gs.fun:.4f}")
    return build_gs_ansatz(n, result_gs.x, gs_layers)


print("Setup complete - helpers defined.")

# Credentials are read from the account saved once via QiskitServerless.save_account(...)
serverless = QiskitServerless()
fn = serverless.load("aqc-dynamics-function")

n = 10
dt = 0.6  # physical time per Trotter step (also the omega-axis unit in POST)
time_steps = 10
center = n // 2 - 1

# MPS-simulator settings, shared by the ground-state prep here and the AQC
# compression inside the function (matches the original tutorial).
mps_max_bond = 32
mps_cutoff = 1e-8

# 1D isotropic Heisenberg (KCuF3) Hamiltonian on n qubits
H = SparsePauliOp.from_sparse_list(
    [(p, [i, i + 1], 0.25) for i in range(n - 1) for p in ("XX", "YY", "ZZ")],
    num_qubits=n,
)

# Ground state (DMRG + fidelity max) + neutron kick baked into the same circuit
gs_circuit = prepare_ground_state(
    n, gs_layers=3, max_bond=mps_max_bond, cutoff=mps_cutoff
)
gs_circuit.rz(
    np.pi / 2, center
)  # exp(-i (pi/2)/2 Z_center): the neutron perturbation
print(
    f"Prepared {n}-qubit ground state with the neutron kick at site {center}."
)

job = fn.run(
    t_steps=time_steps,
    aqc_segments=[
        {
            "n_steps": 3,
            "ansatz_steps": 1,
        },  # early steps -> shallow 1-layer ansatz
        {
            "n_steps": 2,
            "ansatz_steps": 2,
        },  # later steps -> deeper 2-layer ansatz
    ],
    aqc_options={
        "max_bond": mps_max_bond,  # MPS bond dimension for AQC compression
        "cutoff": mps_cutoff,
        "optimizer_settings": {
            "method": "L-BFGS-B",
            "jac": True,
            "options": {"maxiter": 100},
        },
    },
    dt=dt,
    hamiltonian=H,
    initial_state=gs_circuit,  # prepared ground state including the neutron kick
    # observables omitted -> default per-site Z (the neutron sigma_z readout)
    backend="statevector",
)

print(job.status())  # rerun this cell until status says DONE

# The per-site <sigma_z>(t) the function returns is the retarded Green's function
# G(j, j_c, t). The workflow samples t = 1..time_steps, so drop the t = 0 row (the
# prepared+kicked state before any evolution) before post-processing.
result = job.result()
print(
    "AQC fidelities:",
    {k: round(v, 4) for k, v in result["metadata"]["aqc_fidelities"].items()},
)

ev = np.array(result["expectation_values"])
Gjjc = ev[1:]  # shape (time_steps, n)
print("Green's function shape:", Gjjc.shape)

q_res, w_res = 100, 100
spectrum = get_spectrum(n, Gjjc, dt, time_steps, q_res, w_res)
spectrum = -(spectrum + spectrum[:, ::-1]) / 2  # mirror symmetry
spectrum = np.clip(spectrum, a_min=0, a_max=None)  # clip negatives

plot_green(
    n,
    Gjjc,
    time_steps,
    dt,
    title=f"Retarded Green's function - {n} qubits (AQC, statevector)",
)
plot_spectrum(
    spectrum,
    dt,
    q_res,
    w_res,
    lower_bound=True,
    upper_bound=True,
    title=f"Dynamical structure factor - {n} qubits (AQC, statevector)",
)

n = 30
dt = 0.6
time_steps = 20
center = n // 2 - 1

# Same MPS settings as the original large-scale run: a larger bond for the
# longer, more-entangled chain (shared by GS prep and AQC compression).
mps_max_bond = 128
mps_cutoff = 1e-8

# Same KCuF3 Hamiltonian and ground-state prep, on a larger chain
H = SparsePauliOp.from_sparse_list(
    [(p, [i, i + 1], 0.25) for i in range(n - 1) for p in ("XX", "YY", "ZZ")],
    num_qubits=n,
)
gs_circuit = prepare_ground_state(
    n, gs_layers=5, max_bond=mps_max_bond, cutoff=mps_cutoff
)
gs_circuit.rz(np.pi / 2, center)  # neutron kick at the center site
print(
    f"Prepared {n}-qubit ground state with the neutron kick at site {center}."
)

# Steps 2 + 3: the function compresses (varied ansatz) and executes on hardware.
job = fn.run(
    t_steps=time_steps,
    aqc_segments=[
        {
            "n_steps": 6,
            "ansatz_steps": 1,
        },  # early steps -> shallow 1-layer ansatz
        {
            "n_steps": 4,
            "ansatz_steps": 2,
        },  # later steps -> deeper 2-layer ansatz
    ],
    aqc_options={
        "max_bond": mps_max_bond,  # 128 for the longer chain
        "cutoff": mps_cutoff,
        "optimizer_settings": {
            "method": "L-BFGS-B",
            "jac": True,
            "options": {"maxiter": 100},
        },
    },
    dt=dt,
    hamiltonian=H,
    initial_state=gs_circuit,
    backend_name="ibm_pittsburgh",
    # Mitigation settings from the original tutorial. Only the two
    # measure_noise_learning values differ from the function's defaults; the rest
    # restates them, because a caller-supplied estimator_options dict replaces the
    # function's defaults wholesale rather than merging into them.
    estimator_options={
        "environment": {"job_tags": ["TUT-SNS"]},
        "dynamical_decoupling": {"enable": True, "sequence_type": "XY4"},
        "twirling": {
            "enable_gates": True,
            "num_randomizations": 1000,
            "shots_per_randomization": 128,
        },
        "resilience": {
            "measure_mitigation": True,
            "measure_noise_learning": {
                "num_randomizations": 32,
                "shots_per_randomization": 100,
            },
        },
    },
)
print("job ID (save this to reconnect later):", job.job_id)

# Reconnect to a previously submitted job by its ID. Only needed in a NEW kernel
# session; if you are still in the session where you submitted, the `job` handle
# from the preceding cell is already live, so skip this cell. Replace the ID that follows with your own.
job = serverless.get_job_by_id("<your job ID>")

# Check where the job is. Re-run this until it reports DONE before fetching the
# result in the following cell: QUEUED -> INITIALIZING -> RUNNING: OPTIMIZING_FOR_HARDWARE ->
# RUNNING: WAITING_FOR_QPU -> RUNNING: EXECUTING_QPU -> RUNNING: POST_PROCESSING
# -> DONE.
print(job.status())

# Run this only once the preceding status cell reports DONE. result() blocks until
# the job finishes, so calling it earlier just waits (possibly for hours).
result = job.result()
print(
    "AQC fidelities:",
    {k: round(v, 4) for k, v in result["metadata"]["aqc_fidelities"].items()},
)

ev = np.array(result["expectation_values"])
Gjjc = ev[1:]  # drop the t = 0 row -> shape (time_steps, n)

n = result["metadata"]["n"]
q_res, w_res = 100, 100
spectrum = get_spectrum(n, Gjjc, dt, time_steps, q_res, w_res)
spectrum = -(spectrum + spectrum[:, ::-1]) / 2  # mirror symmetry
spectrum = np.clip(spectrum, a_min=0, a_max=None)  # clip negatives

plot_green(
    n,
    Gjjc,
    time_steps,
    dt,
    title=f"Retarded Green's function - {n} qubits (AQC, hardware)",
)
plot_spectrum(
    spectrum,
    dt,
    q_res,
    w_res,
    lower_bound=True,
    upper_bound=True,
    title=f"Dynamical structure factor - {n} qubits (AQC, hardware)",
)