Stochastic Schrödinger equation

Project.toml

@@ -6,6 +6,7 @@ version = "0.12.0"
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
+DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
@@ -18,6 +19,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

src/QuantumToolbox.jl

@@ -7,6 +7,7 @@ import Reexport: @reexport
 @reexport using LinearAlgebra
 @reexport using SparseArrays
 @reexport using OrdinaryDiffEq
+@reexport using StochasticDiffEq
 @reexport using LinearSolve
 
 # other functions in LinearAlgebra
@@ -29,6 +30,7 @@ import Graphs: connected_components, DiGraph
 import IncompleteLU: ilu
 import LinearMaps: LinearMap
 import OrdinaryDiffEq: OrdinaryDiffEqAlgorithm
+import DiffEqNoiseProcess: RealWienerProcess, RealWienerProcess!
 import Pkg
 import Random
 import SpecialFunctions: loggamma
@@ -60,6 +62,7 @@ include("time_evolution/mesolve.jl")
 include("time_evolution/lr_mesolve.jl")
 include("time_evolution/sesolve.jl")
 include("time_evolution/mcsolve.jl")
+include("time_evolution/ssesolve.jl")
 include("time_evolution/time_evolution_dynamical.jl")
 
 # Others

src/qobj/operator_sum.jl

@@ -18,7 +18,7 @@ struct OperatorSum{CT<:Vector{<:Number},OT<:AbstractVector} <: AbstractQuantumOb
         mapreduce(x -> size(x) == size_1, &, operators) ||
             throw(DimensionMismatch("All the operators must have the same dimensions."))
         T = promote_type(
-            mapreduce(x -> eltype(x.data), promote_type, operators),
+            mapreduce(x -> eltype(x), promote_type, operators),
             mapreduce(eltype, promote_type, coefficients),
         )
         coefficients2 = T.(coefficients)

src/time_evolution/mcsolve.jl

@@ -86,7 +86,7 @@ function _mcsolve_output_func(sol, i)
     return (sol, false)
 end
 
-function _mcsolve_generate_statistics(sol, i, times, states, expvals_all, jump_times, jump_which)
+function _mcsolve_generate_statistics!(sol, i, times, states, expvals_all, jump_times, jump_which)
     sol_i = sol[:, i]
     !isempty(sol_i.prob.kwargs[:saveat]) ?
     states[i] = [QuantumObject(sol_i.u[i], dims = sol_i.prob.p.Hdims) for i in 1:length(sol_i.u)] : nothing
