diff --git a/docs/how-to-guides/index.md b/docs/how-to-guides/index.md index 143ddda50..57c22a882 100644 --- a/docs/how-to-guides/index.md +++ b/docs/how-to-guides/index.md @@ -7,4 +7,5 @@ lucj entanglement-forging fermion-operator qiskit-circuits +qiskit-sampler ``` diff --git a/docs/how-to-guides/qiskit-sampler.ipynb b/docs/how-to-guides/qiskit-sampler.ipynb new file mode 100644 index 000000000..cde390ebb --- /dev/null +++ b/docs/how-to-guides/qiskit-sampler.ipynb @@ -0,0 +1,256 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# How to use the ffsim Sampler primitive\n", + "\n", + "In this guide, we show how to use [FfsimSampler](../api/ffsim.qiskit.rst#ffsim.qiskit.FfsimSampler), ffsim's implementation of the Qiskit [Sampler primitive](https://docs.quantum.ibm.com/api/qiskit/qiskit.primitives.BaseSamplerV2), to sample from quantum circuits built from fermionic gates.\n", + "\n", + "## Example of using FfsimSampler\n", + "\n", + "First, let's create an example circuit using gates from the [ffsim.qiskit](../api/ffsim.qiskit.rst) module.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "from qiskit.circuit import QuantumCircuit, QuantumRegister\n", + "\n", + "import ffsim\n", + "\n", + "# Let's use 8 spatial orbitals with 4 alpha electrons and 4 beta electrons.\n", + "norb = 8\n", + "nelec = (4, 4)\n", + "n_alpha, n_beta = nelec\n", + "\n", + "# Generate some random data\n", + "rng = np.random.default_rng(12345)\n", + "orbital_rotation = ffsim.random.random_unitary(norb, seed=rng)\n", + "diag_coulomb_mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng)\n", + "\n", + "# Create an example circuit\n", + "qubits = QuantumRegister(2 * norb)\n", + "circuit = QuantumCircuit(qubits)\n", + "circuit.append(\n", + " ffsim.qiskit.PrepareSlaterDeterminantJW(\n", + " norb,\n", + " occupied_orbitals=[range(n_alpha), range(n_beta)],\n", + " orbital_rotation=orbital_rotation,\n", + " ),\n", + " qubits,\n", + ")\n", + "circuit.append(\n", + " ffsim.qiskit.DiagCoulombEvolutionJW(norb, diag_coulomb_mat, time=1.0), qubits\n", + ")\n", + "circuit.append(ffsim.qiskit.OrbitalRotationJW(norb, orbital_rotation.T.conj()), qubits)\n", + "circuit.measure_all()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next, we initialize the ffsim sampler and use it to sample 10,000 shots from our circuit. The input to the sampler is a list of [primitive unified blocs](https://docs.quantum.ibm.com/api/qiskit/primitives#overview-of-samplerv2), or PUBs. In the cell output we display only the top 10 most commonly encountered bitstrings.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'0000111100001111': 210,\n", + " '0000111100101011': 133,\n", + " '0010101100001111': 101,\n", + " '0000111100011101': 70,\n", + " '0010011100101101': 66,\n", + " '0010110100101011': 65,\n", + " '0010110100100111': 57,\n", + " '0001110100001111': 56,\n", + " '0000111100011011': 48,\n", + " '0010101100101101': 47}" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Initialize ffsim sampler\n", + "sampler = ffsim.qiskit.FfsimSampler(default_shots=10_000, seed=rng)\n", + "\n", + "# Form PUB, submit job, retrieve job result, and extract first (and only) PUB result\n", + "pub = (circuit,)\n", + "job = sampler.run([pub])\n", + "result = job.result()\n", + "pub_result = result[0]\n", + "\n", + "# Get counts\n", + "counts = pub_result.data.meas.get_counts()\n", + "\n", + "# Display the 10 most commonly seen bitstrings and their counts\n", + "{k: v for k, v in sorted(counts.items(), key=lambda x: x[1], reverse=True)[:10]}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "And that's it!\n", + "\n", + "## Criteria for circuits that FfsimSampler can sample\n", + "\n", + "FfsimSampler cannot sample from arbitrary Qiskit circuits. It can handle circuits that satisfy the following criteria:\n", + "\n", + "- The circuit contains only the following types of gates:\n", + " - Any gate from the [ffsim.qiskit](../api/ffsim.qiskit.rst) module.\n", + " - [Barrier](https://docs.quantum.ibm.com/api/qiskit/qiskit.circuit.library.Barrier#barrier).\n", + " - [Measure](https://docs.quantum.ibm.com/api/qiskit/qiskit.circuit.library.Measure#measure).\n", + "- The first gate of the circuit is one of the following gates:\n", + " - [ffsim.qiskit.PrepareHartreeFockJW](../api/ffsim.qiskit.rst#ffsim.qiskit.PrepareHartreeFockJW).\n", + " - [ffsim.qiskit.PrepareHartreeFockSpinlessJW](../api/ffsim.qiskit.rst#ffsim.qiskit.PrepareHartreeFockSpinlessJW).\n", + " - [ffsim.qiskit.PrepareSlaterDeterminantJW](../api/ffsim.qiskit.rst#ffsim.qiskit.PrepareSlaterDeterminantJW).\n", + " - [ffsim.qiskit.PrepareSlaterDeterminantSpinlessJW](../api/ffsim.qiskit.rst#ffsim.qiskit.PrepareSlaterDeterminantSpinlessJW).\n", + "- The state preparation gates listed in the previous item only appear as the first gate of the circuit.\n", + "- All measurements are performed at the end of the circuit. In other words, the circuit does not contain any mid-circuit measurements.\n", + "\n", + "Due to the need to preserve fermionic statistics, some of the supported gates must be applied to consecutive qubits in ascending order.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example of sampling from an LUCJ circuit for a nitrogen molecule\n", + "\n", + "The following code cell demonstrates a possible workflow for sampling from an LUCJ circuit for a nitrogen molecule in the 6-31g basis.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "converged SCF energy = -108.835236570775\n", + "norb = 16\n", + "nelec = (5, 5)\n", + "E(CCSD) = -109.0398256929734 E_corr = -0.2045891221988307\n" + ] + }, + { + "data": { + "text/plain": [ + "{'00000000000111110000000000011111': 9897,\n", + " '00000000010110110000000000011111': 28,\n", + " '00000000000111110000000001011011': 20,\n", + " '00100000000110110000000000011111': 7,\n", + " '10000000000011110000000000011111': 7,\n", + " '00000000000111110000000000110111': 6,\n", + " '00000000000111110010000000011011': 6,\n", + " '00000000001101110000000000110111': 3,\n", + " '01000000000011110000000000011111': 3,\n", + " '00000000000111111000000000001111': 3}" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import pyscf.data.elements\n", + "from pyscf import cc, gto\n", + "\n", + "# Build N2 molecule\n", + "mol = gto.Mole()\n", + "mol.build(\n", + " atom=[[\"N\", (0, 0, 0)], [\"N\", (1.0, 0, 0)]],\n", + " basis=\"6-31g\",\n", + " symmetry=\"Dooh\",\n", + ")\n", + "\n", + "# Define active space\n", + "n_frozen = pyscf.data.elements.chemcore(mol)\n", + "active_space = range(n_frozen, mol.nao_nr())\n", + "\n", + "# Get molecular data and Hamiltonian\n", + "mol_data = ffsim.MolecularData.from_mole(mol, active_space=active_space)\n", + "norb, nelec = mol_data.norb, mol_data.nelec\n", + "mol_hamiltonian = mol_data.hamiltonian\n", + "print(f\"norb = {norb}\")\n", + "print(f\"nelec = {nelec}\")\n", + "\n", + "# Get CCSD t2 amplitudes for initializing the ansatz\n", + "ccsd = cc.CCSD(\n", + " mol_data.scf,\n", + " frozen=[i for i in range(mol.nao_nr()) if i not in active_space],\n", + ")\n", + "_, _, t2 = ccsd.kernel()\n", + "\n", + "# Construct LUCJ operator using indices for a square lattice mapping\n", + "n_reps = 1\n", + "alpha_alpha_indices = [(p, p + 1) for p in range(norb - 1)]\n", + "alpha_beta_indices = [(p, p) for p in range(norb)]\n", + "ucj_op = ffsim.UCJOperator.from_parameters(\n", + " ffsim.UCJOperator.from_t_amplitudes(t2).to_parameters(),\n", + " norb=norb,\n", + " n_reps=n_reps,\n", + " alpha_alpha_indices=alpha_alpha_indices,\n", + " alpha_beta_indices=alpha_beta_indices,\n", + ")\n", + "\n", + "# Construct circuit\n", + "qubits = QuantumRegister(2 * norb)\n", + "circuit = QuantumCircuit(qubits)\n", + "circuit.append(ffsim.qiskit.PrepareHartreeFockJW(norb, nelec), qubits)\n", + "circuit.append(ffsim.qiskit.UCJOperatorJW(ucj_op), qubits)\n", + "circuit.measure_all()\n", + "\n", + "# Sample 10,000 shots from the circuit using FfsimSampler\n", + "sampler = ffsim.qiskit.FfsimSampler(default_shots=10_000, seed=12345)\n", + "pub = (circuit,)\n", + "job = sampler.run([pub])\n", + "result = job.result()\n", + "pub_result = result[0]\n", + "counts = pub_result.data.meas.get_counts()\n", + "\n", + "# Display the 10 most commonly seen bitstrings and their counts\n", + "{k: v for k, v in sorted(counts.items(), key=lambda x: x[1], reverse=True)[:10]}" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "ffsim", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/pyproject.toml b/pyproject.toml index 98895e067..3fdfd27d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "numpy", "opt_einsum", "pyscf >= 2.3", - "qiskit >= 1.0", + "qiskit >= 1.1", "scipy", "typing-extensions", ] diff --git a/python/ffsim/qiskit/__init__.py b/python/ffsim/qiskit/__init__.py index 493f9d0f8..4d25b38d7 100644 --- a/python/ffsim/qiskit/__init__.py +++ b/python/ffsim/qiskit/__init__.py @@ -24,6 +24,7 @@ PrepareSlaterDeterminantSpinlessJW, UCJOperatorJW, ) +from ffsim.qiskit.sampler import FfsimSampler from ffsim.qiskit.transpiler_passes import DropNegligible, MergeOrbitalRotations from ffsim.qiskit.transpiler_stages import pre_init_passes from ffsim.qiskit.util import ffsim_vec_to_qiskit_vec, qiskit_vec_to_ffsim_vec @@ -39,6 +40,7 @@ __all__ = [ "DiagCoulombEvolutionJW", "DropNegligible", + "FfsimSampler", "GivensAnsatzOperatorJW", "GivensAnsatzOperatorSpinlessJW", "MergeOrbitalRotations", diff --git a/python/ffsim/qiskit/sampler.py b/python/ffsim/qiskit/sampler.py new file mode 100644 index 000000000..6a7924252 --- /dev/null +++ b/python/ffsim/qiskit/sampler.py @@ -0,0 +1,196 @@ +# (C) Copyright IBM 2024. +# +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +# +# Any modifications or derivative works of this code must retain this +# copyright notice, and modified files need to carry a notice indicating +# that they have been altered from the originals. + +"""ffsim's implementation of the Qiskit Sampler primitive.""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass + +import numpy as np +from numpy.typing import NDArray +from qiskit.circuit import ClassicalRegister, QuantumCircuit +from qiskit.primitives import ( + BaseSamplerV2, + BitArray, + DataBin, + PrimitiveJob, + PrimitiveResult, + SamplerPubLike, + SamplerPubResult, +) + +# TODO import from qiskit.primitives after qiskit 1.2 +from qiskit.primitives.containers.sampler_pub import SamplerPub + +from ffsim.qiskit.sim import final_statevector, sample_statevector + + +class FfsimSampler(BaseSamplerV2): + """Implementation of the Qiskit Sampler primitive backed by ffsim.""" + + def __init__( + self, + *, + default_shots: int = 1024, + seed: np.random.Generator | int | None = None, + ): + """Initialize the ffsim sampler. + + Args: + default_shots: The default shots to use if not specified during run. + seed: A seed to initialize the pseudorandom number generator. + Should be a valid input to ``np.random.default_rng``. + """ + self._default_shots = default_shots + self._rng = np.random.default_rng(seed) + + def run( + self, pubs: Iterable[SamplerPubLike], *, shots: int | None = None + ) -> PrimitiveJob[PrimitiveResult[SamplerPubResult]]: + if shots is None: + shots = self._default_shots + coerced_pubs = [SamplerPub.coerce(pub, shots) for pub in pubs] + job = PrimitiveJob(self._run, coerced_pubs) + job._submit() + return job + + def _run(self, pubs: Iterable[SamplerPub]) -> PrimitiveResult[SamplerPubResult]: + results = [self._run_pub(pub) for pub in pubs] + return PrimitiveResult(results) + + def _run_pub(self, pub: SamplerPub) -> SamplerPubResult: + circuit, qargs, meas_info = _preprocess_circuit(pub.circuit) + bound_circuits = pub.parameter_values.bind_all(circuit) + arrays = { + item.creg_name: np.zeros( + bound_circuits.shape + (pub.shots, item.num_bytes), dtype=np.uint8 + ) + for item in meas_info + } + for index, bound_circuit in np.ndenumerate(bound_circuits): + final_state = final_statevector(bound_circuit) + if qargs: + samples = sample_statevector( + final_state, indices=qargs, shots=pub.shots, seed=self._rng + ) + else: + samples = [""] * pub.shots + samples_array = np.array( + [np.fromiter(sample, dtype=np.uint8) for sample in samples] + ) + for item in meas_info: + ary = _samples_to_packed_array( + samples_array, item.num_bits, item.qreg_indices + ) + arrays[item.creg_name][index] = ary + + meas = { + item.creg_name: BitArray(arrays[item.creg_name], item.num_bits) + for item in meas_info + } + return SamplerPubResult( + DataBin(**meas, shape=pub.shape), metadata={"shots": pub.shots} + ) + + +@dataclass +class _MeasureInfo: + creg_name: str + num_bits: int + num_bytes: int + qreg_indices: list[int] + + +def _min_num_bytes(num_bits: int) -> int: + """Return the minimum number of bytes needed to store ``num_bits``.""" + return num_bits // 8 + (num_bits % 8 > 0) + + +def _preprocess_circuit(circuit: QuantumCircuit): + num_bits_dict = {creg.name: creg.size for creg in circuit.cregs} + mapping = _final_measurement_mapping(circuit) + qargs = sorted(set(mapping.values())) + qargs_index = {v: k for k, v in enumerate(qargs)} + circuit = circuit.remove_final_measurements(inplace=False) + # num_qubits is used as sentinel to fill 0 in _samples_to_packed_array + sentinel = len(qargs) + indices = {key: [sentinel] * val for key, val in num_bits_dict.items()} + for key, qreg in mapping.items(): + creg, ind = key + indices[creg.name][ind] = qargs_index[qreg] + meas_info = [ + _MeasureInfo( + creg_name=name, + num_bits=num_bits, + num_bytes=_min_num_bytes(num_bits), + qreg_indices=indices[name], + ) + for name, num_bits in num_bits_dict.items() + ] + return circuit, qargs, meas_info + + +def _final_measurement_mapping( + circuit: QuantumCircuit, +) -> dict[tuple[ClassicalRegister, int], int]: + """Return the final measurement mapping for the circuit. + + Parameters: + circuit: Input quantum circuit. + + Returns: + Mapping of classical bits to qubits for final measurements. + """ + active_qubits = set(range(circuit.num_qubits)) + active_cbits = set(range(circuit.num_clbits)) + + # Find final measurements starting in back + mapping = {} + for item in circuit[::-1]: + if item.operation.name == "measure": + loc = circuit.find_bit(item.clbits[0]) + cbit = loc.index + qbit = circuit.find_bit(item.qubits[0]).index + if cbit in active_cbits and qbit in active_qubits: + for creg in loc.registers: + mapping[creg] = qbit + active_cbits.remove(cbit) + elif item.operation.name not in ["barrier", "delay"]: + for qq in item.qubits: + _temp_qubit = circuit.find_bit(qq).index + if _temp_qubit in active_qubits: + active_qubits.remove(_temp_qubit) + + if not active_cbits or not active_qubits: + break + + return mapping + + +def _samples_to_packed_array( + samples: NDArray[np.uint8], num_bits: int, indices: list[int] +) -> NDArray[np.uint8]: + # samples of `Statevector.sample_memory` will be in the order of + # qubit_last, ..., qubit_1, qubit_0. + # reverse the sample order into qubit_0, qubit_1, ..., qubit_last and + # pad 0 in the rightmost to be used for the sentinel introduced by + # _preprocess_circuit. + ary = np.pad(samples[:, ::-1], ((0, 0), (0, 1)), constant_values=0) + # place samples in the order of clbit_last, ..., clbit_1, clbit_0 + ary = ary[:, indices[::-1]] + # pad 0 in the left to align the number to be mod 8 + # since np.packbits(bitorder='big') pads 0 to the right. + pad_size = -num_bits % 8 + ary = np.pad(ary, ((0, 0), (pad_size, 0)), constant_values=0) + # pack bits in big endian order + ary = np.packbits(ary, axis=-1) + return ary diff --git a/python/ffsim/qiskit/sim.py b/python/ffsim/qiskit/sim.py new file mode 100644 index 000000000..9036d795b --- /dev/null +++ b/python/ffsim/qiskit/sim.py @@ -0,0 +1,259 @@ +# (C) Copyright IBM 2024. +# +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +# +# Any modifications or derivative works of this code must retain this +# copyright notice, and modified files need to carry a notice indicating +# that they have been altered from the originals. + +"""Qiskit circuit simulation utilities.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +from qiskit.circuit import CircuitInstruction, QuantumCircuit +from qiskit.circuit.library import Barrier, Measure + +from ffsim import gates, protocols, states +from ffsim.qiskit.gates import ( + DiagCoulombEvolutionJW, + GivensAnsatzOperatorJW, + GivensAnsatzOperatorSpinlessJW, + OrbitalRotationJW, + OrbitalRotationSpinlessJW, + PrepareHartreeFockJW, + PrepareHartreeFockSpinlessJW, + PrepareSlaterDeterminantJW, + PrepareSlaterDeterminantSpinlessJW, + UCJOperatorJW, +) +from ffsim.spin import Spin + + +@dataclass +class Statevector: + """A state vector in the FCI representation. + + Attributes: + vec: Array of state vector coefficients. + norb: The number of spatial orbitals. + nelec: The number of alpha and beta electrons. + """ + + vec: np.ndarray + norb: int + nelec: tuple[int, int] + + +def final_statevector(circuit: QuantumCircuit) -> Statevector: + """Return the final state vector of a fermionic quantum circuit. + + Args: + circuit: The circuit composed of fermionic gates. + + Returns: + The final state vector that results from applying the circuit to the vacuum + state. + """ + if not circuit.data: + raise ValueError("Circuit must contain at least one instruction.") + statevector = _prepare_statevector(circuit.data[0], circuit) + for instruction in circuit.data[1:]: + statevector = _evolve_statevector(statevector, instruction, circuit) + return statevector + + +def _prepare_statevector( + instruction: CircuitInstruction, circuit: QuantumCircuit +) -> Statevector: + op = instruction.operation + qubit_indices = [circuit.find_bit(qubit).index for qubit in instruction.qubits] + consecutive_sorted = qubit_indices == list( + range(min(qubit_indices), max(qubit_indices) + 1) + ) + + if isinstance(op, PrepareHartreeFockJW): + if not consecutive_sorted: + raise ValueError( + f"Gate of type '{op.__class__.__name__}' must be applied to " + "consecutive qubits, in ascending order." + ) + norb = op.norb + nelec = op.nelec + vec = states.hartree_fock_state(norb, nelec) + return Statevector(vec=vec, norb=norb, nelec=nelec) + + if isinstance(op, PrepareHartreeFockSpinlessJW): + if not consecutive_sorted: + raise ValueError( + f"Gate of type '{op.__class__.__name__}' must be applied to " + "consecutive qubits, in ascending order." + ) + norb = op.norb + nelec = (op.nocc, 0) + vec = states.hartree_fock_state(op.norb, nelec) + return Statevector(vec=vec, norb=norb, nelec=nelec) + + if isinstance(op, PrepareSlaterDeterminantJW): + if not consecutive_sorted: + raise ValueError( + f"Gate of type '{op.__class__.__name__}' must be applied to " + "consecutive qubits, in ascending order." + ) + norb = op.norb + occ_a, occ_b = op.occupied_orbitals + nelec = (len(occ_a), len(occ_b)) + vec = states.slater_determinant( + norb, + occupied_orbitals=op.occupied_orbitals, + orbital_rotation=(op.orbital_rotation_a, op.orbital_rotation_b), + ) + return Statevector(vec=vec, norb=norb, nelec=nelec) + + if isinstance(op, PrepareSlaterDeterminantSpinlessJW): + if not consecutive_sorted: + raise ValueError( + f"Gate of type '{op.__class__.__name__}' must be applied to " + "consecutive qubits, in ascending order." + ) + norb = op.norb + nelec = (len(op.occupied_orbitals), 0) + vec = states.slater_determinant( + op.norb, (op.occupied_orbitals, []), orbital_rotation=op.orbital_rotation + ) + return Statevector(vec=vec, norb=norb, nelec=nelec) + + raise ValueError( + "The first instruction of the circuit must be one of the following gates: " + "PrepareHartreeFockJW, PrepareHartreeFockSpinlessJW, " + "PrepareSlaterDeterminantJW, PrepareSlaterDeterminantSpinlessJW." + ) + + +def _evolve_statevector( + statevector: Statevector, instruction: CircuitInstruction, circuit: QuantumCircuit +) -> Statevector: + op = instruction.operation + qubit_indices = [circuit.find_bit(qubit).index for qubit in instruction.qubits] + consecutive_sorted = qubit_indices == list( + range(min(qubit_indices), max(qubit_indices) + 1) + ) + vec = statevector.vec + norb = statevector.norb + nelec = statevector.nelec + + if isinstance(op, DiagCoulombEvolutionJW): + vec = gates.apply_diag_coulomb_evolution( + vec, + op.mat, + op.time, + norb=norb, + nelec=nelec, + mat_alpha_beta=op.mat_alpha_beta, + z_representation=op.z_representation, + copy=False, + ) + return Statevector(vec=vec, norb=norb, nelec=nelec) + + if isinstance(op, (GivensAnsatzOperatorJW, GivensAnsatzOperatorSpinlessJW)): + if not consecutive_sorted: + raise ValueError( + f"Gate of type '{op.__class__.__name__}' must be applied to " + "consecutive qubits, in ascending order." + ) + vec = protocols.apply_unitary( + vec, op.givens_ansatz_operator, norb=norb, nelec=nelec, copy=False + ) + return Statevector(vec=vec, norb=norb, nelec=nelec) + + if isinstance(op, OrbitalRotationJW): + if not consecutive_sorted: + raise ValueError( + f"Gate of type '{op.__class__.__name__}' must be applied to " + "consecutive qubits, in ascending order." + ) + vec = gates.apply_orbital_rotation( + vec, + op.orbital_rotation_a, + norb=norb, + nelec=nelec, + spin=Spin.ALPHA, + copy=False, + ) + vec = gates.apply_orbital_rotation( + vec, + op.orbital_rotation_b, + norb=norb, + nelec=nelec, + spin=Spin.BETA, + copy=False, + ) + return Statevector(vec=vec, norb=norb, nelec=nelec) + + if isinstance(op, OrbitalRotationSpinlessJW): + if not consecutive_sorted: + raise ValueError( + f"Gate of type '{op.__class__.__name__}' must be applied to " + "consecutive qubits, in ascending order." + ) + vec = gates.apply_orbital_rotation( + vec, op.orbital_rotation, norb=norb, nelec=nelec, copy=False + ) + return Statevector(vec=vec, norb=norb, nelec=nelec) + + if isinstance(op, UCJOperatorJW): + if not consecutive_sorted: + raise ValueError( + f"Gate of type '{op.__class__.__name__}' must be applied to " + "consecutive qubits, in ascending order." + ) + vec = protocols.apply_unitary( + vec, op.ucj_operator, norb=norb, nelec=nelec, copy=False + ) + return Statevector(vec=vec, norb=norb, nelec=nelec) + + if isinstance(op, Barrier): + return statevector + + if isinstance(op, Measure): + raise ValueError( + "Encountered a measurement gate, but only unitary operations are allowed." + ) + + raise ValueError(f"Unsupported gate: {op}.") + + +def sample_statevector( + statevector: Statevector, + *, + indices: list[int], + shots: int, + seed: np.random.Generator | int | None = None, +) -> list[str]: + """Sample bitstrings from a state vector. + + Args: + statevector: The state vector to sample from. + indices: The indices of the orbitals to sample from. The indices range from + ``0`` to ``2 * norb - 1``, with the first half of the range indexing the + spin alpha orbitals, and the second half indexing the spin beta orbitals. + shots: The number of bitstrings to sample. + seed: A seed to initialize the pseudorandom number generator. + Should be a valid input to ``np.random.default_rng``. + + Returns: + The sampled bitstrings, as a list of strings of length `shots`. + """ + rng = np.random.default_rng(seed) + probabilities = np.abs(statevector.vec) ** 2 + samples = rng.choice(len(statevector.vec), size=shots, p=probabilities) + bitstrings = states.indices_to_strings(samples, statevector.norb, statevector.nelec) + if indices == list(range(2 * statevector.norb)): + return bitstrings + return [ + "".join(bitstring[-1 - i] for i in indices[::-1]) for bitstring in bitstrings + ] diff --git a/python/ffsim/states/states.py b/python/ffsim/states/states.py index bdc0c8b9f..c25ad7003 100644 --- a/python/ffsim/states/states.py +++ b/python/ffsim/states/states.py @@ -217,7 +217,7 @@ def slater_determinant_rdm( def indices_to_strings( - indices: Sequence[int], norb: int, nelec: tuple[int, int] + indices: Sequence[int] | np.ndarray, norb: int, nelec: tuple[int, int] ) -> list[str]: """Convert statevector indices to bitstrings. diff --git a/tests/python/qiskit/sampler_test.py b/tests/python/qiskit/sampler_test.py new file mode 100644 index 000000000..5fc19457d --- /dev/null +++ b/tests/python/qiskit/sampler_test.py @@ -0,0 +1,269 @@ +# (C) Copyright IBM 2024. +# +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +# +# Any modifications or derivative works of this code must retain this +# copyright notice, and modified files need to carry a notice indicating +# that they have been altered from the originals. + +"""Tests for FfsimSampler.""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest +from qiskit.circuit import ClassicalRegister, QuantumCircuit, QuantumRegister +from qiskit.primitives import StatevectorSampler + +import ffsim +import ffsim.random.random + + +def _fidelity(probs1: dict, probs2: dict) -> float: + result = 0.0 + for bitstring in probs1.keys() | probs2.keys(): + prob1 = probs1.get(bitstring, 0) + prob2 = probs2.get(bitstring, 0) + result += math.sqrt(prob1 * prob2) + return result**2 + + +def _brickwork(norb: int, n_layers: int): + for i in range(n_layers): + for j in range(i % 2, norb - 1, 2): + yield (j, j + 1) + + +# TODO handle norb = 0 +@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(1, 5))) +def test_random_gates_spinful(norb: int, nelec: tuple[int, int]): + """Test sampler with random gates.""" + rng = np.random.default_rng(12285) + + qubits = QuantumRegister(2 * norb) + + orbital_rotation = ffsim.random.random_unitary(norb, seed=rng) + diag_coulomb_mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + ucj_op = ffsim.random.random_ucj_operator( + norb, n_reps=2, with_final_orbital_rotation=True + ) + interaction_pairs = list(_brickwork(norb, norb)) + thetas = rng.uniform(-np.pi, np.pi, size=len(interaction_pairs)) + givens_ansatz_op = ffsim.GivensAnsatzOperator(norb, interaction_pairs, thetas) + + circuit = QuantumCircuit(qubits) + circuit.append(ffsim.qiskit.PrepareHartreeFockJW(norb, nelec), qubits) + circuit.append(ffsim.qiskit.OrbitalRotationJW(norb, orbital_rotation), qubits) + circuit.append( + ffsim.qiskit.DiagCoulombEvolutionJW(norb, diag_coulomb_mat, time=1.0), qubits + ) + circuit.append(ffsim.qiskit.GivensAnsatzOperatorJW(givens_ansatz_op), qubits) + circuit.append(ffsim.qiskit.UCJOperatorJW(ucj_op), qubits) + circuit.measure_all() + + shots = 3000 + + sampler = ffsim.qiskit.FfsimSampler(default_shots=shots, seed=rng) + pub = (circuit,) + job = sampler.run([pub]) + result = job.result() + pub_result = result[0] + counts = pub_result.data.meas.get_counts() + ffsim_probs = {bitstring: count / shots for bitstring, count in counts.items()} + np.testing.assert_allclose(sum(ffsim_probs.values()), 1) + + sampler = StatevectorSampler(default_shots=shots, seed=rng) + pub = (circuit,) + job = sampler.run([pub]) + result = job.result() + pub_result = result[0] + counts = pub_result.data.meas.get_counts() + qiskit_probs = {bitstring: count / shots for bitstring, count in counts.items()} + np.testing.assert_allclose(sum(qiskit_probs.values()), 1) + + assert _fidelity(ffsim_probs, qiskit_probs) > 0.99 + + +# TODO handle norb = 0 +@pytest.mark.parametrize("norb, nocc", ffsim.testing.generate_norb_nocc(range(1, 5))) +def test_random_gates_spinless(norb: int, nocc: int): + """Test sampler with random spinless gates.""" + rng = np.random.default_rng(52622) + + qubits = QuantumRegister(norb) + + orbital_rotation = ffsim.random.random_unitary(norb, seed=rng) + interaction_pairs = list(_brickwork(norb, norb)) + thetas = rng.uniform(-np.pi, np.pi, size=len(interaction_pairs)) + givens_ansatz_op = ffsim.GivensAnsatzOperator(norb, interaction_pairs, thetas) + + circuit = QuantumCircuit(qubits) + circuit.append(ffsim.qiskit.PrepareHartreeFockSpinlessJW(norb, nocc), qubits) + circuit.append( + ffsim.qiskit.OrbitalRotationSpinlessJW(norb, orbital_rotation), qubits + ) + circuit.append( + ffsim.qiskit.GivensAnsatzOperatorSpinlessJW(givens_ansatz_op), qubits + ) + circuit.measure_all() + + shots = 3000 + + sampler = ffsim.qiskit.FfsimSampler(default_shots=shots, seed=rng) + pub = (circuit,) + job = sampler.run([pub]) + result = job.result() + pub_result = result[0] + counts = pub_result.data.meas.get_counts() + ffsim_probs = {bitstring: count / shots for bitstring, count in counts.items()} + np.testing.assert_allclose(sum(ffsim_probs.values()), 1) + + sampler = StatevectorSampler(default_shots=shots, seed=rng) + pub = (circuit,) + job = sampler.run([pub]) + result = job.result() + pub_result = result[0] + counts = pub_result.data.meas.get_counts() + qiskit_probs = {bitstring: count / shots for bitstring, count in counts.items()} + np.testing.assert_allclose(sum(qiskit_probs.values()), 1) + + assert _fidelity(ffsim_probs, qiskit_probs) > 0.99 + + +# TODO handle norb = 0 +@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(1, 5))) +def test_measure_subset_spinful(norb: int, nelec: tuple[int, int]): + """Test measuring a subset of qubits.""" + rng = np.random.default_rng(5332) + + qubits = QuantumRegister(2 * norb, name="q") + clbits = ClassicalRegister(norb, name="meas") + measured_qubits = list(rng.choice(qubits, size=len(clbits), replace=False)) + measured_clbits = list(rng.choice(clbits, size=len(clbits), replace=False)) + + orbital_rotation = ffsim.random.random_unitary(norb, seed=rng) + occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng) + + circuit = QuantumCircuit(qubits, clbits) + circuit.append( + ffsim.qiskit.PrepareSlaterDeterminantJW( + norb, occupied_orbitals, orbital_rotation=orbital_rotation + ), + qubits, + ) + circuit.measure(measured_qubits, measured_clbits) + + shots = 3000 + + sampler = ffsim.qiskit.FfsimSampler(default_shots=shots, seed=rng) + pub = (circuit,) + job = sampler.run([pub]) + result = job.result() + pub_result = result[0] + counts = pub_result.data.meas.get_counts() + ffsim_probs = {bitstring: count / shots for bitstring, count in counts.items()} + np.testing.assert_allclose(sum(ffsim_probs.values()), 1) + + sampler = StatevectorSampler(default_shots=shots, seed=rng) + pub = (circuit,) + job = sampler.run([pub]) + result = job.result() + pub_result = result[0] + counts = pub_result.data.meas.get_counts() + qiskit_probs = {bitstring: count / shots for bitstring, count in counts.items()} + np.testing.assert_allclose(sum(qiskit_probs.values()), 1) + + assert _fidelity(ffsim_probs, qiskit_probs) > 0.99 + + +# TODO handle norb = 0 +@pytest.mark.parametrize("norb, nocc", ffsim.testing.generate_norb_nocc(range(1, 5))) +def test_measure_subset_spinless(norb: int, nocc: int): + """Test measuring a subset of qubits, spinless.""" + rng = np.random.default_rng(5332) + + qubits = QuantumRegister(norb, name="q") + clbits = ClassicalRegister(norb, name="meas") + measured_qubits = list(rng.choice(qubits, size=len(clbits), replace=False)) + measured_clbits = list(rng.choice(clbits, size=len(clbits), replace=False)) + + orbital_rotation = ffsim.random.random_unitary(norb, seed=rng) + occupied_orbitals = ffsim.testing.random_occupied_orbitals( + norb, (nocc, 0), seed=rng + )[0] + + circuit = QuantumCircuit(qubits, clbits) + circuit.append( + ffsim.qiskit.PrepareSlaterDeterminantSpinlessJW( + norb, occupied_orbitals, orbital_rotation=orbital_rotation + ), + qubits, + ) + circuit.measure(measured_qubits, measured_clbits) + + shots = 3000 + + sampler = ffsim.qiskit.FfsimSampler(default_shots=shots, seed=rng) + pub = (circuit,) + job = sampler.run([pub]) + result = job.result() + pub_result = result[0] + counts = pub_result.data.meas.get_counts() + ffsim_probs = {bitstring: count / shots for bitstring, count in counts.items()} + np.testing.assert_allclose(sum(ffsim_probs.values()), 1) + + sampler = StatevectorSampler(default_shots=shots, seed=rng) + pub = (circuit,) + job = sampler.run([pub]) + result = job.result() + pub_result = result[0] + counts = pub_result.data.meas.get_counts() + qiskit_probs = {bitstring: count / shots for bitstring, count in counts.items()} + np.testing.assert_allclose(sum(qiskit_probs.values()), 1) + + assert _fidelity(ffsim_probs, qiskit_probs) > 0.99 + + +# TODO handle norb = 0 +def test_reproducible_with_seed(): + """Test sampler with random gates.""" + rng = np.random.default_rng(14062) + + norb = 4 + nelec = (2, 2) + + qubits = QuantumRegister(2 * norb, name="q") + + orbital_rotation = ffsim.random.random_unitary(norb, seed=rng) + occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng) + + circuit = QuantumCircuit(qubits) + circuit.append( + ffsim.qiskit.PrepareSlaterDeterminantJW( + norb, occupied_orbitals, orbital_rotation=orbital_rotation + ), + qubits, + ) + circuit.measure_all() + + shots = 3000 + + sampler = ffsim.qiskit.FfsimSampler(default_shots=shots, seed=12345) + pub = (circuit,) + job = sampler.run([pub]) + result = job.result() + pub_result = result[0] + counts_1 = pub_result.data.meas.get_counts() + + sampler = ffsim.qiskit.FfsimSampler(default_shots=shots, seed=12345) + pub = (circuit,) + job = sampler.run([pub]) + result = job.result() + pub_result = result[0] + counts_2 = pub_result.data.meas.get_counts() + + assert counts_1 == counts_2