From 757899922963a6a870de6686a92ad2da5e87a0d4 Mon Sep 17 00:00:00 2001 From: termi-official <9196588+termi-official@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:39:39 +0200 Subject: [PATCH 1/2] =?UTF-8?q?=F0=9F=A4=96=20Add=20Configuration=20Tree?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/src/api-reference/index.md | 7 + docs/src/usage/index.md | 97 +++ src/OrdinaryDiffEqOperatorSplitting.jl | 11 + src/config_tree.jl | 407 +++++++++++++ src/integrator.jl | 169 ++++-- src/integrator_node.jl | 792 +++++++++++++++++++++++++ test/config_tree.jl | 395 ++++++++++++ test/qa/qa.jl | 3 + 8 files changed, 1828 insertions(+), 53 deletions(-) create mode 100644 src/config_tree.jl create mode 100644 src/integrator_node.jl create mode 100644 test/config_tree.jl diff --git a/docs/src/api-reference/index.md b/docs/src/api-reference/index.md index 7c630e8..bb2d674 100644 --- a/docs/src/api-reference/index.md +++ b/docs/src/api-reference/index.md @@ -18,3 +18,10 @@ GenericSplitFunction LieTrotterGodunov StrangMarchuk ``` + +## Per-node configuration + +```@docs +SplitNode +TreeOption +``` diff --git a/docs/src/usage/index.md b/docs/src/usage/index.md index 4e3d0e4..94e2bcd 100644 --- a/docs/src/usage/index.md +++ b/docs/src/usage/index.md @@ -64,3 +64,100 @@ for (u, t) in TimeChoiceIterator(integrator, 0.0:0.5:1.0) @show t, u end ``` + +## Configuring individual subintegrators + +`init` takes one value per keyword, which is not enough when the operators want +different treatment -- a stiff reaction term needs small steps while a diffusion term +is happy with large ones, or one operator should be integrated adaptively and another +at a fixed step size. + +A [`TreeOption`](@ref) carries one value per node of the splitting tree instead. It is +built for a given splitting function, so it knows the shape of the tree and rejects +addresses that do not exist: + +```julia +f_reaction = GenericSplitFunction((f_r1, f_r2), (r1dofs, r2dofs)) +f = GenericSplitFunction((f_diffusion, f_reaction), (ddofs, rdofs)) + +dt = TreeOption(f, 1.0e-2) # every node starts with the same value +``` + +Nodes are addressed either by their path or by a [`SplitNode`](@ref) minted from the +splitting function, where `f[2, 1]` is the first operator of the second operator of +`f`. Plain assignment sets a single node, broadcast assignment sets a node and +everything below it: + +```julia +dt[2] = 1.0e-4 # the reaction split node alone +dt[2] .= 1.0e-4 # the reaction split node and both of its operators +dt[f[2]] .= 1.0e-4 # the same, addressed by node +dt[f[2, 1]] = 1.0e-5 # just the first reaction operator + +integrator = init(prob, alg; dt) +``` + +Assignments are applied in order, so a broadcast over a subtree overwrites anything +more specific written before it. + +### Multi-rate integration + +A node's `dt` is the step size it uses to traverse the interval its parent hands it. +Giving the reaction subtree a smaller `dt` therefore makes it subcycle: with an outer +step of `1.0e-2` and a reaction step of `1.0e-4`, the reaction operators take a hundred +steps per splitting step and still land exactly on the synchronization point. + +Two things to keep in mind: + +- Under `StrangMarchuk` a child is handed intervals of `Δt/2`, `Δt` and `Δt/2`. A step + size that does not divide them leaves a short final sub-step in each interval, so the + sub-steps are not all the same length. Step sizes that are exactly representable + (powers of two, say) avoid this. +- A node's `dt` larger than the interval it is handed is clipped to that interval. + +For an adaptive node the configured `dt` is only the initial step size. + +### Mixing adaptive and fixed-step operators + +Adaptivity is configured the same way: + +```julia +adaptive = TreeOption(f, false) +adaptive[f[2, 1]] = true +adaptive[f[2, 2]] = true + +integrator = init(prob, alg; dt = 1.0e-2, adaptive) +``` + +Note that the splitting nodes themselves stay non-adaptive here. Broadcasting +`adaptive[2] .= true` would also mark the reaction *split node* as adaptive, and since +`LieTrotterGodunov` is not an adaptive algorithm that produces a warning. + +Any other keyword accepted by the inner integrators can be given per node as well and +is passed down to the leaves, while keywords a splitting node understands (`dtmin`, +`dtmax`, `failfactor`) are applied at every level: + +```julia +reltol = TreeOption(f, 1.0e-3) +reltol[f[2, 1]] = 1.0e-9 + +integrator = init(prob, alg; dt = 1.0e-2, adaptive, reltol) +``` + +### Reading the tree back + +A `SplitNode` addresses the same position in every tree that mirrors the splitting +function, so it resolves against the algorithm and the integrator too: + +```julia +f[f[2, 1]] # the sub function +alg[f[2, 1]] # the inner algorithm +integrator[f[2, 1]] # the sub integrator +``` + +`integrator[i]` is *not* available for this: SciMLBase already gives integer indexing +of an integrator the meaning "the `i`-th state component". + +Calling `reinit!` without a `dt` restores every node to its configured step size, so a +multi-rate setup survives. Passing a `dt` reconfigures the tree exactly as at `init`: +a single value applies to every node, a `TreeOption` node by node. diff --git a/src/OrdinaryDiffEqOperatorSplitting.jl b/src/OrdinaryDiffEqOperatorSplitting.jl index 1954897..5857b97 100644 --- a/src/OrdinaryDiffEqOperatorSplitting.jl +++ b/src/OrdinaryDiffEqOperatorSplitting.jl @@ -27,6 +27,15 @@ else end _inner_verbose(verbose) = verbose +# `verbose` reaches us either as a Bool or, through DiffEqBase v7's `init`, as a +# DEVerbosity whose first type parameter is the on/off flag. Neither can be used in +# a boolean context directly. +_is_verbose(verbose::Bool) = verbose +_is_verbose(verbose) = true +@static if isdefined(DiffEqBase, :DEVerbosity) + _is_verbose(::DiffEqBase.DEVerbosity{B}) where {B} = B +end + abstract type AbstractOperatorSplitFunction <: SciMLBase.AbstractODEFunction{true} end abstract type AbstractOperatorSplittingAlgorithm end abstract type AbstractOperatorSplittingCache end @@ -35,12 +44,14 @@ abstract type AbstractOperatorSplittingCache end @inline isdtchangeable(alg::AbstractOperatorSplittingAlgorithm) = all(isdtchangeable.(alg.inner_algs)) include("function.jl") +include("config_tree.jl") include("problem.jl") include("integrator.jl") include("solver.jl") include("utils.jl") export GenericSplitFunction, OperatorSplittingProblem, LieTrotterGodunov, StrangMarchuk +export SplitNode, TreeOption include("precompilation.jl") diff --git a/src/config_tree.jl b/src/config_tree.jl new file mode 100644 index 0000000..d118373 --- /dev/null +++ b/src/config_tree.jl @@ -0,0 +1,407 @@ +# --------------------------------------------------------------------------- +# Tree addresses +# --------------------------------------------------------------------------- + +""" + SplitNode + +Address of a node in an operator splitting tree, together with the object that lives +at that address. + +Nodes are minted by indexing a [`GenericSplitFunction`](@ref) with integers. `f[]` +addresses the root, `f[i]` its `i`-th operator and `f[i, j]` the `j`-th operator of +that one: + +```julia +node = f[2, 1] # equivalently f[2][1] +node.path # (2, 1) +node.object # the addressed sub function +``` + +The splitting function defines the shape of the tree, so an address minted from it +denotes the same position in every tree that mirrors it and can be resolved against +each of them: + +```julia +f[node] # the sub function +alg[node] # the inner algorithm +integrator[node] # the sub integrator +``` + +Addresses are also how a [`TreeOption`](@ref) is told which node to configure. +""" +struct SplitNode{N, O} + path::NTuple{N, Int} + object::O +end + +_path(is::Integer...) = map(Int, is) + +# Human readable rendering of a path, used in error messages. +_showpath(::Tuple{}) = "the root operator" +_showpath(path::Tuple) = "operator [" * join(path, ", ") * "]" + +""" + _tree_child(obj, i) + +The `i`-th child of `obj` in whichever splitting tree `obj` belongs to. Every tree +that mirrors a [`GenericSplitFunction`](@ref) implements this so that a +[`SplitNode`](@ref) can be resolved against it. +""" +function _tree_child(obj, i::Int) + return throw( + ArgumentError( + "$(typeof(obj)) is a leaf of the splitting tree and has no operator $i to descend into." + ) + ) +end + +function _tree_child(f::GenericSplitFunction, i::Int) + checkbounds(Bool, 1:length(f.functions), i) || throw( + ArgumentError( + "operator $i is out of range: this splitting node has $(length(f.functions)) operators." + ) + ) + return f.functions[i] +end + +function _tree_child(alg::AbstractOperatorSplittingAlgorithm, i::Int) + checkbounds(Bool, 1:length(alg.inner_algs), i) || throw( + ArgumentError( + "operator $i is out of range: $(nameof(typeof(alg))) has $(length(alg.inner_algs)) inner algorithms." + ) + ) + return alg.inner_algs[i] +end + +_resolve(obj, ::Tuple{}) = obj +_resolve(obj, path::Tuple) = _resolve(_tree_child(obj, first(path)), Base.tail(path)) + +_descend(node::SplitNode, ::Tuple{}) = node +function _descend(node::SplitNode, path::Tuple) + i = first(path) + child = _tree_child(node.object, i) + return _descend(SplitNode((node.path..., i), child), Base.tail(path)) +end + +# --- minting addresses from the splitting function --- +Base.getindex(f::GenericSplitFunction) = SplitNode((), f) +Base.getindex(f::GenericSplitFunction, i::Integer, is::Integer...) = + _descend(SplitNode((), f), _path(i, is...)) +Base.getindex(node::SplitNode, i::Integer, is::Integer...) = _descend(node, _path(i, is...)) + +# --- resolving addresses against the mirroring trees --- +Base.getindex(f::GenericSplitFunction, node::SplitNode) = _resolve(f, node.path) +Base.getindex(alg::AbstractOperatorSplittingAlgorithm, node::SplitNode) = + _resolve(alg, node.path) +Base.getindex(alg::AbstractOperatorSplittingAlgorithm, i::Integer, is::Integer...) = + _resolve(alg, _path(i, is...)) + +get_operator(f::GenericSplitFunction, node::SplitNode) = _resolve(f, node.path) + +function Base.show(io::IO, node::SplitNode) + print(io, "f[", join(node.path, ", "), "]") + return +end +function Base.show(io::IO, ::MIME"text/plain", node::SplitNode) + show(io, node) + print(io, " -> ", nameof(typeof(node.object))) + return +end + +# --------------------------------------------------------------------------- +# Per-node options +# --------------------------------------------------------------------------- + +""" + TreeOption(f::GenericSplitFunction, default) + TreeOption{T}(f::GenericSplitFunction, default) + +A value per node of an operator splitting tree, used to give individual +subintegrators their own settings where `init` only accepts one value for all of +them. + +The option mirrors the shape of `f` and starts out holding `default` everywhere. +Plain assignment sets a single node; broadcast assignment sets a node and everything +below it: + +```julia +dt = TreeOption(f, 1.0e-2) +dt[2] = 1.0e-4 # this node alone +dt[2] .= 1.0e-4 # this node and all operators below it +dt .= 1.0e-4 # every node +dt[f[2]] = 1.0e-4 # addressed by a SplitNode instead of a path +dt[2, 1] # read the value at [2, 1] +``` + +Later assignments overwrite earlier ones, so a broadcast over a subtree clobbers more +specific values written before it. + +The element type is fixed by `default`; use `TreeOption{Any}(f, default)` for an +option whose values are of mixed type, such as a step size controller. +""" +mutable struct TreeOption{T} + value::T + const children::Vector{TreeOption{T}} +end + +TreeOption(f::GenericSplitFunction, default::T) where {T} = _build_option(f, default, T) +TreeOption{T}(f::GenericSplitFunction, default) where {T} = + _build_option(f, _checked_convert(T, default), T) + +function _build_option(f::GenericSplitFunction, value::T, ::Type{T}) where {T} + children = TreeOption{T}[_build_option(fi, value, T) for fi in f.functions] + return TreeOption{T}(value, children) +end +_build_option(_, value::T, ::Type{T}) where {T} = TreeOption{T}(value, TreeOption{T}[]) + +function _checked_convert(::Type{T}, v) where {T} + return try + convert(T, v) + catch err + err isa MethodError || rethrow() + throw( + ArgumentError( + "cannot store a value of type $(typeof(v)) in a TreeOption{$T}. \ + Construct the option as `TreeOption{Any}(f, default)` if it has to hold values of mixed type." + ) + ) + end +end + +_option_node(opt::TreeOption, ::Tuple{}) = opt +function _option_node(opt::TreeOption, path::Tuple) + i = first(path) + isempty(opt.children) && throw( + ArgumentError( + "$(_showpath(path)) does not exist: this option mirrors a leaf operator, which has no operators below it." + ) + ) + checkbounds(Bool, opt.children, i) || throw( + ArgumentError( + "operator $i is out of range: this splitting node has $(length(opt.children)) operators." + ) + ) + return _option_node(opt.children[i], Base.tail(path)) +end + +_store!(opt::TreeOption{T}, v) where {T} = (opt.value = _checked_convert(T, v)) + +function _fill_subtree!(opt::TreeOption, v) + _store!(opt, v) + for child in opt.children + _fill_subtree!(child, v) + end + return opt +end + +Base.getindex(opt::TreeOption) = opt.value +Base.getindex(opt::TreeOption, i::Integer, is::Integer...) = + _option_node(opt, _path(i, is...)).value +Base.getindex(opt::TreeOption, node::SplitNode) = _option_node(opt, node.path).value + +Base.setindex!(opt::TreeOption, v) = _store!(opt, v) +Base.setindex!(opt::TreeOption, v, i::Integer, is::Integer...) = + _store!(_option_node(opt, _path(i, is...)), v) +Base.setindex!(opt::TreeOption, v, node::SplitNode) = + _store!(_option_node(opt, node.path), v) + +""" + TreeOptionSubtree + +The target of a broadcast assignment into a [`TreeOption`](@ref), produced by +`Base.dotview`. Assigning into it writes one value to a node and all of its +descendants. +""" +struct TreeOptionSubtree{T} + node::TreeOption{T} +end + +Base.dotview(opt::TreeOption) = TreeOptionSubtree(opt) +Base.dotview(opt::TreeOption, i::Integer, is::Integer...) = + TreeOptionSubtree(_option_node(opt, _path(i, is...))) +Base.dotview(opt::TreeOption, node::SplitNode) = + TreeOptionSubtree(_option_node(opt, node.path)) + +const _Broadcasted = Base.Broadcast.Broadcasted + +# A TreeOption is not a collection as far as broadcasting is concerned. Without this +# the default `broadcastable` tries to `collect` it, and `opt .op= x` fails with a +# MethodError about iteration instead of the explanation in `_broadcast_value`. +Base.broadcastable(opt::TreeOption) = Ref(opt) + +# Broadcast assignment into a TreeOption only ever means "give every node of this +# subtree the same value". Anything else that lowers to the same call is rejected +# rather than silently given a made up meaning -- see `_broadcast_value`. +Base.materialize!(dest::TreeOptionSubtree, bc::_Broadcasted) = + _fill_subtree!(dest.node, _broadcast_value(bc)) +Base.materialize!(dest::TreeOption, bc::_Broadcasted) = + _fill_subtree!(dest, _broadcast_value(bc)) + +function _broadcast_value(bc::_Broadcasted) + # `opt[i] .op= x` lowers to a broadcast of `op` whose first argument is the value + # read back from the *single* node `opt[i]`, while the assignment target is the + # whole subtree. There is no reading of that which is not surprising, so refuse. + bc.f === identity || throw( + ArgumentError( + "`.$(bc.f)=` is not supported for a TreeOption: it reads the value of a single node \ + but writes to that node and everything below it. Write `opt[...] .= $(bc.f)(opt[...], x)` \ + with an undotted right hand side if that is what you intend." + ) + ) + length(bc.args) == 1 || throw( + ArgumentError("expected a single value to broadcast over the subtree, got $(length(bc.args)).") + ) + v = only(bc.args) + # Values that are not `broadcastable` arrive wrapped in a Ref. + v isa Ref && return v[] + return if v isa Union{AbstractArray, Tuple, _Broadcasted} + throw( + ArgumentError( + "broadcast assignment into a TreeOption expects one scalar value, got $(typeof(v)). \ + Every node of the subtree receives the same value, so there is nothing to distribute." + ) + ) + else + v + end +end + +""" + structure_matches(opt::TreeOption, f) + +Whether `opt` was built for a splitting function of the same shape as `f`. +""" +structure_matches(opt::TreeOption, f::GenericSplitFunction) = + length(opt.children) == length(f.functions) && + all(structure_matches(o, fi) for (o, fi) in zip(opt.children, f.functions)) +structure_matches(opt::TreeOption, _) = isempty(opt.children) + +function Base.show(io::IO, ::MIME"text/plain", opt::TreeOption{T}) where {T} + println(io, "TreeOption{", T, "}:") + return _show_option(io, opt, ()) +end + +function _show_option(io::IO, opt::TreeOption, path::Tuple) + println(io, " [", join(path, ", "), "] => ", repr(opt.value)) + for (i, child) in enumerate(opt.children) + _show_option(io, child, (path..., i)) + end + return +end + +# --------------------------------------------------------------------------- +# Frozen configuration tree +# --------------------------------------------------------------------------- + +""" + ConfigTree + +The settings of every node of a splitting tree, resolved once when the integrator is +built. `values` holds this node's settings and `children` mirrors the operators of +the corresponding [`GenericSplitFunction`](@ref). + +This is what the tree construction in `src/integrator.jl` consumes: unlike a +[`TreeOption`](@ref), which is a mutable sparse thing the user edits, a `ConfigTree` +is immutable and concretely typed, so `config.children[i]` is as inferable as +`alg.inner_algs[i]` and no lookup reaches the stepping loop. +""" +struct ConfigTree{V <: NamedTuple, C <: Tuple} + values::V + children::C +end + +_tree_child(config::ConfigTree, i::Int) = config.children[i] +Base.getindex(config::ConfigTree, node::SplitNode) = _resolve(config, node.path) + +""" + build_config_tree(f::GenericSplitFunction, options::NamedTuple) + +Resolve `options` -- each entry either one value for the whole tree or a +[`TreeOption`](@ref) with a value per node -- into a [`ConfigTree`](@ref) mirroring `f`. +""" +function build_config_tree(f::GenericSplitFunction, options::NamedTuple) + for (name, opt) in pairs(options) + if opt isa TreeOption && !structure_matches(opt, f) + throw( + ArgumentError( + "the TreeOption passed as `$name` mirrors a splitting function of a different \ + shape than the one this problem was built from." + ) + ) + end + end + return _config_tree(f, options) +end + +function _config_tree(f::GenericSplitFunction, options::NamedTuple) + children = ntuple( + i -> _config_tree(f.functions[i], map(o -> _option_child(o, i), options)), + length(f.functions) + ) + return ConfigTree(map(_option_value, options), children) +end +_config_tree(_, options::NamedTuple) = ConfigTree(map(_option_value, options), ()) + +_option_value(opt::TreeOption) = opt.value +_option_value(v) = v +_option_child(opt::TreeOption, i::Int) = opt.children[i] +_option_child(v, ::Int) = v + +function validate_dt_tree(config::ConfigTree, path::Tuple = ()) + dt = config.values.dt + dt > zero(dt) || error("dt must be positive, but $(_showpath(path)) was given $dt.") + for (i, child) in enumerate(config.children) + validate_dt_tree(child, (path..., i)) + end + return +end + +""" + signed_dt_tree(config, tdir, tType) + +Give every node's `dt` the direction of integration and the time type of the problem. +""" +signed_dt_tree(config::ConfigTree, tdir, ::Type{tType}) where {tType} = ConfigTree( + merge(config.values, (; dt = tdir * convert(tType, config.values.dt))), + map(child -> signed_dt_tree(child, tdir, tType), config.children) +) + +""" + warn_non_adaptive(alg, config) + +Warn about every splitting node that was asked to be adaptive although its algorithm +is not. Leaf algorithms are left to the inner integrator, which does its own checking. +""" +function warn_non_adaptive( + alg::AbstractOperatorSplittingAlgorithm, config::ConfigTree, path::Tuple = () + ) + if config.values.adaptive && _is_verbose(config.values.verbose) && + !SciMLBase.isadaptive(alg) + @warn "The algorithm $alg at $(_showpath(path)) is not adaptive." + end + for (i, child) in enumerate(config.children) + warn_non_adaptive(alg.inner_algs[i], child, (path..., i)) + end + return +end +warn_non_adaptive(alg, config::ConfigTree, path::Tuple = ()) = nothing + +# Settings every node handles itself; anything else a user passes to `init` is an +# option of the inner integrators and travels down to the leaves untouched. +const NODE_OPTION_KEYS = (:dt, :adaptive, :verbose, :controller) +inner_values(values::NamedTuple) = + NamedTuple{filter(!in(NODE_OPTION_KEYS), keys(values))}(values) + +# ... of which a splitting node understands these. +const SPLIT_OPTION_KEYS = (:dtmin, :dtmax, :failfactor, :isoutofdomain) + +function split_integrator_options(values::NamedTuple) + inner = inner_values(values) + known = filter(in(SPLIT_OPTION_KEYS), keys(inner)) + return IntegratorOptions(; + verbose = values.verbose, + adaptive = values.adaptive, + NamedTuple{known}(inner)... + ) +end diff --git a/src/integrator.jl b/src/integrator.jl index 6775f02..856474b 100644 --- a/src/integrator.jl +++ b/src/integrator.jl @@ -143,6 +143,7 @@ mutable struct OperatorSplittingIntegrator{ childSyncType, controllerType, optionsType, + configType, } <: SciMLBase.AbstractODEIntegrator{algType, true, uType, tType} const f::fType const alg::algType @@ -178,6 +179,8 @@ mutable struct OperatorSplittingIntegrator{ stats::IntegratorStats tdir::tType next_sync_is_continuous::Bool + # Per-node settings, kept so that `reinit!` can restore them. + config::configType end is_next_sync_continuous(integrator) = false @@ -189,6 +192,18 @@ reset_next_sync_continuous(integrator::OperatorSplittingIntegrator) = integrator const AnySplitIntegrator = Union{SplitSubIntegrator, OperatorSplittingIntegrator} +# DiffEqBase promotes the problem's `tspan` against the `dt` keyword before `__init` +# is ever called. A TreeOption is not a number, so hand it the root value: the step +# size the outermost integrator runs with. +function SciMLBase.promote_tspan(u0, p, tspan, prob::OperatorSplittingProblem, kwargs) + dt = get(kwargs, :dt, nothing) + dt === nothing && return tspan + tspan1, tspan2, _ = promote(tspan..., _root_value(dt)) + return (tspan1, tspan2) +end +_root_value(opt::TreeOption) = opt.value +_root_value(v) = v + # --------------------------------------------------------------------------- # __init # --------------------------------------------------------------------------- @@ -213,13 +228,17 @@ function SciMLBase.__init( (; u0, p) = prob t0, tf = prob.tspan - dt > zero(dt) || error("dt must be positive") - dtcache = dt - dt = tf > t0 ? dt : -dt - tType = typeof(dt) + # Every setting is either one value for the whole tree or a TreeOption carrying a + # value per node. Beyond the four this integrator handles itself, whatever the + # caller passes travels down to the leaf integrators. + config = build_config_tree(prob.f, (; dt, adaptive, verbose, controller, kwargs...)) + validate_dt_tree(config) + tType = typeof(config.values.dt) + config = signed_dt_tree(config, tf > t0 ? one(tType) : -one(tType), tType) + warn_non_adaptive(alg, config) - (!isadaptive(alg) && adaptive && verbose) && - @warn("The algorithm $alg is not adaptive.") + dt = config.values.dt + dtcache = abs(dt) dtchangeable = isdtchangeable(alg) @@ -251,9 +270,9 @@ function SciMLBase.__init( uprev, u, u, # u_master == u at the outermost level 1:length(u), - t0, dt, tf, + t0, tf, tstops, saveat, d_discontinuities, callback, - adaptive, verbose + config ) cache = init_cache( @@ -282,11 +301,12 @@ function SciMLBase.__init( child_solution_indices, child_synchronizers, 0, - controller, - IntegratorOptions(; verbose, adaptive), + config.values.controller, + split_integrator_options(config.values), IntegratorStats(), tType(tstops_internal.ordering isa BinaryHeaps.FasterForward ? 1 : -1), false, + config, ) DiffEqBase.initialize!(callback, u0, t0, integrator) return integrator @@ -302,21 +322,27 @@ function DiffEqBase.reinit!( u0 = integrator.sol.prob.u0; t0 = integrator.sol.prob.tspan[1], tf = integrator.sol.prob.tspan[2], - dt = isadaptive(integrator) ? nothing : integrator.dtcache, + dt = nothing, erase_sol = false, tstops = integrator._tstops, saveat = integrator._saveat, reinit_callbacks = true, reinit_retcode = true ) + # Without a `dt` every node is restored to the step size it was configured with, + # so a multi-rate setup survives a reinit!. Passing one is equivalent to passing + # it to `init`: a scalar reconfigures the whole tree, a TreeOption node by node. + if dt !== nothing + integrator.config = _reconfigure_dt(integrator, dt, t0, tf) + end + config = integrator.config + integrator.u .= u0 integrator.uprev .= u0 integrator.t = t0 integrator.tprev = t0 - if dt !== nothing - integrator.dt = dt - integrator.dtcache = dt - end + integrator.dt = config.values.dt + integrator.dtcache = abs(config.values.dt) integrator.tstops, integrator.saveat = tstops_and_saveat_heaps(t0, tf, tstops, saveat) integrator.iter = 0 @@ -339,26 +365,45 @@ function DiffEqBase.reinit!( _subreinit_tuple!( integrator.f, u0, - integrator.child_subintegrators; - t0, tf, dt, + integrator.child_subintegrators, + config; + t0, tf, erase_sol, tstops, saveat, reinit_callbacks, reinit_retcode ) return nothing end +# Rebuild the step sizes of the stored configuration from a `dt` given to reinit!. +function _reconfigure_dt(integrator::OperatorSplittingIntegrator, dt, t0, tf) + tType = typeof(integrator.dt) + dt_config = build_config_tree(integrator.f, (; dt)) + validate_dt_tree(dt_config) + dt_config = signed_dt_tree(dt_config, tf > t0 ? one(tType) : -one(tType), tType) + return _replace_dt(integrator.config, dt_config) +end + +_replace_dt(config::ConfigTree, dt_config::ConfigTree) = ConfigTree( + merge(config.values, (; dt = dt_config.values.dt)), + ntuple( + i -> _replace_dt(config.children[i], dt_config.children[i]), + length(config.children) + ) +) + # --- subreinit! helpers --- # Iterate over a tuple of children (outermost call from reinit!) @unroll function _subreinit_tuple!( f, u0, - children::Tuple; + children::Tuple, + config::ConfigTree; kwargs... ) i = 1 @unroll for child in children - _subreinit_child!(get_operator(f, i), u0, child; kwargs...) + _subreinit_child!(get_operator(f, i), u0, child, config.children[i]; kwargs...) i += 1 end end @@ -367,14 +412,14 @@ end function _subreinit_child!( f_child, u0, - child::DEIntegrator; - dt, + child::DEIntegrator, + config::ConfigTree; kwargs... ) - if dt !== nothing && child.dtchangeable - SciMLBase.set_proposed_dt!(child, dt) + if child.dtchangeable + SciMLBase.set_proposed_dt!(child, config.values.dt) # Reinit does not touch this, so we reset it manually. - set_dt!(child, dt) + set_dt!(child, config.values.dt) end DiffEqBase.reinit!(child; kwargs...) return nothing @@ -384,17 +429,15 @@ end function _subreinit_child!( f_child, u0, - sub::SplitSubIntegrator; + sub::SplitSubIntegrator, + config::ConfigTree; t0, tf, - dt, kwargs... ) sub.t = t0 - if dt !== nothing - SciMLBase.set_proposed_dt!(sub, dt) - set_dt!(sub, dt) - end + SciMLBase.set_proposed_dt!(sub, config.values.dt) + set_dt!(sub, config.values.dt) sub.iter = 0 sub.force_stepfail = false sub.last_step_failed = false @@ -407,8 +450,9 @@ function _subreinit_child!( _subreinit_tuple!( f_child, u0, - sub.child_subintegrators; - t0, tf, dt, kwargs... + sub.child_subintegrators, + config; + t0, tf, kwargs... ) return nothing end @@ -735,7 +779,7 @@ function SciMLBase.check_error(integrator::OperatorSplittingIntegrator) return integrator.sol.retcode end if DiffEqBase.NAN_CHECK(integrator.dtcache) || DiffEqBase.NAN_CHECK(integrator.dt) - integrator.opts.verbose && + _is_verbose(integrator.opts.verbose) && @warn("NaN dt detected. Likely a NaN value in the state, parameters, or derivative value caused this outcome.") return ReturnCode.DtNaN end @@ -748,7 +792,7 @@ function SciMLBase.check_error(integrator::SplitSubIntegrator) return integrator.status.retcode end if DiffEqBase.NAN_CHECK(integrator.dtcache) || DiffEqBase.NAN_CHECK(integrator.dt) - integrator.opts.verbose && + _is_verbose(integrator.opts.verbose) && @warn("NaN dt detected. Likely a NaN value in the state, parameters, or derivative value caused this outcome.") return ReturnCode.DtNaN end @@ -952,9 +996,9 @@ function build_subintegrators( uouter::AbstractVector, u_master::AbstractVector, solution_indices, - t0, dt, tf, + t0, tf, tstops, saveat, d_discontinuities, callback, - adaptive, verbose + config::ConfigTree ) (; f, p) = prob @@ -966,9 +1010,9 @@ function build_subintegrators( p[i], uprevouter, uouter, u_master, f.solution_indices[i], - t0, dt, tf, + t0, tf, tstops, saveat, d_discontinuities, callback, - adaptive, verbose + config.children[i] ), length(f.functions) ) @@ -987,12 +1031,11 @@ function _build_child( uouter::AbstractVector, u_master::AbstractVector, solution_indices, - t0, dt, tf, + t0, tf, tstops, saveat, d_discontinuities, callback, - adaptive, verbose, - save_end = false, - controller = nothing + config::ConfigTree ) + dt = config.values.dt tType = typeof(dt) # Recurse: build each consecutive child @@ -1004,9 +1047,9 @@ function _build_child( p[i], uprevouter, uouter, u_master, f.solution_indices[i], - t0, dt, tf, + t0, tf, tstops, saveat, d_discontinuities, callback, - adaptive, verbose + config.children[i] ), length(f.functions) ) @@ -1038,7 +1081,7 @@ function _build_child( tstops_internal, 0, # iter EEst_val, - controller, + config.values.controller, false, false, false, # force_stepfail, last_step_failed, u_modified SplitSubIntegratorStatus(), IntegratorStats(), @@ -1047,7 +1090,7 @@ function _build_child( solution_indices, child_solution_indices, child_synchronizers, - IntegratorOptions(; verbose, adaptive), + split_integrator_options(config.values), one(tType), ) @@ -1063,11 +1106,9 @@ function _build_child( uprevouter::S, uouter::S, u_master::S, solution_indices, - t0::T, dt::T, tf::T, + t0::T, tf::T, tstops, saveat, d_discontinuities, callback, - adaptive, verbose, - save_end = false, - controller = nothing + config::ConfigTree ) where {S, T, P, F} u = uouter[solution_indices] u0 = if f isa SciMLBase.AbstractSciMLFunction @@ -1084,18 +1125,40 @@ function _build_child( integrator = SciMLBase.__init( prob2, alg; - dt, + dt = config.values.dt, tstops, saveat = (), d_discontinuities, save_everystep = false, advance_to_tstop = false, - adaptive, controller, - verbose = _inner_verbose(verbose) + adaptive = config.values.adaptive, + controller = config.values.controller, + verbose = _inner_verbose(config.values.verbose), + inner_values(config.values)... ) return integrator end +# --------------------------------------------------------------------------- +# Tree addressing +# +# Only `SplitNode` addresses are accepted here. SciMLBase already gives +# `integrator[i]` the meaning "the i-th state component" via symbolic indexing +# (`Base.getindex(::DEIntegrator, sym)`), so integer indexing is left alone. +# --------------------------------------------------------------------------- +function _tree_child(integrator::AnySplitIntegrator, i::Int) + children = integrator.child_subintegrators + checkbounds(Bool, 1:length(children), i) || throw( + ArgumentError( + "operator $i is out of range: this splitting node has $(length(children)) subintegrators." + ) + ) + return children[i] +end + +Base.getindex(integrator::AnySplitIntegrator, node::SplitNode) = + _resolve(integrator, node.path) + # --------------------------------------------------------------------------- # SciMLBase API # --------------------------------------------------------------------------- diff --git a/src/integrator_node.jl b/src/integrator_node.jl new file mode 100644 index 0000000..6c29f17 --- /dev/null +++ b/src/integrator_node.jl @@ -0,0 +1,792 @@ +""" + OperatorSplittingIntegratorNode <: AbstractODEIntegrator + +A variant of [`ODEIntegrator`](https://github.com/SciML/OrdinaryDiffEq.jl/blob/6ec5a55bda26efae596bf99bea1a1d729636f412/src/integrators/type.jl#L77-L123) to perform opeartor splitting. +""" +mutable struct OperatorSplittingIntegratorNode{ + fType, + algType, + uType, + tType, + pType, + heapType, + tstopsType, + saveatType, + callbackType, + cacheType, + solType, + subintTreeType, + solidxTreeType, + syncTreeType, + controllerType, + optionsType, + } <: DiffEqBase.AbstractODEIntegrator{algType, true, uType, tType} + const f::fType + const alg::algType + u::uType # Master Solution + uprev::uType # Master Solution + tmp::uType # Interpolation buffer + p::pType + t::tType # Current time + tprev::tType + dt::tType # This is the time step length which which we use during time marching + dtcache::tType # This is the proposed time step length + const dtchangeable::Bool # Indicator whether dtcache can be changed + tstops::heapType + _tstops::tstopsType # argument to __init used as default argument to reinit! + saveat::heapType + _saveat::saveatType # argument to __init used as default argument to reinit! + callback::callbackType + advance_to_tstop::Bool + # TODO group these into some internal flag struct + last_step_failed::Bool + force_stepfail::Bool + isout::Bool + u_modified::Bool + # DiffEqBase.initialize! and DiffEqBase.finalize! + cache::cacheType + sol::solType + subintegrator_tree::subintTreeType + iter::Int + controller::controllerType + EEst::Float64 # TODO integrate with controller cache + opts::optionsType + stats::IntegratorStats + tdir::tType +end + +# called by DiffEqBase.init and DiffEqBase.solve +function DiffEqBase.__init( + prob::OperatorSplittingProblem, + alg::AbstractOperatorSplittingAlgorithm, + args...; + dt, + tstops = (), + saveat = (), + d_discontinuities = (), + save_everystep = false, + callback = nothing, + advance_to_tstop = false, + adaptive = DiffEqBase.isadaptive(alg), + controller = nothing, + alias_u0 = false, + verbose = true, + kwargs... + ) + (; u0, p) = prob + t0, tf = prob.tspan + + dt > zero(dt) || error("dt must be positive") + dtcache = dt + dt = tf > t0 ? dt : -dt + tType = typeof(dt) + + # Warn if the algorithm is non-adaptive but the user tries to make it adaptive. + (!DiffEqBase.isadaptive(alg) && adaptive && verbose) && warn("The algorithm $alg is not adaptive.") + + dtchangeable = true # DiffEqBase.isadaptive(alg) + + if tstops isa AbstractArray || tstops isa Tuple || tstops isa Number + _tstops = nothing + else + _tstops = tstops + tstops = () + end + + # Setup tstop logic + tstops_internal = OrdinaryDiffEqCore.initialize_tstops( + tType, tstops, d_discontinuities, prob.tspan + ) + saveat_internal = OrdinaryDiffEqCore.initialize_saveat(tType, saveat, prob.tspan) + d_discontinuities_internal = OrdinaryDiffEqCore.initialize_d_discontinuities( + tType, d_discontinuities, prob.tspan + ) + + u = setup_u(prob, alg, alias_u0) + uprev = setup_u(prob, alg, false) + tmp = setup_u(prob, alg, false) + uType = typeof(u) + + sol = DiffEqBase.build_solution(prob, alg, tType[], uType[]) + + callback = DiffEqBase.CallbackSet(callback) + + subintegrator_tree, + cache = build_subintegrator_tree_with_cache( + prob, alg, + uprev, u, + 1:length(u), + t0, dt, tf, + tstops, saveat, d_discontinuities, callback, + adaptive, verbose + ) + + if controller === nothing && adaptive + controller = default_controller(alg, cache) + end + + integrator = OperatorSplittingIntegrator( + prob.f, + alg, + u, + uprev, + tmp, + p, + t0, + copy(t0), + dt, + dtcache, + dtchangeable, + tstops_internal, + tstops, + saveat_internal, + saveat, + callback, + advance_to_tstop, + false, + false, + false, + false, + cache, + sol, + subintegrator_tree, + build_solution_index_tree(prob.f), + build_synchronizer_tree(prob.f), + 0, + controller, + NaN, + IntegratorOptions(; verbose, adaptive, kwargs...), + IntegratorStats(), + tType(tstops_internal.ordering isa DataStructures.FasterForward ? 1 : -1) + ) + DiffEqBase.initialize!(callback, u0, t0, integrator) # Do I need this? + return integrator +end + +DiffEqBase.has_reinit(integrator::OperatorSplittingIntegrator) = true +function DiffEqBase.reinit!( + integrator::OperatorSplittingIntegrator, + u0 = integrator.sol.prob.u0; + t0 = integrator.sol.prob.tspan[1], + tf = integrator.sol.prob.tspan[2], + erase_sol = false, + tstops = integrator._tstops, + saveat = integrator._saveat, + reinit_callbacks = true, + reinit_retcode = true + ) + integrator.u .= u0 + integrator.uprev .= u0 + integrator.t = t0 + integrator.tprev = t0 + integrator.tstops, integrator.saveat = tstops_and_saveat_heaps(t0, tf, tstops, saveat) + integrator.iter = 0 + if erase_sol + resize!(integrator.sol.t, 0) + resize!(integrator.sol.u, 0) + end + if reinit_callbacks + DiffEqBase.initialize!(integrator.callback, u0, t0, integrator) + else # always reinit the saving callback so that t0 can be saved if needed + saving_callback = integrator.callback.discrete_callbacks[end] + DiffEqBase.initialize!(saving_callback, u0, t0, integrator) + end + if reinit_retcode + integrator.sol = DiffEqBase.solution_new_retcode( + integrator.sol, SciMLBase.ReturnCode.Default + ) + end + + return subreinit!( + integrator.f, + u0, + 1:length(u0), + integrator.subintegrator_tree; + t0, tf, + erase_sol, + tstops, + saveat, + reinit_callbacks, + reinit_retcode + ) +end + +function subreinit!( + f, + u0, + solution_indices, + subintegrator::DiffEqBase.DEIntegrator; + kwargs... + ) + return DiffEqBase.reinit!(subintegrator, u0[solution_indices]; kwargs...) +end + +@unroll function subreinit!( + f, + u0, + solution_indices, + subintegrators::Tuple; + kwargs... + ) + i = 1 + @unroll for subintegrator in subintegrators + subreinit!(get_operator(f, i), u0, f.solution_indices[i], subintegrator; kwargs...) + i += 1 + end +end + +function OrdinaryDiffEqCore.handle_tstop!(integrator::OperatorSplittingIntegrator) + if SciMLBase.has_tstop(integrator) + tdir_t = tdir(integrator) * integrator.t + tdir_tstop = SciMLBase.first_tstop(integrator) + if tdir_t == tdir_tstop + while tdir_t == tdir_tstop #remove all redundant copies + res = SciMLBase.pop_tstop!(integrator) + SciMLBase.has_tstop(integrator) ? + (tdir_tstop = SciMLBase.first_tstop(integrator)) : break + end + notify_integrator_hit_tstop!(integrator) + elseif tdir_t > tdir_tstop + if !integrator.dtchangeable + SciMLBase.change_t_via_interpolation!( + integrator, + tdir(integrator) * + SciMLBase.pop_tstop!(integrator), Val{true} + ) + notify_integrator_hit_tstop!(integrator) + else + error("Something went wrong. Integrator stepped past tstops but the algorithm was dtchangeable. Please report this error.") + end + end + end + return nothing +end + +notify_integrator_hit_tstop!(integrator::OperatorSplittingIntegrator) = nothing + +is_first_iteration(integrator::OperatorSplittingIntegrator) = integrator.iter == 0 +increment_iteration(integrator::OperatorSplittingIntegrator) = integrator.iter += 1 + +# Controller interface +function reject_step!(integrator::OperatorSplittingIntegrator) + OrdinaryDiffEqCore.increment_reject!(integrator.stats) + return reject_step!(integrator, integrator.controller) +end +function reject_step!(integrator::OperatorSplittingIntegrator, controller) + integrator.u .= integrator.uprev + if !integrator.force_stepfail + step_reject_controller!(integrator, controller, integrator.alg) + end + # We need to roll-back the sub-integrators + return prepare_subintegrators_to_redo_step!(integrator) +end +function reject_step!(integrator::OperatorSplittingIntegrator, ::Nothing) + return if length(integrator.uprev) == 0 + error("Cannot roll back integrator. Aborting time integration step at $(integrator.t).") + end +end + +# Solution looping interface +function should_accept_step(integrator::OperatorSplittingIntegrator) + if integrator.force_stepfail || integrator.isout + return false + end + return should_accept_step(integrator, integrator.controller) +end +function should_accept_step(integrator::OperatorSplittingIntegrator, ::Nothing) + return !(integrator.force_stepfail) +end +function accept_step!(integrator::OperatorSplittingIntegrator) + OrdinaryDiffEqCore.increment_accept!(integrator.stats) + return accept_step!(integrator, integrator.cache, integrator.controller) +end +function accept_step!(integrator::OperatorSplittingIntegrator, cache, controller) + return store_previous_info!(integrator) +end +function store_previous_info!(integrator::OperatorSplittingIntegrator) + return if length(integrator.uprev) > 0 # Integrator can rollback + update_uprev!(integrator) + end +end + +function update_uprev!(integrator::OperatorSplittingIntegrator) + SciMLBase.recursivecopy!(integrator.uprev, integrator.u) + return nothing +end + +function step_header!(integrator::OperatorSplittingIntegrator) + # Accept or reject the step + if !is_first_iteration(integrator) + if should_accept_step(integrator) + accept_step!(integrator) + else # Step should be rejected and hence repeated + reject_step!(integrator) + end + elseif integrator.u_modified # && integrator.iter == 0 + update_uprev!(integrator) + end + + # Before stepping we might need to adjust the dt + increment_iteration(integrator) + # OrdinaryDiffEqCore.choose_algorithm!(integrator, integrator.cache) + OrdinaryDiffEqCore.fix_dt_at_bounds!(integrator) + OrdinaryDiffEqCore.modify_dt_for_tstops!(integrator) + return integrator.force_stepfail = false +end + +function footer_reset_flags!(integrator) + return integrator.u_modified = false +end +function setup_validity_flags!(integrator, t_next) + return integrator.isout = false #integrator.opts.isoutofdomain(integrator.u, integrator.p, t_next) +end +function fix_solution_buffer_sizes!(integrator, sol) + resize!(integrator.sol.t, integrator.saveiter) + resize!(integrator.sol.u, integrator.saveiter) + return if !(integrator.sol isa SciMLBase.DAESolution) + resize!(integrator.sol.k, integrator.saveiter_dense) + end +end + +function step_footer!(integrator::OperatorSplittingIntegrator) + ttmp = integrator.t + tdir(integrator) * integrator.dt + + footer_reset_flags!(integrator) + setup_validity_flags!(integrator, ttmp) + + if should_accept_step(integrator) + integrator.last_step_failed = false + integrator.tprev = integrator.t + integrator.t = ttmp #OrdinaryDiffEqCore.fixed_t_for_floatingpoint_error!(integrator, ttmp) + # OrdinaryDiffEqCore.handle_callbacks!(integrator) + step_accept_controller!(integrator) # Noop for non-adaptive algorithms + elseif integrator.force_stepfail # Rejected by solver + if SciMLBase.isadaptive(integrator) + step_reject_controller!(integrator) + OrdinaryDiffEqCore.post_newton_controller!(integrator, integrator.alg) + elseif integrator.dtchangeable # Non-adaptive but can change dt + integrator.dt /= integrator.opts.failfactor + elseif integrator.last_step_failed + return + end + integrator.last_step_failed = true + end + + # integration_monitor_step(integrator) + + return nothing +end + +# called by DiffEqBase.solve +function DiffEqBase.__solve( + prob::OperatorSplittingProblem, + alg::AbstractOperatorSplittingAlgorithm, args...; kwargs... + ) + integrator = DiffEqBase.__init(prob, alg, args...; kwargs...) + return DiffEqBase.solve!(integrator) +end + +# either called directly (after init), or by DiffEqBase.solve (via __solve) +function DiffEqBase.solve!(integrator::OperatorSplittingIntegrator) + while !isempty(integrator.tstops) + while tdir(integrator) * integrator.t < SciMLBase.first_tstop(integrator) + step_header!(integrator) + @timeit_debug "check_error" DiffEqBase.check_error!(integrator) ∉ ( + SciMLBase.ReturnCode.Success, SciMLBase.ReturnCode.Default, + )&&return + __step!(integrator) + step_footer!(integrator) + if !SciMLBase.has_tstop(integrator) + break + end + end + OrdinaryDiffEqCore.handle_tstop!(integrator) + end + OrdinaryDiffEqCore.postamble!(integrator) + if integrator.sol.retcode != SciMLBase.ReturnCode.Default + return integrator.sol + end + return integrator.sol = SciMLBase.solution_new_retcode( + integrator.sol, SciMLBase.ReturnCode.Success + ) +end + +function DiffEqBase.step!(integrator::OperatorSplittingIntegrator) + @timeit_debug "step!" if integrator.advance_to_tstop + tstop = first_tstop(integrator) + while !reached_tstop(integrator, tstop) + step_header!(integrator) + @timeit_debug "check_error" DiffEqBase.check_error!(integrator) ∉ ( + SciMLBase.ReturnCode.Success, SciMLBase.ReturnCode.Default, + )&&return + __step!(integrator) + step_footer!(integrator) + if !SciMLBase.has_tstop(integrator) + break + end + end + else + step_header!(integrator) + @timeit_debug "check_error" DiffEqBase.check_error!(integrator) ∉ ( + SciMLBase.ReturnCode.Success, SciMLBase.ReturnCode.Default, + )&&return + __step!(integrator) + step_footer!(integrator) + while !should_accept_step(integrator) + step_header!(integrator) + @timeit_debug "check_error" DiffEqBase.check_error!(integrator) ∉ ( + SciMLBase.ReturnCode.Success, SciMLBase.ReturnCode.Default, + )&&return + __step!(integrator) + step_footer!(integrator) + end + end + return OrdinaryDiffEqCore.handle_tstop!(integrator) +end + +function SciMLBase.check_error(integrator::OperatorSplittingIntegrator) + if !SciMLBase.successful_retcode(integrator.sol) && + integrator.sol.retcode != SciMLBase.ReturnCode.Default + return integrator.sol.retcode + end + + verbose = true # integrator.opts.verbose + + if DiffEqBase.NAN_CHECK(integrator.dtcache) || DiffEqBase.NAN_CHECK(integrator.dt) # replace with https://github.com/SciML/OrdinaryDiffEq.jl/blob/373a8eec8024ef1acc6c5f0c87f479aa0cf128c3/lib/OrdinaryDiffEqCore/src/iterator_interface.jl#L5-L6 after moving to sciml integrators + if verbose + @warn("NaN dt detected. Likely a NaN value in the state, parameters, or derivative value caused this outcome.") + end + return SciMLBase.ReturnCode.DtNaN + end + + return check_error_subintegrators(integrator, integrator.subintegrator_tree) +end + +function check_error_subintegrators(integrator, subintegrator_tree::Tuple) + for subintegrator in subintegrator_tree + retcode = check_error_subintegrators(integrator, subintegrator) + if !SciMLBase.successful_retcode(retcode) && retcode != SciMLBase.ReturnCode.Default + return retcode + end + end + return integrator.sol.retcode +end + +function check_error_subintegrators(integrator, subintegrator::SciMLBase.DEIntegrator) + return SciMLBase.check_error(subintegrator) +end + +function DiffEqBase.step!(integrator::OperatorSplittingIntegrator, dt, stop_at_tdt = false) + return @timeit_debug "step!" begin + # OridinaryDiffEq lets dt be negative if tdir is -1, but that's inconsistent + dt <= zero(dt) && error("dt must be positive") + stop_at_tdt && !integrator.dtchangeable && + error("Cannot stop at t + dt if dtchangeable is false") + tnext = integrator.t + tdir(integrator) * dt + stop_at_tdt && DiffEqBase.add_tstop!(integrator, tnext) + while !reached_tstop(integrator, tnext, stop_at_tdt) + step_header!(integrator) + @timeit_debug "check_error" DiffEqBase.check_error!(integrator) ∉ ( + SciMLBase.ReturnCode.Success, SciMLBase.ReturnCode.Default, + )&&return + __step!(integrator) + step_footer!(integrator) + end + end +end + +function setup_u(prob::OperatorSplittingProblem, solver, alias_u0) + if alias_u0 + return prob.u0 + else + return OrdinaryDiffEqCore.recursivecopy(prob.u0) + end +end + +# TimeChoiceIterator API +@inline function DiffEqBase.get_tmp_cache(integrator::OperatorSplittingIntegrator) + # DiffEqBase.get_tmp_cache(integrator, integrator.alg, integrator.cache) + return (integrator.tmp,) +end +# @inline function DiffEqBase.get_tmp_cache(integrator::OperatorSplittingIntegrator, ::AbstractOperatorSplittingAlgorithm, cache) +# return (cache.tmp,) +# end +# Interpolation +# TODO via https://github.com/SciML/SciMLBase.jl/blob/master/src/interpolation.jl +function linear_interpolation!(y, t, y1, y2, t1, t2) + return y .= y1 + (t - t1) * (y2 - y1) / (t2 - t1) +end +function (integrator::OperatorSplittingIntegrator)(tmp, t) + return linear_interpolation!( + tmp, t, integrator.uprev, integrator.u, integrator.tprev, integrator.t + ) +end + +""" + stepsize_controller!(::OperatorSplittingIntegrator) + +Updates the controller using the current state of the integrator if the operator splitting algorithm is adaptive. +""" +@inline function stepsize_controller!(integrator::OperatorSplittingIntegrator) + DiffEqBase.isadaptive(integrator) || return nothing + return stepsize_controller!(integrator, integrator.controller, integrator.alg) +end + +""" + step_accept_controller!(::OperatorSplittingIntegrator) + +Updates `dtcache` of the integrator if the step is accepted and the operator splitting algorithm is adaptive. +""" +@inline function step_accept_controller!(integrator::OperatorSplittingIntegrator) + DiffEqBase.isadaptive(integrator) || return nothing + return step_accept_controller!(integrator, integrator.controller, integrator.alg) +end + +""" + step_reject_controller!(::OperatorSplittingIntegrator) + +Updates `dtcache` of the integrator if the step is rejected and the the operator splitting algorithm is adaptive. +""" +@inline function step_reject_controller!(integrator::OperatorSplittingIntegrator) + DiffEqBase.isadaptive(integrator) || return nothing + return step_reject_controller!(integrator, integrator.controller, integrator.alg) +end + +# helper functions for dealing with time-reversed integrators in the same way +# that OrdinaryDiffEq.jl does +tdir(integrator) = integrator.tstops.ordering isa DataStructures.FasterForward ? 1 : -1 +is_past_t(integrator, t) = tdir(integrator) * (t - integrator.t) ≤ zero(integrator.t) +function reached_tstop(integrator, tstop, stop_at_tstop = integrator.dtchangeable) + if stop_at_tstop + integrator.t > tstop && + error("Integrator missed stop at $tstop (current time=$(integrator.t)). Aborting.") + return integrator.t == tstop # Check for exact hit + else #!stop_at_tstop + return is_past_t(integrator, tstop) + end +end + +# Dunno stuff +function SciMLBase.done(integrator::OperatorSplittingIntegrator) + if !( + integrator.sol.retcode in ( + SciMLBase.ReturnCode.Default, SciMLBase.ReturnCode.Success, + ) + ) + return true + elseif isempty(integrator.tstops) + DiffEqBase.postamble!(integrator) + return true + end + return false +end + +function DiffEqBase.postamble!(integrator::OperatorSplittingIntegrator) + return DiffEqBase.finalize!(integrator.callback, integrator.u, integrator.t, integrator) +end + +function __step!(integrator) + tnext = integrator.t + integrator.dt + synchronize_subintegrator_tree!(integrator) + advance_solution_to!(integrator, tnext) + return stepsize_controller!(integrator) +end + +# solvers need to define this interface +function advance_solution_to!(integrator::OperatorSplittingIntegrator, tnext) + return advance_solution_to!(integrator, integrator.cache, tnext) +end + +function advance_solution_to!( + outer_integrator::OperatorSplittingIntegrator, + integrator::DiffEqBase.DEIntegrator, solution_indices, sync, cache, tend + ) + dt = tend - integrator.t + return SciMLBase.step!(integrator, dt, true) +end + +# ----------------------------------- SciMLBase.jl Integrator Interface ------------------------------------ +SciMLBase.has_stats(::OperatorSplittingIntegrator) = true + +SciMLBase.has_tstop(integrator::OperatorSplittingIntegrator) = !isempty(integrator.tstops) +SciMLBase.first_tstop(integrator::OperatorSplittingIntegrator) = first(integrator.tstops) +SciMLBase.pop_tstop!(integrator::OperatorSplittingIntegrator) = pop!(integrator.tstops) + +DiffEqBase.get_dt(integrator::OperatorSplittingIntegrator) = integrator.dt +function set_dt!(integrator::OperatorSplittingIntegrator, dt) + # TODO: figure out interface for recomputing other objects (linear operators, etc) + dt <= zero(dt) && error("dt must be positive") + return integrator.dt = dt +end + +function DiffEqBase.add_tstop!(integrator::OperatorSplittingIntegrator, t) + is_past_t(integrator, t) && + error("Cannot add a tstop at $t because that is behind the current \ + integrator time $(integrator.t)") + return push!(integrator.tstops, t) +end + +function DiffEqBase.add_saveat!(integrator::OperatorSplittingIntegrator, t) + is_past_t(integrator, t) && + error("Cannot add a saveat point at $t because that is behind the \ + current integrator time $(integrator.t)") + return push!(integrator.saveat, t) +end + +# not sure what this should do? +# defined as default initialize: https://github.com/SciML/DiffEqBase.jl/blob/master/src/callbacks.jl#L3 +DiffEqBase.u_modified!(i::OperatorSplittingIntegrator, bool) = nothing + +function synchronize_subintegrator_tree!(integrator::OperatorSplittingIntegrator) + return synchronize_subintegrator!(integrator.subintegrator_tree, integrator) +end + +@unroll function synchronize_subintegrator!( + subintegrator_tree::Tuple, integrator::OperatorSplittingIntegrator + ) + @unroll for subintegrator in subintegrator_tree + synchronize_subintegrator!(subintegrator, integrator) + end +end + +function synchronize_subintegrator!( + subintegrator::SciMLBase.DEIntegrator, integrator::OperatorSplittingIntegrator + ) + @unpack t, dt = integrator + @assert subintegrator.t == t "Integrators out of sync. The outer integrator is at $t, but inner integrator is at $(subintegrator.t)" + return if !DiffEqBase.isadaptive(subintegrator) + SciMLBase.set_proposed_dt!(subintegrator, dt) + end +end + +function advance_solution_to!( + integrator::OperatorSplittingIntegrator, + cache::AbstractOperatorSplittingCache, tnext::Number + ) + return advance_solution_to!( + integrator, integrator.subintegrator_tree, integrator.solution_index_tree, + integrator.synchronizer_tree, cache, tnext + ) +end + +# Dispatch for tree node construction +function build_subintegrator_tree_with_cache( + prob::OperatorSplittingProblem, alg::AbstractOperatorSplittingAlgorithm, + uprevouter::AbstractVector, uouter::AbstractVector, + solution_indices, + t0, dt, tf, + tstops, saveat, d_discontinuities, callback, + adaptive, verbose + ) + (; f, p) = prob + subintegrator_tree_with_caches = ntuple( + i -> build_subintegrator_tree_with_cache( + prob, + alg.inner_algs[i], + get_operator(f, i), + p[i], + uprevouter, uouter, + f.solution_indices[i], + t0, dt, tf, + tstops, saveat, d_discontinuities, callback, + adaptive, verbose + ), + length(f.functions) + ) + + subintegrator_tree = ntuple( + i -> subintegrator_tree_with_caches[i][1], length(f.functions) + ) + caches = ntuple(i -> subintegrator_tree_with_caches[i][2], length(f.functions)) + + # TODO fix mixed device type problems we have to be smarter + return subintegrator_tree, + init_cache( + f, alg; + uprev = uprevouter, u = uouter, alias_u = true, + inner_caches = caches + ) +end + +function build_subintegrator_tree_with_cache( + prob::OperatorSplittingProblem, alg::AbstractOperatorSplittingAlgorithm, + f::GenericSplitFunction, p::Tuple, + uprevouter::AbstractVector, uouter::AbstractVector, + solution_indices, + t0, dt, tf, + tstops, saveat, d_discontinuities, callback, + adaptive, verbose, + save_end = false, + controller = nothing + ) + subintegrator_tree_with_caches = ntuple( + i -> build_subintegrator_tree_with_cache( + prob, + alg.inner_algs[i], + get_operator(f, i), + p[i], + uprevouter, uouter, + f.solution_indices[i], + t0, dt, tf, + tstops, saveat, d_discontinuities, callback, + adaptive, verbose + ), + length(f.functions) + ) + + subintegrator_tree = first.(subintegrator_tree_with_caches) + inner_caches = last.(subintegrator_tree_with_caches) + + # TODO fix mixed device type problems we have to be smarter + uprev = @view uprevouter[solution_indices] + u = @view uouter[solution_indices] + return subintegrator_tree, + init_cache( + f, alg; + uprev = uprev, u = u, + inner_caches = inner_caches + ) +end + +function build_subintegrator_tree_with_cache( + prob::OperatorSplittingProblem, alg::SciMLBase.AbstractODEAlgorithm, + f::F, p::P, + uprevouter::S, uouter::S, + solution_indices, + t0::T, dt::T, tf::T, + tstops, saveat, d_discontinuities, callback, + adaptive, verbose, + save_end = false, + controller = nothing + ) where {S, T, P, F} + uprev = @view uprevouter[solution_indices] + u = @view uouter[solution_indices] + + integrator = DiffEqBase.__init( + SciMLBase.ODEProblem(f, u, (t0, min(t0 + dt, tf)), p), + alg; + dt, + saveat = (), + d_discontinuities, + save_everystep = false, + advance_to_tstop = false, + adaptive, + controller, + verbose + ) + + return integrator, integrator.cache +end + +function forward_sync_subintegrator!( + outer_integrator::OperatorSplittingIntegrator, subintegrator_tree::Tuple, + solution_indices::Tuple, synchronizers::Tuple + ) + return nothing +end +function backward_sync_subintegrator!( + outer_integrator::OperatorSplittingIntegrator, + subintegrator_tree::Tuple, solution_indices::Tuple, synchronizer::Tuple + ) + return nothing +end diff --git a/test/config_tree.jl b/test/config_tree.jl new file mode 100644 index 0000000..020c00b --- /dev/null +++ b/test/config_tree.jl @@ -0,0 +1,395 @@ +using OrdinaryDiffEqOperatorSplitting +import OrdinaryDiffEqOperatorSplitting as OS +using Test + +import DiffEqBase: DiffEqBase, ODEFunction +import SciMLBase: ReturnCode +using OrdinaryDiffEqLowOrderRK +using OrdinaryDiffEqTsit5 + +# --------------------------------------------------------------------------- +# A nested splitting function: f = (f1, (f3, f3)) +# --------------------------------------------------------------------------- +ode1(du, u, p, t) = @. du = -0.1u +function ode3(du, u, p, t) + du[1] = -0.005u[2] + return du[2] = -0.005u[1] +end +f1 = ODEFunction(ode1) +f3 = ODEFunction(ode3) + +f_inner = GenericSplitFunction((f3, f3), ([1, 2], [1, 2])) +f_flat = GenericSplitFunction((f1, f3), ([1, 2, 3], [1, 3])) +f_nested = GenericSplitFunction((f1, f_inner), ([1, 2, 3], [1, 3])) + +@testset "SplitNode" begin + @testset "minting" begin + @test f_nested[].path === () + @test f_nested[].object === f_nested + @test f_nested[1].path === (1,) + @test f_nested[1].object === f1 + @test f_nested[2, 1].path === (2, 1) + @test f_nested[2, 1].object === f3 + + # stepwise descent and varargs descent agree + @test f_nested[2][1].path === f_nested[2, 1].path + @test f_nested[][2, 1].path === f_nested[2, 1].path + end + + @testset "invalid addresses" begin + @test_throws ArgumentError f_nested[3] # out of range + @test_throws ArgumentError f_nested[0] + @test_throws ArgumentError f_nested[1, 1] # descends into a leaf + @test_throws ArgumentError f_nested[2, 3] + end + + @testset "resolution against the mirroring trees" begin + alg = LieTrotterGodunov((Euler(), StrangMarchuk((Tsit5(), Euler())))) + + @test f_nested[f_nested[2, 1]] === f3 + @test f_nested[f_nested[2]] === f_inner + @test f_nested[f_nested[]] === f_nested + @test OS.get_operator(f_nested, f_nested[2, 1]) === f3 + + @test alg[f_nested[1]] === alg.inner_algs[1] + @test alg[f_nested[2]] === alg.inner_algs[2] + @test alg[f_nested[2, 1]] === alg.inner_algs[2].inner_algs[1] + @test alg[2, 1] === alg.inner_algs[2].inner_algs[1] + + u0 = [0.7, 0.9, 0.5] + prob = OperatorSplittingProblem(f_nested, u0, (0.0, 0.1)) + integrator = DiffEqBase.init(prob, alg; dt = 0.01, adaptive = false) + + @test integrator[f_nested[]] === integrator + @test integrator[f_nested[1]] === integrator.child_subintegrators[1] + @test integrator[f_nested[2]] === integrator.child_subintegrators[2] + @test integrator[f_nested[2, 1]] === + integrator.child_subintegrators[2].child_subintegrators[1] + + # SciMLBase's symbolic indexing must keep working: integrator[i] is the + # i-th state component, not the i-th subintegrator. + @test integrator[1] == integrator.u[1] + end + + @test sprint(show, f_nested[2, 1]) == "f[2, 1]" +end + +@testset "TreeOption" begin + @testset "construction mirrors the function tree" begin + opt = TreeOption(f_nested, 1.0e-2) + @test opt[] == 1.0e-2 + @test opt[1] == 1.0e-2 + @test opt[2] == 1.0e-2 + @test opt[2, 1] == 1.0e-2 + @test opt[2, 2] == 1.0e-2 + @test opt isa TreeOption{Float64} + + @test OS.structure_matches(opt, f_nested) + @test !OS.structure_matches(opt, f_flat) + @test OS.structure_matches(TreeOption(f_flat, 1.0e-2), f_flat) + end + + @testset "plain assignment writes a single node" begin + opt = TreeOption(f_nested, 1.0e-2) + opt[2] = 1.0e-4 + @test opt[] == 1.0e-2 + @test opt[1] == 1.0e-2 + @test opt[2] == 1.0e-4 + @test opt[2, 1] == 1.0e-2 # untouched + @test opt[2, 2] == 1.0e-2 + + opt[2, 1] = 1.0e-5 + @test opt[2, 1] == 1.0e-5 + @test opt[2, 2] == 1.0e-2 + + opt[] = 0.5 + @test opt[] == 0.5 + @test opt[1] == 1.0e-2 + end + + @testset "addressing by SplitNode" begin + opt = TreeOption(f_nested, 1.0e-2) + opt[f_nested[2, 1]] = 1.0e-5 + @test opt[2, 1] == 1.0e-5 + @test opt[f_nested[2, 1]] == 1.0e-5 + + opt[f_nested[2]] .= 3.0e-4 + @test opt[2] == 3.0e-4 + @test opt[2, 1] == 3.0e-4 + @test opt[2, 2] == 3.0e-4 + @test opt[1] == 1.0e-2 + end + + @testset "broadcast assignment writes the subtree" begin + opt = TreeOption(f_nested, 1.0e-2) + opt[2] .= 1.0e-4 + @test opt[] == 1.0e-2 + @test opt[1] == 1.0e-2 + @test opt[2] == 1.0e-4 + @test opt[2, 1] == 1.0e-4 + @test opt[2, 2] == 1.0e-4 + + # a more specific write after a subtree write survives + opt[2, 1] = 1.0e-5 + @test opt[2, 1] == 1.0e-5 + @test opt[2, 2] == 1.0e-4 + + # ... and a subtree write after it clobbers it again (order matters) + opt[2] .= 2.0e-4 + @test opt[2, 1] == 2.0e-4 + + opt .= 7.0 + @test opt[] == 7.0 + @test opt[1] == 7.0 + @test opt[2, 2] == 7.0 + end + + @testset "rejected broadcasts" begin + opt = TreeOption(f_nested, 1.0e-2) + + # Update-assignment reads one node but writes a subtree; there is no + # unsurprising reading of that, so it must not silently do something. + @test_throws ArgumentError opt[2] .*= 2 + @test_throws ArgumentError opt[2] .+= 1 + @test_throws ArgumentError opt .*= 2 + @test_throws ArgumentError opt[2] .= [1.0, 2.0] + @test_throws ArgumentError opt[2] .= (1.0, 2.0) + + # nothing was written by the rejected calls + @test opt[2] == 1.0e-2 + @test opt[2, 1] == 1.0e-2 + + # the undotted right hand side the error message suggests does work + opt[2] .= 2 * opt[2] + @test opt[2] == 2.0e-2 + @test opt[2, 1] == 2.0e-2 + end + + @testset "element type" begin + opt = TreeOption(f_nested, true) + @test opt isa TreeOption{Bool} + opt[2] = false + @test opt[2] === false + @test_throws ArgumentError opt[2] = "yes" + + # widening conversions that Julia performs anyway are still fine + num = TreeOption(f_nested, 1.0) + num[1] = 2 + @test num[1] === 2.0 + + # mixed types need the explicit form + ctrl = TreeOption{Any}(f_nested, nothing) + @test ctrl isa TreeOption{Any} + @test ctrl[1] === nothing + ctrl[1] = :something + @test ctrl[1] === :something + ctrl[2] .= nothing + @test ctrl[2, 1] === nothing + end + + @testset "invalid addresses" begin + opt = TreeOption(f_nested, 1.0e-2) + @test_throws ArgumentError opt[3] + @test_throws ArgumentError opt[1, 1] # [1] mirrors a leaf + @test_throws ArgumentError opt[2, 3] + @test_throws ArgumentError opt[3] = 1.0 + @test_throws ArgumentError opt[1, 1] .= 1.0 + end + + @testset "show" begin + opt = TreeOption(f_nested, 1.0e-2) + opt[2] .= 1.0e-4 + str = sprint(show, MIME"text/plain"(), opt) + @test occursin("TreeOption{Float64}", str) + @test occursin("[] => 0.01", str) + @test occursin("[2, 1] => 0.0001", str) + end +end + +# --------------------------------------------------------------------------- +# Per-node settings reaching the integrator tree +# +# f1 + f3 + f3 == f1 + f2 is linear, so the exact solution is available for +# checking that a multi-rate configuration still integrates the right problem. +# --------------------------------------------------------------------------- +trueA = [-0.1 0.0 0.0; 0.0 -0.1 0.0; 0.0 0.0 -0.1] +trueB = [0.0 0.0 -0.01; 0.0 0.0 0.0; -0.01 0.0 0.0] +u0 = [0.7611944793397108, 0.9059606424982555, 0.5755174199139956] +tspan = (0.0, 1.0) +trueu = exp((tspan[2] - tspan[1]) * (trueA + trueB)) * u0 + +prob = OperatorSplittingProblem(f_nested, u0, tspan) +# A power of two, so that the subcycled step sizes below are exactly +# representable and the sub-step counts are deterministic. +dt_outer = 2.0^-7 +nsteps = round(Int, (tspan[2] - tspan[1]) / dt_outer) + +@testset "per-node configuration" begin + @testset "uniform dt is unchanged" begin + alg = LieTrotterGodunov((Euler(), LieTrotterGodunov((Euler(), Euler())))) + scalar = DiffEqBase.init(prob, alg; dt = dt_outer, adaptive = false) + tree = DiffEqBase.init( + prob, alg; dt = TreeOption(f_nested, dt_outer), adaptive = false + ) + DiffEqBase.solve!(scalar) + DiffEqBase.solve!(tree) + @test scalar.u == tree.u + @test scalar.iter == tree.iter + end + + @testset "multi-rate: a subtree subcycles" begin + alg = LieTrotterGodunov((Euler(), LieTrotterGodunov((Euler(), Euler())))) + dt = TreeOption(f_nested, dt_outer) + dt[f_nested[2]] .= dt_outer / 4 # the nested split and its two leaves + + integrator = DiffEqBase.init(prob, alg; dt, adaptive = false) + + @test integrator.dt == dt_outer + @test integrator[f_nested[1]].dt == dt_outer + @test integrator[f_nested[2]].dt == dt_outer / 4 + @test integrator[f_nested[2, 1]].dt == dt_outer / 4 + + DiffEqBase.solve!(integrator) + @test integrator.sol.retcode == ReturnCode.Success + @test integrator.t ≈ tspan[2] + @test integrator.iter == nsteps + + # The outer integrator hands the nested node an interval of dt_outer, which + # it covers in four sub-steps -- and its own children follow along. + @test integrator[f_nested[1]].iter == nsteps + @test integrator[f_nested[2]].iter == 4 * nsteps + @test integrator[f_nested[2, 1]].iter == 4 * nsteps + + # Everything still lands on the same time points. + @test integrator[f_nested[1]].t ≈ tspan[2] + @test integrator[f_nested[2]].t ≈ tspan[2] + @test integrator[f_nested[2, 1]].t ≈ tspan[2] + + @test isapprox(integrator.u, trueu, atol = 1.0e-4) + end + + @testset "multi-rate: a single leaf subcycles" begin + alg = LieTrotterGodunov((Euler(), LieTrotterGodunov((Euler(), Euler())))) + dt = TreeOption(f_nested, dt_outer) + dt[2, 1] = dt_outer / 8 # this leaf only + + integrator = DiffEqBase.init(prob, alg; dt, adaptive = false) + DiffEqBase.solve!(integrator) + + @test integrator.iter == nsteps + @test integrator[f_nested[2]].iter == nsteps + @test integrator[f_nested[2, 1]].iter == 8 * nsteps + @test integrator[f_nested[2, 2]].iter == nsteps + @test isapprox(integrator.u, trueu, atol = 1.0e-4) + end + + @testset "a sub-dt that does not divide the interval still lands on it" begin + # 1/3 of the outer step is not representable, so the leaf covers the last + # sliver of each interval with an extra ulp-sized step. That costs steps but + # must not cost accuracy or leave the tree out of sync. + alg = LieTrotterGodunov((Euler(), LieTrotterGodunov((Euler(), Euler())))) + dt = TreeOption(f_nested, dt_outer) + dt[f_nested[2]] .= dt_outer / 3 + + integrator = DiffEqBase.init(prob, alg; dt, adaptive = false) + DiffEqBase.solve!(integrator) + + @test integrator.sol.retcode == ReturnCode.Success + @test integrator.iter == nsteps + @test integrator.t ≈ tspan[2] + @test integrator[f_nested[2]].t ≈ tspan[2] + @test integrator[f_nested[2, 1]].t ≈ tspan[2] + @test integrator[f_nested[2]].iter ≥ 3 * nsteps + @test isapprox(integrator.u, trueu, atol = 1.0e-4) + end + + @testset "mixed adaptivity" begin + alg = LieTrotterGodunov((Euler(), LieTrotterGodunov((Tsit5(), Tsit5())))) + adaptive = TreeOption(f_nested, false) + adaptive[2, 1] = true # leaves only, the splitting nodes + adaptive[2, 2] = true # stay non-adaptive + + integrator = DiffEqBase.init(prob, alg; dt = dt_outer, adaptive) + + @test integrator.opts.adaptive == false + @test integrator[f_nested[1]].opts.adaptive == false + @test integrator[f_nested[2]].opts.adaptive == false + @test integrator[f_nested[2, 1]].opts.adaptive == true + @test integrator[f_nested[2, 2]].opts.adaptive == true + + DiffEqBase.solve!(integrator) + @test integrator.sol.retcode == ReturnCode.Success + @test integrator.t ≈ tspan[2] + @test integrator.iter == nsteps # the splitting step is untouched + @test isapprox(integrator.u, trueu, atol = 1.0e-4) + end + + @testset "inner integrator options travel to the leaves" begin + alg = LieTrotterGodunov((Euler(), LieTrotterGodunov((Tsit5(), Tsit5())))) + adaptive = TreeOption(f_nested, false) + adaptive[2, 1] = true + adaptive[2, 2] = true + reltol = TreeOption(f_nested, 1.0e-3) + reltol[2, 1] = 1.0e-9 + + integrator = DiffEqBase.init( + prob, alg; dt = dt_outer, adaptive, reltol, dtmin = 1.0e-12 + ) + + # `reltol` means nothing to a splitting node, so it is handed to the leaves + @test integrator[f_nested[2, 1]].opts.reltol == 1.0e-9 + @test integrator[f_nested[2, 2]].opts.reltol == 1.0e-3 + # ... while `dtmin` is understood at every level + @test integrator.opts.dtmin == 1.0e-12 + @test integrator[f_nested[2]].opts.dtmin == 1.0e-12 + + DiffEqBase.solve!(integrator) + @test integrator.sol.retcode == ReturnCode.Success + @test isapprox(integrator.u, trueu, atol = 1.0e-4) + end + + @testset "asking a non-adaptive splitting node to be adaptive warns" begin + alg = LieTrotterGodunov((Euler(), LieTrotterGodunov((Euler(), Euler())))) + adaptive = TreeOption(f_nested, false) + adaptive[2] = true # a splitting node, and LTG is not adaptive + @test_logs (:warn, r"operator \[2\] is not adaptive") DiffEqBase.init( + prob, alg; dt = dt_outer, adaptive + ) + end + + @testset "rejected configurations" begin + alg = LieTrotterGodunov((Euler(), LieTrotterGodunov((Euler(), Euler())))) + + # built for a differently shaped splitting function + wrong = TreeOption(f_flat, dt_outer) + @test_throws ArgumentError DiffEqBase.init(prob, alg; dt = wrong) + + negative = TreeOption(f_nested, dt_outer) + negative[2, 1] = -1.0e-3 + @test_throws ErrorException DiffEqBase.init(prob, alg; dt = negative) + end + + @testset "reinit! keeps the per-node configuration" begin + alg = LieTrotterGodunov((Euler(), LieTrotterGodunov((Euler(), Euler())))) + dt = TreeOption(f_nested, dt_outer) + dt[f_nested[2]] .= dt_outer / 4 + + integrator = DiffEqBase.init(prob, alg; dt, adaptive = false) + DiffEqBase.solve!(integrator) + ufinal = copy(integrator.u) + + DiffEqBase.reinit!(integrator) + @test integrator.dt == dt_outer + @test integrator[f_nested[2]].dt == dt_outer / 4 + @test integrator[f_nested[2, 1]].dt == dt_outer / 4 + DiffEqBase.solve!(integrator) + @test integrator.u ≈ ufinal + @test integrator[f_nested[2]].iter == 4 * nsteps + + # An explicit scalar reconfigures the whole tree, as it would at `init`. + DiffEqBase.reinit!(integrator; dt = dt_outer) + @test integrator.dt == dt_outer + @test integrator[f_nested[2]].dt == dt_outer + @test integrator[f_nested[2, 1]].dt == dt_outer + end +end diff --git a/test/qa/qa.jl b/test/qa/qa.jl index c360a22..901b7c4 100644 --- a/test/qa/qa.jl +++ b/test/qa/qa.jl @@ -29,6 +29,9 @@ run_qa( :fix_dt_at_bounds!, :handle_tstop!, :increment_accept!, # OrdinaryDiffEqCore :increment_reject!, :initialize_d_discontinuities, :initialize_saveat, # OrdinaryDiffEqCore :initialize_tstops, :post_newton_controller!, :timedepentdtmin, # OrdinaryDiffEqCore + :promote_tspan, # SciMLBase + # Broadcast extension points a TreeOption implements (src/config_tree.jl). + :Broadcasted, :broadcastable, :dotview, :materialize!, # Base ), ), all_explicit_imports_are_public = (; From 61a742ed7ce41bf38684026dd96e0c86620b3903 Mon Sep 17 00:00:00 2001 From: termi-official <9196588+termi-official@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:56:03 +0200 Subject: [PATCH 2/2] Remove legacy file --- src/integrator_node.jl | 792 ----------------------------------------- 1 file changed, 792 deletions(-) delete mode 100644 src/integrator_node.jl diff --git a/src/integrator_node.jl b/src/integrator_node.jl deleted file mode 100644 index 6c29f17..0000000 --- a/src/integrator_node.jl +++ /dev/null @@ -1,792 +0,0 @@ -""" - OperatorSplittingIntegratorNode <: AbstractODEIntegrator - -A variant of [`ODEIntegrator`](https://github.com/SciML/OrdinaryDiffEq.jl/blob/6ec5a55bda26efae596bf99bea1a1d729636f412/src/integrators/type.jl#L77-L123) to perform opeartor splitting. -""" -mutable struct OperatorSplittingIntegratorNode{ - fType, - algType, - uType, - tType, - pType, - heapType, - tstopsType, - saveatType, - callbackType, - cacheType, - solType, - subintTreeType, - solidxTreeType, - syncTreeType, - controllerType, - optionsType, - } <: DiffEqBase.AbstractODEIntegrator{algType, true, uType, tType} - const f::fType - const alg::algType - u::uType # Master Solution - uprev::uType # Master Solution - tmp::uType # Interpolation buffer - p::pType - t::tType # Current time - tprev::tType - dt::tType # This is the time step length which which we use during time marching - dtcache::tType # This is the proposed time step length - const dtchangeable::Bool # Indicator whether dtcache can be changed - tstops::heapType - _tstops::tstopsType # argument to __init used as default argument to reinit! - saveat::heapType - _saveat::saveatType # argument to __init used as default argument to reinit! - callback::callbackType - advance_to_tstop::Bool - # TODO group these into some internal flag struct - last_step_failed::Bool - force_stepfail::Bool - isout::Bool - u_modified::Bool - # DiffEqBase.initialize! and DiffEqBase.finalize! - cache::cacheType - sol::solType - subintegrator_tree::subintTreeType - iter::Int - controller::controllerType - EEst::Float64 # TODO integrate with controller cache - opts::optionsType - stats::IntegratorStats - tdir::tType -end - -# called by DiffEqBase.init and DiffEqBase.solve -function DiffEqBase.__init( - prob::OperatorSplittingProblem, - alg::AbstractOperatorSplittingAlgorithm, - args...; - dt, - tstops = (), - saveat = (), - d_discontinuities = (), - save_everystep = false, - callback = nothing, - advance_to_tstop = false, - adaptive = DiffEqBase.isadaptive(alg), - controller = nothing, - alias_u0 = false, - verbose = true, - kwargs... - ) - (; u0, p) = prob - t0, tf = prob.tspan - - dt > zero(dt) || error("dt must be positive") - dtcache = dt - dt = tf > t0 ? dt : -dt - tType = typeof(dt) - - # Warn if the algorithm is non-adaptive but the user tries to make it adaptive. - (!DiffEqBase.isadaptive(alg) && adaptive && verbose) && warn("The algorithm $alg is not adaptive.") - - dtchangeable = true # DiffEqBase.isadaptive(alg) - - if tstops isa AbstractArray || tstops isa Tuple || tstops isa Number - _tstops = nothing - else - _tstops = tstops - tstops = () - end - - # Setup tstop logic - tstops_internal = OrdinaryDiffEqCore.initialize_tstops( - tType, tstops, d_discontinuities, prob.tspan - ) - saveat_internal = OrdinaryDiffEqCore.initialize_saveat(tType, saveat, prob.tspan) - d_discontinuities_internal = OrdinaryDiffEqCore.initialize_d_discontinuities( - tType, d_discontinuities, prob.tspan - ) - - u = setup_u(prob, alg, alias_u0) - uprev = setup_u(prob, alg, false) - tmp = setup_u(prob, alg, false) - uType = typeof(u) - - sol = DiffEqBase.build_solution(prob, alg, tType[], uType[]) - - callback = DiffEqBase.CallbackSet(callback) - - subintegrator_tree, - cache = build_subintegrator_tree_with_cache( - prob, alg, - uprev, u, - 1:length(u), - t0, dt, tf, - tstops, saveat, d_discontinuities, callback, - adaptive, verbose - ) - - if controller === nothing && adaptive - controller = default_controller(alg, cache) - end - - integrator = OperatorSplittingIntegrator( - prob.f, - alg, - u, - uprev, - tmp, - p, - t0, - copy(t0), - dt, - dtcache, - dtchangeable, - tstops_internal, - tstops, - saveat_internal, - saveat, - callback, - advance_to_tstop, - false, - false, - false, - false, - cache, - sol, - subintegrator_tree, - build_solution_index_tree(prob.f), - build_synchronizer_tree(prob.f), - 0, - controller, - NaN, - IntegratorOptions(; verbose, adaptive, kwargs...), - IntegratorStats(), - tType(tstops_internal.ordering isa DataStructures.FasterForward ? 1 : -1) - ) - DiffEqBase.initialize!(callback, u0, t0, integrator) # Do I need this? - return integrator -end - -DiffEqBase.has_reinit(integrator::OperatorSplittingIntegrator) = true -function DiffEqBase.reinit!( - integrator::OperatorSplittingIntegrator, - u0 = integrator.sol.prob.u0; - t0 = integrator.sol.prob.tspan[1], - tf = integrator.sol.prob.tspan[2], - erase_sol = false, - tstops = integrator._tstops, - saveat = integrator._saveat, - reinit_callbacks = true, - reinit_retcode = true - ) - integrator.u .= u0 - integrator.uprev .= u0 - integrator.t = t0 - integrator.tprev = t0 - integrator.tstops, integrator.saveat = tstops_and_saveat_heaps(t0, tf, tstops, saveat) - integrator.iter = 0 - if erase_sol - resize!(integrator.sol.t, 0) - resize!(integrator.sol.u, 0) - end - if reinit_callbacks - DiffEqBase.initialize!(integrator.callback, u0, t0, integrator) - else # always reinit the saving callback so that t0 can be saved if needed - saving_callback = integrator.callback.discrete_callbacks[end] - DiffEqBase.initialize!(saving_callback, u0, t0, integrator) - end - if reinit_retcode - integrator.sol = DiffEqBase.solution_new_retcode( - integrator.sol, SciMLBase.ReturnCode.Default - ) - end - - return subreinit!( - integrator.f, - u0, - 1:length(u0), - integrator.subintegrator_tree; - t0, tf, - erase_sol, - tstops, - saveat, - reinit_callbacks, - reinit_retcode - ) -end - -function subreinit!( - f, - u0, - solution_indices, - subintegrator::DiffEqBase.DEIntegrator; - kwargs... - ) - return DiffEqBase.reinit!(subintegrator, u0[solution_indices]; kwargs...) -end - -@unroll function subreinit!( - f, - u0, - solution_indices, - subintegrators::Tuple; - kwargs... - ) - i = 1 - @unroll for subintegrator in subintegrators - subreinit!(get_operator(f, i), u0, f.solution_indices[i], subintegrator; kwargs...) - i += 1 - end -end - -function OrdinaryDiffEqCore.handle_tstop!(integrator::OperatorSplittingIntegrator) - if SciMLBase.has_tstop(integrator) - tdir_t = tdir(integrator) * integrator.t - tdir_tstop = SciMLBase.first_tstop(integrator) - if tdir_t == tdir_tstop - while tdir_t == tdir_tstop #remove all redundant copies - res = SciMLBase.pop_tstop!(integrator) - SciMLBase.has_tstop(integrator) ? - (tdir_tstop = SciMLBase.first_tstop(integrator)) : break - end - notify_integrator_hit_tstop!(integrator) - elseif tdir_t > tdir_tstop - if !integrator.dtchangeable - SciMLBase.change_t_via_interpolation!( - integrator, - tdir(integrator) * - SciMLBase.pop_tstop!(integrator), Val{true} - ) - notify_integrator_hit_tstop!(integrator) - else - error("Something went wrong. Integrator stepped past tstops but the algorithm was dtchangeable. Please report this error.") - end - end - end - return nothing -end - -notify_integrator_hit_tstop!(integrator::OperatorSplittingIntegrator) = nothing - -is_first_iteration(integrator::OperatorSplittingIntegrator) = integrator.iter == 0 -increment_iteration(integrator::OperatorSplittingIntegrator) = integrator.iter += 1 - -# Controller interface -function reject_step!(integrator::OperatorSplittingIntegrator) - OrdinaryDiffEqCore.increment_reject!(integrator.stats) - return reject_step!(integrator, integrator.controller) -end -function reject_step!(integrator::OperatorSplittingIntegrator, controller) - integrator.u .= integrator.uprev - if !integrator.force_stepfail - step_reject_controller!(integrator, controller, integrator.alg) - end - # We need to roll-back the sub-integrators - return prepare_subintegrators_to_redo_step!(integrator) -end -function reject_step!(integrator::OperatorSplittingIntegrator, ::Nothing) - return if length(integrator.uprev) == 0 - error("Cannot roll back integrator. Aborting time integration step at $(integrator.t).") - end -end - -# Solution looping interface -function should_accept_step(integrator::OperatorSplittingIntegrator) - if integrator.force_stepfail || integrator.isout - return false - end - return should_accept_step(integrator, integrator.controller) -end -function should_accept_step(integrator::OperatorSplittingIntegrator, ::Nothing) - return !(integrator.force_stepfail) -end -function accept_step!(integrator::OperatorSplittingIntegrator) - OrdinaryDiffEqCore.increment_accept!(integrator.stats) - return accept_step!(integrator, integrator.cache, integrator.controller) -end -function accept_step!(integrator::OperatorSplittingIntegrator, cache, controller) - return store_previous_info!(integrator) -end -function store_previous_info!(integrator::OperatorSplittingIntegrator) - return if length(integrator.uprev) > 0 # Integrator can rollback - update_uprev!(integrator) - end -end - -function update_uprev!(integrator::OperatorSplittingIntegrator) - SciMLBase.recursivecopy!(integrator.uprev, integrator.u) - return nothing -end - -function step_header!(integrator::OperatorSplittingIntegrator) - # Accept or reject the step - if !is_first_iteration(integrator) - if should_accept_step(integrator) - accept_step!(integrator) - else # Step should be rejected and hence repeated - reject_step!(integrator) - end - elseif integrator.u_modified # && integrator.iter == 0 - update_uprev!(integrator) - end - - # Before stepping we might need to adjust the dt - increment_iteration(integrator) - # OrdinaryDiffEqCore.choose_algorithm!(integrator, integrator.cache) - OrdinaryDiffEqCore.fix_dt_at_bounds!(integrator) - OrdinaryDiffEqCore.modify_dt_for_tstops!(integrator) - return integrator.force_stepfail = false -end - -function footer_reset_flags!(integrator) - return integrator.u_modified = false -end -function setup_validity_flags!(integrator, t_next) - return integrator.isout = false #integrator.opts.isoutofdomain(integrator.u, integrator.p, t_next) -end -function fix_solution_buffer_sizes!(integrator, sol) - resize!(integrator.sol.t, integrator.saveiter) - resize!(integrator.sol.u, integrator.saveiter) - return if !(integrator.sol isa SciMLBase.DAESolution) - resize!(integrator.sol.k, integrator.saveiter_dense) - end -end - -function step_footer!(integrator::OperatorSplittingIntegrator) - ttmp = integrator.t + tdir(integrator) * integrator.dt - - footer_reset_flags!(integrator) - setup_validity_flags!(integrator, ttmp) - - if should_accept_step(integrator) - integrator.last_step_failed = false - integrator.tprev = integrator.t - integrator.t = ttmp #OrdinaryDiffEqCore.fixed_t_for_floatingpoint_error!(integrator, ttmp) - # OrdinaryDiffEqCore.handle_callbacks!(integrator) - step_accept_controller!(integrator) # Noop for non-adaptive algorithms - elseif integrator.force_stepfail # Rejected by solver - if SciMLBase.isadaptive(integrator) - step_reject_controller!(integrator) - OrdinaryDiffEqCore.post_newton_controller!(integrator, integrator.alg) - elseif integrator.dtchangeable # Non-adaptive but can change dt - integrator.dt /= integrator.opts.failfactor - elseif integrator.last_step_failed - return - end - integrator.last_step_failed = true - end - - # integration_monitor_step(integrator) - - return nothing -end - -# called by DiffEqBase.solve -function DiffEqBase.__solve( - prob::OperatorSplittingProblem, - alg::AbstractOperatorSplittingAlgorithm, args...; kwargs... - ) - integrator = DiffEqBase.__init(prob, alg, args...; kwargs...) - return DiffEqBase.solve!(integrator) -end - -# either called directly (after init), or by DiffEqBase.solve (via __solve) -function DiffEqBase.solve!(integrator::OperatorSplittingIntegrator) - while !isempty(integrator.tstops) - while tdir(integrator) * integrator.t < SciMLBase.first_tstop(integrator) - step_header!(integrator) - @timeit_debug "check_error" DiffEqBase.check_error!(integrator) ∉ ( - SciMLBase.ReturnCode.Success, SciMLBase.ReturnCode.Default, - )&&return - __step!(integrator) - step_footer!(integrator) - if !SciMLBase.has_tstop(integrator) - break - end - end - OrdinaryDiffEqCore.handle_tstop!(integrator) - end - OrdinaryDiffEqCore.postamble!(integrator) - if integrator.sol.retcode != SciMLBase.ReturnCode.Default - return integrator.sol - end - return integrator.sol = SciMLBase.solution_new_retcode( - integrator.sol, SciMLBase.ReturnCode.Success - ) -end - -function DiffEqBase.step!(integrator::OperatorSplittingIntegrator) - @timeit_debug "step!" if integrator.advance_to_tstop - tstop = first_tstop(integrator) - while !reached_tstop(integrator, tstop) - step_header!(integrator) - @timeit_debug "check_error" DiffEqBase.check_error!(integrator) ∉ ( - SciMLBase.ReturnCode.Success, SciMLBase.ReturnCode.Default, - )&&return - __step!(integrator) - step_footer!(integrator) - if !SciMLBase.has_tstop(integrator) - break - end - end - else - step_header!(integrator) - @timeit_debug "check_error" DiffEqBase.check_error!(integrator) ∉ ( - SciMLBase.ReturnCode.Success, SciMLBase.ReturnCode.Default, - )&&return - __step!(integrator) - step_footer!(integrator) - while !should_accept_step(integrator) - step_header!(integrator) - @timeit_debug "check_error" DiffEqBase.check_error!(integrator) ∉ ( - SciMLBase.ReturnCode.Success, SciMLBase.ReturnCode.Default, - )&&return - __step!(integrator) - step_footer!(integrator) - end - end - return OrdinaryDiffEqCore.handle_tstop!(integrator) -end - -function SciMLBase.check_error(integrator::OperatorSplittingIntegrator) - if !SciMLBase.successful_retcode(integrator.sol) && - integrator.sol.retcode != SciMLBase.ReturnCode.Default - return integrator.sol.retcode - end - - verbose = true # integrator.opts.verbose - - if DiffEqBase.NAN_CHECK(integrator.dtcache) || DiffEqBase.NAN_CHECK(integrator.dt) # replace with https://github.com/SciML/OrdinaryDiffEq.jl/blob/373a8eec8024ef1acc6c5f0c87f479aa0cf128c3/lib/OrdinaryDiffEqCore/src/iterator_interface.jl#L5-L6 after moving to sciml integrators - if verbose - @warn("NaN dt detected. Likely a NaN value in the state, parameters, or derivative value caused this outcome.") - end - return SciMLBase.ReturnCode.DtNaN - end - - return check_error_subintegrators(integrator, integrator.subintegrator_tree) -end - -function check_error_subintegrators(integrator, subintegrator_tree::Tuple) - for subintegrator in subintegrator_tree - retcode = check_error_subintegrators(integrator, subintegrator) - if !SciMLBase.successful_retcode(retcode) && retcode != SciMLBase.ReturnCode.Default - return retcode - end - end - return integrator.sol.retcode -end - -function check_error_subintegrators(integrator, subintegrator::SciMLBase.DEIntegrator) - return SciMLBase.check_error(subintegrator) -end - -function DiffEqBase.step!(integrator::OperatorSplittingIntegrator, dt, stop_at_tdt = false) - return @timeit_debug "step!" begin - # OridinaryDiffEq lets dt be negative if tdir is -1, but that's inconsistent - dt <= zero(dt) && error("dt must be positive") - stop_at_tdt && !integrator.dtchangeable && - error("Cannot stop at t + dt if dtchangeable is false") - tnext = integrator.t + tdir(integrator) * dt - stop_at_tdt && DiffEqBase.add_tstop!(integrator, tnext) - while !reached_tstop(integrator, tnext, stop_at_tdt) - step_header!(integrator) - @timeit_debug "check_error" DiffEqBase.check_error!(integrator) ∉ ( - SciMLBase.ReturnCode.Success, SciMLBase.ReturnCode.Default, - )&&return - __step!(integrator) - step_footer!(integrator) - end - end -end - -function setup_u(prob::OperatorSplittingProblem, solver, alias_u0) - if alias_u0 - return prob.u0 - else - return OrdinaryDiffEqCore.recursivecopy(prob.u0) - end -end - -# TimeChoiceIterator API -@inline function DiffEqBase.get_tmp_cache(integrator::OperatorSplittingIntegrator) - # DiffEqBase.get_tmp_cache(integrator, integrator.alg, integrator.cache) - return (integrator.tmp,) -end -# @inline function DiffEqBase.get_tmp_cache(integrator::OperatorSplittingIntegrator, ::AbstractOperatorSplittingAlgorithm, cache) -# return (cache.tmp,) -# end -# Interpolation -# TODO via https://github.com/SciML/SciMLBase.jl/blob/master/src/interpolation.jl -function linear_interpolation!(y, t, y1, y2, t1, t2) - return y .= y1 + (t - t1) * (y2 - y1) / (t2 - t1) -end -function (integrator::OperatorSplittingIntegrator)(tmp, t) - return linear_interpolation!( - tmp, t, integrator.uprev, integrator.u, integrator.tprev, integrator.t - ) -end - -""" - stepsize_controller!(::OperatorSplittingIntegrator) - -Updates the controller using the current state of the integrator if the operator splitting algorithm is adaptive. -""" -@inline function stepsize_controller!(integrator::OperatorSplittingIntegrator) - DiffEqBase.isadaptive(integrator) || return nothing - return stepsize_controller!(integrator, integrator.controller, integrator.alg) -end - -""" - step_accept_controller!(::OperatorSplittingIntegrator) - -Updates `dtcache` of the integrator if the step is accepted and the operator splitting algorithm is adaptive. -""" -@inline function step_accept_controller!(integrator::OperatorSplittingIntegrator) - DiffEqBase.isadaptive(integrator) || return nothing - return step_accept_controller!(integrator, integrator.controller, integrator.alg) -end - -""" - step_reject_controller!(::OperatorSplittingIntegrator) - -Updates `dtcache` of the integrator if the step is rejected and the the operator splitting algorithm is adaptive. -""" -@inline function step_reject_controller!(integrator::OperatorSplittingIntegrator) - DiffEqBase.isadaptive(integrator) || return nothing - return step_reject_controller!(integrator, integrator.controller, integrator.alg) -end - -# helper functions for dealing with time-reversed integrators in the same way -# that OrdinaryDiffEq.jl does -tdir(integrator) = integrator.tstops.ordering isa DataStructures.FasterForward ? 1 : -1 -is_past_t(integrator, t) = tdir(integrator) * (t - integrator.t) ≤ zero(integrator.t) -function reached_tstop(integrator, tstop, stop_at_tstop = integrator.dtchangeable) - if stop_at_tstop - integrator.t > tstop && - error("Integrator missed stop at $tstop (current time=$(integrator.t)). Aborting.") - return integrator.t == tstop # Check for exact hit - else #!stop_at_tstop - return is_past_t(integrator, tstop) - end -end - -# Dunno stuff -function SciMLBase.done(integrator::OperatorSplittingIntegrator) - if !( - integrator.sol.retcode in ( - SciMLBase.ReturnCode.Default, SciMLBase.ReturnCode.Success, - ) - ) - return true - elseif isempty(integrator.tstops) - DiffEqBase.postamble!(integrator) - return true - end - return false -end - -function DiffEqBase.postamble!(integrator::OperatorSplittingIntegrator) - return DiffEqBase.finalize!(integrator.callback, integrator.u, integrator.t, integrator) -end - -function __step!(integrator) - tnext = integrator.t + integrator.dt - synchronize_subintegrator_tree!(integrator) - advance_solution_to!(integrator, tnext) - return stepsize_controller!(integrator) -end - -# solvers need to define this interface -function advance_solution_to!(integrator::OperatorSplittingIntegrator, tnext) - return advance_solution_to!(integrator, integrator.cache, tnext) -end - -function advance_solution_to!( - outer_integrator::OperatorSplittingIntegrator, - integrator::DiffEqBase.DEIntegrator, solution_indices, sync, cache, tend - ) - dt = tend - integrator.t - return SciMLBase.step!(integrator, dt, true) -end - -# ----------------------------------- SciMLBase.jl Integrator Interface ------------------------------------ -SciMLBase.has_stats(::OperatorSplittingIntegrator) = true - -SciMLBase.has_tstop(integrator::OperatorSplittingIntegrator) = !isempty(integrator.tstops) -SciMLBase.first_tstop(integrator::OperatorSplittingIntegrator) = first(integrator.tstops) -SciMLBase.pop_tstop!(integrator::OperatorSplittingIntegrator) = pop!(integrator.tstops) - -DiffEqBase.get_dt(integrator::OperatorSplittingIntegrator) = integrator.dt -function set_dt!(integrator::OperatorSplittingIntegrator, dt) - # TODO: figure out interface for recomputing other objects (linear operators, etc) - dt <= zero(dt) && error("dt must be positive") - return integrator.dt = dt -end - -function DiffEqBase.add_tstop!(integrator::OperatorSplittingIntegrator, t) - is_past_t(integrator, t) && - error("Cannot add a tstop at $t because that is behind the current \ - integrator time $(integrator.t)") - return push!(integrator.tstops, t) -end - -function DiffEqBase.add_saveat!(integrator::OperatorSplittingIntegrator, t) - is_past_t(integrator, t) && - error("Cannot add a saveat point at $t because that is behind the \ - current integrator time $(integrator.t)") - return push!(integrator.saveat, t) -end - -# not sure what this should do? -# defined as default initialize: https://github.com/SciML/DiffEqBase.jl/blob/master/src/callbacks.jl#L3 -DiffEqBase.u_modified!(i::OperatorSplittingIntegrator, bool) = nothing - -function synchronize_subintegrator_tree!(integrator::OperatorSplittingIntegrator) - return synchronize_subintegrator!(integrator.subintegrator_tree, integrator) -end - -@unroll function synchronize_subintegrator!( - subintegrator_tree::Tuple, integrator::OperatorSplittingIntegrator - ) - @unroll for subintegrator in subintegrator_tree - synchronize_subintegrator!(subintegrator, integrator) - end -end - -function synchronize_subintegrator!( - subintegrator::SciMLBase.DEIntegrator, integrator::OperatorSplittingIntegrator - ) - @unpack t, dt = integrator - @assert subintegrator.t == t "Integrators out of sync. The outer integrator is at $t, but inner integrator is at $(subintegrator.t)" - return if !DiffEqBase.isadaptive(subintegrator) - SciMLBase.set_proposed_dt!(subintegrator, dt) - end -end - -function advance_solution_to!( - integrator::OperatorSplittingIntegrator, - cache::AbstractOperatorSplittingCache, tnext::Number - ) - return advance_solution_to!( - integrator, integrator.subintegrator_tree, integrator.solution_index_tree, - integrator.synchronizer_tree, cache, tnext - ) -end - -# Dispatch for tree node construction -function build_subintegrator_tree_with_cache( - prob::OperatorSplittingProblem, alg::AbstractOperatorSplittingAlgorithm, - uprevouter::AbstractVector, uouter::AbstractVector, - solution_indices, - t0, dt, tf, - tstops, saveat, d_discontinuities, callback, - adaptive, verbose - ) - (; f, p) = prob - subintegrator_tree_with_caches = ntuple( - i -> build_subintegrator_tree_with_cache( - prob, - alg.inner_algs[i], - get_operator(f, i), - p[i], - uprevouter, uouter, - f.solution_indices[i], - t0, dt, tf, - tstops, saveat, d_discontinuities, callback, - adaptive, verbose - ), - length(f.functions) - ) - - subintegrator_tree = ntuple( - i -> subintegrator_tree_with_caches[i][1], length(f.functions) - ) - caches = ntuple(i -> subintegrator_tree_with_caches[i][2], length(f.functions)) - - # TODO fix mixed device type problems we have to be smarter - return subintegrator_tree, - init_cache( - f, alg; - uprev = uprevouter, u = uouter, alias_u = true, - inner_caches = caches - ) -end - -function build_subintegrator_tree_with_cache( - prob::OperatorSplittingProblem, alg::AbstractOperatorSplittingAlgorithm, - f::GenericSplitFunction, p::Tuple, - uprevouter::AbstractVector, uouter::AbstractVector, - solution_indices, - t0, dt, tf, - tstops, saveat, d_discontinuities, callback, - adaptive, verbose, - save_end = false, - controller = nothing - ) - subintegrator_tree_with_caches = ntuple( - i -> build_subintegrator_tree_with_cache( - prob, - alg.inner_algs[i], - get_operator(f, i), - p[i], - uprevouter, uouter, - f.solution_indices[i], - t0, dt, tf, - tstops, saveat, d_discontinuities, callback, - adaptive, verbose - ), - length(f.functions) - ) - - subintegrator_tree = first.(subintegrator_tree_with_caches) - inner_caches = last.(subintegrator_tree_with_caches) - - # TODO fix mixed device type problems we have to be smarter - uprev = @view uprevouter[solution_indices] - u = @view uouter[solution_indices] - return subintegrator_tree, - init_cache( - f, alg; - uprev = uprev, u = u, - inner_caches = inner_caches - ) -end - -function build_subintegrator_tree_with_cache( - prob::OperatorSplittingProblem, alg::SciMLBase.AbstractODEAlgorithm, - f::F, p::P, - uprevouter::S, uouter::S, - solution_indices, - t0::T, dt::T, tf::T, - tstops, saveat, d_discontinuities, callback, - adaptive, verbose, - save_end = false, - controller = nothing - ) where {S, T, P, F} - uprev = @view uprevouter[solution_indices] - u = @view uouter[solution_indices] - - integrator = DiffEqBase.__init( - SciMLBase.ODEProblem(f, u, (t0, min(t0 + dt, tf)), p), - alg; - dt, - saveat = (), - d_discontinuities, - save_everystep = false, - advance_to_tstop = false, - adaptive, - controller, - verbose - ) - - return integrator, integrator.cache -end - -function forward_sync_subintegrator!( - outer_integrator::OperatorSplittingIntegrator, subintegrator_tree::Tuple, - solution_indices::Tuple, synchronizers::Tuple - ) - return nothing -end -function backward_sync_subintegrator!( - outer_integrator::OperatorSplittingIntegrator, - subintegrator_tree::Tuple, solution_indices::Tuple, synchronizer::Tuple - ) - return nothing -end