parameterize random_psd_operator with wishart option

toqito/rand/random_psd_operator.py

@@ -7,12 +7,16 @@ def random_psd_operator(
     dim: int,
     is_real: bool = False,
     seed: int | None = None,
+    distribution: str = "uniform",
+    scale: np.ndarray | None = None,
+    num_degrees: int | None = None,
 ) -> np.ndarray:
     r"""Generate a random positive semidefinite operator.
 
     A positive semidefinite operator is a Hermitian operator that has only real and non-negative eigenvalues.
-    This function generates a random positive semidefinite operator by constructing a Hermitian matrix,
-    based on the fact that a Hermitian matrix can have real eigenvalues.
+    This function generates a random PSD operator using one of two sampling strategies: ``"uniform"``
+    constructs a Hermitian matrix via random sampling and eigendecomposition, while ``"wishart"`` samples
+    from the Wishart distribution parameterized by a scale matrix and degrees of freedom.
 
     Examples
     ========
@@ -54,6 +58,7 @@ def random_psd_operator(
     .. jupyter-execute::
 
      from toqito.matrix_props import is_positive_semidefinite
+
      is_positive_semidefinite(real_psd_mat)
 
 
@@ -65,8 +70,6 @@ def random_psd_operator(
 
      seeded = random_psd_operator(2, is_real=True, seed=42)
 
-     seeded
-
 
     References
     ==========
@@ -77,21 +80,41 @@ 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: The sampling strategy to use. Either ``"uniform"`` (default) or ``"wishart"``.
+    :param scale: Scale matrix for the Wishart distribution. Defaults to the identity matrix if not provided.
+                  Only used when ``distribution="wishart"``.
+    :param num_degrees: Degrees of freedom for the Wishart distribution. Defaults to ``dim`` if not provided.
+                        Only used when ``distribution="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 not isinstance(dim, int) or dim < 0:
+        raise ValueError("dim must be a non-negative integer.")
 
-    # 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))
-
-    # Constructing a Hermitian matrix.
-    rand_mat = (rand_mat.conj().T + rand_mat) / 2
-    eigenvals, eigenvecs = np.linalg.eigh(rand_mat)
-
-    Q, R = np.linalg.qr(eigenvecs)
+    gen = np.random.default_rng(seed=seed)
 
-    return Q @ np.diag(np.abs(eigenvals)) @ Q.conj().T
+    if distribution == "uniform":
+        rand_mat = gen.random((dim, dim))
+        if not is_real:
+            rand_mat = rand_mat + 1j * gen.random((dim, dim))
+        rand_mat = (rand_mat.conj().T + rand_mat) / 2
+        eigenvals, eigenvecs = np.linalg.eigh(rand_mat)
+        q_mat, _ = np.linalg.qr(eigenvecs)
+        return q_mat @ np.diag(np.abs(eigenvals)) @ q_mat.conj().T
+
+    if distribution == "wishart":
+        if scale is None:
+            scale = np.eye(dim)
+        if num_degrees is None:
+            num_degrees = dim
+        if is_real:
+            x_mat = gen.multivariate_normal(np.zeros(dim), scale, size=num_degrees).T
+        else:
+            x_mat = (
+                gen.multivariate_normal(np.zeros(dim), scale, size=num_degrees).T
+                + 1j * gen.multivariate_normal(np.zeros(dim), scale, size=num_degrees).T
+            )
+        return x_mat @ x_mat.conj().T
+
+    raise ValueError("Invalid distribution. Supported options are 'uniform' and 'wishart'.")

toqito/rand/tests/test_random_psd_operator.py

@@ -9,27 +9,26 @@
 
 
 @pytest.mark.parametrize(
-    "dim, is_real",
+    "dim, is_real, distribution",
     [
-        # Test with a matrix of dimension 2.
-        (2, True),
-        # Test with a matrix of dimension 4.
-        (4, False),
-        # Test with a matrix of dimension 5.
-        (5, False),
-        # Test with a matrix of dimension 10.
-        (10, True),
+        # Test with a matrix of dimension 2, real, uniform.
+        (2, True, "uniform"),
+        # Test with a matrix of dimension 4, complex, uniform.
+        (4, False, "uniform"),
+        # Test with a matrix of dimension 5, complex, uniform.
+        (5, False, "uniform"),
+        # Test with a matrix of dimension 10, real, uniform.
+        (10, True, "uniform"),
+        # Test with a matrix of dimension 4, complex, wishart.
+        (4, False, "wishart"),
+        # Test with a matrix of dimension 4, real, wishart.
+        (4, True, "wishart"),
     ],
 )
-def test_random_psd_operator(dim, is_real):
+def test_random_psd_operator(dim, is_real, distribution):
     """Test for random_psd_operator function."""
-    # Generate a random positive semidefinite operator.
-    rand_psd_operator = random_psd_operator(dim, is_real)
-
-    # Ensure the matrix has the correct shape.
+    rand_psd_operator = random_psd_operator(dim, is_real, distribution=distribution)
     assert_equal(rand_psd_operator.shape, (dim, dim))
-
-    # Check if the matrix is positive semidefinite.
     assert is_positive_semidefinite(rand_psd_operator)
 
 
@@ -79,3 +78,24 @@ 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)
+
+
+@pytest.mark.parametrize(
+    "dim, distribution",
+    [
+        # Invalid dim type.
+        ("4", "uniform"),
+        # Negative dim.
+        (-2, "uniform"),
+    ],
+)
+def test_random_psd_operator_invalid_dim(dim, distribution):
+    """Test that invalid dim raises ValueError."""
+    with pytest.raises(ValueError):
+        random_psd_operator(dim, distribution=distribution)
+
+
+def test_random_psd_operator_invalid_distribution():
+    """Test that invalid distribution raises ValueError."""
+    with pytest.raises(ValueError):
+        random_psd_operator(4, distribution="invalid")