# Path: epi_forecast_stat_mech/data_model.py
# TIME_ZERO = np.datetime64('2019-12-31')
# def new_dataarray(parameters):
# def new_model(num_locations,
#               num_time_steps,
#               num_static_covariates,
#               num_dynamic_covariates=0):
# def _helper_shift_dataarray(shifts, array_to_shift):
# def _first_true(arr, axis):
# def shift_timeseries_by(data, shifts_all, enforce_positive_shifts: bool = True):
# def shift_timeseries(data, fraction_infected_limits, split_time):
# def validate_data(data,
#                   enable_regex_check=False,
#                   require_dynamics=False,
#                   require_samples=False,
#                   require_no_samples=False,
#                   require_integer_time=True,
#                   require_canonical_split_time=True):
# def validate_data_for_fit(data, **kwargs):
# def calculate_cumulative_infections(new_infections):
# def calculate_total_infections(new_infections):
# def calculate_infection_sums(ds):
# def _get_delta(series):
# def datetime_to_int(date, ds=None):
# def int_to_datetime(int_date, ds=None):
# def compute_integer_time(trajectories):
# def compute_numpy_index_time(trajectories):
# def convert_data_to_integer_time(trajectories, method='days_in_2020'):
# def set_sensible_canonical_split_time(data):
# 
# Path: epi_forecast_stat_mech/vc_sim.py
# SPLIT_TIME=100
# def final_size_poisson_dist(num_locations):
# def new_vc_simulation_model(num_locations,
#                             num_time_steps,
#                             num_static_covariates=1):
# def generate_ground_truth(pred_size, r, p, num_time_steps):
# def generate_simulations(final_size_fn,
#                          num_locations,
#                          num_time_steps=500,
#                          constant_r=1.0,
#                          constant_p=0.6,
#                          fraction_infected_limits=(0.05, 1.0),
#                          shift_timeseries=True):


from absl.testing import absltest
from epi_forecast_stat_mech import data_model
from epi_forecast_stat_mech import vc_sim
"""Tests for epi_forecast_stat_mech/vc_sim ."""



class TestGenerateVCSimulations(absltest.TestCase):

  def setUp(self):
    super(TestGenerateVCSimulations, self).setUp()
    self.num_locations = 50
    self.num_time_steps = 500

  def test_basic_sanity(self):
    trajectories = vc_sim.generate_simulations(

# Path: epi_forecast_stat_mech/data_model.py
# TIME_ZERO = np.datetime64('2019-12-31')
# def new_dataarray(parameters):
# def new_model(num_locations,
#               num_time_steps,
#               num_static_covariates,
#               num_dynamic_covariates=0):
# def _helper_shift_dataarray(shifts, array_to_shift):
# def _first_true(arr, axis):
# def shift_timeseries_by(data, shifts_all, enforce_positive_shifts: bool = True):
# def shift_timeseries(data, fraction_infected_limits, split_time):
# def validate_data(data,
#                   enable_regex_check=False,
#                   require_dynamics=False,
#                   require_samples=False,
#                   require_no_samples=False,
#                   require_integer_time=True,
#                   require_canonical_split_time=True):
# def validate_data_for_fit(data, **kwargs):
# def calculate_cumulative_infections(new_infections):
# def calculate_total_infections(new_infections):
# def calculate_infection_sums(ds):
# def _get_delta(series):
# def datetime_to_int(date, ds=None):
# def int_to_datetime(int_date, ds=None):
# def compute_integer_time(trajectories):
# def compute_numpy_index_time(trajectories):
# def convert_data_to_integer_time(trajectories, method='days_in_2020'):
# def set_sensible_canonical_split_time(data):
# 
# Path: epi_forecast_stat_mech/vc_sim.py
# SPLIT_TIME=100
# def final_size_poisson_dist(num_locations):
# def new_vc_simulation_model(num_locations,
#                             num_time_steps,
#                             num_static_covariates=1):
# def generate_ground_truth(pred_size, r, p, num_time_steps):
# def generate_simulations(final_size_fn,
#                          num_locations,
#                          num_time_steps=500,
#                          constant_r=1.0,
#                          constant_p=0.6,
#                          fraction_infected_limits=(0.05, 1.0),
#                          shift_timeseries=True):


# Path: epi_forecast_stat_mech/tests/test_vc_sim.py
from absl.testing import absltest
from epi_forecast_stat_mech import data_model
from epi_forecast_stat_mech import vc_sim

"""Tests for epi_forecast_stat_mech/vc_sim ."""



class TestGenerateVCSimulations(absltest.TestCase):

  def setUp(self):
    super(TestGenerateVCSimulations, self).setUp()
    self.num_locations = 50
    self.num_time_steps = 500

  def test_basic_sanity(self):
    trajectories = vc_sim.generate_simulations(

        vc_sim.final_size_poisson_dist,