sensitivity solution and downstream is passing

Author: ChrisRackauckas

src/solutions/ode_solutions.jl

@@ -100,22 +100,42 @@ function solution_new_retcode(sol::AbstractODESolution{T,N},retcode) where {T,N}
                      sol.alg,sol.interp,sol.dense,sol.tslocation,sol.destats,retcode)
  end
 
- function solution_new_tslocation(sol::AbstractODESolution{T,N},tslocation) where {T,N}
-   ODESolution{T,N,typeof(sol.u),typeof(sol.u_analytic),typeof(sol.errors),
-                      typeof(sol.t),typeof(sol.k),
-                      typeof(sol.prob),typeof(sol.alg),typeof(sol.interp),typeof(sol.destats)}(
-                      sol.u,sol.u_analytic,sol.errors,sol.t,sol.k,sol.prob,
-                      sol.alg,sol.interp,sol.dense,tslocation,sol.destats,sol.retcode)
+function solution_new_tslocation(sol::AbstractODESolution{T,N},tslocation) where {T,N}
+  ODESolution{T,N,typeof(sol.u),typeof(sol.u_analytic),typeof(sol.errors),
+                    typeof(sol.t),typeof(sol.k),
+                    typeof(sol.prob),typeof(sol.alg),typeof(sol.interp),typeof(sol.destats)}(
+                    sol.u,sol.u_analytic,sol.errors,sol.t,sol.k,sol.prob,
+                    sol.alg,sol.interp,sol.dense,tslocation,sol.destats,sol.retcode)
+end
+
+function solution_slice(sol::AbstractODESolution{T,N},I) where {T,N}
+  ODESolution{T,N,typeof(sol.u),typeof(sol.u_analytic),typeof(sol.errors),
+                     typeof(sol.t),typeof(sol.k),
+                     typeof(sol.prob),typeof(sol.alg),typeof(sol.interp),typeof(sol.destats)}(
+                     sol.u[I],
+                     sol.u_analytic === nothing ? nothing : sol.u_analytic[I],
+                     sol.errors,sol.t[I],
+                     sol.dense ? sol.k[I] : sol.k,
+                     sol.prob,
+                     sol.alg,sol.interp,false,sol.tslocation,sol.destats,sol.retcode)
+ end
+
+function sensitivity_solution(sol::AbstractODESolution,u,t)
+  T = eltype(eltype(u))
+  N = length((size(sol.prob.u0)..., length(u)))
+  interp = if typeof(sol.interp) <: LinearInterpolation
+    LinearInterpolation(t,u)
+  elseif typeof(sol.interp) <: ConstantInterpolation
+    ConstantInterpolation(t,u)
+  else
+    SensitivityInterpolation(t,u)
   end
 
-  function solution_slice(sol::AbstractODESolution{T,N},I) where {T,N}
-    ODESolution{T,N,typeof(sol.u),typeof(sol.u_analytic),typeof(sol.errors),
-                       typeof(sol.t),typeof(sol.k),
-                       typeof(sol.prob),typeof(sol.alg),typeof(sol.interp),typeof(sol.destats)}(
-                       sol.u[I],
-                       sol.u_analytic === nothing ? nothing : sol.u_analytic[I],
-                       sol.errors,sol.t[I],
-                       sol.dense ? sol.k[I] : sol.k,
-                       sol.prob,
-                       sol.alg,sol.interp,false,sol.tslocation,sol.destats,sol.retcode)
-   end
+  ODESolution{T,N,typeof(u),typeof(sol.u_analytic),typeof(sol.errors),
+              typeof(t),Nothing,typeof(sol.prob),typeof(sol.alg),
+              typeof(interp),typeof(sol.destats)}(
+              u,sol.u_analytic,sol.errors,t,nothing,sol.prob,
+              sol.alg,interp,
+              sol.dense,sol.tslocation,
+              sol.destats,sol.retcode)
+end

src/solutions/rode_solutions.jl

@@ -120,3 +120,24 @@ function solution_slice(sol::AbstractRODESolution{T,N},I) where {T,N}
                false,sol.tslocation,sol.destats,
                sol.retcode,sol.seed)
 end
+
+function sensitivity_solution(sol::AbstractRODESolution,u,t)
+  T = eltype(eltype(u))
+  N = length((size(sol.prob.u0)..., length(u)))
+  interp = if typeof(sol.interp) <: LinearInterpolation
+    LinearInterpolation(t,u)
+  elseif typeof(sol.interp) <: ConstantInterpolation
+    ConstantInterpolation(t,u)
+  else
+    SensitivityInterpolation(t,u)
+  end
+
+  RODESolution{T,N,typeof(u),typeof(sol.u_analytic),
+               typeof(sol.errors),typeof(t),
+               typeof(nothing),typeof(sol.prob),typeof(sol.alg),
+               typeof(sol.interp),typeof(sol.destats)}(
+               u,sol.u_analytic,sol.errors,t,nothing,sol.prob,
+               sol.alg,sol.interp,
+               sol.dense,sol.tslocation,sol.destats,
+               sol.retcode,sol.seed)
+end

src/solutions/steady_state_solutions.jl

@@ -18,3 +18,12 @@ function build_solution(prob::AbstractSteadyStateProblem,
 
   SteadyStateSolution{T,N,typeof(u),typeof(resid),typeof(prob),typeof(alg)}(u,resid,prob,alg,retcode)
 end
+
+function sensitivity_solution(sol::AbstractSteadyStateSolution,u)
+  T = eltype(eltype(u))
+  N = length((size(sol.prob.u0)...,))
+
+  SteadyStateSolution{T,N,typeof(u),typeof(sol.resid),
+                      typeof(sol.prob),typeof(sol.alg)}(
+                      u,sol.resid,sol.prob,sol.alg,sol.retcode)
+end

src/solve.jl

@@ -61,8 +61,9 @@ function solve_call(_prob,args...;merge_callbacks = true, kwargs...)
   end
 end
 
+# save_idxs and saveat are here due to https://github.com/FluxML/Zygote.jl/issues/664
 function solve(prob::DEProblem,args...;sensealg=nothing,
-               u0 = nothing, p = nothing,kwargs...)
+               u0 = nothing, p = nothing, kwargs...)
   u0 = u0 !== nothing ? u0 : prob.u0
   p  = p  !== nothing ? p  : prob.p
   solve_up(prob,sensealg,u0,p,args...;kwargs...)

src/zygote.jl

@@ -73,3 +73,8 @@ end
 
 ZygoteRules.@adjoint numargs(f) = (numargs(f),df->(nothing,))
 ChainRulesCore.rrule(::typeof(numargs),f) = (numargs(f),df->(nothing,))
+
+# Until https://github.com/FluxML/Zygote.jl/issues/664 is fixed
+ZygoteRules.@adjoint function Base.pairs(x::NamedTuple)
+  Base.pairs(x), Δ -> (Δ.data,)
+end