standalone geodesic EnsembleProblem that worked fine

ensemble_heisenbug.jl

using DifferentialEquations, LinearAlgebra, ThreadsX, DiffEqGPU, StaticArrays, ThreadsX, BenchmarkTools

struct DnegWormhole
    ρ::Float64
    a::Float64
    M::Float64
end

const DEFAULT_WORMHOLE = Ref{DnegWormhole}(DnegWormhole(1.0, 0.5, 0.5))
default_wormhole() = DEFAULT_WORMHOLE[]

ξ(ℓ::Real, W::DnegWormhole) = 2*(abs(ℓ) - W.a)/(π*W.M)

function r(ℓ::Real, W::DnegWormhole)
    if abs(ℓ) < W.a
        W.ρ
    else
        ξ̂ = ξ(ℓ, W)
        W.ρ + W.M*(ξ̂*atan(ξ̂) - log(1.0 + ξ̂^2)/2.0)
    end
end
r(x::AbstractVector, W::DnegWormhole) = r(x[2], W)

function H(x, p, W::DnegWormhole=defalt_wormhole())
    r̂ = r(x, W)
    h = -1.0 + p[2]^2 + p[3]^2/r̂^2 + p[4]^2 / (r̂*sin(x[3]))^2
    h/2.0
end

function randmomentum()
    @SVector [-1.0, rand(), rand(), rand()]
end
function randposition()
    @SVector [0.0, rand(), π/2, 0.0]
end

# pick simple initial conditions so I can easily tell if it gives the right answer
function makeproblem(Δt::Real=10.0, p₀::AbstractVector=randmomentum(),
                     q₀::AbstractVector=randposition(), W::DnegWormhole=default_wormhole())
    HamiltonianProblem(p₀, q₀, Float32(Δt), SciMLBase.NullParameters()) do p′, q′, Ξ, t
        H(q′, p′, W)
    end
end

generateinit(𝔭) = remake(𝔭, u0=ArrayPartition(randmomentum(), randposition()))

function control(N::Integer=10^4, Δt::Real=10.0)
    ThreadsX.map(1:N) do _
        𝔭 = makeproblem(Δt, randmomentum(), randposition())
        solve(𝔭)
    end
end

function makeensemble(Δt::Real=10.0)
    h = makeproblem(Δt)
    EnsembleProblem(h, safetycopy=false,
                    prob_func=(𝔭, i, _) -> generateinit(𝔭),
                    output_func=(s, i) -> (last(s.u), false),
                   )
end