Merge pull request #1681 from Ericgig/test_simdiag

Add tests for simdiags

qutip/simdiag.py

@@ -5,7 +5,43 @@
 from qutip.qobj import Qobj
 
 
-def simdiag(ops, evals=True):
+def _degen(tol, vecs, ops, i=0):
+    """
+    Private function that finds eigen vals and vecs for degenerate matrices..
+    """
+    if len(ops) == i:
+        return vecs
+
+    # New eigenvectors are sometime not orthogonal.
+    for j in range(1, vecs.shape[1]):
+        for k in range(j):
+            dot = vecs[:, j].dot(vecs[:, k].conj())
+            if np.abs(dot) > tol:
+                vecs[:, j] = ((vecs[:, j] - dot * vecs[:, k])
+                              / (1 - np.abs(dot)**2)**0.5)
+
+    subspace = vecs.conj().T @ ops[i].data @ vecs
+    eigvals, eigvecs = la.eig(subspace)
+
+    perm = np.argsort(eigvals)
+    eigvals = eigvals[perm]
+
+    vecs_new = vecs @ eigvecs[:, perm]
+    for k in range(len(eigvals)):
+        vecs_new[:, k] = vecs_new[:, k] / la.norm(vecs_new[:, k])
+
+    k = 0
+    while k < len(eigvals):
+        ttol = max(tol, tol * abs(eigvals[k]))
+        inds, = np.where(abs(eigvals - eigvals[k]) < ttol)
+        if len(inds) > 1:  # if at least 2 eigvals are degenerate
+            vecs_new[:, inds] = _degen(tol, vecs_new[:, inds], ops, i+1)
+        k = inds[-1] + 1
+    return vecs_new
+
+
+def simdiag(ops, evals: bool = True, *,
+            tol: float = 1e-14, safe_mode: bool = True):
     """Simultaneous diagonalization of commuting Hermitian matrices.
 
     Parameters
@@ -14,6 +50,18 @@ def simdiag(ops, evals=True):
         ``list`` or ``array`` of qobjs representing commuting Hermitian
         operators.
 
+    evals : bool [True]
+        Whether to return the eigenvalues for each ops and eigenvectors or just
+        the eigenvectors.
+
+    tol : float [1e-14]
+        Tolerance for detecting degenerate eigenstates.
+
+    safe_mode : bool [True]
+        Whether to check that all ops are Hermitian and commuting. If set to
+        ``False`` and operators are not commuting, the eigenvectors returned
+        will often be eigenvectors of only the first operator.
+
     Returns
     --------
     eigs : tuple
@@ -22,21 +70,16 @@ def simdiag(ops, evals=True):
         operator.
 
     """
-    tol = 1e-14
-    start_flag = 0
-    if not any(ops):
-        raise ValueError('Need at least one input operator.')
-    if not isinstance(ops, (list, np.ndarray)):
-        ops = np.array([ops])
-    num_ops = len(ops)
+    if not ops:
+        raise ValueError("No input matrices.")
+    N = ops[0].shape[0]
+    num_ops = len(ops) if safe_mode else 0
     for jj in range(num_ops):
         A = ops[jj]
         shape = A.shape
         if shape[0] != shape[1]:
             raise TypeError('Matricies must be square.')
-        if start_flag == 0:
-            s = shape[0]
-        if s != shape[0]:
+        if shape[0] != N:
             raise TypeError('All matrices. must be the same shape')
         if not A.isherm:
             raise TypeError('Matricies must be Hermitian')
@@ -45,75 +88,34 @@ def simdiag(ops, evals=True):
             if (A * B - B * A).norm() / (A * B).norm() > tol:
                 raise TypeError('Matricies must commute.')
 
