-
Notifications
You must be signed in to change notification settings - Fork 32
Adding operator support to time evolution problem #606
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
571abc9
9665194
0e6036a
01baa16
e575f7d
71842a7
41e8216
5a354fe
d8d26a3
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -10,4 +10,5 @@ Manifest.toml | |
| benchmarks/benchmarks_output.json | ||
|
|
||
| .ipynb_checkpoints | ||
| *.ipynb | ||
| .devcontainer/* | ||
| *.ipynb | ||
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -20,9 +20,9 @@ function _mcsolve_output_func(sol, i) | |||||
| return (sol, false) | ||||||
| end | ||||||
|
|
||||||
| function _normalize_state!(u, dims, normalize_states) | ||||||
| function _normalize_state!(u, dims, normalize_states, type) | ||||||
| getVal(normalize_states) && normalize!(u) | ||||||
| return QuantumObject(u, Ket(), dims) | ||||||
| return QuantumObject(u, type, dims) | ||||||
| end | ||||||
|
|
||||||
| function _mcsolve_make_Heff_QobjEvo(H::QuantumObject, c_ops) | ||||||
|
|
@@ -110,15 +110,15 @@ If the environmental measurements register a quantum jump, the wave function und | |||||
| """ | ||||||
| function mcsolveProblem( | ||||||
| H::Union{AbstractQuantumObject{Operator},Tuple}, | ||||||
| ψ0::QuantumObject{Ket}, | ||||||
| ψ0::QuantumObject{ST}, | ||||||
| tlist::AbstractVector, | ||||||
| c_ops::Union{Nothing,AbstractVector,Tuple} = nothing; | ||||||
| e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, | ||||||
| params = NullParameters(), | ||||||
| rng::AbstractRNG = default_rng(), | ||||||
| jump_callback::TJC = ContinuousLindbladJumpCallback(), | ||||||
| kwargs..., | ||||||
| ) where {TJC<:LindbladJumpCallbackType} | ||||||
| ) where {TJC<:LindbladJumpCallbackType,ST<:Union{Ket,Operator}} | ||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
| haskey(kwargs, :save_idxs) && | ||||||
| throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox.")) | ||||||
|
|
||||||
|
|
@@ -221,7 +221,7 @@ If the environmental measurements register a quantum jump, the wave function und | |||||
| """ | ||||||
| function mcsolveEnsembleProblem( | ||||||
| H::Union{AbstractQuantumObject{Operator},Tuple}, | ||||||
| ψ0::QuantumObject{Ket}, | ||||||
| ψ0::QuantumObject{ST}, | ||||||
| tlist::AbstractVector, | ||||||
| c_ops::Union{Nothing,AbstractVector,Tuple} = nothing; | ||||||
| e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, | ||||||
|
|
@@ -234,7 +234,7 @@ function mcsolveEnsembleProblem( | |||||
| prob_func::Union{Function,Nothing} = nothing, | ||||||
| output_func::Union{Tuple,Nothing} = nothing, | ||||||
| kwargs..., | ||||||
| ) where {TJC<:LindbladJumpCallbackType} | ||||||
| ) where {TJC<:LindbladJumpCallbackType,ST<:Union{Ket,Operator}} | ||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I actually put this back to just ST<:Union{Ket}. |
||||||
| _prob_func = isnothing(prob_func) ? _ensemble_dispatch_prob_func(rng, ntraj, tlist, _mcsolve_prob_func) : prob_func | ||||||
| _output_func = | ||||||
| output_func isa Nothing ? | ||||||
|
|
@@ -261,6 +261,7 @@ function mcsolveEnsembleProblem( | |||||
| ensemble_prob = TimeEvolutionProblem( | ||||||
| EnsembleProblem(prob_mc.prob, prob_func = _prob_func, output_func = _output_func[1], safetycopy = false), | ||||||
| prob_mc.times, | ||||||
| ST(), | ||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This was replaced with prob_mc.states_type. |
||||||
| prob_mc.dimensions, | ||||||
| (progr = _output_func[2], channel = _output_func[3]), | ||||||
| ) | ||||||
|
|
@@ -358,7 +359,7 @@ If the environmental measurements register a quantum jump, the wave function und | |||||
| """ | ||||||
| function mcsolve( | ||||||
| H::Union{AbstractQuantumObject{Operator},Tuple}, | ||||||
| ψ0::QuantumObject{Ket}, | ||||||
| ψ0::QuantumObject{ST}, | ||||||
| tlist::AbstractVector, | ||||||
| c_ops::Union{Nothing,AbstractVector,Tuple} = nothing; | ||||||
| alg::AbstractODEAlgorithm = DP5(), | ||||||
|
|
@@ -374,7 +375,7 @@ function mcsolve( | |||||
| keep_runs_results::Union{Val,Bool} = Val(false), | ||||||
| normalize_states::Union{Val,Bool} = Val(true), | ||||||
| kwargs..., | ||||||
| ) where {TJC<:LindbladJumpCallbackType} | ||||||
| ) where {TJC<:LindbladJumpCallbackType} where {ST<:Union{Ket,Operator}} | ||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
| ens_prob_mc = mcsolveEnsembleProblem( | ||||||
| H, | ||||||
| ψ0, | ||||||
|
|
@@ -414,8 +415,11 @@ function mcsolve( | |||||
| _expvals_sol_1 isa Nothing ? nothing : map(i -> _get_expvals(sol[:, i], SaveFuncMCSolve), eachindex(sol)) | ||||||
| expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP | ||||||
|
|
||||||
| # stack to transform Vector{Vector{QuantumObject}} -> Matrix{QuantumObject} | ||||||
| states_all = stack(map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol)), dims = 1) | ||||||
|
|
||||||
| states_all = stack( | ||||||
| map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states, [ens_prob_mc.states_type]), eachindex(sol)), # Unsure why ens_prob_mc.states_type needs to be in an array but the other two arguments don't! | ||||||
| dims = 1, | ||||||
| ) | ||||||
|
|
||||||
| col_times = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.col_times, eachindex(sol)) | ||||||
| col_which = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.col_which, eachindex(sol)) | ||||||
|
|
||||||
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -6,17 +6,17 @@ _mesolve_make_L_QobjEvo(H::Union{QuantumObjectEvolution,Tuple}, c_ops) = liouvil | |||||
| _mesolve_make_L_QobjEvo(H::Nothing, c_ops::Nothing) = throw(ArgumentError("Both H and | ||||||
| c_ops are Nothing. You are probably running the wrong function.")) | ||||||
|
|
||||||
| function _gen_mesolve_solution(sol, times, dimensions, isoperket::Val) | ||||||
| if getVal(isoperket) | ||||||
| ρt = map(ϕ -> QuantumObject(ϕ, type = OperatorKet(), dims = dimensions), sol.u) | ||||||
| function _gen_mesolve_solution(sol, prob::TimeEvolutionProblem{ST}) where {ST<:Union{Operator,OperatorKet,SuperOperator}} | ||||||
| if prob.states_type == Operator | ||||||
| ρt = map(ϕ -> QuantumObject(vec2mat(ϕ), type = prob.states_type, dims = prob.dimensions), sol.u) | ||||||
| else | ||||||
| ρt = map(ϕ -> QuantumObject(vec2mat(ϕ), type = Operator(), dims = dimensions), sol.u) | ||||||
| ρt = map(ϕ -> QuantumObject(ϕ, type = prob.states_type, dims = prob.dimensions), sol.u) | ||||||
| end | ||||||
|
|
||||||
| kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility | ||||||
|
|
||||||
| return TimeEvolutionSol( | ||||||
| times, | ||||||
| prob.times, | ||||||
| sol.t, | ||||||
| ρt, | ||||||
| _get_expvals(sol, SaveFuncMESolve), | ||||||
|
|
@@ -86,8 +86,8 @@ function mesolveProblem( | |||||
| progress_bar::Union{Val,Bool} = Val(true), | ||||||
| inplace::Union{Val,Bool} = Val(true), | ||||||
| kwargs..., | ||||||
| ) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet}} | ||||||
| (isoper(H) && isket(ψ0) && isnothing(c_ops)) && return sesolveProblem( | ||||||
| ) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet,SuperOperator}} | ||||||
| (isoper(H) && (isket(ψ0) || isoper(ψ0)) && isnothing(c_ops)) && return sesolveProblem( | ||||||
| H, | ||||||
| ψ0, | ||||||
| tlist; | ||||||
|
|
@@ -107,11 +107,25 @@ function mesolveProblem( | |||||
| check_dimensions(L_evo, ψ0) | ||||||
|
|
||||||
| T = Base.promote_eltype(L_evo, ψ0) | ||||||
| ρ0 = if isoperket(ψ0) # Convert it to dense vector with complex element type | ||||||
| to_dense(_complex_float_type(T), copy(ψ0.data)) | ||||||
| else | ||||||
| to_dense(_complex_float_type(T), mat2vec(ket2dm(ψ0).data)) | ||||||
| # ρ0 = if isoperket(ψ0) # Convert it to dense vector with complex element type | ||||||
| # to_dense(_complex_float_type(T), copy(ψ0.data)) | ||||||
| # else | ||||||
| # to_dense(_complex_float_type(T), mat2vec(ket2dm(ψ0).data)) | ||||||
| # end | ||||||
| if isoper(ψ0) | ||||||
| ρ0 = to_dense(_complex_float_type(T), mat2vec(ψ0.data)) | ||||||
| state_type = Operator() | ||||||
| elseif isoperket(ψ0) | ||||||
| ρ0 = to_dense(_complex_float_type(T), copy(ψ0.data)) | ||||||
| state_type = OperatorKet() | ||||||
| elseif isket(ψ0) | ||||||
| ρ0 = to_dense(_complex_float_type(T), mat2vec(ket2dm(ψ0).data)) | ||||||
| state_type = Operator() | ||||||
| elseif issuper(ψ0) | ||||||
| ρ0 = to_dense(_complex_float_type(T), copy(ψ0.data)) | ||||||
| state_type = SuperOperator() | ||||||
| end | ||||||
|
|
||||||
| L = cache_operator(L_evo.data, ρ0) | ||||||
|
|
||||||
| kwargs2 = _merge_saveat(tlist, e_ops, DEFAULT_ODE_SOLVER_OPTIONS; kwargs...) | ||||||
|
|
@@ -122,7 +136,7 @@ function mesolveProblem( | |||||
|
|
||||||
| prob = ODEProblem{getVal(inplace),FullSpecialize}(L, ρ0, tspan, params; kwargs4...) | ||||||
|
|
||||||
| return TimeEvolutionProblem(prob, tlist, L_evo.dimensions, (isoperket = Val(isoperket(ψ0)),)) | ||||||
| return TimeEvolutionProblem(prob, tlist, state_type, L_evo.dimensions) | ||||||
| end | ||||||
|
|
||||||
| @doc raw""" | ||||||
|
|
@@ -188,8 +202,8 @@ function mesolve( | |||||
| progress_bar::Union{Val,Bool} = Val(true), | ||||||
| inplace::Union{Val,Bool} = Val(true), | ||||||
| kwargs..., | ||||||
| ) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet}} | ||||||
| (isoper(H) && isket(ψ0) && isnothing(c_ops)) && return sesolve( | ||||||
| ) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet,SuperOperator}} | ||||||
| (isoper(H) && (isket(ψ0) || isoper(ψ0)) && isnothing(c_ops)) && return sesolve( | ||||||
| H, | ||||||
| ψ0, | ||||||
| tlist; | ||||||
|
|
@@ -230,7 +244,7 @@ end | |||||
| function mesolve(prob::TimeEvolutionProblem, alg::AbstractODEAlgorithm = DP5(); kwargs...) | ||||||
| sol = solve(prob.prob, alg; kwargs...) | ||||||
|
|
||||||
| return _gen_mesolve_solution(sol, prob.times, prob.dimensions, prob.kwargs.isoperket) | ||||||
| return _gen_mesolve_solution(sol, prob) | ||||||
| end | ||||||
|
|
||||||
| @doc raw""" | ||||||
|
|
@@ -298,8 +312,8 @@ function mesolve_map( | |||||
| params::Union{NullParameters,Tuple} = NullParameters(), | ||||||
| progress_bar::Union{Val,Bool} = Val(true), | ||||||
| kwargs..., | ||||||
| ) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet}} | ||||||
| (isoper(H) && all(isket, ψ0) && isnothing(c_ops)) && return sesolve_map( | ||||||
| ) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet,SuperOperator}} | ||||||
| (isoper(H) && (all(isket, ψ0) || all(isoper, ψ0)) && isnothing(c_ops)) && return sesolve_map( | ||||||
| H, | ||||||
| ψ0, | ||||||
| tlist; | ||||||
|
|
@@ -315,10 +329,14 @@ function mesolve_map( | |||||
| # Convert to appropriate format based on state type | ||||||
| ψ0_iter = map(ψ0) do state | ||||||
| T = _complex_float_type(eltype(state)) | ||||||
| if isoperket(state) | ||||||
| to_dense(T, copy(state.data)) | ||||||
| else | ||||||
| to_dense(T, mat2vec(ket2dm(state).data)) | ||||||
| if isoper(state) | ||||||
| to_dense(_complex_float_type(T), mat2vec(state.data)) | ||||||
| elseif isoperket(state) | ||||||
| to_dense(_complex_float_type(T), copy(state.data)) | ||||||
| elseif isket(state) | ||||||
| to_dense(_complex_float_type(T), mat2vec(ket2dm(state).data)) | ||||||
| elseif issuper(state) | ||||||
| to_dense(_complex_float_type(T), copy(state.data)) | ||||||
| end | ||||||
| end | ||||||
| if params isa NullParameters | ||||||
|
|
@@ -347,7 +365,7 @@ mesolve_map( | |||||
| tlist::AbstractVector, | ||||||
| c_ops::Union{Nothing,AbstractVector,Tuple} = nothing; | ||||||
| kwargs..., | ||||||
| ) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet}} = | ||||||
| ) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator,OperatorKet,SuperOperator}} = | ||||||
| mesolve_map(H, [ψ0], tlist, c_ops; kwargs...) | ||||||
|
|
||||||
| # this method is for advanced usage | ||||||
|
|
@@ -357,14 +375,14 @@ mesolve_map( | |||||
| # | ||||||
| # Return: An array of TimeEvolutionSol objects with the size same as the given iter. | ||||||
| function mesolve_map( | ||||||
| prob::TimeEvolutionProblem{<:ODEProblem}, | ||||||
| prob::TimeEvolutionProblem{StateOpType, <:AbstractDimensions, <:ODEProblem}, | ||||||
| iter::AbstractArray, | ||||||
| alg::AbstractODEAlgorithm = DP5(), | ||||||
| ensemblealg::EnsembleAlgorithm = EnsembleThreads(); | ||||||
| prob_func::Union{Function,Nothing} = nothing, | ||||||
| output_func::Union{Tuple,Nothing} = nothing, | ||||||
| progress_bar::Union{Val,Bool} = Val(true), | ||||||
| ) | ||||||
| ) where {StateOpType<:Union{Ket,Operator,OperatorKet,SuperOperator}} | ||||||
| # generate ensemble problem | ||||||
| ntraj = length(iter) | ||||||
| _prob_func = isnothing(prob_func) ? (prob, i, repeat) -> _se_me_map_prob_func(prob, i, repeat, iter) : prob_func | ||||||
|
|
@@ -380,14 +398,14 @@ function mesolve_map( | |||||
| ens_prob = TimeEvolutionProblem( | ||||||
| EnsembleProblem(prob.prob, prob_func = _prob_func, output_func = _output_func[1], safetycopy = false), | ||||||
| prob.times, | ||||||
| StateOpType(), | ||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
| prob.dimensions, | ||||||
| (progr = _output_func[2], channel = _output_func[3], isoperket = prob.kwargs.isoperket), | ||||||
| (progr = _output_func[2], channel = _output_func[3]), | ||||||
| ) | ||||||
|
|
||||||
| sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj) | ||||||
|
|
||||||
| # handle solution and make it become an Array of TimeEvolutionSol | ||||||
| sol_vec = | ||||||
| [_gen_mesolve_solution(sol[:, i], prob.times, prob.dimensions, prob.kwargs.isoperket) for i in eachindex(sol)] # map is type unstable | ||||||
| sol_vec = [_gen_mesolve_solution(sol[:, i], prob) for i in eachindex(sol)] # map is type unstable | ||||||
| return reshape(sol_vec, size(iter)) | ||||||
| end | ||||||
Uh oh!
There was an error while loading. Please reload this page.