Gradio APP

app.py

@@ -0,0 +1,98 @@
+import os
+import gradio as gr
+import numpy as np
+import matplotlib
+import matplotlib.pyplot as plt
+from src.sampler import mcmc_sampler
+from src.data_parser import Parser
+matplotlib.use('Agg')
+font = {'size': 30}
+matplotlib.rc('font', **font)
+
+
+def sample(country, d, n_iterations, burnin):
+    P = parser.parse_population(country)
+    start_date = "2020-03-01"
+    end_date = "2020-06-15"
+    i, r = parser.parse_data(start_date, end_date, country)
+    i, r = i.values, r.values
+    s = np.repeat(P, i.shape[0]) - i - r
+
+    p, lam, t, lam_ar, t_ar = mcmc_sampler(s, i, d, P, n_iterations, burnin,
+                                           M=3, sigma=0.01,
+                                           alpha=np.repeat(2, d),
+                                           beta=np.repeat(0.1, d),
+                                           a=1, b=1, phi=0.995)
+
+    lam_estimated = np.average(lam, axis=1)
+    t_estimated = np.average(t, axis=1)
+    p_estimated = np.average(p)
+
+    # Plot the series.
+    fig, axs = plt.subplots(nrows=2)
+    fig.set_figheight(30)
+    fig.set_figwidth(30)
+    ax1_left = axs[0]
+    ax2_left = axs[1]
+    ax1_right = ax1_left.twinx()
+    ax2_right = ax2_left.twinx()
+    ax1_left.plot(s, color='red', label="Susceptible")
+    ax1_right.plot(i, color='blue', label="Infected")
+    ax1_left.legend(loc=2)
+    ax1_right.legend(loc=1)
+    delta_i = -np.diff(s)
+    ax2_left.plot(delta_i, color="blue", label="Newly Infected Individuals")
+    ax2_right.plot(i, color='blue', linestyle='dashed', label="Infected")
+    ax2_left.legend(loc=2)
+    ax2_right.legend(loc=1)
+    # Display obtained breakpoints on plot.
+    for breakpoint in np.average(t, axis=1):
+        ax1_right.axvline(breakpoint, color="green")
+        ax2_right.axvline(breakpoint, color="green")
+
+    # Get output strings
+    lam_string = ""
+    for j, lam_component in enumerate(lam_estimated):
+        lam_string += f"Component {j+1}: {round(lam_component, 4)}\n"
+    lam_string = lam_string.rstrip()
+    t_string = ""
+    for j, t_component in enumerate(t_estimated):
+        t_string += f"Breakpoint {j+1}: {int(round(t_component, 0))}\n"
+    t_string = t_string.rstrip()
+    p_string = f"{round(p_estimated, 4)}"
+
+    return fig, lam_string, t_string, p_string
+
+
+if __name__ == "__main__":
+    confirmed_path = "confirmed.csv"
+    deaths_path = "deaths.csv"
+    recovered_path = "recovered.csv"
+    population_path = "population.csv"
+    data_path = os.path.join(os.getcwd(), "data")
+    parser = Parser(os.path.join(data_path, confirmed_path),
+                    os.path.join(data_path, deaths_path),
+                    os.path.join(data_path, recovered_path),
+                    os.path.join(data_path, population_path))
+    countries = parser.countries
+
+    # Inputs
+    dropdown = gr.Dropdown(choices=countries, value="Germany",
+                           label="Select the Country")
+    slider = gr.Slider(minimum=1, maximum=5, value=3, step=1,
+                       label="Select the Number of Breakpoints")
+    n_iterations = gr.Number(value=10000, precision=0,
+                             label="Select the Number of iterations")
+    burnin = gr.Number(value=1000, precision=0,
+                       label="Select the Number of Burn-In Iterations",
+                       info="Such iterations will be discarded.")
+
+    # Outputs
+    plot = gr.Plot(label="Results")
+    lam = gr.Text(label="Estimated Lambda")
+    t = gr.Text(label="Estimated Breakpoints")
+    p = gr.Text(label="Estimated Recovery Probability")
+    interface = gr.Interface(sample,
+                             inputs=[dropdown, slider, n_iterations, burnin],
+                             outputs=[plot, lam, t, p])
+    interface.launch()

