Skip to content

Revamp MassActionJump initialization#561

Merged
isaacsas merged 26 commits into
SciML:v10_workingfrom
isaacsas:revamp_maj_initialization
Mar 23, 2026
Merged

Revamp MassActionJump initialization#561
isaacsas merged 26 commits into
SciML:v10_workingfrom
isaacsas:revamp_maj_initialization

Conversation

@isaacsas
Copy link
Copy Markdown
Member

@isaacsas isaacsas commented Mar 3, 2026

Immutable MassActionJump Refactor: Move Working Rates to Aggregation Structs

Motivation

MassActionJump.scaled_rates was mutated in-place by update_parameters! (called from remake, reset_aggregated_jumps!, and finalize_parameters_hook!). This caused aliasing bugs where multiple JumpProblems sharing the same MAJ silently corrupted each other's rates, and a class of stale-rate bugs when parameters changed at runtime.

Design

Move the mutable rate working-copy out of MassActionJump into a new maj_rates::Vector{T} field on all aggregation structs. Rate recomputation now happens lazily during initialize! via a new fill_scaled_rates! dispatcher, eliminating update_parameters! entirely.

Key design points:

  • Parameterized MAJs (scaled_rates === nothing): fill_scaled_rates! calls maj.param_mapper(dest, maj, params) — the mapper owns all scaling logic
  • Non-parameterized MAJs (scaled_rates <: AbstractVector): fill_scaled_rates! copies the immutable scaled_rates into maj_rates
  • The JumpProblem constructor no longer eagerly materializes rates — it stores the original {Nothing} MAJ unchanged
  • Rates are materialized on every initialize! (triggered by init, solve, or reset_aggregated_jumps!), always reflecting current parameters

Changes

Core infrastructure:

  • Add in-place 3-arg mapper callable (mapper)(dest, maj, params) to MassActionJumpParamMapper
  • Add fill_scaled_rates! dispatcher that branches on MAJ type parameter
  • Update evalrxrate to 4-arg signature (speciesvec, rxidx, majump, maj_rates) taking an explicit rates vector
  • Propagate maj_rates through calculate_jump_rate, rejectrx, and get_majump_brackets

Aggregation structs:

  • Add maj_rates::Vector{T} field to all 10 aggregation structs (Direct, FRM, DirectCR, SortingDirect, NRM, CCNRM, RSSA, RSSACR, RDirect, Coevolve)
  • Add maj_rates::Vector{F} field to spatial RxRates
  • Add fill_scaled_rates! call to all initialize! methods and spatial fill_rates_and_get_times!

Problem-level simplifications:

  • Remove using_params branch from both JumpProblem constructors (main + PureLeaping) — store original MAJ directly
  • Rewrite reset_aggregated_jumps! — remove update_jump_params kwarg (throws informative error if passed)
  • Remove update_parameters! calls from remake
  • Delete finalize_parameters_hook! entirely
  • Remove dead 1-arg mapper callables from MassActionJumpParamMapper

Safety guards:

  • Error on merging mapper-backed MAJs (in both massaction_jump_combine and JumpSet vector constructor)
  • Error on constructing SpatialMassActionJump from a parameterized MAJ
  • rescale_rates_on_update mismatch check on MAJ merges

Tests:

  • Update scale_rates_field_test.jl: new tests for mapper callable API, immutability after remake, merge error guards
  • Update jprob_symbol_indexing.jl: all rate checks go through discrete_jump_aggregation.maj_rates after init
  • Update ssa_callback_test.jl: scale_rates test uses MAJ-level kwarg
  • Update regular_jumps.jl: PureLeaping test uses fill_scaled_rates! directly

MTK Compatibility

MTKBase's JumpSysMajParamMapper needs to add one method to integrate with this change:

function (mapper::JumpSysMajParamMapper)(dest::AbstractVector, maj::MassActionJump, params)
    # fill dest with scaled rates from params
    # mapper owns all scaling — check maj.rescale_rates_on_update if needed
end

Dependencies:

SciML/SciMLBase.jl#1252 and SciMLBase release
JumpProcesses #557

…structs

Move mutable rate working-copy (maj_rates) out of MassActionJump and into
all 10 aggregation structs + RxRates. Rates are now lazily materialized via
fill_scaled_rates! during initialize! instead of eagerly in the JumpProblem
constructor. This fixes aliasing bugs where multiple JumpProblems sharing
the same MAJ could silently corrupt each others rates.

