|
| 1 | +# taken from https://github.com/SciML/SciMLStructures.jl/pull/28 |
| 2 | +using OrdinaryDiffEq, SciMLSensitivity, Zygote |
| 3 | +using LinearAlgebra |
| 4 | +import SciMLStructures as SS |
| 5 | + |
| 6 | +mutable struct SubproblemParameters{P, Q, R} |
| 7 | + p::P # tunable |
| 8 | + q::Q |
| 9 | + r::R |
| 10 | +end |
| 11 | +mutable struct Parameters{P, C} |
| 12 | + subparams::P |
| 13 | + coeffs::C # tunable matrix |
| 14 | +end |
| 15 | +# the rhs is `du[i] = p[i] * u[i]^2 + q[i] * u[i] + r[i] * t` for i in 1:length(subparams) |
| 16 | +# and `du[length(subparams)+1:end] .= coeffs * u` |
| 17 | +function rhs!(du, u, p::Parameters, t) |
| 18 | + for (i, subpars) in enumerate(p.subparams) |
| 19 | + du[i] = subpars.p * u[i]^2 + subpars.q * u[i] + subpars.r * t |
| 20 | + end |
| 21 | + N = length(p.subparams) |
| 22 | + mul!(view(du, (N + 1):(length(du))), p.coeffs, u) |
| 23 | + return nothing |
| 24 | +end |
| 25 | +u = sin.(0.1:0.1:1.0) |
| 26 | +subparams = [SubproblemParameters(0.1i, 0.2i, 0.3i) for i in 1:5] |
| 27 | +p = Parameters(subparams, cos.([0.1i + 0.33j for i in 1:5, j in 1:10])) |
| 28 | +tspan = (0.0, 1.0) |
| 29 | +prob = ODEProblem(rhs!, u, tspan, p) |
| 30 | +solve(prob, Tsit5()) |
| 31 | + |
| 32 | +# Mark the struct as a SciMLStructure |
| 33 | +SS.isscimlstructure(::Parameters) = true |
| 34 | +# It is mutable |
| 35 | +SS.ismutablescimlstructure(::Parameters) = true |
| 36 | +# Only contains `Tunable` portion |
| 37 | +# We could also add a `Constants` portion to contain the values that are |
| 38 | +# not tunable. The implementation would be similar to this one. |
| 39 | +SS.hasportion(::SS.Tunable, ::Parameters) = true |
| 40 | +function SS.canonicalize(::SS.Tunable, p::Parameters) |
| 41 | + # concatenate all tunable values into a single vector |
| 42 | + buffer = vcat([subpar.p for subpar in p.subparams], vec(p.coeffs)) |
| 43 | + # repack takes a new vector of the same length as `buffer`, and constructs |
| 44 | + # a new `Parameters` object using the values from the new vector for tunables |
| 45 | + # and retaining old values for other parameters. This is exactly what replace does, |
| 46 | + # so we can use that instead. |
| 47 | + repack = let p = p |
| 48 | + function repack(newbuffer) |
| 49 | + SS.replace(SS.Tunable(), p, newbuffer) |
| 50 | + end |
| 51 | + end |
| 52 | + # the canonicalized vector, the repack function, and a boolean indicating |
| 53 | + # whether the buffer aliases values in the parameter object (here, it doesn't) |
| 54 | + return buffer, repack, false |
| 55 | +end |
| 56 | +function SS.replace(::SS.Tunable, p::Parameters, newbuffer) |
| 57 | + N = length(p.subparams) + length(p.coeffs) |
| 58 | + @assert length(newbuffer) == N |
| 59 | + subparams = [SubproblemParameters(newbuffer[i], subpar.q, subpar.r) |
| 60 | + for (i, subpar) in enumerate(p.subparams)] |
| 61 | + coeffs = reshape( |
| 62 | + view(newbuffer, (length(p.subparams) + 1):length(newbuffer)), size(p.coeffs)) |
| 63 | + return Parameters(subparams, coeffs) |
| 64 | +end |
| 65 | +function SS.replace!(::SS.Tunable, p::Parameters, newbuffer) |
| 66 | + N = length(p.subparams) + length(p.coeffs) |
| 67 | + @assert length(newbuffer) == N |
| 68 | + for (subpar, val) in zip(p.subparams, newbuffer) |
| 69 | + subpar.p = val |
| 70 | + end |
| 71 | + copyto!(coeffs, view(newbuffer, (length(p.subparams) + 1):length(newbuffer))) |
| 72 | + return p |
| 73 | +end |
| 74 | + |
| 75 | +Zygote.gradient(0.1ones(length(SS.canonicalize(SS.Tunable(), p)[1]))) do tunables |
| 76 | + newp = SS.replace(SS.Tunable(), p, tunables) |
| 77 | + newprob = remake(prob; p = newp) |
| 78 | + sol = solve(newprob, Tsit5()) |
| 79 | + return sum(sol.u[end]) |
| 80 | +end |
0 commit comments