cirq-core 1.1.0.dev20221220224914__py3-none-any.whl → 1.2.0__py3-none-any.whl

cirq/qis/states_test.py

@@ -475,6 +475,13 @@ def test_dirac_notation_precision():
     assert_dirac_notation_python([sqrt, sqrt], "0.707|0⟩ + 0.707|1⟩", decimals=3)
 
 
+def test_dirac_notation_invalid():
+    with pytest.raises(ValueError, match='state_vector has incorrect size'):
+        _ = cirq.dirac_notation([0.0, 0.0, 1.0])
+    with pytest.raises(ValueError, match='state_vector has incorrect size'):
+        _ = cirq.dirac_notation([1.0, 1.0], qid_shape=(3,))
+
+
 def test_to_valid_state_vector():
     with pytest.raises(ValueError, match='Computational basis state is out of range'):
         cirq.to_valid_state_vector(2, 1)

cirq/sim/clifford/clifford_simulator.py

@@ -11,6 +11,7 @@
 # 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.
+
 """An efficient simulator for Clifford circuits.
 
 Allowed operations include:
@@ -109,7 +110,6 @@ class CliffordSimulator(
         measurements: Dict[str, np.ndarray],
         final_simulator_state: 'cirq.SimulationStateBase[cirq.StabilizerChFormSimulationState]',
     ):
-
         return CliffordTrialResult(
             params=params, measurements=measurements, final_simulator_state=final_simulator_state
         )
@@ -259,8 +259,8 @@ class CliffordState:
     ):
         if not isinstance(op.gate, cirq.MeasurementGate):
             raise TypeError(
-                'apply_measurement only supports cirq.MeasurementGate operations. Found %s instead.'
-                % str(op.gate)
+                f'apply_measurement only supports cirq.MeasurementGate operations. '
+                f'Found {op.gate} instead.'
             )
 
         if collapse_state_vector:

cirq/sim/density_matrix_simulation_state.py

@@ -11,6 +11,7 @@
 # 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.
+
 """Objects and methods for acting efficiently on a density matrix."""
 
 from typing import Any, Callable, List, Optional, Sequence, Tuple, Type, TYPE_CHECKING, Union
@@ -305,7 +306,7 @@ class DensityMatrixSimulationState(SimulationState[_BufferedDensityMatrix]):
         raise TypeError(
             "Can't simulate operations that don't implement "
             "SupportsUnitary, SupportsConsistentApplyUnitary, "
-            "SupportsMixture or SupportsKraus or is a measurement: {!r}".format(action)
+            f"SupportsMixture or SupportsKraus or is a measurement: {action!r}"
         )
 
     def __repr__(self) -> str:

cirq/sim/density_matrix_simulator.py

@@ -303,7 +303,7 @@ class DensityMatrixStepResult(simulator_base.StepResultBase['cirq.DensityMatrixS
         # Dtype doesn't have a good repr, so we work around by invoking __name__.
         return (
             f'cirq.DensityMatrixStepResult(sim_state={self._sim_state!r},'
-            f' dtype=np.{self._dtype.__name__})'
+            f' dtype=np.{np.dtype(self._dtype)!r})'
         )
 
 

cirq/sim/density_matrix_simulator_test.py

@@ -950,10 +950,13 @@ def test_density_matrix_step_result_repr():
                 )
             )
         )
-        == "cirq.DensityMatrixStepResult(sim_state=cirq.DensityMatrixSimulationState("
-        "initial_state=np.array([[(0.5+0j), (0.5+0j)], [(0.5+0j), (0.5+0j)]], dtype=np.complex64), "
+        == "cirq.DensityMatrixStepResult("
+        "sim_state=cirq.DensityMatrixSimulationState("
+        "initial_state=np.array([[(0.5+0j), (0.5+0j)], [(0.5+0j), (0.5+0j)]], "
+        "dtype=np.dtype('complex64')), "
         "qubits=(cirq.LineQubit(0),), "
-        "classical_data=cirq.ClassicalDataDictionaryStore()), dtype=np.complex64)"
+        "classical_data=cirq.ClassicalDataDictionaryStore()), "
+        "dtype=np.dtype('complex64'))"
     )
 
 
@@ -1034,9 +1037,10 @@ def test_density_matrix_trial_result_repr():
     expected_repr = (
         "cirq.DensityMatrixTrialResult("
         "params=cirq.ParamResolver({'s': 1}), "
-        "measurements={'m': np.array([[1]], dtype=np.int32)}, "
+        "measurements={'m': np.array([[1]], dtype=np.dtype('int32'))}, "
         "final_simulator_state=cirq.DensityMatrixSimulationState("
-        "initial_state=np.array([[(0.5+0j), (0.5+0j)], [(0.5+0j), (0.5+0j)]], dtype=np.complex64), "
+        "initial_state=np.array([[(0.5+0j), (0.5+0j)], [(0.5+0j), (0.5+0j)]], "
+        "dtype=np.dtype('complex64')), "
         "qubits=(cirq.LineQubit(0),), "
         "classical_data=cirq.ClassicalDataDictionaryStore()))"
     )

cirq/sim/density_matrix_utils.py

@@ -11,6 +11,7 @@
 # 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.
+
 """Code to handle density matrices."""
 
 from typing import List, Optional, TYPE_CHECKING, Tuple, Sequence
@@ -18,6 +19,7 @@ from typing import List, Optional, TYPE_CHECKING, Tuple, Sequence
 import numpy as np
 
 from cirq import linalg, value
+from cirq.sim import simulation_utils
 
 if TYPE_CHECKING:
     import cirq
@@ -188,33 +190,8 @@ def _probs(
     """Returns the probabilities for a measurement on the given indices."""
     # Only diagonal elements matter.
     all_probs = np.diagonal(np.reshape(density_matrix, (np.prod(qid_shape, dtype=np.int64),) * 2))
-    # Shape into a tensor
-    tensor = np.reshape(all_probs, qid_shape)
-
-    # Calculate the probabilities for measuring the particular results.
-    if len(indices) == len(qid_shape):
-        # We're measuring every qudit, so no need for fancy indexing
-        probs = np.abs(tensor)
-        probs = np.transpose(probs, indices)
-        probs = probs.reshape(-1)
-    else:
-        # Fancy indexing required
-        meas_shape = tuple(qid_shape[i] for i in indices)
-        probs = np.abs(
-            [
-                tensor[
-                    linalg.slice_for_qubits_equal_to(
-                        indices, big_endian_qureg_value=b, qid_shape=qid_shape
-                    )
-                ]
-                for b in range(np.prod(meas_shape, dtype=np.int64))
-            ]
-        )
-        probs = np.sum(probs, axis=tuple(range(1, len(probs.shape))))
 
-    # To deal with rounding issues, ensure that the probabilities sum to 1.
-    probs /= np.sum(probs)
-    return probs
+    return simulation_utils.state_probabilities_by_indices(all_probs.real, indices, qid_shape)
 
 
 def _validate_density_matrix_qid_shape(
@@ -227,10 +204,8 @@ def _validate_density_matrix_qid_shape(
     if len(shape) == 2:
         if np.prod(qid_shape, dtype=np.int64) ** 2 != np.prod(shape, dtype=np.int64):
             raise ValueError(
-                'Matrix size does not match qid shape {!r}. Got matrix with '
-                'shape {!r}. Expected {!r}.'.format(
-                    qid_shape, shape, np.prod(qid_shape, dtype=np.int64)
-                )
+                f'Matrix size does not match qid shape {qid_shape!r}. Got matrix with '
+                f'shape {shape!r}. Expected {np.prod(qid_shape, dtype=np.int64)!r}.'
             )
         return qid_shape
     if len(shape) % 2 != 0:
@@ -257,12 +232,12 @@ def _validate_num_qubits(density_matrix: np.ndarray) -> int:
     if row_size & (row_size - 1):
         raise ValueError(
             'Matrix could not be shaped into a square matrix with dimensions '
-            'that are a power of two. Shape was {}'.format(shape)
+            f'that are a power of two. Shape was {shape}'
         )
     if len(shape) > 2 and not np.allclose(shape, 2):
         raise ValueError(
             'Matrix is a tensor of rank greater than 2, but had dimensions '
-            'that are not powers of two. Shape was {}'.format(shape)
+            f'that are not powers of two. Shape was {shape}'
         )
     return int(row_size).bit_length() - 1
 

cirq/sim/mux.py

@@ -139,7 +139,7 @@ def final_state_vector(
     Returns:
         The state vector resulting from applying the given unitary operations to
         the desired initial state. Specifically, a numpy array containing the
-        the amplitudes in np.kron order, where the order of arguments to kron
+        amplitudes in np.kron order, where the order of arguments to kron
         is determined by the qubit order argument (which defaults to just
         sorting the qubits that are present into an ascending order).
 
@@ -160,7 +160,7 @@ def final_state_vector(
             "because it is not unitary. "
             "Maybe you wanted `cirq.final_density_matrix`?\n"
             "\n"
-            "Program: {!r}".format(circuit_like)
+            f"Program: {circuit_like!r}"
         )
 
     result = sparse_simulator.Simulator(dtype=dtype, seed=seed).simulate(

cirq/sim/simulation_state.py

@@ -23,18 +23,19 @@ from typing import (
     List,
     Optional,
     Sequence,
+    Set,
     TypeVar,
     TYPE_CHECKING,
     Tuple,
 )
+from typing_extensions import Self
 
 import numpy as np
 
-from cirq import protocols, value
-from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
+from cirq import ops, protocols, value
+
 from cirq.sim.simulation_state_base import SimulationStateBase
 
-TSelf = TypeVar('TSelf', bound='SimulationState')
 TState = TypeVar('TState', bound='cirq.QuantumStateRepresentation')
 
 if TYPE_CHECKING:
@@ -146,7 +147,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
             return self._state.sample(self.get_axes(qubits), repetitions, seed)
         raise NotImplementedError()
 
-    def copy(self: TSelf, deep_copy_buffers: bool = True) -> TSelf:
+    def copy(self, deep_copy_buffers: bool = True) -> Self:
         """Creates a copy of the object.
 
         Args:
