Clean up recent changes in sub_state_vector (#8081)

- Verify the maximum eigenvalue is indeed close to 1.

- Enable arbitrary order of `keep_indices` as before.

- Use the same tolerance check as before #8077 with rtol=0.
  This allows to restore previous tests instead of having
  to do them with larger state vector deviation.

- Put back RaiseValueErrorIfNotProvided as a default for
`sub_state_vector`.
  This allows using `None` for the fallback return value.
    
- Add check for negative values in `keep_indices`.

Follow-up to #8077

Author: pavoljuhas

cirq-core/cirq/linalg/transformations.py

@@ -20,12 +20,21 @@
 import functools
 from collections.abc import Sequence
 from types import EllipsisType
-from typing import Any
+from typing import Any, TypeVar
 
 import numpy as np
 
 from cirq import protocols
 
+# This is a special indicator value used by the `sub_state_vector` 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.
+RaiseValueErrorIfNotProvided: np.ndarray = np.array([])
+
+TDefault = TypeVar('TDefault')
+
 
 def reflection_matrix_pow(reflection_matrix: np.ndarray, exponent: float) -> np.ndarray:
     """Raises a matrix with two opposing eigenvalues to a power.
@@ -471,9 +480,9 @@ def sub_state_vector(
     state_vector: np.ndarray,
     keep_indices: list[int],
     *,
-    default: np.ndarray | None = None,
+    default: np.ndarray | TDefault = RaiseValueErrorIfNotProvided,
     atol: float = 1e-6,
-) -> np.ndarray:
+) -> np.ndarray | TDefault:
     r"""Attempts to factor a state vector into two parts and return one of them.
 
     The input `state_vector` must have shape ``(2,) * n`` or ``(2 ** n)`` where
@@ -542,32 +551,30 @@ def sub_state_vector(
         raise ValueError("Input state must be normalized.")
     if len(set(keep_indices)) != len(keep_indices):
         raise ValueError(f"keep_indices were {keep_indices} but must be unique.")
-    if any(ind >= n_qubits for ind in keep_indices):
+    if any(ind < 0 or ind >= n_qubits for ind in keep_indices):
         raise ValueError("keep_indices {} are an invalid subset of the input state vector.")
 
     # The permutation moves the specified qubits to the start of the qubit order.
     keeps = frozenset(keep_indices)
-    remainder = np.array([i for i in range(n_qubits) if i not in keeps], dtype=np.int64)
-    permutation = np.concatenate([keep_indices, remainder])
+    permutation = [*sorted(keep_indices), *(i for i in range(n_qubits) if not i in keeps)]
 
     # Permute qubits and construct the pure-state density matrix.
     raveled = state_vector.reshape([2] * n_qubits)
     raveled = np.transpose(raveled, permutation)
-    num_qubits_out = len(keep_indices)
-    c_psi = raveled.reshape([2**num_qubits_out, -1])
+    c_psi = raveled.reshape(keep_dims, -1)
     rho = c_psi @ c_psi.conj().T
 
     # Return the eigenvector with eigenvalue 1.
     evals, evec = np.linalg.eigh(rho)
-    if np.isclose(evals, 1, atol=atol).any():
-        factor_index = np.argmax(evals)
+    factor_index = np.argmax(evals)
+    if np.isclose(evals[factor_index], 1, atol=atol, rtol=0):
         factored = evec[:, factor_index]
         # Prevent accidental reliance on global phase.
         random_phase = np.exp(2j * np.pi * np.random.random())
         return random_phase * factored.reshape(ret_shape)
 
     # Method did not yield a pure state. Fall back to `default` argument.
-    if default is not None:
+    if default is not RaiseValueErrorIfNotProvided:
         return default
 
     raise EntangledStateError(

cirq-core/cirq/linalg/transformations_test.py

@@ -384,32 +384,31 @@ def test_sub_state_vector() -> None:
     assert cirq.equal_up_to_global_phase(
         cirq.sub_state_vector(reshaped_state, [5, 6, 7, 8], atol=1e-15), c
     )
-
-    # Make an imperfect product state to probe tolerances.
-    rng = np.random.default_rng(0)
-    noise = rng.uniform(-1, 1, size=state.size) + 1j * rng.uniform(-1, 1, state.size)
-    imperfect_state = state + 1e-2 * noise.reshape((2,) * 9)
-    imperfect_state /= np.linalg.norm(imperfect_state)
+    # Output state vector is independent of the order of keep_indices
+    assert cirq.equal_up_to_global_phase(
+        cirq.sub_state_vector(reshaped_state, [8, 5, 7, 6], atol=1e-15), c
+    )
 
     # Reject factoring for very tight tolerance.
     assert (
-        cirq.sub_state_vector(imperfect_state, [0, 1], default=_DEFAULT_ARRAY, atol=1e-16)
-        is _DEFAULT_ARRAY
+        cirq.sub_state_vector(state, [0, 1], default=_DEFAULT_ARRAY, atol=1e-16) is _DEFAULT_ARRAY
     )
     assert (
-        cirq.sub_state_vector(imperfect_state, [2, 3, 4], default=_DEFAULT_ARRAY, atol=1e-16)
+        cirq.sub_state_vector(state, [2, 3, 4], default=_DEFAULT_ARRAY, atol=1e-16)
         is _DEFAULT_ARRAY
     )
     assert (
-        cirq.sub_state_vector(imperfect_state, [5, 6, 7, 8], default=_DEFAULT_ARRAY, atol=1e-16)
+        cirq.sub_state_vector(state, [5, 6, 7, 8], default=_DEFAULT_ARRAY, atol=1e-16)
         is _DEFAULT_ARRAY
     )
 
+    # Ensure None can be passed as the `default` argument
+    assert cirq.sub_state_vector(state, [0, 1], default=None, atol=1e-16) is None
+
     # Permit invalid factoring for loose tolerance.
     for q1 in range(9):
         assert (
-            cirq.sub_state_vector(imperfect_state, [q1], default=_DEFAULT_ARRAY, atol=1)
-            is not _DEFAULT_ARRAY
+            cirq.sub_state_vector(state, [q1], default=_DEFAULT_ARRAY, atol=1) is not _DEFAULT_ARRAY
         )
 
 
@@ -482,6 +481,8 @@ def test_sub_state_vector_invalid_inputs() -> None:
         cirq.sub_state_vector(state, [1, 2, 2], atol=1e-8)
 
     state = np.array([1, 0, 0, 0]).reshape((2, 2))
+    with pytest.raises(ValueError, match='invalid'):
+        cirq.sub_state_vector(state, [-1], atol=1e-8)
     with pytest.raises(ValueError, match='invalid'):
         cirq.sub_state_vector(state, [5], atol=1e-8)
     with pytest.raises(ValueError, match='invalid'):