Improve QuantumObjectEvolution generation function and fix smesolve type instabilities (#548)

albertomercurio · web-flow · commit 64a727c4b1d8 · 2025-09-28T10:01:13.000+08:00
diff --git a/Makefile b/Makefile
@@ -12,10 +12,10 @@ changelog:
 	${JULIA} -e 'using Changelog; Changelog.generate(Changelog.CommonMark(), "CHANGELOG.md"; repo = "qutip/QuantumToolbox.jl")'
 
 test:
-	${JULIA} --project -e 'using Pkg; Pkg.resolve(); Pkg.test()'
+	${JULIA} --project -e 'using Pkg; Pkg.update(); Pkg.test()'
 
 docs:
-	${JULIA} --project=docs -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
+	${JULIA} --project=docs -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.update()'
 	${JULIA} --project=docs docs/make.jl
 
 vitepress:
diff --git a/src/qobj/quantum_object_evo.jl b/src/qobj/quantum_object_evo.jl
@@ -175,19 +175,18 @@ function QuantumObjectEvolution(data::AbstractSciMLOperator; type = Operator(),
 end
 
 @doc raw"""
-    QobjEvo(op_func_list::Union{Tuple,AbstractQuantumObject}, α::Union{Nothing,Number}=nothing; type=nothing)
-    QuantumObjectEvolution(op_func_list::Union{Tuple,AbstractQuantumObject}, α::Union{Nothing,Number}=nothing; type=nothing)
+    QobjEvo(op_func_list::Union{Tuple,AbstractQuantumObject}; type=nothing)
+    QuantumObjectEvolution(op_func_list::Union{Tuple,AbstractQuantumObject}; type=nothing)
 
 Generate [`QuantumObjectEvolution`](@ref).
 
 # Arguments
 - `op_func_list::Union{Tuple,AbstractQuantumObject}`: A tuple of tuples or operators.
-- `α::Union{Nothing,Number}=nothing`: A scalar to pre-multiply the operators.
 
 !!! warning "Beware of type-stability!"
     Please note that, unlike QuTiP, this function doesn't support `op_func_list` as `Vector` type. This is related to the type-stability issue. See the Section [The Importance of Type-Stability](@ref doc:Type-Stability) for more details.
 
-Note that if `α` is provided, all the operators in `op_func_list` will be pre-multiplied by `α`. The `type` parameter is used to specify the type of the [`QuantumObject`](@ref), either `Operator` or `SuperOperator`. The `f` parameter is used to pre-apply a function to the operators before converting them to SciML operators.
+The `type` parameter is used to specify the type of the [`QuantumObject`](@ref), either `Operator` or `SuperOperator`.
 
 !!! note
     `QobjEvo` is a synonym of `QuantumObjectEvolution`.
@@ -267,8 +266,8 @@ Quantum Object:   type=Operator()   dims=[10, 2]   size=(20, 20)   ishermitian=f
 ⎣⠀⠀⠀⠀⠀⠀⠀⠀⠂⡑⎦
 ```
 """
-function QuantumObjectEvolution(op_func_list::Tuple, α::Union{Nothing,Number} = nothing; type = nothing)
-    op, data = _QobjEvo_generate_data(op_func_list, α)
+function QuantumObjectEvolution(op_func_list::Tuple; type = nothing)
+    op, data = _QobjEvo_generate_data(op_func_list)
     dims = op.dimensions
     _check_type(type)
 
@@ -284,17 +283,17 @@ function QuantumObjectEvolution(op_func_list::Tuple, α::Union{Nothing,Number} =
 end
 
 # this is a extra method if user accidentally specify `QuantumObjectEvolution( (op, func) )` or `QuantumObjectEvolution( ((op, func)) )`
-QuantumObjectEvolution(op_func::Tuple{QuantumObject,Function}, α::Union{Nothing,Number} = nothing; type = nothing) =
-    QuantumObjectEvolution((op_func,), α; type = type)
+QuantumObjectEvolution(op_func::Tuple{<:QuantumObject,<:Function}; type = nothing) =
+    QuantumObjectEvolution((op_func,); type = type)
 
 @doc raw"""
-    QuantumObjectEvolution(op::QuantumObject, f::Function, α::Union{Nothing,Number}=nothing; type = nothing)
-    QobjEvo(op::QuantumObject, f::Function, α::Union{Nothing,Number}=nothing; type = nothing)
+    QuantumObjectEvolution(op::QuantumObject, f::Function; type = nothing)
+    QobjEvo(op::QuantumObject, f::Function; type = nothing)
 
 Generate [`QuantumObjectEvolution`](@ref).
 
 # Notes
-- The `f` parameter is used to pre-apply a function to the operators before converting them to SciML operators. The `type` parameter is used to specify the type of the [`QuantumObject`](@ref), either `Operator` or `SuperOperator`.
+- The `f` parameter is time-dependent coefficient that multiplies the operator. The `type` parameter is used to specify the type of the [`QuantumObject`](@ref), either `Operator` or `SuperOperator`.
 - `QobjEvo` is a synonym of `QuantumObjectEvolution`.
 
 # Examples
@@ -318,18 +317,17 @@ Quantum Object Evo.:   type=Operator()   dims=[10, 2]   size=(20, 20)   ishermit
 ScalarOperator(0.0 + 0.0im) * MatrixOperator(20 × 20)
 ```
 """
-QuantumObjectEvolution(op::QuantumObject, f::Function, α::Union{Nothing,Number} = nothing; type = nothing) =
-    QuantumObjectEvolution(((op, f),), α; type = type)
+QuantumObjectEvolution(op::QuantumObject, f::Function; type = nothing) = QuantumObjectEvolution(((op, f),); type = type)
 
-function QuantumObjectEvolution(op::QuantumObject, α::Union{Nothing,Number} = nothing; type = nothing)
+function QuantumObjectEvolution(op::QuantumObject; type = nothing)
     _check_type(type)
     if type isa Nothing
         type = op.type
     end
-    return QuantumObjectEvolution(_make_SciMLOperator(op, α), type, op.dimensions)
+    return QuantumObjectEvolution(_make_SciMLOperator(op), type, op.dimensions)
 end
 
-function QuantumObjectEvolution(op::QuantumObjectEvolution, α::Union{Nothing,Number} = nothing; type = nothing)
+function QuantumObjectEvolution(op::QuantumObjectEvolution; type = nothing)
     _check_type(type)
     if type isa Nothing
         type = op.type
@@ -340,114 +338,97 @@ function QuantumObjectEvolution(op::QuantumObjectEvolution, α::Union{Nothing,Nu
             ),
         )
     end
-    if α isa Nothing
-        return QuantumObjectEvolution(op.data, type, op.dimensions)
-    end
-    return QuantumObjectEvolution(_promote_to_scimloperator(α, op.data), type, op.dimensions)
+    return QuantumObjectEvolution(op.data, type, op.dimensions)
 end
 
 #=
-    _QobjEvo_generate_data(op_func_list::Tuple, α; f::Function=identity)
+    _QobjEvo_generate_data(op_func_list::Tuple)
 
-Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution` object. The `op_func_list` is a tuple of tuples or operators. Each element of the tuple can be a tuple with two elements (operator, function) or an operator. The function is used to generate the time-dependent coefficients for the operators. The `α` parameter is used to pre-multiply the operators by a scalar. The `f` parameter is used to pre-applying a function to the operators before converting them to SciML operators. During the parsing, the dimensions of the operators are checked to be the same, and all the constant operators are summed together.
+Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution` object. The `op_func_list` is a tuple of tuples or operators. Each element of the tuple can be a tuple with two elements (operator, function) or an operator. The function is used to generate the time-dependent coefficients for the operators. During the parsing, the dimensions of the operators are checked to be the same, and all the constant operators are summed together.
 
 # Arguments
 - `op_func_list::Tuple`: A tuple of tuples or operators.
-- `α`: A scalar to pre-multiply the operators.
-- `f::Function=identity`: A function to pre-apply to the operators.
 =#
-@generated function _QobjEvo_generate_data(op_func_list::Tuple, α)
-    op_func_list_types = op_func_list.parameters
-    N = length(op_func_list_types)
-
-    dims_expr = ()
-    func_methods_expr = ()
-    first_op = nothing
-    data_expr = :(0)
-    qobj_expr_const = :(0)
-
-    for i in 1:N
-        op_func_type = op_func_list_types[i]
-        if op_func_type <: Tuple
-            # check the structure of the 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.",
-                ),
-            )
-
-            op = :(op_func_list[$i][1])
-            dims_expr = (dims_expr..., :($op.dimensions))
-            func_methods_expr = (func_methods_expr..., :(methods(op_func_list[$i][2], [Any, Real]).ms)) # [Any, Real] means each func must accept 2 arguments
-            if i == 1
-                first_op = :($op)
-            end
-            data_expr = :($data_expr + _make_SciMLOperator(op_func_list[$i], α))
+function _QobjEvo_generate_data(op_func_list::Tuple)
+    first_op = _QobjEvo_get_first_op(op_func_list[1])
+
+    ops_constant = filter(op_func_list) do x
+        if x isa QuantumObject
+            x.type == first_op.type || throw(ArgumentError("The types of the operators must be the same."))
+            x.dimensions == first_op.dimensions ||
+                throw(ArgumentError("The dimensions of the operators must be the same."))
+            return true
+        elseif x isa Tuple
+            return false
         else
-            op_type = op_func_type
-            (isoper(op_type) || issuper(op_type)) ||
-                throw(ArgumentError("The element must be a Operator or SuperOperator."))
-
-            dims_expr = (dims_expr..., :(op_func_list[$i].dimensions))
-            if i == 1
-                first_op = :(op_func_list[$i])
-            end
-            qobj_expr_const = :($qobj_expr_const + op_func_list[$i])
+            throw(ArgumentError("Each element of the tuple must be either a QuantumObject or a Tuple."))
         end
     end
 
-    quote
-        # check the dims of the operators
-        dims = tuple($(dims_expr...))
-        allequal(dims) || throw(ArgumentError("The dimensions of the operators must be the same."))
-
-        # check if each func accepts 2 arguments
-        func_methods = tuple($(func_methods_expr...))
-        for i in eachindex(func_methods)
-            length(func_methods[i]) == 0 && throw(
-                ArgumentError(
-                    "The following function must only accept two arguments: `$(nameof(op_func_list[i][2]))(p, t)` with t<:Real",
-                ),
-            )
+    ops_time_dep = filter(op_func_list) do x
+        if x isa Tuple
+            _QobjEvo_check_op_func(x)
+            op = x[1]
+            _QobjEvo_check_op(op)
+            op.type == first_op.type || throw(ArgumentError("The types of the operators must be the same."))
+            op.dimensions == first_op.dimensions ||
+                throw(ArgumentError("The dimensions of the operators must be the same."))
+            return true
+        elseif x isa QuantumObject
+            return false
+        else
+            throw(ArgumentError("Each element of the tuple must be either a QuantumObject or a Tuple."))
         end
+    end
 
-        data_expr_const = $qobj_expr_const isa Integer ? $qobj_expr_const : _make_SciMLOperator($qobj_expr_const, α)
+    data_const = isempty(ops_constant) ? zero(eltype(first_op)) : _make_SciMLOperator(sum(ops_constant))
+    data_td =
+        length(ops_time_dep) == 1 ? _make_SciMLOperator(ops_time_dep[1]) :
+        AddedOperator(map(_make_SciMLOperator, ops_time_dep))
+    data = data_const + data_td
 
-        data_expr = data_expr_const + $data_expr
+    return first_op, data
+end
 
-        return $first_op, data_expr
-    end
+function _QobjEvo_check_op(op::AbstractQuantumObject)
+    (isoper(op) || issuper(op)) || throw(ArgumentError("The element must be a Operator or SuperOperator."))
+    return nothing
 end
 
-function _make_SciMLOperator(op_func::Tuple, α)
-    T = eltype(op_func[1])
-    update_func = (a, u, p, t) -> op_func[2](p, t)
-    if α isa Nothing
-        return ScalarOperator(zero(T), update_func) * _promote_to_scimloperator(op_func[1].data)
-    end
-    return ScalarOperator(zero(T), update_func) * _promote_to_scimloperator(α, op_func[1].data)
+function _QobjEvo_check_op_func(op_func::Tuple)
+    length(op_func) == 2 || throw(ArgumentError("The tuple must have two elements."))
+    _QobjEvo_check_op(op_func[1])
+    (op_func[2] isa Function) || throw(ArgumentError("The second element must be a function."))
+    methods(op_func[2], (Any, Real)) |> length == 0 && throw(
+        ArgumentError(
+            "The following function must only accept two arguments: `$(nameof(op_func[2]))(p, t)` with t<:Real",
+        ),
+    )
+    return nothing
 end
 
-function _make_SciMLOperator(op::AbstractQuantumObject, α)
-    if α isa Nothing
-        return _promote_to_scimloperator(op.data)
+_QobjEvo_get_first_op(op_func_list_1::Union{Tuple,AbstractQuantumObject}) =
+    if op_func_list_1 isa Tuple
+        _QobjEvo_check_op_func(op_func_list_1)
+        op = op_func_list_1[1]
+        _QobjEvo_check_op(op)
+        return op
+    else
+        op = op_func_list_1
+        _QobjEvo_check_op(op)
+        return op
     end
-    return _promote_to_scimloperator(α, op.data)
+
+function _make_SciMLOperator(op_func::Tuple)
+    op, coef = op_func
+    T = eltype(op)
+    update_func = (a, u, p, t) -> coef(p, t)
+    return ScalarOperator(zero(T), update_func) * _promote_to_scimloperator(op.data)
 end
+_make_SciMLOperator(op::AbstractQuantumObject) = _promote_to_scimloperator(op.data)
 
 _promote_to_scimloperator(data::AbstractMatrix) = MatrixOperator(data)
 _promote_to_scimloperator(data::AbstractSciMLOperator) = data
-_promote_to_scimloperator(α::Number, data::AbstractMatrix) = MatrixOperator(α * data)
-# We still have to define this for AddedOperator, as it is not present in SciMLOperators.jl
-function _promote_to_scimloperator(α::Number, data::AddedOperator)
-    return AddedOperator(_promote_to_scimloperator.(α, data.ops)) # Try to propagate the rule
-end
-function _promote_to_scimloperator(α::Number, data::AbstractSciMLOperator)
-    return α * data # Going back to the generic case
-end
 
 @doc raw"""
     (A::QuantumObjectEvolution)(ψout, ψin, p, t)
diff --git a/src/time_evolution/mcsolve.jl b/src/time_evolution/mcsolve.jl
@@ -26,16 +26,16 @@ function _normalize_state!(u, dims, normalize_states)
 end
 
 function _mcsolve_make_Heff_QobjEvo(H::QuantumObject, c_ops)
-    c_ops isa Nothing && return QobjEvo(H)
-    return QobjEvo(H - 1im * mapreduce(op -> op' * op, +, c_ops) / 2)
+    c_ops isa Nothing && return QuantumObjectEvolution(H)
+    return QuantumObjectEvolution(H - 1im * mapreduce(op -> op' * op, +, c_ops) / 2)
 end
 function _mcsolve_make_Heff_QobjEvo(H::Tuple, c_ops)
-    c_ops isa Nothing && return QobjEvo(H)
-    return QobjEvo((H..., -1im * mapreduce(op -> op' * op, +, c_ops) / 2))
+    c_ops isa Nothing && return QuantumObjectEvolution(H)
+    return QuantumObjectEvolution((H..., -1im * sum(op -> op' * op, c_ops) / 2))
 end
 function _mcsolve_make_Heff_QobjEvo(H::QuantumObjectEvolution, c_ops)
     c_ops isa Nothing && return H
-    return H + QobjEvo(mapreduce(op -> op' * op, +, c_ops), -1im / 2)
+    return H + QuantumObjectEvolution(sum(op -> -1im * op' * op / 2, c_ops))
 end
 
 @doc raw"""
diff --git a/src/time_evolution/sesolve.jl b/src/time_evolution/sesolve.jl
@@ -1,12 +1,6 @@
 export sesolveProblem, sesolve
 
-_sesolve_make_U_QobjEvo(
-    H::QuantumObjectEvolution{Operator,DimsType,<:MatrixOperator},
-) where {DimsType<:AbstractDimensions} =
-    QobjEvo(MatrixOperator(-1im * H.data.A), dims = H.dimensions, type = Operator())
-_sesolve_make_U_QobjEvo(H::QuantumObject) =
-    QobjEvo(MatrixOperator(-1im * H.data), dims = H.dimensions, type = Operator())
-_sesolve_make_U_QobjEvo(H::Union{QuantumObjectEvolution,Tuple}) = QobjEvo(H, -1im)
+_sesolve_make_U_QobjEvo(H) = -1im * QuantumObjectEvolution(H, type = Operator())
 
 @doc raw"""
     sesolveProblem(
diff --git a/src/time_evolution/smesolve.jl b/src/time_evolution/smesolve.jl
@@ -119,7 +119,7 @@ function smesolveProblem(
         # TODO: # Currently, we are assuming a time-independent MatrixOperator
         # Also, the u state may become non-hermitian, so Tr[Sn * ρ + ρ * Sn'] != real(Tr[Sn * ρ]) / 2
         op_vec = mat2vec(adjoint(op.A))
-        return _spre(op, Id) + _spost(op', Id) + _smesolve_ScalarOperator(op_vec) * Id_op
+        return AddedOperator(_spre(op, Id), _spost(op', Id), _smesolve_ScalarOperator(op_vec) * Id_op)
     end
     D = DiffusionOperator(D_l)
 
diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
@@ -113,7 +113,7 @@ function ssesolveProblem(
         sc_ops_evo_data,
     )
 
-    K = get_data(QobjEvo(H_eff_evo, -1im)) + K_l
+    K = get_data(-1im * QuantumObjectEvolution(H_eff_evo)) + K_l
 
     D_l = map(op -> op + _ScalarOperator_e(op, -) * IdentityOperator(prod(dims)), sc_ops_evo_data)
     D = DiffusionOperator(D_l)
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -26,7 +26,7 @@ const testdir = dirname(@__FILE__)
 if (GROUP == "All") || (GROUP == "Code-Quality")
     Pkg.activate("core-test/code-quality")
     Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
-    Pkg.instantiate()
+    Pkg.update()
 
     using QuantumToolbox
     using Aqua, JET
@@ -39,7 +39,7 @@ end
 if (GROUP == "AutoDiff_Ext")
     Pkg.activate("ext-test/cpu/autodiff")
     Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
-    Pkg.instantiate()
+    Pkg.update()
 
     using QuantumToolbox
     using ForwardDiff
@@ -55,7 +55,7 @@ end
 if (GROUP == "Makie_Ext")
     Pkg.activate("ext-test/cpu/makie")
     Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
-    Pkg.instantiate()
+    Pkg.update()
 
     using QuantumToolbox
     QuantumToolbox.about()
@@ -67,7 +67,7 @@ end
 if (GROUP == "CUDA_Ext")
     Pkg.activate("ext-test/gpu")
     Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
-    Pkg.instantiate()
+    Pkg.update()
 
     using QuantumToolbox
     import LinearAlgebra: Diagonal

Original file line number	Diff line number	Diff line change
`@@ -113,7 +113,7 @@ function ssesolveProblem(`
`113`	`113`	`sc_ops_evo_data,`
`114`	`114`	`)`
`115`	`115`
`116`		`- K = get_data(QobjEvo(H_eff_evo, -1im)) + K_l`
	`116`	`+ K = get_data(-1im * QuantumObjectEvolution(H_eff_evo)) + K_l`
`117`	`117`
`118`	`118`	`D_l = map(op -> op + _ScalarOperator_e(op, -) * IdentityOperator(prod(dims)), sc_ops_evo_data)`
`119`	`119`	`D = DiffusionOperator(D_l)`