refactor: ODESolution dispatches to avoid piracy

Author: sathvikbhagavan

src/NeuralPDE.jl

@@ -29,7 +29,7 @@ using Reexport: @reexport
 using RuntimeGeneratedFunctions: RuntimeGeneratedFunctions, @RuntimeGeneratedFunction
 using SciMLBase: SciMLBase, BatchIntegralFunction, IntegralProblem, NoiseProblem,
                  OptimizationFunction, OptimizationProblem, ReturnCode, discretize,
-                 isinplace, solve, symbolic_discretize, ODEProblem, AbstractODESolution
+                 isinplace, solve, symbolic_discretize, ODEProblem, ODESolution
 using Statistics: Statistics, mean
 using QuasiMonteCarlo: QuasiMonteCarlo, LatinHypercubeSample
 using WeightInitializers: glorot_uniform, zeros32

src/ode_solve.jl

@@ -281,6 +281,12 @@ function (f::NNODEInterpolation)(t::Vector, idxs, ::Type{Val{0}}, p, continuity)
     return DiffEqArray([out[idxs, i] for i in axes(out, 2)], t)
 end
 
+function (sol::ODESolution{T, N, U, U2, D, T2, R, D2, P, A})(
+        t::AbstractVector{<:Number}, ::Type{deriv}, idxs::Nothing,
+        continuity) where {T, N, U, U2, D, T2, R, D2, P, A <: NNODE, deriv}
+    sol.interp(t, idxs, deriv, sol.prob.p, continuity)
+end
+
 SciMLBase.interp_summary(::NNODEInterpolation) = "Trained neural network interpolation"
 SciMLBase.allowscomplex(::NNODE) = true
 

src/pino_ode_solve.jl

@@ -292,9 +292,15 @@ end
 SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation"
 SciMLBase.allowscomplex(::PINOODE) = true
 
-function (sol::AbstractODESolution)(
+function (sol::ODESolution{T, N, U, U2, D, T2, R, D2, P, A})(
         t::AbstractArray, ::Type{deriv}, idxs::Nothing,
-        continuity) where {deriv}
+        continuity) where {T, N, U, U2, D, T2, R, D2, P, A <: PINOODE, deriv}
+    sol.interp(t, idxs, deriv, sol.prob.p, continuity)
+end
+
+function (sol::ODESolution{T, N, U, U2, D, T2, R, D2, P, A})(
+        t::AbstractVector{<:Number}, ::Type{deriv}, idxs::Nothing,
+        continuity) where {T, N, U, U2, D, T2, R, D2, P, A <: PINOODE, deriv}
     sol.interp(t, idxs, deriv, sol.prob.p, continuity)
 end