def test_time_simple(self):
da = xr.open_dataset(self.nc_file).tasmax
lon = -72.4
lat = 46.1
da = xr.open_dataset(self.nc_poslons).tas
da["lon"] -= 360
yr_st = "2050"
yr_ed = "2059"
out = subset.subset_gridpoint(da,
lon=lon,
lat=lat,
start_date=yr_st,
end_date=yr_ed)
np.testing.assert_almost_equal(out.lon, lon, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
np.testing.assert_array_equal(len(np.unique(out.time.dt.year)), 10)
np.testing.assert_array_equal(out.time.dt.year.max(), int(yr_ed))
np.testing.assert_array_equal(out.time.dt.year.min(), int(yr_st))
# test different but equivalent strings
out = subset.subset_gridpoint(da, start_date=yr_st, end_date=yr_ed)
out1 = subset.subset_gridpoint(da,
start_date=f"{yr_st}-01",
end_date=f"{yr_ed}-12")
out2 = subset.subset_gridpoint(da,
start_date=f"{yr_st}-01-01",
end_date=f"{yr_ed}-12-31")
np.testing.assert_array_equal(out, out1)
np.testing.assert_array_equal(out, out2)
np.testing.assert_array_equal(len(np.unique(out.time.dt.year)), 10)
np.testing.assert_array_equal(out.time.dt.year.max(), int(yr_ed))
np.testing.assert_array_equal(out.time.dt.year.min(), int(yr_st))
def test_simple(self):
da = xr.open_dataset(self.nc_file).tasmax
lon = -72.4
lat = 46.1
out = subset.subset_gridpoint(da, lon=lon, lat=lat)
np.testing.assert_almost_equal(out.lon, lon, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
da = xr.open_dataset(self.nc_poslons).tas
da["lon"] -= 360
yr_st = 2050
yr_ed = 2059
out = subset.subset_gridpoint(da,
lon=lon,
lat=lat,
start_yr=yr_st,
end_yr=yr_ed)
np.testing.assert_almost_equal(out.lon, lon, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
np.testing.assert_array_equal(len(np.unique(out.time.dt.year)), 10)
np.testing.assert_array_equal(out.time.dt.year.max(), yr_ed)
np.testing.assert_array_equal(out.time.dt.year.min(), yr_st)
# test time only
out = subset.subset_gridpoint(da, start_yr=yr_st, end_yr=yr_ed)
np.testing.assert_array_equal(len(np.unique(out.time.dt.year)), 10)
np.testing.assert_array_equal(out.time.dt.year.max(), yr_ed)
np.testing.assert_array_equal(out.time.dt.year.min(), yr_st)
def test_time_start_only(self):
da = xr.open_dataset(self.nc_poslons).tas
da["lon"] -= 360
yr_st = "2050"
# start date only
out = subset.subset_gridpoint(
da, start_date="{yr_st}-01".format(yr_st=yr_st))
np.testing.assert_array_equal(out.time.dt.year.min(), int(yr_st))
np.testing.assert_array_equal(out.time.dt.year.max(),
da.time.dt.year.max())
out = subset.subset_gridpoint(
da, start_date="{yr_st}-07".format(yr_st=yr_st))
np.testing.assert_array_equal(out.time.dt.year.min(), int(yr_st))
np.testing.assert_array_equal(out.time.min().dt.month, 7)
np.testing.assert_array_equal(out.time.dt.year.max(),
da.time.dt.year.max())
np.testing.assert_array_equal(out.time.max(), da.time.max())
out = subset.subset_gridpoint(
da, start_date="{yr_st}-07-15".format(yr_st=yr_st))
np.testing.assert_array_equal(out.time.dt.year.min(), int(yr_st))
np.testing.assert_array_equal(out.time.min().dt.month, 7)
np.testing.assert_array_equal(out.time.min().dt.day, 15)
np.testing.assert_array_equal(out.time.dt.year.max(),
da.time.dt.year.max())
np.testing.assert_array_equal(out.time.max(), da.time.max())
def test_raise(self):
da = xr.open_dataset(self.nc_poslons).tas
with pytest.raises(ValueError):
subset.subset_gridpoint(da,
lon=-72.4,
lat=46.1,
start_yr=2056,
end_yr=2055)
def test_tolerance(self):
da = xr.open_dataset(self.nc_poslons).tas
lon = -72.5
lat = 46.2
out = subset.subset_gridpoint(da, lon=lon, lat=lat, tolerance=1)
assert out.isnull().all()
subset.subset_gridpoint(da, lon=lon, lat=lat, tolerance=1e5)
def test_tolerance(self):
da = xr.open_dataset(self.nc_poslons).tas
lon = -72.5
lat = 46.2
with pytest.raises(ValueError):
subset.subset_gridpoint(da, lon=lon, lat=lat, tolerance=1)
subset.subset_gridpoint(da, lon=lon, lat=lat, tolerance=1e5)
def test_type_warn_then_raise(self):
da = xr.open_dataset(self.nc_poslons).tas
with pytest.raises(ValueError):
with pytest.warns(Warning):
subset.subset_gridpoint(da,
lon=-72.4,
lat=46.1,
start_date=2056,
end_date=2055)
def test_positive_lons(self):
da = xr.open_dataset(self.nc_poslons).tas
lon = -72.4
lat = 46.1
out = subset.subset_gridpoint(da, lon=lon, lat=lat)
np.testing.assert_almost_equal(out.lon, lon + 360, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
out = subset.subset_gridpoint(da, lon=lon + 360, lat=lat)
np.testing.assert_almost_equal(out.lon, lon + 360, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
def test_raise(self):
da = xr.open_dataset(self.nc_poslons).tas
with pytest.raises(ValueError):
subset.subset_gridpoint(da,
lon=-72.4,
lat=46.1,
start_yr=2056,
end_yr=2055)
da = xr.open_dataset(self.nc_2dlonlat).tasmax.drop(["lon", "lat"])
with pytest.raises(Exception):
subset.subset_gridpoint(da, lon=-72.4, lat=46.1)
def test_irregular(self):
da = xr.open_dataset(self.nc_2dlonlat).tasmax
lon = -72.4
lat = 46.1
out = subset.subset_gridpoint(da, lon=lon, lat=lat)
np.testing.assert_almost_equal(out.lon, lon, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
def test_simple(self, lat, lon, add_distance):
da = xr.open_dataset(self.nc_file).tasmax
out = subset.subset_gridpoint(da, lon=lon, lat=lat, add_distance=add_distance)
np.testing.assert_almost_equal(out.lon, lon, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
assert ("site" in out.dims) ^ (len(lat) == 1)
assert ("distance" in out.coords) ^ (not add_distance)
def test_irregular(self):
da = xr.open_dataset(self.nc_2dlonlat).tasmax
lon = -72.4
lat = 46.1
out = subset.subset_gridpoint(da, lon=lon, lat=lat)
np.testing.assert_almost_equal(out.lon, lon, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
# test_irregular transposed:
da1 = xr.open_dataset(self.nc_2dlonlat).tasmax
dims = list(da1.dims)
dims.reverse()
daT = xr.DataArray(np.transpose(da1.values), dims=dims)
for d in daT.dims:
args = dict()
args[d] = da1[d]
daT = daT.assign_coords(**args)
daT = daT.assign_coords(lon=(["rlon", "rlat"],
np.transpose(da1.lon.values)))
daT = daT.assign_coords(lat=(["rlon", "rlat"],
np.transpose(da1.lat.values)))
out1 = subset.subset_gridpoint(daT, lon=lon, lat=lat)
np.testing.assert_almost_equal(out1.lon, lon, 1)
np.testing.assert_almost_equal(out1.lat, lat, 1)
np.testing.assert_array_equal(out, out1)
# Dataset with tasmax, lon and lat as data variables (i.e. lon, lat not coords of tasmax)
daT1 = xr.DataArray(np.transpose(da1.values), dims=dims)
for d in daT1.dims:
args = dict()
args[d] = da1[d]
daT1 = daT1.assign_coords(**args)
dsT = xr.Dataset(data_vars=None, coords=daT1.coords)
dsT["tasmax"] = daT1
dsT["lon"] = xr.DataArray(np.transpose(da1.lon.values),
dims=["rlon", "rlat"])
dsT["lat"] = xr.DataArray(np.transpose(da1.lat.values),
dims=["rlon", "rlat"])
out2 = subset.subset_gridpoint(dsT, lon=lon, lat=lat)
np.testing.assert_almost_equal(out2.lon, lon, 1)
np.testing.assert_almost_equal(out2.lat, lat, 1)
np.testing.assert_array_equal(out, out2.tasmax)
def test_dataset(self):
da = xr.open_mfdataset(
[self.nc_file,
self.nc_file.replace("tasmax", "tasmin")])
lon = -72.4
lat = 46.1
out = subset.subset_gridpoint(da, lon=lon, lat=lat)
np.testing.assert_almost_equal(out.lon, lon, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
np.testing.assert_array_equal(out.tasmin.shape, out.tasmax.shape)
def test_time_incomplete_years(self):
da = xr.open_dataset(self.nc_poslons).tas
da["lon"] -= 360
yr_st = "2050"
yr_ed = "2059"
out = subset.subset_gridpoint(da,
start_date=f"{yr_st}-07-01",
end_date=f"{yr_ed}-06-30")
out1 = subset.subset_gridpoint(da,
start_date=f"{yr_st}-07",
end_date=f"{yr_ed}-06")
np.testing.assert_array_equal(out, out1)
np.testing.assert_array_equal(out.time.dt.year.min(), int(yr_st))
np.testing.assert_array_equal(out.time.min().dt.month, 7)
np.testing.assert_array_equal(out.time.min().dt.day, 1)
np.testing.assert_array_equal(out.time.dt.year.max(), int(yr_ed))
np.testing.assert_array_equal(out.time.max().dt.month, 6)
np.testing.assert_array_equal(out.time.max().dt.day, 30)
def test_time_end_only(self):
da = xr.open_dataset(self.nc_poslons).tas
da["lon"] -= 360
yr_ed = "2059"
# end date only
out = subset.subset_gridpoint(
da, end_date="{yr_ed}-01".format(yr_ed=yr_ed))
np.testing.assert_array_equal(out.time.dt.year.max(), int(yr_ed))
np.testing.assert_array_equal(out.time.max().dt.month, 1)
np.testing.assert_array_equal(out.time.max().dt.day, 31)
np.testing.assert_array_equal(out.time.min(), da.time.min())
out = subset.subset_gridpoint(
da, end_date="{yr_ed}-06-15".format(yr_ed=yr_ed))
np.testing.assert_array_equal(out.time.dt.year.max(), int(yr_ed))
np.testing.assert_array_equal(out.time.max().dt.month, 6)
np.testing.assert_array_equal(out.time.max().dt.day, 15)
np.testing.assert_array_equal(out.time.min(), da.time.min())
def _subset_function(dataset):
nonlocal count
count += 1
percentage = start_percentage + int((count - 1) / n_files * (end_percentage - start_percentage))
self.write_log(f"Processing file {count} of {n_files}", response, percentage)
dataset = dataset[variables] if variables else dataset
if lat1 is None and lon1 is None:
return subset_gridpoint(dataset, lon=lon0, lat=lat0, start_yr=y0, end_yr=y1)
else:
return subset_bbox(
dataset, lon_bnds=[lon0, lon1], lat_bnds=[lat0, lat1], start_yr=y0, end_yr=y1
)
def test_time_dates_outofbounds(self):
da = xr.open_dataset(self.nc_poslons).tas
da["lon"] -= 360
yr_st = "1950"
yr_ed = "2099"
with pytest.warns(UserWarning):
out = subset.subset_gridpoint(da,
start_date=f"{yr_st}-01",
end_date=f"{yr_ed}-01")
np.testing.assert_array_equal(out.time.dt.year.min(),
da.time.dt.year.min())
np.testing.assert_array_equal(out.time.dt.year.max(),
da.time.dt.year.max())
def compare_vals(self, ds, sub, vari, flag_2d=False):
# check subsetted values against original
imask = np.where(~np.isnan(sub[vari].isel(time=0)))
if len(imask[0]) > 70:
np.random.RandomState = 42
ii = np.random.randint(0, len(imask[0]), 70)
else:
ii = np.arange(0, len(imask[0]))
for i in zip(imask[0][ii], imask[1][ii]):
if flag_2d:
lat1 = sub.lat[i[0], i[1]]
lon1 = sub.lon[i[0], i[1]]
np.testing.assert_array_equal(
subset.subset_gridpoint(sub, lon=lon1, lat=lat1)[vari],
subset.subset_gridpoint(ds, lon=lon1, lat=lat1)[vari],
)
else:
lat1 = sub.lat.isel(lat=i[0])
lon1 = sub.lon.isel(lon=i[1])
# print(lon1.values, lat1.values)
np.testing.assert_array_equal(
sub[vari].sel(lon=lon1, lat=lat1), ds[vari].sel(lon=lon1,
lat=lat1))
def _subset_function(dataset):
nonlocal count
count += 1
percentage = start_percentage + int(
(count - 1) / n_files * (end_percentage - start_percentage))
self.write_log(f"Processing file {count} of {n_files}", response,
percentage)
dataset = dataset[variables] if variables else dataset
return subset_gridpoint(dataset,
lon=lon,
lat=lat,
start_yr=y0,
end_yr=y1)
def test_time_simple(self):
da = xr.open_dataset(self.nc_poslons).tas
da = da.assign_coords(lon=(da.lon - 360))
lon = -72.4
lat = 46.1
yr_st = "2050"
yr_ed = "2059"
out = subset.subset_gridpoint(
da, lon=lon, lat=lat, start_date=yr_st, end_date=yr_ed
)
np.testing.assert_almost_equal(out.lon, lon, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
np.testing.assert_array_equal(len(np.unique(out.time.dt.year)), 10)
np.testing.assert_array_equal(out.time.dt.year.max(), int(yr_ed))
np.testing.assert_array_equal(out.time.dt.year.min(), int(yr_st))
def test_raise(self):
da = xr.open_dataset(self.nc_poslons).tas
with pytest.raises(ValueError):
subset.subset_gridpoint(
da, lon=-72.4, lat=46.1, start_date="2055-03-15", end_date="2055-03-14"
)
subset.subset_gridpoint(
da, lon=-72.4, lat=46.1, start_date="2055", end_date="2052"
)
da = xr.open_dataset(self.nc_2dlonlat).tasmax.drop_vars(names=["lon", "lat"])
with pytest.raises(Exception):
subset.subset_gridpoint(da, lon=-72.4, lat=46.1)
def test_raise(self):
da = xr.open_dataset(self.nc_2dlonlat).tasmax.drop(["lon", "lat"])
with pytest.raises(Exception):
subset.subset_gridpoint(da, lon=-72.4, lat=46.1)
def test_irregular(self):
da = xr.open_dataset(self.nc_2dlonlat).tasmax
lon = -72.4
lat = 46.1
out = subset.subset_gridpoint(da, lon=lon, lat=lat)
np.testing.assert_almost_equal(out.lon, lon, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
# dask for lon lat
da.lon.chunk({"rlon": 10})
da.lat.chunk({"rlon": 10})
out = subset.subset_gridpoint(da, lon=lon, lat=lat)
np.testing.assert_almost_equal(out.lon, lon, 1)
np.testing.assert_almost_equal(out.lat, lat, 1)
# test_irregular transposed:
da1 = xr.open_dataset(self.nc_2dlonlat).tasmax
dims = list(da1.dims)
dims.reverse()
daT = xr.DataArray(np.transpose(da1.values), dims=dims)
for d in daT.dims:
args = dict()
args[d] = da1[d]
daT = daT.assign_coords(**args)
daT = daT.assign_coords(lon=(["rlon", "rlat"],
np.transpose(da1.lon.values)))
daT = daT.assign_coords(lat=(["rlon", "rlat"],
np.transpose(da1.lat.values)))
out1 = subset.subset_gridpoint(daT, lon=lon, lat=lat)
np.testing.assert_almost_equal(out1.lon, lon, 1)
np.testing.assert_almost_equal(out1.lat, lat, 1)
np.testing.assert_array_equal(out, out1)
# Dataset with tasmax, lon and lat as data variables (i.e. lon, lat not coords of tasmax)
daT1 = xr.DataArray(np.transpose(da1.values), dims=dims)
for d in daT1.dims:
args = dict()
args[d] = da1[d]
daT1 = daT1.assign_coords(**args)
dsT = xr.Dataset(data_vars=None, coords=daT1.coords)
dsT["tasmax"] = daT1
dsT["lon"] = xr.DataArray(np.transpose(da1.lon.values),
dims=["rlon", "rlat"])
dsT["lat"] = xr.DataArray(np.transpose(da1.lat.values),
dims=["rlon", "rlat"])
out2 = subset.subset_gridpoint(dsT, lon=lon, lat=lat)
np.testing.assert_almost_equal(out2.lon, lon, 1)
np.testing.assert_almost_equal(out2.lat, lat, 1)
np.testing.assert_array_equal(out, out2.tasmax)
# Dataset with lon and lat as 1D arrays
lon = -60
lat = -45
da = xr.DataArray(
np.random.rand(5, 4),
dims=("time", "site"),
coords={
"time": np.arange(5),
"site": np.arange(4)
},
)
ds = xr.Dataset(
data_vars={
"da": da,
"lon": ("site", np.linspace(lon, lon + 10, 4)),
"lat": ("site", np.linspace(lat, lat + 5, 4)),
})
gp = subset.subset_gridpoint(ds, lon=lon, lat=lat)
np.testing.assert_almost_equal(gp.lon, lon)
np.testing.assert_almost_equal(gp.lat, lat)
assert gp.site == 0
