Merge pull request #317 from QunaSys/state_vector_sampler

State vector sampler

packages/core/quri_parts/core/sampling/__init__.py

@@ -8,13 +8,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Collection, Iterable, Mapping, NamedTuple, Sequence, Union
-
+from collections import Counter
+from typing import (
+    Callable,
+    Collection,
+    Iterable,
+    Mapping,
+    NamedTuple,
+    Sequence,
+    TypeVar,
+    Union,
+)
+
+import numpy as np
+import numpy.typing as npt
 from typing_extensions import TypeAlias
 
 from quri_parts.backend import SamplingBackend
 from quri_parts.circuit import NonParametricQuantumCircuit
 from quri_parts.core.operator import CommutablePauliSet, Operator
+from quri_parts.core.state import CircuitQuantumState, QuantumStateVector
+
+#: A type variable represents *any* non-parametric quantum state classes.
+#: This is different from :class:`quri_parts.core.state.QuantumStateT`;
+#: ``QuantumStateT`` represents *either one of* the classes, while ``_StateT`` also
+#: covers *a union of* multiple state classes.
+_StateT = TypeVar("_StateT", bound=Union[CircuitQuantumState, QuantumStateVector])
 
 #: MeasurementCounts represents count statistics of repeated measurements of a quantum
 #: circuit. Keys are observed bit patterns encoded in integers and values are counts
@@ -32,6 +51,38 @@
     [Iterable[tuple[NonParametricQuantumCircuit, int]]], Iterable[MeasurementCounts]
 ]
 
+#: StateSampler representes a function that samples a specific (non-parametric) state by
+#: specified times and returns the count statistics. In the case of an ideal
+#: StateSampler, the return value corresponds to probabilities multiplied by shot count.
+StateSampler: TypeAlias = Callable[[_StateT, int], MeasurementCounts]
+
+
+def sample_from_state_vector(
+    state_vector: npt.NDArray[np.complex128], n_shots: int
+) -> MeasurementCounts:
+    """Perform sampling from a state vector."""
+    n_qubits: float = np.log2(state_vector.shape[0])
+    assert n_qubits.is_integer(), "Length of the state vector must be a power of 2."
+    if not np.isclose(np.linalg.norm(state_vector), 1):
+        raise ValueError("probabilities do not sum to 1")
+    probs = np.abs(state_vector) ** 2
+    rng = np.random.default_rng()
+    counts = rng.multinomial(n_shots, probs)
+    return Counter(dict(((i, count) for i, count in enumerate(counts) if count > 0)))
+
+
+def ideal_sample_from_state_vector(
+    state_vector: npt.NDArray[np.complex128], n_shots: int
+) -> MeasurementCounts:
+    """Perform ideal sampling from a state vector."""
+    n_qubits: float = np.log2(state_vector.shape[0])
+    assert n_qubits.is_integer(), "Length of the state vector must be a power of 2."
+    if not np.isclose(np.linalg.norm(state_vector), 1):
+        raise ValueError("probabilities do not sum to 1")
+
+    probs = np.abs(state_vector) ** 2
+    return {i: prob * n_shots for i, prob in enumerate(probs)}
+
 
 def create_sampler_from_sampling_backend(backend: SamplingBackend) -> Sampler:
     """Create a simple :class:`~Sampler` using a :class:`~SamplingBackend`."""

packages/core/tests/core/sampling/test_sample_state_vector.py

@@ -0,0 +1,76 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import Counter
+
+import numpy as np
+import pytest
+
+from quri_parts.core.sampling import (
+    ideal_sample_from_state_vector,
+    sample_from_state_vector,
+)
+
+
+class TestSampleFromStateVector:
+    def test_sample_from_state_vector(self) -> None:
+        n_qubits = 2
+        for i in range(2**n_qubits):
+            phase = np.random.random()
+            state = np.zeros(2**n_qubits, dtype=np.complex128)
+            state[i] = np.exp(1j * phase)
+            assert sample_from_state_vector(state, 1000) == Counter({i: 1000})
+
+    def test_invalid_input(self) -> None:
+        with pytest.raises(
+            AssertionError, match="Length of the state vector must be a power of 2."
+        ):
+            sample_from_state_vector(
+                np.array([0.5, 0.1, np.sqrt(1 - 0.25 - 0.01)]), 1000
+            )
+
+        with pytest.raises(ValueError, match="probabilities do not sum to 1"):
+            sample_from_state_vector(
+                np.random.random(4) + 1j * np.random.random(4), 1000
+            )
+
+
+class TestIdealSampleFromStateVector:
+    def test_ideal_sample_from_state_vector(self) -> None:
+        n_qubits = 2
+        state_vector = np.array(
+            [
+                0.13106223 + 0.70435299j,
+                0.16605566 - 0.36973591j,
+                0.10202236 + 0.48950168j,
+                0.18068102 - 0.19940998j,
+            ]
+        )
+        sampled_cnt = ideal_sample_from_state_vector(state_vector, 1000)
+        expected_cnt = Counter(
+            {0: 513.29044785, 1: 164.27912771, 2: 250.02045389, 3: 72.40997054}
+        )
+
+        for i in range(2**n_qubits):
+            assert np.isclose(sampled_cnt[i], expected_cnt[i])
+        assert np.isclose(sum(expected_cnt.values()), 1000)
+
+    def test_invalid_input(self) -> None:
+        with pytest.raises(
+            AssertionError, match="Length of the state vector must be a power of 2."
+        ):
+            ideal_sample_from_state_vector(
+                np.array([0.5, 0.1, np.sqrt(1 - 0.25 - 0.01)]), 1000
+            )
+
+        with pytest.raises(ValueError, match="probabilities do not sum to 1"):
+            ideal_sample_from_state_vector(
+                np.random.random(4) + 1j * np.random.random(4), 1000
+            )

