support autodiff for SteadyStateODESolver

Author: ytdHuang

src/steadystate.jl

@@ -214,15 +214,16 @@ end
 function _steadystate(L::AbstractQuantumObject{SuperOperator}, solver::SteadyStateODESolver; kwargs...)
     ψ0 = isnothing(solver.ψ0) ? rand_ket(L.dimensions) : solver.ψ0
     ftype = _float_type(ψ0)
-    tlist = [ftype(0), ftype(solver.tmax)]
+    tmax = ftype(solver.tmax)
+    tlist = [ftype(0), tmax]
 
     # overwrite some kwargs and throw warning message to tell the users that we are ignoring these settings
     haskey(kwargs, :progress_bar) && @warn "Ignore keyword argument 'progress_bar' for SteadyStateODESolver"
     haskey(kwargs, :save_everystep) && @warn "Ignore keyword argument 'save_everystep' for SteadyStateODESolver"
     haskey(kwargs, :saveat) && @warn "Ignore keyword argument 'saveat' for SteadyStateODESolver"
     kwargs2 = merge(
         NamedTuple(kwargs), # we convert to NamedTuple just in case if kwargs is empty
-        (progress_bar = Val(false), save_everystep = false, saveat = ftype[]),
+        (progress_bar = Val(false), save_everystep = false, saveat = ftype[tmax]),
     )
 
     # add terminate condition (callback)

test/ext-test/cpu/autodiff/autodiff.jl

@@ -67,6 +67,17 @@ function my_f_mesolve(p)
     return real(expect(a' * a, sol.states[end]))
 end
 
+function my_f_steadystate(p)
+    ρss = steadystate(
+        L,
+        SteadyStateODESolver(ψ0 = ψ0_mesolve, tmax = tlist_mesolve[end]);
+        params = p,
+        sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()),
+    )
+
+    return real(expect(a' * a, ρss))
+end
+
 # Analytical solution
 n_ss(Δ, F, γ) = abs2(F / (Δ + 1im * γ / 2))
 
@@ -113,8 +124,12 @@ n_ss(Δ, F, γ) = abs2(F / (Δ + 1im * γ / 2))
 
         my_f_mesolve_direct(params)
         my_f_mesolve(params)
+        my_f_steadystate(params)
 
+        # calculate exact solution and check if steadystate works
         grad_exact = Zygote.gradient((p) -> n_ss(p[1], p[2], p[3]), params)[1]
+        grad_ss = Zygote.gradient(my_f_steadystate, params)[1]
+        @test grad_ss ≈ grad_exact atol=1e-5
 
         @testset "ForwardDiff.jl" begin
             grad_qt = ForwardDiff.gradient(my_f_mesolve_direct, params)