Skip to content

Commit

Permalink
First working case of sesolve
Browse files Browse the repository at this point in the history
  • Loading branch information
albertomercurio committed Oct 15, 2024
1 parent a1c4c16 commit 213f40b
Show file tree
Hide file tree
Showing 6 changed files with 101 additions and 76 deletions.
6 changes: 6 additions & 0 deletions src/qobj/boolean_functions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ export isunitary
Checks if the [`QuantumObject`](@ref) `A` is a [`BraQuantumObject`](@ref). Default case returns `false` for any other inputs.
"""
isbra(A::QuantumObject) = A.type isa BraQuantumObject
isbra(::Type{QuantumObject{DT,BraQuantumObject,N}}) where {DT,N} = true
isbra(A) = false # default case

@doc raw"""
Expand All @@ -19,6 +20,7 @@ isbra(A) = false # default case
Checks if the [`QuantumObject`](@ref) `A` is a [`KetQuantumObject`](@ref). Default case returns `false` for any other inputs.
"""
isket(A::QuantumObject) = A.type isa KetQuantumObject
isket(::Type{QuantumObject{DT,KetQuantumObject,N}}) where {DT,N} = true
isket(A) = false # default case

@doc raw"""
Expand All @@ -27,6 +29,7 @@ isket(A) = false # default case
Checks if the [`AbstractQuantumObject`](@ref) `A` is a [`OperatorQuantumObject`](@ref). Default case returns `false` for any other inputs.
"""
isoper(A::AbstractQuantumObject) = A.type isa OperatorQuantumObject
isoper(::Type{QuantumObject{DT,OperatorQuantumObject,N}}) where {DT,N} = true
isoper(A) = false # default case

@doc raw"""
Expand All @@ -35,6 +38,7 @@ isoper(A) = false # default case
Checks if the [`QuantumObject`](@ref) `A` is a [`OperatorBraQuantumObject`](@ref). Default case returns `false` for any other inputs.
"""
isoperbra(A::QuantumObject) = A.type isa OperatorBraQuantumObject
isoperbra(::Type{QuantumObject{DT,OperatorBraQuantumObject,N}}) where {DT,N} = true
isoperbra(A) = false # default case

@doc raw"""
Expand All @@ -43,6 +47,7 @@ isoperbra(A) = false # default case
Checks if the [`QuantumObject`](@ref) `A` is a [`OperatorKetQuantumObject`](@ref). Default case returns `false` for any other inputs.
"""
isoperket(A::QuantumObject) = A.type isa OperatorKetQuantumObject
isoperket(::Type{QuantumObject{DT,OperatorKetQuantumObject,N}}) where {DT,N} = true
isoperket(A) = false # default case

@doc raw"""
Expand All @@ -51,6 +56,7 @@ isoperket(A) = false # default case
Checks if the [`AbstractQuantumObject`](@ref) `A` is a [`SuperOperatorQuantumObject`](@ref). Default case returns `false` for any other inputs.
"""
issuper(A::AbstractQuantumObject) = A.type isa SuperOperatorQuantumObject
issuper(::Type{QuantumObject{DT,SuperOperatorQuantumObject,N}}) where {DT,N} = true
issuper(A) = false # default case

@doc raw"""
Expand Down
2 changes: 1 addition & 1 deletion src/qobj/eigsolve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -378,7 +378,7 @@ function eigsolve_al(
L = liouvillian(H, c_ops)
prob = mesolveProblem(
L,
QuantumObject(ρ0, type=Operator, dims = H.dims),
QuantumObject(ρ0, type = Operator, dims = H.dims),
[zero(T), T];
alg = alg,
H_t = H_t,
Expand Down
86 changes: 63 additions & 23 deletions src/qobj/quantum_object_evo.jl
Original file line number Diff line number Diff line change
Expand Up @@ -90,18 +90,14 @@ struct QuantumObjectEvolution{DT<:AbstractSciMLOperator,ObjType<:QuantumObjectTy
dims::SVector{N,Int}
end

function QuantumObjectEvolution(op_func_list::Tuple)
# Make the QuantumObjectEvolution, with the option to pre-multiply by a scalar
function QuantumObjectEvolution(op_func_list::Tuple, α = true)
op = _get_op_func_first(op_func_list)
dims = op.dims
type = op.type
T = eltype(op)
data = sum(op_func_list) do op_func
if op_func isa Tuple
return ScalarOperator(zero(T), update_func = (a, u, p, t) -> op_func[2](p, t)) * MatrixOperator(op_func[1].data)
else
return MatrixOperator(op_func.data)
end
end

data = _generate_data(T, op_func_list, α)

# Preallocate the SciMLOperator cache using a dense vector as a reference
v0 = sparse_to_dense(similar(op.data, size(op, 1)))
Expand All @@ -110,31 +106,75 @@ function QuantumObjectEvolution(op_func_list::Tuple)
return QuantumObjectEvolution(data, type, dims)
end

QuantumObjectEvolution(op::QuantumObject) = QuantumObjectEvolution(MatrixOperator(op.data), op.type, op.dims)

function _get_op_func_first(op_func_list)
dims_type_op = map(op_func_list) do op_func
if op_func isa Tuple
length(op_func) == 2 || throw(ArgumentError("The tuple must have two elements."))
((isoper(op_func[1]) || issuper(op_func[1])) && op_func[2] isa Function) || throw(
QuantumObjectEvolution(op::QuantumObject, α = true) =
QuantumObjectEvolution(MatrixOperator* op.data), op.type, op.dims)

@generated function _get_op_func_first(op_func_list::Tuple)
op_func_list_types = op_func_list.parameters
N = length(op_func_list_types)
T = ()
dims_expr = ()
first_op = nothing
for i in 1:N
op_func_type = op_func_list_types[i]
if op_func_type <: Tuple
length(op_func_type.parameters) == 2 || throw(ArgumentError("The tuple must have two elements."))
op_type = op_func_type.parameters[1]
func_type = op_func_type.parameters[2]
((isoper(op_type) || issuper(op_type)) && func_type <: Function) || throw(
ArgumentError(
"The first element must be a Operator or SuperOperator, and the second element must be a function.",
),
)
return (op_func[1].dims, op_func[1].type, op_func[1])
data_type = op_type.parameters[1]
T = (T..., eltype(data_type))
dims_expr = (dims_expr..., :(op_func_list[$i][1].dims))
if i == 1
first_op = :(op_func_list[$i][1])
end
else
(isoper(op_func) || issuper(op_func)) || throw(ArgumentError("The element must be a QuantumObject."))
return (op_func.dims, op_func.type, op_func)
op_type = op_func_type
(isoper(op_type) || issuper(op_type)) ||
throw(ArgumentError("The element must be a Operator or SuperOperator."))
data_type = op_type.parameters[1]
T = (T..., eltype(data_type))
dims_expr = (dims_expr..., :(op_func_list[$i].dims))
if i == 1
first_op = :(op_func_list[$i])
end
end
end

length(unique(getindex.(dims_type_op, 1))) == 1 ||
throw(ArgumentError("The dimensions of the operators must be the same."))
length(unique(T)) == 1 || throw(ArgumentError("The types of the operators must be the same."))

quote
dims = tuple($(dims_expr...))

length(unique(getindex.(dims_type_op, 2))) == 1 ||
throw(ArgumentError("The types of the operators must be the same."))
length(unique(dims)) == 1 || throw(ArgumentError("The dimensions of the operators must be the same."))

return dims_type_op[1][3]
return $first_op
end
end

@generated function _generate_data(T, op_func_list::Tuple, α)
op_func_list_types = op_func_list.parameters
N = length(op_func_list_types)
data_expr = :(0)
for i in 1:N
op_func_type = op_func_list_types[i]
if op_func_type <: Tuple
data_expr = :(
$data_expr +
ScalarOperator(zero(T), op_func_list[$i][2]) * MatrixOperator* op_func_list[$i][1].data)
)
else
data_expr = :($data_expr + MatrixOperator* op_func_list[$i].data))
end
end

quote
return $data_expr
end
end

function (QO::QuantumObjectEvolution)(p, t)
Expand Down
7 changes: 2 additions & 5 deletions src/qobj/superoperators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -89,16 +89,13 @@ the same function is applied multiple times with a known Hilbert space dimension
See also [`spre`](@ref) and [`spost`](@ref).
"""
function lindblad_dissipator(
O::QuantumObject{DT,OperatorQuantumObject},
Id_cache = I(size(O, 1)),
) where {DT}
function lindblad_dissipator(O::QuantumObject{DT,OperatorQuantumObject}, Id_cache = I(size(O, 1))) where {DT}
Od_O = O' * O
return sprepost(O, O') - spre(Od_O, Id_cache) / 2 - spost(Od_O, Id_cache) / 2
end

# It is already a SuperOperator
lindblad_dissipator(O::QuantumObject{DT,SuperOperatorQuantumObject}, Id_cache=nothing) where {DT} = O
lindblad_dissipator(O::QuantumObject{DT,SuperOperatorQuantumObject}, Id_cache = nothing) where {DT} = O

@doc raw"""
liouvillian(H::QuantumObject, c_ops::Union{Nothing,AbstractVector,Tuple}=nothing, Id_cache=I(prod(H.dims)))
Expand Down
6 changes: 3 additions & 3 deletions src/qobj/synonyms.jl
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,13 @@ Note that this functions is same as `QuantumObject(A; type=type, dims=dims)`.
Qobj(A; kwargs...) = QuantumObject(A; kwargs...)

@doc raw"""
QobjEvo(op_func_list::Union{Tuple,QuantumObject})
QobjEvo(op_func_list::Union{Tuple,QuantumObject}, α::Real=true)
Generate [`QuantumObjectEvolution`](@ref)
Note that this functions is same as `QuantumObjectEvolution(op_func_list)`.
Note that this functions is same as `QuantumObjectEvolution(op_func_list)`. If `α` is provided, all the operators in `op_func_list` will be pre-multiplied by `α`.
"""
QobjEvo(op_func_list::Union{Tuple,QuantumObject}) = QuantumObjectEvolution(op_func_list)
QobjEvo(op_func_list::Union{Tuple,QuantumObject}, α = true) = QuantumObjectEvolution(op_func_list, α)

@doc raw"""
shape(A::QuantumObject)
Expand Down
70 changes: 26 additions & 44 deletions src/time_evolution/sesolve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ function _save_func_sesolve(integrator)
internal_params = integrator.p
progr = internal_params.progr

if !internal_params.is_empty_e_ops
if internal_params.is_empty_e_ops
e_ops = internal_params.e_ops
expvals = internal_params.expvals

Expand Down Expand Up @@ -44,14 +44,13 @@ function _generate_sesolve_kwargs(e_ops, progress_bar::Val{false}, t_l, kwargs)
end

@doc raw"""
sesolveProblem(H::QuantumObject,
ψ0::QuantumObject,
tlist::AbstractVector;
alg::OrdinaryDiffEqAlgorithm=Tsit5()
e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
params::NamedTuple=NamedTuple(),
progress_bar::Union{Val,Bool}=Val(true),
sesolveProblem(H,
ψ0,
tlist;
alg=Tsit5()
e_ops = nothing,
params=NamedTuple(),
progress_bar=Val(true),
kwargs...)
Generates the ODEProblem for the Schrödinger time evolution of a quantum system:
Expand All @@ -62,7 +61,7 @@ Generates the ODEProblem for the Schrödinger time evolution of a quantum system
# Arguments
- `H::QuantumObject`: The Hamiltonian of the system ``\hat{H}``.
- `H::Union{QuantumObject,Tuple}`: The Hamiltonian of the system ``\hat{H}``.
- `ψ0::QuantumObject`: The initial state of the system ``|\psi(0)\rangle``.
- `tlist::AbstractVector`: The time list of the evolution.
- `alg::OrdinaryDiffEqAlgorithm`: The algorithm used for the time evolution.
Expand All @@ -85,47 +84,43 @@ Generates the ODEProblem for the Schrödinger time evolution of a quantum system
- `prob`: The `ODEProblem` for the Schrödinger time evolution of the system.
"""
function sesolveProblem(
H::Union{QuantumObject{MT1,OperatorQuantumObject},Tuple},
ψ0::QuantumObject{<:AbstractVector{T2},KetQuantumObject},
H::Union{QuantumObject{DT1,OperatorQuantumObject},Tuple},
ψ0::QuantumObject{DT2,KetQuantumObject},
tlist::AbstractVector;
alg::OrdinaryDiffEqAlgorithm = Tsit5(),
e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
params::NamedTuple = NamedTuple(),
progress_bar::Union{Val,Bool} = Val(true),
kwargs...,
) where {MT1<:AbstractMatrix,T2}
check_dims(H, ψ0)

) where {DT1,DT2}
haskey(kwargs, :save_idxs) &&
throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))

is_time_dependent = !(H_t isa Nothing)

ϕ0 = sparse_to_dense(_CType(ψ0), get_data(ψ0)) # Convert it to dense vector with complex element type

t_l = convert(Vector{_FType(ψ0)}, tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl

U = -1im * get_data(H)
H_evo = QobjEvo(H, -1im) # pre-multiply by -i
isoper(H_evo) || throw(ArgumentError("The Hamiltonian must be an Operator."))
check_dims(H_evo, ψ0)
U = H_evo.data

progr = ProgressBar(length(t_l), enable = getVal(progress_bar))

if e_ops isa Nothing
expvals = Array{ComplexF64}(undef, 0, length(t_l))
e_ops2 = MT1[]
e_ops_data = ()
is_empty_e_ops = true
else
expvals = Array{ComplexF64}(undef, length(e_ops), length(t_l))
e_ops2 = get_data.(e_ops)
e_ops_data = get_data.(e_ops)
is_empty_e_ops = isempty(e_ops)
end

p = (
U = U,
e_ops = e_ops2,
e_ops = e_ops_data,
expvals = expvals,
progr = progr,
Hdims = H.dims,
H_t = H_t,
Hdims = H_evo.dims,
times = t_l,
is_empty_e_ops = is_empty_e_ops,
params...,
Expand All @@ -136,10 +131,8 @@ function sesolveProblem(
kwargs2 = merge(default_values, kwargs)
kwargs3 = _generate_sesolve_kwargs(e_ops, makeVal(progress_bar), t_l, kwargs2)

dudt! = is_time_dependent ? sesolve_td_dudt! : sesolve_ti_dudt!

tspan = (t_l[1], t_l[end])
return ODEProblem{true,FullSpecialize}(dudt!, ϕ0, tspan, p; kwargs3...)
return ODEProblem{true}(U, ϕ0, tspan, p; kwargs3...)
end

@doc raw"""
Expand Down Expand Up @@ -184,27 +177,16 @@ Time evolution of a closed quantum system using the Schrödinger equation:
- `sol::TimeEvolutionSol`: The solution of the time evolution. See also [`TimeEvolutionSol`](@ref)
"""
function sesolve(
H::QuantumObject{MT1,OperatorQuantumObject},
ψ0::QuantumObject{<:AbstractVector{T2},KetQuantumObject},
H::Union{QuantumObject{DT1,OperatorQuantumObject},Tuple},
ψ0::QuantumObject{DT2,KetQuantumObject},
tlist::AbstractVector;
alg::OrdinaryDiffEqAlgorithm = Tsit5(),
e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
params::NamedTuple = NamedTuple(),
progress_bar::Union{Val,Bool} = Val(true),
kwargs...,
) where {MT1<:AbstractMatrix,T2}
prob = sesolveProblem(
H,
ψ0,
tlist;
alg = alg,
e_ops = e_ops,
H_t = H_t,
params = params,
progress_bar = progress_bar,
kwargs...,
)
) where {DT1,DT2}
prob = sesolveProblem(H, ψ0, tlist; e_ops = e_ops, params = params, progress_bar = progress_bar, kwargs...)

return sesolve(prob, alg)
end
Expand Down

0 comments on commit 213f40b

Please sign in to comment.