Merge pull request #677 from SciML/mtkitize_lowering

fix variable name lowering for non-map parameters and SIR regression

Author: ChrisRackauckas

src/systems/diffeqs/abstractodesystem.jl

@@ -267,14 +267,16 @@ function DiffEqBase.ODEProblem{iip}(sys::AbstractODESystem,u0map,tspan,
                                     kwargs...) where iip
     dvs = states(sys)
     ps = parameters(sys)
-    u0map′ = [lower_varname(value(k), sys.iv) => value(v) for (k, v) in u0map]
+    u0map′ = lower_mapnames(u0map,sys.iv)
     u0 = varmap_to_vars(u0map′,dvs)
+    
     if !(parammap isa DiffEqBase.NullParameters)
-        parammap′ = [value(k) => value(v) for (k, v) in parammap]
+        parammap′ = lower_mapnames(parammap)
         p = varmap_to_vars(parammap′,ps)
     else
         p = ps
     end
+
     f = ODEFunction{iip}(sys,dvs,ps,u0;tgrad=tgrad,jac=jac,checkbounds=checkbounds,
                         linenumbers=linenumbers,parallel=parallel,simplify=simplify,
                         sparse=sparse,eval_expression=eval_expression,kwargs...)
@@ -311,14 +313,16 @@ function ODEProblemExpr{iip}(sys::AbstractODESystem,u0map,tspan,
 
     dvs = states(sys)
     ps = parameters(sys)
-    u0map′ = [lower_varname(value(k), sys.iv) => value(v) for (k, v) in u0map]
-    parammap′ = [value(k) => value(v) for (k, v) in parammap]
+    u0map′ = lower_mapnames(u0map,sys.iv)
     u0 = varmap_to_vars(u0map′,dvs)
+
     if !(parammap isa DiffEqBase.NullParameters)
+        parammap′ = lower_mapnames(parammap)
         p = varmap_to_vars(parammap′,ps)
     else
         p = ps
     end
+
     f = ODEFunctionExpr{iip}(sys,dvs,ps,u0;tgrad=tgrad,jac=jac,checkbounds=checkbounds,
                         linenumbers=linenumbers,parallel=parallel,
                         simplify=simplify,

src/systems/diffeqs/sdesystem.jl

@@ -310,8 +310,17 @@ function DiffEqBase.SDEProblem{iip}(sys::SDESystem,u0map,tspan,parammap=DiffEqBa
 
     dvs = states(sys)
     ps = parameters(sys)
-    u0 = varmap_to_vars(u0map,dvs)
-    p = varmap_to_vars(parammap,ps)
+
+    u0map′ = lower_mapnames(u0map,sys.iv)
+    u0 = varmap_to_vars(u0map′,dvs)
+
+    if !(parammap isa DiffEqBase.NullParameters)
+        parammap′ = lower_mapnames(parammap)
+        p = varmap_to_vars(parammap′,ps)
+    else
+        p = ps
+    end
+
     f = SDEFunction{iip}(sys,dvs,ps,u0;tgrad=tgrad,jac=jac,Wfact=Wfact,
                          checkbounds=checkbounds,
                          linenumbers=linenumbers,parallel=parallel,
@@ -358,8 +367,17 @@ function SDEProblemExpr{iip}(sys::SDESystem,u0map,tspan,
                                     kwargs...) where iip
     dvs = states(sys)
     ps = parameters(sys)
-    u0 = varmap_to_vars(u0map,dvs)
-    p = varmap_to_vars(parammap,ps)
+
+    u0map′ = lower_mapnames(u0map,sys.iv)
+    u0 = varmap_to_vars(u0map′,dvs)
+
+    if !(parammap isa DiffEqBase.NullParameters)
+        parammap′ = lower_mapnames(parammap)
+        p = varmap_to_vars(parammap′,ps)
+    else
+        p = ps
+    end
+    
     f = SDEFunctionExpr{iip}(sys,dvs,ps,u0;tgrad=tgrad,jac=jac,
                         Wfact=Wfact,checkbounds=checkbounds,
                         linenumbers=linenumbers,parallel=parallel,

src/utils.jl

@@ -271,6 +271,15 @@ function lower_varname(t::Term, iv)
 end
 lower_varname(t::Sym, iv) = t
 
+function lower_mapnames(umap::AbstractArray{<:Pair}) where T
+    [value(k) => value(v) for (k, v) in umap]
+end
+function lower_mapnames(umap::AbstractArray{<:Pair},name) where T
+    [lower_varname(value(k), name) => value(v) for (k, v) in umap]
+end
+lower_mapnames(umap::AbstractArray{<:Number}) = umap # Ambiguity
+lower_mapnames(umap::AbstractArray{<:Number},name) = umap
+
 function flatten_differential(O::Term)
     @assert is_derivative(O) "invalid differential: $O"
     is_derivative(O.args[1]) || return (O.args[1], O.op.x, 1)

test/modelingtoolkitize.jl

@@ -1,7 +1,7 @@
 using OrdinaryDiffEq, ModelingToolkit, Test
 using GalacticOptim, Optim
 
-const N = 32
+N = 32
 const xyd_brusselator = range(0,stop=1,length=N)
 brusselator_f(x, y, t) = (((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.
 limit(a, N) = ModelingToolkit.ifelse(a == N+1, 1, ModelingToolkit.ifelse(a == 0, N, a))
@@ -59,3 +59,85 @@ sol = solve(prob,BFGS())
 
 sol = solve(prob,Newton())
 @test sol.minimum < 1e-8
+
+## SIR System Regression Test
+
+β = 0.01# infection rate
+λ_R = 0.05 # inverse of transition time from  infected to recovered
+λ_D = 0.83 # inverse of transition time from  infected to dead
+i₀ = 0.075 # fraction of initial infected people in every age class
+𝒫 = vcat([β, λ_R, λ_D]...)
+
+# regional contact matrix and regional population
+
+## regional contact matrix
+regional_all_contact_matrix = [3.45536   0.485314  0.506389  0.123002 ; 0.597721  2.11738   0.911374  0.323385 ; 0.906231  1.35041   1.60756   0.67411 ; 0.237902  0.432631  0.726488  0.979258] # 4x4 contact matrix
+
+## regional population stratified by age
+N = [723208 , 874150, 1330993, 1411928] # array of 4 elements, each of which representing the absolute amount of population in the corresponding age class.
+
+
+# Initial conditions
+I₀ = repeat([i₀],4)
+S₀ = N.-I₀
+R₀ = [0.0 for n in 1:length(N)]
+D₀ = [0.0 for n in 1:length(N)]
+D_tot₀ = [0.0 for n in 1:length(N)]
+ℬ = vcat([S₀, I₀, R₀, D₀, D_tot₀]...)
+
+# Time
+final_time = 20
+𝒯 = (1.0,final_time);
+
+
+
+
+function SIRD_ac!(du,u,p,t)
+    # Parameters to be calibrated
+    β, λ_R, λ_D = p
+
+    # initialize this parameter (death probability stratified by age, taken from literature)
+
+    δ₁, δ₂, δ₃, δ₄ = [0.003/100, 0.004/100, (0.015+0.030+0.064+0.213+0.718)/(5*100), (2.384+8.466+12.497+1.117)/(4*100)]
+    δ = vcat(repeat([δ₁],1),repeat([δ₂],1),repeat([δ₃],1),repeat([δ₄],4-1-1-1))
+
+
+    C = regional_all_contact_matrix
+
+
+    # State variables
+    S = @view u[4*0+1:4*1]
+    I = @view u[4*1+1:4*2]
+    R = @view u[4*2+1:4*3]
+    D = @view u[4*3+1:4*4]
+    D_tot = @view u[4*4+1:4*5]
+
+    # Differentials
+    dS = @view du[4*0+1:4*1]
+    dI = @view du[4*1+1:4*2]
+    dR = @view du[4*2+1:4*3]
+    dD = @view du[4*3+1:4*4]
+    dD_tot = @view du[4*4+1:4*5]
+
+    # Force of infection
+    Λ = β*[sum([C[i,j]*I[j]/N[j] for j in 1:size(C)[1]]) for i in 1:size(C)[2]]
+
+    # System of equations
+    @. dS = -Λ*S
+    @. dI = Λ*S - ((1-δ)*λ_R + δ*λ_D)*I
+    @. dR = λ_R*(1-δ)*I
+    @. dD = λ_D*δ*I
+    @. dD_tot = dD[1]+dD[2]+dD[3]+dD[4]
+
+
+end;
+
+
+# create problem and check it works
+problem = ODEProblem(SIRD_ac!, ℬ, 𝒯, 𝒫)
+@time solution = solve(problem, Tsit5(), saveat = 1:final_time);
+
+problem = ODEProblem(SIRD_ac!, ℬ, 𝒯, 𝒫)
+sys = modelingtoolkitize(problem)
+fast_problem = ODEProblem(sys,ℬ, 𝒯, 𝒫 )
+@time solution = solve(fast_problem, Tsit5(), saveat = 1:final_time)