def train_qplad(coarse_reference,
fine_reference,
out,
variable,
kind,
selslice=None,
iselslice=None):
"""Train Quantile-Preserving, Localized Analogs Downscaling (QPLAD) model and output to storage"""
sel_slices_d = None
isel_slices_d = None
if selslice:
sel_slices_d = {}
for s in selslice:
k, v = s.split("=")
sel_slices_d[k] = slice(*map(str, v.split(",")))
if iselslice:
isel_slices_d = {}
for s in iselslice:
k, v = s.split("=")
isel_slices_d[k] = slice(*map(int, v.split(",")))
services.train_qplad(
coarse_reference=coarse_reference,
fine_reference=fine_reference,
out=out,
variable=variable,
kind=kind,
sel_slice=sel_slices_d,
isel_slice=isel_slices_d,
)
def test_qplad_train(tmpdir, monkeypatch, kind):
"""Tests that the shape of adjustment factors matches the expected shape"""
monkeypatch.setenv(
"HDF5_USE_FILE_LOCKING", "FALSE"
) # Avoid thread lock conflicts with dask scheduler
# make test data
np.random.seed(0)
lon = [-99.83, -99.32, -99.79, -99.23]
lat = [42.25, 42.21, 42.63, 42.59]
time = xr.cftime_range(start="1994-12-17", end="2015-01-15", calendar="noleap")
data_ref = 15 + 8 * np.random.randn(len(time), 4, 4)
ref_fine = xr.Dataset(
data_vars=dict(
scen=(["time", "lat", "lon"], data_ref),
),
coords=dict(
time=time,
lon=(["lon"], lon),
lat=(["lat"], lat),
),
attrs=dict(description="Weather related data."),
)
# need to set variable units to pass xclim 0.29 check on units
ref_fine["scen"].attrs["units"] = "K"
# take the mean across space to represent coarse reference data for AFs
ds_ref_coarse = ref_fine.mean(["lat", "lon"])
# tile the fine resolution grid with the coarse resolution ref data
ref_coarse = ds_ref_coarse.broadcast_like(ref_fine)
ref_coarse["scen"].attrs["units"] = "K"
# write test data
ref_coarse_url = "memory://test_qplad_downscaling/a/ref_coarse/path.zarr"
ref_fine_url = "memory://test_qplad_downscaling/a/ref_fine/path.zarr"
train_out_url = "memory://test_qplad_downscaling/a/train_output/path.zarr"
repository.write(
ref_coarse_url,
ref_coarse.chunk({"time": -1}),
)
repository.write(ref_fine_url, ref_fine.chunk({"time": -1}))
# now train QPLAD model
train_qplad(ref_coarse_url, ref_fine_url, train_out_url, "scen", kind)
# load adjustment factors
qplad_model = repository.read(train_out_url)
af_expected_shape = (len(lon), len(lat), 365, 620)
assert qplad_model.af.shape == af_expected_shape
def test_train_qplad_isel_slice():
"""Tests that services.train_qplad subsets with isel_slice"""
lon = [-99.83, -99.32, -99.79, -99.23]
lat = [42.25, 42.21, 42.63, 42.59]
time = xr.cftime_range(start="1994-12-17", end="2015-01-15", calendar="noleap")
data_ref = 15 + 8 * np.ones((len(time), 4, 4))
ref_fine = xr.Dataset(
data_vars=dict(
scen=(["time", "lat", "lon"], data_ref),
),
coords=dict(
time=time,
lon=(["lon"], lon),
lat=(["lat"], lat),
),
attrs=dict(description="Weather related data."),
)
# need to set variable units to pass xclim 0.29 check on units
ref_fine["scen"].attrs["units"] = "K"
# take the mean across space to represent coarse reference data for AFs
ds_ref_coarse = ref_fine.mean(["lat", "lon"])
# tile the fine resolution grid with the coarse resolution ref data
ref_coarse = ds_ref_coarse.broadcast_like(ref_fine)
ref_coarse["scen"].attrs["units"] = "K"
# write test data
ref_coarse_url = "memory://train_qplad_isel_slice/a/ref_coarse/path.zarr"
ref_fine_url = "memory://train_qplad_isel_slice/a/ref_fine/path.zarr"
train_out_url = "memory://train_qplad_isel_slice/a/train_output/path.zarr"
repository.write(ref_coarse_url, ref_coarse.chunk({"time": -1}))
repository.write(ref_fine_url, ref_fine.chunk({"time": -1}))
# now train QPLAD model
train_qplad(
coarse_reference=ref_coarse_url,
fine_reference=ref_fine_url,
out=train_out_url,
variable="scen",
kind="additive",
isel_slice={"lat": slice(0, 3)},
)
qplad_model = repository.read(train_out_url)
assert qplad_model["lat"].shape == (3,)
def test_qplad_integration(kind):
"""Integration test of the QDM and QPLAD services"""
lon = [-99.83, -99.32, -99.79, -99.23]
lat = [42.25, 42.21, 42.63, 42.59]
time = xr.cftime_range(start="1994-12-17", end="2015-01-15", calendar="noleap")
data_ref = 15 + 8 * np.random.randn(len(time), 4, 4)
data_train = 15 + 8 * np.random.randn(len(time), 4, 4)
variable = "scen"
ref_fine = xr.Dataset(
data_vars=dict(
scen=(["time", "lat", "lon"], data_ref),
),
coords=dict(
time=time,
lon=(["lon"], lon),
lat=(["lat"], lat),
),
attrs=dict(description="Weather related data."),
)
ds_train = xr.Dataset(
data_vars=dict(
scen=(["time", "lat", "lon"], data_train),
),
coords=dict(
time=time,
lon=(["lon"], lon),
lat=(["lat"], lat),
),
attrs=dict(description="Weather related data."),
)
# take the mean across space to represent coarse reference data for AFs
ds_ref_coarse = ref_fine.mean(["lat", "lon"])
ds_train = ds_train.mean(["lat", "lon"])
# tile the fine resolution grid with the coarse resolution ref data
ref_coarse = ds_ref_coarse.broadcast_like(ref_fine)
ds_bc = ds_train + 3
# need to set variable units to pass xclim 0.29 check on units
ds_train["scen"].attrs["units"] = "K"
ds_bc["scen"].attrs["units"] = "K"
ref_coarse["scen"].attrs["units"] = "K"
ref_fine["scen"].attrs["units"] = "K"
ds_ref_coarse["scen"].attrs["units"] = "K"
# write test data
ref_coarse_coarse_url = (
"memory://test_qplad_downscaling/a/ref_coarse_coarse/path.zarr"
)
ref_coarse_url = "memory://test_qplad_downscaling/a/ref_coarse/path.zarr"
ref_fine_url = "memory://test_qplad_downscaling/a/ref_fine/path.zarr"
qdm_train_url = "memory://test_qplad_downscaling/a/qdm_train/path.zarr"
sim_url = "memory://test_qplad_downscaling/a/sim/path.zarr"
qdm_train_out_url = "memory://test_qplad_downscaling/a/qdm_train_out/path.zarr"
biascorrected_url = "memory://test_qplad_downscaling/a/biascorrected/path.zarr"
sim_biascorrected_key = (
"memory://test_qplad_downscaling/a/biascorrected/sim_biascorrected.zarr"
)
repository.write(ref_coarse_coarse_url, ds_ref_coarse)
repository.write(
ref_coarse_url,
ref_coarse.chunk({"time": -1, "lat": -1, "lon": -1}),
)
repository.write(
ref_fine_url,
ref_fine.chunk({"time": -1, "lat": -1, "lon": -1}),
)
repository.write(qdm_train_url, ds_train)
repository.write(sim_url, ds_bc)
# this is an integration test between QDM and QPLAD, so use QDM services
# for bias correction
target_year = 2005
train_qdm(
historical=qdm_train_url,
reference=ref_coarse_coarse_url,
out=qdm_train_out_url,
variable=variable,
kind=kind,
)
apply_qdm(
simulation=sim_url,
qdm=qdm_train_out_url,
years=[target_year],
variable=variable,
out=sim_biascorrected_key,
)
biascorrected_coarse = repository.read(sim_biascorrected_key)
# make bias corrected data on the fine resolution grid
biascorrected_fine = biascorrected_coarse.broadcast_like(
ref_fine.sel(
time=slice("{}-01-01".format(target_year), "{}-12-31".format(target_year))
)
)
repository.write(
biascorrected_url,
biascorrected_fine.chunk({"time": -1, "lat": -1, "lon": -1}),
)
# write test data
qplad_afs_url = "memory://test_qplad_downscaling/a/qplad_afs/path.zarr"
# Writes NC to local disk, so diff format here:
sim_downscaled_url = "memory://test_qplad_downscaling/a/qplad_afs/downscaled.zarr"
# now train QPLAD model
train_qplad(ref_coarse_url, ref_fine_url, qplad_afs_url, variable, kind)
# downscale
apply_qplad(biascorrected_url, qplad_afs_url, variable, sim_downscaled_url)
# check output
downscaled_ds = repository.read(sim_downscaled_url)
# check that downscaled average equals bias corrected value
bc_timestep = biascorrected_fine[variable].isel(time=100).values[0][0]
qplad_downscaled_mean = downscaled_ds[variable].isel(time=100).mean().values
np.testing.assert_almost_equal(bc_timestep, qplad_downscaled_mean)