/
test_raster.py
115 lines (93 loc) · 3.57 KB
/
test_raster.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import pathlib
import numpy as np
import pandas as pd
import pandas.testing as pdt
import pytest
import xarray as xr
import ptolemy as pt
URL = "https://d2ad6b4ur7yvpq.cloudfront.net/naturalearth-3.3.0/ne_110m_admin_0_countries.geojson"
DATA_PATH = pathlib.Path(__file__).parent / "test_data"
LIKE = xr.tutorial.load_dataset("air_temperature").air.isel(time=0)
LIKE["lon"] = LIKE["lon"] - 360
RASTER_STRATEGIES = ["all_touched", "centroid", "hybrid", "majority", "weighted"]
@pytest.mark.parametrize("as_file", [True, False])
def test_init(as_file):
like = LIKE
if as_file:
like.to_netcdf(DATA_PATH / "foo.nc")
r = pt.Rasterize(like=DATA_PATH / "foo.nc")
else:
r = pt.Rasterize(like=like)
assert r.shape == like.shape
assert r.coords["lat"].equals(like["lat"])
assert r.coords["lon"].equals(like["lon"])
@pytest.mark.parametrize("flatten,exp_size", [(True, 1), (False, 177)])
@pytest.mark.parametrize("idxkey", [None, "iso_a3"])
def test_read_shpf(flatten, exp_size, idxkey):
r = pt.Rasterize(like=LIKE)
r.read_shpf(URL, flatten=flatten, idxkey=idxkey)
assert len(r.geoms_idxs) == exp_size
assert r.idxkey == idxkey
def _do_rasterize(strategy):
r = pt.Rasterize(like=LIKE)
r.read_shpf(URL, idxkey="iso_a3")
idxr = r.rasterize(strategy=strategy, verbose=True)
return idxr
@pytest.mark.parametrize("strategy", RASTER_STRATEGIES)
def test_rasterize(strategy):
obs = _do_rasterize(strategy)
exp = xr.open_dataarray(DATA_PATH / f"{strategy}.nc", mode="r")
assert obs.equals(exp)
obs.close()
exp.close()
def test_df_to_raster_roundtrip():
r = pt.Rasterize(like=LIKE)
r.read_shpf(URL, idxkey="adm0_a3")
idxr = r.rasterize(strategy="hybrid", verbose=True)
df = pd.DataFrame(
{
"adm0_a3": ["USA"] * 2 + ["MEX"] * 2,
"year": [2015, 2020] * 2,
"data": [15, 20, 5, 10],
}
)
idx_map = {v: int(k) for k, v in idxr.attrs.items() if int(k) in np.unique(idxr)}
ds = pt.df_to_raster(df, idxr, "adm0_a3", idx_map, coords=["year"])
assert ds.sum() == 8195
obs_df = pt.raster_to_df(ds, idxr, func="max", idx_map=idx_map, idx_dim="adm0_a3")
pdt.assert_frame_equal(
df.sort_values(by=["adm0_a3", "year"]).reset_index(drop=True),
obs_df.sort_values(by=["adm0_a3", "year"]).reset_index(drop=True)[df.columns],
check_dtype=False,
)
def test_df_to_weighted_raster_roundtrip():
r = pt.Rasterize(like=LIKE)
r.read_shpf(URL, idxkey="adm0_a3")
idxr = r.rasterize(strategy="weighted", verbose=True)
df = pd.DataFrame(
{
"adm0_a3": ["USA"] * 2 + ["MEX"] * 2,
"year": [2015, 2020] * 2,
"data": [15, 20, 5, 10],
}
)
ds = pt.df_to_weighted_raster(df, idxr, extra_coords=["year"])
assert np.isclose(ds.data.sel(year=2015).sum(), 3349.6, rtol=1e-5)
obs_df = pt.raster_to_df(ds, idxr, func="max")
pdt.assert_frame_equal(
df.sort_values(by=["adm0_a3", "year"]).reset_index(drop=True),
obs_df.sort_values(by=["adm0_a3", "year"]).reset_index(drop=True),
check_dtype=False,
)
def test_cell_area():
exp = pt.cell_area_from_file(LIKE)
assert exp.index[0] > exp.index[1]
assert exp.iloc[0] < exp.iloc[1]
assert np.isclose(exp.sum(), 1299347051157.3113)
if __name__ == "__main__":
# save rasterized data for regression
for strategy in RASTER_STRATEGIES:
print(f"Working on {strategy}")
idxr = _do_rasterize(strategy)
idxr.to_netcdf(DATA_PATH / f"{strategy}.nc", mode="w")
idxr.close()