parameterize random_psd_operator with wishart option

toqito/rand/random_psd_operator.py

@@ -7,6 +7,7 @@ def random_psd_operator(
     dim: int,
     is_real: bool = False,
     seed: int | None = None,
+    distribution: str = "uniform",
 ) -> np.ndarray:
     r"""Generate a random positive semidefinite operator.
 
@@ -77,21 +78,38 @@ def random_psd_operator(
     :param is_real: Boolean denoting whether the returned matrix will have all real entries or not.
                     Default is :code:`False`.
     :param seed: A seed used to instantiate numpy's random number generator.
+    :param distribution: Distribution used to generate the PSD operator.
+                         Options are `"uniform"` (default) and `"wishart"`.
     :return: A :code:`dim` x :code:`dim` random positive semidefinite matrix.
 
     """
     # Generate a random matrix of dimension dim x dim.
     gen = np.random.default_rng(seed=seed)
-    rand_mat = gen.random((dim, dim))
 
-    # If is_real is False, add an imaginary component to the matrix.
-    if not is_real:
-        rand_mat = rand_mat + 1j * gen.random((dim, dim))
+    if not isinstance(dim, int) or dim <= 0:
+        raise ValueError("`dim` must be a positive integer.")
 
-    # Constructing a Hermitian matrix.
-    rand_mat = (rand_mat.conj().T + rand_mat) / 2
-    eigenvals, eigenvecs = np.linalg.eigh(rand_mat)
+    if distribution == "uniform":
+        rand_mat = gen.random((dim, dim))
 
-    Q, R = np.linalg.qr(eigenvecs)
+        if not is_real:
+            rand_mat = rand_mat + 1j * gen.random((dim, dim))
 
-    return Q @ np.diag(np.abs(eigenvals)) @ Q.conj().T
+        rand_mat = (rand_mat.conj().T + rand_mat) / 2
+        eigenvals, eigenvecs = np.linalg.eigh(rand_mat)
+        Q, _ = np.linalg.qr(eigenvecs)
+
+        return Q @ np.diag(np.abs(eigenvals)) @ Q.conj().T
+
+    elif distribution == "wishart":
+        if is_real:
+            X = gen.standard_normal((dim, dim))
+        else:
+            X = gen.standard_normal((dim, dim)) + 1j * gen.standard_normal((dim, dim))
+
+        return X @ X.conj().T
+
+    else:
+        raise ValueError(
+            "Invalid distribution. Supported options are 'uniform' and 'wishart'."
+        )

toqito/rand/tests/test_random_psd_operator.py

@@ -79,3 +79,34 @@ def test_random_psd_operator_with_seed(dim, is_real, seed, expected_mat):
     """Test that random_psd_operator function returns the expected output when seeded."""
     matrix = random_psd_operator(dim, is_real, seed)
     assert_allclose(matrix, expected_mat)
+
+def test_random_psd_operator_wishart():
+    """Test Wishart distribution generates PSD matrix."""
+    mat = random_psd_operator(4, distribution="wishart")
+    assert mat.shape == (4, 4)
+    assert is_positive_semidefinite(mat)
+
+
+def test_random_psd_operator_invalid_distribution():
+    """Test invalid distribution raises ValueError."""
+    with pytest.raises(ValueError):
+        random_psd_operator(4, distribution="invalid")
+
+def test_random_psd_operator_invalid_dim_type():
+    """Test invalid dim type raises ValueError."""
+    with pytest.raises(ValueError):
+        random_psd_operator("4")
+
+
+def test_random_psd_operator_invalid_dim_negative():
+    """Test negative dim raises ValueError."""
+    with pytest.raises(ValueError):
+        random_psd_operator(-2)
+
+
+def test_random_psd_operator_wishart_real_branch():
+    """Test Wishart distribution with real sampling branch."""
+    mat = random_psd_operator(4, is_real=True, distribution="wishart")
+    assert mat.shape == (4, 4)
+    assert is_positive_semidefinite(mat)
+