def test_no_method_raises():
clf = ParallelPostFit(LinearRegression())
X, y = make_classification(chunks=50)
clf.fit(X, y)
with pytest.raises(AttributeError) as m:
clf.predict_proba(X)
assert m.match("The wrapped estimator (.|\n)* 'predict_proba' method.")
def test_predict(kind):
X, y = make_classification(chunks=100)
if kind == "numpy":
X, y = dask.compute(X, y)
elif kind == "dask.dataframe":
X = dd.from_dask_array(X)
y = dd.from_dask_array(y)
base = LogisticRegression(random_state=0, n_jobs=1, solver="lbfgs")
wrap = ParallelPostFit(
LogisticRegression(random_state=0, n_jobs=1, solver="lbfgs"))
base.fit(*dask.compute(X, y))
wrap.fit(*dask.compute(X, y))
assert_estimator_equal(wrap.estimator, base)
result = wrap.predict(X)
expected = base.predict(X)
assert_eq_ar(result, expected)
result = wrap.predict_proba(X)
expected = base.predict_proba(X)
assert_eq_ar(result, expected)
result = wrap.predict_log_proba(X)
expected = base.predict_log_proba(X)
assert_eq_ar(result, expected)
def test_multiclass():
X, y = sklearn.datasets.make_classification(n_classes=3, n_informative=4)
X = da.from_array(X, chunks=50)
y = da.from_array(y, chunks=50)
if SK_GE_020:
kwargs = {"multi_class": "auto"}
else:
kwargs = {}
clf = ParallelPostFit(
LogisticRegression(random_state=0, n_jobs=1, solver="lbfgs", **kwargs)
)
clf.fit(X, y)
result = clf.predict(X)
expected = clf.estimator.predict(X)
assert isinstance(result, da.Array)
assert_eq_ar(result, expected)
result = clf.predict_proba(X)
expected = clf.estimator.predict_proba(X)
assert isinstance(result, da.Array)
assert_eq_ar(result, expected)
def test_multiclass():
X, y = sklearn.datasets.make_classification(n_classes=3, n_informative=4)
X = da.from_array(X, chunks=50)
y = da.from_array(y, chunks=50)
clf = ParallelPostFit(
LogisticRegression(random_state=0,
n_jobs=1,
solver="lbfgs",
multi_class="auto"))
clf.fit(*dask.compute(X, y))
result = clf.predict(X)
expected = clf.estimator.predict(X)
assert isinstance(result, da.Array)
assert_eq_ar(result, expected)
result = clf.predict_proba(X)
expected = clf.estimator.predict_proba(X)
assert isinstance(result, da.Array)
assert_eq_ar(result, expected)
result = clf.predict_log_proba(X)
expected = clf.estimator.predict_log_proba(X)
assert_eq_ar(result, expected)
def test_it_works():
clf = ParallelPostFit(GradientBoostingClassifier())
X, y = make_classification(n_samples=1000, chunks=100)
clf.fit(X, y)
assert isinstance(clf.predict(X), da.Array)
assert isinstance(clf.predict_proba(X), da.Array)
def test_auto_rechunk():
clf = ParallelPostFit(GradientBoostingClassifier())
X, y = make_classification(n_samples=1000, n_features=20, chunks=100)
X = X.rechunk({0: 100, 1: 10})
clf.fit(X, y)
assert clf.predict(X).compute().shape == (1000,)
assert clf.predict_proba(X).compute().shape == (1000, 2)
assert clf.score(X, y) == clf.score(X.compute(), y.compute())
def test_warning_on_dask_array_without_array_function():
X, y = make_classification(n_samples=10, n_features=2, chunks=10)
clf = ParallelPostFit(GradientBoostingClassifier())
clf = clf.fit(X, y)
class FakeArray:
def __init__(self, value):
self.value = value
@property
def ndim(self):
return self.value.ndim
@property
def len(self):
return self.value.len
@property
def dtype(self):
return self.value.dtype
@property
def shape(self):
return self.value.shape
ar = FakeArray(np.zeros(shape=(2, 2)))
fake_dask_ar = da.from_array(ar)
fake_dask_ar._meta = FakeArray(np.zeros(shape=(0, 0)))
with pytest.warns(
UserWarning,
match="provide explicit `predict_meta` to the dask_ml.wrapper"):
clf.predict(fake_dask_ar)
with pytest.warns(
UserWarning,
match=
"provide explicit `predict_proba_meta` to the dask_ml.wrapper",
):
clf.predict_proba(fake_dask_ar)
def test_it_works():
clf = ParallelPostFit(GradientBoostingClassifier())
X, y = make_classification(n_samples=1000, chunks=100)
X_, y_ = dask.compute(X, y)
clf.fit(X_, y_)
assert isinstance(clf.predict(X), da.Array)
assert isinstance(clf.predict_proba(X), da.Array)
result = clf.score(X, y)
expected = clf.estimator.score(X_, y_)
assert result == expected
def test_predict_proba_meta_override():
X = pd.DataFrame({"c_0": [1, 2, 3, 4]})
y = np.array([1, 2, 3, 4])
base = CategoricalNB()
base.fit(pd.DataFrame(X), y)
dd_X = dd.from_pandas(X, npartitions=2)
dd_X._meta = pd.DataFrame({"c_0": [5]})
# Failure when not proving predict_proba_meta
# because of value dependent model
wrap = ParallelPostFit(base)
with pytest.raises(ValueError):
wrap.predict_proba(dd_X)
# Success when providing meta over-ride
wrap = ParallelPostFit(base,
predict_proba_meta=np.array([[0.0, 0.1, 0.8, 0.1]]))
result = wrap.predict_proba(dd_X)
expected = base.predict_proba(X)
assert_eq_ar(result, expected)
def test_multiclass():
X, y = make_classification(chunks=50, n_classes=3, n_informative=4)
clf = ParallelPostFit(LogisticRegression(random_state=0))
clf.fit(X, y)
result = clf.predict(X)
expected = clf.estimator.predict(X)
assert isinstance(result, da.Array)
assert_eq_ar(result, expected)
result = clf.predict_proba(X)
expected = clf.estimator.predict_proba(X)
assert isinstance(result, da.Array)
assert_eq_ar(result, expected)
def test_predict(kind):
X, y = make_classification(chunks=100)
if kind == 'numpy':
X, y = dask.compute(X, y)
elif kind == 'dask.dataframe':
X = dd.from_dask_array(X)
y = dd.from_dask_array(y)
base = LogisticRegression(random_state=0)
wrap = ParallelPostFit(LogisticRegression(random_state=0))
base.fit(X, y)
wrap.fit(X, y)
assert_estimator_equal(wrap.estimator, base)
result = wrap.predict(X)
expected = base.predict(X)
assert_eq_ar(result, expected)
result = wrap.predict_proba(X)
expected = base.predict_proba(X)
assert_eq_ar(result, expected)
def predict_xr(
model,
input_xr,
chunk_size=None,
persist=True,
proba=False,
clean=False,
return_input=False,
):
"""
Using dask-ml ParallelPostfit(), runs the parallel
predict and predict_proba methods of sklearn
estimators. Useful for running predictions
on a larger-than-RAM datasets.
Last modified: September 2020
Parameters
----------
model : scikit-learn model or compatible object
Must have a .predict() method that takes numpy arrays.
input_xr : xarray.DataArray or xarray.Dataset.
Must have dimensions 'x' and 'y'
chunk_size : int
The dask chunk size to use on the flattened array. If this
is left as None, then the chunks size is inferred from the
.chunks() method on the `input_xr`
persist : bool
If True, and proba=True, then 'input_xr' data will be
loaded into distributed memory. This will ensure data
is not loaded twice for the prediction of probabilities,
but this will only work if the data is not larger than RAM.
proba : bool
If True, predict probabilities. This only applies if the
model has a .predict_proba() method
clean : bool
If True, remove Infs and NaNs from input and output arrays
return_input : bool
If True, then the data variables in the 'input_xr' dataset will
be appended to the output xarray dataset.
Returns
----------
output_xr : xarray.Dataset
An xarray.Dataset containing the prediction output from model
with input_xr as input, if proba=True then dataset will also contain
the prediciton probabilities. Has the same spatiotemporal structure
as input_xr.
"""
if chunk_size is None:
chunk_size = int(input_xr.chunks["x"][0]) * int(
input_xr.chunks["y"][0])
# convert model to dask predict
model = ParallelPostFit(model)
# with joblib.parallel_backend("dask"):
x, y, crs = input_xr.x, input_xr.y, input_xr.geobox.crs
input_data = []
for var_name in input_xr.data_vars:
input_data.append(input_xr[var_name])
input_data_flattened = []
# TODO: transfer to dask dataframe
for arr in input_data:
data = arr.data.flatten().rechunk(chunk_size)
input_data_flattened.append(data)
# reshape for prediction
input_data_flattened = da.array(input_data_flattened).transpose()
if clean:
input_data_flattened = da.where(da.isfinite(input_data_flattened),
input_data_flattened, 0)
if proba and persist:
# persisting data so we don't require loading all the data twice
input_data_flattened = input_data_flattened.persist()
# apply the classification
print(" predicting...")
out_class = model.predict(input_data_flattened)
# Mask out NaN or Inf values in results
if clean:
out_class = da.where(da.isfinite(out_class), out_class, 0)
# Reshape when writing out
out_class = out_class.reshape(len(y), len(x))
# stack back into xarray
output_xr = xr.DataArray(out_class,
coords={
"x": x,
"y": y
},
dims=["y", "x"])
output_xr = output_xr.to_dataset(name="Predictions")
if proba:
print(" probabilities...")
out_proba = model.predict_proba(input_data_flattened)
# convert to %
out_proba = da.max(out_proba, axis=1) * 100.0
if clean:
out_proba = da.where(da.isfinite(out_proba), out_proba, 0)
out_proba = out_proba.reshape(len(y), len(x))
out_proba = xr.DataArray(out_proba,
coords={
"x": x,
"y": y
},
dims=["y", "x"])
output_xr["Probabilities"] = out_proba
if return_input:
print(" input features...")
# unflatten the input_data_flattened array and append
# to the output_xr containin the predictions
arr = input_xr.to_array()
stacked = arr.stack(z=["y", "x"])
# handle multivariable output
output_px_shape = ()
if len(input_data_flattened.shape[1:]):
output_px_shape = input_data_flattened.shape[1:]
output_features = input_data_flattened.reshape(
(len(stacked.z), *output_px_shape))
# set the stacked coordinate to match the input
output_features = xr.DataArray(
output_features,
coords={
"z": stacked["z"]
},
dims=[
"z",
*[
"output_dim_" + str(idx)
for idx in range(len(output_px_shape))
],
],
).unstack()
# convert to dataset and rename arrays
output_features = output_features.to_dataset(dim="output_dim_0")
data_vars = list(input_xr.data_vars)
output_features = output_features.rename(
{i: j
for i, j in zip(output_features.data_vars, data_vars)} # noqa pylint: disable=unnecessary-comprehension
)
# merge with predictions
output_xr = xr.merge([output_xr, output_features], compat="override")
return assign_crs(output_xr, str(crs))
