Repository: alteryx/evalml (python)

Issue:
Add class imbalance data check to automl
**Background**
In #971 we [implemented a data check for detecting class imbalance](https://github.com/alteryx/evalml/pull/1135/files), but we didn't add it to `DefaultDataChecks`.

**Proposal**
* Add `ClassImbalanceDataCheck` to `DefaultDataChecks`
* It currently has a warning for <10% representation. We should also have that data check include an error if the number of instances of any class is less than `2 * n_cv_folds`, because if not we'll get modeling errors. Please read the discussion in #971.
* Let's make sure our test coverage for this includes both binary and multiclass cases, categorical/int dtypes.

**Testing**
A great dataset to test with is [the abalone dataset](https://www.kaggle.com/rodolfomendes/abalone-dataset), with target "Rings", because that target has a bunch of classes which are heavily imbalanced

```python
df['Rings'].value_counts()
```
outputs
```python
9     689
10    634
8     568
11    487
7     391
12    267
6     259
13    203
14    126
5     115
15    103
16     67
17     58
4      57
18     42
19     32
20     26
3      15
21     14
23      9
22      6
24      2
27      2
1       1
25      1
2       1
26      1
29      1
Name: Rings, dtype: int64
```

Pull request description:
Add ClassImbalanceDataCheck to DefaultDataChecks
fix #1276 

- Throw error when the number of instances of any label is less than `2 * num_cv_folds`
(We default to `num_cv_folds == 3` since that is what our data split techniques use, although we leave it as an input parameter for the user)
- Add ClassImbalanceDataCheck to DefaultDataChecks

** This datacheck is only used in DefaultDataChecks if the problem type is classification, not regression

Patch:
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
index bf0488e8..7ba11b48 100644
--- a/docs/source/release_notes.rst
+++ b/docs/source/release_notes.rst
@@ -12,6 +12,7 @@ Release Notes
         * Added ``HighVarianceCVDataCheck`` and replaced synonymous warning in ``AutoMLSearch`` :pr:`1254`
         * Added ``PCA Transformer`` component for dimensionality reduction :pr:`1270`
         * Updated ``AutoMLSearch`` to support ``Woodwork`` data structures :pr:`1299`
+        * Added cv_folds to ``ClassImbalanceDataCheck`` and added this check to ``DefaultDataChecks`` :pr:`1333`
         * Make ``max_batches`` argument to ``AutoMLSearch.search`` public :pr:`1320`
     * Fixes
         * Fixed ML performance issue with ordered datasets: always shuffle data in automl's default CV splits :pr:`1265`
diff --git a/evalml/data_checks/class_imbalance_data_check.py b/evalml/data_checks/class_imbalance_data_check.py
index 0b4cc15f..95f82735 100644
--- a/evalml/data_checks/class_imbalance_data_check.py
+++ b/evalml/data_checks/class_imbalance_data_check.py
@@ -1,23 +1,28 @@
 import pandas as pd
 
 from .data_check import DataCheck
-from .data_check_message import DataCheckWarning
+from .data_check_message import DataCheckError, DataCheckWarning
 
 
 class ClassImbalanceDataCheck(DataCheck):
     """Checks if any target labels are imbalanced beyond a threshold. Use for classification problems"""
 
-    def __init__(self, threshold=0.1):
-        """Check if any of the features are likely to be ID columns.
+    def __init__(self, threshold=0.1, num_cv_folds=3):
+        """Check if any of the target labels are imbalanced, or if the number of values for each target
+           are below 2 times the number of cv folds
 
         Arguments:
             threshold (float): The minimum threshold allowed for class imbalance before a warning is raised.
                 A perfectly balanced dataset would have a threshold of (1/n_classes), ie 0.50 for binary classes.
                 Defaults to 0.10
+            num_cv_folds (int): The number of cross-validation folds. Must be positive. Choose 0 to ignore this warning.
         """
         if threshold <= 0 or threshold > 0.5:
             raise ValueError("Provided threshold {} is not within the range (0, 0.5]".format(threshold))
         self.threshold = threshold
+        if num_cv_folds < 0:
+            raise ValueError("Provided number of CV folds {} is less than 0".format(num_cv_folds))
+        self.cv_folds = num_cv_folds * 2
 
     def validate(self, X, y):
         """Checks if any target labels are imbalanced beyond a threshold for binary and multiclass problems
@@ -28,18 +33,26 @@ class ClassImbalanceDataCheck(DataCheck):
             y: Target labels to check for imbalanced data.
 
         Returns:
-            list (DataCheckWarning): list with DataCheckWarnings if imbalance in classes is less than the threshold.
+            list (DataCheckWarning, DataCheckError): list with DataCheckWarnings if imbalance in classes is less than the threshold,
+                                                     and DataCheckErrors if the number of values for each target is below 2 * num_cv_folds.
 
         Example:
             >>> X = pd.DataFrame({})
             >>> y = pd.Series([0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
             >>> target_check = ClassImbalanceDataCheck(threshold=0.10)
-            >>> assert target_check.validate(X, y) == [DataCheckWarning("The following labels fall below 10% of the target: [0]", "ClassImbalanceDataCheck")]
+            >>> assert target_check.validate(X, y) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 6 instances: [0]", "ClassImbalanceDataCheck"), DataCheckWarning("The following labels fall below 10% of the target: [0]", "ClassImbalanceDataCheck")]
         """
         if not isinstance(y, pd.Series):
             y = pd.Series(y)
         messages = []
-        counts = y.value_counts(normalize=True)
+        fold_counts = y.value_counts(normalize=False)
+        # search for targets that occur less than twice the number of cv folds first
+        below_threshold_folds = fold_counts.where(fold_counts < self.cv_folds).dropna()
+        if len(below_threshold_folds):
+            warning_msg = "The number of instances of these targets is less than 2 * the number of cross folds = {} instances: {}"
+            messages.append(DataCheckError(warning_msg.format(self.cv_folds, below_threshold_folds.index.tolist()), self.name))
+
+        counts = fold_counts / fold_counts.sum()
         below_threshold = counts.where(counts < self.threshold).dropna()
         # if there are items that occur less than the threshold, add them to the list of messages
         if len(below_threshold):
diff --git a/evalml/data_checks/default_data_checks.py b/evalml/data_checks/default_data_checks.py
index ee4ee275..e7acd177 100644
--- a/evalml/data_checks/default_data_checks.py
+++ b/evalml/data_checks/default_data_checks.py
@@ -1,3 +1,4 @@
+from .class_imbalance_data_check import ClassImbalanceDataCheck
 from .data_checks import DataChecks
 from .highly_null_data_check import HighlyNullDataCheck
 from .id_columns_data_check import IDColumnsDataCheck
@@ -5,10 +6,12 @@ from .invalid_targets_data_check import InvalidTargetDataCheck
 from .no_variance_data_check import NoVarianceDataCheck
 from .target_leakage_data_check import TargetLeakageDataCheck
 
+from evalml.problem_types import ProblemTypes, handle_problem_types
+
 
 class DefaultDataChecks(DataChecks):
     """A collection of basic data checks that is used by AutoML by default.
-    Includes HighlyNullDataCheck, IDColumnsDataCheck, TargetLeakageDataCheck, InvalidTargetDataCheck,
+    Includes HighlyNullDataCheck, IDColumnsDataCheck, TargetLeakageDataCheck, InvalidTargetDataCheck, ClassImbalanceDataCheck,
     and NoVarianceDataCheck."""
 
     _DEFAULT_DATA_CHECK_CLASSES = [HighlyNullDataCheck, IDColumnsDataCheck,
@@ -20,5 +23,9 @@ class DefaultDataChecks(DataChecks):
         Arguments:
             problem_type (str): The problem type that is being validated. Can be regression, binary, or multiclass.
         """
-        super().__init__(self._DEFAULT_DATA_CHECK_CLASSES,
-                         data_check_params={"InvalidTargetDataCheck": {"problem_type": problem_type}})
+        if handle_problem_types(problem_type) == ProblemTypes.REGRESSION:
+            super().__init__(self._DEFAULT_DATA_CHECK_CLASSES,
+                             data_check_params={"InvalidTargetDataCheck": {"problem_type": problem_type}})
+        else:
+            super().__init__(self._DEFAULT_DATA_CHECK_CLASSES + [ClassImbalanceDataCheck],
+                             data_check_params={"InvalidTargetDataCheck": {"problem_type": problem_type}})


Test patch:
diff --git a/evalml/tests/automl_tests/test_automl.py b/evalml/tests/automl_tests/test_automl.py
index 840643b6..a54dd94d 100644
--- a/evalml/tests/automl_tests/test_automl.py
+++ b/evalml/tests/automl_tests/test_automl.py
@@ -1376,16 +1376,11 @@ def test_data_split_binary(X_y_binary):
     X, y = X_y_binary
     y[:] = 0
     y[0] = 1
-
-    automl = AutoMLSearch(problem_type='binary')
-    with pytest.raises(Exception, match="Missing target values in the"):
-        automl.search(X, y)
+    automl = AutoMLSearch(problem_type="binary")
     with pytest.raises(Exception, match="Missing target values in the"):
         automl.search(X, y, data_checks="disabled")
 
     y[1] = 1
-    with pytest.raises(Exception, match="Missing target values in the"):
-        automl.search(X, y)
     with pytest.raises(Exception, match="Missing target values in the"):
         automl.search(X, y, data_checks="disabled")
 
@@ -1399,8 +1394,6 @@ def test_data_split_multi(X_y_multi):
     y[0] = 0
 
     automl = AutoMLSearch(problem_type='multiclass')
-    with pytest.raises(Exception, match="Missing target values"):
-        automl.search(X, y)
     with pytest.raises(Exception, match="Missing target values"):
         automl.search(X, y, data_checks="disabled")
 
@@ -1408,8 +1401,6 @@ def test_data_split_multi(X_y_multi):
     # match based on regex, since data split doesn't have a random seed for reproducibility
     # regex matches the set {} and expects either 2 sets (missing in both train and test)
     #   or 1 set of multiple elements (both missing in train or both in test)
-    with pytest.raises(Exception, match=r"(\{\d?\}.+\{\d?\})|(\{.+\,.+\})"):
-        automl.search(X, y)
     with pytest.raises(Exception, match=r"(\{\d?\}.+\{\d?\})|(\{.+\,.+\})"):
         automl.search(X, y, data_checks="disabled")
 
diff --git a/evalml/tests/data_checks_tests/test_class_imbalance_data_check.py b/evalml/tests/data_checks_tests/test_class_imbalance_data_check.py
index 05ca8950..fd3aa023 100644
--- a/evalml/tests/data_checks_tests/test_class_imbalance_data_check.py
+++ b/evalml/tests/data_checks_tests/test_class_imbalance_data_check.py
@@ -2,13 +2,13 @@ import numpy as np
 import pandas as pd
 import pytest
 
-from evalml.data_checks import DataCheckWarning
+from evalml.data_checks import DataCheckError, DataCheckWarning
 from evalml.data_checks.class_imbalance_data_check import (
     ClassImbalanceDataCheck
 )
 
 
-def test_class_imbalance_invalid_threshold():
+def test_class_imbalance_errors():
     X = pd.DataFrame()
 
     with pytest.raises(ValueError, match="threshold 0 is not within the range"):
@@ -20,38 +20,71 @@ def test_class_imbalance_invalid_threshold():
     with pytest.raises(ValueError, match="threshold -0.5 is not within the range"):
         ClassImbalanceDataCheck(threshold=-0.5).validate(X, y=pd.Series([0, 1, 1]))
 
+    with pytest.raises(ValueError, match="Provided number of CV folds"):
+        ClassImbalanceDataCheck(num_cv_folds=-1).validate(X, y=pd.Series([0, 1, 1]))
+
 
 def test_class_imbalance_data_check_binary():
     X = pd.DataFrame()
-    class_imbalance_check = ClassImbalanceDataCheck()
-
+    class_imbalance_check = ClassImbalanceDataCheck(num_cv_folds=0)
     assert class_imbalance_check.validate(X, y=[0, 1, 1]) == []
     assert class_imbalance_check.validate(X, y=pd.Series([0, 1, 1])) == []
     assert class_imbalance_check.validate(X, y=pd.Series([0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])) == [DataCheckWarning("The following labels fall below 10% of the target: [0]", "ClassImbalanceDataCheck")]
-    assert ClassImbalanceDataCheck(threshold=0.25).validate(X, y=pd.Series([0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])) == [DataCheckWarning("The following labels fall below 25% of the target: [0]", "ClassImbalanceDataCheck")]
+    assert ClassImbalanceDataCheck(threshold=0.25, num_cv_folds=0).validate(X, y=pd.Series([0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])) == [DataCheckWarning("The following labels fall below 25% of the target: [0]", "ClassImbalanceDataCheck")]
+
+    class_imbalance_check = ClassImbalanceDataCheck(num_cv_folds=1)
+    assert class_imbalance_check.validate(X, y=pd.Series([0, 1, 1])) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 2 instances: [0]", "ClassImbalanceDataCheck")]
+    assert class_imbalance_check.validate(X, y=pd.Series([0, 0, 1, 1])) == []
+
+    class_imbalance_check = ClassImbalanceDataCheck()
+    assert class_imbalance_check.validate(X, y=pd.Series([0, 1, 1])) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 6 instances: [1, 0]", "ClassImbalanceDataCheck")]
+    assert class_imbalance_check.validate(X, y=pd.Series(['No', 'No', 'Yes'])) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 6 instances: ['No', 'Yes']", "ClassImbalanceDataCheck")]
 
 
 def test_class_imbalance_data_check_multiclass():
     X = pd.DataFrame()
-    class_imbalance_check = ClassImbalanceDataCheck()
+    class_imbalance_check = ClassImbalanceDataCheck(num_cv_folds=0)
 
     assert class_imbalance_check.validate(X, y=pd.Series([0, 2, 1, 1])) == []
     assert class_imbalance_check.validate(X, y=pd.Series([0, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])) == [DataCheckWarning("The following labels fall below 10% of the target: [0]", "ClassImbalanceDataCheck")]
-    assert ClassImbalanceDataCheck(threshold=0.25).validate(X, y=pd.Series([0, 2, 2, 2, 3, 3, 1, 1, 1, 1])) == [DataCheckWarning("The following labels fall below 25% of the target: [3, 0]", "ClassImbalanceDataCheck")]
+    assert ClassImbalanceDataCheck(threshold=0.25, num_cv_folds=0).validate(X, y=pd.Series([0, 2, 2, 2, 3, 3, 1, 1, 1, 1])) == [DataCheckWarning("The following labels fall below 25% of the target: [3, 0]", "ClassImbalanceDataCheck")]
+
+    class_imbalance_check = ClassImbalanceDataCheck(num_cv_folds=2)
+    assert class_imbalance_check.validate(X, y=pd.Series([0, 0, 1, 2, 2, 1, 1, 1])) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 4 instances: [2, 0]", "ClassImbalanceDataCheck")]
+    y = [0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4]
+    assert class_imbalance_check.validate(X, y=pd.Series(y)) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 4 instances: [1, 0]", "ClassImbalanceDataCheck")]
+
+    class_imbalance_check = ClassImbalanceDataCheck()
+    assert class_imbalance_check.validate(X, y=pd.Series(y)) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 6 instances: [3, 2, 1, 0]", "ClassImbalanceDataCheck")]
 
 
 def test_class_imbalance_empty_and_nan():
     X = pd.DataFrame()
-    class_imbalance_check = ClassImbalanceDataCheck()
+    class_imbalance_check = ClassImbalanceDataCheck(num_cv_folds=0)
+
+    assert class_imbalance_check.validate(X, y=pd.Series([])) == []
+    assert ClassImbalanceDataCheck(threshold=0.5, num_cv_folds=0).validate(X, y=pd.Series([np.nan, np.nan, np.nan, np.nan, 1, 1, 1, 1, 2])) == [DataCheckWarning("The following labels fall below 50% of the target: [2.0]", "ClassImbalanceDataCheck")]
 
+    class_imbalance_check = ClassImbalanceDataCheck(num_cv_folds=1)
     assert class_imbalance_check.validate(X, y=pd.Series([])) == []
-    assert ClassImbalanceDataCheck(threshold=0.5).validate(X, y=pd.Series([np.nan, np.nan, np.nan, np.nan, 1, 1, 1, 1, 2])) == [DataCheckWarning("The following labels fall below 50% of the target: [2.0]", "ClassImbalanceDataCheck")]
+    assert ClassImbalanceDataCheck(threshold=0.5, num_cv_folds=1).validate(X, y=pd.Series([np.nan, np.nan, np.nan, np.nan, 1, 1, 1, 1, 2])) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 2 instances: [2.0]", "ClassImbalanceDataCheck"),
+                                                                                                                                                DataCheckWarning("The following labels fall below 50% of the target: [2.0]", "ClassImbalanceDataCheck")]
 
 
 def test_class_imbalance_nonnumeric():
     X = pd.DataFrame()
-    class_imbalance_check = ClassImbalanceDataCheck(threshold=0.25)
+    class_imbalance_check = ClassImbalanceDataCheck(threshold=0.25, num_cv_folds=0)
 
     assert class_imbalance_check.validate(X, y=[True, False, False, False, False]) == [DataCheckWarning("The following labels fall below 25% of the target: [True]", "ClassImbalanceDataCheck")]
     assert class_imbalance_check.validate(X, y=["yes", "no", "yes", "yes", "yes"]) == [DataCheckWarning("The following labels fall below 25% of the target: ['no']", "ClassImbalanceDataCheck")]
-    assert ClassImbalanceDataCheck(threshold=0.35).validate(X, y=["red", "green", "red", "red", "blue", "green", "red", "blue", "green", "red"]) == [DataCheckWarning("The following labels fall below 35% of the target: ['green', 'blue']", "ClassImbalanceDataCheck")]
+    assert ClassImbalanceDataCheck(threshold=0.35, num_cv_folds=0).validate(X, y=["red", "green", "red", "red", "blue", "green", "red", "blue", "green", "red"]) == [DataCheckWarning("The following labels fall below 35% of the target: ['green', 'blue']", "ClassImbalanceDataCheck")]
+
+    class_imbalance_check = ClassImbalanceDataCheck(num_cv_folds=1)
+    assert class_imbalance_check.validate(X, y=pd.Series(["No", "Yes", "No", "Yes"])) == []
+    assert class_imbalance_check.validate(X, y=pd.Series([True, True, True, False, False])) == []
+    assert class_imbalance_check.validate(X, y=pd.Series(["No", "Maybe", "Maybe", "No", "Yes"])) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 2 instances: ['Yes']", "ClassImbalanceDataCheck")]
+    assert class_imbalance_check.validate(X, y=pd.Series(["red", "green", "red", "red", "blue", "green", "red", "blue", "green", "red"])) == []
+
+    class_imbalance_check = ClassImbalanceDataCheck()
+    assert class_imbalance_check.validate(X, y=pd.Series(["No", "Yes", "No", "Yes", "No"])) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 6 instances: ['No', 'Yes']", "ClassImbalanceDataCheck")]
+    assert class_imbalance_check.validate(X, y=pd.Series(["red", "green", "red", "red", "blue", "green", "red", "blue", "green", "red"])) == [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 6 instances: ['red', 'green', 'blue']", "ClassImbalanceDataCheck")]
diff --git a/evalml/tests/data_checks_tests/test_data_checks.py b/evalml/tests/data_checks_tests/test_data_checks.py
index e1b76a57..81adfd8c 100644
--- a/evalml/tests/data_checks_tests/test_data_checks.py
+++ b/evalml/tests/data_checks_tests/test_data_checks.py
@@ -65,17 +65,19 @@ def test_default_data_checks_classification():
     data_checks = DefaultDataChecks("binary")
 
     leakage = [DataCheckWarning("Column 'has_label_leakage' is 95.0% or more correlated with the target", "TargetLeakageDataCheck")]
+    imbalance = [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 6 instances: [1.0, 0.0]", "ClassImbalanceDataCheck")]
 
-    assert data_checks.validate(X, y) == messages[:3] + leakage + messages[3:]
+    assert data_checks.validate(X, y) == messages[:3] + leakage + messages[3:] + imbalance
 
     data_checks = DataChecks(DefaultDataChecks._DEFAULT_DATA_CHECK_CLASSES,
                              {"InvalidTargetDataCheck": {"problem_type": "binary"}})
     assert data_checks.validate(X, y) == messages[:3] + leakage + messages[3:]
 
+    imbalance = [DataCheckError("The number of instances of these targets is less than 2 * the number of cross folds = 6 instances: [0.0, 2.0, 1.0]", "ClassImbalanceDataCheck")]
     # multiclass
     y = pd.Series([0, 1, np.nan, 2, 0])
     data_checks = DefaultDataChecks("multiclass")
-    assert data_checks.validate(X, y) == messages
+    assert data_checks.validate(X, y) == messages + imbalance
 
     data_checks = DataChecks(DefaultDataChecks._DEFAULT_DATA_CHECK_CLASSES,
                              {"InvalidTargetDataCheck": {"problem_type": "multiclass"}})