data/population.csv

@@ -0,0 +1,266 @@
+Country,Population
+Aruba,106585
+Africa Eastern and Southern,685112705
+Afghanistan,38972230
+Africa Western and Central,466189102
+Angola,33428486
+Albania,2837849
+Andorra,77700
+Arab World,449228296
+United Arab Emirates,9287289
+Argentina,45376763
+Armenia,2805608
+American Samoa,46189
+Antigua and Barbuda,92664
+Australia,25655289
+Austria,8916864
+Azerbaijan,10093121
+Burundi,12220227
+Belgium,11538604
+Benin,12643123
+Burkina Faso,21522626
+Bangladesh,167420951
+Bulgaria,6934015
+Bahrain,1477469
+Bahamas,406471
+Bosnia and Herzegovina,3318407
+Belarus,9379952
+Belize,394921
+Bermuda,63893
+Bolivia,11936162
+Brazil,213196304
+Barbados,280693
+Brunei,441725
+Bhutan,772506
+Botswana,2546402
+Central African Republic,5343020
+Canada,38037204
+Central Europe and the Baltics,102180124
+Switzerland,8638167
+Channel Islands,171113
+Chile,19300315
+China,1411100000
+Cote d'Ivoire,26811790
+Cameroon,26491087
+"Congo, Dem. Rep.",92853164
+"Congo, Rep.",5702174
+Colombia,50930662
+Comoros,806166
+Cabo Verde,582640
+Costa Rica,5123105
+Caribbean small states,7444768
+Cuba,11300698
+Curacao,154947
+Cayman Islands,67311
+Cyprus,1237537
+Czechia,10697858
+Germany,83160871
+Djibouti,1090156
+Dominica,71995
+Denmark,5831404
+Dominican Republic,10999664
+Algeria,43451666
+East Asia & Pacific (excluding high income),2116424876
+Early-demographic dividend,-2147483648
+East Asia & Pacific,-2147483648
+Europe & Central Asia (excluding high income),400811771
+Europe & Central Asia,923103879
+Ecuador,17588595
+Egypt,107465134
+Euro area,342913447
+Eritrea,3555868
+Spain,47365655
+Estonia,1329522
+Ethiopia,117190911
+European Union,447692315
+Fragile and conflict affected situations,979418527
+Finland,5529543
+Fiji,920422
+France,67571107
+Faroe Islands,52415
+Micronesia,112106
+Gabon,2292573
+United Kingdom,67081000
+Georgia,3722716
+Ghana,32180401
+Gibraltar,32709
+Guinea,13205153
+Gambia,2573995
+Guinea-Bissau,2015828
+Equatorial Guinea,1596049
+Greece,10698599
+Grenada,123663
+Greenland,56367
+Guatemala,16858333
+Guam,169231
+Guyana,797202
+High income,1240900955
+"Hong Kong SAR, China",7481000
+Honduras,10121763
+Heavily indebted poor countries (HIPC),838066650
+Croatia,4047680
+Haiti,11306801
+Hungary,9750149
+IBRD only,-2147483648
+IDA & IBRD total,-2147483648
+IDA total,1738306807
+IDA blend,582637127
+Indonesia,271857970
+IDA only,1155669680
+Isle of Man,84046
+India,1396387127
+Ireland,4985382
+Iran,87290193
+Iraq,42556984
+Iceland,366463
+Israel,9215100
+Italy,59438851
+Jamaica,2820436
+Jordan,10928721
+Japan,126261000
+Kazakhstan,18755666
+Kenya,51985780
+Kyrgyzstan,6579900
+Cambodia,16396860
+Kiribati,126463
+Saint Kitts and Nevis,47642
+"Korea, South",51836239
+Kuwait,4360444
+Latin America & Caribbean (excluding high income),588808380
+Laos,7319399
+Lebanon,5662923
+Liberia,5087584
+Libya,6653942
+Saint Lucia,179237
+Latin America & Caribbean,650534967
+Least developed countries: UN classification,1073743450
+Low income,699186538
+Liechtenstein,38756
+Sri Lanka,21919000
+Lower middle income,-2147483648
+Low & middle income,-2147483648
+Lesotho,2254100
+Late-demographic dividend,-2147483648
+Lithuania,2794885
+Luxembourg,630419
+Latvia,1900449
+"Macao SAR, China",676283
+St. Martin (French part),32553
+Morocco,36688772
+Monaco,36922
+Moldova,2635130
+Madagascar,28225177
+Maldives,514438
+Middle East & North Africa,479966649
+Mexico,125998302
+Marshall Islands,43413
+Middle income,-2147483648
+North Macedonia,2072531
+Mali,21224040
+Malta,515332
+Burma,53423198
+Middle East & North Africa (excluding high income),411810124
+Montenegro,621306
+Mongolia,3294335
+Northern Mariana Islands,49587
+Mozambique,31178239
+Mauritania,4498604
+Mauritius,1265740
+Malawi,19377061
+Malaysia,33199993
+North America,369602177
+Namibia,2489098
+New Caledonia,271130
+Niger,24333639
+Nigeria,208327405
+Nicaragua,6755895
+Netherlands,17441500
+Norway,5379475
+Nepal,29348627
+Nauru,12315
+New Zealand,5090200
+OECD members,1370241530
+Oman,4543399
+Other small states,32381190
+Pakistan,227196741
+Panama,4294396
+Peru,33304756
+Philippines,112190977
+Palau,17972
+Papua New Guinea,9749640
+Poland,37899070
+Pre-demographic dividend,984213438
+Puerto Rico,3281538
+"Korea, Dem. People's Rep.",25867467
+Portugal,10297081
+Paraguay,6618695
+West Bank and Gaza,4803269
+Pacific island small states,2566819
+Post-demographic dividend,1117443485
+French Polynesia,301920
+Qatar,2760385
+Romania,19265250
+Russia,144073139
+Rwanda,13146362
+South Asia,1882531620
+Saudi Arabia,35997107
+Sudan,44440486
+Senegal,16436120
+Singapore,5685807
+Solomon Islands,691191
+Sierra Leone,8233970
+El Salvador,6292731
+San Marino,34007
+Somalia,16537016
+Serbia,6899126
+Sub-Saharan Africa (excluding high income),1151203345
+South Sudan,10606227
+Sub-Saharan Africa,1151301807
+Small states,42392777
+Sao Tome and Principe,218641
+Suriname,607065
+Slovakia,5458827
+Slovenia,2102419
+Sweden,10353442
+Eswatini,1180655
+Sint Maarten (Dutch part),42310
+Seychelles,98462
+Syria,20772595
+Turks and Caicos Islands,44276
+Chad,16644701
+East Asia & Pacific (IDA & IBRD countries),2090523535
+Europe & Central Asia (IDA & IBRD countries),462023771
+Togo,8442580
+Thailand,71475664
+Tajikistan,9543207
+Turkmenistan,6250438
+Latin America & the Caribbean (IDA & IBRD countries),634680385
+Timor-Leste,1299995
+Middle East & North Africa (IDA & IBRD countries),407006855
+Tonga,105254
+South Asia (IDA & IBRD),1882531620
+Sub-Saharan Africa (IDA & IBRD countries),1151301807
+Trinidad and Tobago,1518147
+Tunisia,12161723
+Turkey,84135428
+Tuvalu,11069
+Tanzania,61704518
+Uganda,44404611
+Ukraine,44132049
+Upper middle income,-2147483648
+Uruguay,3429086
+US,331501080
+Uzbekistan,34232050
+Saint Vincent and the Grenadines,104632
+Venezuela,28490453
+British Virgin Islands,30910
+Virgin Islands (U.S.),106290
+Vietnam,96648685
+Vanuatu,311685
+World,-2147483648
+Samoa,214929
+Kosovo,1790133
+Yemen,32284046
+South Africa,58801927
+Zambia,18927715
+Zimbabwe,15669666

