feat: new macro @quickloop and @reaction_eq (#6)

* feat: new macro `@quickloop` and `@reaction_eq`

* Bump version to `0.1.4`

Author: wangl-cc

Project.toml

@@ -1,20 +1,23 @@
 name = "EvolutionaryModelingTools"
 uuid = "faf572ed-2b53-4324-a741-daa175e50348"
 authors = ["Long Wang <wangl.cc@outlook.com>"]
-version = "0.1.3"
+version = "0.1.4"
 
 [deps]
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecordedArrays = "1040807a-3a2f-4266-b2c1-805b33c7034a"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 
 [compat]
+LoopVectorization = "0.12.107"
 MacroTools = "0.5"
 RecordedArrays = "0.4.1"
 julia = "1.2"
 
 [extras]
+LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Test", "LoopVectorization"]

example/competition.jl

@@ -1,6 +1,7 @@
 using RecordedArrays
 using RecordedArrays.ResizingTools
 using Random
+using LoopVectorization
 using EvolutionaryModelingTools
 using EvolutionaryModelingTools: sample
 
@@ -13,13 +14,13 @@ const m = ones(1, 1)
 
 # simulate with this package
 ## define reactions
-@reaction growth begin
-    r * (1 - μ) * v
-    v[ind] += 1
-end
 
+# define growth with @reaction_eq, loop vectorization is disabled but inbounds, fastmath and  offset are enable
+@reaction_eq growth r * (1 - μ) v[i] → 2v[i] avx=false inbounds=true fastmath=true offset=true
+
+# define mutation with @reaction and @quickloop with given index
 @reaction mutation begin
-    @. r * μ * v
+    @quickloop r * μ * v[i]
     begin
         push!(v, 1)
         n = length(v)
@@ -30,11 +31,9 @@ end
     end
 end
 
+# define competition with @reaction and @quickloop
 @reaction competition begin
-    begin
-        c = r / K # baseline competition coefficient
-        @. c * v * m * v'
-    end
+    @quickloop (r / K) * v[i] * m[i, j] * v[j] turbo # force enable loop vectorization
     begin
         i = ind[1]
         v[i] -= 1

example/sir.jl

@@ -6,49 +6,89 @@ using EvolutionaryModelingTools.Scalar
 # parameters
 const T = 100.0
 const β = 0.001
-const γ = 0.01
+const ν = 0.01
+const α = 0.2
+const r = 0.5
+const d = 0.1
+const c = 0.001
 const S = 100
 const I = 1
 const R = 0
 
 ## simulation with this package
-@reaction infection begin
-    β * S * I
-    (S[] -= 1; I[] += 1)
-end
-
-@reaction recovery begin
-    γ * I
-    (I[] -= 1; R[] += 1)
+# epidemic dynamics
+@reaction_eq infection β S + I --> I + I
+@reaction_eq recovery ν I --> R
+@reaction_eq virulence α I --> 0 # death of infection host caused by virus
+# demography dynamics
+# generate reactions with @eval
+for sym in (:S, :I, :R)
+    r_name = Symbol(:growth_, sym)
+    d_name = Symbol(:death_, sym)
+    @eval @reaction_eq $r_name r $sym --> S + $sym # growth
+    @eval @reaction_eq $d_name d $sym --> 0 # death
+    for sym′ in (:S, :I, :R)
+        c_name = Symbol(:competition_, sym, sym′)
+        @eval @reaction_eq $c_name c $sym + $sym′ --> $sym′ # competition
+    end
 end
 
-run_gillespie(rng=Random.GLOBAL_RNG, T=T, β=β, γ=γ, S=S, I=I, R=R) = gillespie(
-    rng,
-    T,
-    (; β, γ, S=scalar(S), I=scalar(I), R=scalar(R)),
-    (infection, recovery),
+const REACTIONS = (
+    infection, recovery, virulence,
+    growth_S, growth_I, growth_R,
+    death_S, death_I, death_R,
+    competition_SS, competition_SI, competition_SR,
+    competition_IS, competition_II, competition_IR,
+    competition_RS, competition_RI, competition_RR
 )
 
-function run_gillespie_record(rng=Random.GLOBAL_RNG, T=T, β=β, γ=γ, S=S, I=I, R=R)
+run_gillespie(rng=Random.GLOBAL_RNG, T=T, β=β, ν=ν, α=α, r=r, d=d, c=c, S=S, I=I, R=R) =
+    gillespie(rng, T, (; β, ν, α, r, d, c, S=scalar(S), I=scalar(I), R=scalar(R)), REACTIONS)
+
+function run_gillespie_record(rng=Random.GLOBAL_RNG, T=T, β=β, ν=ν, α=α,
+        r=r, d=d, c=c, S=S, I=I, R=R)
     clock = ContinuousClock(T)
-    ps = (;
-        β,
-        γ,
+    ps = (;β, ν, α, r, d, c,
         S=recorded(DynamicEntry, clock, S),
         I=recorded(DynamicEntry, clock, I),
         R=recorded(DynamicEntry, clock, R),
     )
-    return gillespie(rng, clock, ps, (infection, recovery))
+    return gillespie(rng, clock, ps, REACTIONS)
 end
 
 ## simulation manually
-function run_manually(rng=Random.GLOBAL_RNG, T=T, β=β, γ=γ, S=S, I=I, R=R)
+### most of the code is generated by Github Copilot
+### a useful tool for generating code like this
+function run_manually(rng=Random.GLOBAL_RNG, T=T, β=β, ν=ν, α=α,
+        r=r, d=d, c=c, S=S, I=I, R=R)
     t = 0
     while t <= T
         # calculate rates
         infection_rate = β * S * I
-        recovery_rate = γ * I
-        summed = infection_rate + recovery_rate
+        recovery_rate = ν * I
+        virulence_rate = α * I
+        growth_S_rate = r * S
+        growth_I_rate = r * I
+        growth_R_rate = r * R
+        death_S_rate = d * S
+        death_I_rate = d * I
+        death_R_rate = d * R
+        competition_SS_rate = c * S * S
+        competition_SI_rate = c * S * I
+        competition_SR_rate = c * S * R
+        competition_IS_rate = c * I * S
+        competition_II_rate = c * I * I
+        competition_IR_rate = c * I * R
+        competition_RS_rate = c * R * S
+        competition_RI_rate = c * R * I
+        competition_RR_rate = c * R * R
+        # summed rate
+        summed = infection_rate + recovery_rate + virulence_rate +
+            growth_S_rate + growth_I_rate + growth_R_rate +
+            death_S_rate + death_I_rate + death_R_rate +
+            competition_SS_rate + competition_SI_rate + competition_SR_rate +
+            competition_IS_rate + competition_II_rate + competition_IR_rate +
+            competition_RS_rate + competition_RI_rate + competition_RR_rate
         # break if summed is zero
         summed == 0 && break
         # update current time
@@ -58,24 +98,79 @@ function run_manually(rng=Random.GLOBAL_RNG, T=T, β=β, γ=γ, S=S, I=I, R=R)
         if rn < infection_rate
             S -= 1
             I += 1
-        else
+        elseif (rn -= infection_rate) < recovery_rate
             I -= 1
             R += 1
+        elseif (rn -= recovery_rate) < virulence_rate
+            I -= 1
+        elseif (rn -= virulence_rate) < growth_S_rate
+            S += 1
+        elseif (rn -= growth_S_rate) < growth_I_rate
+            S += 1
+        elseif (rn -= growth_I_rate) < growth_R_rate
+            S += 1
+        elseif (rn -= growth_R_rate) < death_S_rate
+            S -= 1
+        elseif (rn -= death_S_rate) < death_I_rate
+            I -= 1
+        elseif (rn -= death_I_rate) < death_R_rate
+            R -= 1
+        elseif (rn -= death_R_rate) < competition_SS_rate
+            S -= 1
+        elseif (rn -= competition_SS_rate) < competition_SI_rate
+            S -= 1
+        elseif (rn -= competition_SI_rate) < competition_SR_rate
+            S -= 1
+        elseif (rn -= competition_SR_rate) < competition_IS_rate
+            I -= 1
+        elseif (rn -= competition_IS_rate) < competition_II_rate
+            I -= 1
+        elseif (rn -= competition_II_rate) < competition_IR_rate
+            I -= 1
+        elseif (rn -= competition_IR_rate) < competition_RS_rate
+            R -= 1
+        elseif (rn -= competition_RS_rate) < competition_RI_rate
+            R -= 1
+        else
+            R -= 1
         end
     end
     return (; t, S, I, R)
 end
 
-function run_manually_record(rng=Random.GLOBAL_RNG, T=T, β=β, γ=γ, S=S, I=I, R=R)
+function run_manually_record(rng=Random.GLOBAL_RNG, T=T, β=β, ν=ν, α=α,
+        r=r, d=d, c=c, S=S, I=I, R=R)
     clock = ContinuousClock(T)
     S′ = recorded(DynamicEntry, clock, S)
     I′ = recorded(DynamicEntry, clock, I)
     R′ = recorded(DynamicEntry, clock, R)
     for _ in clock
         # calculate rates
         infection_rate = β * S′ * I′
-        recovery_rate = γ * I′
-        summed = infection_rate + recovery_rate
+        recovery_rate = ν * I′
+        virulence_rate = α * I′
+        growth_S_rate = r * S′
+        growth_I_rate = r * I′
+        growth_R_rate = r * R′
+        death_S_rate = d * S′
+        death_I_rate = d * I′
+        death_R_rate = d * R′
+        competition_SS_rate = c * S′ * S′
+        competition_SI_rate = c * S′ * I′
+        competition_SR_rate = c * S′ * R′
+        competition_IS_rate = c * I′ * S′
+        competition_II_rate = c * I′ * I′
+        competition_IR_rate = c * I′ * R′
+        competition_RS_rate = c * R′ * S′
+        competition_RI_rate = c * R′ * I′
+        competition_RR_rate = c * R′ * R′
+        # summed rate
+        summed = infection_rate + recovery_rate + virulence_rate +
+            growth_S_rate + growth_I_rate + growth_R_rate +
+            death_S_rate + death_I_rate + death_R_rate +
+            competition_SS_rate + competition_SI_rate + competition_SR_rate +
+            competition_IS_rate + competition_II_rate + competition_IR_rate +
+            competition_RS_rate + competition_RI_rate + competition_RR_rate
         # break if summed is zero
         summed == 0 && break
         # update current time
@@ -85,9 +180,41 @@ function run_manually_record(rng=Random.GLOBAL_RNG, T=T, β=β, γ=γ, S=S, I=I,
         if rn < infection_rate
             S′[] -= 1
             I′[] += 1
-        else
+        elseif (rn -= infection_rate) < recovery_rate
             I′[] -= 1
             R′[] += 1
+        elseif (rn -= recovery_rate) < virulence_rate
+            I′[] -= 1
+        elseif (rn -= virulence_rate) < growth_S_rate
+            S′[] += 1
+        elseif (rn -= growth_S_rate) < growth_I_rate
+            S′[] += 1
+        elseif (rn -= growth_I_rate) < growth_R_rate
+            S′[] += 1
+        elseif (rn -= growth_R_rate) < death_S_rate
+            S′[] -= 1
+        elseif (rn -= death_S_rate) < death_I_rate
+            I′[] -= 1
+        elseif (rn -= death_I_rate) < death_R_rate
+            R′[] -= 1
+        elseif (rn -= death_R_rate) < competition_SS_rate
+            S′[] -= 1
+        elseif (rn -= competition_SS_rate) < competition_SI_rate
+            S′[] -= 1
+        elseif (rn -= competition_SI_rate) < competition_SR_rate
+            S′[] -= 1
+        elseif (rn -= competition_SR_rate) < competition_IS_rate
+            I′[] -= 1
+        elseif (rn -= competition_IS_rate) < competition_II_rate
+            I′[] -= 1
+        elseif (rn -= competition_II_rate) < competition_IR_rate
+            I′[] -= 1
+        elseif (rn -= competition_IR_rate) < competition_RS_rate
+            R′[] -= 1
+        elseif (rn -= competition_RS_rate) < competition_RI_rate
+            R′[] -= 1
+        else
+            R′[] -= 1
         end
     end
     return (; clock, S=S′, I=I′, R=R′)

src/EvolutionaryModelingTools.jl

@@ -1,15 +1,23 @@
 module EvolutionaryModelingTools
 
 using RecordedArrays
+using RecordedArrays.ArrayInterface: indices
 using Random
 using MacroTools
+using Requires
 
-export @cfunc, @ufunc, @reaction
-export Reaction, gillespie, gillespie!
+export @cfunc, @ufunc, @reaction, @quickloop, @reaction_eq
+export Reaction, gillespie, gillespie!, indices
 
 include("tools.jl")
 include("sample.jl")
 include("model.jl")
 include("scalar.jl")
 
+function __init__()
+    @require LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890" begin
+        DEFAULT_LOOP_MODE[1] = true
+    end
+end
+
 end # module