diff --git a/qadence/measurements/shadow.py b/qadence/measurements/shadow.py index a9c25cda..29978d22 100644 --- a/qadence/measurements/shadow.py +++ b/qadence/measurements/shadow.py @@ -1,40 +1,37 @@ from __future__ import annotations -from collections import Counter -from functools import reduce - import numpy as np import torch from torch import Tensor from qadence.backend import Backend from qadence.backends.pyqtorch import Backend as PyQBackend -from qadence.blocks import AbstractBlock, chain, kron -from qadence.blocks.block_to_tensor import HMAT, IMAT, SDAGMAT, ZMAT, block_to_tensor +from qadence.blocks import AbstractBlock, KronBlock, chain, kron +from qadence.blocks.block_to_tensor import HMAT, IMAT, SDAGMAT from qadence.blocks.composite import CompositeBlock from qadence.blocks.primitive import PrimitiveBlock from qadence.blocks.utils import get_pauli_blocks, unroll_block_with_scaling from qadence.circuit import QuantumCircuit from qadence.engines.differentiable_backend import DifferentiableBackend +from qadence.measurements.utils import get_qubit_indices_for_op from qadence.noise import NoiseHandler -from qadence.operations import X, Y, Z +from qadence.operations import H, I, SDagger, X, Y, Z from qadence.types import Endianness -from qadence.utils import P0_MATRIX, P1_MATRIX pauli_gates = [X, Y, Z] - +pauli_rotations = [ + lambda index: H(index), + lambda index: SDagger(index) * H(index), + lambda index: None, +] UNITARY_TENSOR = [ - ZMAT @ HMAT, - SDAGMAT.squeeze(dim=0) @ HMAT, + HMAT, + HMAT @ SDAGMAT, IMAT, ] -def identity(n_qubits: int) -> Tensor: - return torch.eye(2**n_qubits, dtype=torch.complex128) - - def _max_observable_weight(observable: AbstractBlock) -> int: """ Get the maximal weight for the given observable. @@ -88,27 +85,40 @@ def number_of_samples( return N, K -def local_shadow(sample: Counter, unitary_ids: list) -> Tensor: +def nested_operator_indexing( + idx_array: np.ndarray, +) -> list: + """Obtain the list of rotation operators from indices. + + Args: + idx_array (np.ndarray): Indices for obtaining the operators. + + Returns: + list: Map of rotations. """ - Compute local shadow by inverting the quantum channel for each projector state. + if idx_array.ndim == 1: + return [pauli_rotations[int(ind_pauli)](i) for i, ind_pauli in enumerate(idx_array)] # type: ignore[abstract] + return [nested_operator_indexing(sub_array) for sub_array in idx_array] + + +def kron_if_non_empty(list_operations: list) -> KronBlock | None: + filtered_op: list = list(filter(None, list_operations)) + return kron(*filtered_op) if len(filtered_op) > 0 else None - See https://arxiv.org/pdf/2002.08953.pdf - Supplementary Material 1 and Eqs. (S17,S44). - Expects a sample bitstring in ILO. +def extract_operators(unitary_ids: np.ndarray, n_qubits: int) -> list: + """Sample `shadow_size` rotations of `n_qubits`. + + Args: + unitary_ids (np.ndarray): Indices for obtaining the operators. + n_qubits (int): Number of qubits + Returns: + list: Pauli strings. """ - bitstring = list(sample.keys())[0] - local_density_matrices = [] - for bit, unitary_id in zip(bitstring, unitary_ids): - proj_mat = P0_MATRIX if bit == "0" else P1_MATRIX - unitary_tensor = UNITARY_TENSOR[unitary_id].squeeze(dim=0) - local_density_matrices.append( - 3 * (unitary_tensor.adjoint() @ proj_mat @ unitary_tensor) - identity(1) - ) - if len(local_density_matrices) == 1: - return local_density_matrices[0] - else: - return reduce(torch.kron, local_density_matrices) + operations = nested_operator_indexing(unitary_ids) + if n_qubits > 1: + operations = [kron_if_non_empty(ops) for ops in operations] + return operations def classical_shadow( @@ -119,25 +129,21 @@ def classical_shadow( backend: Backend | DifferentiableBackend = PyQBackend(), noise: NoiseHandler | None = None, endianness: Endianness = Endianness.BIG, -) -> list: - shadow: list = [] - # TODO: Parallelize embarrassingly parallel loop. - for _ in range(shadow_size): - unitary_ids = np.random.randint(0, 3, size=(1, circuit.n_qubits))[0] - random_unitary = [ - pauli_gates[unitary_ids[qubit]](qubit) for qubit in range(circuit.n_qubits) - ] - if len(random_unitary) == 1: - random_unitary_block = random_unitary[0] +) -> tuple[np.ndarray, list[Tensor]]: + unitary_ids = np.random.randint(0, 3, size=(shadow_size, circuit.n_qubits)) + shadow: list = list() + all_rotations = extract_operators(unitary_ids, circuit.n_qubits) + + for i in range(shadow_size): + if all_rotations[i]: + rotated_circuit = QuantumCircuit( + circuit.register, chain(circuit.block, all_rotations[i]) + ) else: - random_unitary_block = kron(*random_unitary) - rotated_circuit = QuantumCircuit( - circuit.n_qubits, - chain(circuit.block, random_unitary_block), - ) + rotated_circuit = circuit # Reverse endianness to get sample bitstrings in ILO. conv_circ = backend.circuit(rotated_circuit) - samples = backend.sample( + batch_samples = backend.sample( circuit=conv_circ, param_values=param_values, n_shots=1, @@ -145,97 +151,61 @@ def classical_shadow( noise=noise, endianness=endianness, ) - batched_shadow = [] - for batch in samples: - batched_shadow.append(local_shadow(sample=batch, unitary_ids=unitary_ids)) - shadow.append(batched_shadow) - - # Reshape the shadow by batches of samples instead of samples of batches. - # FIXME: Improve performance. - return [list(s) for s in zip(*shadow)] - - -def reconstruct_state(shadow: list) -> Tensor: - """Reconstruct the state density matrix for the given shadow.""" - return reduce(torch.add, shadow) / len(shadow) - - -def compute_traces( - qubit_support: tuple, - N: int, - K: int, - shadow: list, - observable: AbstractBlock, - endianness: Endianness = Endianness.BIG, -) -> list: - floor = int(np.floor(N / K)) - traces = [] - # TODO: Parallelize embarrassingly parallel loop. - for k in range(K): - reconstructed_state = reconstruct_state(shadow=shadow[k * floor : (k + 1) * floor]) - # Reshape the observable matrix to fit the density matrix dimensions - # by filling indentites. - # Please note the endianness is also flipped to get results in LE. - # FIXME: Changed below from Little to Big, double-check when Roland is back - # FIXME: Correct these comments. - trace = ( - ( - block_to_tensor( - block=observable, - qubit_support=qubit_support, - endianness=Endianness.BIG, - ).squeeze(dim=0) - @ reconstructed_state - ) - .trace() - .real - ) - traces.append(trace) - return traces + shadow.append(batch_samples) + bitstrings = list() + batchsize = len(batch_samples) + for b in range(batchsize): + bitstrings.append([list(batch[b].keys())[0] for batch in shadow]) + bitstrings_torch = [ + 1 - 2 * torch.stack([torch.tensor([int(b_i) for b_i in sample]) for sample in batch]) + for batch in bitstrings + ] + return unitary_ids, bitstrings_torch def estimators( - qubit_support: tuple, N: int, K: int, - shadow: list, + unitary_shadow_ids: np.ndarray, + shadow_samples: Tensor, observable: AbstractBlock, - endianness: Endianness = Endianness.BIG, ) -> Tensor: """ - Return estimators (traces of observable times mean density matrix). - - for K equally-sized shadow partitions. + Return trace estimators from the samples for K equally-sized shadow partitions. See https://arxiv.org/pdf/2002.08953.pdf Algorithm 1. """ - # If there is no Pauli-Z operator in the observable, - # the sample can't "hit" that measurement. + + obs_qubit_support = observable.qubit_support if isinstance(observable, PrimitiveBlock): - if type(observable) == Z: - traces = compute_traces( - qubit_support=qubit_support, - N=N, - K=K, - shadow=shadow, - observable=observable, - endianness=endianness, - ) - else: - traces = [torch.tensor(0.0)] + if isinstance(observable, I): + return torch.tensor(1.0, dtype=torch.get_default_dtype()) + obs_to_pauli_index = [pauli_gates.index(type(observable))] + elif isinstance(observable, CompositeBlock): - if Z in observable: - traces = compute_traces( - qubit_support=qubit_support, - N=N, - K=K, - shadow=shadow, - observable=observable, - endianness=endianness, - ) + obs_to_pauli_index = [ + pauli_gates.index(type(p)) for p in observable.blocks if not isinstance(p, I) # type: ignore[arg-type] + ] + ind_I = set(get_qubit_indices_for_op((observable, 1.0), I(0))) + obs_qubit_support = tuple([ind for ind in observable.qubit_support if ind not in ind_I]) + + floor = int(np.floor(N / K)) + traces = [] + for k in range(K): + indices_match = np.all( + unitary_shadow_ids[k * floor : (k + 1) * floor, obs_qubit_support] + == obs_to_pauli_index, + axis=1, + ) + if indices_match.sum() > 0: + trace = torch.prod( + shadow_samples[k * floor : (k + 1) * floor][indices_match][:, obs_qubit_support], + axis=-1, + ).sum() / sum(indices_match) + traces.append(trace) else: - traces = [torch.tensor(0.0)] + traces.append(torch.tensor(0.0)) return torch.tensor(traces, dtype=torch.get_default_dtype()) @@ -258,7 +228,7 @@ def estimations( N, K = number_of_samples(observables=observables, accuracy=accuracy, confidence=confidence) if shadow_size is not None: N = shadow_size - shadow = classical_shadow( + unitaries_ids, batch_shadow_samples = classical_shadow( shadow_size=N, circuit=circuit, param_values=param_values, @@ -271,18 +241,17 @@ def estimations( for observable in observables: pauli_decomposition = unroll_block_with_scaling(observable) batch_estimations = [] - for batch in shadow: + for batch in batch_shadow_samples: pauli_term_estimations = [] for pauli_term in pauli_decomposition: # Get the estimators for the current Pauli term. # This is a tensor of size K. estimation = estimators( - qubit_support=circuit.block.qubit_support, N=N, K=K, - shadow=batch, + unitary_shadow_ids=unitaries_ids, + shadow_samples=batch, observable=pauli_term[0], - endianness=endianness, ) # Compute the median of means for the current Pauli term. # Weigh the median by the Pauli term scaling. diff --git a/tests/qadence/test_measurements/test_shadows.py b/tests/qadence/test_measurements/test_shadows.py index d0d8990f..32a26200 100644 --- a/tests/qadence/test_measurements/test_shadows.py +++ b/tests/qadence/test_measurements/test_shadows.py @@ -2,27 +2,20 @@ import json import os -from collections import Counter import pytest import torch -from torch import Tensor from qadence.backends.api import backend_factory from qadence.blocks.abstract import AbstractBlock -from qadence.blocks.block_to_tensor import IMAT from qadence.blocks.utils import add, chain, kron from qadence.circuit import QuantumCircuit from qadence.constructors import ising_hamiltonian, total_magnetization from qadence.execution import expectation from qadence.measurements import Measurements from qadence.measurements.shadow import ( - UNITARY_TENSOR, _max_observable_weight, - classical_shadow, estimations, - estimators, - local_shadow, number_of_samples, ) from qadence.model import QuantumModel @@ -30,7 +23,6 @@ from qadence.parameters import Parameter from qadence.serialization import deserialize from qadence.types import BackendName, DiffMode -from qadence.utils import P0_MATRIX, P1_MATRIX @pytest.mark.parametrize( @@ -61,86 +53,9 @@ def test_number_of_samples( assert K == exp_samples[1] -@pytest.mark.parametrize( - "sample, unitary_ids, exp_shadow", - [ - ( - Counter({"10": 1}), - [0, 2], - torch.kron( - 3 * (UNITARY_TENSOR[0].adjoint() @ P1_MATRIX @ UNITARY_TENSOR[0]) - IMAT, - 3 * (UNITARY_TENSOR[2].adjoint() @ P0_MATRIX @ UNITARY_TENSOR[2]) - IMAT, - ), - ), - ( - Counter({"0111": 1}), - [2, 0, 2, 2], - torch.kron( - torch.kron( - 3 * (UNITARY_TENSOR[2].adjoint() @ P0_MATRIX @ UNITARY_TENSOR[2]) - IMAT, - 3 * (UNITARY_TENSOR[0].adjoint() @ P1_MATRIX @ UNITARY_TENSOR[0]) - IMAT, - ), - torch.kron( - 3 * (UNITARY_TENSOR[2].adjoint() @ P1_MATRIX @ UNITARY_TENSOR[2]) - IMAT, - 3 * (UNITARY_TENSOR[2].adjoint() @ P1_MATRIX @ UNITARY_TENSOR[2]) - IMAT, - ), - ), - ), - ], -) -def test_local_shadow(sample: Counter, unitary_ids: list, exp_shadow: Tensor) -> None: - shadow = local_shadow(sample=sample, unitary_ids=unitary_ids) - assert torch.allclose(shadow, exp_shadow) - - theta = Parameter("theta") -@pytest.mark.skip(reason="Can't fix the seed for deterministic outputs.") -@pytest.mark.parametrize( - "layer, param_values, exp_shadows", - [ - (X(0) @ X(2), {}, []) - # (kron(RX(0, theta), X(1)), {"theta": torch.tensor([0.5, 1.0, 1.5])}, []) - ], -) -def test_classical_shadow(layer: AbstractBlock, param_values: dict, exp_shadows: list) -> None: - circuit = QuantumCircuit(2, layer) - shadows = classical_shadow( - shadow_size=2, - circuit=circuit, - param_values=param_values, - ) - for shadow, exp_shadow in zip(shadows, exp_shadows): - for batch, exp_batch in zip(shadow, exp_shadow): - assert torch.allclose(batch, exp_batch, atol=1.0e-2) - - -@pytest.mark.parametrize( - "N, K, circuit, param_values, observable, exp_traces", - [ - (2, 1, QuantumCircuit(2, kron(X(0), Z(1))), {}, X(1), torch.tensor([0.0])), - ], -) -def test_estimators( - N: int, - K: int, - circuit: QuantumCircuit, - param_values: dict, - observable: AbstractBlock, - exp_traces: Tensor, -) -> None: - shadows = classical_shadow(shadow_size=N, circuit=circuit, param_values=param_values) - estimated_traces = estimators( - qubit_support=circuit.block.qubit_support, - N=N, - K=K, - shadow=shadows[0], - observable=observable, - ) - assert torch.allclose(estimated_traces, exp_traces) - - @pytest.mark.flaky(max_runs=5) @pytest.mark.parametrize( "circuit, observable, values", @@ -209,11 +124,13 @@ def test_estimations_comparison_exact( (QuantumCircuit(2, blocks), values2, DiffMode.GPSR), ], ) +@pytest.mark.parametrize("observable", [Z(0) ^ 2, X(1)]) def test_estimations_comparison_tomo_forward_pass( - circuit: QuantumCircuit, values: dict, diff_mode: DiffMode + circuit: QuantumCircuit, + values: dict, + diff_mode: DiffMode, + observable: AbstractBlock, ) -> None: - observable = Z(0) ^ circuit.n_qubits - pyq_backend = backend_factory(BackendName.PYQTORCH, diff_mode=diff_mode) (conv_circ, conv_obs, embed, params) = pyq_backend.convert(circuit, observable) pyq_exp_exact = pyq_backend.expectation(conv_circ, conv_obs, embed(params, values))