src/data_parser.py

@@ -0,0 +1,63 @@
+import pandas as pd
+import numpy as np
+from datetime import datetime
+
+
+class Parser:
+    def __init__(self, filename_confirmed,
+                 filename_deaths,
+                 filename_recovered,
+                 filename_population):
+        self.confirmed = self.read_csv(filename_confirmed)
+        self.deaths = self.read_csv(filename_deaths)
+        self.recovered = self.read_csv(filename_recovered)
+        self.population = self.read_population(filename_population)
+        self.countries = list(np.intersect1d(self.confirmed.columns.values,
+                                             self.population.index.values))
+
+    def read_csv(self, filename):
+        # Create pandas dataframe from .csv
+        data = pd.read_csv(filename)
+
+        # Manipulate the dataframe to have dates as row indices and country
+        # names as column names
+        data = data.set_index("Country/Region")
+        data = data.T
+        data.index = pd.to_datetime(data.index)
+
+        return data
+
+    def parse_data(self, start_date, end_date, country):
+        self.validate_date(start_date)
+        self.validate_date(end_date)
+        self.validate_country(country)
+
+        delta_i = self.confirmed.loc[:end_date, country].diff().dropna()
+        delta_i = delta_i.astype(int)
+        r = (self.deaths.loc[:end_date, country]
+             + self.recovered.loc[:end_date, country])
+        delta_r = r.diff().dropna().astype(int)
+        i = (delta_i - delta_r).cumsum()
+        return i[start_date:], r[start_date:]
+
+    def read_population(self, filename):
+        # Create pandas dataframe from .csv
+        data = pd.read_csv(filename)
+        data = data.set_index("Country")
+
+        return data
+
+    def parse_population(self, country):
+        population = self.population.loc[country, "Population"]
+
+        return population
+
+    def validate_date(self, date_text):
+        try:
+            datetime.strptime(date_text, '%Y-%m-%d')
+        except ValueError:
+            raise ValueError("Incorrect data format, should be YYYY-MM-DD!")
+
+    def validate_country(self, country):
+        if country not in self.countries:
+            raise ValueError("Country not in list!")

