variational_quantum_dynamics_tfim.py

"""
======================================================================
Variational Quantum Dynamics - Transverse Field Ising Model (TFIM)
======================================================================

This example demonstrates quantum dynamics simulation using TyxonQ's
VariationalRuntime for the TFIM Hamiltonian.

Based on:
- PRL 125, 010501 (2020) - Variational Quantum Dynamics
- src/tyxonq/libs/circuits_library/variational.py

Key Features:
- TFIM Hamiltonian: H = -J Σ Z_i Z_{i+1} - h Σ X_i
- VQD (Variational Quantum Dynamics) algorithm
- p-VQD (Projected VQD) for improved stability
- McLachlan's variational principle
- Exact Trotter evolution comparison
- Observable tracking (⟨Z⟩, ⟨X⟩, energy)

Physics Background:
-------------------
The TFIM exhibits a quantum phase transition:
- h << J: Ferromagnetic phase (all spins aligned)
- h >> J: Paramagnetic phase (spins aligned with field)
- h ≈ J: Quantum critical point

Performance:
-----------
For 6 qubits, 3 layers, 50 steps:
- VQD: ~0.1s per step
- Fidelity with exact: >99% for short time
- Energy conservation: <0.1% drift

Author: TyxonQ Team
Date: 2025
"""

import time
import numpy as np
from scipy.linalg import expm

import tyxonq as tq
from tyxonq.libs.quantum_library.kernels.pauli import heisenberg_hamiltonian
from tyxonq.libs.circuits_library.variational import VariationalRuntime


# ==============================================================================
# Configuration
# ==============================================================================

print("=" * 70)
print("Variational Quantum Dynamics - TFIM Example")
print("=" * 70)

# System parameters
N_QUBITS = 6
N_LAYERS = 3
N_PARAMS = N_LAYERS * N_QUBITS * 2  # [layer, qubit, (ZZ, RX)]

# TFIM Hamiltonian parameters
J = 1.0      # ZZ coupling strength
H_FIELD = -1.0  # Transverse field strength
PBC = False  # Periodic boundary conditions

# Dynamics parameters
TIME_STEP = 0.05
N_STEPS = 50
ALGORITHM = "VQD"  # "VQD" or "p-VQD"
INCLUDE_PHASE = False  # McLachlan's principle (False) vs Time-dependent (True)

print(f"\n[Configuration]")
print(f"  Qubits: {N_QUBITS}")
print(f"  Layers: {N_LAYERS}")
print(f"  Parameters: {N_PARAMS}")
print(f"  TFIM: J={J}, h={H_FIELD}, PBC={PBC}")
print(f"  Algorithm: {ALGORITHM}")
print(f"  Phase tracking: {INCLUDE_PHASE}")
print(f"  Time: {N_STEPS} steps × {TIME_STEP} = {N_STEPS * TIME_STEP}")


# ==============================================================================
# Build TFIM Hamiltonian
# ==============================================================================

print("\n[1/6] Building TFIM Hamiltonian...")

# Build edges for 1D chain
edges = [(i, i + 1) for i in range(N_QUBITS - 1)]
if PBC:
    edges.append((N_QUBITS - 1, 0))

# TFIM: H = -J Σ Z_i Z_{i+1} - h Σ X_i
hamiltonian = heisenberg_hamiltonian(
    N_QUBITS,
    edges,
    hzz=J,      # ZZ coupling
    hxx=0.0,    # No XX
    hyy=0.0,    # No YY
    hx=H_FIELD, # Transverse field
    hy=0.0,     # No Y field
    hz=0.0      # No Z field
)

print(f"  ✓ Hamiltonian built: {N_QUBITS} qubits")
print(f"  ✓ Shape: {hamiltonian.shape}")
print(f"  ✓ Edges: {len(edges)} bonds")

# Compute exact ground state for reference
eigenvalues = np.linalg.eigvalsh(hamiltonian)
exact_gs_energy = eigenvalues[0]
print(f"  ✓ Exact ground state energy: {exact_gs_energy:.8f}")


