Skip to content

Commit 5a354fe

Browse files
Gavin-RockwoodGavin-Rockwood
committed
Updated the state type variable to ST and made the other changes discussed in the PR comments.
1 parent 41e8216 commit 5a354fe

File tree

6 files changed

+42
-46
lines changed

6 files changed

+42
-46
lines changed

‎src/time_evolution/mcsolve.jl‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ end
2222

2323
function _normalize_state!(u, dims, normalize_states, type)
2424
getVal(normalize_states) && normalize!(u)
25-
return QuantumObject(u, type(), dims)
25+
return QuantumObject(u, type, dims)
2626
end
2727

2828
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
110110
"""
111111
function mcsolveProblem(
112112
H::Union{AbstractQuantumObject{Operator},Tuple},
113-
ψ0::QuantumObject{X},
113+
ψ0::QuantumObject{ST},
114114
tlist::AbstractVector,
115115
c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
116116
e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
117117
params = NullParameters(),
118118
rng::AbstractRNG = default_rng(),
119119
jump_callback::TJC = ContinuousLindbladJumpCallback(),
120120
kwargs...,
121-
) where {TJC<:LindbladJumpCallbackType,X<:Union{Ket,Operator}}
121+
) where {TJC<:LindbladJumpCallbackType,ST<:Union{Ket,Operator}}
122122
haskey(kwargs, :save_idxs) &&
123123
throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
124124

@@ -221,7 +221,7 @@ If the environmental measurements register a quantum jump, the wave function und
221221
"""
222222
function mcsolveEnsembleProblem(
223223
H::Union{AbstractQuantumObject{Operator},Tuple},
224-
ψ0::QuantumObject{X},
224+
ψ0::QuantumObject{ST},
225225
tlist::AbstractVector,
226226
c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
227227
e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
@@ -234,7 +234,7 @@ function mcsolveEnsembleProblem(
234234
prob_func::Union{Function,Nothing} = nothing,
235235
output_func::Union{Tuple,Nothing} = nothing,
236236
kwargs...,
237-
) where {TJC<:LindbladJumpCallbackType,X<:Union{Ket,Operator}}
237+
) where {TJC<:LindbladJumpCallbackType,ST<:Union{Ket,Operator}}
238238
_prob_func = isnothing(prob_func) ? _ensemble_dispatch_prob_func(rng, ntraj, tlist, _mcsolve_prob_func) : prob_func
239239
_output_func =
240240
output_func isa Nothing ?
@@ -261,7 +261,7 @@ function mcsolveEnsembleProblem(
261261
ensemble_prob = TimeEvolutionProblem(
262262
EnsembleProblem(prob_mc.prob, prob_func = _prob_func, output_func = _output_func[1], safetycopy = false),
263263
prob_mc.times,
264-
X,
264+
ST(),
265265
prob_mc.dimensions,
266266
(progr = _output_func[2], channel = _output_func[3]),
267267
)
@@ -359,7 +359,7 @@ If the environmental measurements register a quantum jump, the wave function und
359359
"""
360360
function mcsolve(
361361
H::Union{AbstractQuantumObject{Operator},Tuple},
362-
ψ0::QuantumObject{X},
362+
ψ0::QuantumObject{ST},
363363
tlist::AbstractVector,
364364
c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
365365
alg::AbstractODEAlgorithm = DP5(),
@@ -375,7 +375,7 @@ function mcsolve(
375375
keep_runs_results::Union{Val,Bool} = Val(false),
376376
normalize_states::Union{Val,Bool} = Val(true),
377377
kwargs...,
378-
) where {TJC<:LindbladJumpCallbackType} where {X<:Union{Ket,Operator}}
378+
) where {TJC<:LindbladJumpCallbackType} where {ST<:Union{Ket,Operator}}
379379
ens_prob_mc = mcsolveEnsembleProblem(
380380
H,
381381
ψ0,
@@ -415,10 +415,9 @@ function mcsolve(
415415
_expvals_sol_1 isa Nothing ? nothing : map(i -> _get_expvals(sol[:, i], SaveFuncMCSolve), eachindex(sol))
416416
expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP
417417

418-
# stack to transform Vector{Vector{QuantumObject}} -> Matrix{QuantumObject}
419-
# states_all = stack(map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol)), dims = 1)
418+
420419
states_all = stack(
421-
map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states, ens_prob_mc.states_type), eachindex(sol)),
420+
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!
422421
dims = 1,
423422
)
424423

‎src/time_evolution/mesolve.jl‎

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@ _mesolve_make_L_QobjEvo(H::Union{QuantumObjectEvolution,Tuple}, c_ops) = liouvil
66
_mesolve_make_L_QobjEvo(H::Nothing, c_ops::Nothing) = throw(ArgumentError("Both H and
77
c_ops are Nothing. You are probably running the wrong function."))
88

9-
function _gen_mesolve_solution(sol, prob::TimeEvolutionProblem{X}) where {X<:Union{Operator,OperatorKet,SuperOperator}}
10-
if X() == Operator()
11-
ρt = map(ϕ -> QuantumObject(vec2mat(ϕ), type = X(), dims = prob.dimensions), sol.u)
9+
function _gen_mesolve_solution(sol, prob::TimeEvolutionProblem{ST}) where {ST<:Union{Operator,OperatorKet,SuperOperator}}
10+
if ST == Operator
11+
ρt = map(ϕ -> QuantumObject(vec2mat(ϕ), type = prob.states_type, dims = prob.dimensions), sol.u)
1212
else
13-
ρt = map(ϕ -> QuantumObject(ϕ, type = X(), dims = prob.dimensions), sol.u)
13+
ρt = map(ϕ -> QuantumObject(ϕ, type = prob.states_type, dims = prob.dimensions), sol.u)
1414
end
1515

1616
kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
@@ -124,8 +124,6 @@ function mesolveProblem(
124124
elseif issuper(ψ0)
125125
ρ0 = to_dense(_complex_float_type(T), copy(ψ0.data))
126126
state_type = SuperOperator()
127-
else
128-
throw(ArgumentError("Unsupported state type for ψ0 in mesolveProblem."))
129127
end
130128

131129
L = cache_operator(L_evo.data, ρ0)
@@ -138,7 +136,7 @@ function mesolveProblem(
138136

139137
prob = ODEProblem{getVal(inplace),FullSpecialize}(L, ρ0, tspan, params; kwargs4...)
140138

141-
return TimeEvolutionProblem(prob, tlist, state_type, L_evo.dimensions)#, (isoperket = Val(isoperket(ψ0)),))
139+
return TimeEvolutionProblem(prob, tlist, state_type, L_evo.dimensions)
142140
end
143141

144142
@doc raw"""
@@ -246,7 +244,7 @@ end
246244
function mesolve(prob::TimeEvolutionProblem, alg::AbstractODEAlgorithm = DP5(); kwargs...)
247245
sol = solve(prob.prob, alg; kwargs...)
248246

249-
return _gen_mesolve_solution(sol, prob)#, prob.kwargs.isoperket)
247+
return _gen_mesolve_solution(sol, prob)
250248
end
251249

252250
@doc raw"""
@@ -339,8 +337,6 @@ function mesolve_map(
339337
to_dense(_complex_float_type(T), mat2vec(ket2dm(state).data))
340338
elseif issuper(state)
341339
to_dense(_complex_float_type(T), copy(state.data))
342-
else
343-
throw(ArgumentError("Unsupported state type for ψ0 in mesolveProblem."))
344340
end
345341
end
346342
if params isa NullParameters
@@ -379,7 +375,7 @@ mesolve_map(
379375
#
380376
# Return: An array of TimeEvolutionSol objects with the size same as the given iter.
381377
function mesolve_map(
382-
prob::TimeEvolutionProblem{StateOpType,<:ODEProblem},
378+
prob::TimeEvolutionProblem{StateOpType, <:AbstractDimensions, <:ODEProblem},
383379
iter::AbstractArray,
384380
alg::AbstractODEAlgorithm = DP5(),
385381
ensemblealg::EnsembleAlgorithm = EnsembleThreads();

‎src/time_evolution/sesolve.jl‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@ export sesolveProblem, sesolve, sesolve_map
22

33
_sesolve_make_U_QobjEvo(H) = -1im * QuantumObjectEvolution(H, type = Operator())
44

5-
function _gen_sesolve_solution(sol, prob::TimeEvolutionProblem{X}) where {X<:Union{Ket,Operator}}
6-
ψt = map(ϕ -> QuantumObject(ϕ, type = X(), dims = prob.dimensions), sol.u)
5+
function _gen_sesolve_solution(sol, prob::TimeEvolutionProblem{ST}) where {ST<:Union{Ket,Operator}}
6+
ψt = map(ϕ -> QuantumObject(ϕ, type = prob.states_type, dims = prob.dimensions), sol.u)
77

88
kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
99

@@ -61,14 +61,14 @@ Generate the ODEProblem for the Schrödinger time evolution of a quantum system:
6161
"""
6262
function sesolveProblem(
6363
H::Union{AbstractQuantumObject{Operator},Tuple},
64-
ψ0::QuantumObject{X},
64+
ψ0::QuantumObject{ST},
6565
tlist::AbstractVector;
6666
e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
6767
params = NullParameters(),
6868
progress_bar::Union{Val,Bool} = Val(true),
6969
inplace::Union{Val,Bool} = Val(true),
7070
kwargs...,
71-
) where {X<:Union{Ket,Operator}}
71+
) where {ST<:Union{Ket,Operator}}
7272
haskey(kwargs, :save_idxs) &&
7373
throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
7474

