From 20a389e4521d2cca8e3f2a6de4ade95ee74c7814 Mon Sep 17 00:00:00 2001
From: "Kevin J. Sung" <kevjsung@umich.edu>
Date: Sun, 3 Nov 2024 07:24:12 -0500
Subject: [PATCH] sample random states: raise ValueError if dim is zero

---
 python/ffsim/random/random.py | 14 +++++++++++++-
 tests/python/random_test.py   | 29 ++++++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/python/ffsim/random/random.py b/python/ffsim/random/random.py
index 260da8c69..6dc2b811f 100644
--- a/python/ffsim/random/random.py
+++ b/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)
diff --git a/tests/python/random_test.py b/tests/python/random_test.py
index 79e0f30ab..0b57a6fd7 100644
--- a/tests/python/random_test.py
+++ b/tests/python/random_test.py
@@ -137,17 +137,40 @@ 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_raise_errors():
+    """Test errors are raised as expected."""
+    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)