Key changes:
- Add maj_rates field to all aggregation structs and RxRates
- Add fill_scaled_rates! dispatcher for parameterized vs explicit MAJs
- Update evalrxrate to 4-arg signature taking explicit maj_rates
- Propagate maj_rates through calculate_jump_rate, rejectrx, bracketing
- Remove update_parameters!, finalize_parameters_hook!
- Simplify reset_aggregated_jumps! (remove update_jump_params kwarg)
- JumpProblem constructor keeps original parameterized MAJ
- Add merge guards for mapper-backed MAJs
- Add SpatialMassActionJump guard for parameterized MAJs
- Remove dead 1-arg mapper callables
- Update all tests for lazy rate materialization pattern

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas isaacsas changed the base branch from master to v10_working March 3, 2026 20:44
isaacsas and others added 2 commits March 3, 2026 16:36
- Narrow MAJ merge guards to only block custom (non-MassActionJumpParamMapper)
  mapper types. Built-in parameterized MAJ merges via varargs and JumpSet
  vectors work correctly and are tested in ssa_callback_test.jl.
- Fix 3-arg evalrxrate call in test/spatial/reaction_rates.jl to use
  eval_massaction_rate which handles both MassActionJump and
  SpatialMassActionJump correctly.
- Add fill_scaled_rates! call before spatial rate tests to populate maj_rates.
- Update scale_rates_field_test.jl merge test to verify built-in merges
  succeed and custom mapper merges error.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- setup_majump_to_merge now copies arrays before creating the merge
  accumulator, preventing in-place mutation from corrupting the original
  MAJ net_stoch/reactant_stoch arrays
- massaction_jump_combine(::Nothing, ::MassActionJump) now creates a safe
  copy via setup_majump_to_merge instead of returning the original MAJ
- Remove scale_rates and useiszero as JumpProblem kwargs; they are now
  properties of the MassActionJump itself. Passing them raises an
  informative ArgumentError

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas isaacsas closed this Mar 4, 2026
@isaacsas isaacsas reopened this Mar 4, 2026
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas
Copy link
Copy Markdown
Member Author

isaacsas commented Mar 4, 2026

Pre-merge checklist

Before merging this PR, the following temporary changes need to be reverted:

  1. Unpin OrdinaryDiffEqCore — revert =3.15 back to 3.11 in Project.toml
  2. Remove SciMLBase branch source — remove the [sources.SciMLBase] section from Project.toml
  3. Undo CI workflow changes — remove v10_working from the pull_request.branches trigger in Tests.yml, Downgrade.yml, and ThreadSafety.yml

