Support prob_func in ensembles

[ChrisRackauckas]

Works:

from diffeqpy import de

def f(u,p,t):
    x, y, z = u
    sigma, rho, beta = p
    return [sigma * (y - x), x * (rho - z) - y, x * y - beta * z]

u0 = [1.0,0.0,0.0]
tspan = (0., 100.)
p = [10.0,28.0,8/3]
prob = de.ODEProblem(f, u0, tspan, p)
fast_prob = de.jit(prob)
sol = de.solve(fast_prob,saveat=0.01)

import random
def prob_func(prob,i,rep):
  de.remake(prob,u0=[random.uniform(0, 1)*u0[i] for i in range(0,3)],
            p=[random.uniform(0, 1)*p[i] for i in range(0,3)])

ensembleprob = de.EnsembleProblem(fast_prob, safetycopy=False)
sol = de.solve(ensembleprob,de.Tsit5(),de.EnsembleSerial(),trajectories=10000,saveat=0.01)

Fails:

ensembleprob = de.EnsembleProblem(fast_prob, prob_func = prob_func, safetycopy=False)
sol = de.solve(ensembleprob,de.Tsit5(),de.EnsembleSerial(),trajectories=10000,saveat=0.01)

[jodemaey]

Works but then is not perturbing the initial problem... I fell on this while looking for a method to generate ensemble on GPUs.
I tried the example in the README and got troubled by the results.

Example

from diffeqpy import de
import random

def f(u,p,t):
    x, y, z = u
    sigma, rho, beta = p
    return [sigma * (y - x), x * (rho - z) - y, x * y - beta * z]

u0 = [1.0,0.0,0.0]
tspan = (0., 100.)
p = [10.0,28.0,8/3]
prob = de.ODEProblem(f, u0, tspan, p)
fast_prob = de.jit32(prob)

sol = de.solve(fast_prob,saveat=0.01)

def prob_func(prob,i,rep):
    de.remake(prob,u0=[random.uniform(0, 1)*u0[i] for i in range(0,3)],
            p=[random.uniform(0, 1)*p[i] for i in range(0,3)])
    
ensembleprob = de.EnsembleProblem(fast_prob, safetycopy=False)

sol = de.solve(ensembleprob,de.Tsit5(),de.EnsembleSerial(),trajectories=10,saveat=0.01)

then:

and

so all the members are the same because they are all given the same problem (default prob_func does that).
So you need to pass a prob_func which does return a new problem for each member.

How to make it works

This works for me:

from diffeqpy import de
import random

def f(u,p,t):
    x, y, z = u
    sigma, rho, beta = p
    return [sigma * (y - x), x * (rho - z) - y, x * y - beta * z]

u0 = [1.0,0.0,0.0]
tspan = (0., 100.)
p = [10.0,28.0,8/3]
prob = de.ODEProblem(f, u0, tspan, p)
fast_prob = de.jit32(prob)

sol = de.solve(fast_prob,saveat=0.01)

def prob_func(prob,i,rep):
    new_prob = de.remake(prob,u0=[random.uniform(0, 1)*u0[i] for i in range(0,3)],
            p=[random.uniform(0, 1)*p[i] for i in range(0,3)])
    return new_prob

ensembleprob = de.EnsembleProblem(fast_prob, prob_func=prob_func, safetycopy=False)

sol = de.solve(ensembleprob,de.Tsit5(),de.EnsembleSerial(),trajectories=10,saveat=0.01)

then I get :

and

which are indeed different ensemble members. Bottom line: I think you should update the README file.

[jodemaey]

Please also note that when prob_func is correctly set, the GPU computations:

sol = de.solve(ensembleprob,cuda.GPUTsit5(),cuda.EnsembleGPUKernel(cuda.CUDABackend()),trajectories=10000,saveat=0.01)

or

sol = de.solve(ensembleprob,de.Tsit5(),cuda.EnsembleGPUArray(cuda.CUDABackend()),trajectories=10000,saveat=0.01)

does not work !

I get for the first one:

---------------------------------------------------------------------------
JuliaError                                Traceback (most recent call last)
Cell In[3], line 1
----> 1 sol = de.solve(ensembleprob,cuda.GPUTsit5(),cuda.EnsembleGPUKernel(cuda.CUDABackend()),trajectories=1000,saveat=0.01)

