Merge pull request #331 from SciML/qqy/simplify_interp

Author: ErikQQY

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -175,7 +175,6 @@ end
 
 function eval_bc_residual!(
         resid, ::StandardSecondOrderBVProblem, bc!::BC, sol, dsol, p, t) where {BC}
-    M = length(sol[1])
     bc!(resid, dsol, sol, p, t)
 end
 

lib/BoundaryValueDiffEqMIRK/src/adaptivity.jl

@@ -25,6 +25,7 @@ Generate new mesh based on the defect or the global error.
     (; order, errors, mesh, mesh_dt) = cache
     (abstol, _, _), _ = __split_kwargs(; cache.kwargs...)
     N = length(mesh)
+    n = N - 1
 
     safety_factor = T(1.3)
     ρ = T(1.0)
@@ -41,13 +42,12 @@ Generate new mesh based on the defect or the global error.
     r₃ = r₂ / (N - 1)
 
     n_predict = round(Int, (safety_factor * r₂) + 1)
-    n = N - 1
     n_ = T(0.1) * n
     n_predict = ifelse(abs((n_predict - n)) < n_, round(Int, n + n_), n_predict)
 
     if r₁ ≤ ρ * r₃
-        Nsub_star = 2 * (N - 1)
-        if Nsub_star > cache.alg.max_num_subintervals # Need to determine the too large threshold
+        Nsub_star = 2 * n
+        if Nsub_star > cache.alg.max_num_subintervals
             info = ReturnCode.Failure
             meshₒ = mesh
             mesh_dt₀ = mesh_dt
@@ -78,6 +78,7 @@ end
     (; order, errors, mesh, mesh_dt) = cache
     (abstol, _, _), _ = __split_kwargs(; cache.kwargs...)
     N = length(mesh)
+    n = N - 1
 
     safety_factor = T(1.3)
     ρ = T(2.0)
@@ -91,15 +92,14 @@ end
     ŝ .= (ŝ ./ abstol) .^ (T(1) / order)
     r₁ = maximum(ŝ)
     r₂ = sum(ŝ)
-    r₃ = r₂ / (N - 1)
+    r₃ = r₂ / n
 
     n_predict = round(Int, (safety_factor * r₂) + 1)
-    n = N - 1
     n_ = T(0.1) * n
     n_predict = ifelse(abs((n_predict - n)) < n_, round(Int, n + n_), n_predict)
 
     if r₁ ≤ ρ * r₃
-        Nsub_star = 2 * (N - 1)
+        Nsub_star = 2 * n
         # Need to determine the too large threshold
         if Nsub_star > cache.alg.max_num_subintervals
             info = ReturnCode.Failure
@@ -132,6 +132,7 @@ end
     (; order, errors, mesh, mesh_dt) = cache
     (abstol, _, _), _ = __split_kwargs(; cache.kwargs...)
     N = length(mesh)
+    n = N - 1
 
     safety_factor = T(1.3)
     ρ = T(2.0)
@@ -145,15 +146,15 @@ end
     ŝ .= (ŝ ./ abstol) .^ (T(1) / (order + 1))
     r₁ = maximum(ŝ)
     r₂ = sum(ŝ)
-    r₃ = r₂ / (N - 1)
+    r₃ = r₂ / n
 
     n_predict = round(Int, (safety_factor * r₂) + 1)
     n = N - 1
     n_ = T(0.1) * n
     n_predict = ifelse(abs((n_predict - n)) < n_, round(Int, n + n_), n_predict)
 
     if r₁ ≤ ρ * r₃
-        Nsub_star = 2 * (N - 1)
+        Nsub_star = 2 * n
         # Need to determine the too large threshold
         if Nsub_star > cache.alg.max_num_subintervals
             info = ReturnCode.Failure
@@ -186,30 +187,30 @@ end
     (; order, errors, mesh, mesh_dt) = cache
     (abstol, _, _), _ = __split_kwargs(; cache.kwargs...)
     N = length(mesh)
+    n = N - 1
 
     safety_factor = T(1.3)
     ρ = T(2.0)
     Nsub_star = 0
-    Nsub_star_ub = 4 * (N - 1)
+    Nsub_star_ub = 4 * n
     Nsub_star_lb = N ÷ 2
 
     info = ReturnCode.Success
 
