Merge pull request #2050 from Ericgig/spectrum

Update `spectrum` to use the data layer

qutip/solver/spectrum.py

@@ -4,13 +4,11 @@
 import scipy.fftpack
 
 from .steadystate import steadystate
-from ..core import (
-    qeye, Qobj, liouvillian, spre, unstack_columns, stack_columns,
-    tensor, qzero, expect
-)
+from ..core import liouvillian, spre, expect
+from ..core import data as _data
+from qutip.settings import settings
 
-
-def spectrum(H, wlist, c_ops, a_op, b_op, solver="es", use_pinv=False):
+def spectrum(H, wlist, c_ops, a_op, b_op, solver="es"):
     r"""
     Calculate the spectrum of the correlation function
     :math:`\lim_{t \to \infty} \left<A(t+\tau)B(t)\right>`,
@@ -39,10 +37,7 @@ def spectrum(H, wlist, c_ops, a_op, b_op, solver="es", use_pinv=False):
         Operator B.
     solver : str
         Choice of solver (`es` for exponential series and
-        `pi` for psuedo-inverse).
-    use_pinv : bool
-        For use with the `pi` solver: if `True` use numpy's pinv method,
-        otherwise use a generic solver.
+        `pi` for psuedo-inverse, `solve` for generic solver).
 
     Returns
     -------
@@ -51,12 +46,17 @@ def spectrum(H, wlist, c_ops, a_op, b_op, solver="es", use_pinv=False):
         specified in `wlist`.
 
     """
+    if not H.issuper:
+        L = liouvillian(H, c_ops)
+    else:
+        L = H + sum([lindblad_dissipator(c) for c in c_ops])
     if solver == "es":
-        return _spectrum_es(H, wlist, c_ops, a_op, b_op)
-    elif solver == "pi":
-        return _spectrum_pi(H, wlist, c_ops, a_op, b_op, use_pinv)
-    raise ValueError("Unrecognized choice of solver {} (use es or pi)."
-                     .format(solver))
+        return _spectrum_es(L, wlist, a_op, b_op)
+    elif solver in ["pi", "solve"]:
+        return _spectrum_pi(L, wlist, a_op, b_op, use_pinv=solver=="pi")
+    raise ValueError(
+        f"Unrecognized choice of solver {solver} (use 'es', 'pi' or 'solve')."
+    )
 
 
 def spectrum_correlation_fft(tlist, y, inverse=False):
@@ -96,102 +96,93 @@ def spectrum_correlation_fft(tlist, y, inverse=False):
     return 2 * np.pi * f[idx], 2 * dt * np.real(F[idx])
 
 
-def _spectrum_es(H, wlist, c_ops, a_op, b_op):
+def _spectrum_es(L, wlist, a_op, b_op):
     r"""
     Internal function for calculating the spectrum of the correlation function
     :math:`\left<A(\tau)B(0)\right>`.
     """
-    # construct the Liouvillian
-    L = liouvillian(H, c_ops)
     # find the steady state density matrix and a_op and b_op expecation values
     rho0 = steadystate(L)
     a_op_ss = expect(a_op, rho0)
     b_op_ss = expect(b_op, rho0)
     # eseries solution for (b * rho0)(t)
     states, rates = _diagonal_evolution(L, b_op * rho0)
     # correlation
-    ampls = expect(a_op, states)
+    ampls = [_data.expect(a_op.data, state) for state in states]
     # make covariance
-    ampls = np.concatenate([ampls, [-a_op_ss * b_op_ss]])
-    rates = np.concatenate([rates, [0]])
+    ampls += [-a_op_ss * b_op_ss]
+    rates += [0]
     # Tidy up similar rates.
-    uniques = {}
-    for r, a in zip(rates, ampls):
-        for r_ in uniques:
-            if np.abs(r - r_) < 1e-10:
-                uniques[r_] += a
-                break
+    order = np.argsort(rates)
+    clean_rates = []
+    clean_ampls = []
+    prev_rate = np.nan
+    for idx in order:
+        if np.abs(rates[idx] - prev_rate) < settings.core["atol"]:
+            clean_ampls[-1] += ampls[idx]
         else:
-            uniques[r] = a
-    ampls, rates = [], []
-    for r, a in uniques.items():
-        if np.abs(a) > 1e-10:
-            ampls.append(a)
-            rates.append(r)
+            clean_rates.append(rates[idx])
+            clean_ampls.append(ampls[idx])
+            prev_rate = rates[idx]
+    # Remove 0 amplitude
+    rates, ampls = zip(*[
+        (rate, ampl)
+        for rate, ampl in zip(clean_rates, clean_ampls)
+        if np.abs(ampl) > settings.core["atol"]
+    ])
     ampls, rates = np.array(ampls), np.array(rates)
-    return np.array([2 * np.dot(ampls, 1 / (1j * w - rates)).real
-                     for w in wlist])
+    LW = np.subtract.outer(1j * np.array(wlist), rates).T
+    return (ampls @ (2 / LW)).real
 
 
 #
 # pseudo-inverse solvers
