From cf38dcdadaa9bcf67144a930f8636b617fa8656e Mon Sep 17 00:00:00 2001 From: Azercoco <20425137+Azercoco@users.noreply.github.com> Date: Tue, 10 Jun 2025 15:41:21 +0200 Subject: [PATCH 01/10] First commit, draft implementation of RKIP --- lib/OrdinaryDiffEqRKIP/LICENSE.md | 24 +++ lib/OrdinaryDiffEqRKIP/Project.toml | 44 +++++ .../src/OrdinaryDiffEqRKIP.jl | 23 +++ lib/OrdinaryDiffEqRKIP/src/algorithms.jl | 155 ++++++++++++++++++ lib/OrdinaryDiffEqRKIP/src/rkip_cache.jl | 63 +++++++ .../src/rkip_perform_step.jl | 112 +++++++++++++ lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl | 56 +++++++ lib/OrdinaryDiffEqRKIP/test/runtests.jl | 5 + lib/OrdinaryDiffEqRKIP/test/semilinear_fft.jl | 141 ++++++++++++++++ .../test/semilinear_pde_test_cpu.jl | 34 ++++ .../test/semilinear_pde_test_cuda.jl | 35 ++++ lib/OrdinaryDiffEqRKIP/test/type_test.jl | 84 ++++++++++ src/OrdinaryDiffEq.jl | 4 + 13 files changed, 780 insertions(+) create mode 100644 lib/OrdinaryDiffEqRKIP/LICENSE.md create mode 100644 lib/OrdinaryDiffEqRKIP/Project.toml create mode 100644 lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl create mode 100644 lib/OrdinaryDiffEqRKIP/src/algorithms.jl create mode 100644 lib/OrdinaryDiffEqRKIP/src/rkip_cache.jl create mode 100644 lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl create mode 100644 lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl create mode 100644 lib/OrdinaryDiffEqRKIP/test/runtests.jl create mode 100644 lib/OrdinaryDiffEqRKIP/test/semilinear_fft.jl create mode 100644 lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cpu.jl create mode 100644 lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cuda.jl create mode 100644 lib/OrdinaryDiffEqRKIP/test/type_test.jl diff --git a/lib/OrdinaryDiffEqRKIP/LICENSE.md b/lib/OrdinaryDiffEqRKIP/LICENSE.md new file mode 100644 index 0000000000..4a7df96ac5 --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/LICENSE.md @@ -0,0 +1,24 @@ +The OrdinaryDiffEq.jl package is licensed under the MIT "Expat" License: + +> Copyright (c) 2016-2020: ChrisRackauckas, Yingbo Ma, Julia Computing Inc, and +> other contributors: +> +> https://github.com/SciML/OrdinaryDiffEq.jl/graphs/contributors +> +> Permission is hereby granted, free of charge, to any person obtaining a copy +> of this software and associated documentation files (the "Software"), to deal +> in the Software without restriction, including without limitation the rights +> to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +> copies of the Software, and to permit persons to whom the Software is +> furnished to do so, subject to the following conditions: +> +> The above copyright notice and this permission notice shall be included in all +> copies or substantial portions of the Software. +> +> THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +> IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +> FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +> AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +> LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +> OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +> SOFTWARE. diff --git a/lib/OrdinaryDiffEqRKIP/Project.toml b/lib/OrdinaryDiffEqRKIP/Project.toml new file mode 100644 index 0000000000..2c30c83f7a --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/Project.toml @@ -0,0 +1,44 @@ +name = "OrdinaryDiffEqRKIP" +uuid = "a4daff8c-1d43-4ff3-8eff-f78720aeecdc" +authors = ["Azercoco <20425137+Azercoco@users.noreply.github.com>"] +version = "1.0.0" + + +[sources] +OrdinaryDiffEqCore = {path = "../OrdinaryDiffEqCore"} + +[deps] +DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" +DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" +OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8" +SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" + +[compat] +DiffEqBase = "6.175" +DiffEqDevTools = "2.48" +MaybeInplace = "0.1.4" +OrdinaryDiffEqCore = "1.26" +SciMLBase = "2.99" +SciMLOperators = "1.3.1" +StaticArrays = "1.9.13" +UnPack = "1.0.2" +LinearAlgebra = "1.11" +CUDA = "5.5.2" +FFTW = "1.8.0" +SafeTestsets = "0.1.0" + +julia = "1.11" + +[extras] +FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" +SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + +[targets] +test = ["FFTW", "Test", "SafeTestsets", "CUDA"] diff --git a/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl b/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl new file mode 100644 index 0000000000..36aa49f87e --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl @@ -0,0 +1,23 @@ +module OrdinaryDiffEqRKIP + +using LinearAlgebra: ldiv!, exp, axpy!, norm, mul! +using SciMLOperators: AbstractSciMLOperator +using UnPack: @pack!, @unpack +using MaybeInplace: @bb +using SciMLBase: isinplace +using DiffEqBase: ExplicitRKTableau +using DiffEqDevTools: constructDormandPrince6 + +import OrdinaryDiffEqCore: OrdinaryDiffEqAdaptiveExponentialAlgorithm, alg_adaptive_order, + alg_order, alg_cache, @cache, SplitFunction, get_fsalfirstlast, initialize!, perform_step!, + has_dtnew_modification, calculate_residuals, calculate_residuals!, increment_nf!, OrdinaryDiffEqAdaptiveAlgorithm, OrdinaryDiffEqMutableCache, dtnew_modification + + +include("rkip_cache.jl") +include("algorithms.jl") +include("rkip_utils.jl") +include("rkip_perform_step.jl") + +export RKIP + +end \ No newline at end of file diff --git a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl new file mode 100644 index 0000000000..33d9cb2d69 --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl @@ -0,0 +1,155 @@ +using OrdinaryDiffEqCore: OrdinaryDiffEqCore + +mutable struct RKIP_ALG{tableauType<:ExplicitRKTableau,elType,dtType<:AbstractVector{elType}} <: OrdinaryDiffEqAdaptiveAlgorithm + tableau::tableauType + dt_for_expÂ_caching::dtType + clamp_lower_dt::Bool + clamp_higher_dt::Bool + use_ldiv::Bool + cache::Union{Nothing,RKIPCache} +end + +alg_order(alg::RKIP_ALG) = alg.tableau.order +alg_adaptive_order(alg::RKIP_ALG) = alg.tableau.adaptiveorder + +has_dtnew_modification(alg::RKIP_ALG) = true + +function dtnew_modification(alg::RKIP_ALG{tableauType,elType,dtType}, dtnew) where {tableauType,elType,dtType} + @unpack dt_for_expÂ_caching = alg + if first(alg.dt_for_expÂ_caching) > dtnew && alg.clamp_lower_dt + dtnew = first(alg.dt_for_expÂ_caching) + elseif last(alg.dt_for_expÂ_caching) < dtnew && alg.clamp_higher_dt + dtnew = last(alg.dt_for_expÂ_caching) + else + dtnew = alg.dt_for_expÂ_caching[searchsortedfirst(alg.dt_for_expÂ_caching, dtnew)] + end + return dtnew +end + +dtnew_modification(_, alg::RKIP_ALG, dtnew) = dtnew_modification(alg, dtnew) + + +""" +Runge-Kutta in the Interaction Picture (RIKP) method. + +This is suited for solving semilinear problem of the form: + +```math +\frac{du}{dt} = Au + f(u,p,t) +``` +where A is possibly stiff constant linear operator whose scaled exponential exp(Ah) can be calculated efficiently for any h. + +The problem is first transformed into the Interaction Picture. + +```math +u_I(t) = \exp(At) u_0 +\frac{du_I}{dt} = f_I(u_I,p,t) +f_I(u_I,p,t) = f(exp(-At)u_I, p, t) +``` +This new system is then solved with an explicit Runge-Kutta method. + +This type of problem is often encountered in semilinear parabolic PDE: heat diffusion, Schrödinger equation ... + +This method is only implemented for semilinear problem: `SplitODEProblem` when the first function `f1` is a `AbsractSciMLOperator` A. + +The algorithm needs for A to have the following traits of the `AbsractSciMLOperator` implemented + +```julia +LinearAlgebra.exp(A, t) = #... +A(du, u, v, p, t) # for in-place problem +A(u, v p, t) # for out-of-place problem +``` + +Moreover, the resulting ```expA = exp(A, t)``` must also implement ```(expA)(u, v, p, t)``` and/or ```(expA)(du, u, v, p, t)```. +The algorithm only works for constant `A` (does no depend on v, p, t). + +The algorithm will cache and resue the computed operator-exponential for a set of steps. + +# Arguments +- `dtmin::T`: the smallest step `dt` for which `exp(A*dt)` will be cached. Default is `1e-3` +- `dtmax::T`: the largest step `dt` for which `exp(A*dt)` will be cached. Default is `1.0` + +The set of cached steps size `dt_for_expA_caching` will follow a geometric progression such that `first(dt_for_expA_caching) = dtmin` and `last(dt_for_expA_caching) = dtmax`. + +Time stepping can still happen outside the bonds but this will results in no cache (as `exp(A*h)` will be computed each step) degrading the performances. The time step can be clamped within the cache range through the args `clamp_lower_dt` and `clamp_higher_dt`. + +# Optional Arguments +- `nb_of_cache_step::Integer`: the number of steps. Default is `100`. +- `tableau::ExplicitRKTableau`: the Runge-Kutta Tableau to use, default is `constructDormandPrince6()`. +- `clamp_lower_dt::Bool`: weither to clamp proposed step to the smallest cached step in order to force the use of cached exponential, improving performance. + This may prevent reaching the desired tolerance. Default is `false`. +- `clamp_higher_dt::Bool`: weither to clamp proposed step to the largest cached step in order to force the use of cached exponential, improving performance. + This can cause performance degredation if `integrator.dtmax` is too small. Default is `true`. +- `use_ldiv::Bool`: weither, to use `ldiv(exp(A, t), v)` instead of caching `exp(A, -t)*v`. Reduce the memory usage but slightly less efficient. `ldiv` must be implemented. Only works for in-place problem. Default is `false`. + +Cached operator exponentials are stored in the alorithm meaning that in the following code: + +```julia +rkip = RKIP() +solve(ode_prob_1, rkip, t1) +solve(ode_prob_2, rkip, t2) +```` + +will reuse the cache of the first solve. This can be useful when needed to solve several time the same problem but one must be sure that `A` does not change. + +""" +RKIP(dtmin::T=1e-3, dtmax::T=1.0; nb_of_cache_step::Int=100, tableau=constructDormandPrince6(T), clamp_lower_dt::Bool=false, clamp_higher_dt::Bool=true, use_ldiv=false) where {T} = RKIP_ALG{typeof(tableau),T,Vector{T}}( + tableau, + logrange(dtmin, dtmax, nb_of_cache_step), + clamp_lower_dt, + clamp_higher_dt, + use_ldiv, + nothing +) + +function alg_cache(alg::RKIP_ALG, u::uType, rate_prototype, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits, uprev, uprev2, f, t, dt, reltol, p, calck, iip) where {uType} + tmp = zero(u) + utilde = zero(u) + kk = [zero(u) for _ in 1:(alg.tableau.stages)] + + Â = isa(f, SplitFunction) ? f.f1.f : throw(ArgumentError("RKIP is only implemented for semilinear problems")) + opType = typeof(Â) + expOpType = typeof(exp(Â, 1.0)) + + if isnothing(alg.cache) + + is_cached = Vector{Bool}(undef, length(alg.dt_for_expÂ_caching)) + fill!(is_cached, false) + + c_extended = vcat(alg.tableau.c, 1.0) # all the c values of Runge-Kutta and 1 wich is needed for the RKIP step + c_unique = unique(c_extended) # in some tableau, there is duplicate: we only keep the unique value to save on caching time and memory + c_index = [findfirst(==(c), c_unique) for c in c_extended] # index mapping + + exp_cache = ExpCache{expOpType}(Array{expOpType,2}(undef, length(alg.dt_for_expÂ_caching), length(c_unique)), Vector{expOpType}(undef, length(c_unique))) + + if !alg.use_ldiv + exp_cache = ExpCacheNoLdiv(exp_cache, ExpCache{expOpType}(Array{expOpType,2}(undef, length(alg.dt_for_expÂ_caching), length(c_unique)), Vector{expOpType}(undef, length(c_unique)))) + expCacheType = ExpCacheNoLdiv{expOpType} + else + expCacheType = ExpCache{expOpType} + end + + alg.cache = RKIPCache{expOpType,expCacheType,tTypeNoUnits,opType,uType,iip}( + exp_cache, + zero(tTypeNoUnits), + is_cached, + tmp, + utilde, + kk, + c_unique, + c_index + ) + else # cache recycling + alg.cache = RKIPCache{expOpType,typeof(alg.cache.exp_cache),tTypeNoUnits,opType,uType,iip}( + alg.cache.exp_cache, + alg.last_step, + alg.cache.is_cached, + tmp, + utilde, + kk, + alg.cache.c_unique, + alg.cache.c_index + ) + end + return alg.cache +end diff --git a/lib/OrdinaryDiffEqRKIP/src/rkip_cache.jl b/lib/OrdinaryDiffEqRKIP/src/rkip_cache.jl new file mode 100644 index 0000000000..6d3353160e --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/src/rkip_cache.jl @@ -0,0 +1,63 @@ +abstract type AbstractExpCache{expOpType<:AbstractSciMLOperator} end + +struct ExpCache{expOpType} <: AbstractExpCache{expOpType} + expÂ_cached::Array{expOpType,2} + expÂ_for_this_step::Vector{expOpType} +end +struct ExpCacheNoLdiv{expOpType} <: AbstractExpCache{expOpType} + exp_cache::ExpCache{expOpType} + exp_cache_inv::ExpCache{expOpType} +end + +get_op_for_this_step(cache::ExpCache{expOpType}, index::Int) where {expOpType} = cache.expÂ_for_this_step[index] +get_op_for_this_step(cache_no_ldiv::ExpCacheNoLdiv{expOpType}, positive::Bool, index::Int) where {expOpType} = positive ? cache_no_ldiv.exp_cache.expÂ_for_this_step[index] : cache_no_ldiv.exp_cache_inv.expÂ_for_this_step[index] + + +mutable struct RKIPCache{expOpType<:AbstractSciMLOperator,cacheType<:AbstractExpCache{expOpType},tType<:Number,opType<:AbstractSciMLOperator,uType,iip} <: OrdinaryDiffEqMutableCache + exp_cache::cacheType + last_step::tType + cached::Vector{Bool} + tmp::uType + utilde::uType + kk::Vector{uType} + c_unique::Vector{tType} + c_mapping::Vector{Integer} +end + +get_fsalfirstlast(cache::RKIPCache, u) = (zero(cache.tmp), zero(cache.tmp)) + + +@inline function cache_exp!(cache::ExpCache{expOpType}, A::opType, h::T, action::Symbol, step_index::Int, unique_stage_index::Int) where {expOpType<:AbstractSciMLOperator,opType<:AbstractSciMLOperator,T<:Number} + @unpack expÂ_for_this_step, expÂ_cached = cache + expÂ_for_this_step[unique_stage_index] = (action == :use_cached) ? expÂ_cached[step_index, unique_stage_index] : exp(A, h) # fetching or generating exp(Â*c_i*dt) + if action == :cache + expÂ_cached[step_index, unique_stage_index] = expÂ_for_this_step[unique_stage_index] # storing exp(Â*c_i*dt) + end +end + +@inline function cache_exp!(cache::ExpCacheNoLdiv{expOpType}, Â::opType, h::T, action::Symbol, step_index::Int, unique_stage_index::Int) where {expOpType<:AbstractSciMLOperator,opType<:AbstractSciMLOperator,T<:Number} + cache_exp!(cache.exp_cache, Â, h, action, step_index, unique_stage_index) + cache_exp!(cache.exp_cache_inv, Â, -h, action, step_index, unique_stage_index) +end + +""" + Prepare/generate all the needed exp(± A * dt * c[i]) for a step size dt +""" +@inline function cache_exp_op_for_this_step!(cache::RKIPCache{expOpType,cacheType,tType,opType,uType,iip}, Â::opType, dt::tType, alg::algType) where {expOpType,cacheType,tType,opType,uType,algType,iip} + @unpack dt_for_expÂ_caching = alg + + if !iszero(dt) && !(dt ≈ cache.last_step) # we check that new exp(A dt) are needed + dt_abs = abs(dt) # only the positive dt are used for indexing + + step_index = clamp(searchsortedlast(dt_for_expÂ_caching, dt_abs), 1, lastindex(dt_for_expÂ_caching)) # fetching the index corresponding to the step size + action = (dt_for_expÂ_caching[step_index] ≈ dt_abs) ? (cache.cached[step_index] ? :use_cached : :cache) : :single_use # if alreay present, we reuse the cached, otherwise it is generated + + for (unique_stage_index, c) in enumerate(cache.c_unique) # iterating over all unique c_i of the RK tableau + cache_exp!(cache.exp_cache, Â, abs(dt * c), action, step_index, unique_stage_index) # generating and caching + end + cache.cached[step_index] |= (action == :cache) # set the flag that we have already cached exp(Â*c_i*dt) for this dt + end + + cache.last_step = dt +end + diff --git a/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl b/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl new file mode 100644 index 0000000000..575653c3f1 --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl @@ -0,0 +1,112 @@ +""" +Helper function for the (nonlinear) part of the problem maybe in place. +Overwrite compute f(u, p, t) and if possible, overwrite u with the result. +Same principle of operation as `matvec_prod_mip` for mutability/in-place handling. +""" +function nl_part_mip(tmp::uType, f!, u::uType, p, t::tType, ::Val{true}) where {tType,uType} + f!(tmp, u, p, t) + copyto!(u, tmp) + return u +end +nl_part_mip(_::uType, f, u::uType, p, t::tType, ::Val{false}) where {tType,uType} = f(u, p, t) + + +""" +Helper function to compute Au + f(u, p, t) and if in place, store the result in res. +Retunr res if in place, otherwise return Au + f(u, p, t) +""" +function f_mip!(res::uType, tmp::uType, A::opType, f, u::uType, p, t::tType, ::Val{true}) where {tType,uType,opType} + res .= u + res = matvec_prod_mip(tmp, A, res, Val(true)) + f(tmp, u, p, t) + res .+= tmp + return res +end + +f_mip!(_::uType, tmp::uType, A::opType, f, u::uType, p, t::tType, ::Val{false}) where {tType,uType,opType} = matvec_prod_mip(tmp, A, u, Val(false)) + f(u, p, t) + +""" +Helper function for the residual maybe in place. +Same principle of operation as `_safe_matvec_prod` for mutability/in-place handling. +""" +function calculate_residuals_mip(tmp::uType, utilde::uType, uprev::uType, u::uType, abstol, reltol, internalnorm, t, ::Val{true}) where {uType} + calculate_residuals!(tmp, utilde, uprev, u, abstol, reltol, internalnorm, t) + return tmp +end +calculate_residuals_mip(_::uType, utilde::uType, uprev::uType, u::uType, abstol, reltol, internalnorm, t, ::Val{false}) where {uType} = calculate_residuals(utilde, uprev, u, abstol, reltol, internalnorm, t) + +@fastmath function perform_step!(integrator, cache::RKIPCache{expOpType,cacheType,tType,opType,uType,iip}) where {expOpType,cacheType,tType,opType,uType,iip} + @unpack t, dt, uprev, u, f, p, fsalfirst, fsallast, alg = integrator + @unpack c, α, αEEst, stages, A = alg.tableau + @unpack kk, utilde, tmp = cache + @unpack adaptive, abstol, reltol, internalnorm = integrator.opts + + Â::opType = f.f1.f # Linear Operator + + cache_exp_op_for_this_step!(cache, Â, dt, alg) # compute/reuse cached exp(± Â * [dt * cᵢ]) for this dt and cache them if possible + + @bb u .= uprev + + @inbounds for i in 1:(stages) + @bb kk[i] .= uprev + + @inbounds for j in 1:(i-1) + g = A[i, j] * dt + kk[i] = axpy_mip(g, kk[j], kk[i], iip) # kk_i += dt*A[i, j] * kk_j + end + + # if mutable/heaps type, assignment does nothing as the function is in-place, + kk[i] = expmv_rkip_mip(cache, kk[i], dt, i) # kᵢ = exp(Â * [dt * cᵢ])*kᵢ ➡ Change from interaction picture to "true" coordinate + + kk[i] = nl_part_mip(cache.tmp, f.f2, kk[i], p, t + c[i] * dt, iip) # kᵢ = f(u + Σⱼ dt*(Aᵢⱼ kⱼ), t + dt*cᵢ) + + + kk[i] = expmv_rkip_mip(cache, kk[i], -dt, i) # kᵢ = exp(-Â * [dt * cᵢ])*kᵢ ➡ Going back in interaction picture + + integrator.stats.nf += 2 # two exp vec product + integrator.stats.nf2 += 1 # one function evaluation + + u = axpy_mip(α[i] * dt, kk[i], u, iip) # uₙ = uₙ₋₁ + dt Σᵢ * (αᵢ - α*ᵢ) kᵢ + + if adaptive # error estimation ũ = Σᵢ dt * (αᵢ - α*ᵢ) kᵢ + if i == 1 + @bb @. utilde = dt * (α[i] - αEEst[i]) * kk[i] # to avoid filling with zero, we set it to i == 1 + else + utilde = axpy_mip(dt * (α[i] - αEEst[i]), kk[i], utilde, iip) # otherwise we add it + end + end + end + + u = expmv_rkip_mip(cache, u, dt) # stepping forward into the interaction u = exp(Â dt)*u + + if adaptive + utilde = expmv_rkip_mip(cache, utilde, dt) # stepping forward into the interaction ũ = exp(Â dt)*ũ + tmp = calculate_residuals_mip(tmp, utilde, uprev, u, abstol, reltol, internalnorm, t, iip) # error computation maybe in place + integrator.EEst = internalnorm(tmp, t) + end + + fsallast = f_mip!(fsallast, cache.tmp, Â, f.f2, u, p, t + dt, iip) # derivative estimation for interpolation + integrator.stats.nf += 1 + integrator.stats.nf2 += 1 + @pack! integrator = fsalfirst, fsallast, u + + @bb copyto!(integrator.k[1], fsalfirst) + @bb copyto!(integrator.k[2], fsallast) +end + + +function initialize!(integrator, cache::RKIPCache{expOpType,cacheType,tType,opType,uType,iip}) where {expOpType,cacheType,tType,opType,uType,iip} + @unpack f, u, p, t, fsalfirst, fsallast = integrator + + kshortsize = 2 + k = [zero(u) for _ ∈ 1:kshortsize] + + fsalfirst = f_mip!(fsalfirst, cache.tmp, f.f1.f, f.f2, u, p, t, iip) # first derivative for interpolation computation, maybe in place + integrator.stats.nf += 1 + integrator.stats.nf2 += 1 + + @pack! integrator = kshortsize, k, fsalfirst, fsallast + + @bb copyto!(integrator.k[1], fsalfirst) + @bb copyto!(integrator.k[2], fsallast) +end \ No newline at end of file diff --git a/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl b/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl new file mode 100644 index 0000000000..14fba179bc --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl @@ -0,0 +1,56 @@ +# cannot use @bb on axpy due to https://github.com/SciML/MaybeInplace.jl/issues/14 +@inline axpy_mip(α, x::uType, y::uType, ::Val{true}) where {uType} = axpy!(α, x, y) #BLAS call +@inline axpy_mip(α, x::uType, y::uType, ::Val{false}) where {uType} = y .+ α .* x + + +""" +Maybe-In Place "mip" +Safe function for computing the matrix exponential vector product for both mutable and immutable type. +Same principle of operation as `_safe_matvec_prod` for mutability/in-place handling. +""" +@inline expmv_rkip_mip(cache::RKIPCache{expOpType,cacheType,tType,opType,uType,iip}, v::uType, h::tType) where {expOpType,cacheType,tType,opType,uType,iip} = expmv_rkip_mip(cache, v, h, lastindex(cache.c_mapping)) # If i is not precised, this is the final step in the RKIP -> h = dt + +@inline function expmv_rkip_mip(cache::RKIPCache{expOpType,cacheType,tType,opType,uType,iip}, v::uType, h::tType, stage_index::Int) where {expOpType,cacheType,tType,opType,uType,iip} + + if !(h ≈ tType(0.0)) + c = cache.c_mapping[stage_index] + exp_cache = cache.exp_cache + tmp = cache.tmp + return _expmv_rkip_mip(exp_cache, tmp, v, c, h > tType(0.0), iip) + end + return v +end + +@inline function _expmv_rkip_mip(cache::ExpCacheNoLdiv{expOpType}, tmp::vType, v::vType, stage_index::Integer, positive, iip) where {expOpType<:AbstractSciMLOperator,vType} + op = get_op_for_this_step(cache, positive, stage_index) + v = matvec_prod_mip(tmp, op, v, iip) + return v +end + +@inline function _expmv_rkip_mip(cache::ExpCache{expOpType}, tmp::vType, v::vType, stage_index::Integer, positive, ::Val{true}) where {expOpType<:AbstractSciMLOperator,vType} + op = get_op_for_this_step(cache, stage_index) + if positive + matvec_prod_mip(tmp, op, v, Val(true)) + else + ldiv!(tmp, op, v) #ldiv can only be used for in place call + copyto!(v, tmp) + end + return v +end + +""" +Maybe-In Place "mip" +Safe function for matrix vector product for both mutable and immutable type. +For mutable type, overwrite `v` with `mat*v` and return `v`. Otherwise return `v` +Use the dispatch between ::Val{true} (in place) and ::Val{false} immutable to decide +""" +@inline function matvec_prod_mip(tmp::V, mat::M, v::V, ::Val{true}) where {V,M<:AbstractSciMLOperator} + mat(tmp, v, nothing, nothing, zero(eltype(v))) + copyto!(v, tmp) + return v +end + +@inline function matvec_prod_mip(_::V, mat::M, v::V, ::Val{false}) where {V,M<:AbstractSciMLOperator} + v = mat(v, nothing, nothing, zero(eltype(v))) + return v +end \ No newline at end of file diff --git a/lib/OrdinaryDiffEqRKIP/test/runtests.jl b/lib/OrdinaryDiffEqRKIP/test/runtests.jl new file mode 100644 index 0000000000..bb2b5013b1 --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/test/runtests.jl @@ -0,0 +1,5 @@ +using SafeTestsets + +@safetestset "Type Safety Tests" begin include("type_test.jl") end +@safetestset "Fourier Semilinear PDE Tests" begin include("semilinear_pde_test_cpu.jl") end +@safetestset "CUDA Test" begin include("semilinear_pde_test_cuda.jl") end \ No newline at end of file diff --git a/lib/OrdinaryDiffEqRKIP/test/semilinear_fft.jl b/lib/OrdinaryDiffEqRKIP/test/semilinear_fft.jl new file mode 100644 index 0000000000..b9bde1e747 --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/test/semilinear_fft.jl @@ -0,0 +1,141 @@ +import LinearAlgebra +using CUDA +using FFTW + +using OrdinaryDiffEqRKIP: RKIP +using SciMLBase: SplitODEProblem, SplitFunction, solve +using SciMLOperators: AbstractSciMLOperator + +struct DiagonalFourierOperator{T,uType<:AbstractVector{T},fftType,ifftType} <: AbstractSciMLOperator{T} + diag::uType + plan_fft!::fftType + plan_ifft!::ifftType + tmp::uType +end + +function DiagonalFourierOperator(diag, plan_fft!_prototype, plan_ifft!_prototype) + tmp = similar(diag) + return DiagonalFourierOperator(diag, plan_fft!_prototype(tmp), plan_ifft!_prototype(tmp), tmp) +end + +DiagonalFourierOperator(diag::Vector{T}) where {T<:Complex} = DiagonalFourierOperator(diag, FFTW.plan_fft!, FFTW.plan_ifft!) +DiagonalFourierOperator(diag::CuArray{T,1}) where {T<:Complex} = DiagonalFourierOperator(diag, CUDA.CUFFT.plan_fft!, CUDA.CUFFT.plan_ifft!) + +function LinearAlgebra.exp(D::DiagonalFourierOperator, t) + exp_kernel = similar(D.diag) + D.tmp .= D.diag + D.tmp .*= t + exp_kernel .= exp.(D.tmp) + DiagonalFourierOperator(exp_kernel, D.plan_fft!, D.plan_ifft!, D.tmp) +end + +@inline @fastmath function apply_kernel!(du, u, tmp, kernel, op_fft!, op_ifft!) + copyto!(tmp, u) + op_fft! * tmp + tmp .*= kernel + op_ifft! * tmp + copyto!(du, tmp) +end + +(D::DiagonalFourierOperator)(du, u, _, _, _) = apply_kernel!(du, u, D.tmp, D.diag, D.plan_fft!, D.plan_ifft!) + +gpu_init(u0::CuArray, ::Vector{T}) where {T<:Real} = u0 +gpu_init(u0::Vector, ::Vector{T}) where {T<:Real} = CuArray{Complex{T},1}(u0) +gpu_init(u0::Function, τ::Vector{T}) where {T<:Real} = CuArray{Complex{T},1}(u0.(τ)) + +cpu_init(u0::Vector, ::Vector{T}) where {T<:Real} = u0 +cpu_init(u0::Function, τ::Vector{T}) where {T<:Real} = Vector{Complex{T}}(u0.(τ)) + +init(u0, τ, ::Val{true}) = gpu_init(u0, τ) +init(u0, τ, ::Val{false}) = cpu_init(u0, τ) + + +function create_semilinear_problem( + nb_pts::Int, + dτ::T, + integration_limit::T, + u0, + semilinear_fourier_part::Function, + nonlinear_part!::Function, + use_gpu +) where {T<:Real} + + domain_size = dτ * nb_pts + τ = Vector{T}(range(-domain_size / 2.0, domain_size / 2.0, nb_pts)) + + u0_good_type = init(u0, τ, use_gpu) + + ω = 2π .* Vector{T}(FFTW.fftfreq(nb_pts, nb_pts / domain_size)) + diag = typeof(u0_good_type)(semilinear_fourier_part.(ω)) + + D̂ = DiagonalFourierOperator(diag) + N̂ = (du, u, _, t) -> nonlinear_part!(du, u, t) + + splfc = SplitFunction{true}(D̂, N̂) + return SplitODEProblem(splfc, u0_good_type, (0.0, integration_limit)) +end + +@inline @fastmath function nlse_non_linpart(res, u) + res .= u + res .= abs2.(res) + res .*= u + res .*= 1im +end + +@inline @fastmath function lle_non_linpart!(res, u, Δ, S) + nlse_non_linpart(res, u) + LinearAlgebra.axpby!(-1im * Δ, u, 1, res) + res .+= S +end + +create_nlse(nb_pts::Int, dτ::T, integration_limit::T, u0, use_gpu) where {T<:Real} = create_semilinear_problem(nb_pts, dτ, integration_limit, u0, ω -> -0.5im * ω^2, + (du, u, t) -> nlse_non_linpart(du, u), use_gpu) + +create_lle(nb_pts::Int, dτ::T, integration_limit::T, u0, use_gpu, Δ::T, S::T) where {T<:Real} = create_semilinear_problem(nb_pts, dτ, integration_limit, u0, ω -> -1 - 1im * ω^2, + (du, u, t) -> lle_non_linpart!(du, u, Δ, S), use_gpu) + +create_lle_scan(nb_pts::Int, dτ::T, scan_time::T, u0, use_gpu, Δ_min::T, Δ_max::T, S::T) where {T<:Real} = create_semilinear_problem(nb_pts, dτ, scan_time, u0, ω -> -1 - 1im * ω^2, + (du, u, t) -> lle_non_linpart!(du, u, Δ_min + (Δ_max - Δ_min) * t / scan_time, S), use_gpu) + +""" +Test function of the NLSE on a conservative soliton +""" +function nlse_test(::Type{T}, use_gpu; B=5.0, nb_pts=512, dτ=5e-3, propag_dist=10.0, save_everystep=false, reltol=1e-6) where {T<:Real} + B = T(B) + + problem = create_nlse( + nb_pts, + T(dτ), + T(propag_dist), + τ -> B * sech(B * τ), + use_gpu + ) + return solve(problem, RKIP(T(1e-4), T(1.0)); save_everystep, reltol=T(reltol)) +end + +nlse_test(type::Type{T}; kwargs...) where {T<:Real} = nlse_test(type, Val(false); kwargs...) +nlse_test(use_gpu::Val{true}; kwargs...) = nlse_test(Float64, use_gpu; kwargs...) +nlse_test(use_gpu::Val{false}; kwargs...) = nlse_test(Float64, use_gpu; kwargs...) +nlse_test(; kwargs...) = nlse_test(Float64; kwargs...) + +""" +Lugiato-Lefever equation test. Useful for benchmark due to its particular dynamic +""" +function lle_scan_test(::Type{T}, use_gpu; nb_pts=1024, dτ=5e-2, save_everystep=false, reltol=1e-6, S=3, Δ_min=-2.0, Δ_max=12.0, t_max=10) where {T<:Real} + problem = create_lle_scan( + nb_pts, + T(dτ), + T(t_max), + 1e-1 * exp.(2im * π * rand(T, nb_pts)), + use_gpu, + T(Δ_min), + T(Δ_max), + T(S) + ) + return solve(problem, RKIP(T(1e-4), T(1.0)); save_everystep, reltol=T(reltol)) +end + +lle_scan_test(type::Type{T}; kwargs...) where {T<:Real} = lle_scan_test(type, Val(false); kwargs...) +lle_scan_test(use_gpu::Val{true}; kwargs...) = lle_scan_test(Float64, use_gpu; kwargs...) +lle_scan_test(use_gpu::Val{false}; kwargs...) = lle_scan_test(Float64, use_gpu; kwargs...) +lle_scan_test(; kwargs...) = lle_scan_test(Float64; kwargs...) \ No newline at end of file diff --git a/lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cpu.jl b/lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cpu.jl new file mode 100644 index 0000000000..8cbd7efdbe --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cpu.jl @@ -0,0 +1,34 @@ +using Test +using LinearAlgebra: norm + +include("semilinear_fft.jl") + + +@testset "Lugiato-Lefever equation test - Float64" begin + @test !isnothing(lle_scan_test()) +end + +@testset "Lugiato-Lefever equation test - Float32" begin + @test !isnothing(lle_scan_test(Float32)) +end + + +@testset "NLSE test - Float64" begin + sol = nlse_test() + + # NLSE Soliton solution, only the phase change + p1 = abs2.(sol.u[1]) + p2 = abs2.(sol.u[end]) + + @test norm(p1 .- p2) / norm(p1) < 1e-2 # small difference with analytical solution due to boundary periodic condition +end + +@testset "NLSE test - Float32" begin + sol = nlse_test(Float32) + + # NLSE Soliton solution, only the phase change + p1 = abs2.(sol.u[1]) + p2 = abs2.(sol.u[end]) + + @test norm(p1 .- p2) / norm(p1) < 1e-2 # small difference with analytical solution due to boundary periodic condition +end \ No newline at end of file diff --git a/lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cuda.jl b/lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cuda.jl new file mode 100644 index 0000000000..dde5dbc3de --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cuda.jl @@ -0,0 +1,35 @@ +using Test +using LinearAlgebra: norm +using CUDA + +include("semilinear_fft.jl") + +CUDA.allowscalar(false) # ensure no scalar indexing is used (for performance) + +@testset "Lugiato-Lefever equation test on GPU with CUDA - Float64" begin + @test !isnothing(lle_scan_test(Val(true))) +end + +@testset "Lugiato-Lefever equation test on GPU with CUDA - Float32" begin + @test !isnothing(lle_scan_test(Float32, Val(true))) +end + +@testset "NLSE test on GPU with CUDA - Float64" begin + sol = nlse_test(Val(true)) + + # NLSE Soliton solution, only the phase change + p1 = abs2.(sol.u[1]) + p2 = abs2.(sol.u[end]) + + @test norm(p1 .- p2) / norm(p1) < 1.5e-2 # small difference with analytical sech solution due to boundary periodic condition and PDE discretization +end + +@testset "NLSE test on GPU with CUDA - Float32" begin + sol = nlse_test(Float32, Val(true)) + + # NLSE Soliton solution, only the phase change + p1 = abs2.(sol.u[1]) + p2 = abs2.(sol.u[end]) + + @test norm(p1 .- p2) / norm(p1) < 1.5e-2 # small difference with analytical sech solution due to boundary periodic condition and PDE discretization +end \ No newline at end of file diff --git a/lib/OrdinaryDiffEqRKIP/test/type_test.jl b/lib/OrdinaryDiffEqRKIP/test/type_test.jl new file mode 100644 index 0000000000..48b7e36b99 --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/test/type_test.jl @@ -0,0 +1,84 @@ +using Test +using OrdinaryDiffEqRKIP: RKIP +using SciMLBase: SplitODEProblem, SplitFunction, solve +using SciMLOperators: ScalarOperator, DiagonalOperator, MatrixOperator, AbstractSciMLOperator + +using StaticArrays + +import LinearAlgebra + +LinearAlgebra.exp(d::ScalarOperator, t) = ScalarOperator(exp(t * d.val)) # Temporary Fix for missing exp dispatch +LinearAlgebra.exp(d::MatrixOperator, t) = MatrixOperator(exp(t * d.A)) + +function test(A_prototype, u_prototype, iip; use_ldiv=false) + for reltol in [1e-3, 1e-6, 1e-8] # fail for 1e-9, probably due to floating point error + μ = 1.05 + analytic = (u0, _, t) -> u0 .* exp(2μ * t) + + for μₐ in [2μ, μ, 0.0] + + Â = A_prototype(μₐ) + f = (du, u, _, _) -> du .= (2μ - μₐ) .* u + + if !iip + f = (u, _, _) -> (2μ - μₐ) .* u + end + + u0 = u_prototype(1.0) + t = (0.0, 1.0) + + splfc = SplitFunction{iip}(Â, f; analytic=analytic) + spltode = SplitODEProblem(splfc, u0, t) + sol = solve(spltode, RKIP(; use_ldiv=use_ldiv); reltol=reltol) + @test isapprox(sol(t[end]), splfc.analytic(u0, nothing, t[end]); rtol=reltol) + end + end +end + +@testset "In-Place ScalarOperator Vector Adaptative Ldiv" begin + test(μ -> ScalarOperator(μ), u0 -> [u0], true; use_ldiv=true) +end + +@testset "In-Place DiagonalOperator Vector Adaptative Ldiv" begin + test(μ -> DiagonalOperator([μ]), u0 -> [u0], true; use_ldiv=true) +end + + +@testset "In-Place ScalarOperator Vector Adaptative" begin + test(μ -> ScalarOperator(μ), u0 -> [u0], true) +end + +@testset "In-Place DiagonalOperator Vector Adaptative" begin + test(μ -> DiagonalOperator([μ]), u0 -> [u0], true) +end + +@testset "In-Place MatrixOperator Vector Float Adaptative" begin + test(μ -> MatrixOperator([μ;;]), u0 -> [u0], true) +end + +@testset "Out-of-place Scalar Operator Float Adaptative" begin + test(μ -> ScalarOperator(μ), u0 -> u0, false) +end + +@testset "Out-of-place Scalar Operator SVector Adaptative" begin + test(μ -> ScalarOperator(μ), u0 -> SVector(u0), false) +end + +# Failing Test due to missing dispatch in SciMLOperator + + +# fails as calling a MatrixOperator on a StaticVector change its type to Vector, causing a type instability + +# @testset "Out-of-place Diagonal 1x1 Operator SVector Adaptative" begin +# test((μ) -> DiagonalOperator([μ]), u0 -> SVector(u0), false) +# end + +# @testset "Out-of-place Matrix 1x1 Operator SVector Adaptative" begin +# test((μ) -> MatrixOperator([μ;;]), u0 -> SVector(u0), false) +# end + +# # Fails for a strange reason as calling ldiv!(a, B, c) on a MatrixOperator dispatch toward an inexisting method + +# @testset "In-Place MatrixOperator Vector Float Adaptative Ldiv" begin +# test(μ -> MatrixOperator([μ;;]), u0 -> [u0], true; use_ldiv=true) +# end diff --git a/src/OrdinaryDiffEq.jl b/src/OrdinaryDiffEq.jl index e689440fbf..b7c2b057a0 100644 --- a/src/OrdinaryDiffEq.jl +++ b/src/OrdinaryDiffEq.jl @@ -208,6 +208,10 @@ export LawsonEuler, NorsettEuler, ETD1, ETDRK2, ETDRK3, ETDRK4, HochOst4, Exp4, EPIRK4s3B, EPIRK5s3, EXPRB53s3, EPIRK5P1, EPIRK5P2, ETD2, Exprb32, Exprb43 +using OrdinaryDiffEqRKIP +export RKIP + + import PrecompileTools import Preferences import DocStringExtensions From 60016def754d294cc4bef4b1f9b0019799995d2e Mon Sep 17 00:00:00 2001 From: Azercoco <20425137+Azercoco@users.noreply.github.com> Date: Tue, 10 Jun 2025 16:10:18 +0200 Subject: [PATCH 02/10] Reformat Style --- .../src/OrdinaryDiffEqRKIP.jl | 11 ++- lib/OrdinaryDiffEqRKIP/src/rkip_cache.jl | 61 +++++++++---- .../src/rkip_perform_step.jl | 49 ++++++---- lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl | 25 ++++-- lib/OrdinaryDiffEqRKIP/test/runtests.jl | 6 +- lib/OrdinaryDiffEqRKIP/test/semilinear_fft.jl | 89 ++++++++++++------- .../test/semilinear_pde_test_cpu.jl | 4 +- lib/OrdinaryDiffEqRKIP/test/type_test.jl | 22 +++-- 8 files changed, 169 insertions(+), 98 deletions(-) diff --git a/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl b/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl index 36aa49f87e..51257cb75d 100644 --- a/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl +++ b/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl @@ -9,9 +9,12 @@ using DiffEqBase: ExplicitRKTableau using DiffEqDevTools: constructDormandPrince6 import OrdinaryDiffEqCore: OrdinaryDiffEqAdaptiveExponentialAlgorithm, alg_adaptive_order, - alg_order, alg_cache, @cache, SplitFunction, get_fsalfirstlast, initialize!, perform_step!, - has_dtnew_modification, calculate_residuals, calculate_residuals!, increment_nf!, OrdinaryDiffEqAdaptiveAlgorithm, OrdinaryDiffEqMutableCache, dtnew_modification - + alg_order, alg_cache, @cache, SplitFunction, get_fsalfirstlast, + initialize!, perform_step!, + has_dtnew_modification, calculate_residuals, + calculate_residuals!, increment_nf!, + OrdinaryDiffEqAdaptiveAlgorithm, OrdinaryDiffEqMutableCache, + dtnew_modification include("rkip_cache.jl") include("algorithms.jl") @@ -20,4 +23,4 @@ include("rkip_perform_step.jl") export RKIP -end \ No newline at end of file +end diff --git a/lib/OrdinaryDiffEqRKIP/src/rkip_cache.jl b/lib/OrdinaryDiffEqRKIP/src/rkip_cache.jl index 6d3353160e..5b64330b19 100644 --- a/lib/OrdinaryDiffEqRKIP/src/rkip_cache.jl +++ b/lib/OrdinaryDiffEqRKIP/src/rkip_cache.jl @@ -1,7 +1,7 @@ -abstract type AbstractExpCache{expOpType<:AbstractSciMLOperator} end +abstract type AbstractExpCache{expOpType <: AbstractSciMLOperator} end struct ExpCache{expOpType} <: AbstractExpCache{expOpType} - expÂ_cached::Array{expOpType,2} + expÂ_cached::Array{expOpType, 2} expÂ_for_this_step::Vector{expOpType} end struct ExpCacheNoLdiv{expOpType} <: AbstractExpCache{expOpType} @@ -9,11 +9,19 @@ struct ExpCacheNoLdiv{expOpType} <: AbstractExpCache{expOpType} exp_cache_inv::ExpCache{expOpType} end -get_op_for_this_step(cache::ExpCache{expOpType}, index::Int) where {expOpType} = cache.expÂ_for_this_step[index] -get_op_for_this_step(cache_no_ldiv::ExpCacheNoLdiv{expOpType}, positive::Bool, index::Int) where {expOpType} = positive ? cache_no_ldiv.exp_cache.expÂ_for_this_step[index] : cache_no_ldiv.exp_cache_inv.expÂ_for_this_step[index] - +function get_op_for_this_step(cache::ExpCache{expOpType}, index::Int) where {expOpType} + cache.expÂ_for_this_step[index] +end +function get_op_for_this_step(cache_no_ldiv::ExpCacheNoLdiv{expOpType}, + positive::Bool, index::Int) where {expOpType} + positive ? cache_no_ldiv.exp_cache.expÂ_for_this_step[index] : + cache_no_ldiv.exp_cache_inv.expÂ_for_this_step[index] +end -mutable struct RKIPCache{expOpType<:AbstractSciMLOperator,cacheType<:AbstractExpCache{expOpType},tType<:Number,opType<:AbstractSciMLOperator,uType,iip} <: OrdinaryDiffEqMutableCache +mutable struct RKIPCache{ + expOpType <: AbstractSciMLOperator, cacheType <: AbstractExpCache{expOpType}, + tType <: Number, opType <: AbstractSciMLOperator, uType, iip} <: + OrdinaryDiffEqMutableCache exp_cache::cacheType last_step::tType cached::Vector{Bool} @@ -26,38 +34,59 @@ end get_fsalfirstlast(cache::RKIPCache, u) = (zero(cache.tmp), zero(cache.tmp)) - -@inline function cache_exp!(cache::ExpCache{expOpType}, A::opType, h::T, action::Symbol, step_index::Int, unique_stage_index::Int) where {expOpType<:AbstractSciMLOperator,opType<:AbstractSciMLOperator,T<:Number} +@inline function cache_exp!(cache::ExpCache{expOpType}, + A::opType, + h::T, + action::Symbol, + step_index::Int, + unique_stage_index::Int) where { + expOpType <: AbstractSciMLOperator, opType <: AbstractSciMLOperator, T <: Number} @unpack expÂ_for_this_step, expÂ_cached = cache - expÂ_for_this_step[unique_stage_index] = (action == :use_cached) ? expÂ_cached[step_index, unique_stage_index] : exp(A, h) # fetching or generating exp(Â*c_i*dt) + expÂ_for_this_step[unique_stage_index] = (action == :use_cached) ? + expÂ_cached[step_index, unique_stage_index] : + exp(A, h) # fetching or generating exp(Â*c_i*dt) if action == :cache expÂ_cached[step_index, unique_stage_index] = expÂ_for_this_step[unique_stage_index] # storing exp(Â*c_i*dt) end end -@inline function cache_exp!(cache::ExpCacheNoLdiv{expOpType}, Â::opType, h::T, action::Symbol, step_index::Int, unique_stage_index::Int) where {expOpType<:AbstractSciMLOperator,opType<:AbstractSciMLOperator,T<:Number} +@inline function cache_exp!(cache::ExpCacheNoLdiv{expOpType}, + Â::opType, + h::T, + action::Symbol, + step_index::Int, + unique_stage_index::Int) where { + expOpType <: AbstractSciMLOperator, opType <: AbstractSciMLOperator, T <: Number} cache_exp!(cache.exp_cache, Â, h, action, step_index, unique_stage_index) cache_exp!(cache.exp_cache_inv, Â, -h, action, step_index, unique_stage_index) end """ - Prepare/generate all the needed exp(± A * dt * c[i]) for a step size dt + Prepare/generate all the needed exp(± A * dt * c[i]) for a step size dt """ -@inline function cache_exp_op_for_this_step!(cache::RKIPCache{expOpType,cacheType,tType,opType,uType,iip}, Â::opType, dt::tType, alg::algType) where {expOpType,cacheType,tType,opType,uType,algType,iip} +@inline function cache_exp_op_for_this_step!( + cache::RKIPCache{expOpType, cacheType, tType, opType, uType, iip}, + Â::opType, dt::tType, + alg::algType) where {expOpType, cacheType, tType, opType, uType, algType, iip} @unpack dt_for_expÂ_caching = alg if !iszero(dt) && !(dt ≈ cache.last_step) # we check that new exp(A dt) are needed dt_abs = abs(dt) # only the positive dt are used for indexing + action = :single_use # exp(A*dt) is only computed for this step - step_index = clamp(searchsortedlast(dt_for_expÂ_caching, dt_abs), 1, lastindex(dt_for_expÂ_caching)) # fetching the index corresponding to the step size - action = (dt_for_expÂ_caching[step_index] ≈ dt_abs) ? (cache.cached[step_index] ? :use_cached : :cache) : :single_use # if alreay present, we reuse the cached, otherwise it is generated + step_index = clamp(searchsortedlast(dt_for_expÂ_caching, dt_abs), + 1, lastindex(dt_for_expÂ_caching)) # fetching the index corresponding to the step size + + if dt_for_expÂ_caching[step_index] ≈ dt_abs # if dt corresponds to a cahing step + action = (cache.cached[step_index] ? :use_cached : :cache) # if alreay present, we reuse the cached, otherwise it is generated + end for (unique_stage_index, c) in enumerate(cache.c_unique) # iterating over all unique c_i of the RK tableau - cache_exp!(cache.exp_cache, Â, abs(dt * c), action, step_index, unique_stage_index) # generating and caching + cache_exp!( + cache.exp_cache, Â, abs(dt * c), action, step_index, unique_stage_index) # generating and caching end cache.cached[step_index] |= (action == :cache) # set the flag that we have already cached exp(Â*c_i*dt) for this dt end cache.last_step = dt end - diff --git a/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl b/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl index 575653c3f1..6fcee2a47e 100644 --- a/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl +++ b/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl @@ -1,21 +1,24 @@ """ -Helper function for the (nonlinear) part of the problem maybe in place. +Helper function for the (nonlinear) part of the problem maybe in place. Overwrite compute f(u, p, t) and if possible, overwrite u with the result. Same principle of operation as `matvec_prod_mip` for mutability/in-place handling. """ -function nl_part_mip(tmp::uType, f!, u::uType, p, t::tType, ::Val{true}) where {tType,uType} +function nl_part_mip( + tmp::uType, f!, u::uType, p, t::tType, ::Val{true}) where {tType, uType} f!(tmp, u, p, t) copyto!(u, tmp) return u end -nl_part_mip(_::uType, f, u::uType, p, t::tType, ::Val{false}) where {tType,uType} = f(u, p, t) - +function nl_part_mip(_::uType, f, u::uType, p, t::tType, ::Val{false}) where {tType, uType} + f(u, p, t) +end """ Helper function to compute Au + f(u, p, t) and if in place, store the result in res. Retunr res if in place, otherwise return Au + f(u, p, t) """ -function f_mip!(res::uType, tmp::uType, A::opType, f, u::uType, p, t::tType, ::Val{true}) where {tType,uType,opType} +function f_mip!(res::uType, tmp::uType, A::opType, f, u::uType, p, + t::tType, ::Val{true}) where {tType, uType, opType} res .= u res = matvec_prod_mip(tmp, A, res, Val(true)) f(tmp, u, p, t) @@ -23,19 +26,28 @@ function f_mip!(res::uType, tmp::uType, A::opType, f, u::uType, p, t::tType, ::V return res end -f_mip!(_::uType, tmp::uType, A::opType, f, u::uType, p, t::tType, ::Val{false}) where {tType,uType,opType} = matvec_prod_mip(tmp, A, u, Val(false)) + f(u, p, t) +function f_mip!(_::uType, tmp::uType, A::opType, f, u::uType, p, + t::tType, ::Val{false}) where {tType, uType, opType} + matvec_prod_mip(tmp, A, u, Val(false)) + f(u, p, t) +end """ -Helper function for the residual maybe in place. +Helper function for the residual maybe in place. Same principle of operation as `_safe_matvec_prod` for mutability/in-place handling. """ -function calculate_residuals_mip(tmp::uType, utilde::uType, uprev::uType, u::uType, abstol, reltol, internalnorm, t, ::Val{true}) where {uType} +function calculate_residuals_mip(tmp::uType, utilde::uType, uprev::uType, u::uType, abstol, + reltol, internalnorm, t, ::Val{true}) where {uType} calculate_residuals!(tmp, utilde, uprev, u, abstol, reltol, internalnorm, t) return tmp end -calculate_residuals_mip(_::uType, utilde::uType, uprev::uType, u::uType, abstol, reltol, internalnorm, t, ::Val{false}) where {uType} = calculate_residuals(utilde, uprev, u, abstol, reltol, internalnorm, t) +function calculate_residuals_mip(_::uType, utilde::uType, uprev::uType, u::uType, abstol, + reltol, internalnorm, t, ::Val{false}) where {uType} + calculate_residuals(utilde, uprev, u, abstol, reltol, internalnorm, t) +end -@fastmath function perform_step!(integrator, cache::RKIPCache{expOpType,cacheType,tType,opType,uType,iip}) where {expOpType,cacheType,tType,opType,uType,iip} +@fastmath function perform_step!(integrator, + cache::RKIPCache{expOpType, cacheType, tType, opType, uType, iip}) where { + expOpType, cacheType, tType, opType, uType, iip} @unpack t, dt, uprev, u, f, p, fsalfirst, fsallast, alg = integrator @unpack c, α, αEEst, stages, A = alg.tableau @unpack kk, utilde, tmp = cache @@ -47,10 +59,10 @@ calculate_residuals_mip(_::uType, utilde::uType, uprev::uType, u::uType, abstol, @bb u .= uprev - @inbounds for i in 1:(stages) + for i in 1:(stages) @bb kk[i] .= uprev - @inbounds for j in 1:(i-1) + for j in 1:(i - 1) g = A[i, j] * dt kk[i] = axpy_mip(g, kk[j], kk[i], iip) # kk_i += dt*A[i, j] * kk_j end @@ -60,7 +72,6 @@ calculate_residuals_mip(_::uType, utilde::uType, uprev::uType, u::uType, abstol, kk[i] = nl_part_mip(cache.tmp, f.f2, kk[i], p, t + c[i] * dt, iip) # kᵢ = f(u + Σⱼ dt*(Aᵢⱼ kⱼ), t + dt*cᵢ) - kk[i] = expmv_rkip_mip(cache, kk[i], -dt, i) # kᵢ = exp(-Â * [dt * cᵢ])*kᵢ ➡ Going back in interaction picture integrator.stats.nf += 2 # two exp vec product @@ -81,7 +92,8 @@ calculate_residuals_mip(_::uType, utilde::uType, uprev::uType, u::uType, abstol, if adaptive utilde = expmv_rkip_mip(cache, utilde, dt) # stepping forward into the interaction ũ = exp(Â dt)*ũ - tmp = calculate_residuals_mip(tmp, utilde, uprev, u, abstol, reltol, internalnorm, t, iip) # error computation maybe in place + tmp = calculate_residuals_mip( + tmp, utilde, uprev, u, abstol, reltol, internalnorm, t, iip) # error computation maybe in place integrator.EEst = internalnorm(tmp, t) end @@ -94,12 +106,13 @@ calculate_residuals_mip(_::uType, utilde::uType, uprev::uType, u::uType, abstol, @bb copyto!(integrator.k[2], fsallast) end - -function initialize!(integrator, cache::RKIPCache{expOpType,cacheType,tType,opType,uType,iip}) where {expOpType,cacheType,tType,opType,uType,iip} +function initialize!(integrator, + cache::RKIPCache{expOpType, cacheType, tType, opType, uType, iip}) where { + expOpType, cacheType, tType, opType, uType, iip} @unpack f, u, p, t, fsalfirst, fsallast = integrator kshortsize = 2 - k = [zero(u) for _ ∈ 1:kshortsize] + k = [zero(u) for _ in 1:kshortsize] fsalfirst = f_mip!(fsalfirst, cache.tmp, f.f1.f, f.f2, u, p, t, iip) # first derivative for interpolation computation, maybe in place integrator.stats.nf += 1 @@ -109,4 +122,4 @@ function initialize!(integrator, cache::RKIPCache{expOpType,cacheType,tType,opTy @bb copyto!(integrator.k[1], fsalfirst) @bb copyto!(integrator.k[2], fsallast) -end \ No newline at end of file +end diff --git a/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl b/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl index 14fba179bc..ae2078b8d7 100644 --- a/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl +++ b/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl @@ -2,16 +2,17 @@ @inline axpy_mip(α, x::uType, y::uType, ::Val{true}) where {uType} = axpy!(α, x, y) #BLAS call @inline axpy_mip(α, x::uType, y::uType, ::Val{false}) where {uType} = y .+ α .* x - """ Maybe-In Place "mip" Safe function for computing the matrix exponential vector product for both mutable and immutable type. Same principle of operation as `_safe_matvec_prod` for mutability/in-place handling. """ -@inline expmv_rkip_mip(cache::RKIPCache{expOpType,cacheType,tType,opType,uType,iip}, v::uType, h::tType) where {expOpType,cacheType,tType,opType,uType,iip} = expmv_rkip_mip(cache, v, h, lastindex(cache.c_mapping)) # If i is not precised, this is the final step in the RKIP -> h = dt - -@inline function expmv_rkip_mip(cache::RKIPCache{expOpType,cacheType,tType,opType,uType,iip}, v::uType, h::tType, stage_index::Int) where {expOpType,cacheType,tType,opType,uType,iip} +@inline expmv_rkip_mip(cache::RKIPCache{expOpType, cacheType, tType, opType, uType, iip}, v::uType, h::tType) where {expOpType, cacheType, tType, opType, uType, iip} = expmv_rkip_mip( + cache, v, h, lastindex(cache.c_mapping)) # If i is not precised, this is the final step in the RKIP -> h = dt +@inline function expmv_rkip_mip( + cache::RKIPCache{expOpType, cacheType, tType, opType, uType, iip}, v::uType, + h::tType, stage_index::Int) where {expOpType, cacheType, tType, opType, uType, iip} if !(h ≈ tType(0.0)) c = cache.c_mapping[stage_index] exp_cache = cache.exp_cache @@ -21,13 +22,17 @@ Same principle of operation as `_safe_matvec_prod` for mutability/in-place handl return v end -@inline function _expmv_rkip_mip(cache::ExpCacheNoLdiv{expOpType}, tmp::vType, v::vType, stage_index::Integer, positive, iip) where {expOpType<:AbstractSciMLOperator,vType} +@inline function _expmv_rkip_mip( + cache::ExpCacheNoLdiv{expOpType}, tmp::vType, v::vType, stage_index::Integer, + positive, iip) where {expOpType <: AbstractSciMLOperator, vType} op = get_op_for_this_step(cache, positive, stage_index) v = matvec_prod_mip(tmp, op, v, iip) return v end -@inline function _expmv_rkip_mip(cache::ExpCache{expOpType}, tmp::vType, v::vType, stage_index::Integer, positive, ::Val{true}) where {expOpType<:AbstractSciMLOperator,vType} +@inline function _expmv_rkip_mip( + cache::ExpCache{expOpType}, tmp::vType, v::vType, stage_index::Integer, + positive, ::Val{true}) where {expOpType <: AbstractSciMLOperator, vType} op = get_op_for_this_step(cache, stage_index) if positive matvec_prod_mip(tmp, op, v, Val(true)) @@ -44,13 +49,15 @@ Safe function for matrix vector product for both mutable and immutable type. For mutable type, overwrite `v` with `mat*v` and return `v`. Otherwise return `v` Use the dispatch between ::Val{true} (in place) and ::Val{false} immutable to decide """ -@inline function matvec_prod_mip(tmp::V, mat::M, v::V, ::Val{true}) where {V,M<:AbstractSciMLOperator} +@inline function matvec_prod_mip( + tmp::V, mat::M, v::V, ::Val{true}) where {V, M <: AbstractSciMLOperator} mat(tmp, v, nothing, nothing, zero(eltype(v))) copyto!(v, tmp) return v end -@inline function matvec_prod_mip(_::V, mat::M, v::V, ::Val{false}) where {V,M<:AbstractSciMLOperator} +@inline function matvec_prod_mip( + _::V, mat::M, v::V, ::Val{false}) where {V, M <: AbstractSciMLOperator} v = mat(v, nothing, nothing, zero(eltype(v))) return v -end \ No newline at end of file +end diff --git a/lib/OrdinaryDiffEqRKIP/test/runtests.jl b/lib/OrdinaryDiffEqRKIP/test/runtests.jl index bb2b5013b1..497d0316ec 100644 --- a/lib/OrdinaryDiffEqRKIP/test/runtests.jl +++ b/lib/OrdinaryDiffEqRKIP/test/runtests.jl @@ -1,5 +1,5 @@ using SafeTestsets -@safetestset "Type Safety Tests" begin include("type_test.jl") end -@safetestset "Fourier Semilinear PDE Tests" begin include("semilinear_pde_test_cpu.jl") end -@safetestset "CUDA Test" begin include("semilinear_pde_test_cuda.jl") end \ No newline at end of file +@safetestset "Type Safety Tests" include("type_test.jl") +@safetestset "Fourier Semilinear PDE Tests" include("semilinear_pde_test_cpu.jl") +@safetestset "CUDA Test" include("semilinear_pde_test_cuda.jl") diff --git a/lib/OrdinaryDiffEqRKIP/test/semilinear_fft.jl b/lib/OrdinaryDiffEqRKIP/test/semilinear_fft.jl index b9bde1e747..23acb406ca 100644 --- a/lib/OrdinaryDiffEqRKIP/test/semilinear_fft.jl +++ b/lib/OrdinaryDiffEqRKIP/test/semilinear_fft.jl @@ -6,7 +6,8 @@ using OrdinaryDiffEqRKIP: RKIP using SciMLBase: SplitODEProblem, SplitFunction, solve using SciMLOperators: AbstractSciMLOperator -struct DiagonalFourierOperator{T,uType<:AbstractVector{T},fftType,ifftType} <: AbstractSciMLOperator{T} +struct DiagonalFourierOperator{T, uType <: AbstractVector{T}, fftType, ifftType} <: + AbstractSciMLOperator{T} diag::uType plan_fft!::fftType plan_ifft!::ifftType @@ -15,11 +16,16 @@ end function DiagonalFourierOperator(diag, plan_fft!_prototype, plan_ifft!_prototype) tmp = similar(diag) - return DiagonalFourierOperator(diag, plan_fft!_prototype(tmp), plan_ifft!_prototype(tmp), tmp) + return DiagonalFourierOperator( + diag, plan_fft!_prototype(tmp), plan_ifft!_prototype(tmp), tmp) end -DiagonalFourierOperator(diag::Vector{T}) where {T<:Complex} = DiagonalFourierOperator(diag, FFTW.plan_fft!, FFTW.plan_ifft!) -DiagonalFourierOperator(diag::CuArray{T,1}) where {T<:Complex} = DiagonalFourierOperator(diag, CUDA.CUFFT.plan_fft!, CUDA.CUFFT.plan_ifft!) +function DiagonalFourierOperator(diag::Vector{T}) where {T <: Complex} + DiagonalFourierOperator(diag, FFTW.plan_fft!, FFTW.plan_ifft!) +end +function DiagonalFourierOperator(diag::CuArray{T, 1}) where {T <: Complex} + DiagonalFourierOperator(diag, CUDA.CUFFT.plan_fft!, CUDA.CUFFT.plan_ifft!) +end function LinearAlgebra.exp(D::DiagonalFourierOperator, t) exp_kernel = similar(D.diag) @@ -37,29 +43,29 @@ end copyto!(du, tmp) end -(D::DiagonalFourierOperator)(du, u, _, _, _) = apply_kernel!(du, u, D.tmp, D.diag, D.plan_fft!, D.plan_ifft!) +function (D::DiagonalFourierOperator)(du, u, _, _, _) + apply_kernel!(du, u, D.tmp, D.diag, D.plan_fft!, D.plan_ifft!) +end -gpu_init(u0::CuArray, ::Vector{T}) where {T<:Real} = u0 -gpu_init(u0::Vector, ::Vector{T}) where {T<:Real} = CuArray{Complex{T},1}(u0) -gpu_init(u0::Function, τ::Vector{T}) where {T<:Real} = CuArray{Complex{T},1}(u0.(τ)) +gpu_init(u0::CuArray, ::Vector{T}) where {T <: Real} = u0 +gpu_init(u0::Vector, ::Vector{T}) where {T <: Real} = CuArray{Complex{T}, 1}(u0) +gpu_init(u0::Function, τ::Vector{T}) where {T <: Real} = CuArray{Complex{T}, 1}(u0.(τ)) -cpu_init(u0::Vector, ::Vector{T}) where {T<:Real} = u0 -cpu_init(u0::Function, τ::Vector{T}) where {T<:Real} = Vector{Complex{T}}(u0.(τ)) +cpu_init(u0::Vector, ::Vector{T}) where {T <: Real} = u0 +cpu_init(u0::Function, τ::Vector{T}) where {T <: Real} = Vector{Complex{T}}(u0.(τ)) init(u0, τ, ::Val{true}) = gpu_init(u0, τ) init(u0, τ, ::Val{false}) = cpu_init(u0, τ) - function create_semilinear_problem( - nb_pts::Int, - dτ::T, - integration_limit::T, - u0, - semilinear_fourier_part::Function, - nonlinear_part!::Function, - use_gpu -) where {T<:Real} - + nb_pts::Int, + dτ::T, + integration_limit::T, + u0, + semilinear_fourier_part::Function, + nonlinear_part!::Function, + use_gpu +) where {T <: Real} domain_size = dτ * nb_pts τ = Vector{T}(range(-domain_size / 2.0, domain_size / 2.0, nb_pts)) @@ -88,19 +94,30 @@ end res .+= S end -create_nlse(nb_pts::Int, dτ::T, integration_limit::T, u0, use_gpu) where {T<:Real} = create_semilinear_problem(nb_pts, dτ, integration_limit, u0, ω -> -0.5im * ω^2, - (du, u, t) -> nlse_non_linpart(du, u), use_gpu) +function create_nlse( + nb_pts::Int, dτ::T, integration_limit::T, u0, use_gpu) where {T <: Real} + create_semilinear_problem(nb_pts, dτ, integration_limit, u0, ω -> -0.5im * ω^2, + (du, u, t) -> nlse_non_linpart(du, u), use_gpu) +end -create_lle(nb_pts::Int, dτ::T, integration_limit::T, u0, use_gpu, Δ::T, S::T) where {T<:Real} = create_semilinear_problem(nb_pts, dτ, integration_limit, u0, ω -> -1 - 1im * ω^2, - (du, u, t) -> lle_non_linpart!(du, u, Δ, S), use_gpu) +function create_lle( + nb_pts::Int, dτ::T, integration_limit::T, u0, use_gpu, Δ::T, S::T) where {T <: Real} + create_semilinear_problem(nb_pts, dτ, integration_limit, u0, ω -> -1 - 1im * ω^2, + (du, u, t) -> lle_non_linpart!(du, u, Δ, S), use_gpu) +end -create_lle_scan(nb_pts::Int, dτ::T, scan_time::T, u0, use_gpu, Δ_min::T, Δ_max::T, S::T) where {T<:Real} = create_semilinear_problem(nb_pts, dτ, scan_time, u0, ω -> -1 - 1im * ω^2, - (du, u, t) -> lle_non_linpart!(du, u, Δ_min + (Δ_max - Δ_min) * t / scan_time, S), use_gpu) +function create_lle_scan(nb_pts::Int, dτ::T, scan_time::T, u0, use_gpu, + Δ_min::T, Δ_max::T, S::T) where {T <: Real} + create_semilinear_problem(nb_pts, dτ, scan_time, u0, ω -> -1 - 1im * ω^2, + (du, u, t) -> lle_non_linpart!(du, u, Δ_min + (Δ_max - Δ_min) * t / scan_time, S), use_gpu) +end """ Test function of the NLSE on a conservative soliton """ -function nlse_test(::Type{T}, use_gpu; B=5.0, nb_pts=512, dτ=5e-3, propag_dist=10.0, save_everystep=false, reltol=1e-6) where {T<:Real} +function nlse_test( + ::Type{T}, use_gpu; B = 5.0, nb_pts = 512, dτ = 5e-3, propag_dist = 10.0, + save_everystep = false, reltol = 1e-6) where {T <: Real} B = T(B) problem = create_nlse( @@ -110,10 +127,12 @@ function nlse_test(::Type{T}, use_gpu; B=5.0, nb_pts=512, dτ=5e-3, propag_dist= τ -> B * sech(B * τ), use_gpu ) - return solve(problem, RKIP(T(1e-4), T(1.0)); save_everystep, reltol=T(reltol)) + return solve(problem, RKIP(T(1e-4), T(1.0)); save_everystep, reltol = T(reltol)) end -nlse_test(type::Type{T}; kwargs...) where {T<:Real} = nlse_test(type, Val(false); kwargs...) +function nlse_test(type::Type{T}; kwargs...) where {T <: Real} + nlse_test(type, Val(false); kwargs...) +end nlse_test(use_gpu::Val{true}; kwargs...) = nlse_test(Float64, use_gpu; kwargs...) nlse_test(use_gpu::Val{false}; kwargs...) = nlse_test(Float64, use_gpu; kwargs...) nlse_test(; kwargs...) = nlse_test(Float64; kwargs...) @@ -121,7 +140,9 @@ nlse_test(; kwargs...) = nlse_test(Float64; kwargs...) """ Lugiato-Lefever equation test. Useful for benchmark due to its particular dynamic """ -function lle_scan_test(::Type{T}, use_gpu; nb_pts=1024, dτ=5e-2, save_everystep=false, reltol=1e-6, S=3, Δ_min=-2.0, Δ_max=12.0, t_max=10) where {T<:Real} +function lle_scan_test( + ::Type{T}, use_gpu; nb_pts = 1024, dτ = 5e-2, save_everystep = false, + reltol = 1e-6, S = 3, Δ_min = -2.0, Δ_max = 12.0, t_max = 10) where {T <: Real} problem = create_lle_scan( nb_pts, T(dτ), @@ -132,10 +153,12 @@ function lle_scan_test(::Type{T}, use_gpu; nb_pts=1024, dτ=5e-2, save_everystep T(Δ_max), T(S) ) - return solve(problem, RKIP(T(1e-4), T(1.0)); save_everystep, reltol=T(reltol)) + return solve(problem, RKIP(T(1e-4), T(1.0)); save_everystep, reltol = T(reltol)) end -lle_scan_test(type::Type{T}; kwargs...) where {T<:Real} = lle_scan_test(type, Val(false); kwargs...) +function lle_scan_test(type::Type{T}; kwargs...) where {T <: Real} + lle_scan_test(type, Val(false); kwargs...) +end lle_scan_test(use_gpu::Val{true}; kwargs...) = lle_scan_test(Float64, use_gpu; kwargs...) lle_scan_test(use_gpu::Val{false}; kwargs...) = lle_scan_test(Float64, use_gpu; kwargs...) -lle_scan_test(; kwargs...) = lle_scan_test(Float64; kwargs...) \ No newline at end of file +lle_scan_test(; kwargs...) = lle_scan_test(Float64; kwargs...) diff --git a/lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cpu.jl b/lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cpu.jl index 8cbd7efdbe..c322c2e1d1 100644 --- a/lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cpu.jl +++ b/lib/OrdinaryDiffEqRKIP/test/semilinear_pde_test_cpu.jl @@ -3,7 +3,6 @@ using LinearAlgebra: norm include("semilinear_fft.jl") - @testset "Lugiato-Lefever equation test - Float64" begin @test !isnothing(lle_scan_test()) end @@ -12,7 +11,6 @@ end @test !isnothing(lle_scan_test(Float32)) end - @testset "NLSE test - Float64" begin sol = nlse_test() @@ -31,4 +29,4 @@ end p2 = abs2.(sol.u[end]) @test norm(p1 .- p2) / norm(p1) < 1e-2 # small difference with analytical solution due to boundary periodic condition -end \ No newline at end of file +end diff --git a/lib/OrdinaryDiffEqRKIP/test/type_test.jl b/lib/OrdinaryDiffEqRKIP/test/type_test.jl index 48b7e36b99..d0c918946c 100644 --- a/lib/OrdinaryDiffEqRKIP/test/type_test.jl +++ b/lib/OrdinaryDiffEqRKIP/test/type_test.jl @@ -1,7 +1,8 @@ using Test using OrdinaryDiffEqRKIP: RKIP using SciMLBase: SplitODEProblem, SplitFunction, solve -using SciMLOperators: ScalarOperator, DiagonalOperator, MatrixOperator, AbstractSciMLOperator +using SciMLOperators: ScalarOperator, DiagonalOperator, MatrixOperator, + AbstractSciMLOperator using StaticArrays @@ -10,13 +11,12 @@ import LinearAlgebra LinearAlgebra.exp(d::ScalarOperator, t) = ScalarOperator(exp(t * d.val)) # Temporary Fix for missing exp dispatch LinearAlgebra.exp(d::MatrixOperator, t) = MatrixOperator(exp(t * d.A)) -function test(A_prototype, u_prototype, iip; use_ldiv=false) +function test(A_prototype, u_prototype, iip; use_ldiv = false) for reltol in [1e-3, 1e-6, 1e-8] # fail for 1e-9, probably due to floating point error μ = 1.05 analytic = (u0, _, t) -> u0 .* exp(2μ * t) for μₐ in [2μ, μ, 0.0] - Â = A_prototype(μₐ) f = (du, u, _, _) -> du .= (2μ - μₐ) .* u @@ -27,23 +27,22 @@ function test(A_prototype, u_prototype, iip; use_ldiv=false) u0 = u_prototype(1.0) t = (0.0, 1.0) - splfc = SplitFunction{iip}(Â, f; analytic=analytic) + splfc = SplitFunction{iip}(Â, f; analytic = analytic) spltode = SplitODEProblem(splfc, u0, t) - sol = solve(spltode, RKIP(; use_ldiv=use_ldiv); reltol=reltol) - @test isapprox(sol(t[end]), splfc.analytic(u0, nothing, t[end]); rtol=reltol) + sol = solve(spltode, RKIP(; use_ldiv = use_ldiv); reltol = reltol) + @test isapprox(sol(t[end]), splfc.analytic(u0, nothing, t[end]); rtol = reltol) end end end @testset "In-Place ScalarOperator Vector Adaptative Ldiv" begin - test(μ -> ScalarOperator(μ), u0 -> [u0], true; use_ldiv=true) + test(μ -> ScalarOperator(μ), u0 -> [u0], true; use_ldiv = true) end @testset "In-Place DiagonalOperator Vector Adaptative Ldiv" begin - test(μ -> DiagonalOperator([μ]), u0 -> [u0], true; use_ldiv=true) + test(μ -> DiagonalOperator([μ]), u0 -> [u0], true; use_ldiv = true) end - @testset "In-Place ScalarOperator Vector Adaptative" begin test(μ -> ScalarOperator(μ), u0 -> [u0], true) end @@ -66,7 +65,6 @@ end # Failing Test due to missing dispatch in SciMLOperator - # fails as calling a MatrixOperator on a StaticVector change its type to Vector, causing a type instability # @testset "Out-of-place Diagonal 1x1 Operator SVector Adaptative" begin @@ -77,8 +75,8 @@ end # test((μ) -> MatrixOperator([μ;;]), u0 -> SVector(u0), false) # end -# # Fails for a strange reason as calling ldiv!(a, B, c) on a MatrixOperator dispatch toward an inexisting method +# # Fails for a strange reason as calling ldiv!(a, B, c) on a MatrixOperator dispatch toward an inexisting method # @testset "In-Place MatrixOperator Vector Float Adaptative Ldiv" begin # test(μ -> MatrixOperator([μ;;]), u0 -> [u0], true; use_ldiv=true) -# end +# end From 0cb3e46a575b09191bace01db901288864f1afb0 Mon Sep 17 00:00:00 2001 From: Azercoco <20425137+Azercoco@users.noreply.github.com> Date: Tue, 10 Jun 2025 17:05:13 +0200 Subject: [PATCH 03/10] Use generic_solver_docstring for the documentation --- .../src/OrdinaryDiffEqRKIP.jl | 2 +- lib/OrdinaryDiffEqRKIP/src/algorithms.jl | 173 ++++++++++-------- 2 files changed, 98 insertions(+), 77 deletions(-) diff --git a/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl b/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl index 51257cb75d..31515ae1cf 100644 --- a/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl +++ b/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl @@ -14,7 +14,7 @@ import OrdinaryDiffEqCore: OrdinaryDiffEqAdaptiveExponentialAlgorithm, alg_adapt has_dtnew_modification, calculate_residuals, calculate_residuals!, increment_nf!, OrdinaryDiffEqAdaptiveAlgorithm, OrdinaryDiffEqMutableCache, - dtnew_modification + dtnew_modification, generic_solver_docstring include("rkip_cache.jl") include("algorithms.jl") diff --git a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl index 33d9cb2d69..15a4ecc6a5 100644 --- a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl +++ b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl @@ -1,88 +1,46 @@ using OrdinaryDiffEqCore: OrdinaryDiffEqCore -mutable struct RKIP_ALG{tableauType<:ExplicitRKTableau,elType,dtType<:AbstractVector{elType}} <: OrdinaryDiffEqAdaptiveAlgorithm - tableau::tableauType - dt_for_expÂ_caching::dtType - clamp_lower_dt::Bool - clamp_higher_dt::Bool - use_ldiv::Bool - cache::Union{Nothing,RKIPCache} -end - -alg_order(alg::RKIP_ALG) = alg.tableau.order -alg_adaptive_order(alg::RKIP_ALG) = alg.tableau.adaptiveorder - -has_dtnew_modification(alg::RKIP_ALG) = true - -function dtnew_modification(alg::RKIP_ALG{tableauType,elType,dtType}, dtnew) where {tableauType,elType,dtType} - @unpack dt_for_expÂ_caching = alg - if first(alg.dt_for_expÂ_caching) > dtnew && alg.clamp_lower_dt - dtnew = first(alg.dt_for_expÂ_caching) - elseif last(alg.dt_for_expÂ_caching) < dtnew && alg.clamp_higher_dt - dtnew = last(alg.dt_for_expÂ_caching) - else - dtnew = alg.dt_for_expÂ_caching[searchsortedfirst(alg.dt_for_expÂ_caching, dtnew)] - end - return dtnew -end - -dtnew_modification(_, alg::RKIP_ALG, dtnew) = dtnew_modification(alg, dtnew) - - -""" -Runge-Kutta in the Interaction Picture (RIKP) method. +METHOD_DESCRIPTION = """ +Runge-Kutta in the interaction picture. This is suited for solving semilinear problem of the form: ```math \frac{du}{dt} = Au + f(u,p,t) ``` -where A is possibly stiff constant linear operator whose scaled exponential exp(Ah) can be calculated efficiently for any h. -The problem is first transformed into the Interaction Picture. +where A is possibly stiff constant time-indepdant linear operator whose scaled exponential exp(Ah) can be calculated efficiently for any h. +The problem is first transformed in a non-stiff variant (interaction picture) ```math u_I(t) = \exp(At) u_0 \frac{du_I}{dt} = f_I(u_I,p,t) f_I(u_I,p,t) = f(exp(-At)u_I, p, t) ``` -This new system is then solved with an explicit Runge-Kutta method. +and is then solved with an explicit (adaptive) Runge-Kutta method. -This type of problem is often encountered in semilinear parabolic PDE: heat diffusion, Schrödinger equation ... - -This method is only implemented for semilinear problem: `SplitODEProblem` when the first function `f1` is a `AbsractSciMLOperator` A. - -The algorithm needs for A to have the following traits of the `AbsractSciMLOperator` implemented +This solver is only implemented for semilinear problem: `SplitODEProblem` when the first function `f1` is a `AbsractSciMLOperator` A implementing: ```julia -LinearAlgebra.exp(A, t) = #... +LinearAlgebra.exp(A, t) # = exp(A*t) +``` +`A` and the return value of `exp(A, t)` must either also both implement: +``` A(du, u, v, p, t) # for in-place problem -A(u, v p, t) # for out-of-place problem +A(u, v p, t) # for out-of-place problem ``` -Moreover, the resulting ```expA = exp(A, t)``` must also implement ```(expA)(u, v, p, t)``` and/or ```(expA)(du, u, v, p, t)```. -The algorithm only works for constant `A` (does no depend on v, p, t). - -The algorithm will cache and resue the computed operator-exponential for a set of steps. +For performance, the algorithm will cache and reuse the computed operator-exponential for a fixed set of time steps. # Arguments - `dtmin::T`: the smallest step `dt` for which `exp(A*dt)` will be cached. Default is `1e-3` - `dtmax::T`: the largest step `dt` for which `exp(A*dt)` will be cached. Default is `1.0` -The set of cached steps size `dt_for_expA_caching` will follow a geometric progression such that `first(dt_for_expA_caching) = dtmin` and `last(dt_for_expA_caching) = dtmax`. - -Time stepping can still happen outside the bonds but this will results in no cache (as `exp(A*h)` will be computed each step) degrading the performances. The time step can be clamped within the cache range through the args `clamp_lower_dt` and `clamp_higher_dt`. - -# Optional Arguments -- `nb_of_cache_step::Integer`: the number of steps. Default is `100`. -- `tableau::ExplicitRKTableau`: the Runge-Kutta Tableau to use, default is `constructDormandPrince6()`. -- `clamp_lower_dt::Bool`: weither to clamp proposed step to the smallest cached step in order to force the use of cached exponential, improving performance. - This may prevent reaching the desired tolerance. Default is `false`. -- `clamp_higher_dt::Bool`: weither to clamp proposed step to the largest cached step in order to force the use of cached exponential, improving performance. - This can cause performance degredation if `integrator.dtmax` is too small. Default is `true`. -- `use_ldiv::Bool`: weither, to use `ldiv(exp(A, t), v)` instead of caching `exp(A, -t)*v`. Reduce the memory usage but slightly less efficient. `ldiv` must be implemented. Only works for in-place problem. Default is `false`. +The fixed steps will follow a geometric progression. +Time stepping can still happen outside the bonds (for the end step for e.g) but no cache will occur (`exp(A*dt)` getting computed each step) degrading the performances. +The time step can be forcibly clamped within the cache range through the keywords `clamp_lower_dt` and `clamp_higher_dt`. -Cached operator exponentials are stored in the alorithm meaning that in the following code: +The cached operator exponentials are also directly stored in the alorithm such that: ```julia rkip = RKIP() @@ -90,29 +48,85 @@ solve(ode_prob_1, rkip, t1) solve(ode_prob_2, rkip, t2) ```` -will reuse the cache of the first solve. This can be useful when needed to solve several time the same problem but one must be sure that `A` does not change. +will reuse the precomputed exponential cached during the first `solve` call. +This can be useful when needed to solve several time the same problem succesively from a same `A`. """ -RKIP(dtmin::T=1e-3, dtmax::T=1.0; nb_of_cache_step::Int=100, tableau=constructDormandPrince6(T), clamp_lower_dt::Bool=false, clamp_higher_dt::Bool=true, use_ldiv=false) where {T} = RKIP_ALG{typeof(tableau),T,Vector{T}}( - tableau, - logrange(dtmin, dtmax, nb_of_cache_step), - clamp_lower_dt, - clamp_higher_dt, - use_ldiv, - nothing -) - -function alg_cache(alg::RKIP_ALG, u::uType, rate_prototype, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits, uprev, uprev2, f, t, dt, reltol, p, calck, iip) where {uType} +REFERENCE = """Zhongxi Zhang, Liang Chen, and Xiaoyi Bao, "A fourth-order Runge-Kutta in the interaction picture method for numerically solving the coupled nonlinear Schrödinger equation," Opt. Express 18, 8261-8276 (2010)""" + +KEYWORD_DESCRIPTION = """ +- `nb_of_cache_step::Integer`: the number of steps. Default is `100`. +- `tableau::ExplicitRKTableau`: the Runge-Kutta Tableau to use. Default is `constructDormandPrince6()`. +- `clamp_lower_dt::Bool`: weither to clamp proposed step to the smallest cached step in order to force the use of cached exponential, improving performance. + This may prevent reaching the desired tolerance. Default is `false`. +- `clamp_higher_dt::Bool`: weither to clamp proposed step to the largest cached step in order to force the use of cached exponential, improving performance. + This can cause performance degredation if `integrator.dtmax` is too small. Default is `true`. +- `use_ldiv::Bool`: weither, to use `ldiv(exp(A, t), v)` instead of caching `exp(A, -t)*v`. Reduce the memory usage but slightly less efficient. `ldiv` must be implemented. Only works for in-place problem. Default is `false`. +""" + +@doc generic_solver_docstring( + METHOD_DESCRIPTION, "RKIP", "Adaptative Exponential Runge-Kutta", + REFERENCE, KEYWORD_DESCRIPTION, "") +mutable struct RKIP{ + tableauType <: ExplicitRKTableau, elType, dtType <: AbstractVector{elType}} <: + OrdinaryDiffEqAdaptiveAlgorithm + tableau::tableauType + dt_for_expÂ_caching::dtType + clamp_lower_dt::Bool + clamp_higher_dt::Bool + use_ldiv::Bool + cache::Union{Nothing, RKIPCache} +end + +function RKIP(dtmin::T = 1e-3, dtmax::T = 1.0; nb_of_cache_step::Int = 100, + tableau = constructDormandPrince6(T), clamp_lower_dt::Bool = false, + clamp_higher_dt::Bool = true, use_ldiv = false) where {T} + RKIP{ + typeof(tableau), T, Vector{T}}( + tableau, + logrange(dtmin, dtmax, nb_of_cache_step), + clamp_lower_dt, + clamp_higher_dt, + use_ldiv, + nothing + ) +end + + +alg_order(alg::RKIP) = alg.tableau.order +alg_adaptive_order(alg::RKIP) = alg.tableau.adaptiveorder + +has_dtnew_modification(alg::RKIP) = true + +function dtnew_modification(alg::RKIP{tableauType, elType, dtType}, + dtnew) where {tableauType, elType, dtType} + @unpack dt_for_expÂ_caching = alg + if first(alg.dt_for_expÂ_caching) > dtnew && alg.clamp_lower_dt + dtnew = first(alg.dt_for_expÂ_caching) + elseif last(alg.dt_for_expÂ_caching) < dtnew && alg.clamp_higher_dt + dtnew = last(alg.dt_for_expÂ_caching) + else + dtnew = alg.dt_for_expÂ_caching[searchsortedfirst(alg.dt_for_expÂ_caching, dtnew)] + end + return dtnew +end + +dtnew_modification(_, alg::RKIP, dtnew) = dtnew_modification(alg, dtnew) + + +function alg_cache( + alg::RKIP, u::uType, rate_prototype, uEltypeNoUnits, uBottomEltypeNoUnits, + tTypeNoUnits, uprev, uprev2, f, t, dt, reltol, p, calck, iip) where {uType} tmp = zero(u) utilde = zero(u) kk = [zero(u) for _ in 1:(alg.tableau.stages)] - Â = isa(f, SplitFunction) ? f.f1.f : throw(ArgumentError("RKIP is only implemented for semilinear problems")) + Â = isa(f, SplitFunction) ? f.f1.f : + throw(ArgumentError("RKIP is only implemented for semilinear problems")) opType = typeof(Â) expOpType = typeof(exp(Â, 1.0)) if isnothing(alg.cache) - is_cached = Vector{Bool}(undef, length(alg.dt_for_expÂ_caching)) fill!(is_cached, false) @@ -120,16 +134,22 @@ function alg_cache(alg::RKIP_ALG, u::uType, rate_prototype, uEltypeNoUnits, uBot c_unique = unique(c_extended) # in some tableau, there is duplicate: we only keep the unique value to save on caching time and memory c_index = [findfirst(==(c), c_unique) for c in c_extended] # index mapping - exp_cache = ExpCache{expOpType}(Array{expOpType,2}(undef, length(alg.dt_for_expÂ_caching), length(c_unique)), Vector{expOpType}(undef, length(c_unique))) + exp_cache = ExpCache{expOpType}( + Array{expOpType, 2}(undef, length(alg.dt_for_expÂ_caching), length(c_unique)), + Vector{expOpType}(undef, length(c_unique))) if !alg.use_ldiv - exp_cache = ExpCacheNoLdiv(exp_cache, ExpCache{expOpType}(Array{expOpType,2}(undef, length(alg.dt_for_expÂ_caching), length(c_unique)), Vector{expOpType}(undef, length(c_unique)))) + exp_cache = ExpCacheNoLdiv(exp_cache, + ExpCache{expOpType}( + Array{expOpType, 2}( + undef, length(alg.dt_for_expÂ_caching), length(c_unique)), + Vector{expOpType}(undef, length(c_unique)))) expCacheType = ExpCacheNoLdiv{expOpType} else expCacheType = ExpCache{expOpType} end - alg.cache = RKIPCache{expOpType,expCacheType,tTypeNoUnits,opType,uType,iip}( + alg.cache = RKIPCache{expOpType, expCacheType, tTypeNoUnits, opType, uType, iip}( exp_cache, zero(tTypeNoUnits), is_cached, @@ -139,8 +159,9 @@ function alg_cache(alg::RKIP_ALG, u::uType, rate_prototype, uEltypeNoUnits, uBot c_unique, c_index ) - else # cache recycling - alg.cache = RKIPCache{expOpType,typeof(alg.cache.exp_cache),tTypeNoUnits,opType,uType,iip}( + else # cache recycling + alg.cache = RKIPCache{ + expOpType, typeof(alg.cache.exp_cache), tTypeNoUnits, opType, uType, iip}( alg.cache.exp_cache, alg.last_step, alg.cache.is_cached, From a42e3b4597bb39c2e6d4194b7a44f7087f1646fa Mon Sep 17 00:00:00 2001 From: Azercoco <20425137+Azercoco@users.noreply.github.com> Date: Tue, 10 Jun 2025 17:16:23 +0200 Subject: [PATCH 04/10] fix typos and improve docs --- lib/OrdinaryDiffEqRKIP/src/algorithms.jl | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl index 15a4ecc6a5..4ca04096c6 100644 --- a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl +++ b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl @@ -9,25 +9,25 @@ This is suited for solving semilinear problem of the form: \frac{du}{dt} = Au + f(u,p,t) ``` -where A is possibly stiff constant time-indepdant linear operator whose scaled exponential exp(Ah) can be calculated efficiently for any h. +where A is possibly stiff time-independent linear operator whose scaled exponential exp(Ah) can be calculated efficiently for any h. The problem is first transformed in a non-stiff variant (interaction picture) ```math -u_I(t) = \exp(At) u_0 +u_I(t) = \exp(-At) u(t) \frac{du_I}{dt} = f_I(u_I,p,t) f_I(u_I,p,t) = f(exp(-At)u_I, p, t) ``` and is then solved with an explicit (adaptive) Runge-Kutta method. -This solver is only implemented for semilinear problem: `SplitODEProblem` when the first function `f1` is a `AbsractSciMLOperator` A implementing: +This solver is only implemented for semilinear problem: `SplitODEProblem` when the first function `f1` is a `AbstractSciMLOperator` A implementing: ```julia LinearAlgebra.exp(A, t) # = exp(A*t) ``` -`A` and the return value of `exp(A, t)` must either also both implement: -``` +`A` and the return value of `exp(A, t)` must either also both implement the `AbstractSciMLOperator` interface: +```julia A(du, u, v, p, t) # for in-place problem -A(u, v p, t) # for out-of-place problem +A(u, v, p, t) # for out-of-place problem ``` For performance, the algorithm will cache and reuse the computed operator-exponential for a fixed set of time steps. @@ -37,7 +37,7 @@ For performance, the algorithm will cache and reuse the computed operator-expone - `dtmax::T`: the largest step `dt` for which `exp(A*dt)` will be cached. Default is `1.0` The fixed steps will follow a geometric progression. -Time stepping can still happen outside the bonds (for the end step for e.g) but no cache will occur (`exp(A*dt)` getting computed each step) degrading the performances. +Time stepping can still happen outside the bounds (for the end step for e.g) but no cache will occur (`exp(A*dt)` getting computed each step) degrading the performances. The time step can be forcibly clamped within the cache range through the keywords `clamp_lower_dt` and `clamp_higher_dt`. The cached operator exponentials are also directly stored in the alorithm such that: @@ -49,7 +49,7 @@ solve(ode_prob_2, rkip, t2) ```` will reuse the precomputed exponential cached during the first `solve` call. -This can be useful when needed to solve several time the same problem succesively from a same `A`. +This can be useful for solving several times successively problems with a common `A`. """ REFERENCE = """Zhongxi Zhang, Liang Chen, and Xiaoyi Bao, "A fourth-order Runge-Kutta in the interaction picture method for numerically solving the coupled nonlinear Schrödinger equation," Opt. Express 18, 8261-8276 (2010)""" @@ -57,11 +57,11 @@ REFERENCE = """Zhongxi Zhang, Liang Chen, and Xiaoyi Bao, "A fourth-order Runge- KEYWORD_DESCRIPTION = """ - `nb_of_cache_step::Integer`: the number of steps. Default is `100`. - `tableau::ExplicitRKTableau`: the Runge-Kutta Tableau to use. Default is `constructDormandPrince6()`. -- `clamp_lower_dt::Bool`: weither to clamp proposed step to the smallest cached step in order to force the use of cached exponential, improving performance. +- `clamp_lower_dt::Bool`: whether to clamp proposed step to the smallest cached step in order to force the use of cached exponential, improving performance. This may prevent reaching the desired tolerance. Default is `false`. -- `clamp_higher_dt::Bool`: weither to clamp proposed step to the largest cached step in order to force the use of cached exponential, improving performance. - This can cause performance degredation if `integrator.dtmax` is too small. Default is `true`. -- `use_ldiv::Bool`: weither, to use `ldiv(exp(A, t), v)` instead of caching `exp(A, -t)*v`. Reduce the memory usage but slightly less efficient. `ldiv` must be implemented. Only works for in-place problem. Default is `false`. +- `clamp_higher_dt::Bool`: whether to clamp proposed step to the largest cached step in order to force the use of cached exponential, improving performance. + This can cause performance degradation if `integrator.dtmax` is too small. Default is `true`. +- `use_ldiv::Bool`: whether, to use `ldiv(exp(A, t), v)` instead of caching `exp(A, -t)*v`. Reduces the memory usage but is slightly less efficient. `ldiv` must be implemented. Only works for in-place problems. Default is `false`. """ @doc generic_solver_docstring( From b2902e978e480f14e861d97782b8dd5f22562d38 Mon Sep 17 00:00:00 2001 From: Azercoco <20425137+Azercoco@users.noreply.github.com> Date: Tue, 10 Jun 2025 17:20:56 +0200 Subject: [PATCH 05/10] Add aligned in math mode --- lib/OrdinaryDiffEqRKIP/src/algorithms.jl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl index 4ca04096c6..bf01d1f15c 100644 --- a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl +++ b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl @@ -13,9 +13,11 @@ where A is possibly stiff time-independent linear operator whose scaled exponent The problem is first transformed in a non-stiff variant (interaction picture) ```math -u_I(t) = \exp(-At) u(t) -\frac{du_I}{dt} = f_I(u_I,p,t) -f_I(u_I,p,t) = f(exp(-At)u_I, p, t) +\begin{aligned} +u_I(t) &= \exp(-At) u(t) \\ +\frac{du_I}{dt} &= f_I(u_I,p,t) \\ +f_I(u_I,p,t) &= f(exp(-At)u_I, p, t) \\ +\end{aligned} ``` and is then solved with an explicit (adaptive) Runge-Kutta method. From d22b70796411df654aa7103a61c75ff1a27828e6 Mon Sep 17 00:00:00 2001 From: Azercoco <20425137+Azercoco@users.noreply.github.com> Date: Thu, 12 Jun 2025 11:28:49 +0200 Subject: [PATCH 06/10] Fix cache recycling with test and change default tableau to Verner6 --- .../src/OrdinaryDiffEqRKIP.jl | 2 +- lib/OrdinaryDiffEqRKIP/src/algorithms.jl | 12 +++++------ .../test/cache_recycling_test.jl | 21 +++++++++++++++++++ lib/OrdinaryDiffEqRKIP/test/runtests.jl | 1 + lib/OrdinaryDiffEqRKIP/test/type_test.jl | 12 +++++++---- 5 files changed, 36 insertions(+), 12 deletions(-) create mode 100644 lib/OrdinaryDiffEqRKIP/test/cache_recycling_test.jl diff --git a/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl b/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl index 31515ae1cf..f6d1a0ea7e 100644 --- a/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl +++ b/lib/OrdinaryDiffEqRKIP/src/OrdinaryDiffEqRKIP.jl @@ -6,7 +6,7 @@ using UnPack: @pack!, @unpack using MaybeInplace: @bb using SciMLBase: isinplace using DiffEqBase: ExplicitRKTableau -using DiffEqDevTools: constructDormandPrince6 +using DiffEqDevTools: constructVerner6 import OrdinaryDiffEqCore: OrdinaryDiffEqAdaptiveExponentialAlgorithm, alg_adaptive_order, alg_order, alg_cache, @cache, SplitFunction, get_fsalfirstlast, diff --git a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl index bf01d1f15c..f33b13e955 100644 --- a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl +++ b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl @@ -58,7 +58,7 @@ REFERENCE = """Zhongxi Zhang, Liang Chen, and Xiaoyi Bao, "A fourth-order Runge- KEYWORD_DESCRIPTION = """ - `nb_of_cache_step::Integer`: the number of steps. Default is `100`. -- `tableau::ExplicitRKTableau`: the Runge-Kutta Tableau to use. Default is `constructDormandPrince6()`. +- `tableau::ExplicitRKTableau`: the Runge-Kutta Tableau to use. Default is `_`. (TO BE DETERMINED) - `clamp_lower_dt::Bool`: whether to clamp proposed step to the smallest cached step in order to force the use of cached exponential, improving performance. This may prevent reaching the desired tolerance. Default is `false`. - `clamp_higher_dt::Bool`: whether to clamp proposed step to the largest cached step in order to force the use of cached exponential, improving performance. @@ -81,7 +81,7 @@ mutable struct RKIP{ end function RKIP(dtmin::T = 1e-3, dtmax::T = 1.0; nb_of_cache_step::Int = 100, - tableau = constructDormandPrince6(T), clamp_lower_dt::Bool = false, + tableau = constructVerner6(T), clamp_lower_dt::Bool = false, clamp_higher_dt::Bool = true, use_ldiv = false) where {T} RKIP{ typeof(tableau), T, Vector{T}}( @@ -94,7 +94,6 @@ function RKIP(dtmin::T = 1e-3, dtmax::T = 1.0; nb_of_cache_step::Int = 100, ) end - alg_order(alg::RKIP) = alg.tableau.order alg_adaptive_order(alg::RKIP) = alg.tableau.adaptiveorder @@ -115,7 +114,6 @@ end dtnew_modification(_, alg::RKIP, dtnew) = dtnew_modification(alg, dtnew) - function alg_cache( alg::RKIP, u::uType, rate_prototype, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits, uprev, uprev2, f, t, dt, reltol, p, calck, iip) where {uType} @@ -165,13 +163,13 @@ function alg_cache( alg.cache = RKIPCache{ expOpType, typeof(alg.cache.exp_cache), tTypeNoUnits, opType, uType, iip}( alg.cache.exp_cache, - alg.last_step, - alg.cache.is_cached, + alg.cache.last_step, + alg.cache.cached, tmp, utilde, kk, alg.cache.c_unique, - alg.cache.c_index + alg.cache.c_mapping ) end return alg.cache diff --git a/lib/OrdinaryDiffEqRKIP/test/cache_recycling_test.jl b/lib/OrdinaryDiffEqRKIP/test/cache_recycling_test.jl new file mode 100644 index 0000000000..2e51e5ee7f --- /dev/null +++ b/lib/OrdinaryDiffEqRKIP/test/cache_recycling_test.jl @@ -0,0 +1,21 @@ +using SciMLOperators: MatrixOperator +using OrdinaryDiffEqRKIP: RKIP +using SciMLBase: SplitODEProblem, SplitFunction, solve + +import LinearAlgebra + +@testset "Cache Recycling Test" begin + matrix = [1 0.2 -0.2; 0.4 1 -1; 0.2 0.0 1.2] + Â = MatrixOperator(matrix) + u0 = [0.0, 0.0, 0.0] + f! = (dx, x, p, t) -> dx .= cos.(x) + + alg = RKIP(1e-2, 1e1) + + splfc = SplitFunction{true}(Â, f!) + spltode = SplitODEProblem(splfc, u0, (0, 1.0)) + + sol_1 = solve(spltode, alg).u[end] + sol_2 = solve(spltode, alg).u[end] + @test all(sol_1 .≈ sol_2) +end diff --git a/lib/OrdinaryDiffEqRKIP/test/runtests.jl b/lib/OrdinaryDiffEqRKIP/test/runtests.jl index 497d0316ec..e492c53cf3 100644 --- a/lib/OrdinaryDiffEqRKIP/test/runtests.jl +++ b/lib/OrdinaryDiffEqRKIP/test/runtests.jl @@ -1,5 +1,6 @@ using SafeTestsets @safetestset "Type Safety Tests" include("type_test.jl") +@safetestset "Cache Test" include("cache_recycling_test.jl") @safetestset "Fourier Semilinear PDE Tests" include("semilinear_pde_test_cpu.jl") @safetestset "CUDA Test" include("semilinear_pde_test_cuda.jl") diff --git a/lib/OrdinaryDiffEqRKIP/test/type_test.jl b/lib/OrdinaryDiffEqRKIP/test/type_test.jl index d0c918946c..d33c51abc6 100644 --- a/lib/OrdinaryDiffEqRKIP/test/type_test.jl +++ b/lib/OrdinaryDiffEqRKIP/test/type_test.jl @@ -12,8 +12,9 @@ LinearAlgebra.exp(d::ScalarOperator, t) = ScalarOperator(exp(t * d.val)) # Tempo LinearAlgebra.exp(d::MatrixOperator, t) = MatrixOperator(exp(t * d.A)) function test(A_prototype, u_prototype, iip; use_ldiv = false) - for reltol in [1e-3, 1e-6, 1e-8] # fail for 1e-9, probably due to floating point error - μ = 1.05 + for reltol in [1e-3, 1e-6, 1e-8, 1e-10, 1e-12], + μ in 1.05 + analytic = (u0, _, t) -> u0 .* exp(2μ * t) for μₐ in [2μ, μ, 0.0] @@ -29,8 +30,11 @@ function test(A_prototype, u_prototype, iip; use_ldiv = false) splfc = SplitFunction{iip}(Â, f; analytic = analytic) spltode = SplitODEProblem(splfc, u0, t) - sol = solve(spltode, RKIP(; use_ldiv = use_ldiv); reltol = reltol) - @test isapprox(sol(t[end]), splfc.analytic(u0, nothing, t[end]); rtol = reltol) + sol = solve( + spltode, RKIP(; use_ldiv = use_ldiv); reltol = reltol, abstol = 1e-10) + + @test isapprox(sol(t[end]), splfc.analytic(u0, nothing, t[end]); + rtol = 1e2 * reltol, atol = 1e-8) end end end From 5e40016e8467c55b30f2a2e72ebe66af35844ff5 Mon Sep 17 00:00:00 2001 From: Azercoco <20425137+Azercoco@users.noreply.github.com> Date: Thu, 12 Jun 2025 12:27:58 +0200 Subject: [PATCH 07/10] Dispatch p and t to the operator --- .../src/rkip_perform_step.jl | 18 +++++++------ lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl | 25 ++++++++++--------- 2 files changed, 23 insertions(+), 20 deletions(-) diff --git a/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl b/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl index 6fcee2a47e..4a4cf63bf4 100644 --- a/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl +++ b/lib/OrdinaryDiffEqRKIP/src/rkip_perform_step.jl @@ -20,7 +20,7 @@ Retunr res if in place, otherwise return Au + f(u, p, t) function f_mip!(res::uType, tmp::uType, A::opType, f, u::uType, p, t::tType, ::Val{true}) where {tType, uType, opType} res .= u - res = matvec_prod_mip(tmp, A, res, Val(true)) + res = matvec_prod_mip(tmp, A, res, Val(true), p, t) f(tmp, u, p, t) res .+= tmp return res @@ -28,7 +28,7 @@ end function f_mip!(_::uType, tmp::uType, A::opType, f, u::uType, p, t::tType, ::Val{false}) where {tType, uType, opType} - matvec_prod_mip(tmp, A, u, Val(false)) + f(u, p, t) + matvec_prod_mip(tmp, A, u, Val(false), p, t) + f(u, p, t) end """ @@ -59,7 +59,7 @@ end @bb u .= uprev - for i in 1:(stages) + for i in 1:(stages) @bb kk[i] .= uprev for j in 1:(i - 1) @@ -67,12 +67,14 @@ end kk[i] = axpy_mip(g, kk[j], kk[i], iip) # kk_i += dt*A[i, j] * kk_j end + t_ = t + c[i] * dt + # if mutable/heaps type, assignment does nothing as the function is in-place, - kk[i] = expmv_rkip_mip(cache, kk[i], dt, i) # kᵢ = exp(Â * [dt * cᵢ])*kᵢ ➡ Change from interaction picture to "true" coordinate + kk[i] = expmv_rkip_mip(cache, kk[i], dt, i, p, t_) # kᵢ = exp(Â * [dt * cᵢ])*kᵢ ➡ Change from interaction picture to "true" coordinate - kk[i] = nl_part_mip(cache.tmp, f.f2, kk[i], p, t + c[i] * dt, iip) # kᵢ = f(u + Σⱼ dt*(Aᵢⱼ kⱼ), t + dt*cᵢ) + kk[i] = nl_part_mip(cache.tmp, f.f2, kk[i], p, t_, iip) # kᵢ = f(u + Σⱼ dt*(Aᵢⱼ kⱼ), t + dt*cᵢ) - kk[i] = expmv_rkip_mip(cache, kk[i], -dt, i) # kᵢ = exp(-Â * [dt * cᵢ])*kᵢ ➡ Going back in interaction picture + kk[i] = expmv_rkip_mip(cache, kk[i], -dt, i, p, t_) # kᵢ = exp(-Â * [dt * cᵢ])*kᵢ ➡ Going back in interaction picture integrator.stats.nf += 2 # two exp vec product integrator.stats.nf2 += 1 # one function evaluation @@ -88,10 +90,10 @@ end end end - u = expmv_rkip_mip(cache, u, dt) # stepping forward into the interaction u = exp(Â dt)*u + u = expmv_rkip_mip(cache, u, dt, p, t) # stepping forward into the interaction u = exp(Â dt)*u if adaptive - utilde = expmv_rkip_mip(cache, utilde, dt) # stepping forward into the interaction ũ = exp(Â dt)*ũ + utilde = expmv_rkip_mip(cache, utilde, dt, p, t) # stepping forward into the interaction ũ = exp(Â dt)*ũ tmp = calculate_residuals_mip( tmp, utilde, uprev, u, abstol, reltol, internalnorm, t, iip) # error computation maybe in place integrator.EEst = internalnorm(tmp, t) diff --git a/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl b/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl index ae2078b8d7..0eff0d5e69 100644 --- a/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl +++ b/lib/OrdinaryDiffEqRKIP/src/rkip_utils.jl @@ -7,35 +7,36 @@ Maybe-In Place "mip" Safe function for computing the matrix exponential vector product for both mutable and immutable type. Same principle of operation as `_safe_matvec_prod` for mutability/in-place handling. """ -@inline expmv_rkip_mip(cache::RKIPCache{expOpType, cacheType, tType, opType, uType, iip}, v::uType, h::tType) where {expOpType, cacheType, tType, opType, uType, iip} = expmv_rkip_mip( - cache, v, h, lastindex(cache.c_mapping)) # If i is not precised, this is the final step in the RKIP -> h = dt +@inline expmv_rkip_mip(cache::RKIPCache{expOpType, cacheType, tType, opType, uType, iip}, v::uType, h::tType, p, t) where {expOpType, cacheType, tType, opType, uType, iip} = expmv_rkip_mip( + cache, v, h, lastindex(cache.c_mapping), p, t) # If i is not precised, this is the final step in the RKIP -> h = dt @inline function expmv_rkip_mip( cache::RKIPCache{expOpType, cacheType, tType, opType, uType, iip}, v::uType, - h::tType, stage_index::Int) where {expOpType, cacheType, tType, opType, uType, iip} + h::tType, stage_index::Int, p, t) where { + expOpType, cacheType, tType, opType, uType, iip} if !(h ≈ tType(0.0)) c = cache.c_mapping[stage_index] exp_cache = cache.exp_cache tmp = cache.tmp - return _expmv_rkip_mip(exp_cache, tmp, v, c, h > tType(0.0), iip) + return _expmv_rkip_mip(exp_cache, tmp, v, c, h > tType(0.0), iip, p, t) end return v end @inline function _expmv_rkip_mip( cache::ExpCacheNoLdiv{expOpType}, tmp::vType, v::vType, stage_index::Integer, - positive, iip) where {expOpType <: AbstractSciMLOperator, vType} + positive, iip, p, t) where {expOpType <: AbstractSciMLOperator, vType} op = get_op_for_this_step(cache, positive, stage_index) - v = matvec_prod_mip(tmp, op, v, iip) + v = matvec_prod_mip(tmp, op, v, iip, p, t) return v end @inline function _expmv_rkip_mip( cache::ExpCache{expOpType}, tmp::vType, v::vType, stage_index::Integer, - positive, ::Val{true}) where {expOpType <: AbstractSciMLOperator, vType} + positive, ::Val{true}, p, t) where {expOpType <: AbstractSciMLOperator, vType} op = get_op_for_this_step(cache, stage_index) if positive - matvec_prod_mip(tmp, op, v, Val(true)) + matvec_prod_mip(tmp, op, v, Val(true), p, t) else ldiv!(tmp, op, v) #ldiv can only be used for in place call copyto!(v, tmp) @@ -50,14 +51,14 @@ For mutable type, overwrite `v` with `mat*v` and return `v`. Otherwise return ` Use the dispatch between ::Val{true} (in place) and ::Val{false} immutable to decide """ @inline function matvec_prod_mip( - tmp::V, mat::M, v::V, ::Val{true}) where {V, M <: AbstractSciMLOperator} - mat(tmp, v, nothing, nothing, zero(eltype(v))) + tmp::V, mat::M, v::V, ::Val{true}, p, t) where {V, M <: AbstractSciMLOperator} + mat(tmp, v, v, p, t) copyto!(v, tmp) return v end @inline function matvec_prod_mip( - _::V, mat::M, v::V, ::Val{false}) where {V, M <: AbstractSciMLOperator} - v = mat(v, nothing, nothing, zero(eltype(v))) + _::V, mat::M, v::V, ::Val{false}, p, t) where {V, M <: AbstractSciMLOperator} + v = mat(v, v, p, t) return v end From 6c7924a0563bcc7db62ab8fea70c3674f2fdd305 Mon Sep 17 00:00:00 2001 From: Azercoco <20425137+Azercoco@users.noreply.github.com> Date: Fri, 13 Jun 2025 10:39:39 +0200 Subject: [PATCH 08/10] Explicitely use Verner6 as the defaut tableau --- lib/OrdinaryDiffEqRKIP/src/algorithms.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl index f33b13e955..ddf1735ce4 100644 --- a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl +++ b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl @@ -58,7 +58,7 @@ REFERENCE = """Zhongxi Zhang, Liang Chen, and Xiaoyi Bao, "A fourth-order Runge- KEYWORD_DESCRIPTION = """ - `nb_of_cache_step::Integer`: the number of steps. Default is `100`. -- `tableau::ExplicitRKTableau`: the Runge-Kutta Tableau to use. Default is `_`. (TO BE DETERMINED) +- `tableau::ExplicitRKTableau`: the Runge-Kutta Tableau to use. Default is `constructVerner6()`. - `clamp_lower_dt::Bool`: whether to clamp proposed step to the smallest cached step in order to force the use of cached exponential, improving performance. This may prevent reaching the desired tolerance. Default is `false`. - `clamp_higher_dt::Bool`: whether to clamp proposed step to the largest cached step in order to force the use of cached exponential, improving performance. From d9a8e3fe55889cafd4a351480ef04e94097d5ca4 Mon Sep 17 00:00:00 2001 From: Azercoco <20425137+Azercoco@users.noreply.github.com> Date: Wed, 18 Jun 2025 11:14:32 +0200 Subject: [PATCH 09/10] Fix typo Co-authored-by: Christopher Rackauckas --- lib/OrdinaryDiffEqRKIP/src/algorithms.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl index ddf1735ce4..507068ea81 100644 --- a/lib/OrdinaryDiffEqRKIP/src/algorithms.jl +++ b/lib/OrdinaryDiffEqRKIP/src/algorithms.jl @@ -42,7 +42,7 @@ The fixed steps will follow a geometric progression. Time stepping can still happen outside the bounds (for the end step for e.g) but no cache will occur (`exp(A*dt)` getting computed each step) degrading the performances. The time step can be forcibly clamped within the cache range through the keywords `clamp_lower_dt` and `clamp_higher_dt`. -The cached operator exponentials are also directly stored in the alorithm such that: +The cached operator exponentials are also directly stored in the algorithm such that: ```julia rkip = RKIP() From 83439364a6ec00409642d625fcfa50ee98ec7434 Mon Sep 17 00:00:00 2001 From: Azercoco <20425137+Azercoco@users.noreply.github.com> Date: Wed, 18 Jun 2025 11:14:53 +0200 Subject: [PATCH 10/10] Remove direct export Co-authored-by: Christopher Rackauckas --- src/OrdinaryDiffEq.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/OrdinaryDiffEq.jl b/src/OrdinaryDiffEq.jl index b7c2b057a0..64cfcad11d 100644 --- a/src/OrdinaryDiffEq.jl +++ b/src/OrdinaryDiffEq.jl @@ -208,8 +208,6 @@ export LawsonEuler, NorsettEuler, ETD1, ETDRK2, ETDRK3, ETDRK4, HochOst4, Exp4, EPIRK4s3B, EPIRK5s3, EXPRB53s3, EPIRK5P1, EPIRK5P2, ETD2, Exprb32, Exprb43 -using OrdinaryDiffEqRKIP -export RKIP import PrecompileTools