fmmesolve multiple c_ops and rate integration fix

doc/changes/1962.feature

@@ -0,0 +1 @@
+Support for multiple coupling operators in fmmesolve

qutip/solve/floquet.py

@@ -553,6 +553,10 @@ def floquet_master_equation_rates(f_modes_0, f_energies, c_op, H, T,
     Calculate the rates and matrix elements for the Floquet-Markov master
     equation.
 
+    .. note :
+        The number of integration steps (for calculating X) within one period
+        is set to 20 * kmax.
+
     Parameters
     ----------
 
@@ -581,7 +585,7 @@ def floquet_master_equation_rates(f_modes_0, f_energies, c_op, H, T,
     w_th : float
         The temperature in units of frequency.
 
-    k_max : int
+    kmax : int
         The truncation of the number of sidebands (default 5).
 
     f_modes_table_t : nested list of :class:`qutip.qobj` (kets)
@@ -611,7 +615,8 @@ def floquet_master_equation_rates(f_modes_0, f_energies, c_op, H, T,
     Gamma = np.zeros((N, N, M))
     A = np.zeros((N, N))
 
-    nT = 100
+    # time steps for integration of coupling operator
+    nT = int(np.max([20 * kmax, 100]))
     dT = T / nT
     tlist = np.arange(dT, T + dT / 2, dT)
 
@@ -935,10 +940,6 @@ def fmmesolve(H, rho0, tlist, c_ops=[], e_ops=[], spectra_cb=[], T=None,
     """
     Solve the dynamics for the system using the Floquet-Markov master equation.
 
-    .. note::
-
-        This solver currently does not support multiple collapse operators.
-
     Parameters
     ----------
 
@@ -1032,15 +1033,23 @@ def fmmesolve(H, rho0, tlist, c_ops=[], e_ops=[], spectra_cb=[], T=None,
     else:
         w_th = 0
 
-    # TODO: loop over input c_ops and spectra_cb, calculate one R for each set
-
-    # calculate the rate-matrices for the floquet-markov master equation
-    Delta, X, Gamma, Amat = floquet_master_equation_rates(
-        f_modes_0, f_energies, c_ops[0], H, T, args, spectra_cb[0],
-        w_th, kmax, f_modes_table_t)
+    # floquet-markov master equation tensor
+    R = None
+    # loop over input c_ops and spectra_cb, calculate one R for each set
+    for c_op, spectrum in zip(c_ops, spectra_cb):
+        # calculate the rate-matrices for the floquet-markov master equation
+        Delta, X, Gamma, Amat = floquet_master_equation_rates(
+            f_modes_0, f_energies, c_op, H, T, args, spectrum,
+            w_th, kmax, f_modes_table_t)
+
+        # calculate temporary floquet-markov master equation tensor
+        R_tmp = floquet_master_equation_tensor(Amat, f_energies)
+        # add temp tensor to general tensor
+        if not R:
+            R = R_tmp
+        else:
+            R += R_tmp
 
-    # the floquet-markov master equation tensor
-    R = floquet_master_equation_tensor(Amat, f_energies)
 
     return floquet_markov_mesolve(R, rho0, tlist, e_ops,
                                   options=options,

qutip/tests/solve/test_floquet.py

@@ -1,8 +1,9 @@
 import numpy as np
-from qutip import fsesolve, sigmax, sigmaz, rand_ket, num, mesolve
+from qutip import fsesolve, sigmax, sigmaz, rand_ket, num, mesolve, sigmay
 from qutip import sigmap, sigmam, floquet_master_equation_rates, expect, Qobj
 from qutip import floquet_modes, floquet_modes_table, fmmesolve
 from qutip import floquet_modes_t_lookup
+import pytest
 
 
 class TestFloquet:
@@ -191,7 +192,8 @@ def noise_spectrum(omega):
 
         np.testing.assert_allclose(np.real(p_ex), np.real(p_ex_ref), atol=1e-4)
 
-    def testFloquetMasterEquation3(self):
+    @pytest.mark.parametrize("kmax", [5, 100, 300])
+    def testFloquetMasterEquation3(self, kmax):
         """
         Test Floquet-Markov Master Equation for a two-level system
         subject to dissipation with internal transform of fmmesolve
@@ -247,7 +249,86 @@ def noise_spectrum(omega):
         # Solve the floquet-markov master equation
         output1 = fmmesolve(
                             H, psi0, tlist, [c_op_fmmesolve], [num(2)],
-                            [noise_spectrum], T, args, floquet_basis=False)
+                            [noise_spectrum], T, args, floquet_basis=False,
+                            kmax=kmax)
+        p_ex = output1.expect[0]
+        # Compare with mesolve
+        output2 = mesolve(H, psi0, tlist, c_op_mesolve, [num(2)], args)
+        p_ex_ref = output2.expect[0]
+
+        np.testing.assert_allclose(np.real(p_ex), np.real(p_ex_ref),
+                                   atol=5 * 1e-4)
+
+    def testFloquetMasterEquation_multiple_coupling(self):
+        """
+        Test Floquet-Markov Master Equation for a two-level system
+        subject to dissipation with multiple coupling operators
+        """
+
+        delta = 1.0 * 2 * np.pi
+        eps0 = 1.0 * 2 * np.pi
+        A = 0.5 * 2 * np.pi
+        omega = np.sqrt(delta ** 2 + eps0 ** 2)
+        T = (2 * np.pi) / omega
+        tlist = np.linspace(0.0, 2 * T, 101)
+        psi0 = rand_ket(2)
+        H0 = - eps0 / 2.0 * sigmaz() - delta / 2.0 * sigmax()
+        H1 = A / 2.0 * sigmax()
+        args = {'w': omega}
+        H = [H0, [H1, lambda t, args: np.sin(args['w'] * t)]]
+        e_ops = [num(2)]
+        gamma1 = 1
+
+        A = 0. * 2 * np.pi
+        psi0 = rand_ket(2)
+        H1 = A / 2.0 * sigmax()
+        args = {'w': omega}
+        H = [H0, [H1, lambda t, args: np.sin(args['w'] * t)]]
+
+        # Collapse operator for Floquet-Markov Master Equation
+        c_ops_fmmesolve = [sigmax(), sigmay()]
+
+        # Collapse operator for Lindblad Master Equation
+        def noise_spectrum1(omega):
+            if omega > 0:
+                return 0.5 * gamma1 * omega/(2*np.pi)
+            else:
+                return 0
+        def noise_spectrum2(omega):
+            if omega > 0:
+                return 0.5 * gamma1 / (2 * np.pi)
+            else:
+                return 0
+
+        noise_spectra = [noise_spectrum1, noise_spectrum2]
+
+        ep, vp = H0.eigenstates()
+        op0 = vp[0]*vp[0].dag()
+        op1 = vp[1]*vp[1].dag()
+
+        c_op_mesolve = []
+
+
+        for c_op_fmmesolve, noise_spectrum in zip(c_ops_fmmesolve,
+                                                   noise_spectra):
+            gamma = np.zeros([2, 2], dtype=complex)
+            for i in range(2):
+                for j in range(2):
+                    if i != j:
+                        gamma[i][j] = 2*np.pi*c_op_fmmesolve.matrix_element(
+                            vp[j], vp[i])*c_op_fmmesolve.matrix_element(
+                            vp[i], vp[j])*noise_spectrum(ep[j]-ep[i])
+
+            for i in range(2):
+                for j in range(2):
+                    c_op_mesolve.append(
+                        np.sqrt(gamma[i][j])*(vp[i]*vp[j].dag())
+                    )
+
+        # Solve the floquet-markov master equation
+        output1 = fmmesolve(
+                            H, psi0, tlist, c_ops_fmmesolve, [num(2)],
+                            noise_spectra, T, args, floquet_basis=False)
         p_ex = output1.expect[0]
         # Compare with mesolve
         output2 = mesolve(H, psi0, tlist, c_op_mesolve, [num(2)], args)