Skip to content

Specialize more functions on f parameter for trim compatibility #2856

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 21, 2025
Merged

Specialize more functions on f parameter for trim compatibility #2856

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 10 additions & 10 deletions lib/OrdinaryDiffEqCore/src/dense/generic_dense.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -477,17 +477,17 @@ end
return expr
end

function _evaluate_interpolant(f, Θ, dt, timeseries, i₋, i₊,
function _evaluate_interpolant(f::F, Θ, dt, timeseries, i₋, i₊,
cache, idxs,
deriv, ks, ts, p, differential_vars)
deriv, ks, ts, p, differential_vars) where F
_ode_addsteps!(ks[i₊], ts[i₋], timeseries[i₋], timeseries[i₊], dt, f, p,
cache) # update the kcurrent
return ode_interpolant(Θ, dt, timeseries[i₋], timeseries[i₊], ks[i₊],
cache, idxs, deriv, differential_vars)
end
function evaluate_composite_cache(f, Θ, dt, timeseries, i₋, i₊,
function evaluate_composite_cache(f::F, Θ, dt, timeseries, i₋, i₊,
caches::Tuple{C1, C2, Vararg}, idxs,
deriv, ks, ts, p, cacheid, differential_vars) where {C1, C2}
deriv, ks, ts, p, cacheid, differential_vars) where {F, C1, C2}
if (cacheid -= 1) != 0
return evaluate_composite_cache(f, Θ, dt, timeseries, i₋, i₊, Base.tail(caches),
idxs,
Expand All @@ -497,16 +497,16 @@ function evaluate_composite_cache(f, Θ, dt, timeseries, i₋, i₊,
first(caches), idxs,
deriv, ks, ts, p, differential_vars)
end
function evaluate_composite_cache(f, Θ, dt, timeseries, i₋, i₊,
function evaluate_composite_cache(f::F, Θ, dt, timeseries, i₋, i₊,
caches::Tuple{C}, idxs,
deriv, ks, ts, p, _, differential_vars) where {C}
deriv, ks, ts, p, _, differential_vars) where {F, C}
_evaluate_interpolant(f, Θ, dt, timeseries, i₋, i₊,
only(caches), idxs,
deriv, ks, ts, p, differential_vars)
end

function evaluate_default_cache(f, Θ, dt, timeseries, i₋, i₊,
cache::DefaultCache, idxs, deriv, ks, ts, p, cacheid, differential_vars)
function evaluate_default_cache(f::F, Θ, dt, timeseries, i₋, i₊,
cache::DefaultCache, idxs, deriv, ks, ts, p, cacheid, differential_vars) where F
if cacheid == 1
return _evaluate_interpolant(f, Θ, dt, timeseries, i₋, i₊,
cache.cache1, idxs, deriv, ks, ts, p, differential_vars)
Expand All @@ -528,8 +528,8 @@ function evaluate_default_cache(f, Θ, dt, timeseries, i₋, i₊,
end
end

function evaluate_interpolant(f, Θ, dt, timeseries, i₋, i₊, cache, idxs,
deriv, ks, ts, id, p, differential_vars)
function evaluate_interpolant(f::F, Θ, dt, timeseries, i₋, i₊, cache, idxs,
deriv, ks, ts, id, p, differential_vars) where F
if isdiscretecache(cache)
return ode_interpolant(Θ, dt, timeseries[i₋], timeseries[i₊], 0, cache, idxs,
deriv, differential_vars)
Expand Down
4 changes: 2 additions & 2 deletions lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -324,7 +324,7 @@ mutable struct WOperator{IIP, T,
jacvec)
end
end
function WOperator{IIP}(f, u, gamma) where {IIP}
function WOperator{IIP}(f::F, u, gamma) where {IIP, F}
if isa(f, Union{SplitFunction, DynamicalODEFunction})
error("WOperator does not support $(typeof(f)) yet")
end
Expand Down Expand Up @@ -440,7 +440,7 @@ islinearfunction(integrator) = islinearfunction(integrator.f, integrator.alg)

return the tuple `(is_linear_wrt_odealg, islinearodefunction)`.
"""
function islinearfunction(f, alg)::Tuple{Bool, Bool}
function islinearfunction(f::F, alg)::Tuple{Bool, Bool} where F
isode = f isa ODEFunction && islinear(f.f)
islin = isode || (issplit(alg) && f isa SplitFunction && islinear(f.f1.f))
return islin, isode
Expand Down
4 changes: 2 additions & 2 deletions lib/OrdinaryDiffEqDifferentiation/src/derivative_wrappers.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ function jacobian(f::F, x::AbstractArray{<:Number}, integrator) where F
end

# fallback for scalar x, is needed for calc_J to work
function jacobian(f, x, integrator)
function jacobian(f::F, x, integrator) where F
alg = unwrap_alg(integrator, true)

dense = ADTypes.dense_ad(alg_autodiff(alg))
Expand Down Expand Up @@ -393,7 +393,7 @@ function build_grad_config(alg, f::F1, tf::F2, du1, t) where {F1, F2}
end
end

function sparsity_colorvec(f, x)
function sparsity_colorvec(f::F, x) where F
sparsity = f.sparsity

if is_sparse_csc(sparsity)
Expand Down
2 changes: 1 addition & 1 deletion lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -330,7 +330,7 @@ function compute_ustep!(ustep, tmp, γ, z, method)
ustep
end

function _compute_rhs(tmp, γ, α, tstep, invγdt, method::MethodType, p, dt, f, z)
function _compute_rhs(tmp, γ, α, tstep, invγdt, method::MethodType, p, dt, f::F, z) where F
mass_matrix = f.mass_matrix
ustep = compute_ustep(tmp, γ, z, method)
if method === COEFFICIENT_MULTISTEP
Expand Down
2 changes: 1 addition & 1 deletion lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ mutable struct DAEResidualJacobianWrapper{isAD, F, pType, duType, uType, alphaTy
tmp::tmpType
uprev::uprevType
t::tType
function DAEResidualJacobianWrapper(alg, f, p, α, invγdt, tmp, uprev, t)
function DAEResidualJacobianWrapper(alg, f::F, p, α, invγdt, tmp, uprev, t) where F
ad = ADTypes.dense_ad(alg_autodiff(alg))
isautodiff = ad isa AutoForwardDiff
if isautodiff
Expand Down
12 changes: 6 additions & 6 deletions lib/OrdinaryDiffEqSymplecticRK/src/symplectic_perform_step.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,22 +77,22 @@ end
# f.f2(p, q, pa, t) = p which is the Newton/Lagrange equations
# If called with different functions (which are possible in the Hamiltonian case)
# an exception is thrown to avoid silently calculate wrong results.
function verify_f2(f, p, q, pa, t, ::Any,
::C) where {C <: Union{HamiltonConstantCache, VerletLeapfrogConstantCache,
function verify_f2(f::F, p, q, pa, t, ::Any,
::C) where {F, C <: Union{HamiltonConstantCache, VerletLeapfrogConstantCache,
LeapfrogDriftKickDriftConstantCache}}
f(p, q, pa, t)
end
function verify_f2(f, res, p, q, pa, t, ::Any,
::C) where {C <: Union{HamiltonMutableCache, VerletLeapfrogCache,
function verify_f2(f::F, res, p, q, pa, t, ::Any,
::C) where {F, C <: Union{HamiltonMutableCache, VerletLeapfrogCache,
LeapfrogDriftKickDriftCache}}
f(res, p, q, pa, t)
end

function verify_f2(f, p, q, pa, t, integrator, ::C) where {C <: VelocityVerletConstantCache}
function verify_f2(f::F, p, q, pa, t, integrator, ::C) where {F, C <: VelocityVerletConstantCache}
res = f(p, q, pa, t)
res == p ? p : throwex(integrator)
end
function verify_f2(f, res, p, q, pa, t, integrator, ::C) where {C <: VelocityVerletCache}
function verify_f2(f::F, res, p, q, pa, t, integrator, ::C) where {F, C <: VelocityVerletCache}
f(res, p, q, pa, t)
res == p ? res : throwex(integrator)
end
Expand Down
Loading