Merge pull request #1056 from SciML/dg/literal

chore: move `literal_getproperty` adjoint to SciMLSensitivity
SciML · May 26, 2024 · dbb5917 · dbb5917
2 parents da44890 + ac2f59a
commit dbb5917
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/src/SciMLSensitivity.jl b/src/SciMLSensitivity.jl
@@ -37,7 +37,8 @@ import SciMLBase: unwrapped_f, _unwrap_val
 import SciMLBase: AbstractOverloadingSensitivityAlgorithm, AbstractSensitivityAlgorithm,
                   AbstractForwardSensitivityAlgorithm, AbstractAdjointSensitivityAlgorithm,
                   AbstractSecondOrderSensitivityAlgorithm,
-                  AbstractShadowingSensitivityAlgorithm
+                  AbstractShadowingSensitivityAlgorithm,
+                  AbstractTimeseriesSolution
 
 include("parameters_handling.jl")
 include("sensitivity_algorithms.jl")

diff --git a/src/adjoint_common.jl b/src/adjoint_common.jl
@@ -651,3 +651,13 @@ function out_and_ts(_ts, duplicate_iterator_times, sol)
     end
     return out, ts
 end
+
+Zygote.@adjoint function Zygote.literal_getproperty(sol::AbstractTimeseriesSolution,
+        ::Val{:u})
+    function solu_adjoint(Δ)
+        zerou = zero(sol.prob.u0)
+        _Δ = @. ifelse(Δ === nothing, (zerou,), Δ)
+        (SciMLBase.build_solution(sol.prob, sol.alg, sol.t, _Δ),)
+    end
+    sol.u, solu_adjoint
+end