Improve the construction of QobjEvo (qutip#339)

Author: ytdHuang

CHANGELOG.md

@@ -7,7 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main)
 
-
+- Improve the construction of `QobjEvo`. ([#338], [#339])
 
 ## [v0.23.1]
 Release date: 2024-12-06
@@ -53,3 +53,5 @@ Release date: 2024-11-13
 [#324]: https://github.com/qutip/QuantumToolbox.jl/issues/324
 [#330]: https://github.com/qutip/QuantumToolbox.jl/issues/330
 [#335]: https://github.com/qutip/QuantumToolbox.jl/issues/335
+[#338]: https://github.com/qutip/QuantumToolbox.jl/issues/338
+[#339]: https://github.com/qutip/QuantumToolbox.jl/issues/339

src/qobj/quantum_object_evo.jl

@@ -180,8 +180,19 @@ function QuantumObjectEvolution(
     return QuantumObjectEvolution(data, type, dims)
 end
 
-QuantumObjectEvolution(op::QuantumObject, f::Function; type::Union{Nothing,QuantumObjectType} = nothing) =
-    QuantumObjectEvolution(((op, f),); type = type)
+# 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::Union{Nothing,QuantumObjectType} = nothing,
+) = QuantumObjectEvolution((op_func,), α; type = type)
+
+QuantumObjectEvolution(
+    op::QuantumObject,
+    f::Function,
+    α::Union{Nothing,Number} = nothing;
+    type::Union{Nothing,QuantumObjectType} = nothing,
+) = QuantumObjectEvolution(((op, f),), α; type = type)
 
 function QuantumObjectEvolution(
     op::QuantumObject,
@@ -229,13 +240,15 @@ Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution`
     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]
@@ -248,6 +261,7 @@ Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution`
             op = :(op_func_list[$i][1])
             data_type = op_type.parameters[1]
             dims_expr = (dims_expr..., :($op.dims))
+            func_methods_expr = (func_methods_expr..., :(methods(op_func_list[$i][2], [Any, Real]))) # [Any, Real] means each func must accept 2 arguments
             if i == 1
                 first_op = :($op)
             end
@@ -267,10 +281,20 @@ Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution`
     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 f_method in func_methods
+            length(f_method.ms) == 0 && throw(
+                ArgumentError(
+                    "The following function must accept two arguments: `$(f_method.mt.name)(p, t)` with t<:Real",
+                ),
+            )
+        end
+
         data_expr_const = $qobj_expr_const isa Integer ? $qobj_expr_const : _make_SciMLOperator($qobj_expr_const, α)
 
         data_expr = data_expr_const + $data_expr

src/time_evolution/sesolve.jl

@@ -2,7 +2,8 @@ export sesolveProblem, sesolve
 
 _sesolve_make_U_QobjEvo(H::QuantumObjectEvolution{<:MatrixOperator}) =
     QobjEvo(MatrixOperator(-1im * H.data.A), dims = H.dims, type = Operator)
-_sesolve_make_U_QobjEvo(H) = QobjEvo(H, -1im)
+_sesolve_make_U_QobjEvo(H::QuantumObject) = QobjEvo(MatrixOperator(-1im * H.data), dims = H.dims, type = Operator)
+_sesolve_make_U_QobjEvo(H::Union{QuantumObjectEvolution,Tuple}) = QobjEvo(H, -1im)
 
 @doc raw"""
     sesolveProblem(

test/core-test/quantum_objects_evo.jl

@@ -119,16 +119,24 @@
               "Quantum Object Evo.:   type=SuperOperator   dims=$L_dims   size=$L_size   ishermitian=$L_isherm   isconstant=$L_isconst\n$datastring"
     end
 
-    @testset "Type Inference (QuantumObject)" begin
+    @testset "Type Inference (QobjEvo)" begin
+        N = 4
         for T in [ComplexF32, ComplexF64]
-            N = 4
             a = MatrixOperator(rand(T, N, N))
             @inferred QobjEvo(a)
             for type in [Operator, SuperOperator]
                 @inferred QobjEvo(a, type = type)
             end
         end
 
+        a = destroy(N)
+        coef1(p, t) = exp(-t)
+        coef2(p::Vector, t) = sin(p[1] * t)
+        coef3(p::NamedTuple, t) = cos(p.ω * t)
+        @inferred QobjEvo(a, coef1)
+        @inferred QobjEvo((a', coef2))
+        @inferred QobjEvo((a' * a, (a, coef1), (a', coef2), (a + a', coef3)))
+
         @testset "Math Operation" begin
             a = QobjEvo(destroy(20))
             σx = QobjEvo(sigmax())
@@ -182,6 +190,7 @@
         @test isconstant(H_td) == false
         @test isconstant(QobjEvo(a)) == true
         @test isoper(H_td) == true
+        @test QobjEvo(a, coef1) == QobjEvo((a, coef1))
 
         # SuperOperator
         X = a * a'
@@ -205,7 +214,11 @@
         @test isconstant(L_td) == false
         @test issuper(L_td) == true
 
+        coef_wrong1(t) = nothing
+        coef_wrong2(p, t::ComplexF64) = nothing
         @test_logs (:warn,) (:warn,) liouvillian(H_td * H_td) # warnings from lazy tensor
+        @test_throws ArgumentError QobjEvo(a, coef_wrong1)
+        @test_throws ArgumentError QobjEvo(a, coef_wrong2)
         @test_throws MethodError QobjEvo([[a, coef1], a' * a, [a', coef2]])
         @test_throws ArgumentError H_td(ρvec, p, t)
         @test_throws ArgumentError cache_operator(H_td, ρvec)