Squareville: Wastewater Surveillance for COVID-19

breed [persons person]
breed [viruses virus]
breed [plants plant]


persons-own [
  home-x
  home-y
  incubation-period          ;each person is assigned an incubation period at the beginning. This variable is used to determine how soon the people change from exposed to infectious state.
  incubation-period-value    ;this variable stores the original incubation period and is used to check the average incubation period.
  infectious-period          ;This variable is used to track people's disease periods
  contacts                   ;number of susceptible contacts in defined infection radius
  infected-by-me             ;total number of people I infect (numerical results)
  infect-any?                ;whether I have infected any one (true/false)
]

viruses-own [
  viability-period   ;viral viability in treatment plant
  in-plant?          ;check if the virus has been in the plant  (true/false)
]

Patches-own [
  p-type      ;define patch type  0 = town/community, 1 = pipe, 2 = plant, 3 = wild area, 4 = creek, 5 = testing location
]

Globals [
  main-pipe-xcor        ;the most right vertical pipe
  main-pipe-ycor        ;the lowest horizental pipe
  max-daily-cases
  infection-number      ;total of people are infected
  infector-number       ;total of people who infected at least one other person
  virus-yield           ;record the number of viruses that are produced
  start-record?         ;trigger data recording
  prevalence-list       ;This list stores the cumulative infected (include both E and I states) once local survaillence is triggered.
  RNA-copies            ;Coronvirous RNA copies
  individual-RNA-contribution    ;Copies of RNA increase by each person
  R0-list               ;create a list of ongoing R0
  death-toll

]

to-report R0
  if any? persons with [infect-any? = true]
    [report infection-number / infector-number]   ;This is the effective R0. A cross check is the total infected (E and I and R) dvided by total people with [infect-any? = true]
end 

to-report c
  report mean [contacts] of persons with [color = 26]   ;only E people may infect others so only count the susceptibles of E
end 

to-report infectious-period-of-exposed
 report mean [incubation-period-value] of persons with [color = 26]   ;avergage incubation periods of exposed people .
end 

to-report R-t
  report (0.01 * transmission-rate) * c *  infectious-period-of-exposed   ;only E may infect others so use incubation period rather than disease period to calculate Rt
end 

to-report sewer-patches
  report patches with [p-type = 1]
end 

to-report susceptible-people
  report persons with [color = 55]
end 

to-report exposed-people
  report persons with [color = 26]
end 

to-report infectious-people           ;these are symptomatic people. they are quarantined or isolated so they cannot infect others.
  report persons with [color = red]
end 

to-report recovered-people
  report persons with [color = sky]
end 

to-report local-surveillance-infection     ;report the first prevalence in the prevalence list
  if ticks > 0 and not empty? prevalence-list
   [report item 0 prevalence-list]
end 

to-report central-surveillance-infection     ;report the last prevalence in the prevalence list
  if ticks > 0 and length prevalence-list = Central-testing-duration
   [report item (central-testing-duration - 1) prevalence-list]
end 

;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;
;   Setup procedures
;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;

to setup
  clear-all
  set-default-shape viruses "c-virus-0"
  set-default-shape persons "person-0"

  set-patches
  setup-people
  add-an-carrier
  find-contacts

  set start-record? false
  set prevalence-list []
  set R0-list []
  set death-toll 0

  reset-ticks
end 