-def _spectrum_pi(H, wlist, c_ops, a_op, b_op, use_pinv=False):
+def _spectrum_pi(L, wlist, a_op, b_op, use_pinv=False):
     r"""
     Internal function for calculating the spectrum of the correlation function
     :math:`\left<A(\tau)B(0)\right>`.
     """
-    L = H if H.issuper else liouvillian(H, c_ops)
-    tr_mat = tensor([qeye(n) for n in L.dims[0][0]])
-    N = np.prod(L.dims[0][0])
-    A = L.full()
-    b = spre(b_op).full()
-    a = spre(a_op).full()
-
-    tr_vec = np.transpose(stack_columns(tr_mat.full()))
-
+    dtype = type(L.data)
     rho_ss = steadystate(L)
-    rho = np.transpose(stack_columns(rho_ss.full()))
+    tr_mat = _data.identity[dtype](rho_ss.shape[0])
+    tr_vec = _data.column_stack(tr_mat).transpose()
+    rho = _data.column_stack(rho_ss.data)
 
-    I = np.identity(N * N)
-    P = np.kron(np.transpose(rho), tr_vec)
+    A = L.data
+    ket = spre(b_op).data @ rho
+    bra = tr_vec @ spre(a_op).data
+
+    I = _data.identity[dtype](L.shape[0])
+    P = _data.kron(rho, tr_vec)
     Q = I - P
 
     spectrum = np.zeros(len(wlist))
     for idx, w in enumerate(wlist):
-        if use_pinv:
-            MMR = np.linalg.pinv(-1.0j * w * I + A)
+        if use_pinv and np.abs(w) > settings.core["atol"]:
+            # At w == 0., "L - iw" is singular
+            MMR = _data.inv(-1.0j * w * I + A)
         else:
-            MMR = np.dot(Q, np.linalg.solve(-1.0j * w * I + A, Q))
+            MMR = Q @ _data.solve(-1.0j * w * I + A, Q)
 
-        s = np.dot(tr_vec,
-                   np.dot(a, np.dot(MMR, np.dot(b, np.transpose(rho)))))
-        spectrum[idx] = -2 * np.real(s[0, 0])
+        spectrum[idx] = -2 * _data.inner_op(bra, MMR, ket).real
     return spectrum
 
 
-def _diagonal_evolution(L, rho0):
-    """
-    Diagonalise the evolution of density matrix rho0 under a constant
-    Liouvillian L.  Returns a list of `states` and an array of the eigenvalues
-    such that the time evolution of rho0 is represented by
-        sum_k states[k] * exp(evals[k] * t)
-    This is effectively the same as the legacy QuTiP function ode2es, but does
-    not use the removed eseries class.  It exists here because ode2es and
-    essolve were removed.
-    """
-    rho0_full = rho0.full()
-    if np.abs(rho0_full).sum() < 1e-10 + 1e-24:
-        return qzero(rho0.dims[0]), np.array([0])
-    evals, evecs = L.eigenstates()
-    evecs = np.vstack([ket.full()[:, 0] for ket in evecs]).T
-    # evals[i]   = eigenvalue i
-    # evecs[:, i] = eigenvector i
+def _diagonal_evolution(L, rho0, sparse=False):
+    if rho0.norm() < settings.core["atol"]:
+        return [_data.zeros["CSR"](*rho0.shape)], [0]
+    if isinstance(L.data, _data.CSR) and not sparse:
+        L = L.to(_data.Dense)
+    evals, evecs = _data.eigs(L.data)
     size = rho0.shape[0] * rho0.shape[1]
-    r0 = stack_columns(rho0_full)[:, 0]
-    v0 = scipy.linalg.solve(evecs, r0)
-    vv = evecs * v0[None, :]  # product equivalent to `evecs @ np.diag(v0)`
-    states = [Qobj(unstack_columns(vv[:, i]), dims=rho0.dims, type='oper')
-              for i in range(size)]
-    # We don't use QobjEvo because it isn't designed to be efficient when
-    # calculating
-    return states, evals
+    r0 = _data.column_stack(rho0.data)
+    v0 = _data.solve(evecs, r0)
+    vv = evecs @ _data.diag(v0.to_array().flatten(), [0])
+    states = []
+    rates = []
+    for ket, rate in zip(_data.split_columns(vv), evals):
+        if _data.norm.l2(ket) < settings.core["atol"]:
+            continue
+        states.append(_data.column_unstack(ket, rho0.shape[0]))
+        rates.append(rate)
+    return states, rates

qutip/tests/solver/test_correlation.py

@@ -64,6 +64,7 @@ def _spectrum_fft(H, c_ops, a, b):
     pytest.param(_spectrum_fft, id="fft"),
     pytest.param(_spectrum_wrapper("es"), id="es"),
     pytest.param(_spectrum_wrapper("pi"), id="pi"),
+    pytest.param(_spectrum_wrapper("solve"), id="solve"),
 ])
 def test_spectrum_solver_equivalence_to_es(spectrum):
     """Test equivalence of the spectrum solvers to the base "es" method."""