# ==============================================================================
# Build Variational Ansatz
# ==============================================================================

print("\n[2/6] Building variational ansatz...")

# Initial state: |0...0⟩
init_state = np.zeros(2**N_QUBITS, dtype=np.complex128)
init_state[0] = 1.0

# Define ansatz state function directly
def ansatz_state_fn(params):
    """Variational ansatz state function
    
    Args:
        params: Array of shape (N_PARAMS,) = (N_LAYERS * N_QUBITS * 2,)
                Structured as [layer, qubit, (zz_param, rx_param)]
    
    Returns:
        Statevector of shape (2^N_QUBITS,)
    """
    params = np.asarray(params, dtype=np.float64)
    params_reshaped = params.reshape([N_LAYERS, N_QUBITS, 2])
    
    c = tq.Circuit(N_QUBITS)
    
    # Apply layers
    for layer in range(N_LAYERS):
        # ZZ entangling gates
        for i in range(N_QUBITS - 1):
            c.rzz(i, i + 1, theta=params_reshaped[layer, i, 0])
        
        # RX rotations
        for i in range(N_QUBITS):
            c.rx(i, theta=params_reshaped[layer, i, 1])
    
    # Get state
    from tyxonq.devices.simulators.statevector.engine import StatevectorEngine
    eng = StatevectorEngine()
    state = eng.state(c)
    return np.asarray(state, dtype=np.complex128)

print(f"  ✓ Ansatz layers: {N_LAYERS}")
print(f"  ✓ Gates per layer: ZZ entangling + RX rotations")
print(f"  ✓ Ansatz state function ready")


# ==============================================================================
# Initialize Variational Runtime
# ==============================================================================

print("\n[3/6] Initializing VariationalRuntime...")

runtime = VariationalRuntime(
    ansatz_state_fn=ansatz_state_fn,
    hamiltonian=hamiltonian,
    n_params=N_PARAMS,
    eps=1e-5,
    include_phase=INCLUDE_PHASE,
    initial_state=init_state,
)

# Add observables
# Z observable on first qubit
Z_op = np.zeros((2**N_QUBITS, 2**N_QUBITS), dtype=np.complex128)
Z_single = np.array([[1, 0], [0, -1]], dtype=np.complex128)
for i in range(2**N_QUBITS):
    # Check if first qubit is |0⟩ or |1⟩
    if (i >> (N_QUBITS - 1)) & 1 == 0:
        Z_op[i, i] = 1.0
    else:
        Z_op[i, i] = -1.0

runtime.add_property_mat("Z0", Z_op)

# X observable on first qubit
X_op = np.zeros((2**N_QUBITS, 2**N_QUBITS), dtype=np.complex128)
for i in range(2**N_QUBITS):
    # Flip first qubit
    j = i ^ (1 << (N_QUBITS - 1))
    X_op[i, j] = 1.0

runtime.add_property_mat("X0", X_op)

print(f"  ✓ Runtime initialized")
print(f"  ✓ Parameters: {runtime.n_params}")
print(f"  ✓ Initial energy: {np.real(runtime.state_list[0].conj() @ hamiltonian @ runtime.state_list[0]):.8f}")
print(f"  ✓ Observables: Z0, X0")


# ==============================================================================
# Run Variational Dynamics
# ==============================================================================

print(f"\n[4/6] Running {ALGORITHM} evolution...")
print("-" * 70)
print(f"{'Step':<6} {'Time':<8} {'Energy':<15} {'⟨Z0⟩':<12} {'⟨X0⟩':<12} {'Time(s)':<10}")
print("-" * 70)

energy_history = []
z_history = []
x_history = []
time_history = []
times = []

