X-ray Absorption Spectroscopy Simulation in the Time-Domain

Observable: absorption cross section

In XAS experiments, the absorption cross section \(\sigma_A\) as a function of the incident X-ray frequency \(\omega\) is measured for a given material. This is proportional to the rate of absorption of X-ray photons of various energies. For our situation, the electrons in the molecular cluster start in a ground molecular state \(|I\rangle\) with energy \(E_I\). This ground state will be excited to states \(|F\rangle\) with energies \(E_F\) by the radiative field through the action of the dipole operator \(\hat m_\rho,\) where \(\rho\) is any of the Cartesian directions \(x,\ y,\,\) or \(z\).

The absorption cross section is given by the equation

where \(c\) is the speed of light, \(\hbar\) is Planck’s constant, and \(\eta\) is the line broadening, set to match the experimental resolution of the spectroscopy. Below is an illustration of an X-ray absorption spectrum.

The goal of this demo is to implement a quantum algorithm that can calculate this spectrum.

Time-domain simulation

Both the initial state \(|I\rangle\) and the dipole operator acting on the initial state \(\hat m_\rho|I\rangle\) can be determined classically, and we’ll demonstrate how to do that later. With the initial state computed, we will use a mathematical trick called a frequency-domain Green’s function to re-write the absorption cross section as an expectation value of a specific operator. We can write the cross section as the imaginary part of the following Green’s function (see Section IIB in Ref. 2)

With a bit of algebra, we can confirm that this works: inserting a resolution of identity of final states \(\sum_F |F\rangle \langle F|\) and simplifying gives

where the first term is proportional to the absorption cross section. The second term is zero if we centre the frame of reference for our molecular orbitals at the nuclear-charge weighted centre for our molecular cluster of choice.

Okay, so how do we determine \(\mathcal{G}_\rho(\omega)\)? If we are going to evaluate this quantity in a quantum register, it will need to be normalized, so instead we are looking for

Calculating this quantity directly is expensive 1. We need to invert the Hamiltonian matrix, which is hard as it typically requires using linear-solve algorithms like HHL 3. We would also need to calculate the cross section at every frequency, which can quickly run up the cost. Instead, our algorithm will aim to simulate the dynamics of our initial state in the time domain, and then Fourier transform the results classically to obtain the frequency spectrum! This allows us to construct the entire spectrum at once, rather than tediously measure each individual frequency we are interested in. This is especially important if the number of excited states underlying the absorption spectrum is very large, which we expect to be the case for moderately sized clusters.

The time-domain Green’s function \(\tilde G(t_j)\) can be determined using the expectation value of the time-evolution operator (normalized) at times \(t_j = \tau j\), where \(\tau\) is the time step and \(j\) is the time-step index

This quantity can easily be calculated on a quantum computer! We can use a Hadamard test circuit with the unitary of time evolution to measure the expectation value for each time \(t_j\).

Repeating this test a number of times \(N\) and taking the mean of the results gives an estimate for \(G_\rho(t_j)\). The spectrum can then be obtained with a discrete time-domain Fourier transform

The time step \(\tau \sim \mathcal{O}(||\hat H||^{-1})\) should be small enough to resolve the largest frequency components that we are interested in, which correspond to the final states with the largest energy. In practice, this is not the largest eigenvalue of \(\hat H\), but simply the largest energy we want to show in the spectrum.

The circuit we will construct to determine the expectation values is shown above. It has three main components:

• State prep, the state \(\hat m_\rho |I\rangle\) is prepared in the quantum register, and an auxiliary qubit is prepared to control the time evolution.

• Time evolution, the state is evolved under the electronic Hamiltonian.

• Measurement, the time-evolved state is measured to obtain statistics for the expectation value.

Let’s look at how to implement these steps in PennyLane. We will make use of the pennylane.qchem module, as well as initial state preparation methods from PySCF.

Ground state calculation

If you haven’t, check out the demo “Initial state preparation for quantum chemistry”. We will be expanding on that demo with code to import a state from the complete active space configuration interaction (CASCI) methods of PySCF, where we restrict the set of active orbitals used in the calculation.

We start by creating our molecule object using the Gaussian type orbitals module pyscf.gto, and obtaining the Hartree-Fock molecular orbitals with the self-consistent field methods pyscf.scf.

from pyscf import gto, scf
import numpy as np

