act_on_density_matrix_args.py

# Copyright 2021 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, Dict, List, Optional, Tuple, TYPE_CHECKING, Sequence, Type, Union

import numpy as np

from cirq import _compat, protocols, qis, sim
from cirq._compat import proper_repr
from cirq.sim.act_on_args import ActOnArgs, strat_act_on_from_apply_decompose
from cirq.linalg import transformations

if TYPE_CHECKING:
    import cirq


class ActOnDensityMatrixArgs(ActOnArgs):
    """State and context for an operation acting on a density matrix.

    To act on this object, directly edit the `target_tensor` property, which is
    storing the density matrix of the quantum system with one axis per qubit.
    """

    @_compat.deprecated_parameter(
        deadline='v0.15',
        fix='Use classical_data.',
        parameter_desc='log_of_measurement_results and positional arguments',
        match=lambda args, kwargs: 'log_of_measurement_results' in kwargs or len(args) > 5,
    )
    @_compat.deprecated_parameter(
        deadline='v0.15',
        fix='Use cirq.dephase_measurements to transform the circuit before simulating.',
        parameter_desc='ignore_measurement_results',
        match=lambda args, kwargs: 'ignore_measurement_results' in kwargs or len(args) > 7,
    )
    @_compat.deprecated_parameter(
        deadline='v0.15',
        fix='Use initial_state instead and specify all the arguments with keywords.',
        parameter_desc='target_tensor and positional arguments',
        match=lambda args, kwargs: 'target_tensor' in kwargs or len(args) != 1,
    )
    def __init__(
        self,
        target_tensor: Optional[np.ndarray] = None,
        available_buffer: Optional[List[np.ndarray]] = None,
        qid_shape: Optional[Tuple[int, ...]] = None,
        prng: Optional[np.random.RandomState] = None,
        log_of_measurement_results: Optional[Dict[str, List[int]]] = None,
        qubits: Optional[Sequence['cirq.Qid']] = None,
        ignore_measurement_results: bool = False,
        initial_state: Union[np.ndarray, 'cirq.STATE_VECTOR_LIKE'] = 0,
        dtype: Type[np.number] = np.complex64,
        classical_data: Optional['cirq.ClassicalDataStore'] = None,
    ):
        """Inits ActOnDensityMatrixArgs.

        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.
            qubits: Determines the canonical ordering of the qubits. This
                is often used in specifying the initial state, i.e. the
                ordering of the computational basis states.
            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.
            ignore_measurement_results: If True, then the simulation
                will treat measurement as dephasing instead of collapsing
                process. This is only applicable to simulators that can
                model dephasing.
            initial_state: The initial state for the simulation in the
                computational basis.
            dtype: The `numpy.dtype` of the inferred state vector. One of
                `numpy.complex64` or `numpy.complex128`. Only used when
                `target_tenson` is None.
            classical_data: The shared classical data container for this
                simulation.

        Raises:
            ValueError: The dimension of `target_tensor` is not divisible by 2
                and `qid_shape` is not provided.
        """
        if ignore_measurement_results:
            super().__init__(
                prng=prng,
                qubits=qubits,
                log_of_measurement_results=log_of_measurement_results,
                ignore_measurement_results=ignore_measurement_results,
                classical_data=classical_data,
            )
        else:
            super().__init__(
                prng=prng,
                qubits=qubits,
                log_of_measurement_results=log_of_measurement_results,
                classical_data=classical_data,
            )
        if target_tensor is None:
            qubits_qid_shape = protocols.qid_shape(self.qubits)
            initial_matrix = qis.to_valid_density_matrix(
                initial_state, len(qubits_qid_shape), qid_shape=qubits_qid_shape, dtype=dtype
            )
            if np.may_share_memory(initial_matrix, initial_state):
                initial_matrix = initial_matrix.copy()
            target_tensor = initial_matrix.reshape(qubits_qid_shape * 2)
        self.target_tensor = target_tensor

        if available_buffer is None:
            available_buffer = [np.empty_like(target_tensor) for _ in range(3)]
        self.available_buffer = available_buffer

        if qid_shape is None:
            target_shape = target_tensor.shape
            if len(target_shape) % 2 != 0:
                raise ValueError(
                    'The dimension of target_tensor is not divisible by 2.'
                    ' Require explicit qid_shape.'
                )
            qid_shape = target_shape[: len(target_shape) // 2]
        self.qid_shape = qid_shape

    def _act_on_fallback_(
        self,
        action: Union['cirq.Operation', 'cirq.Gate'],
        qubits: Sequence['cirq.Qid'],
        allow_decompose: bool = True,
    ) -> bool:
        strats = [
            _strat_apply_channel_to_state,
        ]
        if allow_decompose:
            strats.append(strat_act_on_from_apply_decompose)  # type: ignore

        # Try each strategy, stopping if one works.
        for strat in strats:
            result = strat(action, self, qubits)
            if result is False:
                break  # coverage: ignore
            if result is True:
                return True
            assert result is NotImplemented, str(result)
        raise TypeError(
            "Can't simulate operations that don't implement "
            "SupportsUnitary, SupportsConsistentApplyUnitary, "
            "SupportsMixture or SupportsKraus or is a measurement: {!r}".format(action)
        )

    def _perform_measurement(self, qubits: Sequence['cirq.Qid']) -> List[int]:
        """Delegates the call to measure the density matrix."""
        bits, _ = sim.measure_density_matrix(
            self.target_tensor,
            self.get_axes(qubits),
            out=self.target_tensor,
            qid_shape=self.qid_shape,
            seed=self.prng,
        )
        return bits

    def _on_copy(self, target: 'cirq.ActOnDensityMatrixArgs', deep_copy_buffers: bool = True):
        target.target_tensor = self.target_tensor.copy()
        if deep_copy_buffers:
            target.available_buffer = [b.copy() for b in self.available_buffer]
        else:
            target.available_buffer = self.available_buffer

    def _on_kronecker_product(
        self, other: 'cirq.ActOnDensityMatrixArgs', target: 'cirq.ActOnDensityMatrixArgs'
    ):
        target_tensor = transformations.density_matrix_kronecker_product(
            self.target_tensor, other.target_tensor
        )
        target.target_tensor = target_tensor
        target.available_buffer = [
            np.empty_like(target_tensor) for _ in range(len(self.available_buffer))
        ]
        target.qid_shape = target_tensor.shape[: int(target_tensor.ndim / 2)]

    def _on_factor(
        self,
        qubits: Sequence['cirq.Qid'],
        extracted: 'cirq.ActOnDensityMatrixArgs',
        remainder: 'cirq.ActOnDensityMatrixArgs',
        validate=True,
        atol=1e-07,
    ):
        axes = self.get_axes(qubits)
        extracted_tensor, remainder_tensor = transformations.factor_density_matrix(
            self.target_tensor, axes, validate=validate, atol=atol
        )
        extracted.target_tensor = extracted_tensor
        extracted.available_buffer = [
            np.empty_like(extracted_tensor) for _ in self.available_buffer
        ]
        extracted.qid_shape = extracted_tensor.shape[: int(extracted_tensor.ndim / 2)]
        remainder.target_tensor = remainder_tensor
        remainder.available_buffer = [
            np.empty_like(remainder_tensor) for _ in self.available_buffer
        ]
        remainder.qid_shape = remainder_tensor.shape[: int(remainder_tensor.ndim / 2)]

    @property
    def allows_factoring(self):
        return True

    def _on_transpose_to_qubit_order(
        self, qubits: Sequence['cirq.Qid'], target: 'cirq.ActOnDensityMatrixArgs'
    ):
        axes = self.get_axes(qubits)
        new_tensor = transformations.transpose_density_matrix_to_axis_order(
            self.target_tensor, axes
        )
        buffer = [np.empty_like(new_tensor) for _ in self.available_buffer]
        target.target_tensor = new_tensor
        target.available_buffer = buffer
        target.qid_shape = new_tensor.shape[: int(new_tensor.ndim / 2)]

    def sample(
        self,
        qubits: Sequence['cirq.Qid'],
        repetitions: int = 1,
        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
    ) -> np.ndarray:
        indices = [self.qubit_map[q] for q in qubits]
        return sim.sample_density_matrix(
            self.target_tensor,
            indices,
            qid_shape=tuple(q.dimension for q in self.qubits),
            repetitions=repetitions,
            seed=seed,
        )

    @property
    def can_represent_mixed_states(self) -> bool:
        return True

    def __repr__(self) -> str:
        return (
            'cirq.ActOnDensityMatrixArgs('
            f'target_tensor={proper_repr(self.target_tensor)},'
            f' available_buffer={proper_repr(self.available_buffer)},'
            f' qid_shape={self.qid_shape!r},'
            f' qubits={self.qubits!r},'
            f' log_of_measurement_results={proper_repr(self.log_of_measurement_results)})'
        )


def _strat_apply_channel_to_state(
    action: Any, args: 'cirq.ActOnDensityMatrixArgs', qubits: Sequence['cirq.Qid']
) -> bool:
    """Apply channel to state."""
    axes = args.get_axes(qubits)
    result = protocols.apply_channel(
        action,
        args=protocols.ApplyChannelArgs(
            target_tensor=args.target_tensor,
            out_buffer=args.available_buffer[0],
            auxiliary_buffer0=args.available_buffer[1],
            auxiliary_buffer1=args.available_buffer[2],
            left_axes=axes,
            right_axes=[e + len(args.qubits) for e in axes],
        ),
        default=None,
    )
    if result is None:
        return NotImplemented
    for i in range(len(args.available_buffer)):
        if result is args.available_buffer[i]:
            args.available_buffer[i] = args.target_tensor
    args.target_tensor = result
    return True