@@ -162,11 +163,43 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         args._state = self._state.copy(deep_copy_buffers=deep_copy_buffers)
         return args
 
-    def create_merged_state(self: TSelf) -> TSelf:
+    def create_merged_state(self) -> Self:
         """Creates a final merged state."""
         return self
 
-    def kronecker_product(self: TSelf, other: TSelf, *, inplace=False) -> TSelf:
+    def add_qubits(self: Self, qubits: Sequence['cirq.Qid']) -> Self:
+        """Add `qubits` in the `|0>` state to a new state space and take the kron product.
+
+        Args:
+            qubits: Sequence of qubits to be added.
+
+        Returns:
+            NotImplemented: If the subclass does not implement this method.
+            Self: A `cirq.SimulationState` with qubits added or `self` if there are no qubits to
+                add.
+        """
+        if not qubits:
+            return self
+        return NotImplemented
+
+    def remove_qubits(self: Self, qubits: Sequence['cirq.Qid']) -> Self:
+        """Remove `qubits` from the state space.
+
+        The qubits to be removed should be untangled from rest of the system and in the |0> state.
+
+        Args:
+            qubits: Sequence of qubits to be removed.
+
+        Returns:
+            NotImplemented: If the subclass does not implement this method.
+            Self: A `cirq.SimulationState` with qubits removed or `self` if there are no qubits to
+                remove.
+        """
+        if not qubits:
+            return self
+        return NotImplemented
+
+    def kronecker_product(self, other: Self, *, inplace=False) -> Self:
         """Joins two state spaces together."""
         args = self if inplace else copy.copy(self)
         args._state = self._state.kron(other._state)