# Create a Mole object.
r = 1.0  # Bond length in Angstrom.
symbols = ["N", "N"]
geometry = np.array([[0.0, 0.0, -r / 2], [0.0, 0.0, r / 2]])
mol = gto.Mole(atom=zip(symbols, geometry), basis="sto3g", symmetry=None)
mol.build(verbose=0)

# Get the molecular orbitals.
hf = scf.RHF(mol)
hf.run(verbose=0)

<pyscf.scf.hf.RHF object at 0x7f5f6c38f970>

Since we will be using PennyLane for other aspects of this calculation, we want to make sure the molecular orbital coefficients are consistent between our PennyLane and PySCF calculations. To do this, we can obtain the molecular orbital coefficients from PennyLane using scf(), and change the coefficients in the hf instance to match.

import pennylane as qml

mole = qml.qchem.Molecule(symbols, geometry, basis_name="sto-3g", unit="angstrom")

# Run self-consistent fields method to get molecular orbital coefficients.
_, coeffs, _, _, _ = qml.qchem.scf(mole)()

# Change MO coefficients in hf object to PennyLane calculated values.
hf.mo_coeff = coeffs

Next, let’s define the active space we will use for our calculation. This will be the number of orbitals \(n_\mathrm{cas}\) and the number of electrons. For \(\mathrm{N}_2\), we will use an active space of five orbitals and four electrons. Running a CASCI instance allows us to calculate the ground state of our system with this selected active space. The CASCI method in PySCF is equivalent to a full-configuration interaction (FCI) procedure on a subset of molecular orbitals.

from pyscf import mcscf

# Define active space of (orbitals, electrons).
n_cas, n_electron_cas = (5, 4)

# Number of core (non-active) orbitals.
n_core = (mol.nelectron - n_electron_cas) // 2

# Initialize CASCI instance of N2 molecule as mycasci.
mycasci = mcscf.CASCI(hf, ncas=n_cas, nelecas=n_electron_cas)
mycasci.run(verbose=0)

# Calculate ground state, and omit small state components.
casci_state = mycasci.ci
casci_state[abs(casci_state) < 1e-6] = 0

To implement this state as a PennyLane state vector, we need to convert the casci_state into a format that is easy to import into PennyLane. One way to do this is to use a sparse matrix representation to turn casci_state into a dictionary, and then use import_state() to import into PennyLane. Here is how you can turn a full-configuration interaction matrix like casci_state into a dictionary.

from scipy.sparse import coo_matrix
from pyscf.fci.cistring import addrs2str

# Convert casci_state into a sparse matrix.
sparse_cascimatr = coo_matrix(casci_state, shape=np.shape(mycasci.ci), dtype=float)
row, col, dat = sparse_cascimatr.row, sparse_cascimatr.col, sparse_cascimatr.data

# Turn indices into strings.
n_cas_a = mycasci.ncas  # Number of alpha spin orbitals.
n_cas_b = n_cas_a  # Number of beta spin orbitals.
n_electron_cas_a, n_electron_cas_b = mycasci.nelecas

strs_row = addrs2str(n_cas_a, n_electron_cas_a, row)
strs_col = addrs2str(n_cas_b, n_electron_cas_b, col)

# Create the FCI matrix as a dict.
wf_casci_dict = dict(zip(list(zip(strs_row, strs_col)), dat))

Lastly, we will use the helper function _sign_chem_to_phys() to adjust the sign of state components to match what they should be for the PennyLane orbital occupation number ordering. Then, we can import the state into PennyLane using _wfdict_to_statevector().

from pennylane.qchem.convert import _sign_chem_to_phys, _wfdict_to_statevector

wf_casci_dict = _sign_chem_to_phys(wf_casci_dict, n_cas)

# Convert dictionary to Pennylane state vector.
wf_casci = _wfdict_to_statevector(wf_casci_dict, n_cas)

Dipole operator action

The electromagnetic field of the incident X-rays can excite the electronic state through the dipole operator. The action of this operator is implemented in PennyLane as dipole_moment(). We can calculate that operator, convert it to a matrix, and apply it to our initial state \(|I\rangle\) to obtain \(\hat m_\rho|I\rangle\).

To generate this operator, we have to specify which molecular orbitals are in our active space. We can obtain the indices of the included and excluded orbitals using active_space() to obtain the lists active and core, respectively.

