Merge pull request #4262 from isaacsas/add_poissonians

Add `@poissonians` macro

Author: AayushSabharwal

lib/ModelingToolkitBase/src/ModelingToolkitBase.jl

@@ -300,7 +300,7 @@ export alg_equations, diff_equations, has_alg_equations, has_diff_equations
 export get_alg_eqs, get_diff_eqs, has_alg_eqs, has_diff_eqs
 
 export @variables, @parameters, @independent_variables, @constants, @brownians, @brownian,
-    @discretes
+    @poissonians, @discretes
 export @named, @nonamespace, @namespace, extend, compose, complete, toggle_namespacing
 export debug_system
 

lib/ModelingToolkitBase/src/parameters.jl

@@ -6,7 +6,7 @@ import SymbolicUtils: symtype, term, hasmetadata, issym
 The type of the declared variable, used for automatic identification of
 variables/parameters/brownians/etc. by the `System` constructor.
 """
-@enum VariableType VARIABLE PARAMETER BROWNIAN
+@enum VariableType VARIABLE PARAMETER BROWNIAN POISSONIAN
 
 """
     $TYPEDEF
@@ -17,6 +17,24 @@ struct MTKVariableTypeCtx end
 
 getvariabletype(x, def = VARIABLE) = safe_getmetadata(MTKVariableTypeCtx, unwrap(x), def)::Union{typeof(def), VariableType}
 
+"""
+    $TYPEDEF
+
+The symbolic metadata key for storing the rate expression of a poissonian variable.
+"""
+struct PoissonianRateCtx end
+
+"""
+    getpoissonianrate(x)
+
+Get the rate expression associated with poissonian variable `x`. Returns `nothing`
+if `x` is not a poissonian or has no rate.
+"""
+getpoissonianrate(x::Union{Num, Symbolics.Arr}) = getpoissonianrate(Symbolics.unwrap(x))
+function getpoissonianrate(x)
+    return Symbolics.getmetadata(unwrap(x), PoissonianRateCtx, nothing)
+end
+
 """
     $TYPEDEF
 

lib/ModelingToolkitBase/src/problems/compatibility.jl

@@ -167,6 +167,19 @@ function check_has_noise(sys::System, T)
     end
 end
 
+function check_no_poissonians(sys::System, T)
+    return if !isempty(poissonians(sys))
+        throw(
+            SystemCompatibilityError(
+                """
+                Systems with unprocessed poissonians cannot be used to construct a `$T`. \
+                Call `mtkcompile` on the system first to convert poissonians to jumps.
+                """
+            )
+        )
+    end
+end
+
 function check_is_discrete(sys::System, T)
     return if !is_discrete_system(sys)
         throw(

lib/ModelingToolkitBase/src/problems/jumpproblem.jl

@@ -120,6 +120,7 @@ function check_compatible_system(T::Union{Type{JumpProblem}}, sys::System)
     check_not_dde(sys)
     check_no_cost(sys, T)
     check_no_constraints(sys, T)
+    check_no_poissonians(sys, T)
     check_has_jumps(sys, T)
     return check_is_continuous(sys, T)
 end

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

@@ -901,6 +901,7 @@ const SYS_PROPS = [
     :ps
     :tspan
     :brownians
+    :poissonians
     :jumps
     :name
     :description
@@ -1417,6 +1418,17 @@ function namespace_brownians(sys::AbstractSystem)
     return bs
 end
 
+function namespace_poissonians(sys::AbstractSystem)
+    ps = poissonians(sys)
+    if ps === get_poissonians(sys)
+        ps = copy(ps)
+    end
+    for i in eachindex(ps)
+        ps[i] = renamespace(sys, ps[i])
+    end
+    return ps
+end
+
 function namespace_assignment(eq::Assignment, sys)
     _lhs = namespace_expr(eq.lhs, sys)
     _rhs = namespace_expr(eq.rhs, sys)
@@ -1999,6 +2011,25 @@ function brownians(sys::AbstractSystem)
     return bs
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Get all of the poissonian variables involved in the system `sys` and all subsystems,
+appropriately namespaced.
+"""
+function poissonians(sys::AbstractSystem)
+    ps = get_poissonians(sys)
+    systems = get_systems(sys)
+    if isempty(systems)
+        return ps
+    end
+    ps = copy(ps)
+    for subsys in systems
+        append!(ps, namespace_poissonians(subsys))
+    end
+    return ps
+end
+
 """
     $(TYPEDSIGNATURES)
 

lib/ModelingToolkitBase/src/systems/system.jl

@@ -102,6 +102,13 @@ struct System <: IntermediateDeprecationSystem
     """
     brownians::Vector{SymbolicT}
     """
+    The poissonian variables of the system, created via `@poissonians`. Each poissonian
+    variable represents an independent Poisson counting process with an associated rate.
+    A system with poissonians cannot be simulated directly. It needs to be compiled using
+    `mtkcompile` which converts poissonians into jump equations.
+    """
+    poissonians::Vector{SymbolicT}
+    """
     The independent variable for a time-dependent system, or `nothing` for a time-independent
     system.
     """
