TypeError: non-boolean (Term{Bool}) used in boolean context when trying to modelingtoolkitize a DifferentialEquations.jl SDEProblem

[ClaudMor]

Hello,

I first posted this issue on the discourse , but maybe this is its natural place.

This issue is maybe related to this one although the solution could be totally different. If this is the wrong place to post it, please tell me and I’ll be happy to move it.

So the goal would be to optimize a DifferentialEquations.jl EnsembleProblem simulation built upon an SDEProblem with noise produced by distributions taken from Distributions.jl.
Here is a simplified version of the model:

using DifferentialEquations, Distributions, DiffEqBayes, AutoOptimize

# Model parameters

β = 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
σ_β = 0.01 
σ_R = 0.01 
σ_D = 0.01 

𝒫 = vcat([β, λ_R, λ_D,σ_β,σ_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₀ = 0.075 # fraction of initial infected people in every age class
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;

# define noise
function SIRD_ac_noise!(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
    Λ = rand(Normal(0.0, σ_β))*[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-δ)*rand(Normal( 0.0,σ_R)) + δ*rand(Normal( 0.0,σ_D)))*I
    @. dR = rand(Normal( 0.0,σ_R))*(1-δ)*I 
    @. dD = rand(Normal( 0.0,σ_D))*δ*I
    @. dD_tot = dD[1]+dD[2]+dD[3]+dD[4]
    

end;

# create problem and check it works
sde_problem = SDEProblem(SIRD_ac!,SIRD_ac_noise!,ℬ, 𝒯, 𝒫 )
solution = @time solve(sde_problem, saveat = 1:final_time);

Attempting to directly auto_optimize the EnsembleProblem using AutoOprimize.jl seems not supported:

ens_problem = EnsembleProblem(sde_problem)

MethodError: no method matching auto_optimize(::EnsembleProblem{SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float64,1},Nothing,SDEFunction{true,typeof(SIRD_ac!),typeof(SIRD_ac_noise!),LinearAlgebra.UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(SIRD_ac_noise!),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing})
Closest candidates are:
  auto_optimize(!Matched::ODEProblem) at C:\Users\claud\.julia\packages\AutoOptimize\29daN\src\AutoOptimize.jl:29
  auto_optimize(!Matched::ODEProblem, !Matched::Any; verbose, stiff, mtkify, sparsify, gpuify, static, gpup) at C:\Users\claud\.julia\packages\AutoOptimize\29daN\src\AutoOptimize.jl:29

Stacktrace:
 [1] top-level scope at In[4]:1
 [2] include_string(::Function, ::Module, ::String, ::String) at .\loading.jl:1091

So hoping to still ger some benefit at ens_problem simulation time, I tried to auto_optimize the sde_problem before building the ens_problem upon it, but this too seems not to be implemented:

auto_sde_problem = auto_optimize(sde_problem)

MethodError: no method matching auto_optimize(::SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float64,1},Nothing,SDEFunction{true,typeof(SIRD_ac!),typeof(SIRD_ac_noise!),LinearAlgebra.UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(SIRD_ac_noise!),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing})
Closest candidates are:
  auto_optimize(!Matched::ODEProblem) at C:\Users\claud\.julia\packages\AutoOptimize\29daN\src\AutoOptimize.jl:29
  auto_optimize(!Matched::ODEProblem, !Matched::Any; verbose, stiff, mtkify, sparsify, gpuify, static, gpup) at C:\Users\claud\.julia\packages\AutoOptimize\29daN\src\AutoOptimize.jl:29

Stacktrace:
 [1] top-level scope at In[5]:1
 [2] include_string(::Function, ::Module, ::String, ::String) at .\loading.jl:1091

Then trying to modelingtoolkitize the problem (as before hoping that doing it before building the ens_problem upon it returns a faster EnsembleProblem that building it directly on sde_problem would), I get:

sys = modelingtoolkitize(sde_problem)

TypeError: non-boolean (Term{Bool}) used in boolean context

Stacktrace:
 [1] macro expansion at C:\Users\claud\.julia\packages\Distributions\xT124\src\utils.jl:5 [inlined]
 [2] Normal(::Num, ::Num; check_args::Bool) at C:\Users\claud\.julia\packages\Distributions\xT124\src\univariate\continuous\normal.jl:37
 [3] Normal at C:\Users\claud\.julia\packages\Distributions\xT124\src\univariate\continuous\normal.jl:37 [inlined]
 [4] Normal at C:\Users\claud\.julia\packages\Distributions\xT124\src\univariate\continuous\normal.jl:42 [inlined]
 [5] SIRD_ac_noise!(::Array{Any,1}, ::Array{Num,1}, ::Array{Num,1}, ::Num) at .\In[1]:111
 [6] modelingtoolkitize(::SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float64,1},Nothing,SDEFunction{true,typeof(SIRD_ac!),typeof(SIRD_ac_noise!),LinearAlgebra.UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(SIRD_ac_noise!),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing}) at C:\Users\claud\.julia\packages\ModelingToolkit\HTjKG\src\systems\diffeqs\modelingtoolkitize.jl:71
 [7] top-level scope at In[2]:1
 [8] include_string(::Function, ::Module, ::String, ::String) at .\loading.jl:1091

I’m particularly puzzled by this last error, since ModelingToolkit seems to support SDEs.

So what am I missing?

environment:

(computationalEpi) pkg> st
Status `E:\IlMIoDrive\magistrale\2anno\primo_periodo\computationalEpi\Project.toml`
  [ff3c4d4f] AutoOptimize v0.1.0 `https://github.com/SciML/AutoOptimize.jl#master`
  [8f4d0f93] Conda v1.5.0
  [071ae1c0] DiffEqGPU v1.8.0
  [0c46a032] DifferentialEquations v6.15.0
  [7073ff75] IJulia v1.23.1
  [961ee093] ModelingToolkit v4.0.8
  [d330b81b] PyPlot v2.9.0

Thank you very much

[ChrisRackauckas]

Auto-optimize only supports ODEs and is very much still a work in progress. I'll look into the MTKitization of the SDE though. That's a real error.

[ChrisRackauckas]

We're just missing registrations of Distributions and random. Adding the function registrations before MTKitization fixes it:

ModelingToolkit.@register Normal(mu,sigma)
ModelingToolkit.@register Base.rand(x)
sys = modelingtoolkitize(sde_problem)

@shashi do you know of a good way to loop over distributions and register them all? @korsbo will need it anyways for distributed delays.

[ChrisRackauckas]

This is fixed in #687