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)}, )