# 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.
"""A protocol for implementing high performance unitary left-multiplies."""
from typing import Any, Union, TypeVar, Tuple, Iterable
import numpy as np
from typing_extensions import Protocol
from cirq import linalg
from cirq.protocols.unitary import unitary
from cirq.type_workarounds import NotImplementedType
# This is a special indicator value used by the apply_unitary method
# to determine whether or not the caller provided a 'default' argument. It must
# be of type np.ndarray to ensure the method has the correct type signature in
# that case. It is checked for using `is`, so it won't have a false positive if
# the user provides a different np.array([]) value.
RaiseTypeErrorIfNotProvided = np.array([]) # type: np.ndarray
TDefault = TypeVar('TDefault')
class ApplyUnitaryArgs:
"""Arguments for performing an efficient left-multiplication by a unitary.
The receiving object is expected to mutate `target_tensor` so that it
contains the state after multiplication, and then return `target_tensor`.
Alternatively, if workspace is required, the receiving object can overwrite
`available_buffer` with the results and return `available_buffer`. Or, if
the receiving object is attempting to be simple instead of fast, it can
create an entirely new array and return that.
Attributes:
target_tensor: The input tensor that needs to be left-multiplied by
the unitary effect of the receiving object. The tensor will
have the shape (2, 2, 2, ..., 2). It usually corresponds to
a multi-qubit superposition, but it could also be a multi-qubit
unitary transformation or some other concept.
available_buffer: Pre-allocated workspace with the same shape and
dtype as the target tensor.
axes: Which axes the unitary effect is being applied to (e.g. the
qubits that the gate is operating on).
"""
def __init__(self,
target_tensor: np.ndarray,
available_buffer: np.ndarray,
axes: Iterable[int]):
"""
Args:
target_tensor: The input tensor that needs to be left-multiplied by
the unitary effect of the receiving object. The tensor will
have the shape (2, 2, 2, ..., 2). It usually corresponds to
a multi-qubit superposition, but it could also be a multi-qubit
unitary transformation or some other concept.
available_buffer: Pre-allocated workspace with the same shape and
dtype as the target tensor.
axes: Which axes the unitary effect is being applied to (e.g. the
qubits that the gate is operating on).
"""
self.target_tensor = target_tensor
self.available_buffer = available_buffer
self.axes = tuple(axes)
def subspace_index(self, little_endian_bits_int: int
) -> Tuple[Union[slice, int, 'ellipsis'], ...]:
"""An index for the subspace where the target axes equal a value.
Args:
little_endian_bits_int: The desired value of the qubits at the
targeted `axes`, packed into an integer. The least significant
bit of the integer is the desired bit for the first axis, and
so forth in increasing order.
Returns:
A value that can be used to index into `target_tensor` and
`available_buffer`, and manipulate only the part of Hilbert space
corresponding to a given bit assignment.
Example:
If `target_tensor` is a 4 qubit tensor and `axes` is `[1, 3]` and
then this method will return the following when given
`little_endian_bits=0b01`:
`(slice(None), 0, slice(None), 1, Ellipsis)`
Therefore the following two lines would be equivalent:
args.target_tensor[args.subspace_index(0b01)] += 1
args.target_tensor[:, 0, :, 1] += 1
"""
return linalg.slice_for_qubits_equal_to(self.axes,
little_endian_bits_int)
[docs]class SupportsApplyUnitary(Protocol):
"""An object that can be efficiently left-multiplied into tensors."""
def _apply_unitary_(self, args: ApplyUnitaryArgs
) -> Union[np.ndarray, None, NotImplementedType]:
"""Left-multiplies a unitary effect onto a tensor with good performance.
This method is given both the target tensor and workspace of the same
shape and dtype. The method then either performs inline modifications of
the target tensor and returns it, or writes its output into the
workspace tensor and returns that. This signature makes it possible to
write specialized simulation methods that run without performing large
allocations, significantly increasing simulation performance.
The target may represent a wave function, a unitary matrix, or some
other tensor. Implementations will work in all of these cases as long as
they correctly focus on only operating on the given axes.
Args:
args: A `cirq.ApplyUnitaryArgs` object with the `args.target_tensor`
to operate on, an `args.available_workspace` buffer to use as
temporary workspace, and the `args.axes` of the tensor to target
with the unitary operation. Note that this method is permitted
(and in fact expected) to mutate `args.target_tensor` and
`args.available_workspace`.
Returns:
If the receiving object is not able to apply its unitary effect,
None or NotImplemented should be returned.
If the receiving object is able to work inline, it should directly
mutate `args.target_tensor` and then return `args.target_tensor`.
The caller will understand this to mean that the result is in
`args.target_tensor`.
If the receiving object is unable to work inline, it can write its
output over `args.available_buffer` and then return
`args.available_buffer`. The caller will understand this to mean
that the result is in `args.available_buffer` (and so what was
`args.available_buffer` will become `args.target_tensor` in the next
call, and vice versa).
The receiving object is also permitted to allocate a new
numpy.ndarray and return that as its result.
"""
[docs]def apply_unitary(unitary_value: Any,
args: ApplyUnitaryArgs,
default: TDefault = RaiseTypeErrorIfNotProvided
) -> Union[np.ndarray, TDefault]:
"""High performance left-multiplication of a unitary effect onto a tensor.
If `unitary_value` defines an `_apply_unitary_` method, that method will be
used to apply `unitary_value`'s unitary effect to the target tensor.
Otherwise, if `unitary_value` defines a `_unitary_` method, its unitary
matrix will be retrieved and applied using a generic method. Otherwise the
application fails, and either an exception is raised or the specified
default value is returned.
Args:
unitary_value: The value with a unitary effect to apply to the target.
args: A mutable `cirq.ApplyUnitaryArgs` object describing the target
tensor, available workspace, and axes to operate on. The attributes
of this object will be mutated as part of computing the result.
default: What should be returned if `unitary_value` doesn't have a
unitary effect. If not specified, a TypeError is raised instead of
returning a default value.
Returns:
If the receiving object is not able to apply its unitary effect,
the specified default value is returned (or a TypeError is raised). If
this occurs, then `target_tensor` should not have been mutated.
If the receiving object was able to work inline, directly
mutating target_tensor it will return target_tensor. The caller is
responsible for checking if the result is target_tensor.
If the receiving object wrote its output over available_buffer, the
result will be available_buffer. The caller is responsible for
checking if the result is available_buffer (and e.g. swapping
the buffer for the target tensor before the next call).
The receiving object may also write its output over a new buffer
that it created, in which case that new array is returned.
Raises:
TypeError: `unitary_value` doesn't have a unitary effect and `default`
wasn't specified.
"""
# Check if the specialized method is present.
func = getattr(unitary_value, '_apply_unitary_', None)
if func is not None:
result = func(args)
if result is not NotImplemented and result is not None:
return result
# Fallback to using the object's _unitary_ matrix.
matrix = unitary(unitary_value, None)
if matrix is not None:
# Special case for single-qubit operations.
if matrix.shape == (2, 2):
zero = args.subspace_index(0)
one = args.subspace_index(1)
return linalg.apply_matrix_to_slices(args.target_tensor,
matrix,
[zero, one],
out=args.available_buffer)
# Fallback to np.einsum for the general case.
return linalg.targeted_left_multiply(
matrix.astype(args.target_tensor.dtype).reshape(
(2,) * (2 * len(args.axes))),
args.target_tensor,
args.axes,
out=args.available_buffer)
# Don't know how to apply. Fallback to specified default behavior.
if default is not RaiseTypeErrorIfNotProvided:
return default
raise TypeError(
"object of type '{}' has no _apply_unitary_ or _unitary_ methods "
"(or they returned None or NotImplemented).".format(
type(unitary_value)))