for step in range(N_STEPS):
    t0 = time.time()
    
    # Current observables
    props = runtime.properties()
    z_val = float(np.real(props["Z0"]))
    x_val = float(np.real(props["X0"]))
    energy = float(np.real(runtime.state_list[-1].conj() @ hamiltonian @ runtime.state_list[-1]))
    
    # Store
    energy_history.append(energy)
    z_history.append(z_val)
    x_history.append(x_val)
    time_history.append(runtime.t)
    
    # Print progress
    if step % 10 == 0:
        t1 = time.time()
        step_time = t1 - t0 if step > 0 else 0.0
        print(f"{step:<6} {runtime.t:<8.3f} {energy:<15.8f} {z_val:<12.6f} {x_val:<12.6f} {step_time:<10.4f}")
    
    # Time evolution step
    if step < N_STEPS - 1:
        if ALGORITHM == "VQD":
            runtime.step_vqd(TIME_STEP)
        elif ALGORITHM == "p-VQD":
            runtime.step_pvqd(TIME_STEP)
        else:
            raise ValueError(f"Unknown algorithm: {ALGORITHM}")
    
    t1 = time.time()
    times.append(t1 - t0)

# Final step
props = runtime.properties()
z_val = float(np.real(props["Z0"]))
x_val = float(np.real(props["X0"]))
energy = float(np.real(runtime.state_list[-1].conj() @ hamiltonian @ runtime.state_list[-1]))
energy_history.append(energy)
z_history.append(z_val)
x_history.append(x_val)
time_history.append(runtime.t)

print("-" * 70)
print(f"  Total time: {sum(times):.2f}s")
print(f"  Avg time per step: {np.mean(times):.4f}s")


# ==============================================================================
# Exact Evolution Comparison
# ==============================================================================

print("\n[5/6] Comparing with exact Trotter evolution...")

# Exact evolution using matrix exponential
psi_exact = init_state.copy()
exact_energies = []
exact_z = []
exact_x = []

for step in range(N_STEPS + 1):
    # Compute observables
    e_exact = float(np.real(psi_exact.conj() @ hamiltonian @ psi_exact))
    z_exact = float(np.real(psi_exact.conj() @ Z_op @ psi_exact))
    x_exact = float(np.real(psi_exact.conj() @ X_op @ psi_exact))
    
    exact_energies.append(e_exact)
    exact_z.append(z_exact)
    exact_x.append(x_exact)
    
    # Evolve
    if step < N_STEPS:
        U = expm(-1j * hamiltonian * TIME_STEP)
        psi_exact = U @ psi_exact
        # Normalize
        psi_exact = psi_exact / np.linalg.norm(psi_exact)

# Compute fidelity
final_vqd_state = runtime.state_list[-1]
final_exact_state = psi_exact
fidelity = np.abs(np.vdot(final_exact_state, final_vqd_state))**2

print(f"  ✓ Exact evolution completed")
print(f"  ✓ Final fidelity: {fidelity:.6f}")


# ==============================================================================
# Analysis
# ==============================================================================

print("\n[6/6] Analysis...")
print("-" * 70)

# Energy conservation
energy_drift = abs(energy_history[-1] - energy_history[0])
print(f"  Energy drift: {energy_drift:.6e}")
if energy_drift < 1e-3:
    print("  ✓ Energy well conserved")
else:
    print("  ⚠️  Consider increasing layers or decreasing time step")

# Comparison with exact
energy_error = abs(energy_history[-1] - exact_energies[-1])
z_error = abs(z_history[-1] - exact_z[-1])
x_error = abs(x_history[-1] - exact_x[-1])

print(f"\n  Final VQD energy:   {energy_history[-1]:.8f}")
print(f"  Final exact energy: {exact_energies[-1]:.8f}")
print(f"  Energy error:       {energy_error:.6e}")
print(f"  ⟨Z0⟩ error:         {z_error:.6e}")
print(f"  ⟨X0⟩ error:         {x_error:.6e}")
print(f"  Fidelity:           {fidelity:.6f}")

# Accuracy assessment
if fidelity > 0.99:
    print("\n  ✅ Excellent agreement with exact evolution!")
