diff --git a/primap2/tests/test_downscale.py b/primap2/tests/test_downscale.py index ca2f59b..a50e4cb 100644 --- a/primap2/tests/test_downscale.py +++ b/primap2/tests/test_downscale.py @@ -10,7 +10,8 @@ from .utils import allclose, assert_equal -def test_downscale_gas_timeseries(empty_ds): +@pytest.fixture +def gas_downscaling_ds(empty_ds): for key in empty_ds: empty_ds[key].pint.magnitude[:] = np.nan empty_ds["CO2"].loc[{"time": "2002"}] = 1 * ureg("Gg CO2 / year") @@ -22,11 +23,42 @@ def test_downscale_gas_timeseries(empty_ds): empty_ds["KYOTOGHG (AR4GWP100)"].loc[{"time": "2020"}] = ( 2 * (1 + sf6 + ch4) * ureg("Gg CO2 / year") ) + return empty_ds - downscaled = empty_ds.pr.downscale_gas_timeseries( + +@pytest.fixture +def dim_downscaling_ds(empty_ds): + for key in empty_ds: + empty_ds[key].pint.magnitude[:] = np.nan + t = empty_ds.loc[{"area (ISO3)": "BOL"}].copy() + t["area (ISO3)"] = ["CAMB"] # here, the sum of COL, ARG, MEX, and BOL + ds = xr.concat([empty_ds, t], dim="area (ISO3)") + da: xr.DataArray = ds["CO2"] + + da.loc[{"area (ISO3)": ["COL", "ARG", "MEX"], "time": "2002"}] = 1 * ureg("Gg CO2 / year") + da.loc[{"area (ISO3)": "BOL", "time": "2002"}] = 3 * ureg("Gg CO2 / year") + da.loc[{"area (ISO3)": "CAMB", "time": "2002"}] = 6 * ureg("Gg CO2 / year") + + da.loc[{"area (ISO3)": ["COL", "ARG", "MEX", "BOL"], "time": "2012"}] = 2 * ureg( + "Gg CO2 / year" + ) + + da.loc[{"area (ISO3)": "CAMB", "source": "RAND2020"}] = np.concatenate( + [np.array([6] * 11), np.stack([8, 8]), np.linspace(8, 10, 8)] + ) * ureg("Gg CO2 / year") + return ds + + +@pytest.fixture +def dim_downscaling_da(dim_downscaling_ds): + return dim_downscaling_ds["CO2"] + + +def test_downscale_gas_timeseries(gas_downscaling_ds): + downscaled = gas_downscaling_ds.pr.downscale_gas_timeseries( basket="KYOTOGHG (AR4GWP100)", basket_contents=["CO2", "SF6", "CH4"] ) - expected = empty_ds.copy() + expected = gas_downscaling_ds.copy() expected["CO2"][:] = 1 * ureg("Gg CO2 / year") expected["SF6"][:] = 1 * ureg("Gg SF6 / year") expected["CH4"][:] = 1 * ureg("Gg CH4 / year") @@ -40,45 +72,26 @@ def test_downscale_gas_timeseries(empty_ds): ValueError, match="Only one of 'skipna' and 'skipna_evaluation_dims' may be supplied, not both.", ): - empty_ds.pr.downscale_gas_timeseries( + gas_downscaling_ds.pr.downscale_gas_timeseries( basket="KYOTOGHG (AR4GWP100)", basket_contents=["CO2", "SF6", "CH4"], skipna_evaluation_dims=["time"], skipna=True, ) - empty_ds["SF6"].loc[{"time": "2002"}] = 2 * ureg("Gg SF6 / year") + gas_downscaling_ds["SF6"].loc[{"time": "2002"}] = 2 * ureg("Gg SF6 / year") with pytest.raises(ValueError, match="To continue regardless, set check_consistency=False"): - empty_ds.pr.downscale_gas_timeseries( + gas_downscaling_ds.pr.downscale_gas_timeseries( basket="KYOTOGHG (AR4GWP100)", basket_contents=["CO2", "SF6", "CH4"] ) -def test_downscale_timeseries(empty_ds): - for key in empty_ds: - empty_ds[key].pint.magnitude[:] = np.nan - t = empty_ds.loc[{"area (ISO3)": "BOL"}].copy() - t["area (ISO3)"] = ["CAMB"] # here, the sum of COL, ARG, MEX, and BOL - ds = xr.concat([empty_ds, t], dim="area (ISO3)") - da: xr.DataArray = ds["CO2"] - - da.loc[{"area (ISO3)": ["COL", "ARG", "MEX"], "time": "2002"}] = 1 * ureg("Gg CO2 / year") - da.loc[{"area (ISO3)": "BOL", "time": "2002"}] = 3 * ureg("Gg CO2 / year") - da.loc[{"area (ISO3)": "CAMB", "time": "2002"}] = 6 * ureg("Gg CO2 / year") - - da.loc[{"area (ISO3)": ["COL", "ARG", "MEX", "BOL"], "time": "2012"}] = 2 * ureg( - "Gg CO2 / year" - ) - - da.loc[{"area (ISO3)": "CAMB", "source": "RAND2020"}] = np.concatenate( - [np.array([6] * 11), np.stack([8, 8]), np.linspace(8, 10, 8)] - ) * ureg("Gg CO2 / year") - - downscaled = da.pr.downscale_timeseries( +def test_downscale_timeseries(dim_downscaling_ds, dim_downscaling_da): + downscaled = dim_downscaling_da.pr.downscale_timeseries( dim="area (ISO3)", basket="CAMB", basket_contents=["COL", "ARG", "MEX", "BOL"] ) - expected = da.copy() + expected = dim_downscaling_da.copy() expected.loc[{"area (ISO3)": ["COL", "ARG", "MEX"], "source": "RAND2020"}] = np.broadcast_to( np.concatenate( @@ -107,7 +120,7 @@ def test_downscale_timeseries(empty_ds): downscaled.loc[{"area (ISO3)": ["COL", "ARG", "MEX", "BOL"]}].sum(dim="area (ISO3)"), ) - downscaled_ds = ds.pr.downscale_timeseries( + downscaled_ds = dim_downscaling_ds.pr.downscale_timeseries( dim="area (ISO3)", basket="CAMB", basket_contents=["COL", "ARG", "MEX", "BOL"] ) assert_equal(downscaled_ds["CO2"], expected, equal_nan=True, atol=0.01) @@ -116,7 +129,7 @@ def test_downscale_timeseries(empty_ds): ValueError, match="Only one of 'skipna' and 'skipna_evaluation_dims' may be supplied, not both.", ): - ds.pr.downscale_timeseries( + dim_downscaling_ds.pr.downscale_timeseries( dim="area (ISO3)", basket="CAMB", basket_contents=["COL", "ARG", "MEX", "BOL"], @@ -124,22 +137,22 @@ def test_downscale_timeseries(empty_ds): skipna=True, ) - da.loc[{"area (ISO3)": "BOL", "time": "2002"}] = 2 * ureg("Gg CO2 / year") + dim_downscaling_da.loc[{"area (ISO3)": "BOL", "time": "2002"}] = 2 * ureg("Gg CO2 / year") with pytest.raises(ValueError, match="To continue regardless, set check_consistency=False"): - da.pr.downscale_timeseries( + dim_downscaling_da.pr.downscale_timeseries( dim="area (ISO3)", basket="CAMB", basket_contents=["COL", "ARG", "MEX", "BOL"], ) - downscaled = da.pr.downscale_timeseries( + downscaled = dim_downscaling_da.pr.downscale_timeseries( dim="area (ISO3)", basket="CAMB", basket_contents=["COL", "ARG", "MEX", "BOL"], check_consistency=False, ) - expected = da.copy() + expected = dim_downscaling_da.copy() expected.loc[{"area (ISO3)": ["COL", "ARG", "MEX"], "source": "RAND2020"}] = np.broadcast_to( np.concatenate( @@ -161,14 +174,14 @@ def test_downscale_timeseries(empty_ds): assert_equal(downscaled, expected, equal_nan=True, atol=0.01) - downscaled = da.pr.downscale_timeseries( + downscaled = dim_downscaling_da.pr.downscale_timeseries( dim="area (ISO3)", basket="CAMB", basket_contents=["COL", "ARG", "MEX", "BOL"], check_consistency=False, sel={"time": slice("2005", "2020")}, ) - expected = da.copy() + expected = dim_downscaling_da.copy() expected.loc[{"area (ISO3)": ["COL", "ARG", "MEX", "BOL"], "source": "RAND2020"}] = ( np.broadcast_to( @@ -201,3 +214,6 @@ def test_downscale_timeseries(empty_ds): expected.loc[{"area (ISO3)": "BOL", "time": "2002"}] = 2 * ureg("Gg CO2 / year") assert_equal(downscaled, expected, equal_nan=True, atol=0.01) + + +# def test_downscale_timeseries_zero(empty_ds):