-    ŝ₁ = [maximum(abs, d) for d in errors.u[1:(N - 1)]]
+    ŝ₁ = [maximum(abs, d) for d in errors.u[1:n]]
     ŝ₂ = [maximum(abs, d) for d in errors.u[N:end]]
     ŝ = similar(ŝ₁)
     ŝ .= (ŝ₁ ./ abstol) .^ (T(1) / (order + 1)) + (ŝ₂ ./ abstol) .^ (T(1) / (order + 1))
     r₁ = maximum(ŝ)
     r₂ = sum(ŝ)
-    r₃ = r₂ / (N - 1)
+    r₃ = r₂ / n
 
     n_predict = round(Int, (safety_factor * r₂) + 1)
-    n = N - 1
     n_ = T(0.1) * n
     n_predict = ifelse(abs((n_predict - n)) < n_, round(Int, n + n_), n_predict)
 
     if r₁ ≤ ρ * r₃
-        Nsub_star = 2 * (N - 1)
+        Nsub_star = 2 * n
         # Need to determine the too large threshold
         if Nsub_star > cache.alg.max_num_subintervals
             info = ReturnCode.Failure
@@ -627,7 +628,8 @@ end
 # Here we should not directly in-place change z in several steps
 # because in final step we actually need to use the original z(which is cache.y₀.u[i])
 # we use fᵢ₂_cache to avoid additional allocations.
