Fix kraus channels for fallbacks to super-operators (#7537)

- There were a few issues with the new strategy to fall back to
super-operator calculations using apply_channels.
- First, the super_operator_to_kraus function is generally not precise
or numerically stable enough to support an atol of 1e-10 so loosened
this to 1e-6 and also allowed the ability to specify this as a
parameter.
- Next, some operators define decomposition, so cirq.kraus should try to
decompose and get a unitary before falling back to using apply_channel.

Fixes: #7536

Author: dstrain115

cirq-core/cirq/circuits/moment_test.py

@@ -885,7 +885,7 @@ def test_kraus() -> None:
 
 
 def test_kraus_too_big() -> None:
-    m = cirq.Moment(cirq.IdentityGate(11).on(*cirq.LineQubit.range(11)))
+    m = cirq.Moment(cirq.IdentityGate(11).with_probability(0.5).on(*cirq.LineQubit.range(11)))
     assert not cirq.has_kraus(m)
     assert not m._has_superoperator_()
     assert m._kraus_() is NotImplemented

cirq-core/cirq/protocols/kraus_protocol.py

@@ -26,6 +26,7 @@
 from cirq import protocols, qis
 from cirq._doc import doc_private
 from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
+from cirq.protocols.has_unitary_protocol import has_unitary
 from cirq.protocols.mixture_protocol import has_mixture
 from cirq.protocols.unitary_protocol import unitary
 
@@ -95,9 +96,14 @@ def _has_kraus_(self) -> bool:
         """
 
 
-def _strat_kraus_from_apply_channel(val: Any) -> tuple[np.ndarray, ...] | None:
+def _strat_kraus_from_apply_channel(val: Any, atol: float) -> tuple[np.ndarray, ...] | None:
     """Attempts to compute a value's Kraus operators via its _apply_channel_ method.
-    This is very expensive (O(16^N)), so only do this as a last resort."""
+    This is very expensive (O(16^N)), so only do this as a last resort.
+
+    Args:
+        val: value to calculate kraus channels from.
+        atol: Absolute tolerance for super-operator calculation.
+            Matrices with all entries less than this will be dropped."""
     method = getattr(val, '_apply_channel_', None)
     if method is None:
         return None
@@ -122,12 +128,15 @@ def _strat_kraus_from_apply_channel(val: Any) -> tuple[np.ndarray, ...] | None:
     if superop is None or superop is NotImplemented:
         return None
     n = np.prod(qid_shape) ** 2
-    kraus_ops = qis.superoperator_to_kraus(superop.reshape((n, n)))
+    # Note that super-operator calculations can be numerically unstable
+    # and we want to avoid returning kraus channels with "almost zero"
+    # components
+    kraus_ops = qis.superoperator_to_kraus(superop.reshape((n, n)), atol=atol)
     return tuple(kraus_ops)
 
 
 def kraus(
-    val: Any, default: Any = RaiseTypeErrorIfNotProvided
+    val: Any, default: Any = RaiseTypeErrorIfNotProvided, atol: float = 1e-6
 ) -> tuple[np.ndarray, ...] | TDefault:
     r"""Returns a list of matrices describing the channel for the given value.
 
@@ -149,6 +158,8 @@ def kraus(
         default: Determines the fallback behavior when `val` doesn't have
             a channel. If `default` is not set, a TypeError is raised. If
             default is set to a value, that value is returned.
+        atol: If calculating Kraus channels from channels, use this tolerance
+            for determining whether a super-operator is all zeros.
 
     Returns:
         If `val` has a `_kraus_` method and its result is not NotImplemented,
@@ -187,6 +198,9 @@ def kraus(
     if unitary_result is not NotImplemented and unitary_result is not None:
         return (unitary_result,)
 
+    if has_unitary(val):
+        return (unitary(val),)
+
     channel_result = NotImplemented if channel_getter is None else channel_getter()
     if channel_result is not NotImplemented:
         return tuple(channel_result)  # pragma: no cover
@@ -195,7 +209,7 @@ def kraus(
     # Note: _apply_channel can lead to kraus being called again, so if default
     # is None, this can trigger an infinite loop.
     if default is not None:
-        result = _strat_kraus_from_apply_channel(val)
+        result = _strat_kraus_from_apply_channel(val, atol)
         if result is not None:
             return result
 

cirq-core/cirq/protocols/kraus_protocol_test.py

@@ -167,7 +167,7 @@ def test_has_kraus(cls) -> None:
     assert cirq.has_kraus(cls())
 
 
-@pytest.mark.parametrize('decomposed_cls', [HasKraus, HasMixture, HasUnitary])
+@pytest.mark.parametrize('decomposed_cls', [HasKraus, HasMixture])
 def test_has_kraus_when_decomposed(decomposed_cls) -> None:
     op = HasKrausWhenDecomposed(decomposed_cls).on(cirq.NamedQubit('test'))
     assert cirq.has_kraus(op)
@@ -243,3 +243,11 @@ def _kraus_(self):
     gate_no_kraus = NoKrausReset()
     # Should still match the original superoperator
     np.testing.assert_allclose(cirq.kraus(gate), cirq.kraus(gate_no_kraus), atol=1e-8)
+
+
+def test_kraus_channel_with_has_unitary():
+    """CZSWAP has no unitary dunder method but has_unitary returns True."""
+    op = cirq.CZSWAP.on(cirq.q(1), cirq.q(2))
+    channels = cirq.kraus(op)
+    assert len(channels) == 1
+    np.testing.assert_allclose(channels[0], cirq.unitary(op))