feat: add adjoint for remake_buffer

Author: AayushSabharwal

ext/MTKChainRulesCoreExt.jl

@@ -2,14 +2,102 @@ module MTKChainRulesCoreExt
 
 import ModelingToolkit as MTK
 import ChainRulesCore
-import ChainRulesCore: NoTangent
+import ChainRulesCore: Tangent, ZeroTangent, NoTangent, zero_tangent, unthunk
 
 function ChainRulesCore.rrule(::Type{MTK.MTKParameters}, tunables, args...)
     function mtp_pullback(dt)
+        dt = unthunk(dt)
         (NoTangent(), dt.tunable[1:length(tunables)],
             ntuple(_ -> NoTangent(), length(args))...)
     end
     MTK.MTKParameters(tunables, args...), mtp_pullback
 end
 
+notangent_or_else(::NoTangent, _, x) = x
+notangent_or_else(_, x, _) = x
+notangent_fallback(x, y) = notangent_or_else(x, x, y)
+reduce_to_notangent(x, y) = notangent_or_else(x, y, x)
+
+function subset_idxs(idxs, portion, template)
+    ntuple(Val(length(template))) do subi
+        [Base.tail(idx.idx) for idx in idxs if idx.portion == portion && idx.idx[1] == subi]
+    end
+end
+
+selected_tangents(::NoTangent, _) = ()
+selected_tangents(::ZeroTangent, _) = ZeroTangent()
+function selected_tangents(
+        tangents::AbstractArray{T}, idxs::Vector{Tuple{Int}}) where {T <: Number}
+    selected_tangents(tangents, map(only, idxs))
+end
+function selected_tangents(tangents::AbstractArray{T}, idxs...) where {T <: Number}
+    newtangents = copy(tangents)
+    view(newtangents, idxs...) .= zero(T)
+    newtangents
+end
+function selected_tangents(
+        tangents::AbstractVector{T}, idxs) where {S <: Number, T <: AbstractArray{S}}
+    newtangents = copy(tangents)
+    for i in idxs
+        j, k... = i
+        if k == ()
+            newtangents[j] = zero(newtangents[j])
+        else
+            newtangents[j] = selected_tangents(newtangents[j], k...)
+        end
+    end
+    newtangents
+end
+function selected_tangents(tangents::AbstractVector{T}, idxs) where {T <: AbstractArray}
+    newtangents = similar(tangents, Union{T, NoTangent})
+    copyto!(newtangents, tangents)
+    for i in idxs
+        j, k... = i
+        if k == ()
+            newtangents[j] = NoTangent()
+        else
+            newtangents[j] = selected_tangents(newtangents[j], k...)
+        end
+    end
+    newtangents
+end
+function selected_tangents(
+        tangents::Union{Tangent{<:Tuple}, Tangent{T, <:Tuple}}, idxs) where {T}
+    ntuple(Val(length(tangents))) do i
+        selected_tangents(tangents[i], idxs[i])
+    end
+end
+
+function ChainRulesCore.rrule(
+        ::typeof(MTK.remake_buffer), indp, oldbuf::MTK.MTKParameters, idxs, vals)
+    if idxs isa AbstractSet
+        idxs = collect(idxs)
+    end
+    idxs = map(idxs) do i
+        i isa MTK.ParameterIndex ? i : MTK.parameter_index(indp, i)
+    end
+    newbuf = MTK.remake_buffer(indp, oldbuf, idxs, vals)
+    tunable_idxs = reduce(
+        vcat, (idx.idx for idx in idxs if idx.portion isa MTK.SciMLStructures.Tunable))
+    disc_idxs = subset_idxs(idxs, MTK.SciMLStructures.Discrete(), oldbuf.discrete)
+    const_idxs = subset_idxs(idxs, MTK.SciMLStructures.Constants(), oldbuf.constant)
+    nn_idxs = subset_idxs(idxs, MTK.NONNUMERIC_PORTION, oldbuf.nonnumeric)
+
+    function remake_buffer_pullback(buf′)
+        buf′ = unthunk(buf′)
+        f′ = NoTangent()
+        indp′ = NoTangent()
+
+        tunable = selected_tangents(buf′.tunable, tunable_idxs)
+        discrete = selected_tangents(buf′.discrete, disc_idxs)
+        constant = selected_tangents(buf′.constant, const_idxs)
+        nonnumeric = selected_tangents(buf′.nonnumeric, nn_idxs)
+        oldbuf′ = Tangent{typeof(oldbuf)}(; tunable, discrete, constant, nonnumeric)
+        idxs′ = NoTangent()
+        vals′ = map(i -> MTK._ducktyped_parameter_values(buf′, i), idxs)
+        return f′, indp′, oldbuf′, idxs′, vals′
+    end
+    newbuf, remake_buffer_pullback
+end
+
 end