@@ -174,8 +207,8 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         return args
 
     def factor(
-        self: TSelf, qubits: Sequence['cirq.Qid'], *, validate=True, atol=1e-07, inplace=False
-    ) -> Tuple[TSelf, TSelf]:
+        self, qubits: Sequence['cirq.Qid'], *, validate=True, atol=1e-07, inplace=False
+    ) -> Tuple[Self, Self]:
         """Splits two state spaces after a measurement or reset."""
         extracted = copy.copy(self)
         remainder = self if inplace else copy.copy(self)
@@ -191,9 +224,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         """Subclasses that allow factorization should override this."""
         return self._state.supports_factor if self._state is not None else False
 
-    def transpose_to_qubit_order(
-        self: TSelf, qubits: Sequence['cirq.Qid'], *, inplace=False
-    ) -> TSelf:
+    def transpose_to_qubit_order(self, qubits: Sequence['cirq.Qid'], *, inplace=False) -> Self:
         """Physically reindexes the state by the new basis.
 
         Args:
@@ -276,7 +307,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         args._set_qubits(qubits)
         return args
 
-    def __getitem__(self: TSelf, item: Optional['cirq.Qid']) -> TSelf:
+    def __getitem__(self, item: Optional['cirq.Qid']) -> Self:
         if item not in self.qubit_map:
             raise IndexError(f'{item} not in {self.qubits}')
         return self
@@ -295,14 +326,23 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
 def strat_act_on_from_apply_decompose(
     val: Any, args: 'cirq.SimulationState', qubits: Sequence['cirq.Qid']
 ) -> bool:
-    operations, qubits1, _ = _try_decompose_into_operations_and_qubits(val)
-    assert len(qubits1) == len(qubits)
-    qubit_map = {q: qubits[i] for i, q in enumerate(qubits1)}
-    if operations is None:
+    if isinstance(val, ops.Gate):
+        decomposed = protocols.decompose_once_with_qubits(val, qubits, flatten=False, default=None)
+    else:
+        decomposed = protocols.decompose_once(val, flatten=False, default=None)
+    if decomposed is None:
         return NotImplemented
-    for operation in operations:
-        operation = operation.with_qubits(*[qubit_map[q] for q in operation.qubits])
+    all_ancilla: Set['cirq.Qid'] = set()
+    for operation in ops.flatten_to_ops(decomposed):
+        curr_ancilla = tuple(q for q in operation.qubits if q not in args.qubits)
+        args = args.add_qubits(curr_ancilla)
+        if args is NotImplemented:
+            return NotImplemented
+        all_ancilla.update(curr_ancilla)
         protocols.act_on(operation, args)
+    args = args.remove_qubits(tuple(all_ancilla))
+    if args is NotImplemented:
+        raise TypeError(f"{type(args)} implements add_qubits but not remove_qubits.")
     return True
 
 

cirq/sim/simulation_state_base.py

@@ -27,6 +27,7 @@ from typing import (
     TypeVar,
     Union,
 )
+from typing_extensions import Self
 
 import numpy as np
 
@@ -37,7 +38,6 @@ if TYPE_CHECKING:
     import cirq
 
 
-TSelfTarget = TypeVar('TSelfTarget', bound='SimulationStateBase')
 TSimulationState = TypeVar('TSimulationState', bound='cirq.SimulationState')
 
 
@@ -98,7 +98,7 @@ class SimulationStateBase(Generic[TSimulationState], metaclass=abc.ABCMeta):
         protocols.act_on(op, self)
 
     @abc.abstractmethod
-    def copy(self: TSelfTarget, deep_copy_buffers: bool = True) -> TSelfTarget:
+    def copy(self, deep_copy_buffers: bool = True) -> Self:
         """Creates a copy of the object.
 
         Args:
@@ -124,11 +124,14 @@ class SimulationStateBase(Generic[TSimulationState], metaclass=abc.ABCMeta):
     ) -> np.ndarray:
         """Samples the state value."""
 
+    @abc.abstractmethod
     def __getitem__(self, item: Optional['cirq.Qid']) -> TSimulationState:
         """Gets the item associated with the qubit."""
 
+    @abc.abstractmethod
     def __len__(self) -> int:
         """Gets the number of items in the mapping."""
 
+    @abc.abstractmethod
     def __iter__(self) -> Iterator[Optional['cirq.Qid']]:
         """Iterates the keys of the mapping."""

cirq/sim/simulation_state_test.py

@@ -19,6 +19,7 @@ import pytest
 
 import cirq
 from cirq.sim import simulation_state
+from cirq.testing import PhaseUsingCleanAncilla, PhaseUsingDirtyAncilla
 
 
 class DummyQuantumState(cirq.QuantumStateRepresentation):
@@ -33,14 +34,43 @@ class DummyQuantumState(cirq.QuantumStateRepresentation):
 
 
 class DummySimulationState(cirq.SimulationState):
-    def __init__(self):
-        super().__init__(state=DummyQuantumState(), qubits=cirq.LineQubit.range(2))
+    def __init__(self, qubits=cirq.LineQubit.range(2)):
+        super().__init__(state=DummyQuantumState(), qubits=qubits)
 
     def _act_on_fallback_(
         self, action: Any, qubits: Sequence['cirq.Qid'], allow_decompose: bool = True
     ) -> bool:
         return True
 