-    A = ops[0]
-    eigvals, eigvecs = la.eig(A.full())
-    zipped = list(zip(-eigvals, range(len(eigvals))))
-    zipped.sort()
-    ds, perm = zip(*zipped)
-    ds = -np.real(np.array(ds))
-    perm = np.array(perm)
-    eigvecs_array = np.array(
-        [np.zeros((A.shape[0], 1), dtype=complex) for k in range(A.shape[0])])
-
-    for kk in range(len(perm)):  # matrix with sorted eigvecs in columns
-        eigvecs_array[kk][:, 0] = eigvecs[:, perm[kk]]
+    eigvals, eigvecs = la.eigh(ops[0].full())
+    perm = np.argsort(eigvals)
+    eigvecs = eigvecs[:, perm]
+    eigvals = eigvals[perm]
+
     k = 0
-    rng = np.arange(len(eigvals))
-    while k < len(ds):
+    while k < N:
         # find degenerate eigenvalues, get indicies of degenerate eigvals
-        inds = np.array(abs(ds - ds[k]) < max(tol, tol * abs(ds[k])))
-        inds = rng[inds]
+        ttol = max(tol, tol * abs(eigvals[k]))
+        inds, = np.where(abs(eigvals - eigvals[k]) < ttol)
         if len(inds) > 1:  # if at least 2 eigvals are degenerate