packages/qulacs/quri_parts/qulacs/sampler.py

@@ -19,47 +19,32 @@
 from quri_parts.circuit import NonParametricQuantumCircuit
 from quri_parts.circuit.noise import NoiseModel
 from quri_parts.core.sampling import ConcurrentSampler, MeasurementCounts, Sampler
+from quri_parts.core.state import GeneralCircuitQuantumState
 from quri_parts.core.utils.concurrent import execute_concurrently
-from quri_parts.qulacs.circuit.compiled_circuit import _QulacsCircuit
 
-from .circuit import convert_circuit
 from .circuit.noise import convert_circuit_with_noise_model
+from .simulator import (
+    create_qulacs_ideal_vector_state_sampler,
+    create_qulacs_vector_state_sampler,
+)
 
 if TYPE_CHECKING:
     from concurrent.futures import Executor
 
+_state_vector_sampler = create_qulacs_vector_state_sampler()
+_ideal_vector_sampler = create_qulacs_ideal_vector_state_sampler()
+
 
 def _sample(circuit: NonParametricQuantumCircuit, shots: int) -> MeasurementCounts:
-    if isinstance(circuit, _QulacsCircuit):
-        qs_circuit = circuit._qulacs_circuit
-    else:
-        qs_circuit = convert_circuit(circuit)
-    qs_state = qulacs.QuantumState(circuit.qubit_count)
-    qs_circuit.update_quantum_state(qs_state)
-
-    if shots > 2 ** max(qs_state.get_qubit_count(), 10):
-        # Use multinomial distribution for faster sampling
-        probs = np.abs(qs_state.get_vector()) ** 2
-        rng = default_rng()
-        counts = rng.multinomial(shots, probs)
-        return dict(((i, count) for i, count in enumerate(counts) if count > 0))
-    else:
-        return Counter(qs_state.sampling(shots))
+    state = GeneralCircuitQuantumState(circuit.qubit_count, circuit)
+    return _state_vector_sampler(state, shots)
 
 
 def _ideal_sample(
     circuit: NonParametricQuantumCircuit, shots: int
 ) -> MeasurementCounts:
-    if isinstance(circuit, _QulacsCircuit):
-        qs_circuit = circuit._qulacs_circuit
-    else:
-        qs_circuit = convert_circuit(circuit)
-    qs_state = qulacs.QuantumState(circuit.qubit_count)
-    qs_circuit.update_quantum_state(qs_state)
-
-    probs = np.abs(qs_state.get_vector()) ** 2
-
-    return {i: prob * shots for i, prob in enumerate(probs)}
+    state = GeneralCircuitQuantumState(circuit.qubit_count, circuit)
+    return _ideal_vector_sampler(state, shots)
 
 
 def create_qulacs_vector_ideal_sampler() -> Sampler:

packages/qulacs/quri_parts/qulacs/simulator.py

@@ -8,6 +8,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections import Counter
 from typing import Union
 
 import numpy as np
@@ -16,18 +17,20 @@
 from numpy.typing import NDArray
 
 from quri_parts.circuit import NonParametricQuantumCircuit
+from quri_parts.core.sampling import (
+    MeasurementCounts,
+    StateSampler,
+    ideal_sample_from_state_vector,
+    sample_from_state_vector,
+)
 from quri_parts.core.state import CircuitQuantumState, QuantumStateVector
 from quri_parts.qulacs.circuit import convert_circuit
 from quri_parts.qulacs.circuit.compiled_circuit import _QulacsCircuit
 