File ~/.julia/packages/PythonCall/wXfah/src/jlwrap/any.jl:208, in __call__(self, *args, **kwargs)
    206     return ValueBase.__dir__(self) + self._jl_callmethod($(pyjl_methodnum(pyjlany_dir)))
    207 def __call__(self, *args, **kwargs):
--> 208     return self._jl_callmethod($(pyjl_methodnum(pyjlany_call)), args, kwargs)
    209 def __bool__(self):
    210     return True

JuliaError: InvalidIRError: compiling MethodInstance for DiffEqGPU.gpu_ode_asolve_kernel(::KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicCheck, Nothing, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}, KernelAbstractions.NDIteration.NDRange{1, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}}}, ::CUDA.CuDeviceVector{DiffEqGPU.ImmutableODEProblem{StaticArraysCore.SVector{3, Float64}, Tuple{Float32, Float32}, true, PyList{Any}, SciMLBase.ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.var"#k#551"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xaab04274, 0x5d0ee958, 0xf8c8e45b, 0xccc3d5f8, 0x3a7c9cab), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x0d447c86, 0xa5ce6e0e, 0x3688e1ff, 0x84dbc3b4, 0xdf6fd07e), Nothing}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Vector{Symbol}, Symbol, Vector{Symbol}, ModelingToolkit.var"#678#generated_observed#561"{Bool, ModelingToolkit.ODESystem, Dict{Any, Any}, Vector{SymbolicUtils.BasicSymbolic{Real}}}, Nothing, ModelingToolkit.ODESystem}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}, 1}, ::DiffEqGPU.GPUTsit5, ::CUDA.CuDeviceMatrix{StaticArraysCore.SVector{3, Float32}, 1}, ::CUDA.CuDeviceMatrix{Float32, 1}, ::Float32, ::SciMLBase.CallbackSet{Tuple{}, Tuple{}}, ::Nothing, ::Float32, ::Float32, ::StepRangeLen{Float32, Float64, Float64, Int64}, ::Val{false}) resulted in invalid LLVM IR
Reason: unsupported call to an unknown function (call to julia.new_gc_frame)
Stacktrace:
  [1] getproperty
    @ ~/.julia/packages/PythonCall/wXfah/src/Py.jl:261
  [2] macro expansion
    @ ~/.julia/packages/PythonCall/wXfah/src/convert.jl:356
  [3] macro expansion
    @ ~/.julia/packages/PythonCall/wXfah/src/Py.jl:131
...

and

--------------------------------------------------------------------------
JuliaError                                Traceback (most recent call last)
Cell In[4], line 1
----> 1 sol = de.solve(ensembleprob,de.Tsit5(),cuda.EnsembleGPUArray(cuda.CUDABackend()),trajectories=1000,saveat=0.01)

File ~/.julia/packages/PythonCall/wXfah/src/jlwrap/any.jl:208, in __call__(self, *args, **kwargs)
    206     return ValueBase.__dir__(self) + self._jl_callmethod($(pyjl_methodnum(pyjlany_dir)))
    207 def __call__(self, *args, **kwargs):
--> 208     return self._jl_callmethod($(pyjl_methodnum(pyjlany_call)), args, kwargs)
    209 def __bool__(self):
    210     return True

JuliaError: CuArray only supports element types that are allocated inline.
Any is not allocated inline

Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] check_eltype(T::Type)
    @ CUDA ~/.julia/packages/CUDA/rXson/src/array.jl:51
...

for the second.

Note that if I use

ensembleprob = de.EnsembleProblem(fast_prob, safetycopy=False)

like you did, then GPU works, but again all the members are the same.

[ChrisRackauckas]

@LilithHafner can you look into this one?

[ChrisRackauckas]

@utkarsh530 don't the prob_funcs resolve before the kernel? The error looks like it's trying to use it in the kernel.

[LilithHafner]

I don't have a cuda machine to reproduce @jodemaey's latest example, but fwiw, when I run

sol = de.solve(ensembleprob,metal.GPUTsit5(),metal.EnsembleGPUKernel(metal.MetalBackend()),trajectories=100,saveat=0.01)

I get

zsh: segmentation fault  python3

If the metal issue and the cuda issue are related it is probably best to debug the cuda one first as it has amore actionable error message.