Allow models to specify a number type?

[ali-ramadhan]

Tried to auto-diff OceanTurb.jl with Zygote.jl but it doesn't support mutating arrays (it know how to auto-diff broadcasts but mutating arrays makes it much harder to auto-diff)

using Zygote, OceanTurb

function g(x) # Basically g(x) = x^2
    model = KPP.Model(N=1, H=1)
    model.solution.T.data .= x^2
    time_step!(model, 1, 1)
    return mean(model.solution.T.data)
end

julia> g'(10)
ERROR: Mutating arrays is not supported

I think you could use ForwardDiff.jl with mutating arrays but I couldn't really try it as it requires OceanTurb.jl to use dual numbers instead of Float64 (and OceanTurb.KPP didn't seem to have a float type kwarg).

[glwagner]

Mutating arrays --- as in, changing the values in an array? This is not supported by Zygote?

[glwagner]

Supporting a number time may not be too hard. Do you know what is needed specifically? Do we need to support dual numbers for all parameters, or do we need to allow different types for every parameter individually (so that just one parameter can be a dual number while the rest are floats?)

We do for some reason restrict KPP.Parameters to AbstractFloat right now, not sure why:

https://github.com/glwagner/OceanTurb.jl/blob/b0d21cddaaf9aa625146108b5ec96a2a544a2f9f/src/models/KPP.jl#L18

Luckily though we currently support typing the boundary conditions and constants.

[glwagner]

For both the ForwardEuler and BackwardEuler time-steppers we use the solution to determine the type of rhs:

https://github.com/glwagner/OceanTurb.jl/blob/b0d21cddaaf9aa625146108b5ec96a2a544a2f9f/src/timesteppers.jl#L173-L183

Though, is set!(fld, 0) an issue if fld is an array of dual numbers?