Merge pull request scikit-learn#5262 from andylamb/lamb-fix-class-weight-check

[MRG + 1] Fix check in `compute_class_weight`.

Author: MechCoder

sklearn/utils/class_weight.py

@@ -71,7 +71,7 @@ def compute_class_weight(class_weight, classes, y):
  " got: %r" % class_weight)
  for c in class_weight:
  i = np.searchsorted(classes, c)
- if classes[i] != c:
+ if i >= len(classes) or classes[i] != c:
  raise ValueError("Class label %d not present." % c)
  else:
  weight[i] = class_weight[c]

sklearn/utils/tests/test_class_weight.py

@@ -9,6 +9,7 @@
 from sklearn.utils.testing import assert_array_almost_equal
 from sklearn.utils.testing import assert_almost_equal
 from sklearn.utils.testing import assert_raises
+from sklearn.utils.testing import assert_raise_message
 from sklearn.utils.testing import assert_true
 from sklearn.utils.testing import assert_equal
 from sklearn.utils.testing import assert_warns
@@ -38,6 +39,28 @@ def test_compute_class_weight_not_present():
  assert_raises(ValueError, compute_class_weight, "balanced", classes, y)
 
 
+def test_compute_class_weight_dict():
+ classes = np.arange(3)
+ class_weights = {0: 1.0, 1: 2.0, 2: 3.0}
+ y = np.asarray([0, 0, 1, 2])
+ cw = compute_class_weight(class_weights, classes, y)
+
+ # When the user specifies class weights, compute_class_weights should just
+ # return them.
+ assert_array_almost_equal(np.asarray([1.0, 2.0, 3.0]), cw)
+
+ # When a class weight is specified that isn't in classes, a ValueError
+ # should get raised
+ msg = 'Class label 4 not present.'
+ class_weights = {0: 1.0, 1: 2.0, 2: 3.0, 4: 1.5}
+ assert_raise_message(ValueError, msg, compute_class_weight, class_weights,
+ classes, y)
+ msg = 'Class label -1 not present.'
+ class_weights = {-1: 5.0, 0: 1.0, 1: 2.0, 2: 3.0}
+ assert_raise_message(ValueError, msg, compute_class_weight, class_weights,
+ classes, y)
+
+
 def test_compute_class_weight_invariance():
  # Test that results with class_weight="balanced" is invariant wrt
  # class imbalance if the number of samples is identical.