Working mesolve

Author: albertomercurio

.github/workflows/CI.yml

@@ -55,19 +55,17 @@ 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:
               os: 'ubuntu-latest'
               arch: 'x64'
             group: 'CairoMakie_Ext'
+          - version: '1'
+            node:
+              os: 'ubuntu-latest'
+              arch: 'x64'
+            group: 'AutoDiff_Ext'
 
     steps:
       - uses: actions/checkout@v4

Project.toml

@@ -31,16 +31,19 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 
 [extensions]
 QuantumToolboxCUDAExt = "CUDA"
 QuantumToolboxCairoMakieExt = "CairoMakie"
 QuantumToolboxGPUArraysExt = ["GPUArrays", "KernelAbstractions"]
+QuantumToolboxChainRulesCore = "ChainRulesCore"
 
 [compat]
 ArrayInterface = "6, 7"
 CUDA = "5"
 CairoMakie = "0.12, 0.13"
+ChainRulesCore = "1"
 DiffEqBase = "6"
 DiffEqCallbacks = "4.2.1 - 4"
 DiffEqNoiseProcess = "5"

ext/QuantumToolboxChainRulesCore.jl

@@ -0,0 +1,14 @@
+module QuantumToolboxChainRulesCore
+
+using LinearAlgebra
+import QuantumToolbox: QuantumObject
+using ChainRulesCore
+
+function ChainRulesCore.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())
+    return obj, f_pullback
+end
+
+end

test/ext-test/autodiff/Project.toml

@@ -0,0 +1,5 @@
+[deps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+QuantumToolbox = "6c2fb7c5-b903-41d2-bc5e-5a7c320b9fab"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

test/ext-test/autodiff/zygote.jl

@@ -0,0 +1,48 @@
+@testset "Zygote.jl Autodiff" begin
+    @testset "mesolve" begin
+        N = 16
+        a = destroy(N)
+
+        coef_Δ(p, t) = p[1]
+        coef_F(p, t) = p[2]
+        coef_γ(p, t) = sqrt(p[3])
+
+        H = QobjEvo(a' * a, coef_Δ) + QobjEvo(a + a', coef_F)
+        c_ops = [QobjEvo(a, coef_γ)]
+        L = liouvillian(H, c_ops)
+
+        ψ0 = fock(N, 0)
+
+        function my_f_mesolve(p)
+            tlist = range(0, 40, 100)
+
+            sol = mesolve(
+                L,
+                ψ0,
+                tlist,
+                progress_bar = Val(false),
+                params = p,
+                sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()),
+            )
+
+            return real(expect(a' * a, sol.states[end]))
+        end
+
+        # Analytical solution
+        n_ss(Δ, F, γ) = abs2(F / (Δ + 1im * γ / 2))
+
+        Δ = 1.0
+        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
+
+        grad_exact = Zygote.gradient((p) -> n_ss(p[1], p[2], p[3]), params)[1]
+
+        @test grad_qt ≈ grad_exact atol=1e-6
+    end
+end

test/runtests.jl

@@ -50,7 +50,20 @@ if (GROUP == "All") || (GROUP == "Code-Quality")
     include(joinpath(testdir, "core-test", "code-quality", "code_quality.jl"))
 end
 
-if (GROUP == "CairoMakie_Ext")# || (GROUP == "All")
+if (GROUP == "AutoDiff_Ext")
+    Pkg.activate("ext-test/cpu/autodiff")
+    Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
+    Pkg.instantiate()
+
+    using QuantumToolbox
+    using Zygote
+    using Enzyme
+    using SciMLSensitivity
+
+    include(joinpath(testdir, "ext-test", "autodiff", "zygote.jl"))
+end
+
+if (GROUP == "CairoMakie_Ext")
     Pkg.activate("ext-test/cpu/cairomakie")
     Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
     Pkg.instantiate()
@@ -62,7 +75,7 @@ if (GROUP == "CairoMakie_Ext")# || (GROUP == "All")
     include(joinpath(testdir, "ext-test", "cpu", "cairomakie", "cairomakie_ext.jl"))
 end
 
-if (GROUP == "CUDA_Ext")# || (GROUP == "All")
+if (GROUP == "CUDA_Ext")
     Pkg.activate("ext-test/gpu")
     Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
     Pkg.instantiate()