Density Matrix Simulator ActOn migration

cirq/__init__.py

@@ -350,6 +350,7 @@
 
 from cirq.sim import (
     ActOnCliffordTableauArgs,
+    ActOnDensityMatrixArgs,
     ActOnStabilizerCHFormArgs,
     ActOnStateVectorArgs,
     StabilizerStateChForm,

cirq/ops/measurement_gate.py

@@ -238,6 +238,19 @@ def _act_on_(self, args: Any) -> bool:
 
             return True
 
+        if isinstance(args, sim.ActOnDensityMatrixArgs):
+            invert_mask = self.full_invert_mask()
+            bits, _ = sim.measure_density_matrix(
+                args.target_tensor,
+                args.axes,
+                out=args.target_tensor,
+                qid_shape=args.qid_shape,
+                seed=args.prng,
+            )
+            corrected = [bit ^ (bit < 2 and mask) for bit, mask in zip(bits, invert_mask)]
+            args.record_measurement_result(self.key, corrected)
+            return True
+
         if isinstance(args, sim.clifford.ActOnCliffordTableauArgs):
             invert_mask = self.full_invert_mask()
             bits = [args.tableau._measure(q, args.prng) for q in args.axes]

cirq/protocols/json_test_data/spec.py

@@ -80,6 +80,7 @@
     should_not_be_serialized=[
         # Intermediate states with work buffers and unknown external prng guts.
         'ActOnCliffordTableauArgs',
+        'ActOnDensityMatrixArgs',
         'ActOnStabilizerCHFormArgs',
         'ActOnStateVectorArgs',
         'ApplyChannelArgs',

cirq/sim/__init__.py

@@ -14,6 +14,10 @@
 
 """Base simulation classes and generic simulators."""
 
+from cirq.sim.act_on_density_matrix_args import (
+    ActOnDensityMatrixArgs,
+)
+
 from cirq.sim.act_on_state_vector_args import (
     ActOnStateVectorArgs,
 )

cirq/sim/act_on_density_matrix_args.py

@@ -0,0 +1,134 @@
+# Copyright 2018 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.
+"""Objects and methods for acting efficiently on a density matrix."""
+
+from typing import Any, Iterable, Dict, List, Sequence, Tuple
+
+import numpy as np
+
+from cirq import protocols, ops
+from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
+
+
+class ActOnDensityMatrixArgs:
+    """State and context for an operation acting on a density matrix.
+
+    There are three common ways to act on this object:
+
+    1. Directly edit the `target_tensor` property, which is storing the density
+        matrix of the quantum system as a numpy array with one axis per qudit.
+    2. Overwrite the `available_buffer` property with the new state vector, and
+        then pass `available_buffer` into `swap_target_tensor_for`.
+    3. Call `record_measurement_result(key, val)` to log a measurement result.
+    """
+
+    def __init__(
+        self,
+        target_tensor: np.ndarray,
+        available_buffer: List[np.ndarray],
+        axes: Iterable[int],
+        qid_shape: Tuple[int, ...],
+        prng: np.random.RandomState,
+        log_of_measurement_results: Dict[str, Any],
+    ):
+        """
+        Args:
+            target_tensor: The state vector to act on, stored as a numpy array
+                with one dimension for each qubit in the system. Operations are
+                expected to perform inplace edits of this object.
+            available_buffer: A workspace with the same shape and dtype as
+                `target_tensor`. Used by operations that cannot be applied to
+                `target_tensor` inline, in order to avoid unnecessary
+                allocations. Passing `available_buffer` into
+                `swap_target_tensor_for` will swap it for `target_tensor`.
+            axes: The indices of axes corresponding to the qubits that the
+                operation is supposed to act upon.
+            qid_shape: The shape of the target tensor.
+            prng: The pseudo random number generator to use for probabilistic
+                effects.
+            log_of_measurement_results: A mutable object that measurements are
+                being recorded into. Edit it easily by calling
+                `ActOnStateVectorArgs.record_measurement_result`.
+        """
+        self.target_tensor = target_tensor
+        self.available_buffer = available_buffer
+        self.axes = tuple(axes)
+        self.qid_shape = qid_shape
+        self.prng = prng
+        self.log_of_measurement_results = log_of_measurement_results
+
+    def record_measurement_result(self, key: str, value: Any):
+        """Adds a measurement result to the log.
+
+        Args:
+            key: The key the measurement result should be logged under. Note
+                that operations should only store results under keys they have
+                declared in a `_measurement_keys_` method.
+            value: The value to log for the measurement.
+        """
+        if key in self.log_of_measurement_results:
+            raise ValueError(f"Measurement already logged to key {key!r}")
+        self.log_of_measurement_results[key] = value
+
+    def _act_on_fallback_(self, action: Any, allow_decompose: bool):
+        """Apply channel to state."""
+        result = protocols.apply_channel(
+            action,
+            args=protocols.ApplyChannelArgs(
+                target_tensor=self.target_tensor,
+                out_buffer=self.available_buffer[0],
+                auxiliary_buffer0=self.available_buffer[1],
+                auxiliary_buffer1=self.available_buffer[2],
+                left_axes=self.axes,
+                right_axes=[e + len(self.qid_shape) for e in self.axes],
+            ),
+            default=None,
+        )
+        if result is not None:
+            for i in range(3):
+                if result is self.available_buffer[i]:
+                    self.available_buffer[i] = self.target_tensor
+            self.target_tensor = result
+            return True
+
+        if allow_decompose:
+            return _strat_act_on_density_matrix_from_apply_decompose(action, self)
+
+        return NotImplemented
+
+
+def _strat_act_on_density_matrix_from_apply_decompose(
+    val: Any,
+    args: ActOnDensityMatrixArgs,
+) -> bool:
+    operations, qubits, _ = _try_decompose_into_operations_and_qubits(val)
+    if operations is None:
+        return NotImplemented
+    return _act_all_on_density_matrix(operations, qubits, args)
+
+
+def _act_all_on_density_matrix(
+    actions: Iterable[Any], qubits: Sequence[ops.Qid], args: ActOnDensityMatrixArgs
+):
+    assert len(qubits) == len(args.axes)
+    qubit_map = {q: args.axes[i] for i, q in enumerate(qubits)}
+
+    old_axes = args.axes
+    try:
+        for action in actions:
+            args.axes = tuple(qubit_map[q] for q in action.qubits)
+            protocols.act_on(action, args)
+    finally:
+        args.axes = old_axes
+    return True

cirq/sim/density_matrix_simulator.py

@@ -12,29 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Simulator for density matrices that simulates noisy quantum circuits."""
-
 import collections
-
 from typing import Any, Dict, Iterator, List, TYPE_CHECKING, Tuple, Type, Union
 
 import numpy as np
 
 from cirq import circuits, ops, protocols, qis, study, value, devices
-from cirq.sim import density_matrix_utils, simulator
+from cirq.ops import flatten_to_ops
+from cirq.sim import density_matrix_utils, simulator, act_on_density_matrix_args
 from cirq.sim.simulator import check_all_resolved, split_into_matching_protocol_then_general
 
 if TYPE_CHECKING:
-    from typing import Tuple
     import cirq
 
 
-class _StateAndBuffers:
-    def __init__(self, num_qubits: int, tensor: np.ndarray):
-        self.num_qubits = num_qubits
-        self.tensor = tensor
-        self.buffers = [np.empty_like(tensor) for _ in range(3)]
-
-
 class DensityMatrixSimulator(
     simulator.SimulatesSamples,
     simulator.SimulatesIntermediateState[
@@ -216,33 +207,13 @@ def _run_sweep_repeat(
                     measurements[k].append(np.array(v, dtype=np.uint8))
         return {k: np.array(v) for k, v in measurements.items()}
 
-    def _apply_op_channel(
-        self, op: ops.Operation, state: _StateAndBuffers, indices: List[int]
-    ) -> None:
-        """Apply channel to state."""
-        result = protocols.apply_channel(
-            op,
-            args=protocols.ApplyChannelArgs(
-                target_tensor=state.tensor,
-                out_buffer=state.buffers[0],
-                auxiliary_buffer0=state.buffers[1],
-                auxiliary_buffer1=state.buffers[2],
-                left_axes=indices,
-                right_axes=[e + state.num_qubits for e in indices],
-            ),
-        )
-        for i in range(3):
-            if result is state.buffers[i]:
-                state.buffers[i] = state.tensor
-        state.tensor = result
-
     def _base_iterator(
         self,
         circuit: circuits.Circuit,
         qubit_order: ops.QubitOrderOrList,
         initial_state: Union[np.ndarray, 'cirq.STATE_VECTOR_LIKE'],
         all_measurements_are_terminal=False,
-    ) -> Iterator:
+    ) -> Iterator['DensityMatrixStepResult']:
         qubits = ops.QubitOrder.as_qubit_order(qubit_order).order_for(circuit.all_qubits())
         qid_shape = protocols.qid_shape(qubits)
         qubit_map = {q: i for i, q in enumerate(qubits)}
@@ -251,71 +222,45 @@ def _base_iterator(
         )
         if np.may_share_memory(initial_matrix, initial_state):
             initial_matrix = initial_matrix.copy()
-        measured = collections.defaultdict(bool)  # type: Dict[Tuple[cirq.Qid, ...], bool]
+
         if len(circuit) == 0:
             yield DensityMatrixStepResult(initial_matrix, {}, qubit_map, self._dtype)
             return
 
-        state = _StateAndBuffers(len(qid_shape), initial_matrix.reshape(qid_shape * 2))
-
-        def on_stuck(bad_op: ops.Operation):
-            return TypeError(
-                "Can't simulate operations that don't implement "
-                "SupportsUnitary, SupportsConsistentApplyUnitary, "
-                "SupportsMixture, SupportsChannel or is a measurement: {!r}".format(bad_op)
-            )
-
-        def keep(potential_op: ops.Operation) -> bool:
-            return protocols.has_channel(potential_op, allow_decompose=False) or isinstance(
-                potential_op.gate, ops.MeasurementGate
-            )
+        tensor = initial_matrix.reshape(qid_shape * 2)
+        sim_state = act_on_density_matrix_args.ActOnDensityMatrixArgs(
+            target_tensor=tensor,
+            available_buffer=[np.empty_like(tensor) for _ in range(3)],
+            axes=[],
+            qid_shape=qid_shape,
+            prng=self._prng,
+            log_of_measurement_results={},
+        )
 
         noisy_moments = self.noise.noisy_moments(circuit, sorted(circuit.all_qubits()))
-
+        measured = collections.defaultdict(bool)  # type: Dict[Tuple[cirq.Qid, ...], bool]
         for moment in noisy_moments:
-            measurements = collections.defaultdict(list)  # type: Dict[str, List[int]]
-
-            channel_ops_and_measurements = protocols.decompose(
-                moment, keep=keep, on_stuck_raise=on_stuck
-            )
-
-            for op in channel_ops_and_measurements:
-                indices = [qubit_map[qubit] for qubit in op.qubits]
-                # TODO: support more general measurements.
-                # Github issue: https://github.com/quantumlib/Cirq/issues/1357
+            for op in flatten_to_ops(moment):
                 if all_measurements_are_terminal and measured[op.qubits]:
                     continue
-                if isinstance(op.gate, ops.MeasurementGate):
+                op_list = [op]
+                if protocols.is_measurement(op):
                     measured[op.qubits] = True
-                    meas = op.gate
                     if all_measurements_are_terminal:
                         continue
                     if self._ignore_measurement_results:
-                        for i, q in enumerate(op.qubits):
-                            self._apply_op_channel(ops.phase_damp(1).on(q), state, [indices[i]])
-                    else:
-                        invert_mask = meas.full_invert_mask()
-                        # Measure updates inline.
-                        bits, _ = density_matrix_utils.measure_density_matrix(
-                            state.tensor,
-                            indices,
-                            qid_shape=qid_shape,
-                            out=state.tensor,
-                            seed=self._prng,
-                        )
-                        corrected = [
-                            bit ^ (bit < 2 and mask) for bit, mask in zip(bits, invert_mask)
-                        ]
-                        key = protocols.measurement_key(meas)
-                        measurements[key].extend(corrected)
-                else:
-                    self._apply_op_channel(op, state, indices)
+                        op_list = [ops.phase_damp(1).on(q) for q in op.qubits]
+                for op in op_list:
+                    sim_state.axes = tuple(qubit_map[qubit] for qubit in op.qubits)
+                    protocols.act_on(op, sim_state)
+
             yield DensityMatrixStepResult(
-                density_matrix=state.tensor,
-                measurements=measurements,
+                density_matrix=sim_state.target_tensor,
+                measurements=dict(sim_state.log_of_measurement_results),
                 qubit_map=qubit_map,
                 dtype=self._dtype,
             )
+            sim_state.log_of_measurement_results.clear()
 
     def _create_simulator_trial_result(
         self,

cirq/sim/density_matrix_utils.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Code to handle density matrices."""
 
-from typing import List, Optional, TYPE_CHECKING, Tuple
+from typing import List, Optional, TYPE_CHECKING, Tuple, Sequence
 
 import numpy as np
 
@@ -36,7 +36,7 @@ def von_neumann_entropy(*args, **kwargs):
 
 def sample_density_matrix(
     density_matrix: np.ndarray,
-    indices: List[int],
+    indices: Sequence[int],
     *,  # Force keyword arguments
     qid_shape: Optional[Tuple[int, ...]] = None,
     repetitions: int = 1,
@@ -103,7 +103,7 @@ def sample_density_matrix(
 
 def measure_density_matrix(
     density_matrix: np.ndarray,
-    indices: List[int],
+    indices: Sequence[int],
     qid_shape: Optional[Tuple[int, ...]] = None,
     out: np.ndarray = None,
     seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
@@ -198,7 +198,7 @@ def measure_density_matrix(
 
 
 def _probs(
-    density_matrix: np.ndarray, indices: List[int], qid_shape: Tuple[int, ...]
+    density_matrix: np.ndarray, indices: Sequence[int], qid_shape: Tuple[int, ...]
 ) -> np.ndarray:
     """Returns the probabilities for a measurement on the given indices."""
     # Only diagonal elements matter.
@@ -280,7 +280,7 @@ def _validate_num_qubits(density_matrix: np.ndarray) -> int:
     return int(row_size).bit_length() - 1
 
 
-def _indices_shape(qid_shape: Tuple[int, ...], indices: List[int]) -> Tuple[int, ...]:
+def _indices_shape(qid_shape: Tuple[int, ...], indices: Sequence[int]) -> Tuple[int, ...]:
     """Validates that the indices have values within range of `len(qid_shape)`."""
     if any(index < 0 for index in indices):
         raise IndexError('Negative index in indices: {}'.format(indices))