How to get the gradient (zygote) calculated by DynamicalSystems

Question

How to get the gradient (zygote) calculated by DynamicalSystems

chooron opened this issue 2 months ago · comments

Describe the bug
Hello, I want to use DynamicalSystems.jl to calculate the gradient, but the following problem occurs, which may be related to mutate array

Minimal Working Example

using DynamicalSystems
using Zygote
using StaticArrays

function f2(pp)
    function henon_rule(u, p, n) # here `n` is "time", but we don't use it.
        x, y = u # system state
        a, b = p # system parameters
        xn = 1.0 - a * x^2 + y
        yn = b * x
        return SVector(xn, yn)
        # return SA[xn, yn]
        # return SVector(1, 2, 3)
        # return @SVector [1, 2, 3]
        # return SVector{2}([xn,yn])
    end

    u0 = [0.2, 0.3]
    p0 = [1.4, 0.3]

    henon = DeterministicIteratedMap(henon_rule, u0, pp)

    total_time = 10_000
    X, t = trajectory(henon, total_time)
    sum(sum(X))
end

gradient(f2, [0.7, 0.6])

output:

ERROR: Mutating arrays is not supported -- called setindex!(Vector{SVector{2, Float64}}, ...)
This error occurs when you ask Zygote to differentiate operations that change
the elements of arrays in place (e.g. setting values with x .= ...)

Possible fixes:
- avoid mutating operations (preferred)
- or read the documentation and solutions for this error
  https://fluxml.ai/Zygote.jl/latest/limitations

Stacktrace:
  [1] error(s::String)
    @ Base .\error.jl:35
  [2] _throw_mutation_error(f::Function, args::Vector{SVector{2, Float64}})
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\lib\array.jl:70
  [3] (::Zygote.var"#539#540"{Vector{SVector{2, Float64}}})(::Nothing)
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\lib\array.jl:82
  [4] (::Zygote.var"#2623#back#541"{Zygote.var"#539#540"{Vector{SVector{2, Float64}}}})(Δ::Nothing)
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\ZygoteRules\M4xmc\src\adjoint.jl:72
  [5] #trajectory_discrete#6
    @ D:\Julia\Julia-1.10.4\packages\packages\DynamicalSystemsBase\zFDur\src\core\trajectory.jl:58 [inlined]
  [6] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Tuple{@NamedTuple{…}, Nothing})
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
  [7] trajectory_discrete
    @ D:\Julia\Julia-1.10.4\packages\packages\DynamicalSystemsBase\zFDur\src\core\trajectory.jl:45 [inlined]
  [8] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Tuple{@NamedTuple{…}, Nothing})
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
  [9] #trajectory#5
    @ D:\Julia\Julia-1.10.4\packages\packages\DynamicalSystemsBase\zFDur\src\core\trajectory.jl:39 [inlined]
 [10] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Tuple{@NamedTuple{…}, Nothing})
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
 [11] trajectory
    @ D:\Julia\Julia-1.10.4\packages\packages\DynamicalSystemsBase\zFDur\src\core\trajectory.jl:33 [inlined]
 [12] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Tuple{@NamedTuple{…}, Nothing})
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
 [13] trajectory
    @ D:\Julia\Julia-1.10.4\packages\packages\DynamicalSystemsBase\zFDur\src\core\trajectory.jl:33 [inlined]
 [14] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Tuple{@NamedTuple{…}, Nothing})
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
 [15] f2
    @ e:\JlCode\LumpedHydro\temp\test_ds.jl:83 [inlined]
 [16] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64)
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
 [17] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface.jl:91
 [18] gradient(f::Function, args::Vector{Float64})
    @ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface.jl:148
 [19] top-level scope
    @ e:\JlCode\LumpedHydro\temp\test_ds.jl:87
Some type information was truncated. Use `show(err)` to see complete types.

Package versions

Please provide the versions you use. To do this, run the code:

using Pkg
Pkg.status([
    "DynamicalSystems",
    "StateSpaceSets", "DynamicalSystemsBase", "RecurrenceAnalysis", "FractalDimensions", "DelayEmbeddings", "ComplexityMeasures", "TimeseriesSurrogates", "PredefinedDynamicalSystems", "Attractors", "ChaosTools"
    ];
    mode = PKGMODE_MANIFEST
)

  [61744808] DynamicalSystems v3.3.17
  [90137ffa] StaticArrays v1.9.7
  [e88e6eb3] Zygote v0.6.70

George Datseris · Answer 1 · Mon Jul 29 2024 22:32:42 GMT+0800 (China Standard Time)

Oh, I wish I knew anything about automagic differentiation to help you with this... But I've never used Zygote.jl so I have no idea what it does. I think it is worth posting a link to this issue to the Julia slack for #diffeq-bridged. It is likely that someone in this channel will have used DynamicalSYstems.jl to do some sort of autodiff stuff.

jing xin · Answer 2 · Mon Jul 29 2024 22:41:09 GMT+0800 (China Standard Time)

My goal is to use zygote to solve the gradient and use it for parameter optimization of the problem, and this problem involves DynamicalSystems, which is similar to partial differential equations. Thank you for your reply. I hope that this package will support gradient solving in the future.

Christopher Rackauckas · Answer 3 · Mon Jul 29 2024 23:50:30 GMT+0800 (China Standard Time)

Did you give Enzyme a try? The issue here is that trajectory internally mutates, which is something that would have to be solved at the DynamicalSystems.jl level. Or just not use Zygote.

