Fix issue 255

src/penn_chime/charts.py

@@ -26,7 +26,7 @@ def new_admissions_chart(
 
     tooltip_dict = {False: "day", True: "date:T"}
     if as_date:
-        projection_admits = add_date_column(projection_admits)
+        projection_admits = add_date_column(projection_admits, parameters.date_first_hospitalized)
         x_kwargs = {"shorthand": "date:T", "title": "Date", "axis": alt.Axis(format=(DATE_FORMAT))}
     else:
         x_kwargs = {"shorthand": "day", "title": "Days from today"}
@@ -66,7 +66,7 @@ def admitted_patients_chart(
     max_y_axis = parameters.max_y_axis
     as_date = parameters.as_date
     if as_date:
-        census = add_date_column(census)
+        census = add_date_column(census, parameters.date_first_hospitalized)
         x_kwargs = {"shorthand": "date:T", "title": "Date", "axis": alt.Axis(format=(DATE_FORMAT))}
         idx = "date:T"
     else:
@@ -113,7 +113,7 @@ def additional_projections_chart(
     max_y_axis = parameters.max_y_axis
 
     if as_date:
-        dat = add_date_column(dat)
+        dat = add_date_column(dat, parameters.date_first_hospitalized)
         x_kwargs = {"shorthand": "date:T", "title": "Date", "axis": alt.Axis(format=(DATE_FORMAT))}
     else:
         x_kwargs = {"shorthand": "day", "title": "Days from today"}

src/penn_chime/defaults.py

@@ -1,4 +1,6 @@
 """Defaults."""
+from typing import Optional
+from datetime import date
 
 from .utils import RateLos
 
@@ -29,9 +31,10 @@ def __init__(
         icu: RateLos,
         ventilated: RateLos,
 
+        date_first_hospitalized: Optional[date] = None,
         as_date: bool = False,
         market_share: float = 1.0,
-        max_y_axis: int = None,
+        max_y_axis: Optional[int] = None,
         n_days: int = 60,
         recovery_days: int = 14,
     ):
@@ -45,6 +48,8 @@ def __init__(
         self.icu = icu
         self.ventilated = ventilated
 
+        self.date_first_hospitalized = date_first_hospitalized
+        self.n_days_since_first_hospitalized = (date.today() - date_first_hospitalized).days
         self.as_date = as_date
         self.market_share = market_share
         self.max_y_axis = max_y_axis

src/penn_chime/models.py

@@ -7,17 +7,17 @@
 
 from __future__ import annotations
 
-from typing import Dict, Generator, Tuple
+from typing import Dict, Generator, Tuple, Optional
 
 import numpy as np  # type: ignore
 import pandas as pd  # type: ignore
 
 from .parameters import Parameters
-
+from .utils import SimSirModelAttributes
 
 class SimSirModel:
 
-    def __init__(self, p: Parameters) -> SimSirModel:
+    def __init__(self, p: Parameters):
         # TODO missing initial non-zero 'recovered' value
         recovered = 0.0
         recovery_days = p.recovery_days
@@ -32,6 +32,9 @@ def __init__(self, p: Parameters) -> SimSirModel:
             for key, d in p.dispositions.items()
         }
 
+        self._rates = rates
+        self._lengths_of_stay = lengths_of_stay
+
         # Note: this should not be an integer.
         # We're appoximating infected from what we do know.
         # TODO market_share > 0, hosp_rate > 0
@@ -45,8 +48,8 @@ def __init__(self, p: Parameters) -> SimSirModel:
             p.known_infected / infected if infected > 1.0e-7 else None
         )
 
-        intrinsic_growth_rate = \
-            (2.0 ** (1.0 / p.doubling_time) - 1.0) if p.doubling_time > 0.0 else 0.0
+        # (2.0 ** (1.0 / p.doubling_time) - 1.0) if p.doubling_time > 0.0 else 0.0
+        intrinsic_growth_rate = self._intrinsic_growth_rate(p.doubling_time)
 
         gamma = 1.0 / recovery_days
 
@@ -75,8 +78,8 @@ def __init__(self, p: Parameters) -> SimSirModel:
             p.n_days,
         )
         dispositions_df = build_dispositions_df(raw_df, rates, p.market_share)
-        admits_df = build_admits_df(dispositions_df)
-        census_df = build_census_df(admits_df, lengths_of_stay)
+        admits_df = build_admits_df(dispositions_df, p.n_days_since_first_hospitalized)
+        census_df = build_census_df(admits_df, lengths_of_stay, p.n_days_since_first_hospitalized)
 
         self.susceptible = susceptible
         self.infected = infected
@@ -94,10 +97,128 @@ def __init__(self, p: Parameters) -> SimSirModel:
         self.dispositions_df = dispositions_df
         self.admits_df = admits_df
         self.census_df = census_df
+
+        if p.n_days_since_first_hospitalized is not None and p.doubling_time is None:
+            # optimize doubling_time
+            argmin_dt = None
+            min_loss = 2.0**99
+            censes = dict()
+            for dt in np.linspace(1,15,29):
+                censes[dt] = self.run_projection(p, dt)
+                self.census_df = censes[dt] # log it into state for loss
+                loss_dt = self.loss_dt(p)
+                if loss_dt < min_loss:
+                    min_loss = loss_dt
+                    argmin_dt = dt
+            self.census_df = censes[dt]
+            p.doubling_time = argmin_dt
+
+            # update all state that is dependent on doubling time.
+            intrinsic_growth_rate = self._intrinsic_growth_rate(p.doubling_time)
+            gamma = 1 / recovery_days
+            beta = self._beta(intrinsic_growth_rate, gamma, susceptible, p.relative_contact_rate)
+            r_t = beta / gamma * susceptible
+            r_naught = (intrinsic_growth_rate + gamma) / gamma
+            doubling_time_t = 1.0 / np.log2(beta * susceptible - gamma + 1)
+            raw_df = sim_sir_df(
+                susceptible,
+                infected,
+                recovered,
+                beta,
+                gamma,
+                p.n_days
+            )
+            dispositions_df = build_dispositions_df(raw_df, rates, p.market_share)
+            admits_df = build_admits_df(dispositions_df, p.n_days_since_first_hospitalized)
+            census_df = build_census_df(admits_df, lengths_of_stay, p.n_days_since_first_hospitalized)
+
+            self.intrinsic_growth_rate = intrinsic_growth_rate
+            self.gamma = gamma
+            self.beta = beta
+            self.r_t = r_t
+            self.r_naught = r_naught
+            self.doubling_time_t = doubling_time_t
+            self.raw_df = raw_df
+            self.dispositions_df = dispositions_df
+            self.admits_df = admits_df
+            self.census_df = census_df
+
         self.daily_growth = daily_growth_helper(p.doubling_time)
         self.daily_growth_t = daily_growth_helper(doubling_time_t)
 
+        return None
+
+    def run_projection(self, p: Parameters, doubling_time: float) -> pd.DataFrame:
+        intrinsic_growth_rate = self._intrinsic_growth_rate(doubling_time)
+
+        recovery_days = p.recovery_days
+        market_share = p.market_share
+        initial_i = 1 / p.hospitalized.rate / market_share
+
+        S, I, R = self.susceptible, self.infected, self.recovered
+
+        # mean recovery rate (inv_recovery_days)
+        gamma = 1 / recovery_days
+
+        # contact rate
+        beta = (intrinsic_growth_rate + gamma) / S
+
+        n_days = p.n_days
+
+        raw_df = sim_sir_df(S,I,R,beta,gamma,n_days)
+
+        # dispositions_df = build_dispositions_df(raw_df, self._rates, p.market_share)
+
+        i_dict_v = get_dispositions(raw_df.infected, self._rates, market_share)
+        r_dict_v = get_dispositions(raw_df.recovered, self._rates, market_share)
+
+        dispositions = {
+            key: value + r_dict_v[key]
+            for key, value in i_dict_v.items()
+        }
+
+        dispositions_df = pd.DataFrame(dispositions)
+        dispositions_df = dispositions_df.assign(day=dispositions_df.index)
 
+        admits_df = build_admits_df(dispositions_df, p.n_days_since_first_hospitalized)
+        census_df = build_census_df(admits_df, self._lengths_of_stay, p.n_days_since_first_hospitalized)
+        return census_df
+
+    def loss_dt(self, p: Parameters) -> float:
+        """Squared error: predicted_current_hospitalized vs. actual current hospitalized
+
+        gets prediction of current hospitalized from a census_df which
+        is dependent on a given doubling_time in state.
+        """
+        # get the predicted number of hospitalized today
+        predicted_current_hospitalized = self.census_df.hospitalized.loc[p.n_days_since_first_hospitalized]
+
+        # compare against actual / user inputted number
+        # we shall optimize squared distance
+        return (p.current_hospitalized - predicted_current_hospitalized) ** 2
+
+
+    @staticmethod
+    def _intrinsic_growth_rate(doubling_time: Optional[float]) -> float:
+        if doubling_time is not None:
+            return (2.0 ** (1.0 / doubling_time) - 1.0) if doubling_time > 0.0 else 0.0
+        return 0.0
+
+    @staticmethod
+    def _beta(
+            intrinsic_growth_rate: float,
+            gamma: float,
+            susceptible: float,
+            relative_contact_rate: float) -> float:
+        return (
+            (intrinsic_growth_rate + gamma)
+            / susceptible
+            * (1.0 - relative_contact_rate)
+        )
+
+###################
+##  MODEL FUNCS  ##
+###################
 def sir(
     s: float, i: float, r: float, beta: float, gamma: float, n: float
 ) -> Tuple[float, float, float]:
@@ -118,7 +239,7 @@ def sir(
 
 def gen_sir(
     s: float, i: float, r: float, beta: float, gamma: float, n_days: int
-) -> Generator[Tuple[float, float, float], None, None]:
+) -> Generator[Tuple[int, float, float, float], None, None]:
     """Simulate SIR model forward in time yielding tuples."""
     s, i, r = (float(v) for v in (s, i, r))
     n = s + i + r
@@ -136,6 +257,19 @@ def sim_sir_df(
         columns=("day", "susceptible", "infected", "recovered"),
     )
 
+
+def get_dispositions(
+    patients: np.ndarray,
+    rates: Dict[str, float],
+    market_share: float,
+) -> Dict[str, np.ndarray]:
+    """Get dispositions of patients adjusted by rate and market_share."""
+    return {
+        key: patients * rate * market_share
+        for key, rate in rates.items()
+    }
+
+
 def build_dispositions_df(
     sim_sir_df: pd.DataFrame,
     rates: Dict[str, float],
@@ -152,18 +286,22 @@ def build_dispositions_df(
     })
 
 
-def build_admits_df(dispositions_df: pd.DataFrame) -> pd.DataFrame:
+def build_admits_df(dispositions_df: pd.DataFrame, n_days_since_first_hospitalized: int) -> pd.DataFrame:
     """Build admits dataframe from dispositions."""
     admits_df = dispositions_df.iloc[:-1, :] - dispositions_df.shift(1)
-    admits_df.day = dispositions_df.day
+    if n_days_since_first_hospitalized is not None:
+        admits_df.day = dispositions_df.day - n_days_since_first_hospitalized
+    else:
+        admits_df.day = dispositions_df.day
     return admits_df
 
 
 def build_census_df(
     admits_df: pd.DataFrame,
     lengths_of_stay: Dict[str, int],
+    n_days_since_first_hospitalized: int
 ) -> pd.DataFrame:
-    """ALOS for each disposition of COVID-19 case (total guesses)"""
+    """Average Length of Stay for each disposition of COVID-19 case (total guesses)"""
     return pd.DataFrame({
         'day': admits_df.day,
         **{
@@ -176,9 +314,13 @@ def build_census_df(
     })
 
 
-def daily_growth_helper(doubling_time):
+
+#############
+##  UTILS  ##
+#############
+def daily_growth_helper(doubling_time: float) -> float:
     """Calculates average daily growth rate from doubling time"""
     result = 0
-    if doubling_time != 0:
+    if doubling_time != 0 and doubling_time is not None:
         result = (np.power(2, 1.0 / doubling_time) - 1) * 100
     return result

src/penn_chime/parameters.py

@@ -4,6 +4,9 @@
 `change_date`, so users can see when results have last changed
 """
 
+from typing import Optional
+from datetime import date
+
 from .utils import RateLos
 
 
@@ -25,7 +28,8 @@ def __init__(
 
         as_date: bool = False,
         market_share: float = 1.0,
-        max_y_axis: int = None,
+        date_first_hospitalized: Optional[date] = None,
+        max_y_axis: Optional[int] = None,
         n_days: int = 60,
         recovery_days: int = 14,
     ):
@@ -44,6 +48,12 @@ def __init__(
         self.max_y_axis = max_y_axis
         self.n_days = n_days
         self.recovery_days = recovery_days
+        if date_first_hospitalized:
+            self.date_first_hospitalized = date_first_hospitalized # needed for utils.add_date_column
+            self.n_days_since_first_hospitalized = (date.today() - date_first_hospitalized).days
+        else:
+            self.date_first_hospitalized = None # needed for utils.add_date_column
+            self.n_days_since_first_hospitalized = None
 
         self.labels = {
             "hospitalized": "Hospitalized",