Add autodiff benchmarks (#564)

Author: albertomercurio

.github/workflows/Benchmarks.yml

@@ -37,7 +37,7 @@ jobs:
       - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
-          version: '1'
+          version: '1.11'
           arch: x64
       - uses: julia-actions/cache@v2
       - name: Run benchmark

benchmarks/Project.toml

@@ -1,4 +1,8 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 QuantumToolbox = "6c2fb7c5-b903-41d2-bc5e-5a7c320b9fab"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

benchmarks/autodiff.jl

@@ -0,0 +1,92 @@
+function benchmark_autodiff!(SUITE)
+    # Use harmonic oscillator system for both sesolve and mesolve
+    N = 20
+    a = destroy(N)
+    ψ0 = fock(N, 0)
+    tlist = range(0, 40, 100)
+
+    # ---- SESOLVE ----
+    # For direct Forward differentiation
+    function my_f_sesolve_direct(p)
+        H = p[1] * a' * a + p[2] * (a + a')
+        sol = sesolve(H, ψ0, tlist, progress_bar = Val(false))
+        return real(expect(a' * a, sol.states[end]))
+    end
+
+    # For SciMLSensitivity.jl (reverse mode with Zygote and Enzyme)
+    coef_Δ(p, t) = p[1]
+    coef_F(p, t) = p[2]
+    H_evo = QobjEvo(a' * a, coef_Δ) + QobjEvo(a + a', coef_F)
+
+    function my_f_sesolve(p)
+        sol = sesolve(
+            H_evo,
+            ψ0,
+            tlist,
+            progress_bar = Val(false),
+            params = p,
+            sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()),
+        )
+        return real(expect(a' * a, sol.states[end]))
+    end
+
+    # ---- MESOLVE ----
+    # For direct Forward differentiation
+    function my_f_mesolve_direct(p)
+        H = p[1] * a' * a + p[2] * (a + a')
+        c_ops = [sqrt(p[3]) * a]
+        sol = mesolve(H, ψ0, tlist, c_ops, progress_bar = Val(false))
+        return real(expect(a' * a, sol.states[end]))
+    end
+
+    # For SciMLSensitivity.jl (reverse mode with Zygote and Enzyme)
+    coef_γ(p, t) = sqrt(p[3])
+    c_ops = [QobjEvo(a, coef_γ)]
+    L = liouvillian(H_evo, c_ops)
+
+    function my_f_mesolve(p)
+        sol = mesolve(
+            L,
+            ψ0,
+            tlist,
+            progress_bar = Val(false),
+            params = p,
+            sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()),
+        )
+        return real(expect(a' * a, sol.states[end]))
+    end
+
+    # Parameters for benchmarks
+    params_sesolve = [1.0, 1.0]
+    params_mesolve = [1.0, 1.0, 1.0]
+
+    # Benchmark sesolve - Forward
+    SUITE["Autodiff"]["sesolve"]["Forward"] = @benchmarkable ForwardDiff.gradient($my_f_sesolve_direct, $params_sesolve)
+
+    # Benchmark sesolve - Reverse (Zygote)
+    SUITE["Autodiff"]["sesolve"]["Reverse (Zygote)"] = @benchmarkable Zygote.gradient($my_f_sesolve, $params_sesolve)
+
+    # Benchmark sesolve - Reverse (Enzyme)
+    SUITE["Autodiff"]["sesolve"]["Reverse (Enzyme)"] = @benchmarkable Enzyme.autodiff(
+        Enzyme.set_runtime_activity(Enzyme.Reverse),
+        Const($my_f_sesolve),
+        Active,
+        Duplicated($params_sesolve, dparams_sesolve),
+    ) setup=(dparams_sesolve = Enzyme.make_zero($params_sesolve))
+
+    # Benchmark mesolve - Forward
+    SUITE["Autodiff"]["mesolve"]["Forward"] = @benchmarkable ForwardDiff.gradient($my_f_mesolve_direct, $params_mesolve)
+
+    # Benchmark mesolve - Reverse (Zygote)
+    SUITE["Autodiff"]["mesolve"]["Reverse (Zygote)"] = @benchmarkable Zygote.gradient($my_f_mesolve, $params_mesolve)
+
+    # Benchmark mesolve - Reverse (Enzyme)
+    SUITE["Autodiff"]["mesolve"]["Reverse (Enzyme)"] = @benchmarkable Enzyme.autodiff(
+        Enzyme.set_runtime_activity(Enzyme.Reverse),
+        Const($my_f_mesolve),
+        Active,
+        Duplicated($params_mesolve, dparams_mesolve),
+    ) setup=(dparams_mesolve = Enzyme.make_zero($params_mesolve))
+
+    return nothing
+end

benchmarks/runbenchmarks.jl

@@ -3,6 +3,10 @@ using LinearAlgebra
 using SparseArrays
 using QuantumToolbox
 using SciMLBase: EnsembleSerial, EnsembleThreads
+using ForwardDiff
+using Zygote
+using Enzyme: Enzyme, Const, Active, Duplicated
+using SciMLSensitivity: BacksolveAdjoint, EnzymeVJP
 
 BLAS.set_num_threads(1)
 
@@ -14,13 +18,15 @@ include("dynamical_shifted_fock.jl")
 include("eigenvalues.jl")
 include("steadystate.jl")
 include("timeevolution.jl")
+include("autodiff.jl")
 
 benchmark_correlations_and_spectrum!(SUITE)
 benchmark_dfd!(SUITE)
 benchmark_dsf!(SUITE)
 benchmark_eigenvalues!(SUITE)
 benchmark_steadystate!(SUITE)
 benchmark_timeevolution!(SUITE)
+benchmark_autodiff!(SUITE)
 
 BenchmarkTools.tune!(SUITE)
 results = BenchmarkTools.run(SUITE, verbose = true)