Scale SQD chemistry workflows with Dice solver

For information on how to install and use qiskit-addon-dice-solver, visit the docs.

For more details on the SQD code used in this example, check out tutorial 1.

%%capture


import numpy as np
import pyscf
import pyscf.cc
import pyscf.mcscf
from qiskit_addon_dice_solver import solve_fermion
from qiskit_addon_sqd.configuration_recovery import recover_configurations
from qiskit_addon_sqd.counts import counts_to_arrays, generate_counts_uniform
from qiskit_addon_sqd.fermion import (
    flip_orbital_occupancies,
)
from qiskit_addon_sqd.subsampling import postselect_and_subsample

# Specify molecule properties
num_orbitals = 16
num_elec_a = num_elec_b = 5
open_shell = False
spin_sq = 0

# Specify molecule properties
open_shell = False
spin_sq = 0

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

# Define active space
n_frozen = 2
active_space = range(n_frozen, mol.nao_nr())

# Get molecular integrals
scf = pyscf.scf.RHF(mol).run()
num_orbitals = len(active_space)
n_electrons = int(sum(scf.mo_occ[active_space]))
num_elec_a = (n_electrons + mol.spin) // 2
num_elec_b = (n_electrons - mol.spin) // 2
cas = pyscf.mcscf.CASCI(scf, num_orbitals, (num_elec_a, num_elec_b))
mo = cas.sort_mo(active_space, base=0)
hcore, nuclear_repulsion_energy = cas.get_h1cas(mo)
eri = pyscf.ao2mo.restore(1, cas.get_h2cas(mo), num_orbitals)

# Compute exact energy
exact_energy = cas.run().e_tot

# Create a seed to control randomness throughout this workflow
rng = np.random.default_rng(24)

# Generate random samples
counts_dict = generate_counts_uniform(10_000, num_orbitals * 2, rand_seed=rng)

# Convert counts into bitstring and probability arrays
bitstring_matrix_full, probs_arr_full = counts_to_arrays(counts_dict)

# SQD options
iterations = 5

# Eigenstate solver options
n_batches = 5
samples_per_batch = 300

# Self-consistent configuration recovery loop
occupancies_bitwise = None  # orbital i corresponds to column i in bitstring matrix
e_hist = np.zeros((iterations, n_batches))
for i in range(iterations):
    # On the first iteration, we have no orbital occupancy information from the
    # solver, so we just post-select from the full bitstring set based on hamming weight.
    if occupancies_bitwise is None:
        bs_mat_tmp = bitstring_matrix_full
        probs_arr_tmp = probs_arr_full

    # In following iterations, we use both the occupancy info and the target hamming
    # weight to refine bitstrings.
    else:
        bs_mat_tmp, probs_arr_tmp = recover_configurations(
            bitstring_matrix_full,
            probs_arr_full,
            occupancies_bitwise,
            num_elec_a,
            num_elec_b,
            rand_seed=rng,
        )

    # Throw out samples with incorrect hamming weight and create batches of subsamples.
    batches = postselect_and_subsample(
        bs_mat_tmp,
        probs_arr_tmp,
        hamming_right=num_elec_a,
        hamming_left=num_elec_b,
        samples_per_batch=samples_per_batch,
        num_batches=n_batches,
        rand_seed=rng,
    )
    # Run eigenstate solvers in a loop. This loop should be parallelized for larger problems.
    int_e = np.zeros(n_batches)
    int_occs = np.zeros((n_batches, 2 * num_orbitals))
    for j, batch in enumerate(batches):
        energy_sci, wf_mags, avg_occs = solve_fermion(
            batch,
            hcore,
            eri,
            mpirun_options=["-n", "8"],
        )
        int_e[j] = energy_sci + nuclear_repulsion_energy
        int_occs[j, :num_orbitals] = avg_occs[0]
        int_occs[j, num_orbitals:] = avg_occs[1]

    # Combine batch results
    avg_occupancy = np.mean(int_occs, axis=0)
    # The occupancies from the solver should be flipped to match the bits in the bitstring matrix.
    occupancies_bitwise = flip_orbital_occupancies(avg_occupancy)

    # Track optimization history
    e_hist[i, :] = int_e

print("Exact energy: {exact_energy}")
print(f"Estimated energy: {np.min(e_hist[-1])}")

Exact energy: -109.10288938
Estimated energy: -109.03013363477585