The action of the dipole operator will be split into the three cartesian directions \(x,\ y\) and \(z\), which we will loop over to obtain the states \(\hat m_{\{x,y,z\}}|I\rangle\). This will give us an operator that describes general absorption of electro-magnetic radiation for our molecule, including for energies outside of the XAS regime 4. This is fine for our simple example, but for more complex instances you may want to modify the dipole operator to restrict the final states (see the Appendix for more details).

# Get core and active orbital indices.
core, active = qml.qchem.active_space(
    mole.n_electrons, mole.n_orbitals,
    active_electrons=n_electron_cas,
    active_orbitals=n_cas
)

m_rho = qml.qchem.dipole_moment(mole, cutoff=1e-8, core=core, active=active)()
rhos = range(len(m_rho))  # [0, 1, 2] are [x, y, z].

wf_dipole = []
dipole_norm = []

# Loop over cartesian coordinates and calculate m_rho|I>.
for rho in rhos:
    dipole_matrix_rho = qml.matrix(m_rho[rho], wire_order=range(2 * n_cas))
    wf = dipole_matrix_rho.dot(wf_casci)  # Multiply state into dipole matrix.

    if np.allclose(wf, np.zeros_like(wf)):  # If wf is zero, then set norm as zero.
        wf_dipole.append(wf)
        dipole_norm.append(0)

    else:  # Normalize the wavefunction.
        dipole_norm.append(np.linalg.norm(wf))
        wf_dipole.append(wf / dipole_norm[rho])

Let’s prepare the circuit that will initialize our qubit register with this state. We will need \(2 n_\mathrm{cas}\) wires, which is twice the number of orbitals in our active space, since we need to account for spin. We will also add one auxiliary wire for the measurement circuit, which we will prepare as the 0 wire with an applied Hadamard gate.

device_type = "lightning.qubit"
dev_prop = qml.device(device_type, wires=int(2*n_cas) + 1, shots=None)


@qml.qnode(dev_prop)
def initial_circuit(wf):
    """Circuit to load initial state and prepare auxiliary qubit."""
    qml.StatePrep(wf, wires=dev_prop.wires.tolist()[1:])
    qml.Hadamard(wires=0)
    return qml.state()

Electronic Hamiltonian

Our electronic Hamiltonian is

where \(a^{(\dagger)}_{p\gamma}\) is the annihilation (creation) operator for a molecular orbital \(p\) with spin \(\gamma\), \(E\) is the nuclear core constant, \(N\) is the number of spatial orbitals, and \((p|\kappa|q)\) and \((pq|rs)\) are the one- and two-electron integrals, respectively 5.

The core constant and the one- and two-electron integrals can be computed in PennyLane using electron_integrals().

core_constant, one, two = qml.qchem.electron_integrals(mole, core=core, active=active)()
core_constant = core_constant[0]

We will have to convert these to chemists’ notation 6.

two_chemist = qml.math.einsum("prsq->pqrs", two)
one_chemist = one - qml.math.einsum("pqrr->pq", two) / 2.0

Next, we will perform compressed double factorization of the Hamiltonian’s two-electron integrals to approximate them as a sum of \(L\) fragments

where \(Z^{(\ell)}\) is symmetric and \(U^{(\ell)}\) is orthogonal. These matrices can be calculated using PennyLane’s factorize() function, with compressed=True. By default, \(L\) is set as twice the number of orbitals in our active space. The Z and U output here are arrays of \(L\) fragment matrices with dimension \(n_\mathrm{cas} \times n_\mathrm{cas}\).

from jax import config

config.update("jax_enable_x64", True)  # Required for factorize consistency.

# Factorize Hamiltonian, producing matrices Z and U for each fragment.
_, Z, U = qml.qchem.factorize(two_chemist, compressed=True)

print("Shape of the factors: ")
print("two_chemist", two_chemist.shape)
print("U", U.shape)
print("Z", Z.shape)

# Compare factorized two-electron fragment sum to the original.
approx_two_chemist = qml.math.einsum("tpk,tqk,tkl,trl,tsl->pqrs", U, U, Z, U, U)
assert qml.math.allclose(approx_two_chemist, two_chemist, atol=0.1)

Shape of the factors:
two_chemist (5, 5, 5, 5)
U (10, 5, 5)
Z (10, 5, 5)