src/sampler.py

@@ -0,0 +1,160 @@
+from src.utility_functions import log_pi_lambda, log_pi_t
+import numpy as np
+
+
+def update_lambda(lam, t, s, i, P, sigma, alpha, beta, phi):
+    # This function updates the parameter vector lambda.
+    # INPUT:
+    # - lam: array of the values of lambda;
+    # - t: array of the breakpoints;
+    # - s: array of susceptible individuals during time;
+    # - i: array of infected individuals during time;
+    # - P: total number of individuals;
+    # - sigma: algorithm parameter for the proposal of a new candidate lambda;
+    # - alpha, beta: hyperparameters of the prior of lambda;
+    # - phi parameter phi of the model.
+    # OUTPUT:
+    # - lam: update array;
+    # - accept: number of acceted candidates.
+    # NOTES: The update is done component-wise in a sequential manner.
+
+    current = np.copy(lam)  # Get the current state of the chain.
+    candidate = np.copy(current)  # Initialize the new candidate.
+
+    # For every component of the parameter vector, we tweak such component
+    # according to the chosen proposal and then update the chain according
+    # to the computed acceptance rate.
+    accepted = 0  # Initialize the count of accepted candidates.
+    for j in range(current.shape[0]):
+        # Tweak the j-th component.
+        candidate[j] = candidate[j] + sigma * np.random.normal()
+        # Compute the acceptance rate.
+        log_alpha = (log_pi_lambda(candidate, t, s, i, P, alpha, beta, phi)
+                     - log_pi_lambda(current, t, s, i, P, alpha, beta, phi))
+
+        # If the candidate is accepted, we move the chain (current = candidate)
+        # and increase the count of accepted candidates. Otherwise, we reject
+        # the candidate and the chain does not move from the current state
+        # (candidate = current).
+        if log_alpha > np.log(np.random.uniform()):
+            current = np.copy(candidate)
+            accepted = accepted + 1
+        else:
+            candidate = np.copy(current)
+
+    return current.reshape(-1, 1), accepted
+
+
+def update_t(lam, t, s, i, P, M, phi):
+    # This function updates the parameter vector lambda.
+    # INPUT:
+    # - lam: array of the values of lambda;
+    # - t: array of the breakpoints;
+    # - s: array of susceptible individuals during time;
+    # - i: array of infected individuals during time;
+    # - P: total number of individuals;
+    # - M: algorithm parameter for the proposal of a new candidate t;
+    # - phi parameter phi of the model.
+    # OUTPUT:
+    # - lam: update array;
+    # - accept: number of acceted candidates.
+    # NOTES: The update is done component-wise in a sequential manner.
+
+    current = np.copy(t)  # Get the current state of the chain.
+    candidate = np.copy(current)  # Initialize the new candidate.
+
+    # For every component of the parameter vector, we tweak such component
+    # according to the chosen proposal and then update the chain according
+    # to the computed acceptance rate.
+    accepted = 0  # Initialize the count of accepted candidates.
+    for j in range(current.shape[0]):
+        # Tweak the j-th component.
+        candidate[j] = candidate[j] + np.random.choice(np.arange(-M, M + 1))
+        # Compute the acceptance rate.
+        log_alpha = (log_pi_t(lam, candidate, s, i, P, phi)
+                     - log_pi_t(lam, current, s, i, P, phi))
+
+        # If the candidate is accepted, we move the chain (current = candidate)
+        # and increase the count of accepted candidates. Otherwise, we reject
+        # the candidate and the chain does not move from the current state
+        # (candidate = current).
+        if log_alpha > np.log(np.random.uniform()):
+            current = np.copy(candidate)
+            accepted = accepted + 1
+        else:
+            candidate = np.copy(current)
+
+    return current.reshape(-1, 1), accepted
+
+
+def mcmc_sampler(s, i, d, P, n_iterations, burnin, M, sigma,
+                 alpha, beta, a, b, phi):
+    # This function implement the hybrid MCMC sampler.
+    # INPUT:
+    # - s: array of susceptible individuals during time;
+    # - i: array of infected individuals during time;
+    # - d: number of breakpoints;
+    # - P: total number of individuals;
+    # - n_iterations: number of iterations for the algorithm;
+    # - burnin: number of burnin iterations to discard;
+    # - M: algorithm parameter for the proposal of a new candidate t;
+    # - sigma: algorithm parameter for the proposal of a new candidate lambda;
+    # - alpha, beta: hyperparameters of the prior of lambda;
+    # - a, b: hyperparameters of the prior of p;
+    # - phi: parameter phi of the model.
+    # OUTPUT:
+    # - p: simulated chain for the probability of removal from
+    # infected population;
+    # - lam: simulated chain for lambda;
+    # - t: simulated chain for the breakpoints.
+
+    T = s.shape[0] - 1  # Index of the final time instant.
+
+    # Initialize the parameters.
+
+    # The initial value of p is drawn from the prior distribution.
+    p = np.random.beta(a, b, size=(1, 1))
+    # Each of the d breakpoints (t_i) is drawn randomly (without replacement)
+    # between 1 and T-1. The obtained vector is then sorted to make sure
+    # that t_1 < t_2 < ... < t_d.
+    t = np.sort(np.random.choice(np.arange(1, T), size=d-1, replace=False))
+    t = t.reshape(-1, 1)
+    # Each of the lambda_i's is drawn independently from
+    # its prior distribution.
+    lam = np.random.gamma(alpha, beta)
+    lam = lam.reshape(-1, 1)
+
+    # Compute the hyperparameters of the posterior of p.
+    a_new = a + i[0] - i[-1] + s[0] - s[-1]
+    b_new = b + np.sum(i[1:]) + s[-1] - s[0]
+
+    # Initialize the count of accepted candidates for lambda and t.
+    a_lam = 0
+    a_t = 0
+
+    # Run the chain
+    for _ in range(n_iterations):
+        # Update p by sampling from its posterior.
+        p = np.hstack((p, np.random.beta(a_new, b_new, size=(1, 1))))
+        # Update lam via Metropolis-Hastings step.
+        new_lam, accepted_lam = update_lambda(lam[:, -1], t[:, -1], s, i, P,
+                                              sigma, alpha, beta, phi)
+        # Update t via Metropolis-Hastings step.
+        new_t, accepted_t = update_t(lam[:, -1], t[:, -1], s, i, P, M, phi)
+        lam = np.hstack((lam, new_lam))
+        t = np.hstack((t, new_t))
+
+        # Update the counts of accepted candidates for lambda and t.
+        a_lam = a_lam + accepted_lam
+        a_t = a_t + accepted_t
+
+    # Compute the acceptance rates for lambda and t.
+    lam_ar = a_lam / n_iterations / d
+    t_ar = a_t / n_iterations / (d-1)
+
+    # Discard burn-in iterations.
+    p = p[:, burnin:]
+    lam = lam[:, burnin:]
+    t = t[:, burnin:]
+
+    return p, lam, t, lam_ar, t_ar

