# Path: epi_forecast_stat_mech/sir_sim.py
# SPLIT_TIME = 100
# def generate_betas_from_single_random_covariate(num_locations):
# def generate_betas_effect_mod(num_locations):
# def generate_betas_many_cov2(num_locations, num_pred=1, num_not_pred=2):
# def gen_dynamic_beta_random_time(num_locations, num_time_steps):
# def gen_social_distancing_weight(num_locations):
# def new_sir_simulation_model(num_locations, num_time_steps,
#                              num_static_covariates, num_dynamic_covariates=0):
# def _helper_ground_truth_setup(population_size,
#                                num_time_steps):
# def _helper_ground_truth_loop(num_susceptible, num_recovered, num_infected,
#                               beta_time_t, gamma, population_size,
#                               prob_infection_constant):
# def generate_ground_truth(population_size,
#                           beta,
#                           gamma,
#                           num_time_steps,
#                           prob_infection_constant=0.2):
# def generate_social_distancing_ground_truth(population_size,
#                                             beta,
#                                             gamma,
#                                             num_time_steps,
#                                             social_distancing_threshold,
#                                             gen_social_distancing_weight_fn,
#                                             prob_infection_constant=0.2
#                                             ):
# def _helper_setup_sir_sim(gen_constant_beta_fn,
#                           num_locations,
#                           num_time_steps=500,
#                           constant_gamma=0.33,
#                           population_size=10000,
#                           gen_dynamic_beta_fn=None):
# def generate_simulations(gen_constant_beta_fn,
#                          num_locations,
#                          num_time_steps=500,
#                          constant_gamma=0.33,
#                          population_size=10000,
#                          gen_dynamic_beta_fn=None,
#                          fraction_infected_limits=(.05, 1.),
#                          prob_infection_constant=0.2):
# def generate_social_distancing_simulations(gen_constant_beta_fn,
#                                            gen_social_distancing_weight_fn,
#                                            num_locations,
#                                            num_time_steps=500,
#                                            constant_gamma=0.33,
#                                            population_size=10000,
#                                            social_distancing_threshold=10000/4,
#                                            fraction_infected_limits=(.05, 1.),
#                                            prob_infection_constant=0.2,
#                                            shift_timeseries=True):


import functools
import numpy as np
import xarray as xr
from absl.testing import absltest
from epi_forecast_stat_mech import sir_sim
        gen_dynamic_beta_fn=dynamic_beta_fn)

    shift_growth_rate = trajectories.growth_rate.shift(time=1)
    num_diff_betas = (trajectories.growth_rate.sel(
        time=slice(None, self.num_time_steps-1)) !=
                      shift_growth_rate.sel(time=slice(1, None))).sum()
    assert num_diff_betas >= self.num_epidemics

  def test_social_distancing_sanity(self):
    beta_fn = functools.partial(sir_sim.generate_betas_many_cov2,
                                num_pred=self.num_important_cov,
                                num_not_pred=self.num_unimportant_cov)
    trajectories = sir_sim.generate_social_distancing_simulations(
        beta_fn,
        sir_sim.gen_social_distancing_weight,
        self.num_epidemics,
        self.num_time_steps)
    assert trajectories.new_infections.shape == (self.num_epidemics,
                                                 self.num_time_steps)

  def test_population_size(self):
    beta_fn = functools.partial(
        sir_sim.generate_betas_many_cov2,
        num_pred=self.num_important_cov,
        num_not_pred=self.num_unimportant_cov)
    population_size = xr.DataArray(
        50000 * np.ones((self.num_epidemics)),
        dims=['location'])
    trajectories = sir_sim.generate_social_distancing_simulations(
        beta_fn,

# Path: epi_forecast_stat_mech/sir_sim.py
# SPLIT_TIME = 100
# def generate_betas_from_single_random_covariate(num_locations):
# def generate_betas_effect_mod(num_locations):
# def generate_betas_many_cov2(num_locations, num_pred=1, num_not_pred=2):
# def gen_dynamic_beta_random_time(num_locations, num_time_steps):
# def gen_social_distancing_weight(num_locations):
# def new_sir_simulation_model(num_locations, num_time_steps,
#                              num_static_covariates, num_dynamic_covariates=0):
# def _helper_ground_truth_setup(population_size,
#                                num_time_steps):
# def _helper_ground_truth_loop(num_susceptible, num_recovered, num_infected,
#                               beta_time_t, gamma, population_size,
#                               prob_infection_constant):
# def generate_ground_truth(population_size,
#                           beta,
#                           gamma,
#                           num_time_steps,
#                           prob_infection_constant=0.2):
# def generate_social_distancing_ground_truth(population_size,
#                                             beta,
#                                             gamma,
#                                             num_time_steps,
#                                             social_distancing_threshold,
#                                             gen_social_distancing_weight_fn,
#                                             prob_infection_constant=0.2
#                                             ):
# def _helper_setup_sir_sim(gen_constant_beta_fn,
#                           num_locations,
#                           num_time_steps=500,
#                           constant_gamma=0.33,
#                           population_size=10000,
#                           gen_dynamic_beta_fn=None):
# def generate_simulations(gen_constant_beta_fn,
#                          num_locations,
#                          num_time_steps=500,
#                          constant_gamma=0.33,
#                          population_size=10000,
#                          gen_dynamic_beta_fn=None,
#                          fraction_infected_limits=(.05, 1.),
#                          prob_infection_constant=0.2):
# def generate_social_distancing_simulations(gen_constant_beta_fn,
#                                            gen_social_distancing_weight_fn,
#                                            num_locations,
#                                            num_time_steps=500,
#                                            constant_gamma=0.33,
#                                            population_size=10000,
#                                            social_distancing_threshold=10000/4,
#                                            fraction_infected_limits=(.05, 1.),
#                                            prob_infection_constant=0.2,
#                                            shift_timeseries=True):


# Path: epi_forecast_stat_mech/tests/test_sir_sim.py
import functools
import numpy as np
import xarray as xr
from absl.testing import absltest
from epi_forecast_stat_mech import sir_sim

        gen_dynamic_beta_fn=dynamic_beta_fn)

    shift_growth_rate = trajectories.growth_rate.shift(time=1)
    num_diff_betas = (trajectories.growth_rate.sel(
        time=slice(None, self.num_time_steps-1)) !=
                      shift_growth_rate.sel(time=slice(1, None))).sum()
    assert num_diff_betas >= self.num_epidemics

  def test_social_distancing_sanity(self):
    beta_fn = functools.partial(sir_sim.generate_betas_many_cov2,
                                num_pred=self.num_important_cov,
                                num_not_pred=self.num_unimportant_cov)
    trajectories = sir_sim.generate_social_distancing_simulations(
        beta_fn,
        sir_sim.gen_social_distancing_weight,
        self.num_epidemics,
        self.num_time_steps)
    assert trajectories.new_infections.shape == (self.num_epidemics,
                                                 self.num_time_steps)

  def test_population_size(self):
    beta_fn = functools.partial(
        sir_sim.generate_betas_many_cov2,
        num_pred=self.num_important_cov,
        num_not_pred=self.num_unimportant_cov)
    population_size = xr.DataArray(
        50000 * np.ones((self.num_epidemics)),
        dims=['location'])
    trajectories = sir_sim.generate_social_distancing_simulations(
        beta_fn,

        sir_sim.gen_social_distancing_weight,