@@ -90,7 +90,7 @@ function sesolveProblem(
9090

9191
prob = ODEProblem{getVal(inplace),FullSpecialize}(U, ψ0, tspan, params; kwargs4...)
9292

93-
return TimeEvolutionProblem(prob, tlist, X(), H_evo.dimensions)
93+
return TimeEvolutionProblem(prob, tlist, ST(), H_evo.dimensions)
9494
end
9595

9696
@doc raw"""
@@ -138,15 +138,15 @@ Time evolution of a closed quantum system using the Schrödinger equation:
138138
"""
139139
function sesolve(
140140
H::Union{AbstractQuantumObject{Operator},Tuple},
141-
ψ0::QuantumObject{X},
141+
ψ0::QuantumObject{ST},
142142
tlist::AbstractVector;
143143
alg::AbstractODEAlgorithm = Vern7(lazy = false),
144144
e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
145145
params = NullParameters(),
146146
progress_bar::Union{Val,Bool} = Val(true),
147147
inplace::Union{Val,Bool} = Val(true),
148148
kwargs...,
149-
) where {X<:Union{Ket,Operator}}
149+
) where {ST<:Union{Ket,Operator}}
150150

151151
# Move sensealg argument to solve for Enzyme.jl support.
152152
# TODO: Remove it when https://github.com/SciML/SciMLSensitivity.jl/issues/1225 is fixed.
@@ -225,15 +225,15 @@ for each combination in the ensemble.
225225
"""
226226
function sesolve_map(
227227
H::Union{AbstractQuantumObject{Operator},Tuple},
228-
ψ0::AbstractVector{<:QuantumObject{X}},
228+
ψ0::AbstractVector{<:QuantumObject{ST}},
229229
tlist::AbstractVector;
230230
alg::AbstractODEAlgorithm = Vern7(lazy = false),
231231
ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
232232
e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
233233
params::Union{NullParameters,Tuple} = NullParameters(),
234234
progress_bar::Union{Val,Bool} = Val(true),
235235
kwargs...,
236-
) where {X<:Union{Ket,Operator}}
236+
) where {ST<:Union{Ket,Operator}}
237237
# mapping initial states and parameters
238238

239239
ψ0 = map(to_dense, ψ0) # Convert all initial states to dense vectors
@@ -260,10 +260,10 @@ function sesolve_map(
260260
end
261261
sesolve_map(
262262
H::Union{AbstractQuantumObject{Operator},Tuple},
263-
ψ0::QuantumObject{X},
263+
ψ0::QuantumObject{ST},
264264
tlist::AbstractVector;
265265
kwargs...,
266-
) where {X<:Union{Ket,Operator}} = sesolve_map(H, [ψ0], tlist; kwargs...)
266+
) where {ST<:Union{Ket,Operator}} = sesolve_map(H, [ψ0], tlist; kwargs...)
267267

268268
# this method is for advanced usage
269269
# User can define their own iterator structure, prob_func and output_func
@@ -272,14 +272,14 @@ sesolve_map(
272272
#
273273
# Return: An array of TimeEvolutionSol objects with the size same as the given iter.
274274
function sesolve_map(
275-
prob::TimeEvolutionProblem{X,<:ODEProblem},
275+
prob::TimeEvolutionProblem{ST, <:AbstractDimensions, <:ODEProblem},
276276
iter::AbstractArray,
277277
alg::AbstractODEAlgorithm = Vern7(lazy = false),
278278
ensemblealg::EnsembleAlgorithm = EnsembleThreads();
279279
prob_func::Union{Function,Nothing} = nothing,
280280
output_func::Union{Tuple,Nothing} = nothing,
281281
progress_bar::Union{Val,Bool} = Val(true),
282-
) where {X<:Union{Ket,Operator}}
282+
) where {ST<:Union{Ket,Operator}}
283283
# generate ensemble problem
284284
ntraj = length(iter)
285285
_prob_func = isnothing(prob_func) ? (prob, i, repeat) -> _se_me_map_prob_func(prob, i, repeat, iter) : prob_func

‎src/time_evolution/smesolve.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ function smesolveProblem(
146146
kwargs4...,
147147
)
148148

149-
return TimeEvolutionProblem(prob, tlist, X(), dims)
149+
return TimeEvolutionProblem(prob, tlist,StateOpType(), dims)
150150
end
151151

152152
@doc raw"""

