I am solving a system of equations with the iterative proportional fitting procedure. That means that at each iteration I solve for a subset of the variables μ0m, then hold those fixed and solve for the other subset of variables μ0f. I am storing the residuals of the system in preallocated arrays resm and resf. I programmed a function sys! that computes the residuals of the system for each of those subsets:
const szA = 9
const szE = 2
const szG = 2
function sys!(resm, resf, μ0m, μ0f; Sm=Sm, Sf=Sf, Π=Π)
μ0 = reshape(vcat(vec(μ0m), vec(μ0f)), (szA, szE, szG))
μ0m′ = μ0[:, :, 1]
μ0f′ = μ0[:, :, 2]
μ0m′[1, szE] = 0.0
μ0f′[1, szE] = 0.0
prodsummand = similar(μ0, (szA, szE, szA, szE))
prodsummand!(prodsummand, μ0m′, μ0f′; Sm=Sm, Sf=Sf, Π=Π)
for idx in CartesianIndices(resm)
am, em = Tuple(idx)
resm[am, em] = Sm[am, em] - abs(μ0m′[am,em]) - sum( prodsummand[am,em,af,ef] for af in 1:szA, ef in 1:szE )
end
for idx in CartesianIndices(resf)
af, ef = Tuple(idx)
resf[af, ef] = Sf[af, ef] - abs(μ0f′[af,ef]) - sum( prodsummand[am,em,af,ef] for am in 1:szA, em in 1:szE )
end
return nothing
end
function prodsummand!(prodsum, μ0m′, μ0f′; Sm=Sm, Sf=Sf, Π=Π)
for idx in CartesianIndices(prodsum)
am, em, af, ef = Tuple(idx)
prodsum[am, em, af, ef] = Π[am,em,af,ef]*sqrt(Sm[am,em]*Sf[af,ef])*prod( auxprod((μ0m′[am+k,em]*μ0f′[af+k,ef])/(Sm[am+k,em]*Sf[af+k,ef]))^(0.5*(β*(1-δ))^k) for k in 0:(maxT(am,af; Z=szA)-1) )
end
return prodsum
end
function maxT(am, af; Z=szA)
Z - max(am, af) + 1
end
auxprod(x) = isfinite(abs(x)) ? abs(x) : 1.0
To solve the system, I am using NLsolve with finite differences, and it finds a solution. I am now wanting to speed up the computation by using automatic differentiation. I am having trouble specifying the Jacobian with AD. This is the algorithm that solves the system with finite diff:
using NLsolve
using LinearAlgebra
function ipfp(init, sys!; Sm=Sm, Sf=Sf, Π=Π, method=:trust_region, autodiff=:finite, outmaxiter=5_000, outtol=1e-8, opts...)
μ0f0 = vec(init)
μ0m1 = similar(init)
μ0f1 = similar(init)
d = Inf
i = 0
while (i <= outmaxiter) && (outtol < d)
solm = nlsolve( (resm, μ0m) -> sys!(resm, resf, μ0m, μ0f0; Sm=Sm, Sf=Sf, Π=Π), init; method=method, autodiff=autodiff, opts... )
if converged(solm)
μ0m1 = abs.(vec(solm.zero))
else
println("solving for μ0m given μ0f did not converge at iteration ", i)
return nothing
end
solf = nlsolve( (resf, μ0f) -> sys!(resm, resf, μ0m1, μ0f; Sm=Sm, Sf=Sf, Π=Π), init; method=method, autodiff=autodiff, opts... )
if converged(solf)
μ0f1 = abs.(vec(solf.zero))
else
println("solving for μ0f given μ0m did not converge at iteration ", i)
return nothing
end
d = norm(μ0f1 - μ0f0)
μ0f0 = copy(μ0f1)
i += 1
end
ret = reshape(vcat(μ0m1, μ0f1), (szA, szE, szG))
ret[1, szE, 1] = 0.0
ret[1, szE, 2] = 0.0
return ret
end
and here are some objects that complete a working example:
Sm = rand(szA, szE)
Sf = rand(szA, szE)
Π = rand(szA, szE, szA, szE)
resm = similar(Sm)
resf = similar(Sf)
init = rand(szA, szE)
I tried adding AD-able functions to the algorithm like
lfunm! = (resm, μ0m) -> sys!(resm, resf, μ0m, μ0f0; Sm=Sm, Sf=Sf, Π=Π)
sys_jacm! = ForwardDiff.jacobian!(lfunm!, resm, init)
but I get an error:
ERROR: MethodError: no method matching jacobian!(::var"#94#96"{Matrix{Float64}, Matrix{Float64}, Array{Float64, 4}, typeof(sys!)}, ::Matrix{Float64}, ::Matrix{Float64})
How can I specify the (partial) Jacobian of sys! using AD to pass it to NLsolve?