Оптимизация кода в стохастическом уравнении Шредингера.Эффективное умножение матриц - PullRequest
0 голосов
/ 14 мая 2019

У меня узкое место в расчете Стохастического уравнения Шредингера (SSE) в Юлии. Основными узкими местами алгоритма являются, конечно, временной шаг, которого для меня достаточно 10e-3, но мне нужно tmax~200; количество реализаций (~100); и тот, который я хотел бы спросить здесь: количество умножений матрицы на вектор, необходимое на каждом временном шаге. Наконец, для Ntot=50 программа работает около получаса, но с 300 это может занять около десяти часов.

Основной код, который выполняет эти вычисления:

function avAPsi(Psi::Array{Complex128}, A::Array{Float64})
    return Psi'*A*Psi
end

function D1(gamma::Float64, Psi::Array{Complex128}, Htot::Array{Float64}, a_mat::Array{Float64}, term1::Complex128, number_operator::Array{Float64})                                                                           
    dummy = term1*a_mat - number_operator
    return -1im*Htot*Psi + 0.5*gamma*dummy*Psi - 0.125*gamma*term1^2*Psi
end

function D2(gamma::Float64, a_mat::Array{Float64}, term1::Complex128, Psi::Array{Complex128})
    return sqrt(gamma)*a_mat*Psi - 0.5*sqrt(gamma)*term1*Psi
end

# ------------------------------------------------------------------------
delta_t = 0.001
wavef_t = Array{Complex128}(Ntot)
dummy = copy(wavef_t)
psiApsi = 0.0
for ii in 1:tsteps
     wavef_t[:] = 0.0
     DW = sqrt(delta_t)*randn()
     # Terms in Runge-Kutta                                                                                   
     psiApsi = avAPsi(init_st, A)
     Psi1 = D1(gamma1, init_st, Htot, a_mat, psiApsi, number_operator)
     dummy = init_st + 0.5*delta_t*Psi1
     psiApsi = avAPsi(dummy, A)
     Psi2 = D1(gamma1, dummy, Htot, a_mat, psiApsi, number_operator)
     dummy = init_st + 0.5*delta_t*Psi2
     psiApsi = avAPsi(dummy, A)
     Psi3 = D1(gamma1, dummy, Htot, a_mat, psiApsi, number_operator)
     dummy = init_st + delta_t*Psi3
     psiApsi = avAPsi(dummy, A)
     Psi4 = D1(gamma1, dummy, Htot, a_mat, psiApsi, number_operator)
     factor2 = D2(gamma1, a_mat, psiApsi, init_st)
     wavef_t = init_st + (1.0/6.0)*(Psi1 + 2.0*Psi2 + 2.0*Psi3 + Psi4)*delta_t + factor2*DW
     wavef_t = wavef_t/norm(wavef_t)                                                   
     init_st = copy(wavef_t)     
end # tsteps 

Таким образом, алгоритм для SSE по существу решает стохастический дифференциал уравнение с использованием 4-го порядка Рунге-Кутта. Так что на каждом временном шаге я должен вычислить 5-член, содержащий матрично-векторное умножение. Как можно улучшить это?

...