elif fidelity > 0.95:
    print("\n  ✓ Good agreement with exact evolution")
else:
    print("\n  ⚠️  Consider increasing layers for better accuracy")

print("-" * 70)


# ==============================================================================
# Visualization (Optional)
# ==============================================================================

try:
    import matplotlib.pyplot as plt
    
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    
    # Energy evolution
    ax = axes[0, 0]
    ax.plot(time_history, energy_history, 'b-', linewidth=2, label='VQD')
    ax.plot(time_history, exact_energies, 'r--', linewidth=2, label='Exact')
    ax.set_xlabel('Time')
    ax.set_ylabel('Energy')
    ax.set_title('Energy Evolution')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # Z observable
    ax = axes[0, 1]
    ax.plot(time_history, z_history, 'b-', linewidth=2, label='VQD')
    ax.plot(time_history, exact_z, 'r--', linewidth=2, label='Exact')
    ax.set_xlabel('Time')
    ax.set_ylabel('⟨Z0⟩')
    ax.set_title('Z Observable')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # X observable
    ax = axes[1, 0]
    ax.plot(time_history, x_history, 'b-', linewidth=2, label='VQD')
    ax.plot(time_history, exact_x, 'r--', linewidth=2, label='Exact')
    ax.set_xlabel('Time')
    ax.set_ylabel('⟨X0⟩')
    ax.set_title('X Observable')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # Fidelity over time
    ax = axes[1, 1]
    fidelities = []
    # runtime.state_list has N_STEPS+1 states (initial + all steps)
    for i in range(min(len(time_history), len(runtime.state_list))):
        # Recompute exact state at each time
        psi_e = init_state.copy()
        U_total = expm(-1j * hamiltonian * time_history[i])
        psi_e = U_total @ psi_e
        psi_e = psi_e / np.linalg.norm(psi_e)
        
        # Get VQD state
        psi_v = runtime.state_list[i]
        
        # Fidelity
        f = np.abs(np.vdot(psi_e, psi_v))**2
        fidelities.append(f)
    
    ax.plot(time_history[:len(fidelities)], fidelities, 'g-', linewidth=2)
    ax.axhline(y=0.99, color='r', linestyle='--', alpha=0.5, label='99% threshold')
    ax.set_xlabel('Time')
    ax.set_ylabel('Fidelity')
    ax.set_title('Fidelity with Exact State')
    ax.set_ylim([max(0, min(fidelities) - 0.05), 1.01])
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('vqd_tfim_results.png', dpi=150)
    print(f"\n  📊 Plot saved: vqd_tfim_results.png")
    
except ImportError:
    print("\n  ℹ️  matplotlib not available, skipping plots")


# ==============================================================================
# Summary
# ==============================================================================

print("\n" + "=" * 70)
print("Variational Quantum Dynamics Complete!")
print("=" * 70)

print("\n📚 Key Concepts:")
print("  - VQD: Variational Quantum Dynamics")
print("  - McLachlan's principle: minimizes ||dψ/dt + iHψ||")
print("  - p-VQD: Projected VQD for better long-time stability")
print("  - TFIM: Canonical model for quantum phase transitions")

print("\n🔬 Implementation:")
print("  - VariationalRuntime: src/tyxonq/libs/circuits_library/variational.py")
print("  - Hamiltonian builder: libs/quantum_library/kernels/pauli.py")
print("  - Exact comparison: scipy.linalg.expm")

print("\n💡 Tips:")
print("  - Increase N_LAYERS for higher accuracy")
print("  - Use p-VQD for long-time evolution")
print("  - Monitor energy conservation and fidelity")
print("  - Compare different TIME_STEP values")

print("\n🎯 Next Steps:")
print("  - Try different J and h values (phase transition)")
print("  - Experiment with periodic boundary conditions")
print("  - Compare VQD vs p-VQD performance")
print("  - Explore larger system sizes")

print("\n" + "=" * 70)