Skip to content

refactor: Convert *Problem constructors to accept SciMLProblemOptions #4901

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AayushSabharwal wants to merge 1 commit into
base: master
Choose a base branch
from
Open

refactor: Convert *Problem constructors to accept SciMLProblemOptions #4901

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 44 additions & 6 deletions lib/ModelingToolkitBase/src/problems/bvproblem.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,42 @@
expression = Val{false}, guesses = Dict(), callback = nothing,
kwargs...
) where {iip, spec}
fn_opts = SciMLFunctionOptions(;
t = tspan !== nothing ? tspan[1] : tspan, eval_expression, eval_module,
checkbounds, check_compatibility, expression, kwargs...
)
opts = SciMLProblemOptions(
sys;
fn_opts, guesses, time_dependent_init = false,
build_initializeprob = supports_initialization(sys),
circular_dependency_max_cycle_length = length(all_symbols(sys)),
kwargs...
)
return BVProblem{iip, spec}(sys, op, tspan, opts; guesses, callback, kwargs...)
end

"""
SciMLBase.BVProblem{iip, spec}(sys::System, op, tspan, opts::SciMLProblemOptions; guesses = Dict(), callback = nothing, kwargs...)

Public entry point that builds a `BVProblem` directly from a pre-assembled
[`SciMLProblemOptions`](@ref), bypassing the `kwargs...` wrapper above.

`opts.fn_opts.check_compatibility` governs the `BVProblem`-level compatibility check. The
inner `ODEFunction` build always uses `check_compatibility = false` regardless (it must
skip its own, unrelated ODE-compatibility check) and derives `t` from `tspan[1]` when
`opts.fn_opts.t` isn't already set — both applied via `ConstructionBase.setproperties` on
a copy of `opts`, so a caller who builds `opts` directly (without going through the
keyword wrapper above) gets the same defaulting. `guesses` is kept separate from
`opts.guesses`: the latter is normalized (into a `SymmapT`) for `process_SciMLProblem`'s
own use, while this method needs the raw, unnormalized value to merge into `_op` exactly
as before.
"""
function SciMLBase.BVProblem{iip, spec}(
sys::System, op, tspan, opts::SciMLProblemOptions{E};
guesses = Dict(), callback = nothing, kwargs...
) where {iip, spec, E}
check_complete(sys, BVProblem)
check_compatibility && check_compatible_system(BVProblem, sys)
opts.fn_opts.check_compatibility && check_compatible_system(BVProblem, sys)
isnothing(callback) || error("BVP solvers do not support callbacks.")

_iip = resolve_iip(iip, op)
Expand All @@ -16,13 +50,17 @@
# for initialization.
_op = has_alg_eqs(sys) ? op : merge(Dict(op), Dict(guesses))

inner_opts = maybe_derive_t_from_tspan(opts, tspan)
inner_opts = setproperties(
inner_opts; fn_opts = setproperties(inner_opts.fn_opts; check_compatibility = false)
)
fode, u0,
p = process_SciMLProblem(
ODEFunction{_iip, spec}, sys, _op; guesses,
t = tspan !== nothing ? tspan[1] : tspan, check_compatibility = false,
checkbounds, time_dependent_init = false, expression, kwargs...
ODEFunction{_iip, spec}, sys, _op, inner_opts; options_struct = Val(true), kwargs...
)

