def test_chop_roi(self, tmp_path):
size_original = (80, 80)
original_ds, _, _ = _make_dataset(size_original)
original_shape = original_ds.VHI.shape
processor = BasePreProcessor(tmp_path)
subset_str = 'east_africa'
new_ds = processor.chop_roi(ds=original_ds, subset_str=subset_str)
output_shape = new_ds.VHI.shape
assert original_shape != output_shape, f"The chop_roi should lead to\
smaller datasets than the original. Expected output_shape: {output_shape}\
to be different from original_shape: {original_shape}"
assert ((new_ds.lat.values.min() >= -11) &
(new_ds.lat.values.max() <= 23)
), f"Expected latitude to be in the range -11 : 23. Currently:\
{new_ds.lat.values.min()} : {new_ds.lat.values.max()}"
assert (
(new_ds.lon.values.min() >= 21) & (new_ds.lon.values.max() <= 51.8)
), f"Expected longitude to be in the range 21 : 51.8. Currently:\
# create forecast_horizon
fh = pd.to_timedelta(ds_new.time.values - ds_new.initialisation_date.values)
ds_new["time"] = fh
ds_new = ds_new.rename({"time": "forecast_horizon"})
# create a new coord
time = ds_new.initialisation_date + ds_new.forecast_horizon
ds_new = ds_new.assign_coords(time=time)
# ------------------------------------------------------------------------------
# Test PREPROCESSING
# ------------------------------------------------------------------------------
from src.preprocess.base import BasePreProcessor
b = BasePreProcessor()
ds1_kenya = b.chop_roi(ds1, inverse_lat=True)
ds2_kenya = b.chop_roi(ds2, inverse_lat=True)
# concat across initialisation dates
ds_kenya = xr.concat([ds1_kenya, ds2_kenya], dim="initialisation_date")
stacked = ds_kenya.stack(time=("initialisation_date", "forecast_horizon"))
# stack each individually
k1 = ds1_kenya.stack(time=("initialisation_date", "forecast_horizon"))
k2 = ds2_kenya.stack(time=("initialisation_date", "forecast_horizon"))
# test selectors
stacked.sel(forecast_horizon=np.timedelta64(28, "D"))
stacked.sel(initialisation_date="1997-01-01")
stacked.swap_dims({"time": "valid_time"}).sel(valid_time="1997-04")