Add sesolve tests

Author: albertomercurio

.github/workflows/CI.yml

@@ -55,6 +55,13 @@ jobs:
           - 'Core'
 
         include:
+          # for core tests (intermediate versions)
+          # - version: '1.x'
+          #   node:
+          #     os: 'ubuntu-latest'
+          #     arch: 'x64'
+          #   group: 'Core'
+
           # for extension tests
           - version: '1'
             node:

Project.toml

@@ -37,7 +37,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 QuantumToolboxCUDAExt = "CUDA"
 QuantumToolboxCairoMakieExt = "CairoMakie"
 QuantumToolboxGPUArraysExt = ["GPUArrays", "KernelAbstractions"]
-QuantumToolboxChainRulesCore = "ChainRulesCore"
+QuantumToolboxChainRulesCoreExt = "ChainRulesCore"
 
 [compat]
 ArrayInterface = "6, 7"

ext/QuantumToolboxChainRulesCore.jl renamed to ext/QuantumToolboxChainRulesCoreExt.jl

@@ -1,10 +1,10 @@
-module QuantumToolboxChainRulesCore
+module QuantumToolboxChainRulesCoreExt
 
 using LinearAlgebra
 import QuantumToolbox: QuantumObject
-using ChainRulesCore
+import ChainRulesCore: rrule, NoTangent, Tangent
 
-function ChainRulesCore.rrule(::Type{QuantumObject}, data, type, dimensions)
+function rrule(::Type{QuantumObject}, data, type, dimensions)
     obj = QuantumObject(data, type, dimensions)
     f_pullback(Δobj) = (NoTangent(), Δobj.data, NoTangent(), NoTangent())
     f_pullback(Δobj_data::AbstractArray) = (NoTangent(), Δobj_data, NoTangent(), NoTangent())

src/QuantumToolbox.jl

@@ -73,6 +73,28 @@ import StaticArraysCore: MVector
 # to achieve better performances for more massive parallelizations
 BLAS.set_num_threads(1)
 
+# TODO: To remove when https://github.com/SciML/SciMLOperators.jl/pull/264 is merged
+SCALINGNUMBERTYPES = (:AbstractSciMLScalarOperator, :Number, :UniformScaling)
+# Special cases for constant scalars. These simplify the structure when applicable
+for T in SCALINGNUMBERTYPES[2:end]
+    @eval function Base.:*(α::$T, L::ScaledOperator)
+        isconstant(L.λ) && return ScaledOperator(α * L.λ, L.L)
+        return ScaledOperator(L.λ, α * L.L) # Try to propagate the rule
+    end
+    @eval function Base.:*(L::ScaledOperator, α::$T)
+        isconstant(L.λ) && return ScaledOperator(α * L.λ, L.L)
+        return ScaledOperator(L.λ, α * L.L) # Try to propagate the rule
+    end
+    @eval function Base.:*(α::$T, L::MatrixOperator)
+        isconstant(L) && return MatrixOperator(α * L.A)
+        return ScaledOperator(α, L) # Going back to the generic case
+    end
+    @eval function Base.:*(L::MatrixOperator, α::$T)
+        isconstant(L) && return MatrixOperator(α * L.A)
+        return ScaledOperator(α, L) # Going back to the generic case
+    end
+end
+
 # Utility
 include("utilities.jl")
 include("versioninfo.jl")

src/qobj/quantum_object_evo.jl

@@ -397,7 +397,6 @@ 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.dimensions))
             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
@@ -409,7 +408,6 @@ Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution`
             (isoper(op_type) || issuper(op_type)) ||
                 throw(ArgumentError("The element must be a Operator or SuperOperator."))
 
-            data_type = op_type.parameters[1]
             dims_expr = (dims_expr..., :(op_func_list[$i].dimensions))
             if i == 1
                 first_op = :(op_func_list[$i])

test/ext-test/autodiff/zygote.jl

@@ -1,6 +1,45 @@
-@testset "Zygote.jl Autodiff" begin
+@testset "Zygote.jl Autodiff" verbose=true begin
+    @testset "sesolve" begin
+        coef_Ω(p, t) = p[1]
+
+        H = QobjEvo(sigmax(), coef_Ω)
+        ψ0 = fock(2, 1)
+        t_max = 10
+
+        function my_f_sesolve(p)
+            tlist = range(0, t_max, 100)
+
+            sol = sesolve(
+                H,
+                ψ0,
+                tlist,
+                progress_bar = Val(false),
+                params = p,
+                sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()),
+            )
+
+            return real(expect(projection(2, 0, 0), sol.states[end]))
+        end
+
+        # Analytical solution
+        my_f_analytic(Ω) = abs2(sin(Ω * t_max))
+        my_f_analytic_deriv(Ω) = 2 * t_max * sin(Ω * t_max) * cos(Ω * t_max)
+
+        Ω = 1.0
+        params = [Ω]
+
+        my_f_analytic(Ω)
+        my_f_sesolve(params)
+
+
+        grad_qt = Zygote.gradient(my_f_sesolve, params)[1] ./ 2
+        grad_exact = [my_f_analytic_deriv(params[1])]
+
+        @test grad_qt ≈ grad_exact atol=1e-6
+    end
+    
     @testset "mesolve" begin
-        N = 16
+        N = 20
         a = destroy(N)
 
         coef_Δ(p, t) = p[1]
@@ -35,8 +74,6 @@
         F = 1.0
         γ = 1.0
         params = [Δ, F, γ]
-        my_f_mesolve(params)
-        n_ss(Δ, F, γ)
 
         # The factor 2 is due to a bug
         grad_qt = Zygote.gradient(my_f_mesolve, params)[1] ./ 2