Add schmidt_decomposition function to quantum_info

qiskit/quantum_info/__init__.py

@@ -171,6 +171,7 @@
     mutual_information,
     partial_trace,
     purity,
+    schmidt_decomposition,
     shannon_entropy,
     state_fidelity,
 )

qiskit/quantum_info/states/__init__.py

@@ -15,7 +15,7 @@
 from .statevector import Statevector
 from .stabilizerstate import StabilizerState
 from .densitymatrix import DensityMatrix
-from .utils import partial_trace, shannon_entropy
+from .utils import partial_trace, schmidt_decomposition, shannon_entropy
 from .measures import (
     state_fidelity,
     purity,

qiskit/quantum_info/states/utils.py

@@ -20,6 +20,7 @@
 from qiskit.quantum_info.states.statevector import Statevector
 from qiskit.quantum_info.states.densitymatrix import DensityMatrix
 from qiskit.quantum_info.operators.channel import SuperOp
+from qiskit.quantum_info.operators.predicates import ATOL_DEFAULT
 
 
 def partial_trace(state, qargs):
@@ -118,6 +119,83 @@ def logfn(x):
     return h_val
 
 
+def schmidt_decomposition(state, qargs):
+    r"""Return the Schmidt Decomposition of a pure quantum state.
+
+    For an arbitrary bipartite state::
+
+    .. math::
+         |\psi\rangle_{AB} = \sum_{i,j} c_{ij}
+                             |x_i\rangle_A \otimes |y_j\rangle_B,
+
+    its Schmidt Decomposition is given by the single-index sum over k:
+
+    .. math::
+        |\psi\rangle_{AB} = \sum_{k} \lambda_{k}
+                            |u_k\rangle_A \otimes |v_k\rangle_B
+
+    where :math:`|u_k\rangle_A` and :math:`|v_k\rangle_B` are an
+    orthonormal set of vectors in their respective spaces :math:`A` and :math:`B`,
+    and the Schmidt coefficients :math:`\lambda_k` are positive real values.
+
+    Args:
+        state (Statevector or DensityMatrix): the input state.
+        qargs (list): the list of Input state positions corresponding to subsystem :math:`B`.
+
+    Returns:
+        list: list of tuples ``(s, u, v)``, where ``s`` (float) are the
+              Schmidt coefficients :math:`\lambda_k`, and ``u`` (Statevector),
+              ``v`` (Statevector) are the Schmidt vectors
+              :math:`|u_k\rangle_A`, :math:`|u_k\rangle_B`, respectively.
+
+    Raises:
+        QiskitError: if Input qargs is not a list of positions of the Input state.
+        QiskitError: if Input qargs is not a proper subset of Input state.
+    """
+    state = _format_state(state, validate=False)
+
+    # convert to statevector if state is density matrix. Errors if state is mixed.
+    if isinstance(state, DensityMatrix):
+        state = state.to_statevector()
+
+    # reshape statevector into state tensor
+    dims = state.dims()
+    state_tens = state._data.reshape(dims[::-1])
+    ndim = state_tens.ndim
+
+    # check if qargs are valid
+    if not isinstance(qargs, (list, np.ndarray)):
+        raise QiskitError("Input qargs is not a list of positions of the Input state")
+    qudits = list(range(ndim))
+    qargs = set(qargs)
+    if qargs == set(qudits) or not qargs.issubset(qudits):
+        raise QiskitError("Input qargs is not a proper subset of Input state")
+
+    # define subsystem A and B qargs and dims
+    qargs_a = list(qargs)
+    qargs_b = [i for i in qudits if i not in qargs_a]
+    dims_a = state.dims(qargs_a)
+    dims_b = state.dims(qargs_b)
+    ndim_a = np.prod(dims_a)
+    ndim_b = np.prod(dims_b)
+
+    # permute state for desired qargs order
+    qargs_axes = [list(qudits)[::-1].index(i) for i in qargs_a + qargs_b][::-1]
+    state_tens = state_tens.transpose(qargs_axes)
+
+    # convert state tensor to matrix of prob amplitudes and perform svd.
+    state_mat = state_tens.reshape([ndim_a, ndim_b])
+    u_mat, s_arr, vh_mat = np.linalg.svd(state_mat, full_matrices=False)
+
+    schmidt_components = [
+        (s, Statevector(u, dims=dims_a), Statevector(v, dims=dims_b))
+        for s, u, v in zip(s_arr, u_mat.T, vh_mat)
+        if s > ATOL_DEFAULT
+    ]
+
+    return schmidt_components
+
+
 def _format_state(state, validate=True):
     """Format input state into class object"""
     if isinstance(state, list):

test/python/quantum_info/states/test_utils.py

@@ -17,7 +17,7 @@
 
 from qiskit.test import QiskitTestCase
 from qiskit.quantum_info.states import Statevector, DensityMatrix
-from qiskit.quantum_info.states import partial_trace, shannon_entropy
+from qiskit.quantum_info.states import partial_trace, shannon_entropy, schmidt_decomposition
 
 
 class TestStateUtils(QiskitTestCase):
@@ -55,6 +55,35 @@ def test_shannon_entropy(self):
         # Base 10
         self.assertAlmostEqual(0.533908120973504, shannon_entropy(input_pvec, 10))
 
+    def test_schmidt_decomposition(self):
+        """Test schmidt_decomposition function"""
+
+        # separable 2-level system without subsystem permutation
+        target = Statevector.from_label("10l")
+        schmidt_comps = schmidt_decomposition(Statevector.from_label("10l"), [0])
+        state = Statevector(sum(suv[0] * np.kron(suv[1], suv[2]) for suv in schmidt_comps))
+        self.assertEqual(state, target)
+
+        # separable 2-level system with subsystem permutation
+        target = Statevector.from_label("0l1")
+        schmidt_comps = schmidt_decomposition(Statevector.from_label("l10"), [2, 1])
+        state = Statevector(sum(suv[0] * np.kron(suv[1], suv[2]) for suv in schmidt_comps))
+        self.assertEqual(state, target)
+
+        # entangled 2-level system
+        target = 1 / np.sqrt(2) * (Statevector.from_label("00") + Statevector.from_label("11"))
+        schmidt_comps = schmidt_decomposition(target, [0])
+        state = Statevector(sum(suv[0] * np.kron(suv[1], suv[2]) for suv in schmidt_comps))
+        self.assertEqual(state, target)
+
+        # entangled 3-level system
+        target = Statevector(np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]) * 1 / np.sqrt(3), dims=(3, 3))
+        schmidt_comps = schmidt_decomposition(target, [0])
+        state = Statevector(
+            sum(suv[0] * np.kron(suv[1], suv[2]) for suv in schmidt_comps), dims=(3, 3)
+        )
+        self.assertEqual(state, target)
+
 
 if __name__ == "__main__":
     unittest.main()