-from . import cast_to_list
+from . import QulacsStateT, cast_to_list
 
 
-def evaluate_state_to_vector(
-    state: Union[CircuitQuantumState, QuantumStateVector]
-) -> QuantumStateVector:
-    """Convert GeneralCircuitQuantumState or QuantumStateVector to
-    QuantumStateVector that only contains the state vector."""
+def _evaluate_qp_state_to_qulacs_state(state: QulacsStateT) -> ql.QuantumState:
     n_qubits = state.qubit_count
 
     if isinstance(state, QuantumStateVector):
@@ -40,32 +43,49 @@ def evaluate_state_to_vector(
             "the input state should be either a GeneralCircuitQuantumState\
              or a QuantumStateVector"
         )
+    return _get_updated_qulacs_state_from_vector(state.circuit, init_state_vector)
 
-    out_state_vector = run_circuit(state.circuit, init_state_vector)
-
-    quri_vector_state = QuantumStateVector(n_qubits=n_qubits, vector=out_state_vector)
-    return quri_vector_state
 
-
-def run_circuit(
-    circuit: NonParametricQuantumCircuit,
+def _get_updated_qulacs_state_from_vector(
+    circuit: Union[NonParametricQuantumCircuit, _QulacsCircuit],
     init_state: NDArray[cfloat],
-) -> NDArray[cfloat]:
-    """Act a NonParametricQuantumCircuit onto a state vector and returns a new
-    state vector."""
-
+) -> ql.QuantumState:
     if len(init_state) != 2**circuit.qubit_count:
         raise ValueError("Inconsistent qubit length between circuit and state")
 
     qulacs_state = ql.QuantumState(circuit.qubit_count)
     qulacs_state.load(cast_to_list(init_state))
+
     if isinstance(circuit, _QulacsCircuit):
         qulacs_cicuit = circuit._qulacs_circuit
     else:
         qulacs_cicuit = convert_circuit(circuit)
 
     qulacs_cicuit.update_quantum_state(qulacs_state)
 
+    return qulacs_state
+
+
+def evaluate_state_to_vector(state: QulacsStateT) -> QuantumStateVector:
+    """Convert GeneralCircuitQuantumState or QuantumStateVector to
+    QuantumStateVector that only contains the state vector."""
+    out_state_vector = _evaluate_qp_state_to_qulacs_state(state)
+
+    # We need to disable type check due to an error in qulacs type annotation
+    # https://github.com/qulacs/qulacs/issues/537
+    return QuantumStateVector(
+        state.qubit_count, out_state_vector.get_vector()  # type: ignore
+    )
+
+
+def run_circuit(
+    circuit: NonParametricQuantumCircuit,
+    init_state: NDArray[cfloat],
+) -> NDArray[cfloat]:
+    """Act a NonParametricQuantumCircuit onto a state vector and returns a new
+    state vector."""
+
+    qulacs_state = _get_updated_qulacs_state_from_vector(circuit, init_state)
     # We need to disable type check due to an error in qulacs type annotation
     # https://github.com/qulacs/qulacs/issues/537
     new_state_vector: NDArray[cfloat] = qulacs_state.get_vector()  # type: ignore
@@ -96,3 +116,30 @@ def get_marginal_probability(
     qulacs_state.load(cast_to_list(state_vector))
     measured = [measured_values.get(i, 2) for i in range(int(n_qubits))]
     return qulacs_state.get_marginal_probability(measured)
+
+
+def create_qulacs_vector_state_sampler() -> StateSampler[QulacsStateT]:
+    """Creates a state sampler based on Qulacs circuit execution."""
+
+    def state_sampler(state: QulacsStateT, n_shots: int) -> MeasurementCounts:
+        if n_shots > 2 ** max(state.qubit_count, 10):
+            # Use multinomial distribution for faster sampling
+            state_vector = evaluate_state_to_vector(state).vector
+            return sample_from_state_vector(state_vector, n_shots)
+
+        qs_state = _evaluate_qp_state_to_qulacs_state(state)
+        return Counter(qs_state.sampling(n_shots))
+
+    return state_sampler
+
+
+def create_qulacs_ideal_vector_state_sampler() -> StateSampler[QulacsStateT]:
+    """Creates an ideal state sampler based on Qulacs circuit execution."""
+
+    def ideal_state_sampler(
+        state: Union[CircuitQuantumState, QuantumStateVector], n_shots: int
+    ) -> MeasurementCounts:
+        state_vector = evaluate_state_to_vector(state).vector
+        return ideal_sample_from_state_vector(state_vector, n_shots)
+
+    return ideal_state_sampler