+    def add_qubits(self, qubits):
+        ret = super().add_qubits(qubits)
+        return self if NotImplemented else ret
+
+
+class DelegatingAncillaZ(cirq.Gate):
+    def __init__(self, exponent=1, measure_ancilla: bool = False):
+        self._exponent = exponent
+        self._measure_ancilla = measure_ancilla
+
+    def num_qubits(self) -> int:
+        return 1
+
+    def _decompose_(self, qubits):
+        a = cirq.NamedQubit('a')
+        yield cirq.CX(qubits[0], a)
+        yield PhaseUsingCleanAncilla(self._exponent).on(a)
+        yield cirq.CX(qubits[0], a)
+        if self._measure_ancilla:
+            yield cirq.measure(a)
+
+
+class Composite(cirq.Gate):
+    def num_qubits(self) -> int:
+        return 1
+
+    def _decompose_(self, qubits):
+        yield cirq.X(*qubits)
+
 
 def test_measurements():
     args = DummySimulationState()
@@ -49,19 +79,24 @@ def test_measurements():
 
 
 def test_decompose():
-    class Composite(cirq.Gate):
-        def num_qubits(self) -> int:
-            return 1
-
-        def _decompose_(self, qubits):
-            yield cirq.X(*qubits)
-
     args = DummySimulationState()
     assert simulation_state.strat_act_on_from_apply_decompose(
         Composite(), args, [cirq.LineQubit(0)]
     )
 
 
+def test_decompose_for_gate_allocating_qubits_raises():
+    class Composite(cirq.testing.SingleQubitGate):
+        def _decompose_(self, qubits):
+            anc = cirq.NamedQubit("anc")
+            yield cirq.CNOT(*qubits, anc)
+
+    args = DummySimulationState()
+
+    with pytest.raises(TypeError, match="add_qubits but not remove_qubits"):
+        simulation_state.strat_act_on_from_apply_decompose(Composite(), args, [cirq.LineQubit(0)])
+
+
 def test_mapping():
     args = DummySimulationState()
     assert list(iter(args)) == cirq.LineQubit.range(2)
@@ -101,3 +136,80 @@ def test_field_getters():
     args = DummySimulationState()
     assert args.prng is np.random
     assert args.qubit_map == {q: i for i, q in enumerate(cirq.LineQubit.range(2))}