-            eigvecs_array[inds] = degen(tol, eigvecs_array[inds], ops[1:])
-        k = max(inds) + 1
-    eigvals_out = np.zeros((num_ops, len(ds)), dtype=np.float64)
-    kets_out = np.empty((len(ds),), dtype=object)
-    kets_out[:] = [
-        Qobj(eigvecs_array[j] / la.norm(eigvecs_array[j]),
+            eigvecs[:, inds] = _degen(tol, eigvecs[:, inds], ops, 1)
+        k = inds[-1] + 1
+
+    for k in range(N):
+        eigvecs[:, k] = eigvecs[:, k] / la.norm(eigvecs[:, k])
+
+    kets_out = [
+        Qobj(eigvecs[:, j],
              dims=[ops[0].dims[0], [1]], shape=[ops[0].shape[0], 1])
-        for j in range(len(ds))
+        for j in range(N)
     ]
+    eigvals_out = np.zeros((len(ops), N), dtype=np.float64)
     if not evals:
         return kets_out
     else:
-        for kk in range(num_ops):
-            for j in range(len(ds)):
-                eigvals_out[kk, j] = np.real(np.dot(
-                    eigvecs_array[j].conj().T,
-                    ops[kk].data * eigvecs_array[j]))
+        for kk in range(len(ops)):
+            for j in range(N):
+                eigvals_out[kk, j] = ops[kk].matrix_element(kets_out[j],
+                                                            kets_out[j]).real
         return eigvals_out, kets_out
-
-
-def degen(tol, in_vecs, ops):
-    """
-    Private function that finds eigen vals and vecs for degenerate matrices..
-    """
-    n = len(ops)
-    if n == 0:
-        return in_vecs
-    A = ops[0]
-    vecs = np.column_stack(in_vecs)
-    eigvals, eigvecs = la.eig(np.dot(vecs.conj().T, A.data.dot(vecs)))
-    zipped = list(zip(-eigvals, range(len(eigvals))))
-    zipped.sort()
-    ds, perm = zip(*zipped)
-    ds = -np.real(np.array(ds))
-    perm = np.array(perm)
-    vecsperm = np.zeros(eigvecs.shape, dtype=complex)
-    for kk in range(len(perm)):  # matrix with sorted eigvecs in columns
-        vecsperm[:, kk] = eigvecs[:, perm[kk]]
-    vecs_new = np.dot(vecs, vecsperm)
-    vecs_out = np.array(
-        [np.zeros((A.shape[0], 1), dtype=complex) for k in range(len(ds))])
-    for kk in range(len(perm)):  # matrix with sorted eigvecs in columns
-        vecs_out[kk][:, 0] = vecs_new[:, kk]
-    k = 0
-    rng = np.arange(len(ds))
-    while k < len(ds):
-        inds = np.array(abs(ds - ds[k]) < max(
-            tol, tol * abs(ds[k])))  # get indicies of degenerate eigvals
-        inds = rng[inds]
-        if len(inds) > 1:  # if at least 2 eigvals are degenerate
-            vecs_out[inds] = degen(tol, vecs_out[inds], ops[1:n])
-        k = max(inds) + 1
-    return vecs_out

qutip/tests/test_simdiag.py

@@ -0,0 +1,87 @@
+import pytest
+import numpy as np
+import qutip
+
+
+@pytest.mark.parametrize('num_mat', [1, 2, 3, 5])
+def test_simdiag(num_mat):
+    N = 10
+
+    U = qutip.rand_unitary(N)
+    commuting_matrices = [U * qutip.qdiags(np.random.rand(N), 0) * U.dag()
+                          for _ in range(num_mat)]
+    all_evals, evecs = qutip.simdiag(commuting_matrices)
+
+    for matrix, evals in zip(commuting_matrices, all_evals):
+        for eval, evec in zip(evals, evecs):
+            assert matrix * evec == evec * eval
+
+
+@pytest.mark.parametrize('num_mat', [1, 2, 3, 5])
+def test_simdiag_no_evals(num_mat):
+    N = 10
+
+    U = qutip.rand_unitary(N)
+    commuting_matrices = [U * qutip.qdiags(np.random.rand(N), 0) * U.dag()
+                          for _ in range(num_mat)]
+    evecs = qutip.simdiag(commuting_matrices, evals=False)
+
+    for matrix in commuting_matrices:
+        for evec in evecs:
+            Mvec = matrix * evec
+            eval = Mvec.norm() / evec.norm()
+            assert matrix * evec == evec * eval
+
+
+def test_simdiag_degen():
+    N = 10
+    U = qutip.rand_unitary(N)
+    commuting_matrices = [
+        U * qutip.qdiags([0, 0, 0, 1, 1, 1, 2, 2, 3, 4], 0) * U.dag(),
+        U * qutip.qdiags([0, 0, 0, 1, 2, 2, 2, 2, 2, 2], 0) * U.dag(),
+        U * qutip.qdiags([0, 0, 2, 1, 1, 2, 2, 3, 3, 4], 0) * U.dag(),
+    ]
+    all_evals, evecs = qutip.simdiag(commuting_matrices)
+
+    for matrix, evals in zip(commuting_matrices, all_evals):
+        for eval, evec in zip(evals, evecs):
+            np.testing.assert_allclose(
+                (matrix * evec).full(),
+                (evec * eval).full()
+            )
+
+
+@pytest.mark.repeat(10)
+def test_simdiag_degen_large():
+    N = 20
+    U = qutip.rand_unitary(N)
+    commuting_matrices = [
+        U * qutip.qdiags(np.random.randint(0, 3, N), 0) * U.dag()
+        for _ in range(5)
+    ]
+    all_evals, evecs = qutip.simdiag(commuting_matrices, tol=1e-12)
+
+    for matrix, evals in zip(commuting_matrices, all_evals):
+        for eval, evec in zip(evals, evecs):
+            np.testing.assert_allclose(
+                (matrix * evec).full(),
+                (evec * eval).full()
+            )
+
+
+def test_simdiag_no_input():
+    with pytest.raises(ValueError):
+        qutip.simdiag([])
+
+
+@pytest.mark.parametrize(['ops', 'error'], [
+    pytest.param([qutip.basis(5, 0)], 'square', id="Not square"),
+    pytest.param([qutip.qeye(5), qutip.qeye(3)], 'shape', id="shape mismatch"),
+    pytest.param([qutip.destroy(5)], 'Hermitian', id="Non Hermitian"),
+    pytest.param([qutip.sigmax(), qutip.sigmay()], 'commute',
+                 id="Not commuting"),
+])
+def test_simdiag_errors(ops, error):
+    with pytest.raises(TypeError) as err:
+        qutip.simdiag(ops)
+    assert error in str(err.value)