def test_loaded_DenseModel_predicts_with_clipped_inputs(tmpdir):
hyperparameters = DenseHyperparameters(
["a", "b"],
["c"],
clip_config=PackerConfig({"a": {
"z": SliceConfig(None, 3)
}}),
)
model = DenseModel(["a", "b"], ["c"], hyperparameters)
nz = 5
dims = ["x", "y", "z"]
shape = (2, 2, nz)
arr = np.arange(np.prod(shape)).reshape(shape).astype(float)
input_data = xr.Dataset({
"a": (dims, arr),
"b": (dims, arr),
"c": (dims, arr + 1)
})
model.fit([input_data])
prediction = model.predict(input_data)
output_path = str(tmpdir.join("trained_model"))
fv3fit.dump(model, output_path)
model_loaded = fv3fit.load(output_path)
loaded_prediction = model_loaded.predict(input_data)
xr.testing.assert_allclose(prediction, loaded_prediction)
def test_SklearnWrapper_fit_predict_with_clipped_input_data():
nz = 5
model = _RegressorEnsemble(
base_regressor=DummyRegressor(strategy="constant",
constant=np.arange(nz)),
n_jobs=1,
)
wrapper = SklearnWrapper(
input_variables=["a", "b"],
output_variables=["c"],
model=model,
packer_config=PackerConfig({"a": {
"z": SliceConfig(2, None)
}}),
)
dims = ["x", "y", "z"]
shape = (2, 2, nz)
arr = np.arange(np.prod(shape)).reshape(shape)
input_data = xr.Dataset({
"a": (dims, arr),
"b": (dims[:-1], arr[:, :, 0]),
"c": (dims, arr + 1)
})
wrapper.fit([input_data])
wrapper.predict(input_data)
def test_Config_asdict():
sl1_kwargs = dict(start=0, stop=10, step=2)
sl2_kwargs = dict(start=None, stop=25, step=None)
sel_map = dict(dummy_in=SliceConfig(**sl1_kwargs),
dummy_out=SliceConfig(**sl2_kwargs))
original = MicrophysicsConfig(
input_variables=["dummy_in"],
direct_out_variables=["dummy_out"],
selection_map=sel_map,
)
config_d = original.asdict()
assert config_d["selection_map"]["dummy_in"] == sl1_kwargs
assert config_d["selection_map"]["dummy_out"] == sl2_kwargs
result = MicrophysicsConfig.from_dict(config_d)
assert result == original
def test_clip_differing_slices():
ds = get_dataset(["var1", "var2"],
[[SAMPLE_DIM, FEATURE_DIM], [SAMPLE_DIM, FEATURE_DIM]])
clip_config = {
"var1": {
FEATURE_DIM: SliceConfig(4, 8)
},
"var2": {
FEATURE_DIM: SliceConfig(None, 6, 2)
},
}
clipped_data = clip(ds, clip_config)
for name in ds:
expected_da = ds[name].assign_coords(
{dim: range(ds.sizes[dim])
for dim in ds[name].dims})
for dim, slice_config in clip_config[name].items():
expected_da = expected_da.isel({dim: slice_config.slice})
xr.testing.assert_identical(clipped_data[name], expected_da)
def test_clip(dataset: xr.Dataset):
name = list(dataset.data_vars)[0]
if FEATURE_DIM in dataset[name].dims:
indices = {name: {FEATURE_DIM: SliceConfig(4, 8)}}
clipped_data = clip(dataset, indices)
expected_da = dataset[name].assign_coords(
{dim: range(dataset.sizes[dim])
for dim in dataset[name].dims})
for dim, slice_config in indices[name].items():
expected_da = expected_da.isel({dim: slice_config.slice})
xr.testing.assert_identical(clipped_data[name], expected_da)
def test_sklearn_pack_unpack_with_clipping(dataset: xr.Dataset):
name = list(dataset.data_vars)[0]
if FEATURE_DIM in dataset[name].dims:
pack_config = PackerConfig({name: {FEATURE_DIM: SliceConfig(3, None)}})
packed_array, feature_index = pack(dataset, [SAMPLE_DIM], pack_config)
unpacked_dataset = unpack(packed_array, [SAMPLE_DIM], feature_index)
expected = {}
for k in dataset:
da = dataset[k].copy(deep=True)
indices = pack_config.clip.get(k, {})
for dim, slice_config in indices.items():
da = da.isel({dim: slice_config.slice})
expected[k] = da
xr.testing.assert_allclose(unpacked_dataset, xr.Dataset(expected))
def test_DenseModel_raises_not_implemented_error_with_clipped_output_data():
hyperparameters = DenseHyperparameters(
["a", "b"],
["c"],
clip_config=ClipConfig({"c": {
"z": SliceConfig(None, 3)
}}),
)
with pytest.raises(NotImplementedError):
DenseModel(
["a", "b"],
["c"],
hyperparameters,
)
def test_array_packer_dump_and_load(tmpdir):
dataset = get_dataset(["var1"], [[SAMPLE_DIM, FEATURE_DIM]])
packer_config = PackerConfig({"var1": {"z": SliceConfig(None, 2)}})
packer = ArrayPacker(SAMPLE_DIM, list(dataset.data_vars), packer_config)
packer.to_array(dataset)
with open(str(tmpdir.join("packer.yaml")), "w") as f:
packer.dump(f)
with open(str(tmpdir.join("packer.yaml"))) as f:
loaded_packer = ArrayPacker.load(f)
assert packer._pack_names == loaded_packer._pack_names
assert packer._n_features == loaded_packer._n_features
assert packer._sample_dim_name == loaded_packer._sample_dim_name
assert packer._config == packer_config
for orig, loaded in zip(packer._feature_index,
loaded_packer._feature_index):
assert orig == loaded
def test_SklearnWrapper_raises_not_implemented_error_with_clipped_output_data(
):
nz = 5
model = _RegressorEnsemble(
base_regressor=DummyRegressor(strategy="constant",
constant=np.arange(nz)),
n_jobs=1,
)
with pytest.raises(NotImplementedError):
SklearnWrapper(
input_variables=["a", "b"],
output_variables=["c"],
model=model,
packer_config=PackerConfig({"c": {
"z": SliceConfig(2, None)
}}),
)
def test_DenseModel_clipped_inputs():
hyperparameters = DenseHyperparameters(
["a", "b"],
["c"],
clip_config=PackerConfig({"a": {
"z": SliceConfig(None, 3)
}}),
)
model = DenseModel(["a", "b"], ["c"], hyperparameters)
nz = 5
dims = ["x", "y", "z"]
shape = (2, 2, nz)
arr = np.arange(np.prod(shape)).reshape(shape).astype(float)
input_data = xr.Dataset({
"a": (dims, arr),
"b": (dims, arr),
"c": (dims, arr + 1)
})
slice_filled_input = xr.Dataset({
"a":
input_data["a"].where(input_data.z < 3).fillna(1.0),
"b":
input_data["b"]
})
model.fit([input_data])
prediction_clipped = model.predict(input_data)
assert model.X_packer._n_features["a"] == 3
assert model.X_packer._n_features["b"] == 5
prediction_nan_filled = model.predict(slice_filled_input)
xr.testing.assert_allclose(prediction_nan_filled,
prediction_clipped,
rtol=1e-3)
def test_SliceConfig(start, stop, step):
expected = slice(start, stop, step)
config = SliceConfig(start=start, stop=stop, step=step)
assert config.slice == expected
def _get_config() -> TransformConfig:
return TransformConfig(
antarctic_only=False,
vertical_subselections={"a": SliceConfig(start=5)},
)