Annotate metrics (dask#630)

* Add annotations for metrics module

Author: prasunanand

ci/code_checks.sh

@@ -13,4 +13,8 @@ MSG='Checking isort... ' ; echo $MSG
 isort --recursive --check-only .
 RET=$(($RET + $?)) ; echo $MSG "DONE"
 
+MSG='Checking mypy... ' ; echo $MSG
+mypy dask_ml/metrics
+RET=$(($RET + $?)) ; echo $MSG "DONE"
+
 exit $RET

ci/environment-3.6.yaml

@@ -12,6 +12,7 @@ dependencies:
   - isort
   - msgpack-python ==0.6.2
   - multipledispatch
+  - mypy
   - numba
   - numpy ==1.17.3
   - numpydoc

ci/environment-3.7.yaml

@@ -12,6 +12,7 @@ dependencies:
   - flake8
   - isort
   - multipledispatch >=0.4.9
+  - mypy
   - numba
   - numpy >=1.16.3
   - numpydoc

ci/environment-docs.yaml

@@ -12,6 +12,7 @@ dependencies:
   - ipython
   - isort
   - multipledispatch
+  - mypy
   - nbsphinx
   - nomkl
   - nose

ci/windows.yaml

@@ -42,4 +42,7 @@ jobs:
       echo "[codecov]"
       codecov
 
+      echo "[mypy]"
+      mypy dask_ml/metrics
+
     displayName: "Lint"

dask_ml/_typing.py

@@ -0,0 +1,10 @@
+from typing import TypeVar
+
+import numpy as np
+from dask.array import Array
+from pandas import Index, Series
+
+# array-like
+
+AnyArrayLike = TypeVar("AnyArrayLike", Index, Series, Array, np.ndarray)
+ArrayLike = TypeVar("ArrayLike", Array, np.ndarray)

dask_ml/metrics/classification.py

@@ -1,11 +1,21 @@
+from typing import Optional
+
 import dask
 import dask.array as da
 import numpy as np
 import sklearn.metrics
 import sklearn.utils.multiclass
 
+from .._typing import ArrayLike
+
 
-def accuracy_score(y_true, y_pred, normalize=True, sample_weight=None, compute=True):
+def accuracy_score(
+    y_true: ArrayLike,
+    y_pred: ArrayLike,
+    normalize: bool = True,
+    sample_weight: Optional[ArrayLike] = None,
+    compute: bool = True,
+) -> ArrayLike:
     """Accuracy classification score.
 
     In multilabel classification, this function computes subset accuracy:
@@ -84,7 +94,9 @@ def accuracy_score(y_true, y_pred, normalize=True, sample_weight=None, compute=T
     return score
 
 
-def _log_loss_inner(x, y, sample_weight, **kwargs):
+def _log_loss_inner(
+    x: ArrayLike, y: ArrayLike, sample_weight: Optional[ArrayLike], **kwargs
+):
     # da.map_blocks wasn't able to concatenate together the results
     # when we reduce down to a scalar per block. So we make an
     # array with 1 element.
@@ -110,7 +122,7 @@ def log_loss(
 
     if y_pred.ndim > 1 and y_true.ndim == 1:
         y_true = y_true.reshape(-1, 1)
-        drop_axis = 1
+        drop_axis: Optional[int] = 1
         if sample_weight is not None:
             sample_weight = sample_weight.reshape(-1, 1)
     else:

dask_ml/metrics/pairwise.py

@@ -2,6 +2,7 @@
 Daskified versions of sklearn.metrics.pairwise
 """
 import warnings
+from typing import Any, Callable, Dict, Optional, Tuple, Union
 
 import dask.array as da
 import numpy as np
@@ -10,11 +11,17 @@
 from sklearn import metrics
 from sklearn.metrics.pairwise import KERNEL_PARAMS
 
+from .._typing import ArrayLike
 from ..utils import row_norms
 
 
 def pairwise_distances_argmin_min(
-    X, Y, axis=1, metric="euclidean", batch_size=None, metric_kwargs=None
+    X: ArrayLike,
+    Y: ArrayLike,
+    axis: int = 1,
+    metric: Union[str, Callable[[ArrayLike, ArrayLike], float]] = "euclidean",
+    batch_size: Optional[int] = None,
+    metric_kwargs: Optional[Dict[str, Any]] = None,
 ):
     if batch_size is not None:
         msg = "'batch_size' is deprecated. Use sklearn.config_context instead.'"
@@ -39,7 +46,13 @@ def pairwise_distances_argmin_min(
     return argmins, mins
 
 
-def pairwise_distances(X, Y, metric="euclidean", n_jobs=None, **kwargs):
+def pairwise_distances(
+    X: ArrayLike,
+    Y: ArrayLike,
+    metric: Union[str, Callable[[ArrayLike, ArrayLike], float]] = "euclidean",
+    n_jobs: Optional[int] = None,
+    **kwargs: Any
+):
     if isinstance(Y, da.Array):
         raise TypeError("`Y` must be a numpy array")
     chunks = (X.chunks[0], (len(Y),))
@@ -54,8 +67,12 @@ def pairwise_distances(X, Y, metric="euclidean", n_jobs=None, **kwargs):
 
 
 def euclidean_distances(
-    X, Y=None, Y_norm_squared=None, squared=False, X_norm_squared=None
-):
+    X: ArrayLike,
+    Y: Optional[ArrayLike] = None,
+    Y_norm_squared: Optional[ArrayLike] = None,
+    squared: bool = False,
+    X_norm_squared: Optional[ArrayLike] = None,
+) -> ArrayLike:
     if Y is None:
         Y = X
 
@@ -87,7 +104,9 @@ def euclidean_distances(
     return distances if squared else da.sqrt(distances)
 
 
-def check_pairwise_arrays(X, Y, precomputed=False):
+def check_pairwise_arrays(
+    X: ArrayLike, Y: ArrayLike, precomputed: bool = False
+) -> Tuple[ArrayLike, ArrayLike]:
     # XXX
     if Y is None:
         Y = X
@@ -113,13 +132,15 @@ def check_pairwise_arrays(X, Y, precomputed=False):
 
 
 @derived_from(metrics.pairwise)
-def linear_kernel(X, Y=None):
+def linear_kernel(X: ArrayLike, Y: Optional[ArrayLike] = None) -> ArrayLike:
     X, Y = check_pairwise_arrays(X, Y)
     return da.dot(X, Y.T)
 
 
 @derived_from(metrics.pairwise)
-def rbf_kernel(X, Y=None, gamma=None):
+def rbf_kernel(
+    X: ArrayLike, Y: Optional[ArrayLike] = None, gamma: Optional[float] = None
+) -> ArrayLike:
     X, Y = check_pairwise_arrays(X, Y)
     if gamma is None:
         gamma = 1.0 / X.shape[1]
@@ -130,7 +151,13 @@ def rbf_kernel(X, Y=None, gamma=None):
 
 
 @derived_from(metrics.pairwise)
-def polynomial_kernel(X, Y=None, degree=3, gamma=None, coef0=1):
+def polynomial_kernel(
+    X: ArrayLike,
+    Y: Optional[ArrayLike] = None,
+    degree: int = 3,
+    gamma: Optional[float] = None,
+    coef0: float = 1,
+) -> ArrayLike:
     X, Y = check_pairwise_arrays(X, Y)
     if gamma is None:
         gamma = 1.0 / X.shape[1]
@@ -140,7 +167,12 @@ def polynomial_kernel(X, Y=None, degree=3, gamma=None, coef0=1):
 
 
 @derived_from(metrics.pairwise)
-def sigmoid_kernel(X, Y=None, gamma=None, coef0=1):
+def sigmoid_kernel(
+    X: ArrayLike,
+    Y: Optional[ArrayLike] = None,
+    gamma: Optional[float] = None,
+    coef0: float = 1,
+) -> ArrayLike:
     X, Y = check_pairwise_arrays(X, Y)
     if gamma is None:
         gamma = 1.0 / X.shape[1]
@@ -165,7 +197,14 @@ def sigmoid_kernel(X, Y=None, gamma=None, coef0=1):
 }
 
 
-def pairwise_kernels(X, Y=None, metric="linear", filter_params=False, n_jobs=1, **kwds):
+def pairwise_kernels(
+    X: ArrayLike,
+    Y: Optional[ArrayLike] = None,
+    metric: Union[str, Callable[[ArrayLike, ArrayLike], float]] = "linear",
+    filter_params: bool = False,
+    n_jobs: Optional[int] = 1,
+    **kwds
+):
     from sklearn.gaussian_process.kernels import Kernel as GPKernel
 
     if metric == "precomputed":
@@ -176,6 +215,7 @@ def pairwise_kernels(X, Y=None, metric="linear", filter_params=False, n_jobs=1,
     elif metric in PAIRWISE_KERNEL_FUNCTIONS:
         if filter_params:
             kwds = dict((k, kwds[k]) for k in kwds if k in KERNEL_PARAMS[metric])
+        assert isinstance(metric, str)
         func = PAIRWISE_KERNEL_FUNCTIONS[metric]
     elif callable(metric):
         raise NotImplementedError()

dask_ml/metrics/regression.py

@@ -1,15 +1,21 @@
+from typing import Optional
+
 import dask.array as da
 import numpy as np
 import sklearn.metrics
 from dask.utils import derived_from
 
+from .._typing import ArrayLike
+
 
-def _check_sample_weight(sample_weight):
+def _check_sample_weight(sample_weight: Optional[ArrayLike]):
     if sample_weight is not None:
         raise ValueError("'sample_weight' is not supported.")
 
 
-def _check_reg_targets(y_true, y_pred, multioutput):
+def _check_reg_targets(
+    y_true: ArrayLike, y_pred: ArrayLike, multioutput: Optional[str]
+):
     if multioutput != "uniform_average":
         raise NotImplementedError("'multioutput' must be 'uniform_average'")
 
@@ -24,8 +30,12 @@ def _check_reg_targets(y_true, y_pred, multioutput):
 
 @derived_from(sklearn.metrics)
 def mean_squared_error(
-    y_true, y_pred, sample_weight=None, multioutput="uniform_average", compute=True
-):
+    y_true: ArrayLike,
+    y_pred: ArrayLike,
+    sample_weight: Optional[ArrayLike] = None,
+    multioutput: Optional[str] = "uniform_average",
+    compute: bool = True,
+) -> ArrayLike:
     _check_sample_weight(sample_weight)
     output_errors = ((y_pred - y_true) ** 2).mean(axis=0)
 
@@ -45,8 +55,12 @@ def mean_squared_error(
 
 @derived_from(sklearn.metrics)
 def mean_absolute_error(
-    y_true, y_pred, sample_weight=None, multioutput="uniform_average", compute=True
-):
+    y_true: ArrayLike,
+    y_pred: ArrayLike,
+    sample_weight: Optional[ArrayLike] = None,
+    multioutput: Optional[str] = "uniform_average",
+    compute: bool = True,
+) -> ArrayLike:
     _check_sample_weight(sample_weight)
     output_errors = abs(y_pred - y_true).mean(axis=0)
 
@@ -66,8 +80,12 @@ def mean_absolute_error(
 
 @derived_from(sklearn.metrics)
 def r2_score(
-    y_true, y_pred, sample_weight=None, multioutput="uniform_average", compute=True
-):
+    y_true: ArrayLike,
+    y_pred: ArrayLike,
+    sample_weight: Optional[ArrayLike] = None,
+    multioutput: Optional[str] = "uniform_average",
+    compute: bool = True,
+) -> ArrayLike:
     _check_sample_weight(sample_weight)
     _, y_true, y_pred, multioutput = _check_reg_targets(y_true, y_pred, multioutput)
     weight = 1.0

dask_ml/metrics/scorer.py

@@ -1,11 +1,13 @@
+from typing import Any, Callable, Tuple, Union
+
 from sklearn.metrics import check_scoring as sklearn_check_scoring, make_scorer
 
 from . import accuracy_score, log_loss, mean_squared_error, r2_score
 
 # Scorers
-accuracy_scorer = (accuracy_score, {})
+accuracy_scorer: Tuple[Any, Any] = (accuracy_score, {})
 neg_mean_squared_error_scorer = (mean_squared_error, dict(greater_is_better=False))
-r2_scorer = (r2_score, {})
+r2_scorer: Tuple[Any, Any] = (r2_score, {})
 neg_log_loss_scorer = (log_loss, dict(greater_is_better=False, needs_proba=True))
 
 
@@ -17,7 +19,7 @@
 )
 
 
-def get_scorer(scoring, compute=True):
+def get_scorer(scoring: Union[str, Callable], compute: bool = True) -> Callable:
     """Get a scorer from string
 
     Parameters

setup.cfg

@@ -25,6 +25,15 @@ skip=
 [coverage:run]
 source=dask_ml
 
+[mypy]
+ignore_missing_imports=True
+no_implicit_optional=True
+check_untyped_defs=True
+strict_equality=True
+
+[mypy-dask_ml.metrics]
+check_untyped_defs=False
+
 [tool:pytest]
 addopts = -rsx -v --durations=10
 minversion = 3.2