text: import numpy.random as npr
  import numpy as np
  from scipy.stats import expon
  from scipy.stats._distn_infrastructure import rv_frozen
  from covid19.utils import make_lognormal_from_interval, EmpiricalDistribution
  
  
  class SEIRBayes:
      ''' Model with Susceptible, Exposed, Infectious and Recovered compartments.
  
      This class implements the SEIR model with stochastic incubation and
      infectious periods as well as the basic reproduction number R0. 
  
      This model is an implicit density function on 4 time series S(t), E(t), 
      I(t) and R(t) for t = 0 to t_max-1. Sampling is done via numerical 
      resolution of a system of stochastic differential equations with 6 
      degrees of randomness: alpha, gamma, r0 and the number of subjects 
      transitioning between compartments; S -> E, E -> I, I -> R.
  
      Infectious (1/gamma) and incubation (1/alpha) periods, as well as basic 
      reproduction number r0, can be specified in 3 ways: 
          * 4-tuple as (lower bound, upper bound, density value, dist. family);
          * SciPy distribution objects, taken from scipy.stats;
          * array-like containers such as lists, numpy arrays and pandas Series.
      Se the __init__ method for greater detail.
  
      The probability of an individual staying in a compartment up to time t
      is proportional to exp(-p*t), therefore the probability of leaving is
      1 - exp(-p*t). The rate p is different for each pair of source and 
      destination compartments. They are as follows
  
          ======== ============= ========== ============== 
           Source   Destination     Rate        Period     
          ======== ============= ========== ============== 
           S        E             beta*I/N   1/(beta*I/N)  
           E        I             alpha      1/alpha       
           I        R             gamma      1/gamma       
          ======== ============= ========== ============== 
  
      Since the above discussion is for a single individual, a Binomial 
      distribution with success rate 1 - exp(-p*t) is used to scale to the
      total number of individuals in each compartment.
  
      Attributes:
          params (dict): Summary of model parameter values. Is compatible with
              the constructor; SEIRBayes(*params).
          _params (dict): Similar to params, but with modifications to
              facilitate usage internally. Isn't compatible with the constructor.
  
      Examples:
          Default init.
          
          >>> np.random.seed(0)
          >>> model = SEIRBayes(t_max=5)
          >>> S, E, I, R, t_space = model.sample(3)
          >>> I
          array([[10., 10., 10.],
                 [15., 10.,  9.],
                 [17., 15.,  8.],
                 [22., 17., 12.],
                 [24., 18., 15.]])
          >>> model.params['r0_dist'].interval(0.95)
          (2.5, 6.0)
  
  
      '''
      def __init__(self, 
                   NEIR0=(100, 20, 10, 0),
                   r0_dist=(2.5, 6.0, 0.95, 'lognorm'),
                   gamma_inv_dist=(7, 14, 0.95, 'lognorm'),
                   alpha_inv_dist=(4.1, 7, 0.95, 'lognorm'),
                   fator_subr=10,
                   t_max=30):
          '''Default constructor method.
  
  
          Args:
              NEIR0 (tuple): Initial conditions in the form of 
                  (population size, exposed, infected, recovered). Notice that
                  S0, the initial susceptible population, is not needed as it 
                  can be calculated as S0 = N - fator_subr*(E0 + I0 + R0).
              fator_subr (float): Multiplicative factor of I0 and E0 to take
                  into account sub-reporting.
              t_max (int): Length of the time-series.
  
              r0_dist, alpha_inv_dist, and gamma_inv_dist can be specified as
              a tuple, scipy distribution, or array-like object.
                  tuple: (lower bound, upper bound, density, dist. family)
                  scipy dist: object from scipy.stats with rvs method
                  array-like: the i-th value will be used for the i-th sample
  
              r0_dist (object): basic reproduction number.
              alpha_inv_dist (object): incubation period.
              gamma_inv_dist (object): infectious period.
  
          Examples:
              >>> np.random.seed(0)
              >>> model = SEIRBayes(fator_subr=2)
              >>> model.params['fator_subr']
              2
              >>> model.params['r0_dist'].rvs(10)
              array([5.74313347, 4.23505111, 4.81923138, 6.3885136 , 5.87744241,
                     3.11354468, 4.7884938 , 3.74424985, 3.78472191, 4.24493851])
              >>> model.params['NEIR0']
              (100, 20, 10, 0)
              >>> model._params['init_conditions']
              (40, 40, 20, 0)
              >>> model._params['total_population']
              100
  
              >>> np.random.seed(0)
              >>> model = SEIRBayes(r0_dist=(1.96, 5.1, 0.99, 'lognorm'),
              ...                   alpha_inv_dist=[5.1, 4.9, 6.0])
              >>> model.params['r0_dist'].rvs(10)
              array([4.38658658, 3.40543675, 3.79154822, 4.79256981, 4.4716837 ,
                     2.63710471, 3.7714376 , 3.07405105, 3.10164345, 3.41204358])
              >>> model.params['alpha_inv_dist'].rvs(3)
              array([5.1, 4.9, 6. ])
              
          '''
          r0_dist = self.init_param_dist(r0_dist)
          alpha_inv_dist = self.init_param_dist(alpha_inv_dist)
          gamma_inv_dist = self.init_param_dist(gamma_inv_dist)
  
          self.params = {
              'NEIR0': NEIR0,
              'r0_dist': r0_dist,
              'gamma_inv_dist': gamma_inv_dist,
              'alpha_inv_dist': alpha_inv_dist,
              'fator_subr': fator_subr,
              't_max': t_max
          }
  
          N, E0, I0, R0 = NEIR0
          S0 = N - (100/fator_subr)*(I0 + E0 + R0)
  
          self._params = {
              'init_conditions': (S0, (100/fator_subr)*E0, (100/fator_subr)*I0, (100/fator_subr)*R0),
              'fator_subr': fator_subr,
              'total_population': N,
              'alpha_inv_dist': alpha_inv_dist,
              'gamma_inv_dist': gamma_inv_dist,
              'r0_dist': r0_dist,
              't_max': t_max,
              'param_samples': {}
          }
  
      @classmethod
      def init_param_dist(cls, param_init):
          '''Initialize distribution from tuple, scipy or array-like object.
  
          Args:
              param_init (tuple, scipy.stats dist., or array-like)
  
          Examples:
              >>> np.random.seed(0)
              >>> dist = SEIRBayes.init_param_dist((1, 2, .9, 'lognorm'))
              >>> dist.interval(0.9)
              (1.0, 2.0)
  
              >>> dist = SEIRBayes.init_param_dist([0.1, 0.2, 0.3])
              >>> dist.rvs(2)
              array([0.1, 0.2])
  
              >>> from scipy.stats import lognorm
              >>> dist = SEIRBayes.init_param_dist(lognorm(s=.1, scale=1))
              >>> dist.mean()
              1.005012520859401
  
          '''
          if isinstance(param_init, tuple):
              lb, ub, density, family = param_init
              if family != 'lognorm':
                  raise NotImplementedError('Only family lognorm '
                                            'is implemented')
              dist = make_lognormal_from_interval(lb, ub, density)
          elif isinstance(param_init, rv_frozen):
              dist = param_init
          else:
              dist = EmpiricalDistribution(param_init)
          return dist
  
  
      def sample(self, size=1, return_param_samples=False):
          '''Sample from model.
          Args:
              size (int): Number of samples.
              return_param_samples (bool): If true, returns the parameter
                  samples (taken from {r0,gamma,alpha}_dist) used.
  
          Examples:
  
              >>> np.random.seed(0)
              >>> model = SEIRBayes(t_max=5)
              >>> S, E, I, R, t_space = model.sample(3)
              >>> S.shape, E.shape, I.shape, R.shape
              ((5, 3), (5, 3), (5, 3), (5, 3))
              >>> I
              array([[10., 10., 10.],
                     [15., 10.,  9.],
                     [17., 15.,  8.],
                     [22., 17., 12.],
                     [24., 18., 15.]])
              >>> t_space
              array([0, 1, 2, 3, 4])
  
              Return parameter samples for analysis.
  
              >>> np.random.seed(0)
              >>> model = SEIRBayes(t_max=5)
              >>> (S, E, I, R, t_space, r0,
              ...  alpha, gamma, beta) = model.sample(5, True)
              >>> r0
              array([5.74313347, 4.23505111, 4.81923138, 6.3885136 , 5.87744241])
              >>> alpha
              array([0.18303002, 0.15306351, 0.16825044, 0.18358956, 0.17569263])
              >>> gamma
              array([0.12007063, 0.08539356, 0.10375533, 0.1028759 , 0.09394099])
              >>> np.isclose(r0, beta/gamma).all()
              True
              >>> t_space
              array([0, 1, 2, 3, 4])
              
          '''
          t_space = np.arange(0, self._params['t_max'])
          N = self._params['total_population']
          S, E, I, R = [np.zeros((self._params['t_max'], size))
                        for _ in range(4)]
          S[0, ], E[0, ], I[0, ], R[0, ] = self._params['init_conditions']
  
          r0 = self._params['r0_dist'].rvs(size)
          gamma = 1/self._params['gamma_inv_dist'].rvs(size)
          alpha = 1/self._params['alpha_inv_dist'].rvs(size)
          beta = r0*gamma
  
          for t in t_space[1:]:
              SE = npr.binomial(S[t-1, ].astype(int),
                                expon(scale=1/(beta*I[t-1, ]/N)).cdf(1))
              EI = npr.binomial(E[t-1, ].astype(int),
                                expon(scale=1/alpha).cdf(1))
              IR = npr.binomial(I[t-1, ].astype(int),
                                expon(scale=1/gamma).cdf(1))
  
              dS =  0 - SE
              dE = SE - EI
              dI = EI - IR
              dR = IR - 0
  
              S[t, ] = S[t-1, ] + dS
              E[t, ] = E[t-1, ] + dE
              I[t, ] = I[t-1, ] + dI
              R[t, ] = R[t-1, ] + dR
  
          if return_param_samples:
              return S, E, I, R, t_space, r0, alpha, gamma, beta
          else:
              return S, E, I, R, t_space
  