George Datseris · Answer 4 · Tue Jul 30 2024 00:54:10 GMT+0800 (China Standard Time)

Thanks Chris!

Besides trying Enzyme, can you also tell us which function from Dynamical Systems you are interested in? I doubt it is trajectory, this is just a convenience function. Maybe the function you care about can become non mutating or is already.

jing xin · Answer 5 · Tue Jul 30 2024 09:06:10 GMT+0800 (China Standard Time)

My issue involves calculating flood progression using the Muskingum-Cunge method in hydrological science, incorporating two variables: outflow, inflow, outflow(t+1) = f(outflow(t), inflow(t+1), inflow(t)), outflow[0] = inflow[0], where inflow is a known variable. The code implementation is as follows:

inflow = Float64[2, 2, 3, 5, 1, 2]
itp = LinearInterpolation(inflow, 1:length(inflow), extrapolate=true)

function msk(inflow, p)
    c0, c1, c2 = p
    outflow = zeros(length(inflow))
    outflow[1] = inflow[1]
    for i in 1:length(outflow)-1
        outflow[i+1] = (c0 * inflow[i+1]) + (c1 * inflow[i]) + (c2 * outflow[i])
    end
    return outflow
end

k, x, dt = 0.5, 0.2, 3.0
c0 = ((dt / k) - (2 * x)) / ((2 * (1 - x)) + (dt / k))
c1 = ((dt / k) + (2 * x)) / ((2 * (1 - x)) + (dt / k))
c2 = ((2 * (1 - x)) - (dt / k)) / ((2 * (1 - x)) + (dt / k))
outflow = msk(inflow, [c0,c1,c2])

Due to this code involving array mutation, it cannot compute Zygote.gradient. Therefore, I thought DynamicSystems.jl might support Zygote.gradient computation like DifferentialEquations.jl, so I converted it into a Dynamic Systems problem:

function ds(u, p, t)
    q = u[1]
    c0, c1, c2 = p
    qn = c0 * itp(t + 1) + c1 * itp(t) + c2 * q
    SVector(qn)
end

u0 = inflow[1]
prob = DeterministicIteratedMap(ds, [u0], p)
total_time = length(inflow)
X, t = trajectory(prob, total_time)

Although the computed results are consistent, functions built using DynamicSystems still face the array mutation limitation:

function f1(p)
    prob = DeterministicIteratedMap(ds, [u0], p)
    total_time = length(inflow)
    X, t = trajectory(prob, total_time)
    sum(sum(X))
end

gradient(f1, [0.7, 0.6, -0.5])

Indeed, Enzyme.jl is also a good autodiff package that supports mutating arrays, but the Optimization.jl package I'm using has only tried AutoZygote and AutoForwarddiff, and AutoEnzyme may need further exploration.

In fact, the Muskingum method is derived from simplifying an ODE problem, so perhaps I can directly construct an ODE problem for solution.

Finally, thank you for your thoughtful and reliable responses.

George Datseris · Answer 6 · Wed Jul 31 2024 04:45:46 GMT+0800 (China Standard Time)

But if DifferentialEquations.jl supports Zygote.gradient, why don't you use the discrete map version of DifferentialEquations.jl? (formally called a discrete problem https://docs.sciml.ai/DiffEqDocs/stable/types/discrete_types/ )

If the only thing you need is to integrate forwards in time a dynamical system, my advice is to use DiffEq directly; DynamicalSystems.jl is primarily for doing some analysis that is happening either on top, or after, evolving the system.

p.s.: I admit I don't understand how DiffEq can support Zygote.gradient; in the end of the day we have an ODEIntegrator, and this object constantly mutates its own state. So why isn't that this is not a "mutability problem" for the gradient estimation? In the sense that it is a problem to iteratively create a trajectory because you mutate a vector of vectors, but mutating the ODE integrator itself is not a problem? I think autodiff stuff are currently beyond my knowledge level :D

jing xin · Answer 7 · Wed Jul 31 2024 09:15:02 GMT+0800 (China Standard Time)

But if DifferentialEquations.jl supports Zygote.gradient, why don't you use the discrete map version of DifferentialEquations.jl? (formally called a discrete problem https://docs.sciml.ai/DiffEqDocs/stable/types/discrete_types/ )

If the only thing you need is to integrate forwards in time a dynamical system, my advice is to use DiffEq directly; DynamicalSystems.jl is primarily for doing some analysis that is happening either on top, or after, evolving the system.

Yes, it seems better to construct a discrete problem, I will try to convert this problem

p.s.: I admit I don't understand how DiffEq can support Zygote.gradient; in the end of the day we have an ODEIntegrator, and this object constantly mutates its own state. So why isn't that this is not a "mutability problem" for the gradient estimation? In the sense that it is a problem to iteratively create a trajectory because you mutate a vector of vectors, but mutating the ODE integrator itself is not a problem? I think autodiff stuff are currently beyond my knowledge level :D

Indeed, DiffEq and DynamicalSystems are very similar in writing, but DiffEq can support Zygote.gradient. I think it may be related to the SciMLSensitivity.jl package, which provides support for some autodiff libraries.

Finally, thank you for your reply. I will look forward to the support of gradient calculation in this package in the future, because adding a neural network to DynamicalSystems seems to be a very interesting idea.

George Datseris · Answer 8 · Wed Jul 31 2024 14:43:55 GMT+0800 (China Standard Time)

Okay, let's keep this issue open for future reference. I don't do autodiff so I can't solve it, but maybe someone in the future will contribute a solution!