Fix and test DPRKN interpolation with idxs

Author: ChrisRackauckas

src/dense/generic_dense.jl

@@ -591,6 +591,7 @@ function ode_interpolant(Θ, dt, y₀, y₁, k, cache::OrdinaryDiffEqMutableCach
         T::Type{Val{TI}}, differential_vars) where {TI}
     if idxs isa Number || y₀ isa Union{Number, SArray}
         # typeof(y₀) can be these if saveidxs gives a single value
+        @show Main.@which _ode_interpolant(Θ, dt, y₀, y₁, k, cache, idxs, T, differential_vars)
         _ode_interpolant(Θ, dt, y₀, y₁, k, cache, idxs, T, differential_vars)
     elseif idxs isa Nothing
         if y₁ isa Array{<:Number}

src/dense/interpolants.jl

@@ -2947,6 +2947,28 @@ end
           b4Θ * k4[idxs] + b5Θ * k5[idxs] + b6Θ * k6[idxs])))
 end
 
+@muladd function _ode_interpolant(Θ, dt, y₀, y₁, k,
+        cache::Union{DPRKN6ConstantCache, DPRKN6Cache}, idxs::Number,
+        T::Type{Val{0}}, differential_vars::Nothing)
+    @dprkn6pre0
+    halfsize = length(y₀) ÷ 2
+    if idxs <= halfsize
+        duprev[idxs] +
+        dt * Θ *
+        (bp1Θ * k1[idxs] + bp3Θ * k3[idxs] +
+         bp4Θ * k4[idxs] + bp5Θ * k5[idxs] + bp6Θ * k6[idxs])
+    else
+        idxs = idxs - halfsize
+        uprev[idxs] +
+        dt * Θ *
+        (duprev[idxs] +
+         dt * Θ *
+         (b1Θ * k1[idxs] +
+          b3Θ * k3[idxs] +
+          b4Θ * k4[idxs] + b5Θ * k5[idxs] + b6Θ * k6[idxs]))
+    end
+end
+
 @muladd function _ode_interpolant!(out, Θ, dt, y₀, y₁, k,
         cache::Union{DPRKN6ConstantCache, DPRKN6Cache},
         idxs::Nothing, T::Type{Val{0}}, differential_vars::Nothing)
@@ -2976,25 +2998,25 @@ end
         cache::Union{DPRKN6ConstantCache, DPRKN6Cache}, idxs,
         T::Type{Val{0}}, differential_vars::Nothing)
     @dprkn6pre0
-    @inbounds @.. broadcast=false out.x[2]=uprev[idxs] +
+    halfsize = length(y₀) ÷ 2
+    isfirsthalf = idxs .<= halfsize
+    secondhalf =  idxs .> halfsize
+    firstidxs = idxs[isfirsthalf]
+    secondidxs_shifted = idxs[secondhalf]
+    secondidxs = secondidxs_shifted .- halfsize
+
+    @views @.. broadcast=false out[secondhalf]=uprev[secondidxs] +
                                            dt * Θ *
-                                           (duprev[idxs] +
+                                           (duprev[secondidxs] +
                                             dt * Θ *
-                                            (b1Θ * k1[idxs] +
-                                             b3Θ * k3[idxs] +
-                                             b4Θ * k4[idxs] + b5Θ * k5[idxs] +
-                                             b6Θ * k6[idxs]))
-    @inbounds @.. broadcast=false out.x[1]=duprev[idxs] +
+                                            (b1Θ * k1[secondidxs] +
+                                             b3Θ * k3[secondidxs] +
+                                             b4Θ * k4[secondidxs] + b5Θ * k5[secondidxs] +
+                                             b6Θ * k6[secondidxs]))
+    @views @.. broadcast=false out[isfirsthalf]=duprev[firstidxs] +
                                            dt * Θ *
-                                           (bp1Θ * k1[idxs] + bp3Θ * k3[idxs] +
-                                            bp4Θ * k4[idxs] + bp5Θ * k5[idxs] +
-                                            bp6Θ * k6[idxs])
-    #for (j,i) in enumerate(idxs)
-    #  out.x[2][j]  = uprev[i] + dt*Θ*(duprev[i] + dt*Θ*(b1Θ*k1[i] +
-    #                                                    b3Θ*k3[i] +
-    #                                                    b4Θ*k4[i] + b5Θ*k5[i] + b6Θ*k6[i]))
-    #  out.x[1][j] = duprev[i] + dt*Θ*(bp1Θ*k1[i] + bp3Θ*k3[i] +
-    #                                  bp4Θ*k4[i] + bp5Θ*k5[i] + bp6Θ*k6[i])
-    #end
+                                           (bp1Θ * k1[firstidxs] + bp3Θ * k3[firstidxs] +
+                                            bp4Θ * k4[firstidxs] + bp5Θ * k5[firstidxs] +
+                                            bp6Θ * k6[firstidxs])
     out
 end

test/interface/interpolation_output_types.jl

@@ -17,10 +17,10 @@ res1 = solve(prob, Vern8(), dt = 1 / 10, saveat = 1 / 10)
 res3 = solve(prob, CalvoSanz4(), dt = 1 / 10, saveat = 1 / 10)
 
 sol = solve(prob, CalvoSanz4(), dt = 1 / 10)
-@test sol(0.32) isa OrdinaryDiffEq.ArrayPartition
-@test sol(0.32, Val{1}) isa OrdinaryDiffEq.ArrayPartition
-@test sol(0.32, Val{2}) isa OrdinaryDiffEq.ArrayPartition
-@test sol(0.32, Val{3}) isa OrdinaryDiffEq.ArrayPartition
+@test sol(0.32) isa RecursiveArrayTools.ArrayPartition
+@test sol(0.32, Val{1}) isa RecursiveArrayTools.ArrayPartition
+@test sol(0.32, Val{2}) isa RecursiveArrayTools.ArrayPartition
+@test sol(0.32, Val{3}) isa RecursiveArrayTools.ArrayPartition
 
 function f(du, u, p, t)
     du .= u
@@ -40,3 +40,30 @@ sol(0:0.1:100; idxs = [1, 2])
 @test sol(0:0.1:100) isa DiffEqArray
 @test length(sol(0:0.1:100)) == length(0:0.1:100)
 @test length(sol(0:0.1:100).u[1]) == 3
+
+## Test DPRKN Interpolation
+
+#Parameters
+ω = 1
+
+#Initial Conditions
+x₀ = [0.0]
+dx₀ = [π / 2]
+tspan = (0.0, 2π)
+
+ϕ = atan((dx₀[1] / ω) / x₀[1])
+A = √(x₀[1]^2 + dx₀[1]^2)
+
+function harmonicoscillator(ddu, du, u, ω, t)
+    ddu .= -ω^2 * u
+end
+
+prob = SecondOrderODEProblem(harmonicoscillator, dx₀, x₀, tspan, ω)
+sol = solve(prob, DPRKN6())
+@test sol(0.5) isa RecursiveArrayTools.ArrayPartition
+@test sol(0.5; idxs = 1) isa Number
+@test sol(0.5; idxs = [1]) isa Vector
+@test sol(0.5; idxs = [1,2]) isa Vector
+@test Vector(sol(0.5)) == sol(0.5; idxs = [1,2])
+@test reverse(Vector(sol(0.5))) == sol(0.5; idxs = [2,1])
+@test Vector(sol(0.5)) == [sol(0.5; idxs = 1); sol(0.5; idxs = 2)]