Closed evolbio closed 2 years ago
@frankschae can you look at this one/?
Yes, I can have a look. @evolbio can you post an MWE? I am a bit confused by the stack trace. Are you differentiating the solve call (did you set a sensealg
?)?
@frankschae: I apologize, I do not yet have an MWE, which I realize would help a lot. For now, I think the following code pieces show most of the key points. In essence, I am using Optimization.jl solve, which requires setting up an OptimizationProblem based on an OptimizationFunction. If you have any questions, let me know. If this is not sufficient, I will try to put together a MWE.
I think the main point is that the error arises from line 60 of implicit_split_step.jl. I generated the same error when using ForwardDiff in another case, so I don't think it is Zygote, but rather something related to differentiating through ISSEM. I can't get any other SDE solver to work either but the issues seem to differ. Previously I was successfully using ISSEM for very similar code, so maybe some recent changes?
result = solve(opt_prob(p,S,L), ADAM(S.adm_learn), callback = callback,
maxiters=S.max_it)
opt_func(S,L) = OptimizationFunction((p,u) -> loss(p,S,L), Optimization.AutoZygote())
opt_prob(p,S,L) = OptimizationProblem(opt_func(S,L), p, L.u0)
function loss(p, S, L)
G = S.f_data(S; init_on=L.init_on, rand_offset=L.rand_offset, noise_wait=L.noise_wait)
diffeq = (u, p, t) -> node(u, p, t, S, L.tf, L.re, L.state, G)
prob = SDEProblem(diffeq, ode_noise, L.u0, L.prob.tspan, p,
reltol = S.rtol, abstol = S.atol,
callback=nothing, saveat = values(L.prob.kwargs).saveat,
maxiters=1e6)
pred_all = L.predict(p, prob)
# pick out tracked protein, first protein in system output
pred = @view pred_all[S.n+1,:]
pred_length = length(pred)
input = @view G.input_true[1:pred_length]
hill_input = hill.(0.5,L.hill_k,input)
hill_pred = hill.(S.switch_level,L.hill_k,pred)
@assert pred_length == length(L.w)
loss = sum(abs2, L.w .* (hill_input .- hill_pred))
return loss, S, L, G, pred_all
end
sqrt_abs(x) = (x > 16.) ? sqrt(x) : abs(x) / 4.
ode_noise(u, p, t) = sqrt_abs.(u)
# this is condensed but has main points, solver is ISSEM()
L.predict(p,prob) = solve(prob, S.solver, u0=u_init, p=p[S.ddim+1:end])
function node(u, p, t, S, tf, re, state, G)
n = S.n
u_m = @view u[1:n] # mRNA
u_p = @view u[n+1:2n] # protein
p_rates = @view p[1:4n]
p_nn = @view p[4n+1:end]
# first 4n parameters are rates
pp = linear_sigmoid.(p_rates, S.d, S.k1, S.k2) # normalizes on [0,1] w/linear_sigmoid
# set min on rates m_a, m_d, p_a, p_d, causes top to be p_max + p_min
ppp = (pp .* S.p_mult) .+ S.p_min
m_a = @view ppp[1:n]
m_d = @view ppp[n+1:2n]
p_a = @view ppp[2n+1:3n]
p_d = @view ppp[3n+1:4n]
# remainder of p is for NN
f_val = tf(u_p,re(p_nn),state)[1]
du_m = m_a .* f_val .- m_d .* u_m # mRNA level
du_p = p_a .* u_m .- p_d .* u_p # protein level
light = G.circadian_val(G,t) # noisy circadian input
# fast extra production rate by post-translation modification or allostery
du_p_2 = du_p[2] + S.light_prod_rate * intensity(light)
return [du_m; du_p[1]; du_p_2; du_p[3:end]]
end
I am using 1.8.0-rc1 on Mac Arm. 1.7 does not work fully on Mac M1 (arm).
julia> versioninfo()
Julia Version 1.8.0-rc1
Commit 6368fdc6565 (2022-05-27 18:33 UTC)
Platform Info:
OS: macOS (arm64-apple-darwin21.3.0)
CPU: 20 × Apple M1 Ultra
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-13.0.1 (ORCJIT, apple-m1)
Threads: 6 on 16 virtual cores
Environment:
JULIA_DEPOT_PATH = /opt/julia:
JULIA_NUM_THREADS = 6
LD_LIBRARY_PATH = /Users/steve/sim/simlib/lib_osx:/opt/local/lib:/usr/local/lib:/usr/lib:/usr/X11/lib
JULIA_WARN_COLOR = red
I noticed additional problems with, for example, loading CUDA when trying 1.8.0-rc3, so maybe julia version is related to this issue??
Previously I was successfully using ISSEM for very similar code, so maybe some recent changes?
It seems to be a bug in the out-of-place code for ISSEM
with adaptive steps. The in-place formulation works for me. I'll fix it and link the PR here. Perhaps, your other code was written in the in-place form?
MWE:
using StochasticDiffEq
fIto(u, p, t) = p[1] * u
σ(u, p, t) = p[2] * u
# set initial state
u0 = [1 / 6]
tspan = (0.0, 1.0)
dt = 1e-5
# non-exploding initialization.
α = 1 / (exp(-randn()) + 1)
β = -α^2 - 1 / (exp(-randn()) + 1)
p = [α, β]
# define problem in Ito sense
probIto = SDEProblem(fIto, σ, u0, tspan, p)
# solve Ito sense
solIto = solve(probIto, ISSEM())
Thanks very much. I did have success with a different problem and in-place form with ISSEM. Not sure about all of the variations I have tried, but it seems likely that your comments are correct.
I updated StochasticDiffEq to 6.50.1 and also updated other packages to 7/24/2022 date. Prior error solved. However, I get a new error
ERROR: MethodError: no method matching vec(::Nothing)
Closest candidates are:
vec(::Tuple{Vararg{MultivariatePolynomials.AbstractVariable}}) at /opt/julia/packages/MultivariatePolynomials/1bIGc/src/operators.jl:348
vec(::StrideArraysCore.PtrArray{S, D, T, 1, C, 0}) where {S, D, T, C} at /opt/julia/packages/StrideArraysCore/VQxXL/src/reshape.jl:1
vec(::StrideArraysCore.PtrArray{S, D, T, N, C, 0}) where {S, D, T, N, C} at /opt/julia/packages/StrideArraysCore/VQxXL/src/reshape.jl:2
...
Stacktrace:
[1] _jacNoise!(λ::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}
Full stacktrace attached. Thanks. trace.txt
NOTES: My code works with ForwardDiff and also works for an ODE problem when using Zygote, so it seems to be an interaction between Zygote and the SDE form. The switch from ODE to SDE requires (1) adding a noise function, which in this case is
sqrt_abs(x) = (x > 16.) ? sqrt(x) : abs(x) / 4.
ode_noise(u, p, t) = sqrt_abs.(u)
and (2) changing from Tstit5() to ISSEM() when using SDEProblem instead of ODEProblem. The following small example works, so the problem is associated with the more extensive form used in my code, but I still have not gotten it down to a MWE.
# This succeeds, not clear what needs to be added to trigger error
using Optimization, DifferentialEquations, DiffEqSensitivity, OptimizationOptimisers
n = 3
function ode1(u, p, t)
du = p[1:n] .* u .- p[n+1:2n]
end
ode = ode1
sqrt_abs(x) = (x > 16.) ? sqrt(x) : abs(x) / 4.
ode_noise(u, p, t) = sqrt_abs.(u)
function loss(p, u0)
prob = SDEProblem(ode, ode_noise, u0, (0.,1.), p, saveat=0:0.1:1)
x = solve(prob, ISSEM(), p=p)
y = [sin(2*pi*z) for z = 0:0.1:1]
sum(abs2.(x[1,:] - y))
end
opt_func(p,u0) = OptimizationFunction((p,u0) -> loss(p,u0), Optimization.AutoZygote())
opt_prob(p,u0) = OptimizationProblem(opt_func(p,u0), p, u0)
p = randn(2n);
u0 = randn(n);
result = solve(opt_prob(p, u0), ADAM(), maxiters=50)
This might be more difficult to solve than the previous one because it's related to the sensitivity analysis tools and not only to the primal evaluation of the SDE solution.
so the problem is associated with the more extensive form used in my code, but I still have not gotten it down to a MWE.
Can you try to replace
x = solve(prob, ISSEM(), p=p)
by
x = solve(prob, ISSEM(), p=p, sensealg=BacksolveAdjoint())
I think your example might have a too small number of parameters. For small numbers of parameters, we're using tangent/forward sensitivity methods per default to compute the backward pass through the SDE solution.
I wanted to look into some SDE sensitivity issues. So I hope to get to it soon.
Using the option sensealg=BacksolveAdjoint()
in the above short example code leads to the warning:
julia> result = solve(opt_prob(p, u0), ADAM(), maxiters=50)
┌ Warning: Automatic AD choice of autojacvec failed in ODE adjoint, failing back to ODE adjoint + numerical vjp
└ @ DiffEqSensitivity /opt/julia/packages/DiffEqSensitivity/Pn9H4/src/sensitivity_interface.jl:290
I raised n=1000 to invoke internal algorithms for a larger parameter vector and then tried some other combinations. All either gave an error or were so slow that I gave up before getting any return.
Can you make a new issue to SciMLSensitivity.jl? What's left is unrelated to the SDE solver itself.
From solve
(prob, S.solver, u0=u_init, p=p[S.ddim+1:end])
for prob as SDEProblem with no special args, and solver as ISSEM().