henhen724
commented
2 months ago I made an attempt to add support for non-diagonal noise in a fork. Ultimately, I concluded that it would require refactoring other other packages, namely StochasticDiffEq.jl, which seems like too large a change for a pull request. https://github.com/henhen724/DiffEqGPU.jl/tree/non_diagonal_array_ensembles Instead, I'll leave what I learned here.
The problem with implementing non-diagonal noise in ArrayEnsembles is that the solve functions expects to be able to matrix multiply the noise term by a vector of Wiener increments and add that to the solution, but in ArrayEnsembles the solution is a matrix and not a vector. Here is a link to an arbitrarily chosen SDE algorithm: https://github.com/SciML/StochasticDiffEq.jl/blob/0c03d8b6378f133a1f819ebc2be8f0bee5d69f06/src/perform_step/sra.jl#L144
integrator.g(g1,uprev,p,t+c11*dt)
integrator.f(k1,uprev,p,t)
if is_diagonal_noise(integrator.sol.prob)
@.. H01 = uprev + dt*a21*k1 + chi2*b21*g1
else
mul!(E₁,g1,chi2)
@.. H01 = uprev + dt*a21*k1 + b21*E₁
endintegrator.g is the noise function. It's output is stored in g1 which is then matrix multiplied by chi2 which is the vector of Wiener increaments. This means the expression b21*E is a vector. This of course makes sense for single copies of the problem.
The problem is that H01 and k1 which would both be vectors for an individual problem are nxm matrices in the ArrayEnsembles where n is the dimension of the underlying problem and m is the number of trajectories.
There is then a second larger problem which is that StochasticDiffEq.jl assumes that noise_rate_prototype is a matrix and sets the dimensions of the g1 matrix handed to integrator.g according to the dimensions of noise_rate_prototype. If noise_rate_prototype is a 3 tensor instead of matrix, then the solve function in StochasticDiffEq.jl will throw an error when trying to find the dimension for the Wiener process: https://github.com/SciML/StochasticDiffEq.jl/blob/0c03d8b6378f133a1f819ebc2be8f0bee5d69f06/src/solve.jl#L302
rand_prototype = false .* noise_rate_prototype[1,:]BoundsError: attempt to access 4×2×10 Array{Float32, 3} at index [1, 1:2]
As far as I can see, the solution would be to refactor StochasticDiffEq.jl to allow non-matrix "noise_rate_prototype"s, and then add a definition of mul! for 3 tensor to the KernelAbstractions.jl library, as well as the changes in my fork to DiffEqGPU.jl.
Describe the bug 🐞
Expected behavior The expected behavior is for the solver to finish without throwing an error and return an accurate solution.
Minimal Reproducible Example 👇
Error & Stacktrace ⚠️
Environment (please complete the following information):
using Pkg; Pkg.status()using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)versioninfo()Additional context This error is caused by the assumption in the file src > ensemblegpuarray > kernels.jl that the time series for du can be written as a matrix (with one index for ODE coordinate and the next index for time). When a problem has non-diagonal noise you need two coordinate indexes and a time index, so the du time series needs to be a three index tensor or include some flatten and resize adaptor when evaluating the noise function.