(; eval_expression, eval_module) = opts.fn_opts.codegen
checkbounds = opts.fn_opts.codegen.codegen.checkbounds
fcost = generate_bvp_cost(
sys,
GeneratedFunctionOptions(;
Expand Down Expand Up @@ -52,10 +90,10 @@
@warn "The BVProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by op) exceeds the total number of states. The BVP solvers will default to doing a nonlinear least-squares optimization."
end

kwargs = process_kwargs(sys; expression, tspan, kwargs...)
kwargs = process_kwargs(sys; expression = Val{E}, tspan, kwargs...)
args = (; bvpfn, u0, tspan, p)

return maybe_codegen_scimlproblem(expression, BVProblem{_iip}, args; kwargs...)
return maybe_codegen_scimlproblem(Val{E}, BVProblem{_iip}, args; kwargs...)
end

function check_compatible_system(T::Type{BVProblem}, sys::System)
Expand Down
38 changes: 32 additions & 6 deletions lib/ModelingToolkitBase/src/problems/daeproblem.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,20 +86,46 @@ end
eval_module = @__MODULE__, check_compatibility = true,
expression = Val{false}, kwargs...
) where {iip, spec}
fn_opts = SciMLFunctionOptions(;
t = tspan !== nothing ? tspan[1] : tspan, eval_expression, eval_module,
check_compatibility, expression, kwargs...
)
opts = SciMLProblemOptions(
sys;
fn_opts, check_length, implicit_dae = true,
build_initializeprob = supports_initialization(sys),
time_dependent_init = is_time_dependent(sys),
circular_dependency_max_cycle_length = length(all_symbols(sys)),
kwargs...
)
return DAEProblem{iip, spec}(sys, op, tspan, opts; callback, kwargs...)
end

"""
SciMLBase.DAEProblem{iip, spec}(sys::System, op, tspan, opts::SciMLProblemOptions; callback = nothing, kwargs...)

Public entry point that builds a `DAEProblem` directly from a pre-assembled
[`SciMLProblemOptions`](@ref), bypassing the `kwargs...` wrapper above.
"""
function SciMLBase.DAEProblem{iip, spec}(
sys::System, op, tspan, opts::SciMLProblemOptions{E};
callback = nothing, kwargs...
) where {iip, spec, E}
check_complete(sys, DAEProblem)
check_compatibility && check_compatible_system(DAEProblem, sys)
opts.fn_opts.check_compatibility && check_compatible_system(DAEProblem, sys)

opts = maybe_derive_t_from_tspan(opts, tspan)

_iip = resolve_iip(iip, op)
f, du0,
u0,
p = process_SciMLProblem(
DAEFunction{_iip, spec}, sys, op;
t = tspan !== nothing ? tspan[1] : tspan, check_length, eval_expression,
eval_module, check_compatibility, implicit_dae = true, expression, kwargs...
DAEFunction{_iip, spec}, sys, op, opts; options_struct = Val(true), kwargs...
)

(; eval_expression, eval_module) = opts.fn_opts.codegen
kwargs = process_kwargs(
sys; expression, callback, eval_expression, eval_module,
sys; expression = Val{E}, callback, eval_expression, eval_module,
op, tspan, kwargs...
)

Expand All @@ -110,5 +136,5 @@ end
args = (; f, du0, u0, tspan, p)
kwargs = (; differential_vars, kwargs...)

return maybe_codegen_scimlproblem(expression, DAEProblem{_iip}, args; kwargs...)
return maybe_codegen_scimlproblem(Val{E}, DAEProblem{_iip}, args; kwargs...)
end
57 changes: 47 additions & 10 deletions lib/ModelingToolkitBase/src/problems/ddeproblem.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,29 +62,63 @@ end
sys::System, op, tspan;
callback = nothing, check_length = true, checkbounds = false,
eval_expression = false, eval_module = @__MODULE__, check_compatibility = true,
u0_constructor = identity, expression = Val{false}, kwargs...
u0_constructor = identity, expression = Val{false}, constant_lags = missing,
kwargs...
) where {iip, spec}
fn_opts = SciMLFunctionOptions(;
t = tspan !== nothing ? tspan[1] : tspan, eval_expression, eval_module,
checkbounds, check_compatibility, expression, kwargs...
)
opts = SciMLProblemOptions(
sys;
fn_opts, check_length, symbolic_u0 = true, u0_constructor,
build_initializeprob = supports_initialization(sys),
time_dependent_init = is_time_dependent(sys),
circular_dependency_max_cycle_length = length(all_symbols(sys)),
kwargs...
)
return DDEProblem{iip, spec}(sys, op, tspan, opts; callback, constant_lags, kwargs...)
end

"""
SciMLBase.DDEProblem{iip, spec}(sys::System, op, tspan, opts::SciMLProblemOptions; callback = nothing, kwargs...)

Public entry point that builds a `DDEProblem` directly from a pre-assembled
[`SciMLProblemOptions`](@ref), bypassing the `kwargs...` wrapper above.

`constant_lags` (default `missing`, meaning "not explicitly provided by the caller") is an
explicit keyword here — not part of `kwargs` — since it's relevant only to the final
`SciMLBase.DDEProblem` construction, not the inner `DDEFunction` build; the opts-accepting
`DDEFunction` method has no generic `kwargs...` sink to harmlessly absorb it the way its
keyword-based wrapper does.
"""
function SciMLBase.DDEProblem{iip, spec}(
sys::System, op, tspan, opts::SciMLProblemOptions{E};
callback = nothing, constant_lags = missing, kwargs...
) where {iip, spec, E}
check_complete(sys, DDEProblem)
check_compatibility && check_compatible_system(DDEProblem, sys)
opts.fn_opts.check_compatibility && check_compatible_system(DDEProblem, sys)

