Issue with PDMP and Forwardiff - DifferentialEquation

General Usage Performance
11

Hi,

It seems I have fixed the allocation issue in PDMP using the cache trick from DiffEqBase and a profound rewrite of the library. For scalar vector fields, autodiff seems to work fine. But then I still see the error The resulting array would have non-integral first dimension.

A MWE would be:

using PiecewiseDeterministicMarkovProcesses, Random, DifferentialEquations
const PDMP = PiecewiseDeterministicMarkovProcesses

function F_tcp!(ẋ, xc, xd, parms, t)
	if mod(xd[1],2)==0
		 ẋ[1] =  1.0
		 ẋ[2] = -1.0 * xc[2]
	else
		 ẋ[1] = -1.0 * xc[1]
		 ẋ[2] =  1.0
	end
	nothing
end

R(x) = x

function R_tcp!(rate, xc, xd, parms, t, issum::Bool)
	if issum == false
		rate[1] = R(xc[1]) +  R(xc[2])
		rate[2] = parms[1]
		return 0.
	else
		return R(xc[1]) +  R(xc[2]) + parms[1]
	end
end

xc0 = vec([0.05,0.075])
xd0 = vec([0, 1])

nu_tcp = [[1 0];[0 -1]]
parms = vec([0.1])
tf = 10000.0
nj = 1000

Random.seed!(1234)
# define a PDMP problem `a la` ODEProblem
problem = PDMP.PDMPProblem(F_tcp!, R_tcp!, nu_tcp, xc0, xd0, parms, (0.0, tf))

and then, the solve part which calls autodiff

result4 = PDMP.solve(problem, CHV(Rodas5()))
ERROR: ArgumentError: cannot reinterpret an `ForwardDiff.Dual{nothing,Float64,2}` array to `ForwardDiff.Dual{ForwardDiff.Tag{DiffEqDiffTools.UJacobianWrapper{ODEFunction{true,getfield(PiecewiseDeterministicMarkovProcesses, Symbol("##34#37")){PDMPProblem{Float64,Int64,Array{Float64,1},Array{Int64,1},PiecewiseDeterministicMarkovProcesses.PDMPCaracteristics{typeof(F_tcp!),typeof(R_tcp!),PiecewiseDeterministicMarkovProcesses.RateJump{Int64,Array{Int64,2},typeof(PiecewiseDeterministicMarkovProcesses.Delta_dummy)},Array{Float64,1},Array{Int64,1},PiecewiseDeterministicMarkovProcesses.DiffCache{Array{Float64,1},Array{ForwardDiff.Dual{nothing,Float64,2},1}},Array{Float64,1}}},CHV{Rodas5{0,true,DefaultLinSolve,DataType}}},LinearAlgebra.UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Float64,DiffEqBase.NullParameters},Float64},Float64,3}` whose first dimension has size `2`.
The resulting array would have non-integral first dimension.

I can isolate the issue a bit more by extracting the function to be differentiated:

algo = CHV(Tsit5())

# vectors to be sent to the function vf
xd1 = zeros(Float64, length(xc0)+1)
xdd1 = similar(xd1)

# for hoilding the jacobian
J = zeros(3,3)

# this is the vector field, use with an ODEProblem
vf = (dx,x) -> algo(dx, x, problem.caract, 0.)

# it works!
vf(xdd1, xd1)

ForwardDiff.jacobian!(J, vf, xdd1, xd1)

which gives

ERROR: ArgumentError: cannot reinterpret an `ForwardDiff.Dual{nothing,Float64,2}` array to `ForwardDiff.Dual{ForwardDiff.Tag{getfield(Main, Symbol("##9#10")),Float64},Float64,3}` whose first dimension has size `2`.
The resulting array would have non-integral first dimension.

I have no idea why this pops out :expressionless:

If anybody has an idea, I am all ears

Thank you for your help