Implement iterative steady-state solver (#308)

* Implement iterative steady-state solver * Restructure source code * Apply suggestions from code review Co-authored-by: Filippo Vicentini <filippovicentini@gmail.com> * Fix default method * Ensure simd alignment * Larger error margin for 32 bit test * Use similar to conserve container type Co-authored-by: Filippo Vicentini <filippovicentini@gmail.com>
qojulia · Jul 5, 2021 · 45df5b4 · 45df5b4
1 parent 08e1cca
commit 45df5b4
Show file tree

Hide file tree

Showing 6 changed files with 207 additions and 24 deletions.
diff --git a/Project.toml b/Project.toml
@@ -6,7 +6,9 @@ version = "v0.8.9"
 Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 QuantumOpticsBase = "4f57444f-1401-5e15-980d-4471b28d5678"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -20,6 +22,8 @@ WignerSymbols = "9f57e263-0b3d-5e2e-b1be-24f2bb48858b"
 Arpack = "0.5.1"
 DiffEqCallbacks = "2"
 FFTW = "1"
+IterativeSolvers = "0.9"
+LinearMaps = "3"
 OrdinaryDiffEq = "5"
 QuantumOpticsBase = "0.2.7"
 RecursiveArrayTools = "2"

diff --git a/src/QuantumOptics.jl b/src/QuantumOptics.jl
@@ -24,7 +24,15 @@ module timeevolution
     include("mcwf.jl")
     include("bloch_redfield_master.jl")
 end
-include("steadystate.jl")
+module steadystate
+    using QuantumOpticsBase
+    using ..timeevolution
+    using Arpack, LinearAlgebra
+    import LinearMaps
+    import IterativeSolvers
+    include("steadystate.jl")
+    include("steadystate_iterative.jl")
+end
 include("timecorrelations.jl")
 include("spectralanalysis.jl")
 include("semiclassical.jl")

diff --git a/src/master.jl b/src/master.jl
@@ -84,20 +84,7 @@ function master(tspan, rho0::T, H::AbstractOperator{B,B}, J::Vector;
         dmaster_h_(t, rho::T, drho::T) = dmaster_h(rho, H, rates, J, Jdagger, drho, tmp)
         return integrate_master(tspan, dmaster_h_, rho0, fout; kwargs...)
     else
-        Hnh = copy(H)
-        if isa(rates, Matrix)
-            for i=1:length(J), j=1:length(J)
-                Hnh -= complex(float(eltype(H)))(0.5im*rates[i,j])*Jdagger[i]*J[j]
-            end
-        elseif isa(rates, Vector)
-            for i=1:length(J)
-                Hnh -= complex(float(eltype(H)))(0.5im*rates[i])*Jdagger[i]*J[i]
-            end
-        else
-            for i=1:length(J)
-                Hnh -= complex(float(eltype(H)))(0.5im)*Jdagger[i]*J[i]
-            end
-        end
+        Hnh = nh_hamiltonian(H,J,Jdagger,rates)
         Hnhdagger = dagger(Hnh)
         tmp = copy(rho0)
         dmaster_nh_(t, rho::T, drho::T) = dmaster_nh(rho, Hnh, Hnhdagger, rates, J, Jdagger, drho, tmp)
@@ -196,6 +183,28 @@ master_nh(tspan, psi0::Ket{B}, Hnh::AbstractOperator{B,B}, J::Vector; kwargs...)
 master_dynamic(tspan, psi0::Ket{B}, f::Function; kwargs...) where B<:Basis = master_dynamic(tspan, dm(psi0), f; kwargs...)
 master_nh_dynamic(tspan, psi0::Ket{B}, f::Function; kwargs...) where B<:Basis = master_nh_dynamic(tspan, dm(psi0), f; kwargs...)
 
+# Non-hermitian Hamiltonian
+function nh_hamiltonian(H,J,Jdagger,::Nothing)
+    Hnh = copy(H)
+    for i=1:length(J)
+        Hnh -= complex(float(eltype(H)))(0.5im)*Jdagger[i]*J[i]
+    end
+    return Hnh
+end
+function nh_hamiltonian(H,J,Jdagger,rates::AbstractVector)
+    Hnh = copy(H)
+    for i=1:length(J)
+        Hnh -= complex(float(eltype(H)))(0.5im*rates[i])*Jdagger[i]*J[i]
+    end
+    return Hnh
+end
+function nh_hamiltonian(H,J,Jdagger,rates::AbstractMatrix)
+    Hnh = copy(H)
+    for i=1:length(J), j=1:length(J)
+        Hnh -= complex(float(eltype(H)))(0.5im*rates[i,j])*Jdagger[i]*J[j]
+    end
+    return Hnh
+end
 
 # Recasting needed for the ODE solver is just providing the underlying data
 function recast!(x::T, rho::Operator{B,B,T}) where {B<:Basis,T}

diff --git a/src/steadystate.jl b/src/steadystate.jl
@@ -1,9 +1,3 @@
-module steadystate
-
-using QuantumOpticsBase
-using ..timeevolution
-using Arpack, LinearAlgebra
-
 """
     steadystate.master(H, J; <keyword arguments>)
 
@@ -115,6 +109,3 @@ function eigenvector(L::SuperOperator; tol::Real = 1e-9, nev::Int = 2, which::Sy
 end
 
 eigenvector(H::AbstractOperator, J::Vector; rates::Union{Vector, Matrix}=ones(Float64, length(J)), kwargs...) = eigenvector(liouvillian(H, J; rates=rates); kwargs...)
-
-
-end # module
diff --git a/src/steadystate_iterative.jl b/src/steadystate_iterative.jl
@@ -0,0 +1,127 @@
+using ..timeevolution: nh_hamiltonian, dmaster_h, dmaster_nh, check_master
+
+"""
+    iterative!(rho0, H, J, [method!], args...; [log=false], kwargs...) -> rho[, log]
+
+Compute the steady state density matrix of master equation defined by a
+Hamiltonian and a set of jump operators by solving `L rho = 0` via an iterative
+method provided as argument.
+
+# Arguments
+* `rho0`: Initial density matrix. Note that this gets mutated in-place.
+* `H`: Operator specifying the Hamiltonian.
+* `J`: Vector containing all jump operators.
+* `method!`: The iterative method to be used. Defaults to `IterativeSolvers.bicgstabl!`.
+* `rates=nothing`: Vector or matrix specifying the coefficients (decay rates) for
+    the jump operators. If nothing is specified all rates are assumed to be 1.
+* `Jdagger=dagger.(J)`: Vector containing the hermitian conjugates of the jump
+    operators. If they are not given they are calculated automatically.
+* `args...`: Further arguments are passed on to the iterative solver.
+* `kwargs...`: Further keyword arguments are passed on to the iterative solver.
+See also: [`iterative`](@ref)
+
+Credit for this implementation goes to Z. Denis and F. Vicentini.
+See also https://github.com/Z-Denis/SteadyState.jl
+"""
+function iterative!(rho0::Operator, H::AbstractOperator, J,
+                    method! = IterativeSolvers.bicgstabl!, args...;
+                    rates=nothing, Jdagger=dagger.(J), kwargs...)
+
+    # Solution x must satisfy L*x = y with y[end] = tr(x) = 1 and y[j≠end] = 0.
+    M = length(rho0.basis_l)
+    x0 = similar(rho0.data, M^2+1)
+    x0[1:end-1] .= reshape(rho0.data, M^2)
+    x0[end] = zero(eltype(rho0))
+
+    y = similar(rho0.data, M^2+1)
+    y[1:end-1] .= zero(eltype(rho0))
+    y[end] = one(eltype(rho0))
+
+    # Define the linear map lm: rho ↦ L(rho)
+    lm = _linmap_liouvillian(rho0,H,J,Jdagger,rates)
+
+    log = get(kwargs,:log,false)
+
+    # Solve the linear system with the iterative solver, then devectorize rho
+    if !log
+        rho0.data .= @views reshape(method!(x0,lm,y,args...;kwargs...)[1:end-1],(M,M))
+        return rho0
+    else
+        R, history = method!(x0,lm,y,args...;kwargs...)
+        rho0.data .= @views reshape(R[1:end-1],(M,M))
+        return rho0, history
+    end
+end
+
+"""
+    iterative(H, J, [method!], args...; [log=false], kwargs...) -> rho[, log]
+
+Compute the steady state density matrix of master equation defined by a
+Hamiltonian and a set of jump operators by solving `L rho = 0` via an iterative
+method provided as argument.
+
+# Arguments
+* `rho0`: Initial density matrix. Note that this gets mutated in-place.
+* `H`: Operator specifying the Hamiltonian.
+* `J`: Vector containing all jump operators.
+* `method!`: The iterative method to be used. Defaults to `IterativeSolvers.bicgstabl!`
+    or `IterativeSolvers.idrs!` depending on arguments' types.
+* `rates=nothing`: Vector or matrix specifying the coefficients (decay rates) for
+    the jump operators. If nothing is specified all rates are assumed to be 1.
+* `Jdagger=dagger.(J)`: Vector containing the hermitian conjugates of the jump
+    operators. If they are not given they are calculated automatically.
+* `rho0=nothing`: Initial density operator.
+* `args...`: Further arguments are passed on to the iterative solver.
+* `kwargs...`: Further keyword arguments are passed on to the iterative solver.
+See also: [`iterative!`](@ref)
+
+Credit for this implementation goes to Z. Denis and F. Vicentini.
+See also https://github.com/Z-Denis/SteadyState.jl
+"""
+function iterative(H::AbstractOperator, args...;
+                rho0=nothing, kwargs...)
+    if rho0 === nothing
+        rho = DenseOperator(H.basis_l, H.basis_r)
+        rho.data[1,1] = 1
+    else
+        rho = deepcopy(rho0)
+    end
+    return iterative!(rho, H, args...; kwargs...)
+end
+
+
+function _linmap_liouvillian(rho,H,J,Jdagger,rates)
+    bl = rho.basis_l
+    br = rho.basis_r
+    M = length(bl)
+
+    # Cache stuff
+    drho = copy(rho)
+    # rho = copy(rho)
+    Jrho_cache = copy(rho)
+
+    # Check reducibility
+    isreducible = check_master(rho,H,J,Jdagger,rates)
+    if isreducible
+        Hnh = nh_hamiltonian(H,J,Jdagger,rates)
+        Hnhdagger = dagger(Hnh)
+        dmaster_ = (drho,rho) -> dmaster_nh(rho,Hnh,Hnhdagger,rates,J,Jdagger,drho,Jrho_cache)
+    else
+        dmaster_ = (drho,rho) -> dmaster_h(rho,H,rates,J,Jdagger,drho,Jrho_cache)
+    end
+
+    # Linear mapping
+    function f!(y,x)
+        # Reshape
+        drho.data .= @views reshape(y[1:end-1], M, M)
+        rho.data .= @views reshape(x[1:end-1], M, M)
+        # Apply function
+        dmaster_(drho,rho)
+        # Recast data
+        @views y[1:end-1] .= reshape(drho.data, M^2)
+        y[end] = tr(rho)
+        return y
+    end
+
+    return LinearMaps.LinearMap{eltype(rho)}(f!,M^2+1;ismutating=true,issymmetric=false,ishermitian=false,isposdef=false)
+end
diff --git a/test/test_steadystate.jl b/test/test_steadystate.jl
@@ -71,6 +71,40 @@ ev, ops = steadystate.liouvillianspectrum(H, sqrt(2).*J; rates=0.5.*ones(length(
 @test tracedistance(ρss, ops[1]/tr(ops[1])) < 1e-12
 @test ev[sortperm(abs.(ev))] == ev
 
+# Test iterative solvers
+ρss, h = steadystate.iterative(Hdense, Jdense; log=true)
+@test tracedistance(ρss, ρt[end]) < 1e-3
+
+ρss = copy(ρ₀)
+steadystate.iterative!(ρss, H, J)
+@test tracedistance(ρss, ρt[end]) < 1e-3
+
+ρss = steadystate.iterative(H, sqrt(2) .* J; rates=0.5.*ones(length(J)))
+@test tracedistance(ρss, ρt[end]) < 1e-3
+
+ρss = steadystate.iterative(H, sqrt(2) .* J; rates=diagm(0=>[0.5,0.5]))
+@test tracedistance(ρss, ρt[end]) < 1e-3
+
+# Test different float types
+ρss32 = Operator(basis, basis, Matrix{ComplexF32}(ρ₀.data))
+Hdense32 = Operator(basis, basis, Matrix{ComplexF32}(H.data))
+Jdense32 = [Operator(basis, basis, Matrix{ComplexF32}(j.data)) for j ∈ J]
+steadystate.iterative!(ρss32, Hdense32, Jdense32)
+ρt32 = DenseOperator(ρt[end].basis_l, ρt[end].basis_r, Matrix{ComplexF32}(ρt[end].data))
+@test tracedistance(ρss32, ρt[end]) < 2*1e-3
+
+# Iterative methods with lazy operators
+Hlazy = LazySum(Ha, Hc, Hint)
+Ja_lazy = LazyTensor(basis, 1, sqrt(γ)*sm)
+Jc_lazy = LazyTensor(basis, 2, sqrt(κ)*destroy(fockbasis))
+Jlazy = [Ja_lazy, Jc_lazy]
+
+ρss = steadystate.iterative(Hlazy, Jlazy)
+@test tracedistance(ρss, ρt[end]) < 1e-3
+
+# Make sure initial state never got mutated
+@test ρ₀ == dm(Ψ₀)
+
 # Compute steady-state photon number of a driven cavity (analytically: η^2/κ^2)
 Hp = η*(destroy(fockbasis) + create(fockbasis))
 Jp = [sqrt(2κ)*destroy(fockbasis)]
@@ -88,6 +122,16 @@ nss = expect(create(fockbasis)*destroy(fockbasis), ρss)
 nss = expect(create(fockbasis)*destroy(fockbasis), ρss)
 @test n_an - real(nss) < 1e-3
 
+# Test iterative solvers
+ρss = steadystate.iterative(dense(Hp), dense.(Jp))
+nss = expect(create(fockbasis)*destroy(fockbasis), ρss)
+@test n_an - real(nss) < 1e-3
+
+ρss = copy(ρ0_p)
+steadystate.iterative!(ρss, Hp, Jp)
+nss = expect(create(fockbasis)*destroy(fockbasis), ρss)
+@test n_an - real(nss) < 1e-3
+
 
 # Test error messages
 #@test_throws ErrorException steadystate.eigenvector(sx, [sm])