OrdinaryDiffEq.jl/lib/ImplicitDiscreteSolve/src/cache.jl at 4b597fb82a9ea3469049ec41c348d83581c36640 · SciML/OrdinaryDiffEq.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
struct ImplicitDiscreteState{uType, pType, tType}
    u::uType
    p::pType
    t::tType
end

mutable struct IDSolveCache{uType, cType, thetaType} <: OrdinaryDiffEqMutableCache
    u::uType
    uprev::uType
    z::uType
    nlcache::cType
    Θks::thetaType
end

function alg_cache(alg::IDSolve, u, rate_prototype, ::Type{uEltypeNoUnits},
        ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
        dt, reltol, p, calck,
        ::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
    state = ImplicitDiscreteState(isnothing(u) ? nothing : zero(u), p, t)
    f_nl = (resid, u_next, p) -> f(resid, u_next, p.u, p.p, p.t)

    u_len = isnothing(u) ? 0 : length(u)
    nlls = !isnothing(f.resid_prototype) && (length(f.resid_prototype) != u_len)
    unl = isnothing(u) ? Float64[] : u # FIXME nonlinear solve cannot handle nothing for u
    prob = if nlls
        NonlinearLeastSquaresProblem{isinplace(f)}(
            NonlinearFunction(f_nl; resid_prototype = f.resid_prototype),
            unl, state)
    else
        NonlinearProblem{isinplace(f)}(f_nl, unl, state)
    end

    nlcache = init(prob, alg.nlsolve)

    IDSolveCache(u, uprev, state.u, nlcache, uBottomEltypeNoUnits[])
end

isdiscretecache(cache::IDSolveCache) = true

function alg_cache(alg::IDSolve, u, rate_prototype, ::Type{uEltypeNoUnits},
        ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
        dt, reltol, p, calck,
        ::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
    @assert !isnothing(u) "Empty u not supported with out of place functions yet."

    state = ImplicitDiscreteState(isnothing(u) ? nothing : zero(u), p, t)
    f_nl = (u_next, p) -> f(u_next, p.u, p.p, p.t)

    u_len = isnothing(u) ? 0 : length(u)
    nlls = !isnothing(f.resid_prototype) && (length(f.resid_prototype) != u_len)
    unl = isnothing(u) ? Float64[] : u # FIXME nonlinear solve cannot handle nothing for u
    prob = if nlls
        NonlinearLeastSquaresProblem{isinplace(f)}(
            NonlinearFunction(f_nl; resid_prototype = f.resid_prototype),
            unl, state)
    else
        NonlinearProblem{isinplace(f)}(f_nl, unl, state)
    end

    nlcache = init(prob, alg.nlsolve)

    # FIXME Use IDSolveConstantCache?
    IDSolveCache(u, uprev, state.u, nlcache, uBottomEltypeNoUnits[])
end

get_fsalfirstlast(cache::IDSolveCache, rate_prototype) = (nothing, nothing)