diff --git a/src/pykrige/ck.py b/src/pykrige/ck.py index 4199c87..89a3b7e 100644 --- a/src/pykrige/ck.py +++ b/src/pykrige/ck.py @@ -132,7 +132,7 @@ def fit(self, p, x, y): ml_pred_ilr = ilr_transformation(ml_pred) self.onehotencode = OneHotEncoder(categories=[self.classes_]) - y_ohe = np.array(self.onehotencode.fit_transform(y).todense()) + y_ohe = np.array(self.onehotencode.fit_transform(y.reshape(-1, 1)).todense()) y_ohe_ilr = ilr_transformation(y_ohe) for i in range(len(self.classes_) - 1): diff --git a/tests/test_classification_krige.py b/tests/test_classification_krige.py index 9110a97..a41674f 100644 --- a/tests/test_classification_krige.py +++ b/tests/test_classification_krige.py @@ -62,6 +62,38 @@ def test_classification_krige(): assert class_model.score(X_test, lon_lat_test, y_test) > 0.25 +@pytest.mark.skipif(not SKLEARN_INSTALLED, reason="requires scikit-learn") +@pytest.mark.parametrize("method", ["ordinary", "universal"]) +@pytest.mark.parametrize("column_target", [False, True]) +def test_classification_krige_target_shape(method, column_target): + rng = np.random.RandomState(4) + p = rng.uniform(size=(12, 2)) + x = rng.uniform(size=(12, 2)) + y = np.arange(12) % 3 + + def model(): + return ClassificationKriging( + classification_model=RandomForestClassifier(n_estimators=8, random_state=0), + method=method, + n_closest_points=4, + variogram_parameters=[1.0, 0.0], + ) + + actual = model() + actual.fit(p, x, y[:, None] if column_target else y) + reference = model() + reference.fit(p, x, y[:, None]) + + query = np.array([[0.3, 0.7], [0.8, 0.2]]) + np.testing.assert_array_equal( + actual.predict(p[:2], query), reference.predict(p[:2], query) + ) + np.testing.assert_allclose( + actual.krige_residual(query), reference.krige_residual(query) + ) + assert actual.score(p, x, y) == 1.0 + + @pytest.mark.skipif(not SKLEARN_INSTALLED, reason="requires scikit-learn") def test_krige_classification_housing(): import ssl