src/systems/parameter_buffer.jl

@@ -436,8 +436,12 @@ function validate_parameter_type(ic::IndexCache, p, idx::ParameterIndex, val)
 end
 
 function validate_parameter_type(ic::IndexCache, idx::ParameterIndex, val)
+    stype = get_buffer_template(ic, idx).type
+    if idx.portion == SciMLStructures.Tunable() && !(idx.idx isa Int)
+        stype = AbstractArray{<:stype}
+    end
     validate_parameter_type(
-        ic, get_buffer_template(ic, idx).type, Symbolics.Unknown(), nothing, idx, val)
+        ic, stype, Symbolics.Unknown(), nothing, idx, val)
 end
 
 function validate_parameter_type(ic::IndexCache, stype, sz, sym, index, val)
@@ -489,6 +493,9 @@ function indp_to_system(indp)
 end
 
 function SymbolicIndexingInterface.remake_buffer(indp, oldbuf::MTKParameters, idxs, vals)
+    _remake_buffer(indp, oldbuf, idxs, vals)
+end
+function _remake_buffer(indp, oldbuf::MTKParameters, idxs, vals; validate = true)
     newbuf = @set oldbuf.tunable = similar(oldbuf.tunable, Any)
     @set! newbuf.discrete = Tuple(similar(buf, Any) for buf in newbuf.discrete)
     @set! newbuf.constant = Tuple(similar(buf, Any) for buf in newbuf.constant)

test/extensions/Project.toml

@@ -1,5 +1,7 @@
 [deps]
 BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"

test/extensions/ad.jl

@@ -6,6 +6,9 @@ using SciMLStructures
 using OrdinaryDiffEq
 using SciMLSensitivity
 using ForwardDiff
+using ChainRulesCore
+using ChainRulesCore: NoTangent
+using ChainRulesTestUtils: test_rrule
 
 @variables x(t)[1:3] y(t)
 @parameters p[1:3, 1:3] q
@@ -51,3 +54,46 @@ end
 
     @test ForwardDiff.gradient(x_at_0, [0.3, 0.7]) == zeros(2)
 end
+
+@parameters a b[1:3] c(t) d::Integer e[1:3] f[1:3, 1:3]::Int g::Vector{AbstractFloat} h::String
+@named sys = ODESystem(
+    Equation[], t, [], [a, b, c, d, e, f, g, h],
+    continuous_events = [[a ~ 0] => [c ~ 0]])
+sys = complete(sys)
+
+ivs = Dict(c => 3a, b => ones(3), a => 1.0, d => 4, e => [5.0, 6.0, 7.0],
+    f => ones(Int, 3, 3), g => [0.1, 0.2, 0.3], h => "foo")
+
+ps = MTKParameters(sys, ivs)
+
+varmap = Dict(a => 1.0f0, b => 3ones(Float32, 3), c => 2.0,
+    e => Float32[0.4, 0.5, 0.6], g => ones(Float32, 4))
+get_values = getp(sys, [a, b..., c, e...])
+get_g = getp(sys, g)
+for (_idxs, vals) in [
+    # all portions
+    (collect(keys(varmap)), collect(values(varmap))),
+    # non-arrays
+    (keys(varmap), values(varmap)),
+    # tunable only
+    ([a], [varmap[a]]),
+    ([a, b], (varmap[a], varmap[b])),
+    ([a, b[2]], (varmap[a], varmap[b][2]))
+]
+    for idxs in [_idxs, map(i -> parameter_index(sys, i), collect(_idxs))]
+        loss = function (p)
+            newps = remake_buffer(sys, ps, idxs, p)
+            return sum(get_values(newps)) + sum(get_g(newps))
+        end
+
+        grad = Zygote.gradient(loss, vals)[1]
+        for (val, g) in zip(vals, grad)
+            @test eltype(val) == eltype(g)
+            if val isa Number
+                @test isone(g)
+            else
+                @test all(isone, g)
+            end
+        end
+    end
+end