Merge pull request #2872 from SciML/enzyme_discrete

Add tests for enzyme discrete adjoints

Author: ChrisRackauckas

.github/workflows/CI.yml

@@ -61,6 +61,7 @@ jobs:
 
           - ImplicitDiscreteSolve
           - SimpleImplicitDiscreteSolve
+          - Enzyme
 
         version:
           - 'lts'

test/enzyme/discrete_adjoints.jl

@@ -0,0 +1,51 @@
+using Enzyme, OrdinaryDiffEqTsit5, StaticArrays, DiffEqBase, ForwardDiff, Test
+
+function lorenz!(du, u, p, t)
+    du[1] = 10.0(u[2] - u[1])
+    du[2] = u[1] * (28.0 - u[3]) - u[2]
+    du[3] = u[1] * u[2] - (8 / 3) * u[3]
+end
+
+const _saveat =  SA[0.0,0.25,0.5,0.75,1.0,1.25,1.5,1.75,2.0,2.25,2.5,2.75,3.0]
+
+function f_dt(y::Array{Float64}, u0::Array{Float64})
+    tspan = (0.0, 3.0)
+    prob = ODEProblem{true, SciMLBase.FullSpecialize}(lorenz!, u0, tspan)
+    sol = DiffEqBase.solve(prob, Tsit5(), saveat = _saveat, sensealg = DiffEqBase.SensitivityADPassThrough(), abstol=1e-12, reltol=1e-12)
+    y .= sol[1,:]
+    return nothing
+end;
+
+function f_dt(u0)
+    tspan = (0.0, 3.0)
+    prob = ODEProblem{true, SciMLBase.FullSpecialize}(lorenz!, u0, tspan)
+    sol = DiffEqBase.solve(prob, Tsit5(), saveat = _saveat, sensealg = DiffEqBase.SensitivityADPassThrough(), abstol=1e-12, reltol=1e-12)
+    sol[1,:]
+end;
+
+u0 = [1.0; 0.0; 0.0]
+fdj = ForwardDiff.jacobian(f_dt, u0)
+
+ezj = stack(map(1:3) do i
+    d_u0 = zeros(3)
+    dy = zeros(13)
+    y  = zeros(13)
+    d_u0[i] = 1.0
+    Enzyme.autodiff(Forward, f_dt,  Duplicated(y, dy), Duplicated(u0, d_u0));
+    dy
+end)
+
+@test ezj ≈ fdj
+
+function f_dt2(u0)
+    tspan = (0.0, 3.0)
+    prob = ODEProblem{true, SciMLBase.FullSpecialize}(lorenz!, u0, tspan)
+    sol = DiffEqBase.solve(prob, Tsit5(), dt=0.1, saveat = _saveat, sensealg = DiffEqBase.SensitivityADPassThrough(), abstol=1e-12, reltol=1e-12)
+    sum(sol[1,:])
+end
+
+fdg = ForwardDiff.gradient(f_dt2, u0)
+d_u0 = zeros(3)
+Enzyme.autodiff(Reverse, f_dt2,  Active, Duplicated(u0, d_u0));
+
+@test d_u0 ≈ fdg

test/runtests.jl

@@ -179,6 +179,7 @@ end
     if !is_APPVEYOR && GROUP == "Enzyme" && isempty(VERSION.prerelease)
         activate_enzyme_env()
         @time @safetestset "Autodiff Events Tests" include("enzyme/autodiff_events.jl")
+        @time @safetestset "Discrete Adjoint Tests" include("enzyme/discrete_adjoints.jl")
     end
 
     # Don't run ODEInterface tests on prerelease