Computing Magnetization With The Suzuki-Trotter Approximation

A key application for quantum computers is the simulation of quantum systems. This tutorial will demonstrate a CUDA-Q implementation of a Suzuki-Trotter simulation of a Heisenberg Model Hamiltonian to compute the magnetization of a spin chain.

This example takes advantage of CUDA-Q’s state handling abilities to perform the recursive Trotter simulations.

Problem Setup

import cudaq
import time
import numpy as np
from typing import List

The goal of this problem is to estimate the time evolution of an initial quantum state, governed by a Hamiltonian, and then compute the average magnetization fo the final state. The time evolution can be approximated using the Trotter method:

\[e^{-iHt} \approx \prod_{n=0}^N e^{\frac{-iHt}{n}}\]

The Heisenberg Hamiltonian is defined below, parameterized by these predefined constants: \(g\), \(Jx\), \(Jy\), and \(\omega = 2\pi\). The time step \(dt\) is selected for this problem as well as n_steps and n_spins (\(N\)), determining the size of the simulation. The default setup considers 11 spins and 10 steps, which can easily be simulated on a CPU. If you have access to at least one GPU, the problem can be increased to around 25 steps and 100 time steps.

\[H = \sum_{j=1}^{N}(J_x\sigma^x_j\sigma^x_{j+1} + J_x\sigma^y_j\sigma^y_{j+1} + J_x\sigma^z_j\sigma^z_{j+1} + \cos(\omega * t)\sigma^x_j )\]

g = 1.0
Jx = 1.0
Jy = 1.0
Jz = g
dt = 0.05
n_steps = 10
n_spins = 11
omega = 2 * np.pi


def heisenbergModelHam(t: float) -> cudaq.SpinOperator:
    tdOp = cudaq.SpinOperator(num_qubits=n_spins)
    for i in range(0, n_spins - 1):
        tdOp += (Jx * cudaq.spin.x(i) * cudaq.spin.x(i + 1))
        tdOp += (Jy * cudaq.spin.y(i) * cudaq.spin.y(i + 1))
        tdOp += (Jz * cudaq.spin.z(i) * cudaq.spin.z(i + 1))
    for i in range(0, n_spins):
        tdOp += (np.cos(omega * t) * cudaq.spin.x(i))
    return tdOp

Next, two CUDA-Q kernels are defined. The first, getInitState, prepares an initial state where an \(X\) gate is applied to the first and then every other qubit to initialize a chain of alternating spins.

@cudaq.kernel
def getInitState(numSpins: int):
    q = cudaq.qvector(numSpins)
    for qId in range(0, numSpins, 2):
        x(q[qId])

The second, trotter, performs a single Trotter step given some provided initial state and then returns the resulting state. Two notes should be made about this kernel. 1. It takes advantage of CUDA-Q’s state handling abilities, allowing each step to proceed from the previous, still stored in GPU memory during the simulation. This provides a significant speedup which is demonstrated in the plot at the end of this tutorial. 2. This kernel takes advantage of the pauli_word object which allows a list of Pauli words and their coefficients to be passed into the kernel and applied as an exponentiated matrix operation using exp_pauli.

@cudaq.kernel
def trotter(state: cudaq.State, coefficients: List[complex],
            words: List[cudaq.pauli_word], dt: float):
    q = cudaq.qvector(state)
    for i in range(len(coefficients)):
        exp_pauli(coefficients[i].real * dt, q, words[i])

The functions below are used to strip the Hamiltonian spin operator into a list of coefficients and Pauli words.

def termCoefficients(op: cudaq.SpinOperator) -> List[complex]:
    result = []
    ham.for_each_term(lambda term: result.append(term.get_coefficient()))
    return result


def termWords(op: cudaq.SpinOperator) -> List[str]:
    result = []
    ham.for_each_term(lambda term: result.append(term.to_string(False)))
    return result

Finally, a second spin operator is defined which will be used to compute the expectation value corresponding to the average magnetization.

average_magnetization = cudaq.SpinOperator(num_qubits=n_spins)
for i in range(0, n_spins):
    average_magnetization += ((1.0 / n_spins) * cudaq.spin.z(i))
average_magnetization -= 1.0

Running the Simulation

Before looping through the Trotter steps, the initial state is constructed.

state = cudaq.get_state(getInitState, n_spins)

Next, the time steps are looped through. At each step, the time dependent Hamiltonian is defined. This Hamiltonian is used to construct the Trotter kernel for that time step, which is then used to compute the average magnetization expectation value. The state is saved, and used as the initial state for the next time step.

results = []
times = []
for i in range(0, n_steps):
    start_time = time.time()
    ham = heisenbergModelHam(i * dt)
    coefficients = termCoefficients(ham)
    words = termWords(ham)
    magnetization_exp_val = cudaq.observe(trotter, average_magnetization, state,
                                          coefficients, words, dt)
    result = magnetization_exp_val.expectation()
    results.append(result)
    state = cudaq.get_state(trotter, state, coefficients, words, dt)
    stepTime = time.time() - start_time
    times.append(stepTime)
    print(f"Step {i}: time [s]: {stepTime}, result: {result}")

print(f"Step times: {times}")
print(f"Results: {results}")

Step 0: time [s]: 0.03444695472717285, result: -0.09042024163828166
Step 1: time [s]: 0.0026793479919433594, result: -0.08898564687193886
Step 2: time [s]: 0.002758026123046875, result: -0.08698024360923415
Step 3: time [s]: 0.002524852752685547, result: -0.08507694741170907
Step 4: time [s]: 0.0026259422302246094, result: -0.08394118068746997
Step 5: time [s]: 0.002542734146118164, result: -0.08394076573115139
Step 6: time [s]: 0.0027430057525634766, result: -0.08502222139504187
Step 7: time [s]: 0.0025305747985839844, result: -0.08677832064885871
Step 8: time [s]: 0.003045797348022461, result: -0.08863390649349775
Step 9: time [s]: 0.0025949478149414062, result: -0.09005513983609514
Step times: [0.03444695472717285, 0.0026793479919433594, 0.002758026123046875, 0.002524852752685547, 0.0026259422302246094, 0.002542734146118164, 0.0027430057525634766, 0.0025305747985839844, 0.003045797348022461, 0.0025949478149414062]
Results: [-0.09042024163828166, -0.08898564687193886, -0.08698024360923415, -0.08507694741170907, -0.08394118068746997, -0.08394076573115139, -0.08502222139504187, -0.08677832064885871, -0.08863390649349775, -0.09005513983609514]

CUDA-Q’s state handling capabilities provide a massive performance boost for this algorithm. Rather than resimulate all previous operation at any given time step, saving the previous state in GPU memory allows completion of the simulation with fewer operations. The figure below demonstrates the 24X speedup realized by a 100 step Trotter simulation.