|
1 | | -struct SDEDiffusionTermWrapper{iip,G,H} |
2 | | - g::G |
3 | | - h::H |
| 1 | +struct SDEDiffusionTermWrapper{iip, G, H} |
| 2 | + g::G |
| 3 | + h::H |
4 | 4 | end |
5 | | -function SDEDiffusionTermWrapper(g::G,h::H) where {G,H} |
6 | | - SDEDiffusionTermWrapper{isinplace(g,5),G,H}(g,h) |
| 5 | +function SDEDiffusionTermWrapper(g::G, h::H) where {G, H} |
| 6 | + SDEDiffusionTermWrapper{isinplace(g, 5), G, H}(g, h) |
7 | 7 | end |
8 | 8 | (g::SDEDiffusionTermWrapper{true})(du, u, p, t) = g.g(du, u, g.h, p, t) |
9 | 9 | (g::SDEDiffusionTermWrapper{false})(u, p, t) = g.g(u, g.h, p, t) |
10 | 10 |
|
11 | | -struct SDEFunctionWrapper{iip,F,G,H,TMM,Ta,Tt,TJ,JVP,VJP,JP,SP,TW,TWt,TPJ,GG,S,TCV} <: DiffEqBase.AbstractRODEFunction{iip} |
12 | | - f::F |
13 | | - g::G |
14 | | - h::H |
15 | | - mass_matrix::TMM |
16 | | - analytic::Ta |
17 | | - tgrad::Tt |
18 | | - jac::TJ |
19 | | - jvp::JVP |
20 | | - vjp::VJP |
21 | | - jac_prototype::JP |
22 | | - sparsity::SP |
23 | | - Wfact::TW |
24 | | - Wfact_t::TWt |
25 | | - paramjac::TPJ |
26 | | - ggprime::GG |
27 | | - syms::S |
28 | | - colorvec::TCV |
| 11 | +struct SDEFunctionWrapper{iip, F, G, H, TMM, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TWt, TPJ, GG, |
| 12 | + S, TCV} <: DiffEqBase.AbstractRODEFunction{iip} |
| 13 | + f::F |
| 14 | + g::G |
| 15 | + h::H |
| 16 | + mass_matrix::TMM |
| 17 | + analytic::Ta |
| 18 | + tgrad::Tt |
| 19 | + jac::TJ |
| 20 | + jvp::JVP |
| 21 | + vjp::VJP |
| 22 | + jac_prototype::JP |
| 23 | + sparsity::SP |
| 24 | + Wfact::TW |
| 25 | + Wfact_t::TWt |
| 26 | + paramjac::TPJ |
| 27 | + ggprime::GG |
| 28 | + syms::S |
| 29 | + colorvec::TCV |
29 | 30 | end |
30 | 31 |
|
31 | | - |
32 | 32 | (f::SDEFunctionWrapper{true})(du, u, p, t) = f.f(du, u, f.h, p, t) |
33 | 33 | (f::SDEFunctionWrapper{false})(u, p, t) = f.f(u, f.h, p, t) |
34 | 34 |
|
35 | 35 | function wrap_functions_and_history(f::SDDEFunction, g, h) |
36 | | - gwh = SDEDiffusionTermWrapper{isinplace(g,5),typeof(g),typeof(h)}(g,h) |
| 36 | + gwh = SDEDiffusionTermWrapper{isinplace(g, 5), typeof(g), typeof(h)}(g, h) |
37 | 37 |
|
38 | | - if f.jac === nothing |
39 | | - jac = nothing |
40 | | - else |
41 | | - if isinplace(f) |
42 | | - jac = let f_jac = f.jac, h = h |
43 | | - (J, u, p, t) -> f_jac(J, u, h, p, t) |
44 | | - end |
| 38 | + if f.jac === nothing |
| 39 | + jac = nothing |
45 | 40 | else |
46 | | - jac = let f_jac = f.jac, h = h |
47 | | - (u, p, t) -> f_jac(u, h, p, t) |
48 | | - end |
| 41 | + if isinplace(f) |
| 42 | + jac = let f_jac = f.jac, h = h |
| 43 | + (J, u, p, t) -> f_jac(J, u, h, p, t) |
| 44 | + end |
| 45 | + else |
| 46 | + jac = let f_jac = f.jac, h = h |
| 47 | + (u, p, t) -> f_jac(u, h, p, t) |
| 48 | + end |
| 49 | + end |
49 | 50 | end |
50 | | - end |
51 | 51 |
|
52 | | - SDEFunctionWrapper{isinplace(f),typeof(f.f),typeof(gwh),typeof(h),typeof(f.mass_matrix), |
53 | | - typeof(f.analytic),typeof(f.tgrad),typeof(jac),typeof(f.jvp), |
54 | | - typeof(f.vjp),typeof(f.jac_prototype),typeof(f.sparsity), |
55 | | - typeof(f.Wfact),typeof(f.Wfact_t),typeof(f.paramjac), |
56 | | - typeof(f.ggprime),typeof(f.syms),typeof(f.colorvec)}( |
57 | | - f.f, gwh, h, f.mass_matrix, f.analytic, f.tgrad, jac, |
58 | | - f.jvp, f.vjp, |
59 | | - f.jac_prototype, f.sparsity, f.Wfact, f.Wfact_t, |
60 | | - f.paramjac, f.ggprime, f.syms, f.colorvec), gwh |
| 52 | + SDEFunctionWrapper{isinplace(f), typeof(f.f), typeof(gwh), typeof(h), |
| 53 | + typeof(f.mass_matrix), |
| 54 | + typeof(f.analytic), typeof(f.tgrad), typeof(jac), typeof(f.jvp), |
| 55 | + typeof(f.vjp), typeof(f.jac_prototype), typeof(f.sparsity), |
| 56 | + typeof(f.Wfact), typeof(f.Wfact_t), typeof(f.paramjac), |
| 57 | + typeof(f.ggprime), typeof(f.syms), typeof(f.colorvec)}(f.f, gwh, h, |
| 58 | + f.mass_matrix, |
| 59 | + f.analytic, |
| 60 | + f.tgrad, jac, |
| 61 | + f.jvp, f.vjp, |
| 62 | + f.jac_prototype, |
| 63 | + f.sparsity, |
| 64 | + f.Wfact, |
| 65 | + f.Wfact_t, |
| 66 | + f.paramjac, |
| 67 | + f.ggprime, |
| 68 | + f.syms, |
| 69 | + f.colorvec), |
| 70 | + gwh |
61 | 71 | end |
0 commit comments