Simulate neutron scattering with an AQC + Trotter dynamics Serverless workflow

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
# above is already live, so skip this cell. Replace the id below 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 below: OPTIMIZING_FOR_HARDWARE -> WAITING_FOR_QPU -> EXECUTING_QPU ->
# POST_PROCESSING -> DONE.
print(job.status())

# Run this only once the status cell above 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 = job.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)",
)