to set-patches
  ask patches [set pcolor white set p-type 0 ]  ;set background

  ask patches with [pycor <= 40] [set pcolor 9 set p-type 3] ;set wild field

  ask patches with [pycor <= min-pxcor + 4] [set pcolor 105 set p-type 4]  ;set a creek

  ;set horizontal pipe
  ask patches
   [if remainder (pycor - 2) 25 = 0 and pycor > 40 and pycor <= max-pycor and pxcor >= 2 and pxcor <= max-pxcor - 2
   [set pcolor 89 set p-type 1]]

  ;set vertical pipes
  ask patches
  [if remainder (pxcor - 2) 35 = 0 and pxcor >= 10 and pxcor <= max-pxcor and pycor <= max-pycor - 2 and pycor >= 45
   [set pcolor 89 set p-type 1]]

  ;connect to plant
  if any? patches with [p-type = 1]
  [let y min [pycor] of patches with [p-type = 1]
   let lowest-patch patches with [p-type = 1 and pycor = y]
   ;let connecting-patch lowest-patch with-max [pxcor]
   let x max [pxcor] of lowest-patch

   ask patches with [pxcor = x and pycor = y] [set pcolor 88]   ;label the correct connecting point

   ask patches with [pxcor = x and pycor < y and pycor >= min-pycor + 20 ] [set pcolor 89 set p-type 1]   ;extend the pipline


   ;set the plant
   create-plants 1
   [set shape "w-plant"
    set size 14
    set color 89
    setxy x 11.5]

   set main-pipe-xcor x    ;record the main-pipe xcor
   set main-pipe-ycor y    ;record the lowest pipe's ycor

    ask patches with [pxcor = main-pipe-xcor and pycor = 20] [set pcolor red]

  ]
end 

to setup-people
  ask n-of population patches with [p-type = 0]
  [sprout-persons 1
    [set size 2.75
     set color 55
     set shape "person-0"
     set home-x xcor
     set home-y ycor
     set incubation-period precision (random-normal 6 2) 0  ;set an incubation period based on normal distribution a mean of 6 and a SD of 2
     set incubation-period-value incubation-period   ;set an incubation period value for tracking the change in incubation period
     set infectious-period 0
     set infected-by-me 0
     set infect-any? false
     ]
  ]

  ;set individual RNA shed contribution based on population size. Equation from Data from Marc Johnson in MO 9The 10^9 copies of RNA per infected person gets diluted into the bulk wastewater.)
  set individual-RNA-contribution (1000 / (1.1355 * population - 900))  ;900 is used to adjust the function for running a smaller population
end 

to add-an-carrier
  ask one-of persons
  [set color 26 ]
end 



;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;
;   Go procedures
;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;

to go

  if ticks > 0 and (not any? exposed-people and not any? infectious-people) [stop]

  set RNA-copies 0       ;reset the number of RNA-copies at the testing site each tick
  move-viruses    ;put move-virus before produce-viruses make the newly produced viruses arrive the plant the next day or later.
  move-people

 (ifelse
    Issue-lockdown? = "none"    ;This is for comparing local and certral surveillance
      [SEIR]

   Issue-lockdown? = "local"    ;This is for gathering LOCAL alert prevalence
    [if length prevalence-list = 1
      [set Transmission-rate 1
        set %-travelers 5]

        SEIR]

    Issue-lockdown? = "central"    ;This is for gathering CENTRAL alert prevalence
    [if length prevalence-list = Central-testing-duration
      [set Transmission-rate 1
        set %-travelers 5]

       SEIR])

  find-contacts
  find-max-daily-cases
  calculate-average-R0

  surveillance
  disinfected
  back-home

  plot-infection-over-R0

 tick
end 

to SEIR
  incubation    ;E->I     ;put incubation before infect to update incubation period in the next day
  infect        ;S->E
  produce-viruses
  recover       ;I->E
end 

to move-people
 ask up-to-n-of ((count persons with [color != red]) * %-travelers * 0.01) persons with [color != red]   ;let up to 10% of the people who are not symptomatics to randomly move
  [setxy random-xcor (41 + random (max-pycor - 41))]
end 

to move-viruses              ;viruses arrive plant in a day
  ask viruses with [xcor = main-pipe-xcor and ycor = 20]   ;move the virues at the testing site down one patch to empty the testing site
    [setxy main-pipe-xcor 19]

  ask viruses with [xcor != main-pipe-xcor or ycor > 20]   ;move the newly produced viruses to the testing site
    [setxy main-pipe-xcor 20 set in-plant? true]
end 

to produce-viruses

    ask exposed-people                ;ask exposed-infected people
    [let sewer min-one-of sewer-patches [distance myself]  ;find the closest sewer
     if random 100 < 43
     [ hatch-viruses 1                               ;produce a new virus at a chance of 43%
       [set size 0.75                                 ;set the size of a virus
        set color 35
        setxy [pxcor] of sewer  [pycor] of sewer
        set viability-period 1
        set virus-yield virus-yield + 1
        set in-plant? false
       ]
     ]]
