Ensemble rng redesign

[isaacsas]

Ensemble-level RNG management for reproducible per-trajectory seeding

This PR adds deterministic, thread-count-independent reproducibility to ensemble solves. Passing seed or rng to an ensemble solve pre-generates per-trajectory seeds from a master RNG, then seeds each trajectory's TaskLocalRNG before solving. Results are identical regardless of thread count, batch size, or ensemble algorithm.

New kwargs on ensemble solve

• seed: Master seed (integer). Pre-generates deterministic per-trajectory seeds.
• rng: Master RNG (e.g. Xoshiro). Takes priority over seed.
• rng_func: Custom per-trajectory RNG factory (ctx::EnsembleContext) -> AbstractRNG. Defaults to default_rng_func which seeds the TaskLocalRNG.

New exported type

• EnsembleContext{S, R}: Passed to rng_func and the optional 5-arg prob_func. Fields: global_trajectory_id, worker_id, trajectory_seed, master_rng.

Usage examples

Basic reproducible ensemble solve:

using OrdinaryDiffEq

prob = ODEProblem((u, p, t) -> 1.01u, 0.5, (0.0, 1.0))
eprob = EnsembleProblem(prob;
    prob_func = (prob, i, repeat) -> remake(prob; u0 = rand() * 0.1 + 0.5))

# Same seed → identical results, regardless of thread count or ensemble algorithm
sim1 = solve(eprob, Tsit5(), EnsembleThreads(); seed = UInt64(42), trajectories = 100)
sim2 = solve(eprob, Tsit5(), EnsembleSerial();  seed = UInt64(42), trajectories = 100)
@assert [sol[end] for sol in sim1] == [sol[end] for sol in sim2]

5-arg prob_func with explicit RNG access:

eprob = EnsembleProblem(prob;
    prob_func = (prob, i, repeat, rng, ctx) -> remake(prob; u0 = rand(rng) * 0.1 + 0.5))

sim = solve(eprob, Tsit5(); seed = UInt64(42), trajectories = 50)

Custom rng_func with StableRNG:

using StableRNGs

eprob = EnsembleProblem(prob;
    prob_func = (prob, i, repeat, rng, ctx) -> remake(prob; u0 = rand(rng) * 0.1 + 0.5))

sim = solve(eprob, Tsit5();
    seed = UInt64(42),
    rng_func = ctx -> StableRNG(ctx.trajectory_seed),
    trajectories = 50)

Master RNG instead of seed:

using Random

sim = solve(eprob, Tsit5(); rng = Xoshiro(99), trajectories = 50)

Implementation details

• Should be fully non-breaking in all non-broken cases (i.e. outside of JumpProblems with safety copies disabled). I don't think this impacts DiffEqGPU but that should be reviewed (i.e. DiffEqGPU would just not support the new kwargs).
• Ensemble layer is the single source of truth for RNG: rng_func is always called to seed each trajectory. The rng kwarg is forwarded to inner solve only when supports_solve_rng(prob, alg) returns true. Non-DE solvers that don't accept rng still benefit from the already-seeded TaskLocalRNG.
• 5-arg prob_func detection: When any method of prob.prob_func has ≥5 arguments (detected via numargs), the form (prob, i, repeat, rng, ctx) is used.
• Per-task JumpProblem isolation: Replaces the vestigial threadid()-indexed deepcopy array with task_local_storage()-based per-task copies, which is safe under Julia's M:N threading model. Only active when safetycopy = false resulting in one JumpProblem deepcopy per thread.
• Distributed master_rng sanitization: master_rng is set to nothing before pmap captures the closure in EnsembleDistributed and EnsembleSplitThreads, avoiding serialization of potentially non-serializable RNG objects.
• worker_id propagation: EnsembleDistributed and EnsembleSplitThreads pass worker_id = myid() through solve_batch to batch_func, so ctx.worker_id correctly reflects the distributed worker pid.

Prerequisite PRs (should be merged and released in order before this PR for testing purposes)

1. SciMLBase, NON-BREAKING — Add supports_solve_rng trait for ensemble RNG forwarding guard #1250
2. OrdinaryDiffEq, NON-BREAKING — Add supports_solve_rng trait for ODE/DAE problem+algorithm paths OrdinaryDiffEq.jl#3081, requires release of 1
3. StochasticDiffEq, NON-BREAKING — Add integrator-level RNG support (Phase 1, Step 3) StochasticDiffEq.jl#681, requires release of 1
4. JumpProcesses, BREAKING — use integrator rng API JumpProcesses.jl#557, requires release of 1-3
5. This PR (SciMLBase ensemble_rng_redesign), NON-BREAKING: Merge after steps 2–4 are released.

[isaacsas]

This is currently marked WIP as I need to carefully review all the new tests still, and because it shouldn't be merged till the listed prerequisite PRs are handled with package releases.

[isaacsas]

@ChrisRackauckas assuming tests pass now this is good with me. I think it would be better to merge and release this before making the JumpProcesses v10 release. This should be setup to work correctly/safely with JumpProcesses v9 too.

[ChrisRackauckas]

DiffEqGPU.jl should be aligned with this as well, since it hooks into the non-public API part here so I think this'll break it.

[isaacsas]

I tried to add a shim to keep the old behavior for DiffEqGPU. I don't have any real GPU coding knowledge, so wouldn't feel comfortable that I could handle that ok at this point (I would have to essentially just trust an AI, which I don't think is a good approach here). I will open a PR to DiffEqGPU to see if I can get it to run CI vs. this branch to confirm its tests still pass with it.

[isaacsas]

@ChrisRackauckas how does this look now?

[isaacsas]

I can followup to have DiffEqGPU support the new prob_func and output_func signatures once this gets released.

[isaacsas]

@ChrisRackauckas should be set to merge to v3.