Adding operator sinkhorn operator.

toqito/channels/__init__.py

@@ -1,9 +1,9 @@
 """A number of widely-studied quantum channels."""
 
+from toqito.channels.reduction import reduction
+from toqito.channels.partial_trace import partial_trace
 from toqito.channels.choi import choi
 from toqito.channels.dephasing import dephasing
 from toqito.channels.depolarizing import depolarizing
-from toqito.channels.partial_trace import partial_trace
 from toqito.channels.partial_transpose import partial_transpose
 from toqito.channels.realignment import realignment
-from toqito.channels.reduction import reduction

toqito/matrices/__init__.py

@@ -1,5 +1,7 @@
 """Matrices."""
 
+from toqito.matrices.operator_sinkhorn import operator_sinkhorn
+from toqito.matrices.standard_basis import standard_basis
 from toqito.matrices.gen_pauli_x import gen_pauli_x
 from toqito.matrices.gen_pauli_z import gen_pauli_z
 from toqito.matrices.cnot import cnot
@@ -9,5 +11,4 @@
 from toqito.matrices.gen_pauli import gen_pauli
 from toqito.matrices.hadamard import hadamard
 from toqito.matrices.pauli import pauli
-from toqito.matrices.standard_basis import standard_basis
 from toqito.matrices.cyclic_permutation_matrix import cyclic_permutation_matrix

toqito/matrices/operator_sinkhorn.py