@@ -286,7 +293,7 @@ struct System <: IntermediateDeprecationSystem
 
     function System(
             tag, eqs, noise_eqs, jumps, constraints, costs, consolidate, unknowns, ps,
-            brownians, iv, observed, var_to_name, name, description, bindings,
+            brownians, poissonians, iv, observed, var_to_name, name, description, bindings,
             initial_conditions, guesses, systems, initialization_eqs, continuous_events,
             discrete_events, connector_type, assertions = Dict{SymbolicT, String}(),
             metadata = MetadataT(), gui_metadata = nothing, is_dde = false, tstops = [],
@@ -350,7 +357,7 @@ struct System <: IntermediateDeprecationSystem
         end
         return new(
             tag, eqs, noise_eqs, jumps, constraints, costs,
-            consolidate, unknowns, ps, brownians, iv,
+            consolidate, unknowns, ps, brownians, poissonians, iv,
             observed, var_to_name, name, description, bindings, initial_conditions,
             guesses, systems, initialization_eqs, continuous_events, discrete_events,
             connector_type, assertions, metadata, gui_metadata, is_dde,
@@ -433,6 +440,7 @@ All other keyword arguments are named identically to the corresponding fields in
 """
 function System(
         eqs::Vector{Equation}, iv, dvs, ps, brownians = SymbolicT[];
+        poissonians = SymbolicT[],
         constraints = Union{Equation, Inequality}[], noise_eqs = nothing, jumps = JumpType[],
         costs = SymbolicT[], consolidate = default_consolidate,
         # `@nospecialize` is only supported on the first 32 arguments. Keep this early.
@@ -477,6 +485,7 @@ function System(
         filter!(!Base.Fix2(_is_unknown_delay_or_evalat, iv), dvs)
     end
     brownians = unwrap_vars(brownians)
+    poissonians = unwrap_vars(poissonians)
 
     if noise_eqs !== nothing
         noise_eqs = unwrap_vars(noise_eqs)
@@ -613,7 +622,7 @@ function System(
     jumps = Vector{JumpType}(jumps)
     return System(
         __get_new_tag(), eqs, noise_eqs, jumps, constraints,
-        costs, consolidate, dvs, ps, brownians, iv, observed,
+        costs, consolidate, dvs, ps, brownians, poissonians, iv, observed,
         var_to_name, name, description, bindings, initial_conditions, guesses, systems, initialization_eqs,
         continuous_events, discrete_events, connector_type, assertions, metadata, gui_metadata, is_dde,
         tstops, inputs, outputs, tearing_state, true, false,
@@ -704,13 +713,25 @@ function System(eqs::Vector{Equation}, iv; kwargs...)
     eqs = [diffeqs; othereqs]
 
     brownians = Set{SymbolicT}()
+    poissonians = Set{SymbolicT}()
     for x in allunknowns
         x = unwrap(x)
         if getvariabletype(x) == BROWNIAN
             push!(brownians, x)
+        elseif getvariabletype(x) == POISSONIAN
+            push!(poissonians, x)
         end
     end
     setdiff!(allunknowns, brownians)
+    setdiff!(allunknowns, poissonians)
+
+    # Extract variables and parameters from poissonian rate expressions
+    for p in poissonians
+        rate = getpoissonianrate(p)
+        if rate !== nothing
+            collect_vars!(allunknowns, ps, rate, iv)
+        end
+    end
 
     cstrs = Vector{Union{Equation, Inequality}}(get(kwargs, :constraints, []))
     _cstrunknowns, cstrps = process_constraint_system(cstrs, allunknowns, ps, iv)
@@ -764,7 +785,8 @@ function System(eqs::Vector{Equation}, iv; kwargs...)
     end
 
     return System(
-        eqs, iv, collect(allunknowns), collect(new_ps), collect(brownians); kwargs...
+        eqs, iv, collect(allunknowns), collect(new_ps), collect(brownians);
+        poissonians = collect(poissonians), kwargs...
     )
 end
 
@@ -995,6 +1017,7 @@ function flatten(sys::System, noeqs = false)
     return System(
         noeqs ? Equation[] : equations(sys), get_iv(sys), unknowns(sys),
         parameters(sys; initial_parameters = true), brownians(sys);
+        poissonians = poissonians(sys),
         jumps = jumps(sys), constraints = constraints(sys), costs = costs,
         consolidate = default_consolidate, observed = observed(sys),
         bindings = bindings(sys), initial_conditions = initial_conditions(sys),
@@ -1424,6 +1447,7 @@ function Base.isapprox(sysa::System, sysb::System)
         issetequal(get_unknowns(sysa), get_unknowns(sysb)) &&
         issetequal(get_ps(sysa), get_ps(sysb)) &&
         issetequal(get_brownians(sysa), get_brownians(sysb)) &&
+        issetequal(get_poissonians(sysa), get_poissonians(sysb)) &&
         issetequal(get_observed(sysa), get_observed(sysb)) &&
         isequal(get_description(sysa), get_description(sysb)) &&
         isequal(get_bindings(sysa), get_bindings(sysb)) &&
@@ -1452,7 +1476,8 @@ function Base.copy(sys::System)
     return System(
         __get_new_tag(), copy(get_eqs(sys)), _maybe_copy(get_noise_eqs(sys)), copy(get_jumps(sys)),
         copy(get_constraints(sys)), copy(get_costs(sys)), get_consolidate(sys),
-        copy(get_unknowns(sys)), copy(get_ps(sys)), copy(get_brownians(sys)), get_iv(sys),
+        copy(get_unknowns(sys)), copy(get_ps(sys)), copy(get_brownians(sys)),
+        copy(get_poissonians(sys)), get_iv(sys),
         copy(get_observed(sys)), copy(get_var_to_name(sys)), nameof(sys), get_description(sys),
         copy(get_bindings(sys)), copy(get_initial_conditions(sys)), copy(get_guesses(sys)),
         map(copy, get_systems(sys)), copy(get_initialization_eqs(sys)),