# Path: epi_forecast_stat_mech/datasets/nanhandling.py
# def _assume_2d(array):
# def fillna_zero(array):
# def median_filter(array, filter_days=5):
# def mean_filter(array, filter_days=5):
# def fillna_beginning(array, value=0):
#   def fillna_beginning_1d(array):
# def fillna_ffill(array):
# def fillna_bfill(array):
# def fillna_interp(array):
#   def interp_1d(array):
# def plot_nans(array, title, ax=None):
# def longest_nans(array):
#   def longest_nans_1d(array):


from absl.testing import absltest
from epi_forecast_stat_mech.datasets import nanhandling
import numpy as np
                           [np.nan, np.nan, np.nan, 3, np.nan]])

  def test_fillna_zero(self):
    result = nanhandling.fillna_zero(self.array)
    expected = np.array([[5, 0, 0, 7, 2], [3, 0, 1, 8, 0], [0, 0, 0, 3, 1],
                         [0, 0, 0, 0, 0], [0, 0, 0, 3, 0]])
    self.assertTrue(np.array_equal(expected, result))

  def test_median_filter(self):
    result = nanhandling.median_filter(self.array)
    expected = np.array([[5, 5, 5, 2, 2], [3, 3, 3, 4.5, 4.5],
                         [np.nan, 3, 2, 1, 1],
                         [np.nan, np.nan, np.nan, np.nan, np.nan],
                         [np.nan, 3, 3, 3, 3]])
    np.testing.assert_array_equal(result, expected)

  def test_mean_filter(self):
    result = nanhandling.mean_filter(self.array)
    expected = np.array(
        [[
            5, (5 + 5 + 7) / 3, (5 + 7 + 2) / 3, (7 + 2 + 2) / 3,
            (7 + 2 + 2 + 2) / 4
        ],
         [(3 + 3 + 3 + 1) / 4, (3 + 3 + 1 + 8) / 4, (3 + 1 + 8) / 3,
          (1 + 8) / 2, (1 + 8) / 2],
         [np.nan, 3, (3 + 1) / 2, (3 + 1 + 1) / 3, (3 + 1 + 1 + 1) / 4],
         [np.nan, np.nan, np.nan, np.nan, np.nan], [np.nan, 3, 3, 3, 3]])
    np.testing.assert_array_equal(result, expected)

  def test_fillna_beginning(self):

# Path: epi_forecast_stat_mech/datasets/nanhandling.py
# def _assume_2d(array):
# def fillna_zero(array):
# def median_filter(array, filter_days=5):
# def mean_filter(array, filter_days=5):
# def fillna_beginning(array, value=0):
#   def fillna_beginning_1d(array):
# def fillna_ffill(array):
# def fillna_bfill(array):
# def fillna_interp(array):
#   def interp_1d(array):
# def plot_nans(array, title, ax=None):
# def longest_nans(array):
#   def longest_nans_1d(array):


# Path: epi_forecast_stat_mech/datasets/nanhandling_test.py
from absl.testing import absltest
from epi_forecast_stat_mech.datasets import nanhandling
import numpy as np

                           [np.nan, np.nan, np.nan, 3, np.nan]])

  def test_fillna_zero(self):
    result = nanhandling.fillna_zero(self.array)
    expected = np.array([[5, 0, 0, 7, 2], [3, 0, 1, 8, 0], [0, 0, 0, 3, 1],
                         [0, 0, 0, 0, 0], [0, 0, 0, 3, 0]])
    self.assertTrue(np.array_equal(expected, result))

  def test_median_filter(self):
    result = nanhandling.median_filter(self.array)
    expected = np.array([[5, 5, 5, 2, 2], [3, 3, 3, 4.5, 4.5],
                         [np.nan, 3, 2, 1, 1],
                         [np.nan, np.nan, np.nan, np.nan, np.nan],
                         [np.nan, 3, 3, 3, 3]])
    np.testing.assert_array_equal(result, expected)

  def test_mean_filter(self):
    result = nanhandling.mean_filter(self.array)
    expected = np.array(
        [[
            5, (5 + 5 + 7) / 3, (5 + 7 + 2) / 3, (7 + 2 + 2) / 3,
            (7 + 2 + 2 + 2) / 4
        ],
         [(3 + 3 + 3 + 1) / 4, (3 + 3 + 1 + 8) / 4, (3 + 1 + 8) / 3,
          (1 + 8) / 2, (1 + 8) / 2],
         [np.nan, 3, (3 + 1) / 2, (3 + 1 + 1) / 3, (3 + 1 + 1 + 1) / 4],
         [np.nan, np.nan, np.nan, np.nan, np.nan], [np.nan, 3, 3, 3, 3]])
    np.testing.assert_array_equal(result, expected)

  def test_fillna_beginning(self):

    result = nanhandling.fillna_beginning(self.array)