@@ -0,0 +1,87 @@
+"""Sinkhorn operator."""
+
+import numpy as np
+import scipy.linalg
+
+from toqito.channels import partial_trace
+
+
+def operator_sinkhorn(rho, dim=None, tol=np.sqrt(np.finfo(float).eps)):
+    """Perform the operator Sinkhorn iteration.
+
+    This function implements the Sinkhorn algorithm to find a density matrix that
+    is locally equivalent to a given bipartite or multipartite density matrix
+    `rho`, but has both of its partial traces proportional to the identity.
+    """
+    rho = rho.astype(np.complex128)  # Ensure complex128 type for rho
+
+    dX = len(rho)
+    sdX = round(np.sqrt(dX))
+    tr_rho = np.trace(rho)
+
+    # set optional argument defaults: dim=sqrt(length(rho)), tol=sqrt(eps).
+    if dim is None:
+        dim = [sdX, sdX]
+    num_sys = len(dim)
+
+    # allow the user to enter a single number for dim.
+    if num_sys == 1:
+        dim.append(dX // dim[0])
+        if abs(dim[1] - round(dim[1])) >= 2 * dX * np.finfo(float).eps:
+            raise ValueError("If DIM is a scalar, X must be square and DIM must evenly divide length(X); "
+                             "please provide the DIM array containing the dimensions of the subsystems.")
+        dim[1] = round(dim[1])
+        num_sys = 2
+
+    tr_rho_p = tr_rho ** (1 / (2 * num_sys))
+
+    # Prepare the iteration.
+    Prho = [np.eye(d, dtype=np.complex128) / d for d in dim]  # Force complex128 type for Prho
+    F = [np.eye(d, dtype=np.complex128) * tr_rho_p for d in dim]  # Force complex128 for F
+    ldim = [np.prod(dim[:j]) for j in range(num_sys)]
+    rdim = [np.prod(dim[j+1:]) for j in range(num_sys)]
+
+    # Perform the operator Sinkhorn iteration.
+    it_err = 1
+    max_cond = 0
+
+    while it_err > tol:
+        it_err = 0
+        max_cond = 0
+
+        try:
+            for j in range(num_sys):
+                # Compute the reduced density matrix on the j-th system.
+                Prho_tmp = partial_trace(rho, list(set(range(num_sys)) - {j}), dim)
+                Prho_tmp = (Prho_tmp + Prho_tmp.T) / 2  # for numerical stability
+                it_err += np.linalg.norm(Prho[j] - Prho_tmp)
+                Prho[j] = Prho_tmp.astype(np.complex128)  # Force complex128 for Prho_tmp
+
+                # Apply the filter with explicit complex128 conversion
+                T_inv = np.linalg.inv(Prho[j]).astype(np.complex128)
+                T = scipy.linalg.sqrtm(T_inv) / np.sqrt(dim[j])
+                T = T.astype(np.complex128)
+
+                # Construct the Kronecker product
+                eye_ldim = np.eye(int(ldim[j]), dtype=np.complex128)
+                eye_rdim = np.eye(int(rdim[j]), dtype=np.complex128)
+                Tk = np.kron(eye_ldim, np.kron(T, eye_rdim))
+
+                rho = Tk @ rho @ Tk.T
+                F[j] = T @ F[j]
+
+                max_cond = max(max_cond, np.linalg.cond(F[j]))
+
+        except np.linalg.LinAlgError:
+            it_err = 1
+
+        # Check the condition number to ensure invertibility.
+        if it_err == 1 or max_cond >= 1 / tol:
+            raise ValueError("The operator Sinkhorn iteration does not converge for RHO. "
+                             "This is often the case if RHO is not of full rank.")
+
+    # Stabilize the output for numerical reasons.
+    sigma = (rho + rho.T) / 2
+    sigma = tr_rho * sigma / np.trace(sigma)
+
+    return sigma, F

toqito/matrices/tests/test_operator_sinkhorn.py

@@ -0,0 +1,85 @@
+"""Tests for operator_sinkhorn."""
+
+import numpy as np
+
+from toqito.channels import partial_trace
+from toqito import matrices
+from toqito.rand import random_density_matrix
+
+
+def test_operator_sinkhorn_bipartite_partial_trace():
+    """Test operator Sinkhorn partial trace on a bipartite system."""
+    # Generate a random density matrix for a 3x3 system (9-dimensional).
+    rho = random_density_matrix(9)
+    sigma, F = matrices.operator_sinkhorn(rho)
+
+    # Expected partial trace should be proportional to identity matrix.
+    expected_identity = np.eye(3) * (1 / 3)
+
+    # Partial trace on the first subsystem.
+    pt = partial_trace(sigma, 0, [3, 3])
+    pt_rounded = np.around(pt, decimals=2)
+
+    # Check that partial trace matches the expected identity.
+    np.testing.assert_array_almost_equal(pt_rounded, expected_identity, decimal=2)
+
+
+def test_operator_sinkhorn_tripartite_partial_trace():
+    """Test operator Sinkhorn partial trace on a tripartite system."""
+    # Generate a random density matrix for a 2x2x2 system (8-dimensional).
+    rho = random_density_matrix(8)
+    sigma, _ = matrices.operator_sinkhorn(rho, [2, 2, 2])
+
+    # Expected partial trace should be proportional to identity matrix.
+    expected_identity = np.eye(2) * (1 / 2)
+
+    # Partial trace on the first and third subsystems.
+    pt = partial_trace(sigma, [0, 2], [2, 2, 2])
+    pt_rounded = np.around(pt, decimals=2)
+
+    # Check that partial trace matches the expected identity.
+    np.testing.assert_array_almost_equal(pt_rounded, expected_identity, decimal=2)
+
+
+def test_operator_sinkhorn_singular_matrix():
+    """Test operator Sinkhorn with a singular matrix that triggers LinAlgError."""
+    # Create a valid 4x4 singular matrix (non-invertible).
+    rho = np.array([[1, 0, 0, 0],
+                    [0, 1, 0, 0],
+                    [0, 0, 0, 0],
+                    [0, 0, 0, 0]])  # This matrix is singular
+
+    try:
+        matrices.operator_sinkhorn(rho, dim=[2, 2])
+    except ValueError as e:
+        expected_msg = (
+            "The operator Sinkhorn iteration does not converge for RHO. "
+            "This is often the case if RHO is not of full rank."
+        )
+        assert str(e) == expected_msg
+
+
+def test_operator_sinkhorn_invalid_single_dim():
+    """Test operator Sinkhorn when a single number is passed as `dim` and it is invalid."""
+    rho = random_density_matrix(8)  # 8-dimensional density matrix
+
+    # The dimension `4` does not divide evenly into `8`, so we expect an error
+    try:
+        matrices.operator_sinkhorn(rho, dim=[4])
+    except ValueError as e:
+        expected_msg = (
+            "If DIM is a scalar, X must be square and DIM must evenly divide length(X); "
+            "please provide the DIM array containing the dimensions of the subsystems."
+        )
+        assert str(e) == expected_msg
+
+
+def test_operator_sinkhorn_valid_single_dim():
+    """Test operator Sinkhorn when a single valid number is passed as `dim`."""
+    rho = random_density_matrix(9)  # 9-dimensional density matrix
+
+    # The dimension `3` divides evenly into `9`, so no error should occur
+    sigma, F = matrices.operator_sinkhorn(rho, dim=[3])
+
+    # Check that sigma is a valid density matrix with trace equal to 1
+    np.testing.assert_almost_equal(np.trace(sigma), 1)

toqito/matrix_props/trace_norm.py

@@ -32,11 +32,12 @@ def trace_norm(rho: np.ndarray) -> float:
 
     It can be observed using :code:`toqito` that :math:`||\rho||_1 = 1` as follows.
 
+    >>> import numpy as np
     >>> from toqito.states import bell
     >>> from toqito.matrix_props import trace_norm
     >>> rho = bell(0) @ bell(0).conj().T
-    >>> trace_norm(rho)
-    np.float64(0.9999999999999999)
+    >>> np.around(trace_norm(rho), decimals=2)
+    np.float64(1.0)
 
     References
     ==========

toqito/state_metrics/fidelity.py

@@ -49,8 +49,8 @@ def fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
     ...      [1, 0, 0, 1]]
     ... )
     >>> sigma = rho
-    >>> fidelity(rho, sigma)
-    np.float64(1.0000000000000002)
+    >>> np.around(fidelity(rho, sigma), decimals=2)
+    np.float64(1.0)
 
     References
     ==========

toqito/states/__init__.py

@@ -1,5 +1,6 @@
 """Quantum states."""
 
+from toqito.states.max_entangled import max_entangled
 from toqito.states.basis import basis
 from toqito.states.bb84 import bb84
 from toqito.states.bell import bell
@@ -9,7 +10,6 @@
 from toqito.states.ghz import ghz
 from toqito.states.gisin import gisin
 from toqito.states.horodecki import horodecki
-from toqito.states.max_entangled import max_entangled
 from toqito.states.max_mixed import max_mixed
 from toqito.states.isotropic import isotropic
 from toqito.states.tile import tile