Merge pull request #800 from SciML/noise_scaling_update

Noise scaling update

Author: TorkelE

src/reactionsystem.jl

@@ -496,7 +496,8 @@ function make_ReactionSystem_internal(rxs_and_eqs::Vector, iv, us_in, ps_in; spa
             (sym isa Symbolic) && findvars!(ps, us, sym, ivs, vars)
         end
 
-        # Will appear here: add stuff from nosie scaling.
+        # Extract all quantities encountered in relevant `Reaction` metadata.
+        has_noise_scaling(rx) && findvars!(ps, us, get_noise_scaling(rx), ivs, vars)
     end
 
     # Extracts any species, variables, and parameters that occur in (non-reaction) equations.
@@ -1233,7 +1234,26 @@ function set_default_metadata(rs::ReactionSystem;  default_reaction_metadata = [
     updated_equations = map(eqtransform, get_eqs(rs))
     @set! rs.eqs = updated_equations
     @set! rs.rxs = Reaction[rx for rx in updated_equations if rx isa Reaction]
-    
+
+    # Special routine to handle `Reaction` metadata that can contain new symbolic variables.
+    # Currently, `noise_scaling` is the only relevant metadata supported this way.
+    drm_dict = Dict(default_reaction_metadata)
+    if haskey(drm_dict, :noise_scaling)
+        # Finds parameters, species, and variables in the noise scaling term.       
+        ns_expr = drm_dict[:noise_scaling] 
+        ns_syms = [Symbolics.unwrap(sym) for sym in get_variables(ns_expr)]
+        ns_ps = Iterators.filter(ModelingToolkit.isparameter, ns_syms)
+        ns_sps = Iterators.filter(Catalyst.isspecies, ns_syms)
+        ns_vs = Iterators.filter(sym -> !Catalyst.isspecies(sym) &&
+                                        !ModelingToolkit.isparameter(sym), ns_syms)
+        # Adds parameters, species, and variables to the `ReactionSystem`.
+        @set! rs.ps = union(get_ps(rs), ns_ps)
+        sps_new = union(get_species(rs), ns_sps)
+        @set! rs.species = sps_new
+        vs_old = @view get_unknowns(rs)[length(get_species(rs))+1 : end]            
+        @set! rs.unknowns = union(sps_new, vs_old, ns_vs)
+    end
+
     # Updates reaction metadata for all its subsystems.
     new_sub_systems = similar(get_systems(rs))
     for (i, sub_system) in enumerate(get_systems(rs))

test/simulation_and_solving/simulate_SDEs.jl

@@ -7,6 +7,9 @@ using Catalyst, Statistics, StochasticDiffEq, Test
 using StableRNGs
 rng = StableRNG(123456)
 
+# Sets the default `t` to use.
+t = default_t()
+
 # Fetch test functions and networks.
 include("../test_functions.jl")
 include("../test_networks.jl")
@@ -315,30 +318,61 @@ let
     end
 end
 
+# Tests noise scaling is added properly for programmatically created system.
+# Tests noise scaling for interpolation into the `@reaction macro`.
+# Tests that species, variables, and parameters are extracted from noise scaling metadata when 
+# `ReactionSystems`s are declared programmatically.
+let
+    # Programmatically creates a model with noise scaling.
+    t = default_t()
+    @species X(t) H(t)
+    @variables h(t)
+    @parameters p d η
+    rx1 = Reaction(p, nothing, [X]; metadata = [:noise_scaling => η*H + 1])
+    rx2 = @reaction d, X --> 0, [noise_scaling=$h]
+    @named rs = ReactionSystem([rx1, rx2], t)
+
+    # Checks that noise scaling has been added correctly.
+    @test issetequal([X, H], species(rs))
+    @test issetequal([X, H, h], unknowns(rs))
+    @test issetequal([p, d, η], parameters(rs))
+    @test isequal(get_noise_scaling(reactions(rs)[1]), η*H + 1)
+    @test isequal(get_noise_scaling(reactions(rs)[2]), h)
+end
+
 # Tests the `remake_noise_scaling` function.
+# Checks a parameter part of the initial system is added properly.
+# Checks that parameters, variables, and species not part of the original system are added properly.
 let
     # Creates noise scaling networks.
     noise_scaling_network1 = @reaction_network begin
+        @parameters η1
         p, 0 --> X, [noise_scaling=2.0]
         d, X --> 0
     end
-    noise_scaling_network2 = set_default_noise_scaling(noise_scaling_network1, 0.5)
-
-    # Checks that the two networks' reactions have the correct metadata.
-    @test reactions(noise_scaling_network1)[1].metadata == [:noise_scaling => 2.0]
-    @test reactions(noise_scaling_network1)[2].metadata == []
-    @test reactions(noise_scaling_network2)[1].metadata == [:noise_scaling => 2.0]
-    @test reactions(noise_scaling_network2)[2].metadata == [:noise_scaling => 0.5]
+    @unpack p, d, η1, X = noise_scaling_network1
+    @parameters η2
+    @species Y(t)
+    @variables V(t)
+    
+    noise_scaling_network2 = set_default_noise_scaling(noise_scaling_network1, η1 + η2 + Y + V)
+    @test isequal(reactions(noise_scaling_network2)[1].metadata, [:noise_scaling => 2.0])
+    @test isequal(reactions(noise_scaling_network2)[2].metadata, [:noise_scaling => η1 + η2 + Y + V])
+    @test issetequal([p, d, η1, η2], parameters(noise_scaling_network2))
+    @test issetequal([X, Y], species(noise_scaling_network2))
+    @test issetequal([X, Y, V], unknowns(noise_scaling_network2))
+    @test issetequal([V], nonspecies(noise_scaling_network2))
 end
 
 # Tests the `remake_noise_scaling` function on a hierarchical model.
-let
+# Checks that species, variables, and parameters not part of the original system is added properly.
+let        
     # Creates hierarchical model.
-    rn1 = @reaction_network begin
+    rn1 = @reaction_network rn1 begin
         p, 0 --> X, [noise_scaling=2.0]
         d, X --> 0
     end
-    rn2 = @reaction_network begin
+    rn2 = @reaction_network rn2 begin
         k1, X1 --> X2, [noise_scaling=5.0]
         k2, X2 --> X1
     end
@@ -355,6 +389,7 @@ let
 end
 
 
+
 ### Other Tests ###
 
 # Tests simulating a network without parameters.