diff --git a/lib/ModelingToolkitBase/src/problems/jumpproblem.jl b/lib/ModelingToolkitBase/src/problems/jumpproblem.jl index 6f3cd9117d..b681bfd129 100644 --- a/lib/ModelingToolkitBase/src/problems/jumpproblem.jl +++ b/lib/ModelingToolkitBase/src/problems/jumpproblem.jl @@ -9,7 +9,7 @@ has_vrjs = any(x -> x isa VariableRateJump, jumps(sys)) has_eqs = !isempty(equations(sys)) - has_noise = get_noise_eqs(sys) !== nothing + has_noise = get_noise_eqs(sys) !== nothing || !isempty(brownians(sys)) if (has_vrjs || has_eqs) if has_eqs && has_noise diff --git a/lib/ModelingToolkitBase/src/systems/system.jl b/lib/ModelingToolkitBase/src/systems/system.jl index c8f0e2c0a8..036c10f2ad 100644 --- a/lib/ModelingToolkitBase/src/systems/system.jl +++ b/lib/ModelingToolkitBase/src/systems/system.jl @@ -326,6 +326,13 @@ struct System <: IntermediateDeprecationSystem error() end N1 == Neq || throw(IllFormedNoiseEquationsError(N1, Neq)) + if noise_eqs !== nothing && !isempty(brownians) + throw(ArgumentError( + "A system cannot have both `noise_eqs` and `brownians` specified. " * + "Use either `noise_eqs` (a matrix of noise coefficients) or " * + "`brownians` (symbolic brownian variables in equations), but not both." + )) + end check_equations(equations(continuous_events), iv) check_subsystems(systems) end diff --git a/lib/ModelingToolkitBase/src/systems/systems.jl b/lib/ModelingToolkitBase/src/systems/systems.jl index b04b429520..292f22b353 100644 --- a/lib/ModelingToolkitBase/src/systems/systems.jl +++ b/lib/ModelingToolkitBase/src/systems/systems.jl @@ -120,13 +120,21 @@ function scalarized_vars(vars) end function _mtkcompile(sys::AbstractSystem; kwargs...) - # TODO: convert noise_eqs to brownians for simplification - if has_noise_eqs(sys) && get_noise_eqs(sys) !== nothing - sys = noise_to_brownians(sys; names = :αₘₜₖ) - end + # For systems with jumps, skip full structural simplification to preserve + # variables that only appear in jumps. if !isempty(jumps(sys)) + # If brownians are present, extract them to noise_eqs for SDEProblem construction. + # If noise_eqs is already set, return as-is (no need to convert). + if !isempty(brownians(sys)) + return extract_brownians_to_noise_eqs(sys) + end return sys end + + # For non-jump systems, convert noise_eqs to brownians for simplification + if has_noise_eqs(sys) && get_noise_eqs(sys) !== nothing + sys = noise_to_brownians(sys; names = :αₘₜₖ) + end if isempty(equations(sys)) && !is_time_dependent(sys) && !_iszero(cost(sys)) return simplify_optimization_system(sys; kwargs...)::System end @@ -524,15 +532,20 @@ function add_array_observed!(obseqs::Vector{Equation}) return end -function simplify_sde_system(sys::AbstractSystem; kwargs...) - brown_vars = brownians(sys) - @set! sys.brownians = SymbolicT[] - sys = __mtkcompile(sys; kwargs...) +""" + _brownians_to_noise_eqs(eqs::Vector{Equation}, brown_vars::Vector) - new_eqs = copy(equations(sys)) +Extract brownian coefficients from equations and return (new_eqs, noise_eqs). +The brownian terms are removed from the equations and collected into a noise matrix. +This is a helper function used by both `extract_brownians_to_noise_eqs` and +`simplify_sde_system`. +""" +function _brownians_to_noise_eqs(eqs::Vector{Equation}, brown_vars::Vector) + new_eqs = copy(eqs) Is = Int[] Js = Int[] vals = SymbolicT[] + for (i, eq) in enumerate(new_eqs) resid = eq.rhs for (j, bvar) in enumerate(brown_vars) @@ -557,24 +570,48 @@ function simplify_sde_system(sys::AbstractSystem; kwargs...) end g = Matrix(sparse(Is, Js, vals, length(new_eqs), length(brown_vars))) - @set! sys.eqs = new_eqs + + # Determine noise type (scalar, diagonal, or general) # Fix for https://github.com/SciML/ModelingToolkit.jl/issues/2490 - if size(g, 2) == 1 - # If there's only one brownian variable referenced across all the equations, - # we get a Nx1 matrix of noise equations, which is a special case known as scalar noise - noise_eqs = reshape(g[:, 1], (:, 1)) - is_scalar_noise = true + noise_eqs = if size(g, 2) == 1 + # Scalar noise: Nx1 matrix + reshape(g[:, 1], (:, 1)) elseif __num_isdiag_noise(g) - # If each column of the noise matrix has either 0 or 1 non-zero entry, then this is "diagonal noise". - # In this case, the solver just takes a vector column of equations and it interprets that to - # mean that each noise process is independent - noise_eqs = __get_num_diag_noise(g) - is_scalar_noise = false + # Diagonal noise: each column has 0 or 1 non-zero entry + __get_num_diag_noise(g) else - noise_eqs = g - is_scalar_noise = false + g end + return new_eqs, noise_eqs +end + +""" + extract_brownians_to_noise_eqs(sys::AbstractSystem) + +Extract brownian variables from equations and convert them to a noise_eqs matrix, +without performing structural simplification. This is used for systems with both +jumps and brownians, where full simplification could eliminate variables that +only appear in jumps. +""" +function extract_brownians_to_noise_eqs(sys::AbstractSystem) + brown_vars = brownians(sys) + new_eqs, noise_eqs = _brownians_to_noise_eqs(equations(sys), brown_vars) + + @set! sys.eqs = new_eqs + @set! sys.noise_eqs = noise_eqs + @set! sys.brownians = SymbolicT[] + + return sys +end + +function simplify_sde_system(sys::AbstractSystem; kwargs...) + brown_vars = brownians(sys) + @set! sys.brownians = SymbolicT[] + sys = __mtkcompile(sys; kwargs...) + + new_eqs, noise_eqs = _brownians_to_noise_eqs(equations(sys), brown_vars) + dummy_sub = Dict{SymbolicT, SymbolicT}() for eq in new_eqs isdiffeq(eq) || continue diff --git a/lib/ModelingToolkitBase/test/jumpsystem.jl b/lib/ModelingToolkitBase/test/jumpsystem.jl index db52182485..47d359cb72 100644 --- a/lib/ModelingToolkitBase/test/jumpsystem.jl +++ b/lib/ModelingToolkitBase/test/jumpsystem.jl @@ -1,7 +1,7 @@ using ModelingToolkitBase, DiffEqBase, JumpProcesses, Test, LinearAlgebra using SymbolicIndexingInterface, OrderedCollections using Random, StableRNGs, NonlinearSolve -using OrdinaryDiffEq +using OrdinaryDiffEq, StochasticDiffEq, Statistics using ModelingToolkitBase: t_nounits as t, D_nounits as D using BenchmarkTools using Symbolics: SymbolicT, unwrap @@ -831,3 +831,274 @@ end @test jprob.discrete_jump_aggregation.save_positions == (true, true) end end + +# Test that JumpProblem correctly detects brownians and creates SDEProblem +# Issue: JumpProblem was only checking get_noise_eqs(sys), not brownians(sys) +# Also tests that mtkcompile properly processes brownians for systems with jumps +@testset "JumpProblem with brownians creates SDEProblem" begin + # Test 1: System with brownians and a mass action jump + @testset "Brownians + MassActionJump" begin + @variables X(t) = 10.0 + @parameters k = 1.0 + @brownians B + + # Equation with Brownian noise: dX = -k*X*dt + sqrt(k)*dB + eqs = [D(X) ~ -k * X + sqrt(k) * B] + + # A simple mass action jump: X -> 0 with rate k + jump = MassActionJump(k, [X => 1], [X => -1]) + + # Build the system with @mtkcompile - this properly processes brownians + @mtkcompile sys = System(eqs, t; jumps = [jump]) + + # After mtkcompile, brownians are converted to noise_eqs + @test MT.get_noise_eqs(sys) !== nothing + + # Create JumpProblem - should create SDEProblem + op = [X => 10.0, k => 1.0] + tspan = (0.0, 1.0) + jprob = JumpProblem(sys, op, tspan; rng) + + # The underlying problem should be SDEProblem, not ODEProblem + @test jprob.prob isa SDEProblem + + # Should be solvable without error + sol = solve(jprob, SOSRI()) + @test SciMLBase.successful_retcode(sol) + end + + # Test 2: System with brownians and a constant rate jump + @testset "Brownians + ConstantRateJump" begin + @variables X(t) = 5.0 + @parameters k = 0.5 + @brownians B + + eqs = [D(X) ~ k + 0.1 * B] + crj = ConstantRateJump(k * X, [X ~ Pre(X) - 1]) + + @mtkcompile sys = System(eqs, t; jumps = [crj]) + + @test MT.get_noise_eqs(sys) !== nothing + + op = [X => 5.0, k => 0.5] + tspan = (0.0, 1.0) + jprob = JumpProblem(sys, op, tspan; rng) + + @test jprob.prob isa SDEProblem + + sol = solve(jprob, SOSRI()) + @test SciMLBase.successful_retcode(sol) + end + + # Test 3: System with brownians and a variable rate jump + @testset "Brownians + VariableRateJump" begin + @variables X(t) = 5.0 + @parameters k = 0.5 + @brownians B + + eqs = [D(X) ~ k + 0.1 * B] + vrj = VariableRateJump(k * (1 + sin(t)), [X ~ Pre(X) + 1]) + + @mtkcompile sys = System(eqs, t; jumps = [vrj]) + + @test MT.get_noise_eqs(sys) !== nothing + + op = [X => 5.0, k => 0.5] + tspan = (0.0, 1.0) + jprob = JumpProblem(sys, op, tspan; rng) + + @test jprob.prob isa SDEProblem + + sol = solve(jprob, SOSRI()) + @test SciMLBase.successful_retcode(sol) + end + + # Test 4: System with brownians and multiple jump types + @testset "Brownians + mixed jump types" begin + @variables X(t) = 10.0 Y(t) = 5.0 + @parameters k1 = 1.0 k2 = 0.5 + @brownians B + + eqs = [D(X) ~ -k1 * X + 0.1 * B, D(Y) ~ k2] + maj = MassActionJump(k1, [X => 1], [X => -1]) + crj = ConstantRateJump(k2 * Y, [Y ~ Pre(Y) - 1]) + + @mtkcompile sys = System(eqs, t; jumps = [maj, crj]) + + @test MT.get_noise_eqs(sys) !== nothing + + op = [X => 10.0, Y => 5.0, k1 => 1.0, k2 => 0.5] + tspan = (0.0, 1.0) + jprob = JumpProblem(sys, op, tspan; rng) + + @test jprob.prob isa SDEProblem + + sol = solve(jprob, SOSRI()) + @test SciMLBase.successful_retcode(sol) + end + + # Test 5: Ensure systems WITHOUT brownians still work correctly + # (i.e., VRJ-only systems should create ODEProblem, not SDEProblem) + @testset "No brownians, VRJ only -> ODEProblem" begin + @variables X(t) = 5.0 + @parameters k = 0.5 + + # No brownians, but has equations and variable rate jump + eqs = [D(X) ~ k] + vrj = VariableRateJump(k * (1 + sin(t)), [X ~ Pre(X) + 1]) + + @mtkcompile sys = System(eqs, t; jumps = [vrj]) + + @test isempty(MT.brownians(sys)) + @test MT.get_noise_eqs(sys) === nothing + + op = [X => 5.0, k => 0.5] + tspan = (0.0, 1.0) + jprob = JumpProblem(sys, op, tspan; rng) + + # Should be ODEProblem since there are no brownians + @test jprob.prob isa ODEProblem + + sol = solve(jprob, Tsit5()) + @test SciMLBase.successful_retcode(sol) + end +end + +# Correctness tests: verify symbolic SDE+jump solutions match analytical/direct expectations +@testset "Brownians + Jumps correctness" begin + # Test 1: Pure diffusion + constant rate jump + # dX = sig*dB, X(0) = 0, with jumps X → X + delta at rate lam + # E[X(T)] = lam*delta*T (diffusion has zero mean) + @testset "Diffusion + CRJ mean" begin + @variables X(t) = 0.0 + @parameters sig = 0.3 lam = 2.0 delta = 1.0 + @brownians B + + eqs = [D(X) ~ sig * B] + crj = ConstantRateJump(lam, [X ~ Pre(X) + delta]) + + # Must pass all parameters explicitly since System doesn't auto-collect from jumps + @mtkcompile sys = System(eqs, t, [X], [sig, lam, delta], [B]; jumps = [crj]) + + T = 2.0 + Nsims = 4000 + sig_val, lam_val, delta_val = 0.3, 2.0, 1.0 + E_X = lam_val * delta_val * T # = 4.0 + + # Create JumpProblem once, use seed parameter to vary randomness + jprob = JumpProblem(sys, [X => 0.0, sig => sig_val, lam => lam_val, delta => delta_val], + (0.0, T); rng, save_positions = (false, false)) + + seed = 1111 + Xfinal = zeros(Nsims) + for i in 1:Nsims + sol = solve(jprob, SOSRI(); save_everystep = false, seed) + Xfinal[i] = sol[X, end] + seed += 1 + end + + sample_mean = mean(Xfinal) + rel_error = abs(sample_mean - E_X) / E_X + @test rel_error < 0.05 # 5% relative error + end + + # Test 2: Compare symbolic vs direct JumpProcesses construction + # Verifies that the symbolic system produces the same statistics as manual construction + @testset "Symbolic vs Direct JumpProcesses" begin + sig_val = 0.2 + lam_val = 3.0 + delta_val = 0.5 + X0 = 1.0 + T = 1.5 + Nsims = 3000 + + # Build symbolically + @variables X(t) = X0 + @parameters sig = sig_val lam = lam_val delta = delta_val + @brownians B + + eqs = [D(X) ~ sig * B] + crj = ConstantRateJump(lam, [X ~ Pre(X) + delta]) + + # Must pass all parameters explicitly since System doesn't auto-collect from jumps + @mtkcompile sys = System(eqs, t, [X], [sig, lam, delta], [B]; jumps = [crj]) + + # Create JumpProblem once for symbolic version + jprob_sym = JumpProblem(sys, [X => X0, sig => sig_val, lam => lam_val, delta => delta_val], + (0.0, T); rng, save_positions = (false, false)) + + seed = 2222 + Xfinal_sym = zeros(Nsims) + for i in 1:Nsims + sol = solve(jprob_sym, SOSRI(); save_everystep = false, seed) + Xfinal_sym[i] = sol[X, end] + seed += 1 + end + + # Build directly with JumpProcesses + f_direct(du, u, p, t) = (du[1] = 0.0) + g_direct(du, u, p, t) = (du[1] = sig_val) + sprob = SDEProblem(f_direct, g_direct, [X0], (0.0, T)) + rate_direct(u, p, t) = lam_val + affect_direct!(integ) = (integ.u[1] += delta_val) + crj_direct = ConstantRateJump(rate_direct, affect_direct!) + + jprob_direct = JumpProblem(sprob, Direct(), crj_direct; rng, save_positions = (false, false)) + + seed = 2222 # Use same seeds for comparison + Xfinal_direct = zeros(Nsims) + for i in 1:Nsims + sol = solve(jprob_direct, SOSRI(); save_everystep = false, seed) + Xfinal_direct[i] = sol[end][1] + seed += 1 + end + + # Expected mean: X0 + lam*delta*T = 1.0 + 3.0*0.5*1.5 = 3.25 + E_X = X0 + lam_val * delta_val * T + + mean_sym = mean(Xfinal_sym) + mean_direct = mean(Xfinal_direct) + + # Both should match each other and the analytical value within 5% + @test abs(mean_sym - mean_direct) / E_X < 0.05 + @test abs(mean_sym - E_X) / E_X < 0.05 + @test abs(mean_direct - E_X) / E_X < 0.05 + end + + # Test 3: Drift + diffusion + MassActionJump (birth-death with noise) + # dX = (alph - bet*X)*dt + sig*dB + # Birth: ∅ → X at rate gam + # At steady state (long time), E[X] ≈ (alph + gam) / bet + @testset "Drift + diffusion + MAJ steady state" begin + @variables X(t) = 5.0 + @parameters alph = 2.0 bet = 0.5 gam = 3.0 sig = 0.1 + @brownians B + + # ODE part drives toward alph/bet, MAJ adds gam births per unit time + eqs = [D(X) ~ alph - bet * X + sig * B] + birth = MassActionJump(gam, [0 => 1], [X => 1]) + + # Must pass all parameters explicitly since System doesn't auto-collect from jumps + @mtkcompile sys = System(eqs, t, [X], [alph, bet, gam, sig], [B]; jumps = [birth]) + + T = 20.0 # Long enough to reach steady state + Nsims = 2000 + alph_val, bet_val, gam_val, sig_val = 2.0, 0.5, 3.0, 0.1 + E_X_ss = (alph_val + gam_val) / bet_val # = 10 + + jprob = JumpProblem(sys, [X => 5.0, alph => alph_val, bet => bet_val, gam => gam_val, sig => sig_val], + (0.0, T); rng, save_positions = (false, false)) + + seed = 3333 + Xfinal = zeros(Nsims) + for i in 1:Nsims + sol = solve(jprob, SOSRI(); save_everystep = false, seed) + Xfinal[i] = sol[X, end] + seed += 1 + end + + sample_mean = mean(Xfinal) + rel_error = abs(sample_mean - E_X_ss) / E_X_ss + @test rel_error < 0.05 # 5% relative error + end +end