+
+
+@pytest.mark.parametrize('exp', np.linspace(0, 2 * np.pi, 10))
+def test_delegating_gate_unitary(exp):
+    q = cirq.LineQubit(0)
+
+    test_circuit = cirq.Circuit()
+    test_circuit.append(cirq.H(q))
+    test_circuit.append(DelegatingAncillaZ(exp).on(q))
+
+    control_circuit = cirq.Circuit(cirq.H(q))
+    control_circuit.append(cirq.ZPowGate(exponent=exp).on(q))
+
+    assert_test_circuit_for_dm_simulator(test_circuit, control_circuit)
+    assert_test_circuit_for_sv_simulator(test_circuit, control_circuit)
+
+
+@pytest.mark.parametrize('exp', np.linspace(0, 2 * np.pi, 10))
+def test_delegating_gate_channel(exp):
+    q = cirq.LineQubit(0)
+
+    test_circuit = cirq.Circuit()
+    test_circuit.append(cirq.H(q))
+    test_circuit.append(DelegatingAncillaZ(exp, True).on(q))
+
+    control_circuit = cirq.Circuit(cirq.H(q))
+    control_circuit.append(cirq.ZPowGate(exponent=exp).on(q))
+
+    with pytest.raises(TypeError, match="DensityMatrixSimulator doesn't support"):
+        # TODO: This test should pass once we extend support to DensityMatrixSimulator.
+        assert_test_circuit_for_dm_simulator(test_circuit, control_circuit)
+
+
+@pytest.mark.parametrize('num_ancilla', [1, 2, 3])
+def test_phase_using_dirty_ancilla(num_ancilla: int):
+    q = cirq.LineQubit(0)
+    anc = cirq.NamedQubit.range(num_ancilla, prefix='anc')
+
+    u = cirq.MatrixGate(cirq.testing.random_unitary(2 ** (num_ancilla + 1)))
+    test_circuit = cirq.Circuit(
+        u.on(q, *anc), PhaseUsingDirtyAncilla(ancilla_bitsize=num_ancilla).on(q)
+    )
+    control_circuit = cirq.Circuit(u.on(q, *anc), cirq.Z(q))
+    assert_test_circuit_for_dm_simulator(test_circuit, control_circuit)
+    assert_test_circuit_for_sv_simulator(test_circuit, control_circuit)
+
+
+@pytest.mark.parametrize('num_ancilla', [1, 2, 3])
+@pytest.mark.parametrize('theta', np.linspace(0, 2 * np.pi, 10))
+def test_phase_using_clean_ancilla(num_ancilla: int, theta: float):
+    q = cirq.LineQubit(0)
+    u = cirq.MatrixGate(cirq.testing.random_unitary(2))
+    test_circuit = cirq.Circuit(
+        u.on(q), PhaseUsingCleanAncilla(theta=theta, ancilla_bitsize=num_ancilla).on(q)
+    )
+    control_circuit = cirq.Circuit(u.on(q), cirq.ZPowGate(exponent=theta).on(q))
+    assert_test_circuit_for_dm_simulator(test_circuit, control_circuit)
+    assert_test_circuit_for_sv_simulator(test_circuit, control_circuit)
+
+
+def assert_test_circuit_for_dm_simulator(test_circuit, control_circuit) -> None:
+    # Density Matrix Simulator: For unitary gates, this fallbacks to `cirq.apply_channel`
+    # which recursively calls to `cirq.apply_unitary(decompose=True)`.
+    for split_untangled_states in [True, False]:
+        sim = cirq.DensityMatrixSimulator(split_untangled_states=split_untangled_states)
+        control_sim = sim.simulate(control_circuit).final_density_matrix
+        test_sim = sim.simulate(test_circuit).final_density_matrix
+        assert np.allclose(test_sim, control_sim)
+
+
+def assert_test_circuit_for_sv_simulator(test_circuit, control_circuit) -> None:
+    # State Vector Simulator.
+    for split_untangled_states in [True, False]:
+        sim = cirq.Simulator(split_untangled_states=split_untangled_states)
+        control_sim = sim.simulate(control_circuit).final_state_vector
+        test_sim = sim.simulate(test_circuit).final_state_vector
+        assert np.allclose(test_sim, control_sim)