src/utility_functions.py

@@ -0,0 +1,134 @@
+import numpy as np
+from scipy.special import gammaln
+
+
+def lambda_time(lam, t, time):
+    # This function computes the value of lambda.
+    # INPUT:
+    # - lam, t: arrays defining the (piecewise constant) function lambda(t);
+    # - time: time instants at which we want to evaluate lambda.
+    # OUTPUT:
+    # - lambda_time: value of lambda.
+    # NOTES: The function is vectorized in the array time. Indeed,
+    # it allows to compute kappa for all the time instants in the array time.
+
+    lambda_time = lam[np.searchsorted(t, time, side="right")]
+    return lambda_time
+
+
+def compute_kappa_time(s_t, i_t, lam, t, time, phi, P):
+    # This function computes the value of kappa.
+    # INPUT:
+    # - s_t: number of susceptible individuals;
+    # - i_t: number of infected individuals;
+    # - lam, t: arrays defining the (piecewise constant) function lambda(t);
+    # - time: time instants at which we want to evaluate kappa;
+    # - phi: parameter phi of the model;
+    # - P: total number of individuals.
+    # OUTPUT:
+    # - kappa_time: value of kappa.
+    # NOTES: The function is vectorized in the arrays time, s_t and i_t.
+    # Indeed, it allows to compute kappa for all the time instants in
+    # the array time. It is required that time, s_t and i_t
+    # have the same dimension.
+
+    p_si_t = 1 - np.exp(-np.multiply(lambda_time(lam, t, time), i_t) / P)
+    kappa_time = (1/phi - 1) * np.multiply(s_t, p_si_t)
+    return kappa_time
+
+
+def log_pi_lambda(lam, t, s, i, P, alpha, beta, phi):
+    # This function computes the (log-) full-conditional
+    # of the parameter vector lambda.
+    # INPUT:
+    # - lam: array of the values of lambda;
+    # - t: array of the breakpoints;
+    # - s: array of susceptible individuals during time;
+    # - i: array of infected individuals during time;
+    # - P: total number of individuals;
+    # - alpha, beta: hyperparameters of the prior of lambda;
+    # - phi: parameter phi of the model.
+    # OUTPUT:
+    # - result: (log-) full-conditional of lambda evaluated at lam.
+
+    T = s.shape[0] - 1  # Index of the final time instant.
+    time = np.arange(T + 1)  # Array of all time instants.
+
+    # First, we initialize the result to -inf. If all the components of the
+    # vector lam are positive (i.e., the vector is admissible) we compute the
+    # (log-) full-conditional of lambda evaluated at lam and return the result.
+    result = np.NINF
+    if all(lam > 0):
+        kappa_vec = compute_kappa_time(s[:-1], i[:-1], lam,
+                                       t, time[:-1], phi, P)
+        result = (np.sum(gammaln(-np.diff(s) + kappa_vec)
+                         + kappa_vec * np.log(1 - phi)
+                         - gammaln(kappa_vec))
+                  + np.sum(np.log(np.power(lam, alpha-1))
+                           - np.multiply(beta, lam)))
+    return result
+
+
+def log_pi_t(lam, t, s, i, P, phi):
+    # This function computes the (log-) full-conditional
+    # of the parameter vector t.
+    # INPUT:
+    # - lam: array of the values of lambda;
+    # - t: array of the breakpoints;
+    # - s: array of susceptible individuals during time;
+    # - i: array of infected individuals during time;
+    # - P: total number of individuals;
+    # - phi parameter phi of the model.
+    # OUTPUT:
+    # - result: (log-) full-conditional of t evaluated at t.
+
+    T = s.shape[0] - 1  # Index of the final time instant.
+    time = np.arange(T + 1)  # Array of all time instants.
+
+    # First, we initialize the result to -inf. If we have that
+    # 0 < t_1 < t_2 < ... < t_(d-1) < T (i.e., the vector is admissible)
+    # we compute the (log-) full-conditional of t evaluated at t
+    # and return the result.
+    result = np.NINF
+    if np.all(np.diff(t) > 0) and t[0] > 0 and t[-1] < T:
+        kappa_vec = compute_kappa_time(s[:-1], i[:-1], lam,
+                                       t, time[:-1], phi, P)
+        result = np.sum(gammaln(-np.diff(s) + kappa_vec)
+                        + kappa_vec * np.log(1 - phi)
+                        - gammaln(kappa_vec))
+    return result
+
+
+def simulate_data(T, lam, t, s_0, i_0, p_r, phi):
+    # This function simulates data according to the process described above.
+    # INPUT:
+    # - T: index of the final time instant;
+    # - lam, t: (true) arrays defining the function lambda(t);
+    # - s_0: initial number of susceptible individuals;
+    # - i_0: initial number of infected individuals;
+    # - p_r: (true) probability of removal from infected population;
+    # - phi: parameter phi of the model.
+    # OUTPUT:
+    # - s: array of susceptible individuals during time;
+    # - i: array of infected individuals during time.
+
+    # Compute the total number of individuals.
+    P = s_0 + i_0
+
+    # Initialize the arrays s and t.
+    s = np.array([s_0])
+    i = np.array([i_0])
+
+    time = np.arange(T + 1)
+    for t_i in time:
+        # Draw a realization of delta_r.
+        delta_r = np.random.binomial(i[-1], p_r)
+        # Compute the kappa parameter at time t_i.
+        kappa = compute_kappa_time(s[-1], i[-1], lam, t, t_i, phi, P)
+        # Draw a realization of delta_i.
+        delta_i = np.random.negative_binomial(kappa, 1 - phi)
+
+        # Update s and i according to the model.
+        s = np.append(s, s[-1] - delta_i)
+        i = np.append(i, i[-1] + delta_i - delta_r)
+    return s, i