opts = maybe_derive_t_from_tspan(opts, tspan)

_iip = resolve_iip(iip, op)
f, u0,
p = process_SciMLProblem(
DDEFunction{_iip, spec}, sys, op;
t = tspan !== nothing ? tspan[1] : tspan, check_length, checkbounds,
eval_expression, eval_module, check_compatibility, symbolic_u0 = true,
expression, u0_constructor, kwargs...
DDEFunction{_iip, spec}, sys, op, opts; options_struct = Val(true), kwargs...
)

(; u0_constructor) = opts
(; eval_expression, eval_module) = opts.fn_opts.codegen
checkbounds = opts.fn_opts.codegen.codegen.checkbounds
h = generate_history(
sys, u0,
GeneratedFunctionOptions(;
expression, wrap_gfw = Val{true}, eval_expression, eval_module,
expression = Val{E}, wrap_gfw = Val{true}, eval_expression, eval_module,
codegen_function_options = Symbolics.CodegenFunctionOptions(; checkbounds)
)
)

if expression == Val{true}
if E
if u0 !== nothing
u0 = :($u0_constructor($map($float, h(p, tspan[1]))))
end
Expand All @@ -95,11 +129,14 @@ end
end

kwargs = process_kwargs(
sys; expression, callback, eval_expression, eval_module, op, tspan, kwargs...
sys; expression = Val{E}, callback, eval_expression, eval_module, op, tspan, kwargs...
)
args = (; f, u0, h, tspan, p)
constant_lags = resolve_constant_lags(sys, constant_lags, p)
constant_lags_kw = constant_lags === missing ? (;) : (; constant_lags)
kwargs = (; constant_lags_kw..., kwargs...)

return maybe_codegen_scimlproblem(expression, DDEProblem{_iip}, args; kwargs...)
return maybe_codegen_scimlproblem(Val{E}, DDEProblem{_iip}, args; kwargs...)
end

function check_compatible_system(T::Union{Type{DDEFunction}, Type{DDEProblem}}, sys::System)
Expand Down
32 changes: 26 additions & 6 deletions lib/ModelingToolkitBase/src/problems/discreteproblem.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,21 +58,41 @@ end
sys::System, op, tspan;
check_compatibility = true, expression = Val{false}, kwargs...
) where {iip, spec}
fn_opts = SciMLFunctionOptions(;
t = tspan !== nothing ? tspan[1] : tspan, check_compatibility, expression, kwargs...
)
opts = SciMLProblemOptions(
sys;
fn_opts, build_initializeprob = supports_initialization(sys),
time_dependent_init = is_time_dependent(sys),
circular_dependency_max_cycle_length = length(all_symbols(sys)),
kwargs...
)
return DiscreteProblem{iip, spec}(sys, op, tspan, opts; kwargs...)
end

"""
SciMLBase.DiscreteProblem{iip, spec}(sys::System, op, tspan, opts::SciMLProblemOptions; kwargs...)

Public entry point that builds a `DiscreteProblem` directly from a pre-assembled
[`SciMLProblemOptions`](@ref), bypassing the `kwargs...` wrapper above.
"""
function SciMLBase.DiscreteProblem{iip, spec}(
sys::System, op, tspan, opts::SciMLProblemOptions{E}; kwargs...
) where {iip, spec, E}
check_complete(sys, DiscreteProblem)
check_compatibility && check_compatible_system(DiscreteProblem, sys)
opts.fn_opts.check_compatibility && check_compatible_system(DiscreteProblem, sys)

_iip = resolve_iip(iip, op)
dvs = unknowns(sys)
op = to_varmap(op, dvs)
add_toterms!(op; replace = true)
f, u0,
p = process_SciMLProblem(
DiscreteFunction{_iip, spec}, sys, op;
t = tspan !== nothing ? tspan[1] : tspan, check_compatibility, expression,
kwargs...
DiscreteFunction{_iip, spec}, sys, op, opts; options_struct = Val(true), kwargs...
)

if expression == Val{true}
if E
u0 = :(f($u0, p, tspan[1]))
else
u0 = f(u0, p, tspan[1])
Expand All @@ -81,7 +101,7 @@ end
kwargs = process_kwargs(sys; kwargs...)
args = (; f, u0, tspan, p)

