Julia DifferentialEquations.jl: EnsembleProblem where tspan is varied

I am currently working on some ensemble problems where I vary quite a few parameters of an ODE. The ODE in code is given by sys! below. It certainly is quite easy to vary the parameters, w, plateu_cycle, px and py using the Parallel Ensemble Interface (which even is doable on GPU). However I also need to vary tspan = [tstart, tend], specifically the tend. How can I do this? It seems that the EnsembleProblem does not accept variation of tspan. The idea will be to paralelly execute all Ensemble Problems on GPU using the DiffEqGPU.jl library.

How can I vary tspan with parallel execution capability? Is there a specific problem type that I might be missing? Written code:

using DifferentialEquations, PhysicalConstants.CODATA2018

e::Float64 = ElementaryCharge.val
function E_0(w::Float64)
    return w/e
end

function E(px::Float64, py::Float64)
    return sqrt(1+px^2+py^2)
end

function Vectorpotential(t, w::Float64, plateu_cycle::Float64)
    return -E_0(w)/w * sin.(t) * sineSquaredWindow(t,w,plateu_cycle)
end

function sineSquaredWindow(t, w::Float64, plateu_cycle::Float64)
    T = 2pi/w
    if t <= 0
        return 0
    elseif t/w >= 0 && (t/w <= T/2)
        return sin(t/2)^2
    elseif (t / w >= T / 2) && (t / w <= T * (plateu_cycle + 1 / 2))
        return 1
    elseif (t / w >= T * (plateu_cycle + 1 / 2)) & (t / w <= T * (plateu_cycle + 1))
        return 1-sin(w/2 * (t/w - T*(plateu_cycle+1/2)))^2
    end

    return 0
end

function sys!(du, u, params, t)
    w, plateu_cycle, px, py = params
    Energy = E(px,py)
    function nu(t, plateu_cycle::Float64, w::Float64, px::Float64, py::Float64)
        return -1im * e * Vectorpotential(w*t,w,plateu_cycle)*exp(2im*Energy*t) * ( px*py /(Energy*(Energy+1))+1im*(1-py^2 / (Energy*(Energy+1))))
    end
    function kappa(t, plateu_cycle::Float64, w::Float64, py::Float64)
        return 1im*e*Vectorpotential(w*t,w,plateu_cycle)*py/Energy
    end
    du[1] = kappa(t,plateu_cycle,w,py)*u[1] + nu(t, plateu_cycle,w, px,py) *u[2]
    du[2] = -conj(nu(t,plateu_cycle,w,px,py))*u[1] + conj(kappa(t,plateu_cycle,w, py))*u[2]
    return nothing
end 
w = 0.84::Float64
plateu_cycle = 8.0::Float64
px = 0.0
py = 0.0
params=[w, plateu_cycle, px, py]
f0=complex(0.0,0.0)
g0=complex(1.0,0.0)
init_cond = [f0, g0]
tstart = 0
tend = 2pi/w * (plateu_cycle+1)+1
tspan = (tstart, tend)
# this is just en example solve. It does not do the ensemble problems yet
ts = collect(range(tstart, tend, length=1000))
prob = ODEProblem(sys!,init_cond,tspan, params)
sol = solve(prob, Tsit5(), saveat=ts)

How can I do this? It seems that the EnsembleProblem does not accept variation of tspan.

Have you tried it yet? When you define a prob_func you choose the problem per trajectory, and since each problem has a tspan, you just give it a different tspan per trajectory. It's no different from the normal ensemble code.

I cannot modify your code for it because you didn't share an ensemble code, but it would just be exactly like what's shown in the doc page here. For example:

function prob_func(prob,i,repeat)
  remake(prob, tspan = (new_tstart, new_tend))
end