‎src/time_evolution/ssesolve.jl‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is th
7676
"""
7777
function ssesolveProblem(
7878
H::Union{AbstractQuantumObject{Operator},Tuple},
79-
ψ0::QuantumObject{X},
79+
ψ0::QuantumObject{ST},
8080
tlist::AbstractVector,
8181
sc_ops::Union{Nothing,AbstractVector,Tuple,AbstractQuantumObject} = nothing;
8282
e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
@@ -85,7 +85,7 @@ function ssesolveProblem(
8585
progress_bar::Union{Val,Bool} = Val(true),
8686
store_measurement::Union{Val,Bool} = Val(false),
8787
kwargs...,
88-
) where {X<:Union{Ket,Operator}}
88+
) where {ST<:Union{Ket,Operator}}
8989
haskey(kwargs, :save_idxs) &&
9090
throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
9191

@@ -142,7 +142,7 @@ function ssesolveProblem(
142142
kwargs4...,
143143
)
144144

145-
return TimeEvolutionProblem(prob, tlist, X(), dims)
145+
return TimeEvolutionProblem(prob, tlist, ST(), dims)
146146
end
147147

148148
@doc raw"""
@@ -218,7 +218,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is t
218218
"""
219219
function ssesolveEnsembleProblem(
220220
H::Union{AbstractQuantumObject{Operator},Tuple},
221-
ψ0::QuantumObject{X},
221+
ψ0::QuantumObject{ST},
222222
tlist::AbstractVector,
223223
sc_ops::Union{Nothing,AbstractVector,Tuple,AbstractQuantumObject} = nothing;
224224
e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
@@ -231,7 +231,7 @@ function ssesolveEnsembleProblem(
231231
progress_bar::Union{Val,Bool} = Val(true),
232232
store_measurement::Union{Val,Bool} = Val(false),
233233
kwargs...,
234-
) where {X<:Union{Ket,Operator}}
234+
) where {ST<:Union{Ket,Operator}}
235235
_prob_func =
236236
isnothing(prob_func) ?
237237
_ensemble_dispatch_prob_func(
@@ -268,7 +268,7 @@ function ssesolveEnsembleProblem(
268268
ensemble_prob = TimeEvolutionProblem(
269269
EnsembleProblem(prob_sme, prob_func = _prob_func, output_func = _output_func[1], safetycopy = true),
270270
prob_sme.times,
271-
X(),
271+
ST(),
272272
prob_sme.dimensions,
273273
(progr = _output_func[2], channel = _output_func[3]),
274274
)
@@ -356,7 +356,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is th
356356
"""
357357
function ssesolve(
358358
H::Union{AbstractQuantumObject{Operator},Tuple},
359-
ψ0::QuantumObject{X},
359+
ψ0::QuantumObject{ST},
360360
tlist::AbstractVector,
361361
sc_ops::Union{Nothing,AbstractVector,Tuple,AbstractQuantumObject} = nothing;
362362
alg::Union{Nothing,AbstractSDEAlgorithm} = nothing,
@@ -371,7 +371,7 @@ function ssesolve(
371371
keep_runs_results::Union{Val,Bool} = Val(false),
372372
store_measurement::Union{Val,Bool} = Val(false),
373373
kwargs...,
374-
) where {X<:Union{Ket,Operator}}
374+
) where {ST<:Union{Ket,Operator}}
375375
ens_prob = ssesolveEnsembleProblem(
376376
H,
377377
ψ0,
@@ -418,7 +418,7 @@ function ssesolve(
418418
expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP
419419

420420
# stack to transform Vector{Vector{QuantumObject}} -> Matrix{QuantumObject}
421-
states_all = stack(map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol)), dims = 1)
421+
states_all = stack(map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states, [ens_prob.states_type]), eachindex(sol)), dims = 1)
422422

