Skip to content

use integrator rng API#557

Open
isaacsas wants to merge 22 commits intoSciML:masterfrom
isaacsas:add_integrator_rngs
Open

use integrator rng API#557
isaacsas wants to merge 22 commits intoSciML:masterfrom
isaacsas:add_integrator_rngs

Conversation

@isaacsas
Copy link
Member

@isaacsas isaacsas commented Feb 24, 2026 •
edited
Loading

Needs SciML/StochasticDiffEq.jl#681 and a StochasticDiffEq release before we can fully run CI here.

Summary

Makes the integrator the single source of truth for RNG in JumpProcesses.

Previously, each aggregator struct, JumpProblem, and variable-rate aggregator stored its own rng field. This caused data races with EnsembleThreads (fixed in #556 with a workaround) and made RNG control fragmented across multiple objects.

Now all RNG access goes through get_rng(integrator) at runtime, using the SciMLBase has_rng/get_rng/set_rng! interface.

Key changes

  • Aggregator structs: Remove rng field and RNG type parameter from AbstractSSAJumpAggregator and all 10 concrete aggregators (Direct, FRM, NRM, CCNRM, DirectCR, RDirect, SortingDirect, RSSA, RSSACR, Coevolve) plus 2 spatial aggregators (NSM, DirectCRDirect). All aggregate() signatures drop the rng positional arg.
  • SSAIntegrator: Add rng field with has_rng/get_rng/set_rng! implementation, so the SSAStepper pathway has the same RNG interface as ODE/SDE integrators.
  • Variable-rate callbacks: Remove rng from VR_FRMEventCallback and VR_DirectEventCache; they now read get_rng(integrator) in their affect/initialize functions.
  • JumpProblem: Remove rng struct field. The rng kwarg no gives an error that it should instead be passed via solve.
  • solve.jl: Add resolve_rng(rng, seed) utility. Simplify __jump_init to always forward rng. Remove _derive_jump_seed and seed-based reseeding.
  • SimpleTauLeaping / SimpleExplicitTauLeaping: Accept rng/seed kwargs directly instead of reading from jump_prob.rng.
  • Docs: Add RNG control section to API docs and FAQ entry.
  • Tests: Add test/rng_kwarg_tests.jl covering all three solve pathways (SSAStepper, ODE+VR, SDE+VR) and SimpleTauLeaping. Update ensemble tests.

Dependencies

Requires:

  • SciMLBase >= 2.144 (for has_rng/get_rng/set_rng!)
  • OrdinaryDiffEqCore >= 3.9 (for ODEIntegrator.rng)
  • StochasticDiffEq, see above

isaacsas and others added 11 commits February 23, 2026 22:05
@isaacsas @claude
Move RNG from aggregators/JumpProblem to integrator
e0498e3 
Make the integrator the single source of truth for RNG (Phase 1 Step 4).
Remove rng field from all aggregator structs, JumpProblem, and variable-rate
callback structs. All code now reads get_rng(integrator) at runtime.

Key changes:
- Remove RNG type param from AbstractSSAJumpAggregator and all concrete aggregators
- Remove rng from build_jump_aggregation and all aggregate() signatures
- Add rng field to SSAIntegrator with has_rng/get_rng/set_rng! interface
- Add resolve_rng utility for rng/seed/fallback priority resolution
- Remove rng from VR_FRMEventCallback and VR_DirectEventCache
- JumpProblem rng kwarg now forwards to solver via solkwargs
- Simplify solve.jl: remove _derive_jump_seed, simplify __jump_init
- Update SimpleTauLeaping/SimpleExplicitTauLeaping to use rng kwarg
- Add rng_kwarg_tests.jl covering all three solve pathways
- Update docs with RNG control section and FAQ entry
- Bump version to 9.23.0, SciMLBase compat to 2.144

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Revert version bump to 9.22.1
8861d6e 
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Remove rng kwarg from JumpProblem; pass rng to solve/init only
e37938b 
Breaking: JumpProblem no longer accepts an `rng` keyword argument.
Passing `rng` now throws an ArgumentError directing users to pass
it to `solve` or `init` instead. This fixes inconsistent RNG
priority ordering across solver pathways (SSAStepper, ODE/SDE,
tau-leaping) and makes the integrator the single source of truth
for RNG state.

- Simplify resolve_rng to 2-arg form (remove fallback_rng)
- Remove jump_prob.kwargs rng fallback from tau-leaping solvers
- Update all tests to pass rng to solve/init instead of JumpProblem
- Update docs and HISTORY.md with breaking change notes

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Clean up code quality: revert unnecessary diffs, add set_rng! type ch…
a7de080 
…ecking

- Revert <:Any back to named type params in ssajump.jl (minimize diff from master)
- Revert solkwargs to simpler ternary form in problem.jl (rng removal made filter unnecessary)
- Add type checking to SSAIntegrator set_rng! to match StochasticDiffEq pattern

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Add tests for set_rng! type mismatch error on SSAIntegrator
0ded524 
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Remove unused StableRNG import and variable from GPU test
20376a9 
The GPU tau-leaping kernel uses PassthroughRNG internally; the
StableRNG was never consumed by the EnsembleGPUKernel path.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Simplify test solve calls: remove unnecessary isnothing(rng) checks
0c32cce 
solve(; rng = nothing) falls through to default_rng() via resolve_rng,
so the ternary guards are redundant.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Use jump times instead of final state for trajectory uniqueness tests
33b86e9 
Comparing continuous-valued jump times is more robust than discrete
final state values, which could collide by chance.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Strengthen VR_FRM threshold uniqueness test
2b6d3c6 
Check jump_u thresholds directly via init, and verify both event times
and post-event thresholds from ensemble solve.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Merge master into add_integrator_rngs (includes VR_FRM u_modified fix)
42cc20d 
Resolves conflicts by combining integrator RNG (get_rng) with
u_modified!(true) fix, and using robust first_jump_time helpers
in tests that account for the initialization save.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Move EnsembleThreads tests to thread_safety.jl, fix Project.toml
44ced84 
- Remove StochasticDiffEq from [deps] (test-only, belongs in [extras])
- Move EnsembleThreads data race tests from ensemble_problems.jl to
  thread_safety.jl with stronger assertions (400 trajectories, allunique
  on first reaction times, all 3 VR aggregators, adaptive SDE solver)
- Renumber remaining ensemble_problems.jl sections

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
isaacsas
isaacsas commented Feb 24, 2026
View reviewed changes
@isaacsas @claude
Address PR review: strengthen tests, add seed coverage
55e130b 
- ensemble_problems.jl: check initial jump_u threshold (u[2]) instead
  of post-event (u[3]) now that u_modified!(true) produces init save
- extended_jump_array_remake.jl: rework to solve → remake → init →
  compare jump_u, verifying fresh thresholds after remake
- geneexpr_test.jl: remove unnecessary isnothing(rng) branch
- linearreaction_test.jl: make rng const for type stability
- variable_rate.jl: switch getmean from seed integer to StableRNG
- rng_kwarg_tests.jl: add SDE+VR different seeds test, SDE+VR no-rng
  functional test, seed kwarg reproducibility and different-seeds tests
  for all pathways (SSAStepper, ODE+VR, SDE+VR), use first_jump_time
  helper for robust event time extraction

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas isaacsas changed the title WIP: use integrator rng API use integrator rng API Feb 24, 2026
@isaacsas
Copy link
Member Author

isaacsas commented Feb 24, 2026

This should be good to go once there is a StochasticDiffEq release and we can run tests to check CI here.

isaacsas and others added 3 commits February 24, 2026 14:21
@isaacsas @claude
Address review round 3: fix docs example, add tau-leaping seed tests,…
9686436 
… strengthen thread_safety assertions

- docs/advanced_point_process.md: move rng from JumpProblem to solve
- rng_kwarg_tests.jl: add seed kwarg tests for SimpleTauLeaping (tests 13-14), renumber 15-20
- thread_safety.jl: SSAStepper EnsembleThreads now runs 400 trajectories and asserts allunique first jump times

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Add dedicated threaded CI for thread_safety tests
ec3b661 
- Create ThreadSafety.yml workflow: runs thread_safety.jl with --threads=4
- Add ThreadSafety group to Tests.yml matrix (serial) and runtests.jl
- Move thread_safety.jl from InterfaceII to ThreadSafety group
- Fix shared rng data race: SSAStepper ensemble tests now use task-local
  default_rng() instead of a shared StableRNG across threads

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Add supports_solve_rng trait for JumpProblem solve paths
201e70e 
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas isaacsas mentioned this pull request Feb 25, 2026
Open
isaacsas and others added 5 commits February 28, 2026 07:24
@isaacsas
Update package versions in Project.toml
613611e
@isaacsas @claude
Add rescale_rates_on_update field to MassActionJump to prevent double…
a434959 
…-scaling

When MTKBase/Catalyst constructs MassActionJumps, the symbolic rate expressions
already include stoichiometric scaling (e.g. k/3! for 3X → Y), so they pass
scale_rates=false. However, update_parameters! defaulted to scale_rates=true,
causing every subsequent parameter update path (reset_aggregated_jumps!, remake,
finalize_parameters_hook!) to double-scale the rates.

The fix stores the intended scaling behavior on the struct as
rescale_rates_on_update, which update_parameters! now reads as its default.
This field is propagated through all merge/combine paths and validated for
consistency when merging multiple MassActionJumps.

Also fixes SpatialMassActionJump conversion constructor to default
scale_rates=false since it receives already-scaled rates.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas @claude
Fix PreScaledMapper test to actually catch double-scaling regression
83352c5 
The mapper now conditionally applies scalerates! based on scale_rates,
matching MTKBase JumpSysMajParamMapper behavior. Previously it ignored
scale_rates entirely, so the test would pass even with the old buggy
default of scale_rates=true.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas
Update project version to 9.23
4a52abe
@isaacsas @claude
add thread-safety warning to JumpProblem docstring
d1d6beb 
Document that JumpProblem contains mutable state and is not thread-safe
for concurrent solve calls. Note that EnsembleProblem handles isolation
automatically via SciMLBase per-task deepcopy and per-trajectory RNG.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas isaacsas force-pushed the add_integrator_rngs branch from 870ae0b to d1d6beb Compare February 28, 2026 17:32
isaacsas and others added 2 commits March 1, 2026 17:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@isaacsas