def test_big_mean(self, tmp_path, monkeypatch): def mockiter(self): class MockIterator: def __init__(self): self.idx = 0 self.max_idx = 10 def __iter__(self): return self def __next__(self): if self.idx < self.max_idx: # batch_size = 10, timesteps = 2, num_features = 1 self.idx += 1 return (np.ones( (10, 2, 1)), np.ones( (10, ), dtype=np.int8), np.ones( (10, 2)), np.ones((10, 2)), np.ones( (10, 2)), np.ones((10, 2))), None else: raise StopIteration() return MockIterator() def do_nothing(self, data_path, batch_file_size, shuffle_data, mode, pred_months, surrounding_pixels, ignore_vars): pass monkeypatch.setattr(DataLoader, '__iter__', mockiter) monkeypatch.setattr(DataLoader, '__init__', do_nothing) model = LinearRegression(tmp_path) calculated_mean = model._calculate_big_mean() # 1 for the 2 features and for the first month, 0 for the rest expected_mean = np.array([ 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1. ]) # np.isclose because of rounding assert np.isclose(calculated_mean, expected_mean).all()
def test_big_mean(self, tmp_path, monkeypatch): def mockiter(self): class MockIterator: def __init__(self): self.idx = 0 self.max_idx = 10 def __iter__(self): return self def __next__(self): if self.idx < self.max_idx: # batch_size = 10, timesteps = 2, num_features = 1 self.idx += 1 return ( ( np.ones((10, 2, 1)), np.ones((10,), dtype=np.int8), np.ones((10, 2)), np.ones((10, 2)), np.ones((10, 2)), np.ones((10, 2)), np.ones((10, 1)), ), None, ) else: raise StopIteration() return MockIterator() def do_nothing( self, data_path, batch_file_size, mode, shuffle_data, clear_nans, normalize, experiment, mask, pred_months, to_tensor, surrounding_pixels, ignore_vars, monthly_aggs, static, device, predict_delta, spatial_mask, normalize_y, incl_yearly_aggs, ): pass monkeypatch.setattr(DataLoader, "__iter__", mockiter) monkeypatch.setattr(DataLoader, "__init__", do_nothing) model = LinearRegression(tmp_path, normalize_y=False) calculated_mean = model._calculate_big_mean() # 1 for the 2 features and for the first month, 0 for the rest expected_mean = np.array( [ 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, ] ) # np.isclose because of rounding assert np.isclose(calculated_mean, expected_mean).all()