Skip to content

Commit

Permalink
Moved data preparation for downscaling function tests to individual f…
Browse files Browse the repository at this point in the history 
…unctions.
  • Loading branch information
@JGuetschow
JGuetschow committed Oct 22, 2024
1 parent 5a5deb0 commit bcfac31
Showing 1 changed file with 52 additions and 36 deletions.
88 changes: 52 additions & 36 deletions primap2/tests/test_downscale.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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")
Expand All @@ -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(
Expand Down Expand Up @@ -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)
Expand All @@ -116,30 +129,30 @@ 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"],
skipna_evaluation_dims=["time"],
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(
Expand All @@ -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(
Expand Down Expand Up @@ -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):

0 comments on commit bcfac31

Please sign in to comment.