# Path: epi_forecast_stat_mech/mask_time.py
# class MaskTime():
#   def __init__(self, data):
#   def mask_min_infections(self, min_value):
#   def mask_to_recent(self, recent_day_limit):
#   def mask_dynamic_cov(self, data, cov_name, mask_value):
# def make_mask(data,
#               min_value=None,
#               recent_day_limit=None,
#               dynamic_cov_name=None,
#               dynamic_cov_value=0):


from absl.testing import absltest
from epi_forecast_stat_mech import mask_time
import numpy as np
import xarray as xr
  def test_basic_sanity(self):
    data = self.data.copy(deep=True)
    data.new_infections[dict(time=0)] = np.nan
    tm = mask_time.MaskTime(data)
    mask = tm.mask
    expected_mask = np.ones((self.num_locations, self.num_time_steps))
    expected_mask[:, 0] = 0
    np.testing.assert_array_equal(expected_mask.astype(bool), mask)

  def test_min_value(self):
    min_value = 3
    tm = mask_time.MaskTime(self.data)
    tm.mask_min_infections(min_value=min_value)
    mask = tm.mask
    expected_mask = np.ones((self.num_locations, self.num_time_steps))
    expected_mask[:, :min_value-1] = 0
    np.testing.assert_array_equal(expected_mask.astype(bool), mask)

  def test_dynamic_cov_ones(self):
    cov_name = '1'
    cov_value = 1
    tm = mask_time.MaskTime(self.data)
    tm.mask_dynamic_cov(self.data, cov_name, cov_value)
    mask = tm.mask
    expected_mask = np.zeros((self.num_locations, self.num_time_steps))
    np.testing.assert_array_equal(expected_mask.astype(bool), mask)

  def test_dynamic_cov_zeros(self):
    cov_name = '1'
    cov_value = 0

# Path: epi_forecast_stat_mech/mask_time.py
# class MaskTime():
#   def __init__(self, data):
#   def mask_min_infections(self, min_value):
#   def mask_to_recent(self, recent_day_limit):
#   def mask_dynamic_cov(self, data, cov_name, mask_value):
# def make_mask(data,
#               min_value=None,
#               recent_day_limit=None,
#               dynamic_cov_name=None,
#               dynamic_cov_value=0):


# Path: epi_forecast_stat_mech/tests/test_mask_time.py
from absl.testing import absltest
from epi_forecast_stat_mech import mask_time
import numpy as np
import xarray as xr

  def test_basic_sanity(self):
    data = self.data.copy(deep=True)
    data.new_infections[dict(time=0)] = np.nan
    tm = mask_time.MaskTime(data)
    mask = tm.mask
    expected_mask = np.ones((self.num_locations, self.num_time_steps))
    expected_mask[:, 0] = 0
    np.testing.assert_array_equal(expected_mask.astype(bool), mask)

  def test_min_value(self):
    min_value = 3
    tm = mask_time.MaskTime(self.data)
    tm.mask_min_infections(min_value=min_value)
    mask = tm.mask
    expected_mask = np.ones((self.num_locations, self.num_time_steps))
    expected_mask[:, :min_value-1] = 0
    np.testing.assert_array_equal(expected_mask.astype(bool), mask)

  def test_dynamic_cov_ones(self):
    cov_name = '1'
    cov_value = 1
    tm = mask_time.MaskTime(self.data)
    tm.mask_dynamic_cov(self.data, cov_name, cov_value)
    mask = tm.mask
    expected_mask = np.zeros((self.num_locations, self.num_time_steps))
    np.testing.assert_array_equal(expected_mask.astype(bool), mask)

  def test_dynamic_cov_zeros(self):
    cov_name = '1'
    cov_value = 0

    tm = mask_time.MaskTime(self.data)