423423
_m_expvals =
424424
_m_expvals_sol_1 isa Nothing ? nothing : map(i -> _get_m_expvals(sol[:, i], SaveFuncSSESolve), eachindex(sol))

‎src/time_evolution/time_evolution.jl‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,14 @@ A Julia constructor for handling the `ODEProblem` of the time evolution of quant
1919
2020
- `prob::AbstractSciMLProblem`: The `ODEProblem` of the time evolution.
2121
- `times::AbstractVector`: The time list of the evolution.
22+
- `states_type::QuantumObjectType`: The type of the quantum states during the evolution (e.g., `Ket`, `Operator`, `OperatorKet` or `SuperOperator`).
2223
- `dimensions::AbstractDimensions`: The dimensions of the Hilbert space.
2324
- `kwargs::KWT`: Generic keyword arguments.
2425
2526
!!! note "`dims` property"
2627
For a given `prob::TimeEvolutionProblem`, `prob.dims` or `getproperty(prob, :dims)` returns its `dimensions` in the type of integer-vector.
2728
"""
28-
struct TimeEvolutionProblem{ST,PT<:AbstractSciMLProblem,TT<:AbstractVector,DT<:AbstractDimensions,KWT}
29+
struct TimeEvolutionProblem{ST<:QuantumObjectType, DT<:AbstractDimensions,PT<:AbstractSciMLProblem,TT<:AbstractVector,KWT}
2930
prob::PT
3031
times::TT
3132
states_type::ST

0 commit comments

Comments
 (0)