# Path: epi_forecast_stat_mech/high_level.py
# def get_simple_estimator_dict():
# def get_meta_estimator_dict(validation_times=(14, 28, 42)):
# def get_estimator_dict():
# 
# 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/evaluation/run_on_data.py
# def train_test_split_time(data, first_test_day):
# def shard_locations_randomly(data, num_shards=5, seed=0):


import collections
import functools
import numpy as np
import sklearn
from absl.testing import absltest
from absl.testing import parameterized
from epi_forecast_stat_mech import high_level
from epi_forecast_stat_mech import sir_sim
from epi_forecast_stat_mech.evaluation import run_on_data
from jax.config import config
# Lint as: python3
"""Tests for epi_forecast_stat_mech.high_level."""





config.parse_flags_with_absl()  # Necessary for running on TPU.


def create_synthetic_dataset(
    seed=0,
    num_epidemics=50,
    num_important_cov=1,
    num_unimportant_cov=2,
    train_length=100,
    prediction_length=10,
):
  """Creates synthetic data."""
  np.random.seed(seed)  # TODO(shoyer): use np.random.RandomState
  beta_fn = functools.partial(sir_sim.generate_betas_many_cov2,
                              num_pred=num_important_cov,
                              num_not_pred=num_unimportant_cov)


# Path: epi_forecast_stat_mech/high_level.py
# def get_simple_estimator_dict():
# def get_meta_estimator_dict(validation_times=(14, 28, 42)):
# def get_estimator_dict():
# 
# 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/evaluation/run_on_data.py
# def train_test_split_time(data, first_test_day):
# def shard_locations_randomly(data, num_shards=5, seed=0):


# Path: epi_forecast_stat_mech/tests/test_high_level.py
import collections
import functools
import numpy as np
import sklearn
from absl.testing import absltest
from absl.testing import parameterized
from epi_forecast_stat_mech import high_level
from epi_forecast_stat_mech import sir_sim
from epi_forecast_stat_mech.evaluation import run_on_data
from jax.config import config

# Lint as: python3
"""Tests for epi_forecast_stat_mech.high_level."""





config.parse_flags_with_absl()  # Necessary for running on TPU.


def create_synthetic_dataset(
    seed=0,
    num_epidemics=50,
    num_important_cov=1,
    num_unimportant_cov=2,
    train_length=100,
    prediction_length=10,
):
  """Creates synthetic data."""
  np.random.seed(seed)  # TODO(shoyer): use np.random.RandomState
  beta_fn = functools.partial(sir_sim.generate_betas_many_cov2,
                              num_pred=num_important_cov,
                              num_not_pred=num_unimportant_cov)


  trajectories = sir_sim.generate_simulations(