diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml index 9c793591..320e0c07 100644 --- a/.JuliaFormatter.toml +++ b/.JuliaFormatter.toml @@ -1,2 +1,3 @@ style = "sciml" -format_markdown = true \ No newline at end of file +format_markdown = true +annotate_untyped_fields_with_any = false diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml index cf17c8d8..48d3b873 100644 --- a/.github/workflows/Downgrade.yml +++ b/.github/workflows/Downgrade.yml @@ -22,7 +22,6 @@ jobs: with: version: ${{ matrix.version }} - uses: julia-actions/julia-downgrade-compat@v1 -# if: ${{ matrix.version == '1.6' }} with: skip: Pkg,TOML - uses: julia-actions/julia-buildpkg@v1 diff --git a/Project.toml b/Project.toml index 770e6937..407671e7 100644 --- a/Project.toml +++ b/Project.toml @@ -1,57 +1,54 @@ name = "DiffEqCallbacks" uuid = "459566f4-90b8-5000-8ac3-15dfb0a30def" authors = ["Chris Rackauckas "] -version = "3.9.1" +version = "4.0.0" [deps] +ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" -NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" -Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a" RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" -[weakdeps] -OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" -OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8" -Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4" - [compat] +ADTypes = "1.9.0" Aqua = "0.8" -DataInterpolations = "4.6" +ConcreteStructs = "0.2.3" +DataInterpolations = "4.6, 5, 6" DataStructures = "0.18.13" -DiffEqBase = "6.154" +DiffEqBase = "6.155.3" +DifferentiationInterface = "0.6.1" ForwardDiff = "0.10.36" Functors = "0.4" LinearAlgebra = "1.10" Markdown = "1.10" -NonlinearSolve = "3.7.2" -ODEProblemLibrary = "0.1.5" +NonlinearSolve = "3.14" +ODEProblemLibrary = "0.1.8" OrdinaryDiffEq = "6.88" -OrdinaryDiffEqCore = "1" -Parameters = "0.12" QuadGK = "2.9" RecipesBase = "1.3.4" -RecursiveArrayTools = "3.9" -SciMLBase = "2.26" -SciMLSensitivity = "7.56" -StaticArrays = "1.8" +RecursiveArrayTools = "3.27" +SciMLBase = "2.54" +SciMLSensitivity = "7.68" +StaticArrays = "1.9.7" StaticArraysCore = "1.4" Sundials = "4.19.2" -Test = "1" -Tracker = "0.2.26" +Test = "1.10" +Tracker = "0.2.35" Zygote = "0.6.69" julia = "1.10" [extras] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" @@ -64,4 +61,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["Aqua", "DataInterpolations", "OrdinaryDiffEq", "ODEProblemLibrary", "Test", "QuadGK", "SciMLSensitivity", "StaticArrays", "Tracker", "Zygote", "NonlinearSolve"] +test = ["ADTypes", "Aqua", "DataInterpolations", "ForwardDiff", "NonlinearSolve", "ODEProblemLibrary", "OrdinaryDiffEq", "QuadGK", "SciMLSensitivity", "StaticArrays", "Test", "Tracker", "Zygote"] diff --git a/docs/Project.toml b/docs/Project.toml index 759ff0db..8688ff38 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,4 +1,5 @@ [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -6,7 +7,8 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" [compat] -DiffEqCallbacks = "3" +ADTypes = "1.9.0" +DiffEqCallbacks = "4" Documenter = "1" OrdinaryDiffEq = "6.88" Plots = "1.36" diff --git a/docs/src/projection.md b/docs/src/projection.md index 31b08bc2..53289c50 100644 --- a/docs/src/projection.md +++ b/docs/src/projection.md @@ -12,7 +12,7 @@ ManifoldProjection Here we solve the harmonic oscillator: ```@example manifold -using OrdinaryDiffEq, DiffEqCallbacks, Plots +using OrdinaryDiffEq, DiffEqCallbacks, Plots, ADTypes u0 = ones(2) function f(du, u, p, t) @@ -28,14 +28,13 @@ to conserve the sum of squares: ```@example manifold function g(resid, u, p, t) resid[1] = u[2]^2 + u[1]^2 - 2 - resid[2] = 0 end ``` To build the callback, we just call ```@example manifold -cb = ManifoldProjection(g) +cb = ManifoldProjection(g; autodiff = AutoForwardDiff(), resid_prototype = zeros(1)) ``` Using this callback, the Runge-Kutta method `Vern7` conserves energy. Note that the diff --git a/src/DiffEqCallbacks.jl b/src/DiffEqCallbacks.jl index 84967de1..e483795f 100644 --- a/src/DiffEqCallbacks.jl +++ b/src/DiffEqCallbacks.jl @@ -1,17 +1,22 @@ module DiffEqCallbacks -using DiffEqBase, RecursiveArrayTools, DataStructures, RecipesBase, LinearAlgebra, - StaticArraysCore, NonlinearSolve, ForwardDiff, Functors - -import Base.Iterators - -using Markdown - -using Parameters: @unpack - -import SciMLBase - -using DiffEqBase: get_tstops, get_tstops_array, get_tstops_max +using ConcreteStructs: @concrete +using DataStructures: DataStructures, BinaryMaxHeap, BinaryMinHeap +using DiffEqBase: DiffEqBase, get_tstops, get_tstops_array, get_tstops_max +using DifferentiationInterface: DifferentiationInterface, Constant +using Functors: fmap +using LinearAlgebra: LinearAlgebra, adjoint, axpy!, copyto!, mul!, ldiv! +using Markdown: @doc_str +using RecipesBase: @recipe +using RecursiveArrayTools: RecursiveArrayTools, DiffEqArray, copyat_or_push! +using SciMLBase: SciMLBase, CallbackSet, DiscreteCallback, NonlinearFunction, + NonlinearLeastSquaresProblem, NonlinearProblem, RODEProblem, + ReturnCode, SDEProblem, add_tstop!, check_error, get_du, + get_proposed_dt, get_tmp_cache, init, reinit!, + set_proposed_dt!, solve!, terminate!, u_modified! +using StaticArraysCore: StaticArraysCore + +const DI = DifferentiationInterface include("functor_helpers.jl") include("autoabstol.jl") diff --git a/src/domain.jl b/src/domain.jl index d9879886..581a18f8 100644 --- a/src/domain.jl +++ b/src/domain.jl @@ -2,33 +2,35 @@ abstract type AbstractDomainAffect{T, S, uType} end +(f::AbstractDomainAffect)(integrator) = affect!(integrator, f) + struct PositiveDomainAffect{T, S, uType} <: AbstractDomainAffect{T, S, uType} abstol::T scalefactor::S u::uType end -struct GeneralDomainAffect{autonomous, F, T, S, uType} <: AbstractDomainAffect{T, S, uType} +struct GeneralDomainAffect{F <: AbstractNonAutonomousFunction, T, S, uType, A} <: + AbstractDomainAffect{T, S, uType} g::F abstol::T scalefactor::S u::uType resid::uType + autonomous::A +end - function GeneralDomainAffect{autonomous}(g::F, abstol::T, scalefactor::S, u::uType, - resid::uType) where {autonomous, F, T, S, uType - } - new{autonomous, F, T, S, uType}(g, abstol, scalefactor, u, resid) +function initialize_general_domain_affect(cb, u, t, integrator) + return initialize_general_domain_affect(cb.affect!, u, t, integrator) +end +function initialize_general_domain_affect(affect!::GeneralDomainAffect, u, t, integrator) + if affect!.autonomous === nothing + autonomous = maximum(SciMLBase.numargs(affect!.g.f)) == + 2 + SciMLBase.isinplace(integrator.f) + affect!.g.autonomous = autonomous end end -# definitions of callback functions - -# Workaround since it is not possible to add methods to an abstract type: -# https://github.com/JuliaLang/julia/issues/14919 -(f::PositiveDomainAffect)(integrator) = affect!(integrator, f) -(f::GeneralDomainAffect)(integrator) = affect!(integrator, f) - # general method definitions for domain callbacks """ @@ -41,6 +43,8 @@ function affect!(integrator, f::AbstractDomainAffect{T, S, uType}) where {T, S, throw(ArgumentError("domain callback can only be applied to adaptive algorithms")) end + iip = Val(SciMLBase.isinplace(integrator.f)) + # define array of next time step, absolute tolerance, and scale factor if uType <: Nothing if integrator.u isa Union{Number, StaticArraysCore.SArray} @@ -55,7 +59,7 @@ function affect!(integrator, f::AbstractDomainAffect{T, S, uType}) where {T, S, scalefactor = S <: Nothing ? 1 // 2 : f.scalefactor # setup callback and save additional arguments for checking next time step - args = setup(f, integrator) + args = setup(f, integrator, iip) # obtain proposed next time step dt = get_proposed_dt(integrator) @@ -80,7 +84,7 @@ function affect!(integrator, f::AbstractDomainAffect{T, S, uType}) where {T, S, end # check whether time step is accepted - isaccepted(u, p, t, abstol, f, args...) && break + isaccepted(u, p, t, abstol, f, iip, args...) && break # reduce time step dtcache = dt @@ -120,12 +124,12 @@ was modified. modify_u!(integrator, ::AbstractDomainAffect) = false """ - setup(f::AbstractDomainAffect, integrator) + setup(f::AbstractDomainAffect, integrator, ::Val{iip}) where {iip} Setup callback `f` and return an arbitrary tuple whose elements are used as additional arguments in checking whether time step is accepted. """ -setup(::AbstractDomainAffect, integrator) = () +setup(::AbstractDomainAffect, integrator, ::Val{iip}) where {iip} = () """ isaccepted(u, abstol, f::AbstractDomainAffect, args...) @@ -133,7 +137,7 @@ setup(::AbstractDomainAffect, integrator) = () Return whether `u` is an acceptable state vector at the next time point given absolute tolerance `abstol`, callback `f`, and other optional arguments. """ -isaccepted(u, p, t, tolerance, ::AbstractDomainAffect, args...) = true +isaccepted(u, p, t, tolerance, ::AbstractDomainAffect, ::Val{iip}, args...) where {iip} = true # specific method definitions for positive domain callback @@ -175,27 +179,30 @@ function _set_neg_zero!(integrator, u::StaticArraysCore.SArray) end # state vector is accepted if its entries are greater than -abstol -isaccepted(u, p, t, abstol::Number, ::PositiveDomainAffect) = all(ui -> ui > -abstol, u) -function isaccepted(u, p, t, abstol, ::PositiveDomainAffect) +function isaccepted(u, p, t, abstol::Number, ::PositiveDomainAffect, ::Val{iip}) where {iip} + return all(ui -> ui > -abstol, u) +end +function isaccepted(u, p, t, abstol, ::PositiveDomainAffect, ::Val{iip}) where {iip} length(u) == length(abstol) || throw(DimensionMismatch("numbers of states and tolerances do not match")) - all(ui > -tol for (ui, tol) in zip(u, abstol)) + return all(ui > -tol for (ui, tol) in zip(u, abstol)) end # specific method definitions for general domain callback # create array of residuals -function setup(f::GeneralDomainAffect, integrator) - f.resid isa Nothing ? (similar(integrator.u),) : (f.resid,) +setup(f::GeneralDomainAffect, integrator, ::Val{false}) = (nothing,) +function setup(f::GeneralDomainAffect, integrator, ::Val{true}) + return f.resid === nothing ? (similar(integrator.u),) : (f.resid,) end -function isaccepted(u, p, t, abstol, f::GeneralDomainAffect{autonomous, F, T, S, uType}, - resid) where {autonomous, F, T, S, uType} +function isaccepted(u, p, t, abstol, f::GeneralDomainAffect, ::Val{iip}, resid) where {iip} # calculate residuals - if autonomous + f.g.t = t + if iip f.g(resid, u, p) else - f.g(resid, u, p, t) + resid = f.g(u, p) end # accept time step if residuals are smaller than the tolerance @@ -214,26 +221,32 @@ end """ GeneralDomain( g, u = nothing; save = true, abstol = nothing, scalefactor = nothing, - autonomous = maximum(SciMLBase.numargs(g)) == 3, nlsolve_kwargs = (; - abstol = 10 * eps()), kwargs...) + autonomous = nothing, domain_jacobian = nothing, + nlsolve_kwargs = (; abstol = 10 * eps()), kwargs...) A `GeneralDomain` callback in DiffEqCallbacks.jl generalizes the concept of -a `PositiveDomain` callback to arbitrary domains. Domains are specified by -in-place functions `g(resid, u, p)` or `g(resid, u, p, t)` that calculate residuals of a -state vector `u` at time `t` relative to that domain, with `p` the parameters of the -corresponding integrator. As for `PositiveDomain`, steps are accepted if residuals -of the extrapolated values at the next time step are below -a certain tolerance. Moreover, this callback is automatically coupled with a -`ManifoldProjection` that keeps all calculated state vectors close to the desired -domain, but in contrast to a `PositiveDomain` callback the nonlinear solver in a -`ManifoldProjection` cannot guarantee that all state vectors of the solution are -actually inside the domain. Thus, a `PositiveDomain` callback should generally be -preferred. +a `PositiveDomain` callback to arbitrary domains. + +Domains are specified by + - in-place functions `g(resid, u, p)` or `g(resid, u, p, t)` if the corresponding + ODEProblem is an inplace problem, or + - out-of-place functions `g(u, p)` or `g(u, p, t)` if the corresponding ODEProblem is + an out-of-place problem. + +The function calculates residuals of a state vector `u` at time `t` relative to that domain, +with `p` the parameters of the corresponding integrator. + +As for `PositiveDomain`, steps are accepted if residuals of the extrapolated values at the +next time step are below a certain tolerance. Moreover, this callback is automatically +coupled with a `ManifoldProjection` that keeps all calculated state vectors close to the +desired domain, but in contrast to a `PositiveDomain` callback the nonlinear solver in a +`ManifoldProjection` cannot guarantee that all state vectors of the solution are actually +inside the domain. Thus, a `PositiveDomain` callback should generally be preferred. ## Arguments - - `g`: the implicit definition of the domain as a function `g(resid, u, p)` or - `g(resid, u, p, t)` which is zero when the value is in the domain. + - `g`: the implicit definition of the domain as a function as described above which is + zero when the value is in the domain. - `u`: A prototype of the state vector of the integrator. A copy of it is saved and extrapolated values are written to it. If it is not specified, every application of the callback allocates a new copy of the state vector. @@ -248,9 +261,13 @@ preferred. specified, time steps are halved. - `autonomous`: Whether `g` is an autonomous function of the form `g(resid, u, p)`. If it is not specified, it is determined automatically. - - `kwargs`: All other keyword arguments are passed to `ManifoldProjection`. + - `kwargs`: All other keyword arguments are passed to [`ManifoldProjection`](@ref). - `nlsolve_kwargs`: All keyword arguments are passed to the nonlinear solver in `ManifoldProjection`. The default is `(; abstol = 10 * eps())`. + - `domain_jacobian`: The Jacobian of the domain (wrt the state). This has the same + signature as `g` and the first argument is the Jacobian if inplace. This corresponds to + the `manifold_jacobian` argument of [`ManifoldProjection`](@ref). Note that passing + a `manifold_jacobian` is not supported for `GeneralDomain` and results in an error. ## References @@ -260,20 +277,27 @@ Non-negative solutions of ODEs. Applied Mathematics and Computation 170 """ function GeneralDomain( g, u = nothing; save = true, abstol = nothing, scalefactor = nothing, - autonomous = maximum(SciMLBase.numargs(g)) == 3, nlsolve_kwargs = (; - abstol = 10 * eps()), kwargs...) - _autonomous = SciMLBase._unwrap_val(autonomous) - if u isa Nothing - affect! = GeneralDomainAffect{_autonomous}(g, abstol, scalefactor, nothing, nothing) + autonomous = nothing, domain_jacobian = nothing, manifold_jacobian = missing, + nlsolve_kwargs = (; abstol = 10 * eps()), kwargs...) + if manifold_jacobian !== missing + throw(ArgumentError("`manifold_jacobian` is not supported for `GeneralDomain`. \ + Use `domain_jacobian` instead.")) + end + manifold_projection = ManifoldProjection( + g; save = false, autonomous, manifold_jacobian = domain_jacobian, + kwargs..., nlsolve_kwargs...) + domain = wrap_autonomous_function(autonomous, g) + domain_jacobian = wrap_autonomous_function(autonomous, domain_jacobian) + affect! = if u === nothing + GeneralDomainAffect(domain, abstol, scalefactor, nothing, nothing, autonomous) else - affect! = GeneralDomainAffect{_autonomous}(g, abstol, scalefactor, deepcopy(u), - deepcopy(u)) + GeneralDomainAffect( + domain, abstol, scalefactor, deepcopy(u), deepcopy(u), autonomous) end - condition = (u, t, integrator) -> true - CallbackSet( - ManifoldProjection( - g; save = false, autonomous, isinplace = Val(true), kwargs..., nlsolve_kwargs...), - DiscreteCallback(condition, affect!; save_positions = (false, save))) + domain_cb = DiscreteCallback( + Returns(true), affect!; initialize = initialize_general_domain_affect, + save_positions = (false, save)) + return CallbackSet(manifold_projection, domain_cb) end @doc doc""" diff --git a/src/iterative_and_periodic.jl b/src/iterative_and_periodic.jl index 9e66de88..05166174 100644 --- a/src/iterative_and_periodic.jl +++ b/src/iterative_and_periodic.jl @@ -89,7 +89,7 @@ function (S::PeriodicCallbackAffect)(integrator) end function add_next_tstop!(integrator, S) - @unpack Δt, t0, index = S + (; Δt, t0, index) = S # Schedule next call to `f` using `add_tstops!`, but be careful not to keep integrating forever tnew = t0[] + (index[] + 1) * Δt @@ -177,7 +177,8 @@ end # Checking for floating point equality is OK here as `DifferentialEquations.jl` # sets the time exactly to the final time in the last iteration return integrator.t == last(integrator.sol.prob.tspan) || - (hasfield(typeof(integrator), :iter) && (integrator.iter == integrator.opts.maxiters)) + (hasfield(typeof(integrator), :iter) && + (integrator.iter == integrator.opts.maxiters)) end export PeriodicCallback diff --git a/src/manifold.jl b/src/manifold.jl index 891114b7..6b593da7 100644 --- a/src/manifold.jl +++ b/src/manifold.jl @@ -1,20 +1,7 @@ -# wrapper for non-autonomous functions -mutable struct NonAutonomousFunction{iip, F, autonomous} - f::F - t::Any -end - -(f::NonAutonomousFunction{true, F, true})(res, u, p) where {F} = f.f(res, u, p) -(f::NonAutonomousFunction{true, F, false})(res, u, p) where {F} = f.f(res, u, p, f.t) - -(f::NonAutonomousFunction{false, F, true})(u, p) where {F} = f.f(u, p) -(f::NonAutonomousFunction{false, F, false})(u, p) where {F} = f.f(u, p, f.t) - -SciMLBase.isinplace(::NonAutonomousFunction{iip}) where {iip} = iip - """ - ManifoldProjection(g; nlsolve = nothing, save = true, nlls = Val(true), - isinplace = Val(true), autonomous = nothing, resid_prototype = nothing, kwargs...) + ManifoldProjection( + manifold; nlsolve = missing, save = true, autonomous = nothing, + manifold_jacobian = nothing, autodiff = nothing, kwargs...) In many cases, you may want to declare a manifold on which a solution lives. Mathematically, a manifold `M` is defined by a function `g` as the set of points where `g(u) = 0`. An @@ -33,30 +20,34 @@ properties. ## Arguments - - `g`: The residual function for the manifold. - - * This is an inplace function of form `g(resid, u, p)` or `g(resid, u, p, t)` which - writes to the residual `resid` the difference from the manifold components. Here, it - is assumed that `resid` is of the same shape as `u`. - * If `isinplace = Val(false)`, then `g` should be a function of the form `g(u, p)` or - `g(u, p, t)` which returns the residual. + - `manifold`: The residual function for the manifold. If the ODEProblem is an inplace + problem, then `g` must be an inplace function of form `g(resid, u, p)` or + `g(resid, u, p, t)` and similarly if the ODEProblem is an out-of-place problem then `g` + must be a function of form `g(u, p)` or `g(u, p, t)`. ## Keyword Arguments - - `nlsolve`: A nonlinear solver as defined in the - [NonlinearSolve.jl format](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/). - Defaults to a PolyAlgorithm. + - `nlsolve`: Defaults to a special single-factorize algorithm (denoted by `missing`) that + would work in most cases (See [1] for details). Alternatively, a nonlinear solver as + defined in the + [NonlinearSolve.jl format](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/) + can be specified. Additionally if NonlinearSolve.jl is loaded and `nothing` is specified + a polyalgorithm is used. - `save`: Whether to do the standard saving (applied after the callback) - - `nlls`: If the problem is a nonlinear least squares problem. `nlls = Val(false)` - generates a `NonlinearProblem` which is typically faster than - `NonlinearLeastSquaresProblem`, but is only applicable if the residual size is same as - the state size. - `autonomous`: Whether `g` is an autonomous function of the form `g(resid, u, p)` or - `g(u, p)`. Specify it as `Val(::Bool)` to ensure this function call is type stable. - - `residual_prototype`: This needs to be specified if `nlls = Val(true)` and - `inplace = Val(true)` are specified together, else it is taken to be same as `u`. + `g(u, p)`. Specify it as `Val(::Bool)` to disable runtime branching. If `nothing`, + we attempt to infer it from the signature of `g`. + - `residual_prototype`: The size of the manifold residual. If it is not specified, + we assume it to be same as `u` in the inplace case. Else we run a single evaluation of + `manifold` to determine the size. - `kwargs`: All other keyword arguments are passed to - [NonlinearSolve.jl](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/). + [NonlinearSolve.jl](https://docs.sciml.ai/NonlinearSolve/stable/basics/solve/) if + `nlsolve` is not `missing`. + - `autodiff`: The autodifferentiation algorithm to use to compute the Jacobian if + `manifold_jacobian` is not specified. This must be specified if `manifold_jacobian` is + not specified. + - `manifold_jacobian`: The Jacobian of the manifold (wrt the state). This has the same + signature as `manifold` and the first argument is the Jacobian if inplace. ### Saveat Warning @@ -72,78 +63,374 @@ order of the integrator even with the ManifoldProjection. ## References -Ernst Hairer, Christian Lubich, Gerhard Wanner. Geometric Numerical Integration: +[1] Ernst Hairer, Christian Lubich, Gerhard Wanner. Geometric Numerical Integration: Structure-Preserving Algorithms for Ordinary Differential Equations. Berlin ; New York :Springer, 2002. """ -mutable struct ManifoldProjection{iip, nlls, autonomous, F, NL, R, K} - g::F +@concrete mutable struct ManifoldProjection + manifold + manifold_jacobian + autodiff nlcache::Any - nlsolve::NL - resid_prototype::R - kwargs::K - - function ManifoldProjection{iip, nlls, autonomous}( - g, nlsolve, resid_prototype, kwargs) where {iip, nlls, autonomous} - # replace residual function if it is time-dependent - _g = NonAutonomousFunction{iip, typeof(g), autonomous}(g, 0) - return new{iip, nlls, autonomous, typeof(_g), typeof(nlsolve), - typeof(resid_prototype), typeof(kwargs)}( - _g, nothing, nlsolve, resid_prototype, kwargs) - end + nlsolve + kwargs + autonomous +end + +function ManifoldProjection( + manifold; nlsolve = missing, save = true, autonomous = nothing, + manifold_jacobian = nothing, autodiff = nothing, kwargs...) + affect! = ManifoldProjection( + manifold, autodiff, manifold_jacobian, nlsolve, kwargs, autonomous) + return DiscreteCallback( + Returns(true), affect!; initialize = initialize_manifold_projection, + save_positions = (false, save)) +end + +function ManifoldProjection( + manifold, autodiff, manifold_jacobian, nlsolve, kwargs, autonomous) + wrapped_manifold = wrap_autonomous_function(autonomous, manifold) + wrapped_manifold_jacobian = wrap_autonomous_function(autonomous, manifold_jacobian) + return ManifoldProjection(wrapped_manifold, wrapped_manifold_jacobian, + autodiff, nothing, nlsolve, kwargs, autonomous) end -# Now make `affect!` for this: -function (p::ManifoldProjection{ - iip, nlls, autonomous, NL})(integrator) where {iip, nlls, - autonomous, NL} +function (proj::ManifoldProjection)(integrator) # update current time if residual function is time-dependent - !autonomous && (p.g.t = integrator.t) + proj.manifold.t = integrator.t + proj.manifold_jacobian !== nothing && (proj.manifold_jacobian.t = integrator.t) - # solve the nonlinear problem - reinit!(p.nlcache, integrator.u; p = integrator.p) - sol = solve!(p.nlcache) + SciMLBase.reinit!(proj.nlcache, integrator.u; integrator.p) + _, u, retcode = SciMLBase.solve!(proj.nlcache) - if !SciMLBase.successful_retcode(sol) - SciMLBase.terminate!(integrator, sol.retcode) + if !SciMLBase.successful_retcode(retcode) + SciMLBase.terminate!(integrator, retcode) return end - copyto!(integrator.u, sol.u) + SciMLBase.copyto!(integrator.u, u) end -function Manifold_initialize(cb, u, t, integrator) - return Manifold_initialize(cb.affect!, u, t, integrator) +function initialize_manifold_projection(cb, u, t, integrator) + return initialize_manifold_projection(cb.affect!, u, t, integrator) end -function Manifold_initialize( - affect!::ManifoldProjection{iip, nlls}, u, t, integrator) where {iip, nlls} - nlfunc = NonlinearFunction{iip}(affect!.g; affect!.resid_prototype) - nlprob = if nlls - NonlinearLeastSquaresProblem(nlfunc, u, integrator.p) +function initialize_manifold_projection(affect!::ManifoldProjection, u, t, integrator) + if affect!.autonomous === nothing + autonomous = maximum(SciMLBase.numargs(affect!.manifold.f)) == + 2 + SciMLBase.isinplace(integrator.f) + affect!.manifold.autonomous = autonomous + affect!.manifold_jacobian !== nothing && + (affect!.manifold_jacobian.autonomous = autonomous) + end + + affect!.manifold.t = t + affect!.manifold_jacobian !== nothing && (affect!.manifold_jacobian.t = t) + + if affect!.nlsolve === missing + cache = init_manifold_projection( + Val(SciMLBase.isinplace(integrator.f)), affect!.manifold, affect!.autodiff, + affect!.manifold_jacobian, u, integrator.p; affect!.kwargs...) else - NonlinearProblem(nlfunc, u, integrator.p) + cache = init_manifold_projection_nonlinear_problem( + Val(SciMLBase.isinplace(integrator.f)), affect!.manifold, affect!.autodiff, + affect!.manifold_jacobian, u, integrator.p, affect!.nlsolve; affect!.kwargs...) end - affect!.nlcache = init(nlprob, affect!.nlsolve; affect!.kwargs...) + affect!.nlcache = cache u_modified!(integrator, false) end -function ManifoldProjection(g; nlsolve = nothing, save = true, nlls = Val(true), - isinplace = Val(true), autonomous = nothing, resid_prototype = nothing, - kwargs...) - _nlls = SciMLBase._unwrap_val(nlls) - iip = SciMLBase._unwrap_val(isinplace) - if autonomous === nothing - if iip - autonomous = maximum(SciMLBase.numargs(g)) == 3 +export ManifoldProjection + +# wrapper for non-autonomous functions +function wrap_autonomous_function(autonomous::Union{Val{true}, Val{false}}, g) + g === nothing && return nothing + return TypedNonAutonomousFunction{SciMLBase._unwrap_val(autonomous)}(g, nothing) +end +function wrap_autonomous_function(autonomous::Union{Bool, Nothing}, g) + g === nothing && return nothing + autonomous = autonomous === nothing ? false : autonomous + return UntypedNonAutonomousFunction(autonomous, g, nothing) +end + +abstract type AbstractNonAutonomousFunction end + +@concrete mutable struct TypedNonAutonomousFunction{autonomous} <: + AbstractNonAutonomousFunction + f + t::Any +end + +(f::TypedNonAutonomousFunction{false})(res, u, p) = f.f(res, u, p, f.t) +(f::TypedNonAutonomousFunction{true})(res, u, p) = f.f(res, u, p) + +(f::TypedNonAutonomousFunction{false})(u, p) = f.f(u, p, f.t) +(f::TypedNonAutonomousFunction{true})(u, p) = f.f(u, p) + +@concrete mutable struct UntypedNonAutonomousFunction <: AbstractNonAutonomousFunction + autonomous::Bool + f + t::Any +end + +function (f::UntypedNonAutonomousFunction)(res, u, p) + return f.autonomous ? f.f(res, u, p) : f.f(res, u, p, f.t) +end +(f::UntypedNonAutonomousFunction)(u, p) = f.autonomous ? f.f(u, p) : f.f(u, p, f.t) + +# This is solving the langrange multiplier formulation. This is more accurate but at the +# same time significantly more expensive. +@concrete mutable struct NonlinearSolveManifoldProjectionCache{iip} + manifold + p + λ + z + ũ + gu_cache + nlcache + + first_call::Bool + J + manifold_jacobian + autodiff + di_extras +end + +function SciMLBase.reinit!( + cache::NonlinearSolveManifoldProjectionCache{iip}, u; p = cache.p) where {iip} + if !cache.first_call || (cache.ũ !== u || cache.p !== p) + compute_manifold_jacobian!(cache.J, cache.manifold_jacobian, cache.autodiff, + Val(iip), cache.manifold, cache.gu_cache, u, p, cache.di_extras) + end + cache.first_call = false + cache.ũ = u + cache.p = p + + cache.z[1:length(cache.λ)] .= false + cache.z[(length(cache.λ) + 1):end] .= vec(u) + SciMLBase.reinit!(cache.nlcache, cache.z; p = (u, cache.J, p)) +end + +function init_manifold_projection_nonlinear_problem( + IIP::Val{iip}, manifold, autodiff, manifold_jacobian, ũ, p, alg; + resid_prototype = nothing, kwargs...) where {iip} + if iip + if resid_prototype !== nothing + gu = similar(resid_prototype) + λ = similar(resid_prototype) else - autonomous = maximum(SciMLBase.numargs(g)) == 2 + @warn "`resid_prototype` not provided for in-place problem. Assuming size of \ + output is the same as input. This might be incorrect." maxlog=1 + gu = similar(ũ) + λ = similar(ũ) end + else + gu = nothing + λ = manifold(ũ, p) end - affect! = ManifoldProjection{iip, _nlls, SciMLBase._unwrap_val(autonomous)}( - g, nlsolve, resid_prototype, kwargs) - condition = (u, t, integrator) -> true - return DiscreteCallback(condition, affect!; initialize = Manifold_initialize, - save_positions = (false, save)) + + J, di_extras = setup_manifold_jacobian(manifold_jacobian, autodiff, IIP, manifold, + gu, ũ, p) + z = vcat(vec(λ), vec(ũ)) + + nlfunc = if iip + let λlen = length(λ), λsz = size(λ), zsz = size(ũ) + @views (resid, u, ps) -> begin + ũ2, J2, p2 = ps + λ2, z2 = u[1:λlen], u[(λlen + 1):end] + manifold(reshape(resid[1:λlen], λsz), reshape(z2, zsz), p2) + resid[(λlen + 1):end] .= z2 .- vec(ũ2) .+ vec(vec(J2' * λ2)) + end + end + else + let λlen = length(λ), zsz = size(ũ) + @views (u, ps) -> begin + ũ2, J2, p2 = ps + λ2, z2 = u[1:λlen], u[(λlen + 1):end] + gz = vec(manifold(reshape(z2, zsz), p2)) + resid = z2 .- vec(ũ2) .+ vec(J2' * λ2) + return vcat(gz, resid) + end + end + end + + nlprob = NonlinearProblem(NonlinearFunction{iip}(nlfunc), z, (ũ, J, p)) + nlcache = SciMLBase.init(nlprob, alg; kwargs...) + + return NonlinearSolveManifoldProjectionCache{iip}( + manifold, p, λ, z, ũ, gu, nlcache, true, J, manifold_jacobian, autodiff, di_extras) end -export ManifoldProjection +@views function SciMLBase.solve!(cache::NonlinearSolveManifoldProjectionCache{iip}) where {iip} + sol = SciMLBase.solve!(cache.nlcache) + (; u, retcode) = sol + λ = reshape(u[1:length(cache.λ)], size(cache.λ)) + ũ = reshape(u[(length(cache.λ) + 1):end], size(cache.ũ)) + return λ, ũ, retcode +end + +# This is the algorithm described in Hairer III. +@concrete mutable struct SingleFactorizeManifoldProjectionCache{iip} + manifold + p + abstol + maxiters::Int + + ũ + JJᵀfact::Any # LU might fail and we might end up doing QR + u_cache + λ_cache + gu_cache + + first_call::Bool + J + JJᵀ + manifold_jacobian + autodiff + di_extras +end + +function SciMLBase.reinit!( + cache::SingleFactorizeManifoldProjectionCache{iip}, u; p = cache.p) where {iip} + if !cache.first_call || (cache.ũ !== u || cache.p !== p) + compute_manifold_jacobian!(cache.J, cache.manifold_jacobian, cache.autodiff, + Val(iip), cache.manifold, cache.gu_cache, u, p, cache.di_extras) + mul!(cache.JJᵀ, cache.J, cache.J') + cache.JJᵀfact = safe_factorize!(cache.JJᵀ) + end + cache.first_call = false + cache.ũ = u + cache.p = p +end + +default_abstol(::Type{T}) where {T} = real(oneunit(T)) * (eps(real(one(T))))^(4 // 5) + +function init_manifold_projection(IIP::Val{iip}, manifold, autodiff, manifold_jacobian, ũ, + p; abstol = default_abstol(eltype(ũ)), maxiters = 1000, + resid_prototype = nothing, kwargs...) where {iip} + if iip + if resid_prototype !== nothing + gu = similar(resid_prototype) + λ = similar(resid_prototype) + else + @warn "`resid_prototype` not provided for in-place problem. Assuming size of \ + output is the same as input. This might be incorrect." maxlog=1 + gu = similar(ũ) + λ = similar(ũ) + end + else + gu = nothing + λ = manifold(ũ, p) + end + + J, di_extras = setup_manifold_jacobian(manifold_jacobian, autodiff, IIP, manifold, + gu, ũ, p) + JJᵀ = J * J' + JJᵀfact = safe_factorize!(JJᵀ) + + return SingleFactorizeManifoldProjectionCache{iip}( + manifold, p, abstol, maxiters, ũ, JJᵀfact, similar(ũ), λ, gu, + true, J, JJᵀ, manifold_jacobian, autodiff, di_extras) +end + +function SciMLBase.solve!(cache::SingleFactorizeManifoldProjectionCache{iip}) where {iip} + fill!(cache.λ_cache, false) + ũ = cache.ũ + gu = cache.gu_cache + + internal_solve_failed = true + + if cache.gu_cache !== nothing + cache.manifold(gu, ũ, cache.p) + else + gu = cache.manifold(ũ, cache.p) + end + + for _ in 1:(cache.maxiters) + if maximum(abs, gu) < cache.abstol + internal_solve_failed = false + break + end + + δλ = cache.JJᵀfact \ gu + @. cache.λ_cache -= δλ + + mul!(vec(cache.u_cache), cache.J', vec(cache.λ_cache)) + cache.u_cache += ũ + if cache.gu_cache !== nothing + cache.manifold(gu, cache.u_cache, cache.p) + else + gu = cache.manifold(cache.u_cache, cache.p) + end + end + + return (cache.λ_cache, cache.u_cache, + ifelse(internal_solve_failed, ReturnCode.ConvergenceFailure, ReturnCode.Success)) +end + +function setup_manifold_jacobian( + manifold_jacobian::M, autodiff, ::Val{iip}, manifold, gu, ũ, p) where {M, iip} + if iip + J = similar(ũ, promote_type(eltype(gu), eltype(ũ)), (length(gu), length(ũ))) + manifold_jacobian(J, ũ, p) + else + J = manifold_jacobian(ũ, p) + end + return J, nothing +end + +function setup_manifold_jacobian( + ::Nothing, autodiff, ::Val{iip}, manifold, gu, ũ, p) where {iip} + if iip + di_extras = DI.prepare_jacobian(manifold, gu, autodiff, ũ, Constant(p)) + J = DI.jacobian(manifold, gu, di_extras, autodiff, ũ, Constant(p)) + else + di_extras = DI.prepare_jacobian(manifold, autodiff, ũ, Constant(p)) + J = DI.jacobian(manifold, di_extras, autodiff, ũ, Constant(p)) + end + return J, di_extras +end + +function setup_manifold_jacobian( + ::Nothing, ::Nothing, ::Val{iip}, manifold, gu, ũ, p) where {iip} + error("`autodiff` is set to `nothing` and analytic manifold jacobian is not provided.") +end + +function compute_manifold_jacobian!(J, manifold_jacobian, autodiff, ::Val{iip}, + manifold, gu, ũ, p, di_extras) where {iip} + if iip + manifold_jacobian(J, ũ, p) + else + J = manifold_jacobian(ũ, p) + end + return J +end + +function compute_manifold_jacobian!(J, ::Nothing, autodiff, ::Val{iip}, manifold, gu, + ũ, p, di_extras) where {iip} + if iip + DI.jacobian!(manifold, gu, J, di_extras, autodiff, ũ, Constant(p)) + else + DI.jacobian!(manifold, J, di_extras, autodiff, ũ, Constant(p)) + end + return J +end + +function safe_factorize!(A::AbstractMatrix) + if issquare(A) + fact = LinearAlgebra.cholesky(A; check = false) + fact_successful(fact) && return fact + elseif size(A, 1) > size(A, 2) + fact = LinearAlgebra.qr(A) + fact_successful(fact) && return fact + end + return LinearAlgebra.qr!(A, LinearAlgebra.ColumnNorm()) +end + +function fact_successful(F::LinearAlgebra.QRCompactWY) + m, n = size(F) + U = view(F.factors, 1:min(m, n), 1:n) + return all(!iszero, Iterators.reverse(@view U[diagind(U)])) +end +function fact_successful(F::FT) where {FT} + return hasmethod(LinearAlgebra.issuccess, (FT,)) ? LinearAlgebra.issuccess(F) : true +end diff --git a/test/domain_tests.jl b/test/domain_tests.jl index 9aaee626..6777aa07 100644 --- a/test/domain_tests.jl +++ b/test/domain_tests.jl @@ -1,4 +1,4 @@ -using DiffEqCallbacks, OrdinaryDiffEq, Test +using DiffEqCallbacks, OrdinaryDiffEq, Test, ADTypes, NonlinearSolve # Non-negative ODE examples # @@ -39,7 +39,11 @@ naive_sol_absval = solve(prob_absval, BS3()) function g(resid, u, p) resid[1] = u[1] < 0 ? -u[1] : 0 end -general_sol_absval = solve(prob_absval, BS3(); callback = GeneralDomain(g, [1.0]), +general_sol_absval = solve( + prob_absval, BS3(); + callback = GeneralDomain(g, [1.0]; + autodiff = AutoForwardDiff(), + nlsolve = NewtonRaphson(; autodiff = AutoForwardDiff())), save_everystep = false) @test all(x -> x[1] ≥ 0, general_sol_absval.u) @test general_sol_absval.errors[:l∞] < 9.9e-5 @@ -49,7 +53,11 @@ general_sol_absval = solve(prob_absval, BS3(); callback = GeneralDomain(g, [1.0] # test "non-autonomous" function g_t(resid, u, p, t) = g(resid, u, p) -general_t_sol_absval = solve(prob_absval, BS3(); callback = GeneralDomain(g_t, [1.0]), +general_t_sol_absval = solve( + prob_absval, BS3(); + callback = GeneralDomain(g_t, [1.0]; + autodiff = AutoForwardDiff(), + nlsolve = NewtonRaphson(; autodiff = AutoForwardDiff())), save_everystep = false) @test general_sol_absval.t ≈ general_t_sol_absval.t @test general_sol_absval.u ≈ general_t_sol_absval.u diff --git a/test/manifold_tests.jl b/test/manifold_tests.jl index 67e8a994..40fc6633 100644 --- a/test/manifold_tests.jl +++ b/test/manifold_tests.jl @@ -1,4 +1,5 @@ using OrdinaryDiffEq, Test, DiffEqBase, DiffEqCallbacks, RecursiveArrayTools, NonlinearSolve +using ForwardDiff, ADTypes u0 = ones(2, 2) f = function (du, u, p, t) @@ -14,46 +15,44 @@ end g_t(resid, u, p, t) = g(resid, u, p) -function isautonomous(::ManifoldProjection{ - iip, nlls, autonomous, NL}) where {iip, nlls, autonomous, NL} - return autonomous -end - sol = solve(prob, Vern7()) @test !(sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2) # autodiff=true @inferred ManifoldProjection(g; autonomous = Val(false), resid_prototype = zeros(2)) -cb = ManifoldProjection(g; resid_prototype = zeros(2)) -@test isautonomous(cb.affect!) +cb = ManifoldProjection(g; resid_prototype = zeros(2), autodiff = AutoForwardDiff()) solve(prob, Vern7(), callback = cb) @time sol = solve(prob, Vern7(), callback = cb) @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2 -cb_t = ManifoldProjection(g_t; resid_prototype = zeros(2)) -@test !isautonomous(cb_t.affect!) +cb_t = ManifoldProjection(g_t; resid_prototype = zeros(2), autodiff = AutoForwardDiff()) solve(prob, Vern7(), callback = cb_t) @time sol_t = solve(prob, Vern7(), callback = cb_t) @test sol_t.u == sol.u && sol_t.t == sol.t # autodiff=false cb_false = ManifoldProjection( - g; nlsolve = GaussNewton(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2)) -@test isautonomous(cb_false.affect!) + g; nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2), + autodiff = AutoFiniteDiff()) solve(prob, Vern7(), callback = cb_false) sol = solve(prob, Vern7(), callback = cb_false) @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2 cb_t_false = ManifoldProjection(g_t, - nlsolve = GaussNewton(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2)) -@test !isautonomous(cb_t_false.affect!) + nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), resid_prototype = zeros(2), + autodiff = AutoFiniteDiff()) solve(prob, Vern7(), callback = cb_t_false) sol_t = solve(prob, Vern7(), callback = cb_t_false) @test sol_t.u == sol.u && sol_t.t == sol.t # test array partitions +function f_ap!(du, u, p, t) + du[1:2] .= u[3:4] + du[3:4] .= u[1:2] +end + u₀ = ArrayPartition(ones(2), ones(2)) -prob = ODEProblem(f, u₀, (0.0, 100.0)) +prob = ODEProblem(f_ap!, u₀, (0.0, 100.0)) sol = solve(prob, Vern7(), callback = cb) @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2 @@ -73,7 +72,14 @@ function g_unsat(resid, u, p) resid[2] = u[2]^2 + u[1]^2 - 20 end -cb_unsat = ManifoldProjection(g_unsat; resid_prototype = zeros(2)) +cb_unsat = ManifoldProjection( + g_unsat; resid_prototype = zeros(2), autodiff = AutoForwardDiff()) +sol = solve(prob, Vern7(), callback = cb_unsat) +@test !SciMLBase.successful_retcode(sol) +@test last(sol.t) != 100.0 + +cb_unsat = ManifoldProjection( + g_unsat; resid_prototype = zeros(2), autodiff = AutoForwardDiff(), nlsolve = NewtonRaphson()) sol = solve(prob, Vern7(), callback = cb_unsat) @test !SciMLBase.successful_retcode(sol) @test last(sol.t) != 100.0 @@ -86,40 +92,41 @@ end g_oop_t(u, p, t) = g_oop(u, p) -prob = ODEProblem(f, u0, (0.0, 100.0)) +f_oop = function (u, p, t) + return stack((u[2, :], -u[1, :])) +end +prob = ODEProblem(f_oop, u0, (0.0, 100.0)) # autodiff=true -@inferred ManifoldProjection(g_oop; autonomous = Val(false), isinplace = Val(false)) -cb = ManifoldProjection(g_oop; isinplace = Val(false)) -@test isautonomous(cb.affect!) +@inferred ManifoldProjection(g_oop; autonomous = Val(false)) +cb = ManifoldProjection(g_oop; autodiff = AutoForwardDiff()) solve(prob, Vern7(), callback = cb) @time sol = solve(prob, Vern7(), callback = cb) @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2 -cb_t = ManifoldProjection(g_oop_t; isinplace = Val(false)) -@test !isautonomous(cb_t.affect!) +cb_t = ManifoldProjection(g_oop_t; autodiff = AutoForwardDiff()) solve(prob, Vern7(), callback = cb_t) @time sol_t = solve(prob, Vern7(), callback = cb_t) @test sol_t.u == sol.u && sol_t.t == sol.t # autodiff=false cb_false = ManifoldProjection( - g_oop; nlsolve = GaussNewton(; autodiff = AutoFiniteDiff()), isinplace = Val(false)) -@test isautonomous(cb_false.affect!) + g_oop; nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), autodiff = AutoFiniteDiff()) solve(prob, Vern7(), callback = cb_false) sol = solve(prob, Vern7(), callback = cb_false) @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2 cb_t_false = ManifoldProjection(g_oop_t, - nlsolve = GaussNewton(; autodiff = AutoFiniteDiff()), isinplace = Val(false)) -@test !isautonomous(cb_t_false.affect!) + nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff()), autodiff = AutoFiniteDiff()) solve(prob, Vern7(), callback = cb_t_false) sol_t = solve(prob, Vern7(), callback = cb_t_false) @test sol_t.u == sol.u && sol_t.t == sol.t # test array partitions +f_ap(u, p, t) = ArrayPartition(u[3:4], u[1:2]) + u₀ = ArrayPartition(ones(2), ones(2)) -prob = ODEProblem(f, u₀, (0.0, 100.0)) +prob = ODEProblem(f_ap, u₀, (0.0, 100.0)) sol = solve(prob, Vern7(), callback = cb) @test sol.u[end][1]^2 + sol.u[end][2]^2 ≈ 2 diff --git a/test/preset_time.jl b/test/preset_time.jl index bf9b4a23..be082135 100644 --- a/test/preset_time.jl +++ b/test/preset_time.jl @@ -91,4 +91,5 @@ sol2 = solve(prob2, Tsit5(), callback = cb1) @test sol2(48.0 + eps(48.0)) ≈ [10.0] _some_test_func(integrator) = u_modified!(integrator, false) -@inferred PresetTimeCallback(collect(range(0, 10, 100)), _some_test_func, save_positions=(false, false)) +@inferred PresetTimeCallback( + collect(range(0, 10, 100)), _some_test_func, save_positions = (false, false))