Note there are some terms in this decomposition of the two-electron integrals that are exactly diagonalizable, and can be combined with the one-electron integrals to simplify how the Hamiltonian time evolution is implemented. We call these the “one-electron extra” terms and add them to the one-electron integrals. We can then diagonalize their sum into the matrix \(Z^{(0)}\) with basis rotation matrix \(U^{(0)}\) using np.linalg.eigh, for easy implementation in the simulation circuit.

# Calculate the one-electron extra terms.
Z_prime = np.stack([np.diag(np.sum(Z[i], axis=-1)) for i in range(Z.shape[0])], axis=0)
one_electron_extra = qml.math.einsum("tpk,tkk,tqk->pq", U, Z_prime, U)

# Diagonalize the one-electron integral matrix while adding the one-electron extra.
eigenvals, U0 = np.linalg.eigh(one_chemist + one_electron_extra)
Z0 = np.diag(eigenvals)

Constructing the time-propagation circuit

The main work of our algorithm will be to implement the time evolution of the Hamiltonian fragments by using a Trotter product formula.

The trick when time evolving a compressed double-factorized Hamiltonian is to use Thouless’s theorem 7 to construct a size \(2^{n_\mathrm{cas}} \times 2^{n_\mathrm{cas}}\) unitary \({\bf U}^{(\ell)}\) that is induced by a the single-particle basis transformation \(U^{(\ell)}\) (of size \(n_\mathrm{cas} \times n_\mathrm{cas}\)). The Jordan-Wigner transform can then turn the number operators \(a^\dagger_p a_p = n_{p}\) into Pauli \(Z\) rotations, via \(n_p = (1-\sigma_{z,p})/2\). Note the \(1/2\) term will affect the global phase, and we will have to keep track of that carefully. The resulting Hamiltonian looks like the following (for a derivation see Appendix A of Ref. 2)

The first term is a sum of the core constant and constant factors that arise from the Jordan-Wigner transform and is just a global phase. The second and third terms are the one- and two-electron fragments, respectively. Below is an illustration of the circuit we will use to implement the one- and two-electron fragments in a time-evolution of the factorized Hamiltonian.

We can use BasisRotation to generate a Givens decomposition for the large unitary \({\bf U}^{(\ell)}\) that is generated by the single-particle basis rotation \(U^{(\ell)}\) via Thouless’ theorem. Note that we can do this separately for both spin-halves of the register, since our Hamiltonian does not mix spin sectors: this is cheaper than having one large unitary for the entire register.

def U_rotations(U, control_wires):
    """Circuit implementing the basis rotations of the Hamiltonian fragments."""
    U_spin = qml.math.kron(U, qml.math.eye(2))  # Apply to both spins.
    qml.BasisRotation(
        unitary_matrix=U_spin, wires=[int(i + control_wires) for i in range(2 * n_cas)]
    )

Next we write a function to perform the Z rotations. Controlled arbitrary-angle rotations are expensive. To reduce the cost of having to implement many controlled Z rotations at angles determined by the matrices \(Z^{(\ell)}\), we instead implement uncontrolled Z rotations sandwiched by CNOT gates.

For the one-electron terms, we loop over both the spin and orbital indices, and apply the Z rotations using this double-phase trick. The two-electron fragments are implemented the same way, except the two-qubit rotations use MultiRZ. Notice that this is also where we keep track of part of the global phase from the constant term in the Hamiltonian.

from itertools import product


def Z_rotations(Z, step, is_one_electron_term, control_wires):
    """Circuit implementing the Z rotations of the Hamiltonian fragments."""
    if is_one_electron_term:
        for sigma in range(2):
            for i in range(n_cas):
                qml.ctrl(
                    qml.X(wires=int(2*i + sigma + control_wires)),
                    control=range(control_wires),
                    control_values=0,
                )
                qml.RZ(-Z[i, i] * step / 2, wires=int(2*i + sigma + control_wires))
                qml.ctrl(
                    qml.X(wires=int(2*i + sigma + control_wires)),
                    control=range(control_wires),
                    control_values=0,
                )
        globalphase = np.sum(Z) * step  # Phase from 1/2 term in Jordan-Wigner transform.

    else:  # It's a two-electron fragment.
        for sigma, tau in product(range(2), repeat=2):
            for i, k in product(range(n_cas), repeat=2):
                if i != k or sigma != tau:  # Skip the one-electron correction terms.
                    qml.ctrl(qml.X(wires=int(2*i + sigma + control_wires)),
                            control=range(control_wires), control_values=0)
                    qml.MultiRZ(Z[i, k] / 8.0 * step,
                        wires=[int(2*i + sigma + control_wires),
                               int(2*k + tau + control_wires)])
                    qml.ctrl(qml.X(wires=int(2 * i + sigma + control_wires)),
                        control=range(control_wires), control_values=0)
        globalphase = np.trace(Z)/4.0*step - np.sum(Z)*step/2.0

    qml.PhaseShift(-globalphase, wires=0)

