! Pandemic/Epidemiology model in LINGO.                           (Pandemic.lng)
! Model the effect that the number of new infections in a population
  segment depend upon the infection fraction in other population segments.
  The possible states are: susceptible, infected, recovered, deceased;
! Keywords: Bass model, Chain reaction, Contagion, COVID-19, Diffusion,
   Epidemic, Multi-period, Pandemic, SIR model;
SETS:
  segment: susceptbl0, infected0, recoverd0, decsd0,
           recrate, dratei, drate, spnrate;
  sxs( segment, segment) : prifct;
  period: totalinf, tindex;
  pxs( period, segment) : susceptbl, infected, recoverd, decsd;
ENDSETS
DATA:
! Note, you get essentially the Bass model if either:
  there is only 1 segment, or they are all identical, i.e.,
    spnrate( i) = spnrate( j) for all i and j, the spontaneous,
                  infection rate,
    dratei = drate = 0, i.e., no one leaves a population by death,
    prifct( i, j) = prifct( i1, j1) for all i, j, i1, j1, and
    recrate = 0 for all populations, 
                i.e, once in infected state you stay there,;

  period = 1..500;
!             0-5  6-18  19-30 31-60 a61-above ;
  segment =  age05 age18 age30 age60 age99;
! Initial population;
  susceptbl0= 100   200   200   250    50;
! Initial infections;
  infected0=   0     0     1     0     0; 
! Initial population of people that have recovered;
  recoverd0=   0     0     0     0     0; 
! Initial population of number deceased;
  decsd0 =    0     0     0     0     0;
! Recovery rate/period for infected;
   recrate = 0.04  0.07   0.07   0.07   0.05;
! Death rate for infected. Lower for young;
   dratei = 0.0001  0.001  0.001  0.001  0.01;
! Death rate for not infected;
   drate = 0.0001 0.0001 0.0001 0.0001 0.0001;
! Spontaneous infection rate;
   spnrate =   0     0     0      0      0;
! Infection probability matrix.
  prifct( i, j) = Prob( susceptbl person  j gets infected in a period by infected person in group i
    if every person in i is infected.
   Expect that infection rate within a segment > rate between segments;
  prifct = 
        0.05   0.002  0.002  0.002  0.002 
        0.002  0.05   0.002  0.002  0.002
        0.001  0.03   0.04   0.03   0.002
        0.002  0.002  0.003  0.05   0.003
        0.002  0.002  0.002  0.002  0.03 ;
ENDDATA
CALC:
! Initialization;
 @FOR( segment( s):
   susceptbl( 1, s) = susceptbl0;
   infected( 1, s) = infected0;
   recoverd( 1, s) = recoverd0;
   decsd( 1, s)    = decsd0;
     );
! Compute state in t+1 from state in t;
 @FOR( period( t) | t #lt# @SIZE( period) :
   @FOR( segment( s):
! Compute new infections. It is proportional to 
    number susceptible and number of infections. Probability of infection
    from group i depends upon what fraction of group i is infected.
    Prob( infection) = 1 - Prob( not getting infected).
    prnot = Prob( not getting infected) =
    Pr(not infected by Grp1)*Pr(not infected by Grp2) ...
    Prob( infected by i) = prifct(i,s)*(fraction of i infected);
    prnot = 1;
     @FOR( segment( i):
! Prob( s not infected by any of segment 1, 2, ...).
   This is similar to the Bass Marketing diffusion model;
       prnot = prnot*( 1 - prifct(i, s)*infected( t, i)/
         ( susceptbl( t, i) + infected( t, i) +  recoverd( t, i) +  decsd( t, i));
         );
        );
! Compute new still susceptible;
    susceptbl( t+1, s) = susceptbl( t, s) 
         - drate( s) * susceptbl( t, s)   ! Deaths;
         - spnrate( s) * susceptbl( t, s)  ! Spontaneous infections;
         - ( 1- prnot) * susceptbl( t, s) ;! Infections by infected;
! Compute new infected;
    infected( t+1, s) = infected( t, s)
         - dratei( s) * infected( t, s)  ! Deaths of infected;
         - recrate( s) * infected( t, s) ! Recoveries;
         + spnrate( s) * susceptbl( t, s) ! Spontaneous infections;
         + ( 1- prnot) * susceptbl( t, s);! Infections by infected;
! Compute new recovered;
   recoverd( t+1, s) = recoverd( t, s) 
         + recrate( s) * infected( t, s) ! Recoveries;
         - drate( s) * recoverd( t, s); ! Deaths of recovered;
! Compute new deceased;
   decsd( t+1, s) = decsd( t, s) 
         + drate( s) * susceptbl( t, s)  ! Deaths of susceptbl;
         + dratei( s) * infected( t, s)  ! Deaths of infected;
         + drate( s) * recoverd( t, s); ! Deaths of recovered;
     );
   );

! Display a little report;
  @FOR( period( t):
    total = 0;
    @WRITE( @FORMAT( t, '4.0f'),') ');
    @FOR( segment( s):
      total = total + infected( t, s);
      @WRITE( @FORMAT( infected( t, s), '10.1f'));
        );
      @WRITE(' Total= ', @FORMAT( total, '10.1f'), @NEWLINE( 1));
      totalinf( t) = total;
      tindex( t) = t;  ! Give ourselves a x/time index;
      );
! Generate some charts of number infected over time;
   @CHARTCURVE( 'Number infected vs. Time, All Segments', 'Time', 'Current Infections', '', tindex, totalinf);
   @FOR( segment( s):
     @FOR( period( t): totalinf(t) = infected( t, s));
     @CHARTCURVE( 'Number infected in segment '+segment(s)+' vs. Time', 'Time', 'Current Infections in Segment ', '', tindex, totalinf);
         );  
ENDCALC