%pylab inline

from qutip import *

import time

reps = 1

from IPython.display import HTML

def show_bm_mat(bm_mat, solvers):
    
    m = bm_mat[bm_mat > 0].min()
    
    html = "<table>"

    html += "<tr><td><b>Solver</b></td><td><b>Elapsed time</b></td><td><b>Ratio</b></td></tr>"
    
    for idx, (desc, func) in enumerate(solvers):
    
        if bm_mat[idx] == m:
            html += "<tr><td><b>%s</b></td><td><b>%.8f</b></td><td><b>%.2f</b></td></tr>" % \
                (desc, bm_mat[idx], bm_mat[idx]/m)
        else:
            html += "<tr><td>%s</td><td>%.8f</td><td>%.2f</td></tr>" % \
                (desc, bm_mat[idx], bm_mat[idx]/m)
            
    
    html += "</table>"

    return HTML(html)

def benchmark_steadystate_solvers(args, solvers, problem_func):

    bm_mat = zeros(len(solvers))

    H, c_ops = problem_func(args)
    
    for sol_idx, solver in enumerate(solvers):
        solver_name, solver_func = solver
        try:
            t1 = time.time()
            for r in range(reps):
                rhoss = solver_func(H, c_ops)
            t2 = time.time()
            bm_mat[sol_idx] = (t2 - t1)/reps
    
        except:
            bm_mat[sol_idx] = nan
            
    return bm_mat

solvers = [["power use_umfpack=True",          
            lambda H, c_ops: steadystate(H, c_ops, method='power', use_umfpack=True)],
           ["power use_umfpack=False",         
            lambda H, c_ops: steadystate(H, c_ops, method='power', use_umfpack=False)],
           ["direct sparse use_umfpack=True",  
            lambda H, c_ops: steadystate(H, c_ops, use_umfpack=True, sparse=True)],
           ["direct sparse use_umfpack=False", 
            lambda H, c_ops: steadystate(H, c_ops, use_umfpack=False, sparse=True)],
           ["iterative use_precond=True",      
            lambda H, c_ops: steadystate(H, c_ops, method='iterative', use_precond=True)],
           ["iterative use_precond=False",     
            lambda H, c_ops: steadystate(H, c_ops, method='iterative', use_precond=False)],
           ["iterative-bicg use_precond=True",  
            lambda H, c_ops: steadystate(H, c_ops, method='iterative-bicg', use_precond=True)],
           ["iterative-bicg use_precond=False", 
            lambda H, c_ops: steadystate(H, c_ops, method='iterative-bicg', use_precond=False)],
           ["direct dense use_umfpack=True",   
            lambda H, c_ops: steadystate(H, c_ops, use_umfpack=True, sparse=False)],
           ["direct dense use_umfpack=False",  
            lambda H, c_ops: steadystate(H, c_ops, use_umfpack=False, sparse=False)],
           ["svd_dense",                       
            lambda H, c_ops: steadystate(H, c_ops, method='svd')],
           ["lu",                              
            lambda H, c_ops: steadystate(H, c_ops, method='lu')],
          ]

large_solvers = [
    ["power use_umfpack=True",          
     lambda H, c_ops: steadystate(H, c_ops, method='power', use_umfpack=True)],
    ["power use_umfpack=False",         
     lambda H, c_ops: steadystate(H, c_ops, method='power', use_umfpack=False)],
    ["direct sparse use_umfpack=True",  
     lambda H, c_ops: steadystate(H, c_ops, use_umfpack=True, sparse=True)],
    ["direct sparse use_umfpack=False", 
     lambda H, c_ops: steadystate(H, c_ops, use_umfpack=False, sparse=True)],
    ["iterative use_precond=True",      
     lambda H, c_ops: steadystate(H, c_ops, method='iterative', use_precond=True)],
    ["iterative-bicg use_precond=True",  
     lambda H, c_ops: steadystate(H, c_ops, method='iterative-bicg', use_precond=True)],
   ]

def bm_problem1(N):

    a = tensor(destroy(N), identity(2))
    b = tensor(identity(N), destroy(2))

    H = a.dag() * a + b.dag() * b + 0.25 * (a + a.dag()) * (b + b.dag())

    c_ops = [sqrt(0.1) * a, sqrt(0.075) * a.dag(), sqrt(0.1) * b]
    
    return H, c_ops

bm_mat = benchmark_steadystate_solvers(10, solvers, bm_problem1)
show_bm_mat(bm_mat, solvers)

bm_mat = benchmark_steadystate_solvers(50, large_solvers, bm_problem1)
show_bm_mat(bm_mat, large_solvers)

def bm_problem2(N):
    
    a = destroy(N)
    H = a.dag() * a
    nth = N / 4
    gamma = 0.05
    c_ops = [sqrt(gamma * (nth + 1)) * a, sqrt(gamma * nth) * a.dag()]

    return H, c_ops

bm_mat = benchmark_steadystate_solvers(50, solvers, bm_problem2)
show_bm_mat(bm_mat, solvers)

def bm_problem3(N):
    
    a = tensor(destroy(N), identity(N))
    b = tensor(identity(N), destroy(N))
    
    H = a.dag() * a + 0.25 * b.dag() * b + 0.05 * a.dag() * a * (b + b.dag()) + 0.1 * (a + a.dag()) 

    c_ops = [sqrt(0.05) * a, sqrt(0.015) * a.dag(), sqrt(0.1) * b, sqrt(0.075) * b.dag()]

    return H, c_ops

bm_mat = benchmark_steadystate_solvers(10, large_solvers, bm_problem3)
show_bm_mat(bm_mat, large_solvers)

def bm_problem4(args=None):

    sz = sigmaz()    
    sx = sigmax()
    
    H = sz
    c_ops = [sqrt(0.05) * sx]

    return H, c_ops

bm_mat = benchmark_steadystate_solvers(None, solvers, bm_problem4)
show_bm_mat(bm_mat, solvers)

def bm_problem5(N=1):

    H = 0
    for m in range(N):
        H += tensor([sigmaz() if n == m else identity(2) for n in range(N)])

    for m in range(N-1):
        H += tensor([sigmax() if n in [m,m+1] else identity(2) for n in range(N)])      
        
    c_ops = [sqrt(0.05) * tensor([sigmam() if n == m else identity(2) for n in range(N)])
             for m in range(N)]
   
    return H, c_ops

bm_mat = benchmark_steadystate_solvers(2, solvers, bm_problem5)
show_bm_mat(bm_mat, solvers)

bm_mat = benchmark_steadystate_solvers(4, solvers, bm_problem5)
show_bm_mat(bm_mat, solvers)

bm_mat = benchmark_steadystate_solvers(6, large_solvers, bm_problem5)
show_bm_mat(bm_mat, large_solvers)

from qutip.ipynbtools import version_table

version_table()