Improve lazy tensor warning for SciMLOperators

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Change the structure of block diagonalization functions, using `BlockDiagonalForm` struct and changing the function name from `bdf` to `block_diagonal_form`. ([#349])
 - Add **GPUArrays** compatibility for `ptrace` function, by using **KernelAbstractions.jl**. ([#350])
 - Introduce `Space`, `Dimensions`, `GeneralDimensions` structures to support wider definitions and operations of `Qobj/QobjEvo`, and potential functionalities in the future. ([#271], [#353], [#360])
+- Improve lazy tensor warning for `SciMLOperators`. ([#370])
 
 ## [v0.24.0]
 Release date: 2024-12-13
@@ -75,3 +76,4 @@ Release date: 2024-11-13
 [#350]: https://github.com/qutip/QuantumToolbox.jl/issues/350
 [#353]: https://github.com/qutip/QuantumToolbox.jl/issues/353
 [#360]: https://github.com/qutip/QuantumToolbox.jl/issues/360
+[#370]: https://github.com/qutip/QuantumToolbox.jl/issues/370

src/qobj/functions.jl

@@ -184,6 +184,7 @@ function LinearAlgebra.kron(
     B::AbstractQuantumObject{DT2,OpType,<:Dimensions},
 ) where {DT1,DT2,OpType<:Union{KetQuantumObject,BraQuantumObject,OperatorQuantumObject}}
     QType = promote_op_type(A, B)
+    _lazy_tensor_warning(A.data, B.data)
     return QType(kron(A.data, B.data), A.type, Dimensions((A.dimensions.to..., B.dimensions.to...)))
 end
 
@@ -197,6 +198,7 @@ for ADimType in (:Dimensions, :GeneralDimensions)
                     B::AbstractQuantumObject{DT2,OperatorQuantumObject,<:$BDimType},
                 ) where {DT1,DT2}
                     QType = promote_op_type(A, B)
+                    _lazy_tensor_warning(A.data, B.data)
                     return QType(
                         kron(A.data, B.data),
                         Operator,
@@ -221,6 +223,7 @@ for AOpType in (:KetQuantumObject, :BraQuantumObject, :OperatorQuantumObject)
                     B::AbstractQuantumObject{DT2,$BOpType},
                 ) where {DT1,DT2}
                     QType = promote_op_type(A, B)
+                    _lazy_tensor_warning(A.data, B.data)
                     return QType(
                         kron(A.data, B.data),
                         Operator,

src/qobj/superoperators.jl

@@ -26,16 +26,17 @@ _spre(A::MatrixOperator, Id::AbstractMatrix) = MatrixOperator(_spre(A.A, Id))
 _spre(A::ScaledOperator, Id::AbstractMatrix) = ScaledOperator(A.λ, _spre(A.L, Id))
 _spre(A::AddedOperator, Id::AbstractMatrix) = AddedOperator(map(op -> _spre(op, Id), A.ops))
 function _spre(A::AbstractSciMLOperator, Id::AbstractMatrix)
-    _lazy_tensor_warning("spre", A)
+    _lazy_tensor_warning(Id, A)
     return kron(Id, A)
 end
 
 _spost(B::MatrixOperator, Id::AbstractMatrix) = MatrixOperator(_spost(B.A, Id))
 _spost(B::ScaledOperator, Id::AbstractMatrix) = ScaledOperator(B.λ, _spost(B.L, Id))
 _spost(B::AddedOperator, Id::AbstractMatrix) = AddedOperator(map(op -> _spost(op, Id), B.ops))
 function _spost(B::AbstractSciMLOperator, Id::AbstractMatrix)
-    _lazy_tensor_warning("spost", B)
-    return kron(transpose(B), Id)
+    B_T = transpose(B)
+    _lazy_tensor_warning(B_T, Id)
+    return kron(B_T, Id)
 end
 
 ## intrinsic liouvillian 

src/utilities.jl

@@ -153,8 +153,21 @@
           join(arg, ", ") *
           ")` instead of `$argname = $arg`." maxlog = 1
 
-_lazy_tensor_warning(func_name::String, data::AbstractSciMLOperator) =
-    @warn "The function `$func_name` uses lazy tensor (which can hurt performance) for data type: $(get_typename_wrapper(data))"
+# lazy tensor warning
+for AType in (:AbstractArray, :AbstractSciMLOperator)
+    for BType in (:AbstractArray, :AbstractSciMLOperator)
+        if AType == BType == :AbstractArray
+            @eval begin
+                _lazy_tensor_warning(::$AType, ::$BType) = nothing
+            end
+        else
+            @eval begin
+                _lazy_tensor_warning(A::$AType, B::$BType) =
+                    @warn "using lazy tensor (which can hurt performance) between data types: $(get_typename_wrapper(A)) and $(get_typename_wrapper(B))"
+            end
+        end
+    end
+end
 
 # functions for getting Float or Complex element type
 _FType(::AbstractArray{T}) where {T<:Number} = _FType(T)

test/core-test/quantum_objects_evo.jl

@@ -160,25 +160,19 @@
             @inferred a * a
             @inferred a * a'
 
-            # TODO: kron is currently not supported
-            # @inferred kron(a)
-            # @inferred kron(a, σx)
-            # @inferred kron(a, eye(2))
+            @inferred kron(a)
+            @test_logs (:warn,) @inferred kron(a, σx)
+            @test_logs (:warn,) @inferred kron(a, eye(2))
+            @test_logs (:warn,) (:warn,) @inferred kron(a, eye(2), eye(2))
         end
     end
 
-    # TODO: tensor is currently not supported
-    # @testset "tensor" begin
-    #     σx = sigmax()
-    #     X3 = kron(σx, σx, σx)
-    #     @test tensor(σx) == kron(σx)
-    #     @test tensor(fill(σx, 3)...) == X3
-    #     X_warn = @test_logs (
-    #         :warn,
-    #         "`tensor(A)` or `kron(A)` with `A` is a `Vector` can hurt performance. Try to use `tensor(A...)` or `kron(A...)` instead.",
-    #     ) tensor(fill(σx, 3))
-    #     @test X_warn == X3
-    # end
+    @testset "tensor" begin
+        σx = QobjEvo(sigmax())
+        X3 = @test_logs (:warn,) (:warn,) tensor(σx, σx, σx)
+        X_warn = @test_logs (:warn,) (:warn,) (:warn,) tensor(fill(σx, 3))
+        @test X_warn(0) == X3(0) == tensor(sigmax(), sigmax(), sigmax())
+    end
 
     @testset "Time Dependent Operators and SuperOperators" begin
         N = 10