JumpProblem(...; save_positions=...)` does not propagate to `VR_Direct`'s `ContinuousCallback

[RomanSahakyan03]

When aggregating multiple VariableRateJumps with vr_aggregator = VR_Direct(), the save_positions keyword argument passed to JumpProblem(...) is silently ignored. The per-VariableRateJump default of (false, true) is used instead, inserting an extra save after every fired jump. With saveat, this produces ensemble trajectories of non-uniform length along the canonical grid a silent ensemble-average bias rather than an error.

I hit this while validating an MCWF unraveling of a Lindblad master equation. The fix turned out to be one line per jump, but the failure mode was non-obvious because the kwarg is accepted without warning.

using DifferentialEquations, JumpProcesses

const γ = (1.0, 1.0)
const H = ComplexF64[0 2 0; 2 0 2; 0 2 0]   # constant Ω = 2, Δ = 0

real_to_ψ(u) = ComplexF64[u[1]+1im*u[2], u[3]+1im*u[4], u[5]+1im*u[6]]
ψ_to_real(ψ) = Float64[real(ψ[1]),imag(ψ[1]),real(ψ[2]),imag(ψ[2]),real(ψ[3]),imag(ψ[3])]

function nojump!(du, u, p, t)
  ψ  = real_to_ψ(u)
  P2 = abs2(ψ[2]); γt = γ[1] + γ[2]
  dψ = -1im .* (H * ψ)
  dψ[1] +=  0.5 * γt * P2       * ψ[1]
  dψ[2] += -0.5 * γt * (1 - P2) * ψ[2]
  dψ[3] +=  0.5 * γt * P2       * ψ[3]
  du .= ψ_to_real(dψ)
end

rate1(u,p,t) = γ[1] * abs2(real_to_ψ(u)[2])
rate2(u,p,t) = γ[2] * abs2(real_to_ψ(u)[2])
affect1!(i)  = (fill!(i.u, 0.0); i.u[1] = 1.0)
affect2!(i)  = (fill!(i.u, 0.0); i.u[5] = 1.0)

ode = ODEProblem(nojump!, ψ_to_real(ComplexF64[1, 0, 0]), (0.0, 5.0))
ts  = collect(0:0.05:5)
N   = 500

jp_0 = JumpProblem(ode, Direct(),
      VariableRateJump(rate1, affect1!),
      VariableRateJump(rate2, affect2!);
      vr_aggregator = VR_Direct())

jp_a = JumpProblem(ode, Direct(),
      VariableRateJump(rate1, affect1!),
      VariableRateJump(rate2, affect2!);
      vr_aggregator  = VR_Direct(),
      save_positions = (false, false))   # ← silently ignored

jp_b = JumpProblem(ode, Direct(),
      VariableRateJump(rate1, affect1!; save_positions = (false, false)),
      VariableRateJump(rate2, affect2!; save_positions = (false, false));
      vr_aggregator  = VR_Direct(),
      save_positions = (false, false))

run_ens(jp) = solve(EnsembleProblem(jp), Tsit5(), EnsembleSerial();
                  trajectories = N, saveat = ts,
                  abstol = 1e-9, reltol = 1e-9)

for (label, jp) in [("(0) bare defaults", jp_0),
                  ("(A) JumpProblem kwarg only", jp_a),
                  ("(B) per-jump save_positions=(F,F)", jp_b)]
  sol   = run_ens(jp)
  n_off = count(traj -> length(traj.t) != length(ts), sol.u)
  if n_off > 0
      @warn "$label: $n_off / $N trajectories off-grid"
  else
      println("$label: clean — all $N trajectories match the saveat grid")
  end
end

Output

┌ Warning: (0) bare defaults: 261 / 500 trajectories off-grid
└ @ Main d:\Git repos\jumpAD\stage1\save_positions_repro.jl:58
┌ Warning: (A) JumpProblem kwarg only: 266 / 500 trajectories off-grid
└ @ Main d:\Git repos\jumpAD\stage1\save_positions_repro.jl:58
(B) per-jump save_positions=(F,F): clean — all 500 trajectories match the saveat grid

• VariableRateJump's default is save_positions = (false, true) at src/jumps.jl:178, which inserts a save after each jump.
• build_variable_integcallback (src/variable_rate.jl:380) constructs the VR_Direct ContinuousCallback by OR-reducing only the per-VariableRateJump save_positions fields.
• JumpProblem(...; save_positions=...) is consumed by a different (constant-rate) callback path and never reaches VR_Direct.

The failure is silent and scales with jump frequency: ensemble averaging accumulates a bias that's small at low rates but dominates the signal in jump-rich regimes. There's no error or warning unless you explicitly check trajectory lengths, so it's easy to miss without a deterministic reference to compare against.

---

As a suggested fix I reccomend to make JumpProblem's save_positions kwarg, when explicitly passed, override the per-VariableRateJump defaults at construction or change the per-VariableRateJump default to (false, false) to match user intuition.

Happy to send a PR. Let me know which is preferred.

[isaacsas]

Yeah, the current (and long standing) setup is that for VariableRateJumps you control this at the jump level, while for MassActionJumps and ConstantRateJumps this is controlled via save_positions to JumpProblem. This could definitely use a cleaner interface, but it would be a breaking change and probably needs a bit of thought (we'd probably want a uniform way to control it globally, or at a per-jump level for every individual jump, and for the variable rate solvers I believe it is actually handled via the callback interface in OrdinaryDiffEq so we'd still need to pass a reasonable value for it for generated callbacks).

For now we should probably just better document that it only controls what happens for non-VariableRateJumps.

[RomanSahakyan03]

Thanks for clarifying @isaacsas . I can open a documentation PR noting that JumpProblem(...; save_positions=...) only controls non-VariableRateJump cases, and that VariableRateJump users should set save_positions on each jump directly.

[isaacsas]

Doc updates would always be welcome! I tweaked the doc string on master, but any other updates would be helpful.