Skip to content

Second attempt at issue #995 #1385

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

Open
arn97 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Second attempt at issue #995 #1385

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
48 changes: 42 additions & 6 deletions src/gauss_adjoint.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -457,8 +457,17 @@ function GaussIntegrand(sol, sensealg, checkpoints, dgdp = nothing)
u0 = state_values(prob)
p = parameter_values(prob)

if p === nothing || p isa SciMLBase.NullParameters
tunables, repack = p, identity
# if p === nothing || p isa SciMLBase.NullParameters
# tunables, repack = p, identity
# ----------------------------------------------
# fix 0:
if p === nothing
tunables = SciMLBase.NullParameters()
repack = _ -> nothing
elseif p isa SciMLBase.NullParameters
tunables = p
repack = identity
# ----------------------------------------------
elseif isscimlstructure(p)
tunables, repack, _ = canonicalize(Tunable(), p)
elseif isfunctor(p)
Expand All @@ -477,7 +486,11 @@ function GaussIntegrand(sol, sensealg, checkpoints, dgdp = nothing)
)
end

numparams = length(tunables)
#--------------------------------------------
# fix 1:
numparams = (tunables === nothing || tunables isa SciMLBase.NullParameters) ? 0 : length(tunables)
#--------------------------------------------
# numparams = length(tunables)
y = zero(state_values(prob))
λ = zero(state_values(prob))
# we need to alias `y`
Expand Down Expand Up @@ -558,6 +571,12 @@ end

# out = λ df(u, p, t)/dp at u=y, p=p, t=t
function vec_pjac!(out, λ, y, t, S::GaussIntegrand)
#---------------------------------------------
# fix 5:
if S.tunables isa SciMLBase.NullParameters
return out # should already be length-0 / zeros
end
#---------------------------------------------
(; pJ, pf, p, f_cache, dgdp_cache, paramjac_config, sensealg, sol, tunables, repack) = S
_odef = sol.prob.f
f = unwrapped_f(_odef)
Expand Down Expand Up @@ -695,8 +714,17 @@ function _adjoint_sensitivities(
throw(SciMLStructuresCompatibilityError())
end

if p === nothing || p isa SciMLBase.NullParameters
tunables, repack = p, identity
# if p === nothing || p isa SciMLBase.NullParameters
# tunables, repack = p, identity
# ----------------------------------------------
# fix 0:
if p === nothing
tunables = SciMLBase.NullParameters()
repack = _ -> nothing
elseif p isa SciMLBase.NullParameters
tunables = p
repack = identity
# ----------------------------------------------
elseif isscimlstructure(p)
tunables, repack, _ = canonicalize(Tunable(), p)
elseif isfunctor(p)
Expand All @@ -714,15 +742,23 @@ function _adjoint_sensitivities(
integrand = GaussIntegrand(sol, sensealg, checkpoints, dgdp_continuous)
integrand_values = IntegrandValuesSum(allocate_zeros(tunables))
if sensealg isa GaussAdjoint
# cb = IntegratingSumCallback(
# (out, u, t, integrator) -> integrand(out, t, u),
# integrand_values, allocate_vjp(tunables)
# )
#----------------------------------------------
# fix 4:
cb = IntegratingSumCallback(
(out, u, t, integrator) -> integrand(out, t, u),
(out, t, integrator) -> integrand(out, t, integrator.u),
integrand_values, allocate_vjp(tunables)
)
#----------------------------------------------
elseif sensealg isa GaussKronrodAdjoint
cb = IntegratingGKSumCallback(
(out, u, t, integrator) -> integrand(out, t, u),
integrand_values, allocate_vjp(tunables)
)

end
rcb = nothing
cb2 = nothing
Expand Down
16 changes: 16 additions & 0 deletions src/parameters_handling.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,13 @@ function allocate_vjp(λ::AbstractArray, x::NamedTuple{F}) where {F}
end
allocate_vjp(λ::AbstractArray, x) = fmap(Base.Fix1(allocate_vjp, λ), x)

# ---------------------------------------------
# fix 3: make allocate_vjp safe on "no params"
allocate_vjp(x::Nothing) = nothing
allocate_vjp(x::SciMLBase.NullParameters) = nothing
allocate_vjp(::AbstractArray, ::Nothing) = nothing
allocate_vjp(::AbstractArray, ::SciMLBase.NullParameters) = nothing
# ---------------------------------------------
allocate_vjp(x::AbstractArray) = zero(x) # similar(x)
allocate_vjp(x::Tuple) = allocate_vjp.(x)
allocate_vjp(x::NamedTuple{F}) where {F} = NamedTuple{F}(allocate_vjp.(values(x)))
Expand All @@ -83,6 +90,11 @@ allocate_vjp(x) = fmap(allocate_vjp, x)

`zero.(x)` for generic `x`. This is used to handle non-array parameters!
"""
#---------------------------------------------
# fix 2:
allocate_zeros(::Nothing) = nothing
allocate_zeros(::SciMLBase.NullParameters) = SciMLBase.NullParameters()
#---------------------------------------------
allocate_zeros(x::AbstractArray) = zero.(x)
allocate_zeros(x::Tuple) = allocate_zeros.(x)
allocate_zeros(x::NamedTuple{F}) where {F} = NamedTuple{F}(allocate_zeros.(values(x)))
Expand All @@ -93,6 +105,10 @@ allocate_zeros(x) = fmap(allocate_zeros, x)

`adjoint(y)` for generic `y`. This is used to handle non-array parameters!
"""
#---------------------------------------------
# fix 3: make recursive_adjoint safe on "no params"
recursive_adjoint(::Nothing) = nothing
#---------------------------------------------
recursive_adjoint(y::AbstractArray) = adjoint(y)
recursive_adjoint(y::Tuple) = recursive_adjoint.(y)
recursive_adjoint(y::NamedTuple{F}) where {F} = NamedTuple{F}(recursive_adjoint.(values(y)))
Expand Down