def test_r_list_to_frame(): example_r_list = robjects.r("simulated = list(n_cases = 1:10)") as_frame = r_list_to_frame(example_r_list, ["n_cases"]) expected_frame = pd.DataFrame({ "n_cases": list(range(1, 11)), "timestep": list(range(1, 11)) }) pd.testing.assert_frame_equal(as_frame, expected_frame)
def simulate( self, length: int, state_weight: Optional[float] = None, state: Optional[Sequence[int]] = None, ) -> pd.DataFrame: r""" Simulate outbreaks. Parameters ---------- length Number of weeks to model. ``length`` is ignored if ``state`` is given. In this case the length of ``state`` is used. state Use a state chain to define the status at this time point (outbreak or not). If not given, a Markov chain is generated automatically. state_weight Additional weight for an outbreak which influences the distribution parameter :math:`\mu`. Returns ------- A ``DataFrame`` of an endemic time series where each row contains the case counts of this week. It also contains the mean case count value based on the underlying sinus model. """ if self.seed: base = robjects.packages.importr("base") base.set_seed(self.seed) simulated = surveillance.sim_seasonalNoise( A=self.amplitude, alpha=self.alpha, beta=self.trend, phi=self.seasonal_move, length=length, frequency=self.frequency, state=robjects.NULL if state is None else robjects.IntVector(state), K=robjects.NULL if state_weight is None else state_weight, ) simulated = r_list_to_frame(simulated, ["mu", "seasonalBackground"]) simulated = simulated.pipe(add_date_time_index_to_frame).rename( columns={ "mu": "mean", "seasonalBackground": "n_cases" }) return simulated
def simulate( self, length: int, state_weight: float = 0, state: Optional[Sequence[int]] = None, ) -> pd.DataFrame: """Simulate outbreaks. Parameters ---------- length Number of weeks to model. ``length`` is ignored if ``state`` is given. In this case, the length of ``state`` is used. state Use a state chain to define the status at this time point (outbreak or not). If not given, a Markov chain is generated automatically. state_weight Additional weight for an outbreak which influences the distribution parameter mu. Returns ------- A ``DataFrame`` of simulated case counts per week, separated into baseline and outbreak cases. """ if self.seed: base.set_seed(self.seed) simulated = surveillance.sim_pointSource( p=self.p, r=self.r, length=length, A=self.amplitude, alpha=self.alpha, beta=self.trend, phi=self.seasonal_move, frequency=self.frequency, state=robjects.NULL if state is None else robjects.IntVector(state), K=state_weight, ) simulated_as_frame = r_list_to_frame(simulated, ["observed", "state"]) return simulated_as_frame.pipe(add_date_time_index_to_frame).rename( columns={ "observed": "n_cases", "state": "is_outbreak" })