cirq/sim/simulation_utils.py

@@ -0,0 +1,59 @@
+# Copyright 2023 The Cirq Developers
+#
+# 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
+#
+#     https://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 typing import Sequence, Tuple
+
+import numpy as np
+
+from cirq import linalg
+
+
+def state_probabilities_by_indices(
+    state_probability: np.ndarray, indices: Sequence[int], qid_shape: Tuple[int, ...]
+) -> np.ndarray:
+    """Returns the probabilities for a state/measurement on the given indices.
+
+    Args:
+        state_probability: The multi-qubit state probability vector. This is an
+            array of 2 to the power of the number of real numbers, and
+            so state must be of size ``2**integer``.  The `state_probability` can be
+            a vector of size ``2**integer`` or a tensor of shape
+            ``(2, 2, ..., 2)``.
+        indices: Which qubits are measured. The `state_probability` is assumed to be
+            supplied in big endian order. That is the xth index of v, when
+            expressed as a bitstring, has its largest values in the 0th index.
+        qid_shape: The qid shape of the `state_probability`.
+
+    Returns:
+        State probabilities.
+    """
+    probs = state_probability.reshape((-1,))
+    not_measured = [i for i in range(len(qid_shape)) if i not in indices]
+    if linalg.can_numpy_support_shape(qid_shape):
+        # Use numpy transpose if we can since it's more efficient.
+        probs = probs.reshape(qid_shape)
+        probs = np.transpose(probs, list(indices) + not_measured)
+        probs = probs.reshape((-1,))
+    else:
+        # If we can't use numpy due to numpy/numpy#5744, use a slower method.
+        probs = linalg.transpose_flattened_array(probs, qid_shape, list(indices) + not_measured)
+
+    if len(not_measured):
+        # Not all qudits are measured.
+        volume = np.prod([qid_shape[i] for i in indices])
+        # Reshape into a 2D array in which each of the measured states correspond to a row.
+        probs = probs.reshape((volume, -1))
+        probs = np.sum(probs, axis=-1)
+
+    # To deal with rounding issues, ensure that the probabilities sum to 1.
+    return probs / np.sum(probs)

cirq/sim/simulation_utils_test.py

@@ -0,0 +1,32 @@
+# Copyright 2023 The Cirq Developers
+#
+# 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
+#
+#     https://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.
+import pytest
+
+import numpy as np
+
+from cirq.sim import simulation_utils
+from cirq import testing
+
+
+@pytest.mark.parametrize('n,m', [(n, m) for n in range(1, 4) for m in range(1, n + 1)])
+def test_state_probabilities_by_indices(n: int, m: int):
+    np.random.seed(0)
+    state = testing.random_superposition(1 << n)
+    d = (state.conj() * state).real
+    desired_axes = list(np.random.choice(n, m, replace=False))
+    not_wanted = [i for i in range(n) if i not in desired_axes]
+    got = simulation_utils.state_probabilities_by_indices(d, desired_axes, (2,) * n)
+    want = np.transpose(d.reshape((2,) * n), desired_axes + not_wanted)
+    want = np.sum(want.reshape((1 << len(desired_axes), -1)), axis=-1)
+    np.testing.assert_allclose(want, got)

cirq/sim/simulator.py

@@ -452,6 +452,7 @@ class SimulatesExpectationValues(metaclass=value.ABCMetaImplementAnyOneOf):
             ValueError if 'program' has terminal measurement(s) and
             'permit_terminal_measurements' is False.
         """
+        raise NotImplementedError
 
 
 class SimulatesFinalState(
@@ -959,7 +960,7 @@ def check_all_resolved(circuit):
         unresolved = [op for moment in circuit for op in moment if protocols.is_parameterized(op)]
         raise ValueError(
             'Circuit contains ops whose symbols were not specified in '
-            'parameter sweep. Ops: {}'.format(unresolved)
+            f'parameter sweep. Ops: {unresolved}'
         )