@@ -524,7 +524,7 @@ function mcsolve(
     jump_which = Vector{Vector{Int16}}(undef, length(sol))
 
     foreach(
-        i -> _mcsolve_generate_statistics(sol, i, times, states, expvals_all, jump_times, jump_which),
+        i -> _mcsolve_generate_statistics!(sol, i, times, states, expvals_all, jump_times, jump_which),
         eachindex(sol),
     )
     expvals = dropdims(sum(expvals_all, dims = 1), dims = 1) ./ length(sol)

src/time_evolution/ssesolve.jl

@@ -0,0 +1,219 @@
+export ssesolveProblem, ssesolveEnsembleProblem, ssesolve
+
+#TODO: Check if works in GPU
+function _ssesolve_update_coefficients!(ψ, coefficients, c_ops)
+    _get_en = op -> real(dot(ψ, op, ψ)) #this is en/2: <Sn + Sn'>/2 = Re<Sn>
+    @. coefficients[2:end-1] = _get_en(c_ops) #coefficients of the OperatorSum: Σ Sn * en/2
+    coefficients[end] = -sum(x -> x^2, coefficients[2:end-1]) / 2 #this last coefficient is -Σen^2/8
+    return nothing
+end
+
+function ssesolve_drift!(du, u, p, t)
+    _ssesolve_update_coefficients!(u, p.K.coefficients, p.c_ops)
+
+    mul!(du, p.K, u)
+
+    return nothing
+end
+
+function ssesolve_diffusion!(du, u, p, t)
+    @inbounds en = @view(p.K.coefficients[2:end-1])
+
+    # du:(H,W). du_reshaped:(H*W,). 
+    # H:Hilbert space dimension, W: number of c_ops
+    du_reshaped = reshape(du, :)
+    mul!(du_reshaped, p.D, u) #du[:,i] = D[i] * u
+
+    du .-= u .* reshape(en, 1, :) #du[:,i] -= en[i] * u
+
+    return nothing
+end
+
+function _ssesolve_prob_func(prob, i, repeat)
+    internal_params = prob.p
+
+    noise = RealWienerProcess!(prob.tspan[1], zeros(length(internal_params.c_ops)), zeros(length(internal_params.c_ops)), save_everystep = false)
+
+    prm = merge(
+        internal_params,
+        (
+            expvals = similar(internal_params.expvals),
+            progr = ProgressBar(size(internal_params.expvals, 2), enable = false),
+        ),
+    )
+
+    return remake(prob, p = prm, noise = noise)
+end
+
+_ssesolve_output_func(sol, i) = (sol, false)
+
+function _ssesolve_generate_statistics!(sol, i, states, expvals_all)
+    sol_i = sol[:, i]
+    !isempty(sol_i.prob.kwargs[:saveat]) ?
+    states[i] = [QuantumObject(sol_i.u[i], dims = sol_i.prob.p.Hdims) for i in 1:length(sol_i.u)] : nothing
+
+    copyto!(view(expvals_all, i, :, :), sol_i.prob.p.expvals)
+    return nothing
+end
+
+function ssesolveProblem(
+    H::QuantumObject{MT1,OperatorQuantumObject},
+    ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+    tlist::AbstractVector,
+    c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
+    e_ops::Union{Nothing,AbstractVector} = nothing,
+    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
+    params::NamedTuple = NamedTuple(),
+    progress_bar::Union{Val,Bool} = Val(true),
+    kwargs...,
+) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
+    H.dims != ψ0.dims && throw(DimensionMismatch("The two quantum objects are not of the same Hilbert dimension."))
+
+    haskey(kwargs, :save_idxs) &&
+        throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
+
+    !(H_t isa Nothing) && throw(ArgumentError("Time-dependent Hamiltonians are not currently supported in ssesolve."))
+
+    progress_bar_val = makeVal(progress_bar)
+
+    t_l = convert(Vector{Float64}, tlist) # Convert it into Float64 to avoid type instabilities for OrdinaryDiffEq.jl
+
+    ϕ0 = get_data(ψ0)
+
+    H_eff = get_data(H - T2(0.5im) * mapreduce(op -> op' * op, +, c_ops))
+    c_ops2 = get_data.(c_ops)
+
+    coefficients = [1.0, fill(0.0, length(c_ops) + 1)...]
+    operators = [-1im * H_eff, c_ops2..., I(prod(H.dims))]
+    K = OperatorSum(coefficients, operators)
+    _ssesolve_update_coefficients!(ϕ0, K.coefficients, c_ops2)
+
+    D = vcat(c_ops2...)
+
+    progr = ProgressBar(length(t_l), enable = getVal(progress_bar_val))
+
+    if e_ops isa Nothing
+        expvals = Array{ComplexF64}(undef, 0, length(t_l))
+        e_ops2 = MT1[]
+        is_empty_e_ops = true
+    else
+        expvals = Array{ComplexF64}(undef, length(e_ops), length(t_l))
+        e_ops2 = get_data.(e_ops)
+        is_empty_e_ops = isempty(e_ops)
+    end
+
+    p = (
+        K = K,
+        D = D,
+        e_ops = e_ops2,
+        c_ops = c_ops2,
+        expvals = expvals,
+        progr = progr,
+        Hdims = H.dims,
+        H_t = H_t,
+        is_empty_e_ops = is_empty_e_ops,
+        params...,
+    )
+
+    saveat = e_ops isa Nothing ? t_l : [t_l[end]]
+    default_values = (DEFAULT_SDE_SOLVER_OPTIONS..., saveat = saveat)
+    kwargs2 = merge(default_values, kwargs)
+    kwargs3 = _generate_sesolve_kwargs(e_ops, progress_bar_val, t_l, kwargs2)
+
+    tspan = (t_l[1], t_l[end])
+    noise = RealWienerProcess!(t_l[1], zeros(length(c_ops)), zeros(length(c_ops)), save_everystep = false)
+    noise_rate_prototype = similar(ϕ0, length(ϕ0), length(c_ops))
+    return SDEProblem{true}(
+        ssesolve_drift!,
+        ssesolve_diffusion!,
+        ϕ0,
+        tspan,
+        p;
+        noise_rate_prototype = noise_rate_prototype,
+        noise = noise,
+        kwargs3...,
+    )
+end
+
+function ssesolveEnsembleProblem(
+    H::QuantumObject{MT1,OperatorQuantumObject},
+    ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+    tlist::AbstractVector,
+    c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
+    e_ops::Union{Nothing,AbstractVector} = nothing,
+    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
+    params::NamedTuple = NamedTuple(),
+    prob_func::Function = _ssesolve_prob_func,
+    output_func::Function = _ssesolve_output_func,
+    kwargs...,
+) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
+    prob_sse = ssesolveProblem(H, ψ0, tlist, c_ops; alg = alg, e_ops = e_ops, H_t = H_t, params = params, kwargs...)
+
+    ensemble_prob = EnsembleProblem(prob_sse, prob_func = prob_func, output_func = output_func)
+
+    return ensemble_prob
+end
+
+function ssesolve(
+    H::QuantumObject{MT1,OperatorQuantumObject},
+    ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+    tlist::AbstractVector,
+    c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
+    e_ops::Union{Nothing,AbstractVector} = nothing,
+    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
+    params::NamedTuple = NamedTuple(),
+    n_traj::Int = 1,
+    ensemble_method = EnsembleThreads(),
+    prob_func::Function = _ssesolve_prob_func,
+    output_func::Function = _ssesolve_output_func,
+    kwargs...,
+) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
+    ens_prob = ssesolveEnsembleProblem(
+        H,
+        ψ0,
+        tlist,
+        c_ops;
+        alg = alg,
+        e_ops = e_ops,
+        H_t = H_t,
+        params = params,
+        prob_func = prob_func,
+        output_func = output_func,
+        kwargs...,
+    )
+
+    return ssesolve(ens_prob; alg = alg, n_traj = n_traj, ensemble_method = ensemble_method)
+end
+
+function ssesolve(
+    ens_prob::EnsembleProblem;
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
+    n_traj::Int = 1,
+    ensemble_method = EnsembleThreads(),
+)
+    sol = solve(ens_prob, alg, ensemble_method, trajectories = n_traj)
+    _sol_1 = sol[:, 1]
+
+    expvals_all = Array{ComplexF64}(undef, length(sol), size(_sol_1.prob.p.expvals)...)
+    states =
+        isempty(_sol_1.prob.kwargs[:saveat]) ? fill(QuantumObject[], length(sol)) :
+        Vector{Vector{QuantumObject}}(undef, length(sol))
+
+    foreach(i -> _ssesolve_generate_statistics!(sol, i, states, expvals_all), eachindex(sol))
+    expvals = dropdims(sum(expvals_all, dims = 1), dims = 1) ./ length(sol)
+
+    return TimeEvolutionSSESol(
+        n_traj,
+        _sol_1.t,
+        states,
+        expvals,
+        expvals_all,
+        sol.converged,
+        _sol_1.alg,
+        _sol_1.prob.kwargs[:abstol],
+        _sol_1.prob.kwargs[:reltol],
+    )
+end

src/time_evolution/time_evolution.jl

@@ -4,6 +4,7 @@ export TimeEvolutionSol, TimeEvolutionMCSol
 export liouvillian, liouvillian_floquet, liouvillian_generalized
 
 const DEFAULT_ODE_SOLVER_OPTIONS = (abstol = 1e-8, reltol = 1e-6, save_everystep = false, save_end = true)
+const DEFAULT_SDE_SOLVER_OPTIONS = (abstol = 1e-2, reltol = 1e-2, save_everystep = false, save_end = true)
 
 @doc raw"""
     struct TimeEvolutionSol
@@ -95,6 +96,37 @@ function Base.show(io::IO, sol::TimeEvolutionMCSol)
     return nothing
 end
 
+struct TimeEvolutionSSESol{
+    TT<:Vector{<:Real},
+    TS<:AbstractVector,
+    TE<:Matrix{ComplexF64},
+    TEA<:Array{ComplexF64,3},
+    T<:Real,
+}
+    n_traj::Int
+    times::TT
+    states::TS
+    expect::TE
+    expect_all::TEA
+    converged::Bool
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm
+    abstol::T
+    reltol::T
+end
+
+function Base.show(io::IO, sol::TimeEvolutionSSESol)
+    print(io, "Solution of quantum trajectories\n")
+    print(io, "(converged: $(sol.converged))\n")
+    print(io, "--------------------------------\n")
+    print(io, "num_trajectories = $(sol.n_traj)\n")
+    # print(io, "num_states = $(length(sol.states[1]))\n")
+    print(io, "num_expect = $(size(sol.expect, 1))\n")
+    print(io, "SDE alg.: $(sol.alg)\n")
+    print(io, "abstol = $(sol.abstol)\n")
+    print(io, "reltol = $(sol.reltol)\n")
+    return nothing
+end
+
 abstract type LindbladJumpCallbackType end
 
 struct ContinuousLindbladJumpCallback <: LindbladJumpCallbackType