add sol(t)

SciML · Oct 28, 2024 · e8ac7f5 · e8ac7f5
1 parent 501496e
commit e8ac7f5
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/pino_ode_solve.jl b/src/pino_ode_solve.jl
@@ -216,7 +216,7 @@ function generate_loss(
     end
 end
 
-struct PINOODEInterpolation{T <: PINOPhi, T2}
+@concrete struct PINOODEInterpolation{T <: PINOPhi, T2}
     phi::T
     θ::T2
 end
@@ -226,6 +226,10 @@ end
 SciMLBase.interp_summary(::PINOODEInterpolation) = "Trained neural network interpolation"
 SciMLBase.allowscomplex(::PINOODE) = true
 
+function (sol::SciMLBase.AbstractODESolution)(t::AbstractArray)
+    sol.interp(t)
+end
+
 function SciMLBase.__solve(prob::SciMLBase.AbstractODEProblem,
         alg::PINOODE,
         args...;