dependabot Bot and others added 20 commits March 9, 2026 08:35
Updates the requirements on [Catalyst](https://github.com/SciML/Catalyst.jl) to permit the latest version.

Updates `Catalyst` to 16.0.0
- [Release notes](https://github.com/SciML/Catalyst.jl/releases)
- [Changelog](https://github.com/SciML/Catalyst.jl/blob/master/HISTORY.md)
- [Commits](https://github.com/SciML/Catalyst.jl/commits/v16.0.0)

---
updated-dependencies:
- dependency-name: Catalyst
  dependency-version: 16.0.0
  dependency-type: direct:production
  dependency-group: all-julia-packages
...

Signed-off-by: dependabot[bot] <support@github.com>
…ia-packages-d0679da742

Update Catalyst requirement from 14.0, 15 to 14.0, 15, 16.0 in /docs in the all-julia-packages group across 1 directory
…n /docs in the all-julia-packages group across 1 directory"
…ocs/all-julia-packages-d0679da742

Revert "Update Catalyst requirement from 14.0, 15 to 14.0, 15, 16.0 in /docs in the all-julia-packages group across 1 directory"
The SSAIntegrator-specific method did not accept kwargs, so calling
reset_aggregated_jumps!(int; update_jump_params=true) dispatched to the
generic method which uses integrator.opts.callback (empty for SSAStepper)
instead of integrator.cb. Forward kwargs through the SSAIntegrator method.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
OrdinaryDiffEqCore added AbstractJumpProblem to its __init Union type,
creating a method ambiguity with JumpProcesses' __init. Add a package
extension that defines a more specific method for OrdinaryDiffEqAlgorithm
and DAEAlgorithm to resolve the dispatch conflict.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
OrdinaryDiffEq v6.106+ requires OrdinaryDiffEqCore v3. The old compat
bound (1.32.0) prevented the resolver from picking up v3.22.0, which is
the version that introduced the AbstractJumpProblem union in __init.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Add lmul!(::AbstractQ, ::ExtendedJumpArray) needed by LinearSolve's QR
factorization path (broke in LinearSolve v3.66.0). Fix existing ldiv!
which created a temporary but never wrote results back to the
ExtendedJumpArray fields. Update docs Catalyst compat to v16.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Replace DiscreteProblem(rn, ...) + JumpProblem(rn, prob, ...) with the
new JumpProblem(rn, u0, tspan, p) constructor. Update Catalyst compat
to v16 in both docs Project.toml files.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Fix reset_aggregated_jumps! kwarg dispatch for SSAIntegrator (SciML#562)
OrdinaryDiffEqCore now defines __init with:
  Union{OrdinaryDiffEqAlgorithm, DAEAlgorithm, StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}

This causes an ambiguity with JumpProcesses' __init(::AbstractJumpProblem, ::DEAlgorithm).

The fix adds StochasticDiffEqAlgorithm and StochasticDiffEqRODEAlgorithm to the
extension's Union type to match the exact signature from OrdinaryDiffEqCore.
…rithm-ambiguity

Fix SDE/RODE algorithm ambiguity in OrdinaryDiffEqCore extension
Adds a regression test for PR SciML#567 that checks integrator.opts.callback
to ensure jump callbacks appear exactly once when using SDE algorithms
with JumpProblems. Tests both ConstantRateJump (discrete callback) and
VariableRateJump (continuous callback) with EM() solver.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…back-duplication-test

Add test for SDE+jump callback duplication (PR SciML#567 follow-up)
Merge master into revamp_maj_initialization, bringing in:
- OrdinaryDiffEqCore package extension for __init ambiguity fix
- ExtendedJumpArray lmul!/ldiv! bug fix
- SDE+jump callback duplication test
- Doc updates for Catalyst v16

Reverted master update_jump_params kwarg addition to SSAIntegrator
reset_aggregated_jumps! and removed associated test, since this branch
removes that kwarg entirely (throws ArgumentError if passed).

Dropped OrdinaryDiffEqCore =3.15 pin in favor of master compat bound.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- Move rng kwarg from JumpProblem to solve/init in merged test files
  (extended_jump_array.jl, callbacks.jl)
- Remove stale seed! import from JumpProcesses.jl (QA failure)
- Add fill_scaled_rates! in spatial bracketing test to populate maj_rates
- Remove old prob_func(prob, i, repeat) signature in thread_safety.jl
  (use default EnsembleProblem which already deepcopies)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
… update bracketing test

- Add fill_scaled_rates! no-op for SpatialMassActionJump (rates stored in
  struct, no working copy needed). Placed in spatial_massaction_jump.jl since
  the type is defined after jumps.jl in the include order.
- Update test/bracketing.jl: pass maj_rates to get_majump_brackets, add
  maj_rates field to DummyAggregator, call fill_scaled_rates! to init rates.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas isaacsas changed the title WIP: Revamp MassActionJump initialization Revamp MassActionJump initialization Mar 21, 2026
Bug 1: A single parameterized MassActionJump with scalar param_idxs
(e.g. param_idxs = 1) would segfault when passed directly to JumpProblem
without merging, because the scalar MassActionJumpParamMapper{Int} had no
callable method. Fix: convert scalar param_idxs to a one-element vector at
construction time, and add a JumpSet constructor for single-reaction
parameterized MAJs to normalize them to collection format.

Bug 2: Merging two single-reaction parameterized MAJs with vector-backed
param_idxs would mutate the first MAJ's mapper via aliasing in
to_collection. Fix: to_collection now copies the param_idxs vector.

Also: add docstring for the mapper callable API, remove dead code for
Int-typed mappers, and add tests for both bugs.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
…TORY.md

- Move mapper callable API docstring from the method to
  MassActionJumpParamMapper struct where it belongs
- Add HISTORY.md entries for v10 breaking changes: scale_rates/useiszero
  removal from JumpProblem, immutable parameterized MAJ design,
  update_parameters! removal with reset_aggregated_jumps! guidance,
  custom mapper 3-arg callable requirement, scalar param_idxs normalization

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@isaacsas isaacsas merged commit 5917594 into SciML:v10_working Mar 23, 2026
16 of 17 checks passed
@isaacsas isaacsas deleted the revamp_maj_initialization branch March 23, 2026 15:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants