AlgebraicJulia / StockFlow.jl

https://algebraicjulia.github.io/StockFlow.jl/
MIT License
65 stars 6 forks source link

Stock and flow macro for easier definition of Stock and Flow models #15

Closed Saityi closed 1 year ago

Saityi commented 1 year ago

Description

In the target branch of this pull request, functionStructureInSchema, a new data type, StockAndFlowF, has been added to allow for the flow equations of a Stock and Flow model to be captured in the model itself instead of as separate Julia equations. This gives rise to a new difficulty in definition of models: it is very tedious to specify the equations in this new form. One must manually write out the equations as symbolic expressions of tuples, like so:

SIR = StockAndFlowF(
    # stocks
    (:S => (:F_NONE, :inf, :N), :I => (:inf, :rec, :N), :R => (:rec, :F_NONE, :N)),
    # parameters
    (:c, :beta, :tRec),
    # dynamical variables
    (   :v_prevalence => ((:I, :N) => :/),
        :v_meanInfectiousContactsPerS => ((:c, :v_prevalence) => :*),
        :v_perSIncidenceRate => ((:beta, :v_meanInfectiousContactsPerS) => :*),
        :v_newInfections => ((:S, :v_perSIncidenceRate) => :*),
        :v_newRecovery => ((:I, :tRec) => :/),
    ),
    # flows
    (:inf => :v_newInfections, :rec => :v_newRecovery),
    # sum dynamical variables
    (:N),
)

Notable difficulties with this syntax are:

Changes

Add a macro stock_and_flow in a new module Syntax which allows Stock and Flow models to be defined using a more configuration language style syntax, e.g.

SIR = @stock_and_flow begin
    :stocks
      S
      I
      R

    :parameters
      c
      beta
      tRec

    :dynamic_variables
      v_prevalence = I / N
      v_meanInfectiousContactsPerS = c * v_prevalence
      v_perSIncidenceRate = beta * v_meanInfectiousContactsPerS
      v_newInfections = S * v_perSIncidenceRate
      v_newRecovery = I / tRec

    :flows
      S => inf(v_newInfections) => I
      I => rec(v_newRecovery) => R

    :sums
      N = [S, I, R]
end

This new syntax has a couple 'syntactic sugar' features:

As an example, see this version of the model which is the same as the other two above:

SIR = @stock_and_flow begin
    :stocks
      S
      I
      R

    :parameters
      c
      beta
      tRec

    # We can leave out dynamic variables and 
    # let them be inferred from flows entirely!

    :flows
      S => inf(S * beta * (c * (I / N))) => I
      I => rec(I / tRec) => R

    :sums
      N = [S, I, R]
end

Testing

When written in the longer form (see the second example) so that all of the variable names match, you will find the first two models are ==, partially confirming the transformations performed.

When written in the shorter form, the variables are all inferred, so it's less clear it's the same model. There may be bugs in the implementation.

I would like to request testing suggestions from @Xiaoyan-Li especially because this has not been thoroughly tested and may contain many bugs.

~TODOs / Discussion~

Note on differences

You may notice in the two Graph diagrams above a strange arrow in the new version not in the original: this is an artifact of the 'binop' stage, but should be the same model. Specifically:

That is, v_newInfections = S * v_perSIncidenceRate == v_newInfections = S * beta * c * (I / N).

As far as I can tell, despite the differences in the appearance of the models when Graph'd, they are the same model encoded slightly differently.

Further examples

Dr Nathaniel Osgood gave, as an example of the desired syntax, this model definition:

SIR_3 = @stock_and_flow begin
    :stocks
    S
    I
    R

    :parameters
    c
    beta
    tRec
    omega
    alpha

    :dynamic_variables
    v_prevalence = I / totalPopulation
    v_forceOfInfection = c * v_prevalence * beta

    :flows
    S => inf(S * v_forceOfInfection) => I
    ☁ => births(totalPopulation * alpha) => S
    S => deathsS(S * omega) => ☁
    I => rec(I / tRec) => R
    I => deathsI(I * omega) => ☁
    R => deathsR(R * omega) => ☁

    :sums
    totalPopulation = [S, I, R]
end

This syntax now defines an SIR model when used this macro, and when Graph'd, looks like:

sir_3

This may not be exactly right but I'm not sure how to verify the changes (see the testing section of this PR).

FYI

@Xiaoyan-Li @ndo885

Saityi commented 1 year ago

Dr Osgood (@ndo885) suggested (in a 1-1 meeting) that generated dynamic variables for flows use the flow name as a base for gensym. This has been added in 4d63d06 ; I plan to do the same for dynamic variable expansion. Interim variables will use the original dynamic variable's name.

Saityi commented 1 year ago

d026566 addresses the dynamic variables also being used as a base for generated variable names.

Updated screenshots of SIR models via Graph:

Screenshot from 2023-02-13 13-25-50

(Inferred dyvars example)

Screenshot from 2023-02-13 13-15-26

(Idea # 3 example)

fyi @ndo885 @Xiaoyan-Li