I have a loss function such as below:
function loss(params)
sol_ = solve(prob_seird, Tsit5(),p=params, abstol = 1e-8, reltol = 1e-6,
saveat=tsteps)
sol = convert(Array,VectorOfArray(sol_.u))
l1 = sum((diff(sol[:,1,:],dims=2)-epi_array[:,1,2:end]).^2)
l2 = sum((diff(sol[:,10,:],dims=2)-epi_array[:,2,2:end]).^2)
println(size(l1))
return l1+l2
end
where sol_.u is Array{Array{Float64,2},1} and epic_array is Array{Float32,3}. I use convert to turn sol_.u into Array{Float32,3}.
Solution is from:
function SEIRD_mobility_coupled_outer(mobility_, coupling_matrix_, nn_)
function SEIRD_mobility_coupled_inner(du, u, p_, t)
s =@view u[:,1]
e =@view u[:,2]
id1 =@view u[:,3]
id2 =@view u[:,4]
id3 =@view u[:,5]
id4 =@view u[:,6]
id5 =@view u[:,7]
id6 =@view u[:,8]
id7 =@view u[:,9]
d =@view u[:,10]
ir1 =@view u[:,11]
ir2 =@view u[:,12]
ir3 =@view u[:,13]
ir4 =@view u[:,14]
ir5 =@view u[:,15]
r =@view u[:,16]
ds =@view du[:,1]
de =@view du[:,2]
did1 =@view du[:,3]
did2 =@view du[:,4]
did3 =@view du[:,5]
did4 =@view du[:,6]
did5 =@view du[:,7]
did6 =@view du[:,8]
did7 =@view du[:,9]
dd =@view du[:,10]
dir1 =@view du[:,11]
dir2 =@view du[:,12]
dir3 =@view du[:,13]
dir4 =@view du[:,14]
dir5 =@view du[:,15]
dr =@view du[:,16]
ฮบ, ฮฑ, ฮณ = softplus.(p_[1:3])
# ฮบ*ฮฑ and ฮณ*ฮท are not independent. The probablibility of transition from e to Ir and Id has to add up to 1
ฮท = - log(-expm1(-ฮบ*ฮฑ))/(ฮณ+1.0e-8)
n_c = size(coupling_matrix_,1)
scaler_ = softplus.(p_[4:3+n_c])
cm_ = scaler_ .* coupling_matrix_[:,:,Int32(round(t+1.0))]
p_nnet = p_[4+n_c:end]
ฮฒ = nn_(mobility_[:,:,Int32(round(t+1.0))], p_nnet)[1,:]
i = id1+id2+id3+ir1+ir2+ir3+ir4+ir5
c1 = (reshape(i,1,:)*transpose(cm_))[1,:]
c2 = cm_*i
a = ฮฒ .* s .* i + ฮฒ .* s .* (c1+c2)
@. ds = -a
@. de = a - ฮบ*ฮฑ*e - ฮณ*ฮท*e
@. did1 = ฮบ*(ฮฑ*e-id1)
@. did2 = ฮบ*(id1-id2)
@. did3 = ฮบ*(id2-id3)
@. did4 = ฮบ*(id3-id4)
@. did5 = ฮบ*(id4-id5)
@. did6 = ฮบ*(id5-id6)
@. did7 = ฮบ*(id6-id7)
@. dd = ฮบ*id7
@. dir1 = ฮณ*(ฮท*e-ir1)
@. dir2 = ฮณ*(ir1-ir2)
@. dir3 = ฮณ*(ir2-ir3)
@. dir4 = ฮณ*(ir3-ir4)
@. dir5 = ฮณ*(ir4-ir5)
@. dr = ฮณ*ir5
end
end
where each component i.e. ds, de, dir1 etc has a dimension of 104. Initial conditions are constructed:
using RecursiveArrayTools
ifr = 0.007
n_counties = size(coupling_matrix,1)
n = ones(n_counties)
ic0 = epi_array[1,:,1]
d0 = epi_array[1,:,2]
r0 = d0/ifr
s0 = n-ic0-r0-d0
e0 = ic0
id10=id20=id30=id40=id50=id60=id70=ic0*ifr/7.0
ir10=ir20=ir30=ir40=ir50=ic0*(1.0-ifr)/5.0
u0 = [s0,
e0,
id10,id20,id30,id40,id50, id60, id70, d0,
ir10,ir20,ir30,ir40,ir50,r0];
u0 = VectorOfArray([s0,
e0,
id10,id20,id30,id40,id50, id60, id70, d0,
ir10,ir20,ir30,ir40,ir50, r0]);
u0 = convert(Array,u0);
When I try to calculate the gradient
Zygote.gradient(loss,p_init)
I am getting:
Need an adjoint for constructor VectorOfArray{Float64,3,Array{Array{Float64,2},1}}