%pylab inline

from qutip import *
import time

def system_ops2(N):
    N1 = N
    N2 = N + 1
    a1 = tensor(rand_dm(N1, density=0.75), identity(N2))
    a2 = tensor(identity(N1), rand_dm(N2, density=0.75))
    H = a1.dag() * a1 + a2.dag() * a2
    c_ops = [sqrt(0.01) * a1, sqrt(0.025) * a2]
    
    return H, c_ops

def system_ops(N):
    N1 = N
    N2 = N + 1
    N3 = N + 2
    a1 = tensor(rand_dm(N1, density=0.75), identity(N2), identity(N3))
    a2 = tensor(identity(N1), rand_dm(N2, density=0.75), identity(N3))
    a3 = tensor(identity(N1), identity(N2), rand_dm(N3, density=0.75))
    H = a1.dag() * a1 + a2.dag() * a2 + a3.dag() * a3
    c_ops = [sqrt(0.01) * a1, sqrt(0.025) * a2, sqrt(0.05) * a3]
    
    return H, c_ops

H, c_ops = system_ops(3)

L1 = liouvillian(H, c_ops)

L1

L2 = liouvillian_fast(H, c_ops)

L2

timeit liouvillian(H, c_ops)

timeit liouvillian_fast(H, c_ops)

N1 = 2
N2 = 2

a1 = tensor(rand_dm(N1, density=1.0), identity(N2))
a2 = tensor(identity(N1), rand_dm(N2, density=1.0))

H = a1.dag() * a1 + a2.dag() * a2

c_ops = [sqrt(1.0) * a1, sqrt(2.0) * a2]

N1 = 2

a1 = rand_dm(N1, density=1.0)

H = a1.dag() * a1

c_ops = [sqrt(1.0) * a1]

L1 = liouvillian(H, c_ops)

L1

L2 = liouvillian_fast(H, c_ops)

L2

L1-L2

def benchmark_liouvillian(N_vec):
    
    res1 = zeros(shape(N_vec))
    res2 = zeros(shape(N_vec))
    
    for idx, N in enumerate(N_vec):
        H, c_ops = system_ops(N)
        
        t0 = time.time()
        L = liouvillian(H, c_ops)
        t1 = time.time()
        res1[idx] = t1-t0

        t0 = time.time()
        L = liouvillian_fast(H, c_ops)
        t1 = time.time()
        res2[idx] = t1-t0
    
    return res1, res2

N_vec = array([2,3,4,5,6,7,8])

t0 = time.time()
res1, res2 = benchmark_liouvillian(N_vec)
t1 = time.time()

t1-t0

fig, axes = subplots(1, 2, figsize=(14, 6))

axes[0].plot(N_vec, res1, lw=2, label='liouvillian')
axes[0].plot(N_vec, res2, lw=2, label='liouvillian_fast')
axes[0].legend(loc=0)

axes[1].plot(N_vec, res1 / res2, lw=2, label='speedup')
axes[1].legend(loc=0)

H, c_ops = system_ops2(3)

liouvillian(H, c_ops) - liouvillian_fast(H, c_ops)

liouvillian(H, []) - liouvillian_fast(H, [])

liouvillian(None, c_ops) - liouvillian_fast(None, c_ops)

L_c_ops = liouvillian(None, c_ops)

liouvillian(H, [L_c_ops]) - liouvillian_fast(H, [L_c_ops])

liouvillian(None, [L_c_ops]) - liouvillian_fast(None, [L_c_ops])

L_H = liouvillian_fast(H, [])

liouvillian(H, []) - liouvillian_fast(L_H, [])

liouvillian(H, c_ops) - liouvillian_fast(L_H, c_ops)

liouvillian(H, c_ops) - liouvillian_fast(L_H, [L_c_ops])

from qutip.ipynbtools import version_table; version_table()