Implementing Trotter simulation of the double-factorized Hamiltonian

In this tutorial, we’ll write a function to implement approximate time evolution of the double-factorized Hamiltonian via a Trotter-Suzuki formula. See Double-factorized representation of the molecular Hamiltonian for background information. We’ll compare our implementation with exact time evolution computed using direct operator exponentiation, as well as ffsim’s built-in implementation simulate_trotter_double_factorized.

Build the Hamiltonian

We begin by building a molecular Hamiltonian to test our code on. We’ll create a nitrogen molecule in an active space of 8 orbitals and 10 electrons. We use ffsim’s MolecularData class, which implements a simplistic wrapper around PySCF to compute a representation of the Hamiltonian as an instance of MolecularHamiltonian, which we store in the mol_hamiltonian variable.

import pyscf

import ffsim

# Build N2 molecule
mol = pyscf.gto.Mole()
mol.build(
    atom=[["N", (0, 0, 0)], ["N", (1.0, 0, 0)]],
    basis="sto-6g",
    symmetry="Dooh",
)

# Define active space
n_frozen = pyscf.data.elements.chemcore(mol)
active_space = range(n_frozen, mol.nao_nr())

# Get molecular data and Hamiltonian
scf = pyscf.scf.RHF(mol).run()
mol_data = ffsim.MolecularData.from_scf(scf, active_space=active_space)
norb, nelec = mol_data.norb, mol_data.nelec
mol_hamiltonian = mol_data.hamiltonian

print(f"norb = {norb}")
print(f"nelec = {nelec}")

converged SCF energy = -108.464957764796
norb = 8
nelec = (5, 5)

Next, we compute the double-factorized representation of the Hamiltonian. In ffsim, the DoubleFactorizedHamiltonian class is used to store this Hamiltonian representation.

# Get the Hamiltonian in the double-factorized representation
df_hamiltonian = ffsim.DoubleFactorizedHamiltonian.from_molecular_hamiltonian(
    mol_hamiltonian
)

To get a sense of how the two different Hamiltonian representations differ, let’s print out the shapes of the tensors describing the original and double-factorized representations.

print("Original representation")
print("-----------------------")
print("One-body tensor shape:")
print(mol_hamiltonian.one_body_tensor.shape)
print()
print("Two-body tensor shape:")
print(mol_hamiltonian.two_body_tensor.shape)
print()

print("Double-factorized representation")
print("--------------------------------")
print("One-body tensor shape:")
print(df_hamiltonian.one_body_tensor.shape)
print()
print("Diagonal Coulomb matrices shape:")
print(df_hamiltonian.diag_coulomb_mats.shape)
print()
print("Orbital rotations shape:")
print(df_hamiltonian.orbital_rotations.shape)

Original representation
-----------------------
One-body tensor shape:
(8, 8)

Two-body tensor shape:
(8, 8, 8, 8)

Double-factorized representation
--------------------------------
One-body tensor shape:
(8, 8)

Diagonal Coulomb matrices shape:
(35, 8, 8)

Orbital rotations shape:
(35, 8, 8)

We see that instead of an \(N \times N \times N \times N\) two-body tensor (\(N\) is the number of spatial orbitals), the double-factorized representation stores a list of \(L\) diagonal Coulomb matrices as well as a list of \(L\) orbital rotations. Here, \(L = 35\). The value of \(L\) depends on the per-tensor-entry error tolerance allowed in the double factorization, which defaults to \(10^{-8}\). Setting a higher error tolerance may yield a more compact representation with a smaller \(L\). Let’s see this in action by setting the error tolerance to \(10^{-3}\).

df_hamiltonian_alt = ffsim.DoubleFactorizedHamiltonian.from_molecular_hamiltonian(
    mol_hamiltonian, tol=1e-3
)
print(f"Number of terms: {len(df_hamiltonian_alt.diag_coulomb_mats)}")

Number of terms: 26

With an error tolerance of \(10^{-3}\), the factorization results in \(L = 26\).

In addition to setting the error tolerance, you can also specify a maximum value for \(L\) via the max_vecs argument. Bear in mind that setting a low value for the maximum number of terms may introduce significant error in the decomposition. The max_vecs argument is always respected, so the resulting decomposition may exceed the error tolerance specified by the tol argument. Let’s try setting the maximum number of terms to 10.

df_hamiltonian_alt = ffsim.DoubleFactorizedHamiltonian.from_molecular_hamiltonian(
    mol_hamiltonian, max_vecs=10
)
print(f"Number of terms: {len(df_hamiltonian_alt.diag_coulomb_mats)}")

Number of terms: 10

The error in the decomposition can be computed by reconstructing the two-body tensor from its factorized form. The following code cell performs this reconstruction using np.einsum, then prints out the maximum error in one of the tensor entries.

import numpy as np

reconstructed = np.einsum(
    "kij,kpi,kqi,krj,ksj->pqrs",
    df_hamiltonian_alt.diag_coulomb_mats,
    df_hamiltonian_alt.orbital_rotations,
    df_hamiltonian_alt.orbital_rotations,
    df_hamiltonian_alt.orbital_rotations,
    df_hamiltonian_alt.orbital_rotations,
)
max_error = np.max(np.abs(reconstructed - mol_hamiltonian.two_body_tensor))

print(f"Maximum error in a tensor entry: {max_error}")

Maximum error in a tensor entry: 0.03668541730983732

