Skip to content

WIP: Revamp MassActionJump initialization#561

Open
isaacsas wants to merge 4 commits intoSciML:v10_workingfrom
isaacsas:revamp_maj_initialization
Open

WIP: Revamp MassActionJump initialization#561
isaacsas wants to merge 4 commits intoSciML:v10_workingfrom
isaacsas:revamp_maj_initialization

Conversation

@isaacsas
Copy link
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
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

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.

1 participant