end 

to infect    ; S->E

  ask exposed-people; with [contacts > 0]
      [if any? susceptible-people in-radius infection-radius           ;if any susceptible people nearby
       [repeat count susceptible-people in-radius infection-radius     ;repeat the times of susceptible contacts to contact all these contacts
          [if random 1000 < (Transmission-Rate  * 10)                      ;at the chance of transmission rate
             [let infected-person one-of susceptible-people in-radius infection-radius      ;pick one of the susceptible people in the infection radius
              ask infected-person                               ;infect him/her
                 [set color 26]                                 ;change color. No need to set the infection-period as it is updated in recover procedure

              set infected-by-me infected-by-me + 1             ;update how many people I have infected
              if infect-any? = false                            ;if it is new infector
                [set infect-any? true                           ;update infection record
                 set infector-number infector-number + 1]       ;update total number of of infectors


              set infection-number infection-number + 1]]]]     ;update total number of infection
end 

to incubation     ;E->I

  ask exposed-people
     [ifelse incubation-period > 0
        [set incubation-period (incubation-period - 1)]
        [set color red]]
end 

to recover   ;I->R
  ask infectious-people
  [ifelse infectious-period >= 14    ; set disease period as 2 weeks. This period does not affect R0 as it is assumed that infected people are well isolated and do not transmit the disease.
    [ifelse random 100 < 3           ; mortality equals 3%
      [set death-toll death-toll + 1 die]
      [set color sky set infectious-period 0]]
    [set infectious-period infectious-period + 1]
  ]
end 



;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;
;   supporting procedures
;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;

to find-max-daily-cases
  if count infectious-people > max-daily-cases        ;Count the infectious.If it is greater than the current record of max daily cases
  [set max-daily-cases count infectious-people]       ;update the max daily case
end 

to disinfected

  ask viruses with [xcor = main-pipe-xcor and ycor = 19]    ; ask viruses in the treatment plant to die after 1 days ro save compute power.
    [ifelse viability-period > 0
      [set viability-period viability-period - 1]
      [die]]
end 

to find-contacts
  ask persons
  [set contacts count other persons with [color = 55] in-radius infection-radius]
end 

to surveillance
  if (count viruses-on patch main-pipe-xcor 20) * individual-RNA-contribution  >= Surveillance-threshold
    [set start-record? true]

 if start-record? = true
    [if length prevalence-list < Central-testing-duration
        [ set prevalence-list lput (precision (infection-number * 100 / Population) 2) prevalence-list]]
end 

to plot-infection-over-R0   ;;; this creates creates the scotter plot graph

    set-current-plot "Prevalence under Local vs. Central Surveillance"
    clear-plot
    plot-pen-down
  set-current-plot-pen "Day 1(Local)"
    if any? persons with [infect-any? = true] and length prevalence-list >= 1 [plotxy 1 item 0 prevalence-list]
  set-current-plot-pen "Day 2"
    if any? persons with [infect-any? = true] and length prevalence-list >= 2 [plotxy 2 item 1 prevalence-list]
  set-current-plot-pen "Day 3"
    if any? persons with [infect-any? = true] and length prevalence-list >= 3 [plotxy 3 item 2 prevalence-list]
  set-current-plot-pen "Day 4(Central)"
    if any? persons with [infect-any? = true] and length prevalence-list >= 4 [plotxy 4 item 3 prevalence-list]
  set-current-plot-pen "Day 5(Central)"
    if any? persons with [infect-any? = true] and length prevalence-list >= 5 [plotxy 5 item 4 prevalence-list]
  set-current-plot-pen "Day 6(Central)"
    if any? persons with [infect-any? = true] and length prevalence-list >= 6 [plotxy 6 item 5 prevalence-list]
  set-current-plot-pen "Day 7(Central)"
    if any? persons with [infect-any? = true] and length prevalence-list >= 7 [plotxy 7 item 6 prevalence-list]
end 

to calculate-average-R0
 if any? exposed-people with [infect-any? = true ] [set R0-list lput (precision R0 2) R0-list]
end 

to back-home
  ask persons with [xcor != home-x or ycor != home-y]
  [setxy home-x home-y]
End