return maybe_codegen_scimlproblem(expression, DiscreteProblem{_iip}, args; kwargs...)
return maybe_codegen_scimlproblem(Val{E}, DiscreteProblem{_iip}, args; kwargs...)
end

function check_compatible_system(
Expand Down
50 changes: 43 additions & 7 deletions lib/ModelingToolkitBase/src/problems/homotopyproblem.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ be wrong mid-sweep.
expression = Val{false}, λspan = (0.0, 1.0),
check_length = true, check_compatibility = true,
eval_expression = false, eval_module = @__MODULE__,
checkbounds = false, cse = true, kwargs...
checkbounds = false, kwargs...
) where {iip, spec}
if expression !== Val{false}
throw(
Expand All @@ -45,11 +45,46 @@ be wrong mid-sweep.
)
)
end
if is_time_dependent(sys)
sys = NonlinearSystem(sys)
end
fn_opts = SciMLFunctionOptions(;
eval_expression, eval_module, checkbounds, check_compatibility, expression, kwargs...
)
opts = SciMLProblemOptions(
sys;
fn_opts, check_length, build_initializeprob = supports_initialization(sys),
time_dependent_init = is_time_dependent(sys),
circular_dependency_max_cycle_length = length(all_symbols(sys)),
kwargs...
)
return HomotopyProblem{iip, spec}(sys, op, opts; λspan, kwargs...)
end

"""
SciMLBase.HomotopyProblem{iip, spec}(sys::System, op, opts::SciMLProblemOptions; λspan = (0.0, 1.0), kwargs...)

Public entry point that builds a `HomotopyProblem` directly from a pre-assembled
[`SciMLProblemOptions`](@ref), bypassing the `kwargs...` wrapper above.
"""
function SciMLBase.HomotopyProblem{iip, spec}(
sys::System, op, opts::SciMLProblemOptions{E};
λspan = (0.0, 1.0), kwargs...
) where {iip, spec, E}
if E
throw(
ArgumentError(
"`HomotopyProblem(sys, op)` does not yet support " *
"`expression = Val{true}`; build the problem directly " *
"(the default `expression = Val{false}`)."
)
)
end
check_complete(sys, SciMLBase.HomotopyProblem)
if is_time_dependent(sys)
sys = NonlinearSystem(sys)
end
check_compatibility && check_compatible_system(SciMLBase.NonlinearProblem, sys)
opts.fn_opts.check_compatibility && check_compatible_system(SciMLBase.NonlinearProblem, sys)
if !has_any_homotopy(sys)
throw(
ArgumentError(
Expand All @@ -63,18 +98,19 @@ be wrong mid-sweep.
_iip = resolve_iip(iip, op)
f, u0,
p = process_SciMLProblem(
SciMLBase.NonlinearFunction{_iip, spec}, sys, op;
check_length, check_compatibility, expression,
eval_expression, eval_module, checkbounds, cse, kwargs...
SciMLBase.NonlinearFunction{_iip, spec}, sys, op, opts; options_struct = Val(true),
kwargs...
)

# Swap the opaque-`actual` residual for the homotopy-swept `f(u, p, λ)`. The
# observed function and residual prototype carry over; `initialization_data`,
# jacobian, and sparsity are deliberately not carried (the latter two encode
# the `λ = 1` system and would be wrong mid-sweep).
(; eval_expression, eval_module) = opts.fn_opts.codegen
checkbounds = opts.fn_opts.codegen.codegen.checkbounds
shadow, λ = lower_homotopy(sys)
hf = generate_homotopy_residual(
shadow, λ; eval_expression, eval_module, checkbounds, cse
shadow, λ; eval_expression, eval_module, checkbounds
)
swept_f = SciMLBase.NonlinearFunction{_iip}(
hf; sys = f.sys, observed = f.observed, resid_prototype = f.resid_prototype
Expand All @@ -83,6 +119,6 @@ be wrong mid-sweep.
kwargs = process_kwargs(sys; kwargs...)
args = (; f = swept_f, u0, p)
return maybe_codegen_scimlproblem(
expression, SciMLBase.HomotopyProblem{_iip}, args; λspan, kwargs...
Val{E}, SciMLBase.HomotopyProblem{_iip}, args; λspan, kwargs...
)
end
Loading
Loading