Implement Trotter simulation

As explained in Double-factorized representation of the molecular Hamiltonian, the doubled-factorized Hamiltonian can be expressed as a sum of \(L + 1\) terms,

\[H = \sum_{k=0}^L H_k,\]

where

• \(H_0\) is a quadratic Hamiltonian, and

• \(H_k\) is a rotated diagonal Coulomb operator for \(k = 1, \ldots, L\).

Let’s write a function to simulate a single Trotter step of the Hamiltonian. Our function will perform the following steps:

1. Diagonalize \(H_0\) to obtain the orbital energies and rotation needed to simulate it.

2. Apply time evolution by \(H_0\) using the function apply_num_op_sum_evolution.

3. For \(k = 1, \ldots, L\), apply time evolution by \(H_k\) using the function apply_diag_coulomb_evolution.

import numpy as np


def simulate_trotter_step_double_factorized(
    vec: np.ndarray,
    hamiltonian: ffsim.DoubleFactorizedHamiltonian,
    time: float,
    norb: int,
    nelec: tuple[int, int],
) -> np.ndarray:
    # Diagonalize the one-body term
    one_body_energies, one_body_basis_change = np.linalg.eigh(
        hamiltonian.one_body_tensor
    )
    # Simulate the one-body term
    vec = ffsim.apply_num_op_sum_evolution(
        vec,
        one_body_energies,
        time,
        norb=norb,
        nelec=nelec,
        orbital_rotation=one_body_basis_change,
    )
    # Simulate the two-body terms
    for diag_coulomb_mat, orbital_rotation in zip(
        hamiltonian.diag_coulomb_mats, hamiltonian.orbital_rotations
    ):
        vec = ffsim.apply_diag_coulomb_evolution(
            vec,
            diag_coulomb_mat,
            time,
            norb=norb,
            nelec=nelec,
            orbital_rotation=orbital_rotation,
        )
    return vec

To finish, we need to write a higher-level function that handles splitting the total time evolution into multiple Trotter steps and simulates each Trotter step using the function we just wrote.

def simulate_trotter_double_factorized(
    vec: np.ndarray,
    hamiltonian: ffsim.DoubleFactorizedHamiltonian,
    time: float,
    norb: int,
    nelec: tuple[int, int],
    n_steps: int = 1,
) -> np.ndarray:
    step_time = time / n_steps
    for _ in range(n_steps):
        vec = simulate_trotter_step_double_factorized(
            vec,
            hamiltonian,
            step_time,
            norb=norb,
            nelec=nelec,
        )
    return vec

To test our implementation, let’s apply time evolution to the Hartree-Fock state. Before calling our Trotter simulation function, let’s first compute the exact result of time evolution by directly exponentiating the Hamiltonian using SciPy. Later, we’ll compare the result of our approximate time evolution with this exact result. In order to perform the operator exponentiation, we convert the Hamiltonian to a Scipy LinearOperator (see Hamiltonians).

import scipy.sparse.linalg

# Construct the initial state.
initial_state = ffsim.hartree_fock_state(norb, nelec)

# Set the evolution time.
time = 1.0

# Convert the Hamiltonian to a LinearOperator
linop = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)

# Compute the exact result of time evolution
exact_state = scipy.sparse.linalg.expm_multiply(
    -1j * time * linop,
    initial_state,
    traceA=-1j * time * ffsim.trace(mol_hamiltonian, norb=norb, nelec=nelec),
)

Now, let’s test our implementation. First, let’s evolve the initial state using a single Trotter step.

final_state = simulate_trotter_double_factorized(
    initial_state,
    df_hamiltonian,
    time,
    norb=norb,
    nelec=nelec,
    n_steps=1,
)

fidelity = abs(np.vdot(final_state, exact_state))
print(f"Fidelity of Trotter-evolved state with exact state: {fidelity}")

Fidelity of Trotter-evolved state with exact state: 0.9402387920715227

The fidelity of the final result can be improved by increasing the number of Trotter steps.

final_state = simulate_trotter_double_factorized(
    initial_state,
    df_hamiltonian,
    time,
    norb=norb,
    nelec=nelec,
    n_steps=5,
)

fidelity = abs(np.vdot(final_state, exact_state))
print(f"Fidelity of Trotter-evolved state with exact state: {fidelity}")

Fidelity of Trotter-evolved state with exact state: 0.9985210983148279

In the code cell below, we reproduce the results of our manually implemented function using ffsim’s built-in implementation.

final_state = ffsim.simulate_trotter_double_factorized(
    initial_state,
    df_hamiltonian,
    time,
    norb=norb,
    nelec=nelec,
    n_steps=5,
    order=0,
)

fidelity = abs(np.vdot(final_state, exact_state))
print(f"Fidelity of Trotter-evolved state with exact state: {fidelity}")

Fidelity of Trotter-evolved state with exact state: 0.9985210983145961

A higher order formula achieves a higher fidelity with fewer Trotter steps:

final_state = ffsim.simulate_trotter_double_factorized(
    initial_state,
    df_hamiltonian,
    time,
    norb=norb,
    nelec=nelec,
    n_steps=2,
    order=1,
)

fidelity = abs(np.vdot(final_state, exact_state))
print(f"Fidelity of Trotter-evolved state with exact state: {fidelity}")

Fidelity of Trotter-evolved state with exact state: 0.9996731173230682

You’ve made it to the end of this tutorial!