Add function to generate random density matrix (#2780)

* add random_density_matrix

* format

* add to docs/api.rst

cirq/testing/__init__.py

@@ -57,6 +57,7 @@
 
 from cirq.testing.lin_alg_utils import (
     assert_allclose_up_to_global_phase,
+    random_density_matrix,
     random_orthogonal,
     random_special_orthogonal,
     random_special_unitary,

cirq/testing/lin_alg_utils.py

@@ -11,7 +11,6 @@
 # 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 testing class with utilities for checking linear algebra."""
 
 from typing import Optional
@@ -44,6 +43,30 @@ def random_superposition(dim: int,
     return state_vector
 
 
+def random_density_matrix(dim: int,
+                          *,
+                          random_state: value.RANDOM_STATE_LIKE = None
+                         ) -> np.ndarray:
+    """Returns a random density matrix distributed with Hilbert-Schmidt measure.
+
+    Args:
+        dim: The width and height of the matrix.
+        random_state: A seed to use for random number generation.
+
+    Returns:
+        The sampled density matrix.
+
+    Reference:
+        'Random Bures mixed states and the distribution of their purity'
+        https://arxiv.org/abs/0909.5094
+    """
+    random_state = value.parse_random_state(random_state)
+
+    mat = random_state.randn(dim, dim) + 1j * random_state.randn(dim, dim)
+    mat = mat @ mat.T.conj()
+    return mat / np.trace(mat)
+
+
 def random_unitary(dim: int, *,
                    random_state: value.RANDOM_STATE_LIKE = None) -> np.ndarray:
     """Returns a random unitary matrix distributed with Haar measure.
@@ -163,11 +186,10 @@ def assert_allclose_up_to_global_phase(
     # pylint: enable=unused-variable
 
     actual, desired = linalg.match_global_phase(actual, desired)
-    np.testing.assert_allclose(
-        actual=actual,
-        desired=desired,
-        rtol=rtol,
-        atol=atol,
-        equal_nan=equal_nan,
-        err_msg=err_msg,
-        verbose=verbose)
+    np.testing.assert_allclose(actual=actual,
+                               desired=desired,
+                               rtol=rtol,
+                               atol=atol,
+                               equal_nan=equal_nan,
+                               err_msg=err_msg,
+                               verbose=verbose)

cirq/testing/lin_alg_utils_test.py

@@ -16,18 +16,19 @@
 import pytest
 
 from cirq.testing import (
-    random_superposition,
-    random_unitary,
+    assert_allclose_up_to_global_phase,
+    random_density_matrix,
     random_orthogonal,
-    random_special_unitary,
     random_special_orthogonal,
-    assert_allclose_up_to_global_phase,
+    random_special_unitary,
+    random_superposition,
+    random_unitary,
 )
-from cirq.linalg import (is_unitary, is_orthogonal,
-                         is_special_unitary, is_special_orthogonal)
+from cirq.linalg import (is_unitary, is_orthogonal, is_special_unitary,
+                         is_special_orthogonal)
 
 
-@pytest.mark.parametrize('dim',  range(1, 10))
+@pytest.mark.parametrize('dim', range(1, 10))
 def test_random_superposition(dim):
     state = random_superposition(dim)
 
@@ -42,13 +43,32 @@ def test_random_superposition_deterministic_given_seed():
     np.testing.assert_equal(state1, state2)
 
 
+@pytest.mark.parametrize('dim', range(1, 10))
+def test_random_density_matrix(dim):
+    state = random_density_matrix(dim)
+
+    assert state.shape == (dim, dim)
+    np.testing.assert_allclose(np.trace(state), 1)
+    np.testing.assert_allclose(state, state.T.conj())
+    eigs, _ = np.linalg.eigh(state)
+    assert np.all(eigs >= 0)
+
+
+def test_random_density_matrix_deterministic_given_seed():
+    state1 = random_density_matrix(10, random_state=1234)
+    state2 = random_density_matrix(10, random_state=1234)
+
+    np.testing.assert_equal(state1, state2)
+
+
 def test_random_unitary():
     u1 = random_unitary(2)
     u2 = random_unitary(2)
     assert is_unitary(u1)
     assert is_unitary(u2)
     assert not np.allclose(u1, u2)
 
+
 def test_random_orthogonal():
     o1 = random_orthogonal(2)
     o2 = random_orthogonal(2)
@@ -96,29 +116,24 @@ def test_random_special_orthogonal_deterministic_given_seed():
 
 
 def test_assert_allclose_up_to_global_phase():
-    assert_allclose_up_to_global_phase(
-        np.array([[1]]),
-        np.array([[1j]]),
-        atol=0)
+    assert_allclose_up_to_global_phase(np.array([[1]]),
+                                       np.array([[1j]]),
+                                       atol=0)
 
     with pytest.raises(AssertionError):
-        assert_allclose_up_to_global_phase(
-            np.array([[1]]),
-            np.array([[2]]),
-            atol=0)
+        assert_allclose_up_to_global_phase(np.array([[1]]),
+                                           np.array([[2]]),
+                                           atol=0)
 
-    assert_allclose_up_to_global_phase(
-        np.array([[1e-8, -1, 1e-8]]),
-        np.array([[1e-8, 1, 1e-8]]),
-        atol=1e-6)
+    assert_allclose_up_to_global_phase(np.array([[1e-8, -1, 1e-8]]),
+                                       np.array([[1e-8, 1, 1e-8]]),
+                                       atol=1e-6)
 
     with pytest.raises(AssertionError):
-        assert_allclose_up_to_global_phase(
-            np.array([[1e-4, -1, 1e-4]]),
-            np.array([[1e-4, 1, 1e-4]]),
-            atol=1e-6)
-
-    assert_allclose_up_to_global_phase(
-        np.array([[1, 2], [3, 4]]),
-        np.array([[-1, -2], [-3, -4]]),
-        atol=0)
+        assert_allclose_up_to_global_phase(np.array([[1e-4, -1, 1e-4]]),
+                                           np.array([[1e-4, 1, 1e-4]]),
+                                           atol=1e-6)
+
+    assert_allclose_up_to_global_phase(np.array([[1, 2], [3, 4]]),
+                                       np.array([[-1, -2], [-3, -4]]),
+                                       atol=0)

docs/api.rst

@@ -572,6 +572,7 @@ operation.
     cirq.testing.highlight_text_differences
     cirq.testing.nonoptimal_toffoli_circuit
     cirq.testing.random_circuit
+    cirq.testing.random_density_matrix
     cirq.testing.random_orthogonal
     cirq.testing.random_special_orthogonal
     cirq.testing.random_special_unitary