code: import pandas as pd
  import numpy as np
  
  from scipy.integrate import odeint
  from sklearn.linear_model import LinearRegression
  from sklearn.metrics import mean_absolute_error, mean_squared_error
  
  
  def logistic_R0(t, R_0_start, k, x0, R_0_end):
      """
      R0 moduled as logistic function
      """
  
      return (R_0_start - R_0_end) / (1 + np.exp(-k * (-t + x0))) + R_0_end
  
  
  def beta(t, R_0_start, k, x0, R_0_end, gamma):
      """
      Computes beta at a given time `t`
      """
  
      return logistic_R0(t, R_0_start, k, x0, R_0_end) * gamma
  
  
  def sird_calc(y, t, N, gamma, alpha, R_0_start, k, x0, R_0_end, beta):
      """
      Computes SIRD model
      """
  
      S, I, R, D = y
      dSdt = -beta(t, R_0_start, k, x0, R_0_end, gamma) * S * I / N
      dIdt = -dSdt - (1 - alpha) * gamma * I - alpha * I
      dRdt = (1 - alpha) * gamma * I
      dDdt = alpha * I
      return dSdt, dIdt, dRdt, dDdt
  
  
  def sird(
      province,
      pop_prov_df,
      gamma=1 / 7,
      alpha=0.01,
      days=101,
      R_0_start=2,
      k=0.2,
      x0=40,
      R_0_end=0.3,
      prov_list_df=None,
  ):
      """
      Create and compute a SIRD model
  
      Parameters
      ----------
  
      province : str
          The province name.
  
      pop_prov_df : pandas DataFrame
          The DataFrame with demographic data.
  
      gamma : float (default=1/7)
          Inverse of how many days the infection lasts.
  
      alpha : float (default=0.01)
          Death rate.
  
      days : int (default=101)
          Total number of days to predict + 1.
  
      R_0_start : float (default=2)
          Starting value of RO
  
      k : float (default=0.2)
          How quickly R0 declines. Lower values of k will
          let R0 need more time to become lower.
  
      x0 : int (default=40)
          Value on the x-axis of the inflection point of R0.
          This can be interpreted as the day in which lockdown
          comes into effect.
  
      R_0_end : float (default=0.3)
          Final value of RO
  
      Returns
      -------
      A numpy array of shape (4, days).
      """
  
      # Population
      if prov_list_df is not None:
          prov_list = prov_list_df[prov_list_df.Region == province]["Province"].values
  
          N = 0
          for prov in prov_list:
              N += pop_prov_df.loc[
                  (pop_prov_df.Territorio == prov) & (pop_prov_df.Eta == "Total")
              ]["Value"].values[0]
      else:
          N = pop_prov_df.loc[
              (pop_prov_df.Territorio == province) & (pop_prov_df.Eta == "Total")
          ]["Value"].values[0]
  
      times = range(days)
  
      # S0, I0, R0, D0: initial conditions vector
      init = N - 1, 1, 0, 0
  
      # Solve the model
      sirsol = odeint(
          sird_calc, init, times, args=(N, gamma, alpha, R_0_start, k, x0, R_0_end, beta)
      )
  
      return sirsol.T
  
  
  def Model(days, N, R_0_start, k, x0, R_0_end, alpha, gamma):
      y0 = (
          N - 1.0,
          1.0,
          0.0,
          0.0,
      )
      times = range(0, days)
  
      sirsol = odeint(
          sird_calc, y0, times, args=(N, gamma, alpha, R_0_start, k, x0, R_0_end, beta)
      )
  
      S, I, R, D = sirsol.T
      R0_over_time = [
          beta(i, R_0_start, k, x0, R_0_end, gamma) / gamma for i in range(len(times))
      ]
  
      return times, S, I, R, D, R0_over_time
  
  
  class DeterministicSird:
      def __init__(
          self,
          data_df,
          pop_prov_df,
          prov_list_df,
          area,
          group_column,
          data_column,
          data_filter,
          lag,
          days_to_predict,
          is_regional=True,
          pcm_data=None,
      ):
  
          self.data_df = data_df
          self.pop_prov_df = pop_prov_df
          self.prov_list_df = prov_list_df
          self.area = area
          self.group_column = group_column
          self.data_column = data_column
          self.data_filter = data_filter
          self.lag = lag
          self.days_to_predict = days_to_predict
          self.is_regional = is_regional
          self.pcm_data = pcm_data
  
      def get_region_pop(self, region, pop_df, prov_df):
          """
          Computes the total population for a region
          starting from the provinces' popuplation
  
          Parameters
          ----------
          region : str
              The region whose population we need
  
          pop_df : pandas DataFrame
              Data for provinces population
  
          prov_df : pandas DataFrame
              Data that associates each province with
              its region
  
          Returns
          -------
          N : int
              The population of the region
          """
  
          prov_list = prov_df[prov_df.Region == region]["Province"].values
  
          N = 0
          for prov in prov_list:
              N += pop_df.loc[(pop_df.Territorio == prov) & (pop_df.Eta == "Total")][
                  "Value"
              ].values[0]
  
          return N
  
      def get_prov_pop(self):
          return self.pop_prov_df.loc[
              (self.pop_prov_df.Territorio == self.area)
              & (self.pop_prov_df.Eta == "Total")
          ]["Value"].values[0]
  
      def fix_arr(self, arr):
          arr[arr < 0] = 0
          arr[np.isinf(arr)] = 0
          return np.nan_to_num(arr)
  
      def lag_data(self, data, lag=7, return_all=False):
          if isinstance(data, np.ndarray):
              N = data.shape[0]
          else:
              N = len(data)
  
          X = np.empty(shape=(N - lag, lag + 1))
  
          for i in range(lag, N):
              X[
                  i - lag,
              ] = [data[i - j] for j in range(lag + 1)]
  
          if not return_all:
              return X[-1, 1:]
  
          return X[:, 1:], X[:, 0]
  
      def _prepare_data_regional(self):
          data_df = self.data_df.loc[
              (self.data_df[self.group_column] == self.area),
              [
                  "data",
                  "totale_positivi",
                  "dimessi_guariti",
                  "deceduti",
                  "totale_casi",
                  "nuovi_positivi",
              ],
          ]
  
          data_df = data_df.query(self.data_filter + " > " + self.data_column)
  
          data_df["suscettibili"] = self.pop - data_df["totale_casi"]
  
          data_df = data_df.loc[
              :,
              [
                  "data",
                  "totale_positivi",
                  "dimessi_guariti",
                  "deceduti",
                  "suscettibili",
                  "nuovi_positivi",
              ],
          ]
  
          return data_df
  
      def _prepare_data_provincial(self):
          regione = self.data_df[self.data_df[self.group_column] == self.area][
              "Region"
          ].values[0]
  
          pop = self.get_prov_pop()
  
          pcm_data = self.pcm_data.loc[
              (self.pcm_data["denominazione_regione"] == regione),
              [
                  "data",
                  "totale_positivi",
                  "dimessi_guariti",
                  "deceduti",
                  "totale_casi",
                  "nuovi_positivi",
              ],
          ].reset_index(drop=True)
  
          data_df = self.data_df.loc[
              (self.data_df[self.group_column] == self.area),
              ["Date", "New_cases", "Curr_pos_cases", "Tot_deaths"],
          ].reset_index(drop=True)
  
          recov_rate = (pcm_data["dimessi_guariti"] / pcm_data["totale_casi"])[
              : data_df.shape[0]
          ]
  
          recov_rate = self.fix_arr(recov_rate)
  
          recov = recov_rate * data_df["Curr_pos_cases"].values
          data_df["dimessi_guariti"] = recov
  
          infected = (
              data_df["Curr_pos_cases"].values
              - data_df["Tot_deaths"].values
              - data_df["dimessi_guariti"]
          )
  
          data_df["totale_positivi"] = infected
  
          data_df["suscettibili"] = pop - data_df["Curr_pos_cases"]
  
          query = (
              (
                  pd.Timestamp(self.data_filter) + pd.DateOffset(self.days_to_predict)
              ).strftime("%Y%m%d")
              + " > "
              + self.data_column
          )
  
          real_df = data_df.query(query)
  
          real_df.rename(
              columns={
                  "New_cases": "nuovi_positivi",
                  "Curr_pos_cases": "totale_casi",
                  "Tot_deaths": "deceduti",
                  "Date": "data",
              },
              inplace=True,
          )
  
          real_df = real_df.loc[
              :,
              [
                  "data",
                  "totale_positivi",
                  "dimessi_guariti",
                  "deceduti",
                  "suscettibili",
                  "nuovi_positivi",
              ],
          ]
  
          self._realdf = real_df
  
          data_df = data_df.query(self.data_filter + " > " + self.data_column)
  
          data_df.rename(
              columns={
                  "New_cases": "nuovi_positivi",
                  "Curr_pos_cases": "totale_casi",
                  "Tot_deaths": "deceduti",
                  "Date": "data",
              },
              inplace=True,
          )
  
          data_df = data_df.astype(
              {
                  "totale_positivi": "int32",
                  "dimessi_guariti": "int32",
                  "deceduti": "int32",
                  "suscettibili": "int32",
                  "nuovi_positivi": "int32",
              }
          )
  
          data_df = data_df.loc[
              :,
              [
                  "data",
                  "totale_positivi",
                  "dimessi_guariti",
                  "deceduti",
                  "suscettibili",
                  "nuovi_positivi",
              ],
          ]
  
          return data_df
  
      def fit(self):
          if self.is_regional:
              self.pop = self.get_region_pop(
                  region=self.area, pop_df=self.pop_prov_df, prov_df=self.prov_list_df
              )
  
              data_df = self._prepare_data_regional()
          else:
              self.pop = self.get_prov_pop()
              data_df = self._prepare_data_provincial()
  
          n = data_df.shape[0]
  
          gamma = (
              np.diff(data_df["dimessi_guariti"].values)
              / data_df.iloc[: n - 1]["totale_positivi"].values
          )
  
          alpha = (
              np.diff(data_df["deceduti"].values)
              / data_df.iloc[: n - 1]["totale_positivi"].values
          )
  
          beta = (
              (self.pop / data_df.iloc[: n - 1]["suscettibili"].values)
              * (
                  np.diff(data_df["totale_positivi"].values)
                  + np.diff(data_df["dimessi_guariti"].values)
                  + np.diff(data_df["deceduti"].values)
              )
              / data_df.iloc[: n - 1]["totale_positivi"].values
          )
          R0 = beta / (gamma + alpha)
  
          gamma = self.fix_arr(gamma)
          alpha = self.fix_arr(alpha)
          beta = self.fix_arr(beta)
          R0 = self.fix_arr(R0)
  
          reg_beta = LinearRegression().fit(*self.lag_data(beta, self.lag, True))
          reg_gamma = LinearRegression().fit(*self.lag_data(gamma, self.lag, True))
          reg_alpha = LinearRegression().fit(*self.lag_data(alpha, self.lag, True))
  
          S = np.zeros(self.days_to_predict + 2)
          I = np.zeros(self.days_to_predict + 2)
          R = np.zeros(self.days_to_predict + 2)
          D = np.zeros(self.days_to_predict + 2)
          S[0] = data_df.iloc[-1]["suscettibili"]
          I[0] = data_df.iloc[-1]["totale_positivi"]
          R[0] = data_df.iloc[-1]["dimessi_guariti"]
          D[0] = data_df.iloc[-1]["deceduti"]
  
          for i in range(self.days_to_predict + 1):
              _beta = self.fix_arr(
                  reg_beta.predict(self.lag_data(beta, self.lag).reshape(1, -1))
              )
              _gamma = self.fix_arr(
                  reg_gamma.predict(self.lag_data(gamma, self.lag).reshape(1, -1))
              )
              _alpha = self.fix_arr(
                  reg_alpha.predict(self.lag_data(alpha, self.lag).reshape(1, -1))
              )
  
              beta = np.append(beta, _beta, axis=0)
              gamma = np.append(gamma, _gamma, axis=0)
              alpha = np.append(alpha, _alpha, axis=0)
  
              dIdt = np.round((1 + _beta * (S[i] / self.pop) - _gamma - _alpha) * I[i])
              dRdt = np.round(R[i] + _gamma * I[i])
              dDdt = np.round(D[i] + _alpha * I[i])
              dSdt = self.pop - dIdt[0] - dRdt[0] - dDdt[0]
  
              S[i + 1] = dSdt
              I[i + 1] = dIdt
              R[i + 1] = dRdt
              D[i + 1] = dDdt
  
          S = S[1:]
          I = I[1:]
          R = R[1:]
          D = D[1:]
  
          dates = pd.date_range(
              start=(data_df.iloc[-1]["data"] + pd.DateOffset(1)).strftime("%Y-%m-%d"),
              periods=self.days_to_predict + 1,
          )
  
          tmp_df = pd.DataFrame(
              np.column_stack([np.zeros(self.days_to_predict + 1), I, R, D, S]),
              columns=[
                  "data",
                  "totale_positivi",
                  "dimessi_guariti",
                  "deceduti",
                  "suscettibili",
              ],
          )
  
          tmp_df["data"] = dates
  
          data_df = pd.concat([data_df, tmp_df], ignore_index=True)
  
          data_df["nuovi_positivi"] = [0] + list(
              np.diff(data_df["totale_positivi"].values)
              + np.diff(data_df["dimessi_guariti"].values)
              + np.diff(data_df["deceduti"].values)
          )
  
          data_df["nuovi_positivi"] = data_df["nuovi_positivi"].apply(
              lambda x: 0 if x < 0 else x
          )
  
          beta = np.append(beta, np.zeros((1,)), axis=0)
          gamma = np.append(gamma, np.zeros((1,)), axis=0)
          alpha = np.append(alpha, np.zeros((1,)), axis=0)
  
          data_df["beta"] = beta
          data_df["gamma"] = gamma
          data_df["alpha"] = alpha
          data_df["R0"] = self.fix_arr(beta / (gamma + alpha))
          data_df = data_df[:-1]
  
          data_df = data_df.astype(
              {
                  "totale_positivi": "int32",
                  "dimessi_guariti": "int32",
                  "deceduti": "int32",
                  "suscettibili": "int32",
                  "nuovi_positivi": "int32",
              }
          )
  
          self._fdf = data_df
          self._realdf = self._get_real_data()
  
          return data_df
  
      @property
      def fitted_df(self):
          return self._fdf
  
      @property
      def real_df(self):
          return self._realdf
  
      def _get_real_data(self):
          if not self.is_regional:
              return self._realdf
  
          real_df = self.data_df[self.data_df[self.group_column] == self.area][
              [
                  "data",
                  "totale_positivi",
                  "dimessi_guariti",
                  "deceduti",
                  "totale_casi",
                  "nuovi_positivi",
              ]
          ]
  
          query = (
              pd.Timestamp(self.data_filter) + pd.DateOffset(self.days_to_predict)
          ).strftime("%Y%m%d") + " > data"
  
          real_df = real_df.query(query)
          real_df["suscettibili"] = self.pop - real_df["totale_casi"]
          real_df = real_df[
              [
                  "data",
                  "totale_positivi",
                  "dimessi_guariti",
                  "deceduti",
                  "suscettibili",
                  "nuovi_positivi",
              ]
          ]
  
          self._realdf = real_df
  
          return real_df
  
      def extract_ys(self, y_true, y_pred, compart):
          if y_true is None or y_pred is None:
              y_true = self.real_df[compart].values
              y_pred = self.fitted_df[compart].values
  
          return y_true, y_pred
  
      def mae(self, y_true=None, y_pred=None, compart=None):
          return mean_absolute_error(*self.extract_ys(y_true, y_pred, compart))
  
      def mse(self, y_true=None, y_pred=None, compart=None):
          return mean_squared_error(*self.extract_ys(y_true, y_pred, compart))
  
      def rmse(self, y_true=None, y_pred=None, compart=None):
          return mean_squared_error(
              *self.extract_ys(y_true, y_pred, compart), squared=False
          )
  
apis: scipy.integrate.odeint numpy.empty numpy.zeros numpy.diff numpy.append numpy.isinf sklearn.linear_model.LinearRegression numpy.exp pandas.concat pandas.DateOffset numpy.round numpy.nan_to_num pandas.Timestamp