Sensitivity of an SDE with non-diagonal noise

Specific Domains Modelling & Simulations
, ,
1

Hi everyone,

I’m simulating an SDE with non-diagonal noise and I’m trying to compute the gradient (using ForwardDiff) of a loss function that uses the SDE solution. When I try to compute the gradient the following error happens:

    nested task error: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{typeof(loss_function), Float64}, Float64, 6})
    The type `Float64` exists, but no method is defined for this combination of argument types when trying to construct it.
    
    Closest candidates are:
      (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat
       @ Base rounding.jl:265
      (::Type{T})(::T) where T<:Number
       @ Core boot.jl:900
      Float64(::IrrationalConstants.Quartπ)
       @ IrrationalConstants ~/.julia/packages/IrrationalConstants/lWTip/src/macro.jl:131

   Stacktrace:
      [1] convert(::Type{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{typeof(loss_function), Float64}, Float64, 6})
        @ Base ./number.jl:7
      [2] setindex!
        @ ./array.jl:994 [inlined]
      [3] diffusion!(du::Matrix{Float64}, u::Vector{ForwardDiff.Dual{…}}, p::Vector{ForwardDiff.Dual{…}}, t::Float64)

Which indicates that there’s an error when trying to convert the diffusion du.

The code I’m using is as follows:


using DifferentialEquations, SciMLSensitivity, ForwardDiff

# Define drift function
function drift!(du, u, p, t)
    du[1] = p[1] * u[1]
    du[2] = p[2] * u[2]
end

# Define diffusion function with non-diagonal noise
function diffusion!(du, u, p, t)
    du[1, 1] = p[3] * u[1]
    du[1, 2] = p[4] * u[2]
    du[2, 1] = p[5] * u[1]
    du[2, 2] = p[6] * u[2]
end

# Initial conditions and parameters
u0    = [1.0, 1.0]
p     = [0.5, -0.3, 0.1, 0.2, 0.3, 0.1]
tspan = (0.0, 1.0)

noise_prototype = zeros(eltype(p), 2, 2)

# Define SDE problem
sde_prob = SDEProblem(drift!, diffusion!, u0, tspan, p, noise_rate_prototype=noise_prototype)

# Define the number of ensemble realizations
num_trajectories = 100

# Function to compute loss for a given parameter set
function loss_function(p)
    prob = remake(sde_prob, p=p)
    
    # Solve ensemble of SDEs
    ensemble_prob = EnsembleProblem(prob)
    ensemble_sol = solve(ensemble_prob, SRA1(), EnsembleThreads(), trajectories=num_trajectories, saveat=0.1)

    return sum(sum(sol[1, :] for sol in ensemble_sol.u))
end

# Compute gradients using ForwardDiff
grad = ForwardDiff.gradient(loss_function, p)
println("Gradients: ", grad)

Any ideas on how to solve this issue?

Thanks in advance!

2

Your noise rate prototype matrix isn’t converted. We should do that automatically in the library, but until we do this should work:

function loss_function(p)
    noise_prototype = zeros(eltype(p), 2, 2)
    sde_prob = SDEProblem(drift!, diffusion!, u0, tspan, p, noise_rate_prototype=noise_prototype)
    
    # Solve ensemble of SDEs
    ensemble_prob = EnsembleProblem(prob)
    ensemble_sol = solve(ensemble_prob, SRA1(), EnsembleThreads(), trajectories=num_trajectories, saveat=0.1)

    return sum(sum(sol[1, :] for sol in ensemble_sol.u))
end
3

Amazing, thank you very much, Chris!