Operator averaging for molecular Hamiltonians

Download Notebook - operator_averaging.ipynb

"""Example showing operator averaging on the H2 STO-3G Hartree-Fock state."""

from __future__ import annotations

from typing import no_type_check

import zixy.qubit.pauli as zqp
from zixy.fermion.mappings import JordanWignerMapper
from zixy.fermion.operator.general import String as FermionString
from guppylang import guppy
from guppylang.std.builtins import array, comptime
from guppylang.std.quantum import discard_array, qubit, x

from guppyalgos.primitives.measurement import (
    estimate_expectation_from_binary_samples,
    estimate_pauli_observable_expectation_from_binary_samples,
    make_hadamard_test_pauli,
)
from guppyalgos.utils import qarray

H2_STO3G_HF_ENERGY = -1.1175058842043306
H2_STO3G_XY_SHOTS = 512
H2_STO3G_SHOTS = 64
H2_HAM_QUBITS = 4 # for parsing paulis
H2_TOTAL_QUBITS = H2_HAM_QUBITS + 1  # for emulator

H2_STO3G_FERMION_OPERATOR_ZIXY = [
    # is mode-indexed
    (0.7430177069924181, ()),
    (-1.2702927243904383, ((0, 1), (0, 0))),  # F0^ F0
    (-1.2702927243904380, ((1, 1), (1, 0))),
    (-0.45680735030941033, ((2, 1), (2, 0))),
    (-0.45680735030941020, ((3, 1), (3, 0))),
    (0.6800618575841273, ((0, 1), (0, 0), (1, 1), (1, 0))),  # F0^ F0 F1^ F1
    (0.48890859745047305, ((0, 1), (0, 0), (2, 1), (2, 0))),
    (0.6685772770134887, ((0, 1), (0, 0), (3, 1), (3, 0))),
    (0.6685772770134887, ((1, 1), (1, 0), (2, 1), (2, 0))),
    (0.48890859745047305, ((1, 1), (1, 0), (3, 1), (3, 0))),
    (0.7028135332762804, ((2, 1), (2, 0), (3, 1), (3, 0))),
    (-0.35933735912603115, ((0, 1), (1, 1), (2, 0), (3, 0))),  # F0^ F1^ F2 F3
    (-0.35933735912603115, ((0, 1), (1, 0), (3, 1), (2, 0))),  # F0^ F1 F3^ F2
]

@guppy
@no_type_check
def prepare_h2_sto3g_hf(qreg: array[qubit, comptime(4)]) -> None:
    """Prepare the H2 STO-3G Hartree-Fock reference state.

    The spin-orbitals are ordered from lowest to highest orbital energy,
    matching `H2_STO3G_FERMION_OPERATOR_ZIXY`. Modes 0 and 1 are occupied.
    """
    x(qreg[0])
    x(qreg[1])


def build_hadamard_program(pauli_string: zqp.String, size: int, prepare_state):
    hadamard_test_pauli = make_hadamard_test_pauli(pauli_string, size)

    @guppy
    @no_type_check
    def program() -> None:
        qreg = qarray(comptime(size))
        prepare_state(qreg)
        hadamard_test_pauli(qreg)
        discard_array(qreg)

    return program

def _shots_for_term(term: zqp.RealTerm) -> int:
    """Use a larger shot budget for the off-diagonal exchange terms."""
    if any(label in (zqp.X, zqp.Y) for label in term.string.get_dict().values()):
        return H2_STO3G_XY_SHOTS
    return H2_STO3G_SHOTS

"""Run the operator averaging demo."""
single_pauli = zqp.String.from_str("Z0", H2_HAM_QUBITS)
single_program = build_hadamard_program(
    single_pauli,
    H2_HAM_QUBITS,
    prepare_h2_sto3g_hf,
)
result = single_program.emulator(n_qubits=H2_TOTAL_QUBITS).with_shots(10).run()
ancilla_samples = [bool(sample["ancilla"][0]) for sample in result.collated_shots()]
single_estimate = estimate_expectation_from_binary_samples(ancilla_samples)

print(
    f"<{single_pauli}> on HF 1100 = {single_estimate.expectation:.6f} "
    f"+/- {single_estimate.standard_error:.6f}"
)

<Z0> on HF 1100 = -1.000000 +/- 0.000000

jw = JordanWignerMapper(H2_HAM_QUBITS, mode_ordering=None)
h2_ham_op = zqp.RealTermSum.from_str("(0.0, I0)", H2_HAM_QUBITS)

for coeff, ops in H2_STO3G_FERMION_OPERATOR_ZIXY:
    if len(ops) == 0:
        h2_ham_op += zqp.RealTermSum.from_str(f"({coeff}, I0)", H2_HAM_QUBITS)
    else:
        fermion_string = FermionString(
            H2_HAM_QUBITS, [(mode, bool(adj)) for mode, adj in ops]
        )
        h2_ham_op += jw.encode_real(fermion_string, coeff)

operator = h2_ham_op

#print(operator)

operator_samples: dict[str, list[bool]] = {}

for term in operator.to_terms():
    if term.string.is_identity():
        continue

    term_program = build_hadamard_program(
        term.string,
        H2_HAM_QUBITS,
        prepare_h2_sto3g_hf,
    )
    result = term_program.emulator(n_qubits=H2_TOTAL_QUBITS).with_shots(_shots_for_term(term)).run()
    operator_samples[str(term.string)] = [bool(sample["ancilla"][0]) for sample in result.collated_shots()]


operator_estimate = estimate_pauli_observable_expectation_from_binary_samples(operator, operator_samples)
operator_estimate.print_terms()
print(
    f"\nH2 STO-3G energy on HF 1100 = {operator_estimate.expectation:.6f} "
    f"+/- {operator_estimate.standard_error:.6f}"
)
print(f"Reference HF energy = {H2_STO3G_HF_ENERGY:.12f}")

term                exp         var      stderr  shots     +     -
I             +1.000000    0.000000    0.000000      0     0     0
X0 X1 Y2 Y3   +0.023438    0.001952    0.044182    512   262   250
X0 Y1 Y2 X3   -0.027344    0.001952    0.044178    512   249   263
Y0 X1 X2 Y3   +0.019531    0.001952    0.044186    512   261   251
Y0 Y1 X2 X3   +0.031250    0.001951    0.044173    512   264   248
Z0            -1.000000    0.000000    0.000000     64     0    64
Z0 Z1         +1.000000    0.000000    0.000000     64    64     0
Z0 Z2         -1.000000    0.000000    0.000000     64     0    64
Z0 Z3         -1.000000    0.000000    0.000000     64     0    64
Z1            -1.000000    0.000000    0.000000     64     0    64
Z1 Z2         -1.000000    0.000000    0.000000     64     0    64
Z1 Z3         -1.000000    0.000000    0.000000     64     0    64
Z2            +1.000000    0.000000    0.000000     64    64     0
Z2 Z3         +1.000000    0.000000    0.000000     64    64     0
Z3            +1.000000    0.000000    0.000000     64    64     0

H2 STO-3G energy on HF 1100 = -1.120313 +/- 0.003969
Reference HF energy = -1.117505884204