sample random states: raise ValueError if dim is zero

python/ffsim/random/random.py

@@ -30,7 +30,13 @@ def random_state_vector(dim: int, *, seed=None, dtype=complex) -> np.ndarray:
 
     Returns:
         The sampled state vector.
+
+    Raises:
+        ValueError: Dimension must be at least one.
     """
+    if dim < 1:
+        raise ValueError("Dimension must be at least one.")
+
     rng = np.random.default_rng(seed)
     vec = rng.standard_normal(dim).astype(dtype, copy=False)
     if np.issubdtype(dtype, np.complexfloating):
@@ -42,7 +48,7 @@ def random_state_vector(dim: int, *, seed=None, dtype=complex) -> np.ndarray:
 def random_density_matrix(dim: int, *, seed=None, dtype=complex) -> np.ndarray:
     """Returns a random density matrix distributed with Hilbert-Schmidt measure.
 
-    A density matrix is Hermitian and has eigenvalues between 0 and 1.
+    A density matrix is positive semi-definite and has trace equal to one.
 
     Args:
         dim: The width and height of the matrix.
@@ -52,11 +58,17 @@ def random_density_matrix(dim: int, *, seed=None, dtype=complex) -> np.ndarray:
     Returns:
         The sampled density matrix.
 
+    Raises:
+        ValueError: Dimension must be at least one.
+
     References:
         - `arXiv:0909.5094`_
 
     .. _arXiv:0909.5094: https://arxiv.org/abs/0909.5094
     """
+    if dim < 1:
+        raise ValueError("Dimension must be at least one.")
+
     rng = np.random.default_rng(seed)
 
     mat = rng.standard_normal((dim, dim)).astype(dtype, copy=False)

tests/python/random_test.py

@@ -137,17 +137,39 @@ def test_random_antihermitian_matrix(dim: int):
     assert ffsim.linalg.is_antihermitian(mat)
 
 
-@pytest.mark.parametrize("dim", range(10))
+@pytest.mark.parametrize("dim", range(1, 10))
+def test_random_state_vector(dim: int):
+    """Test random state vector."""
+    vec = ffsim.random.random_state_vector(dim, seed=rng)
+    assert vec.dtype == complex
+    np.testing.assert_allclose(np.linalg.norm(vec), 1)
+
+    vec = ffsim.random.random_state_vector(dim, seed=rng, dtype=float)
+    assert vec.dtype == float
+    np.testing.assert_allclose(np.linalg.norm(vec), 1)
+
+
+@pytest.mark.parametrize("dim", range(1, 10))
 def test_random_density_matrix(dim: int):
     """Test random density matrix."""
     mat = ffsim.random.random_density_matrix(dim, seed=rng)
     assert mat.dtype == complex
     assert ffsim.linalg.is_hermitian(mat)
     eigs, _ = np.linalg.eigh(mat)
-    assert all(0 <= e <= 1 for e in eigs)
+    assert all(eigs >= 0)
+    np.testing.assert_allclose(np.trace(mat), 1)
 
     mat = ffsim.random.random_density_matrix(dim, seed=rng, dtype=float)
     assert mat.dtype == float
     assert ffsim.linalg.is_hermitian(mat)
     eigs, _ = np.linalg.eigh(mat)
-    assert all(0 <= e <= 1 for e in eigs)
+    assert all(eigs >= 0)
+    np.testing.assert_allclose(np.trace(mat), 1)
+
+
+def test_bad_dimension():
+    with pytest.raises(ValueError, match="Dimension"):
+        _ = ffsim.random.random_state_vector(0, seed=rng)
+
+    with pytest.raises(ValueError, match="Dimension"):
+        _ = ffsim.random.random_density_matrix(0, seed=rng)