-function sum_stages!(z::AbstractArray, cache::MIRKCache{iip, T, use_both, DiffCacheNeeded},
+@views function sum_stages!(
+        z::AbstractArray, cache::MIRKCache{iip, T, use_both, DiffCacheNeeded},
         w, i::Int, dt = cache.mesh_dt[i]) where {iip, T, use_both}
     (; stage, k_discrete, k_interp, fᵢ₂_cache) = cache
     (; s_star) = cache.ITU
@@ -640,7 +642,7 @@ function sum_stages!(z::AbstractArray, cache::MIRKCache{iip, T, use_both, DiffCa
 
     return nothing
 end
-function sum_stages!(
+@views function sum_stages!(
         z::AbstractArray, cache::MIRKCache{iip, T, use_both, NoDiffCacheNeeded},
         w, i::Int, dt = cache.mesh_dt[i]) where {iip, T, use_both}
     (; stage, k_discrete, k_interp, fᵢ₂_cache) = cache

lib/BoundaryValueDiffEqMIRK/src/interpolation.jl

@@ -67,23 +67,22 @@ end
     i = interval(mesh, t)
     dt = mesh_dt[i]
     τ = (t - mesh[i]) / dt
-    w, w′ = interp_weights(τ, cache.alg)
-    sum_stages!(z, id, cache, w, i)
+    w, _ = interp_weights(τ, cache.alg)
+    sum_stages!(z, id, cache, w, i, T)
 end
 
 @inline function interpolant!(dz::AbstractArray, id::MIRKInterpolation,
         cache::MIRKCache, t, mesh, mesh_dt, T::Type{Val{1}})
     i = interval(mesh, t)
     dt = mesh_dt[i]
     τ = (t - mesh[i]) / dt
-    w, w′ = interp_weights(τ, cache.alg)
-    z = similar(dz)
-    sum_stages!(z, dz, id, cache, w, w′, i)
+    _, w′ = interp_weights(τ, cache.alg)
+    sum_stages!(dz, id, cache, w′, i, T)
 end
 
-function sum_stages!(z::AbstractArray, id::MIRKInterpolation,
+@views function sum_stages!(z::AbstractArray, id::MIRKInterpolation,
         cache::MIRKCache{iip, T, use_both, DiffCacheNeeded},
-        w, i::Int) where {iip, T, use_both}
+        w, i::Int, ::Type{Val{0}}) where {iip, T, use_both}
     (; stage, k_discrete, k_interp) = cache
     (; s_star) = cache.ITU
     dt = cache.mesh_dt[i]
@@ -93,11 +92,11 @@ function sum_stages!(z::AbstractArray, id::MIRKInterpolation,
         z, k_interp.u[i][:, 1:(s_star - stage)], w[(stage + 1):s_star], true, true)
     z .= z .* dt .+ id.u[i]
 
-    return z
+    return nothing
 end
-function sum_stages!(z::AbstractArray, id::MIRKInterpolation,
+@views function sum_stages!(z::AbstractArray, id::MIRKInterpolation,
         cache::MIRKCache{iip, T, use_both, NoDiffCacheNeeded},
-        w, i::Int) where {iip, T, use_both}
+        w, i::Int, ::Type{Val{0}}) where {iip, T, use_both}
     (; stage, k_discrete, k_interp) = cache
     (; s_star) = cache.ITU
     dt = cache.mesh_dt[i]
@@ -107,44 +106,33 @@ function sum_stages!(z::AbstractArray, id::MIRKInterpolation,
         z, k_interp.u[i][:, 1:(s_star - stage)], w[(stage + 1):s_star], true, true)
     z .= z .* dt .+ id.u[i]
 
-    return z
+    return nothing
 end
 
 @views function sum_stages!(
-        z, z′, id::MIRKInterpolation, cache::MIRKCache{iip, T, use_both, DiffCacheNeeded},
-        w, w′, i::Int) where {iip, T, use_both}
+        z′, id::MIRKInterpolation, cache::MIRKCache{iip, T, use_both, DiffCacheNeeded},
+        w′, i::Int, ::Type{Val{1}}) where {iip, T, use_both}
     (; stage, k_discrete, k_interp) = cache
     (; s_star) = cache.ITU
-    dt = cache.mesh_dt[i]
-    z .= zero(z)
-    __maybe_matmul!(z, k_discrete[i].du[:, 1:stage], w[1:stage])
-    __maybe_matmul!(
-        z, k_interp.u[i][:, 1:(s_star - stage)], w[(stage + 1):s_star], true, true)
     z′ .= zero(z′)
     __maybe_matmul!(z′, k_discrete[i].du[:, 1:stage], w′[1:stage])
     __maybe_matmul!(
         z′, k_interp.u[i][:, 1:(s_star - stage)], w′[(stage + 1):s_star], true, true)
-    z .= z .* dt[1] .+ id.u[i]
 
-    return z, z′
+    return nothing
 end
 @views function sum_stages!(
-        z, z′, id::MIRKInterpolation, cache::MIRKCache{iip, T, use_both, NoDiffCacheNeeded},
-        w, w′, i::Int, dt = cache.mesh_dt[i]) where {iip, T, use_both}
+        z′, id::MIRKInterpolation, cache::MIRKCache{iip, T, use_both, NoDiffCacheNeeded},
+        w′, i::Int, ::Type{Val{1}}) where {iip, T, use_both}
     (; stage, k_discrete, k_interp) = cache
     (; s_star) = cache.ITU
 
-    z .= zero(z)
-    __maybe_matmul!(z, k_discrete[i][:, 1:stage], w[1:stage])
-    __maybe_matmul!(
-        z, k_interp.u[i][:, 1:(s_star - stage)], w[(stage + 1):s_star], true, true)
     z′ .= zero(z′)
     __maybe_matmul!(z′, k_discrete[i][:, 1:stage], w′[1:stage])
     __maybe_matmul!(
         z′, k_interp.u[i][:, 1:(s_star - stage)], w′[(stage + 1):s_star], true, true)
-    z .= z .* dt[1] .+ id.u[i]
 
-    return z, z′
+    return nothing
 end
 
 @inline __build_interpolation(cache::MIRKCache, u::AbstractVector) = MIRKInterpolation(
@@ -163,9 +151,9 @@ function (s::EvalSol{C})(tval::Number) where {C <: MIRKCache}
     dt = cache.mesh_dt[ii]
     τ = (tval - t[ii]) / dt
     w, _ = evalsol_interp_weights(τ, alg)
-    K = __needs_diffcache(alg.jac_alg) ? k_discrete[ii].du[:, 1:stage] :
-        k_discrete[ii][:, 1:stage]
-    __maybe_matmul!(z, K, w[1:stage])
+    K = __needs_diffcache(alg.jac_alg) ? @view(k_discrete[ii].du[:, 1:stage]) :
+        @view(k_discrete[ii][:, 1:stage])
+    __maybe_matmul!(z, K, @view(w[1:stage]))
     z .= z .* dt .+ u[ii]
     return z
 end