For the factorized Hamiltonian \(H_\mathrm{CDF} = \sum_j^N H_j\), with non-commuting fragments \(H_j\), a second-order product formula approximates the time evolution for a time step \(\tau\) as

Note in the formula above, the second product of time-evolution fragments is reversed in order. The function first_order_trotter will implement the \(U\) rotations and \(Z\) rotations, and adjust the global phase from the core constant term. It will also be able to reverse the order of applied fragments so we can later construct the second-order Trotter step. Finally, we can further reduce cost by tracking the last \(U\) rotation used: that way we can implement two consecutive rotations at once as \(V^{(\ell)} = U^{(\ell-1)}(U^{(\ell)})^T\), halving the number of rotations required per Trotter step.

def first_order_trotter(step, prior_U, final_rotation, reverse=False):
    """Implements a Trotter step for the CDF Hamiltonian."""
    # Prepend the one-electron fragment to the two-electron fragment array.
    _U0 = np.expand_dims(U0, axis=0)
    _Z0 = np.expand_dims(Z0, axis=0)
    _U = np.concatenate((_U0, U), axis=0)
    _Z = np.concatenate((_Z0, Z), axis=0)

    num_two_electron_fragments = U.shape[0]  # Number of fragments L.
    is_one_body = np.array([True] + [False] * num_two_electron_fragments)
    order = list(range(len(_Z)))

    if reverse:
        order = order[::-1]

    for fragment in order:
        U_rotations(prior_U @ _U[fragment], 1)
        Z_rotations(_Z[fragment], step, is_one_body[fragment], 1)
        prior_U = _U[fragment].T

    if final_rotation:
        U_rotations(prior_U, 1)

    # Global phase adjustment from core constant.
    qml.PhaseShift(-core_constant * step, wires=0)

    return prior_U

Our function second_order_trotter implements a second-order Trotter step, returning a QNode. The returned circuit applies StatePrep to prepare the register in the previous quantum state, and then two first_order_trotter evolutions for time step/2, first in the nominal order and then reversed. The total step size is step.

def second_order_trotter(dev, state, step):
    """Returns a second-order Trotter product circuit."""
    qubits = dev.wires.tolist()

    def circuit():
        # State preparation -- set as the final state from the previous iteration.
        qml.StatePrep(state, wires=qubits)

        prior_U = np.eye(n_cas)  # No initial prior U, so set as identity matrix.
        prior_U = first_order_trotter(step / 2, prior_U=prior_U,
                                      final_rotation=False, reverse=False)
        prior_U = first_order_trotter(step / 2, prior_U=prior_U,
                                      final_rotation=True, reverse=True)

        return qml.state()

    return qml.QNode(circuit, dev)

Core-valence separation approximation

For larger molecular instances, a simplifying approximation is to restrict the excited states in the calculation to only include those of relevance for XAS, which are states where a core electron is excited, i.e. there exists a hole in the core orbitals. These are known as core-excited states, and lie significantly above the valence-excited states in energy. Typically the frequency range is focused on a target atom in a molecular cluster, and also near a transition energy, such as targeting core \(1s\) electrons.

To make the restriction to core-excited states in practice, we need to turn off the two-electron terms that couple core-excited and valence-excited states. These terms are in general small, but by setting them to zero the coupling is removed entirely from the time evolution. To implement the core-valence separation approximation in an XAS simulation algorithm, there are two steps:

• Before performing compressed double factorization on the two-electron integrals, remove from the two-electron integrals the terms that include at least one core orbital and at least one valence orbital. This decouples the core-excited and valence-excited state manifolds.

• Remove all the matrix elements from the dipole operator that do not include exactly one core orbital. This allows the quantum algorithm to start directly in the core-excited state manifold.

This approximation would be useful when simulating a complex molecular instance, such as a cluster in a lithium-excess material.