ENH Speedup confusion_matrix
@@ -309,7 +309,7 @@ def confusion_matrix(y_true, y_pred, *, labels=None, sample_weight=None,
raise ValueError("'labels' should contains at least one label.")
elif y_true.size == 0:
return np.zeros((n_labels, n_labels), dtype=int)
- elif np.all([l not in y_true for l in labels]):
+ elif len(np.intersect1d(y_true, labels)) == 0:
raise ValueError("At least one label specified must be in y_true")
if sample_weight is None:
@@ -324,13 +324,21 @@ def confusion_matrix(y_true, y_pred, *, labels=None, sample_weight=None,
"'all', None}")
n_labels = labels.size
+ # If labels are not consecutive integers starting from zero, then
+ # y_true and y_pred must be converted into index form
+ need_index_conversion = not (
+ labels.dtype.kind in {'i', 'u', 'b'} and
+ np.all(labels == np.arange(n_labels)) and
+ y_true.min() >= 0 and y_pred.min() >= 0
+ )
+ if need_index_conversion:
label_to_ind = {y: x for x, y in enumerate(labels)}
- # convert yt, yp into index
y_pred = np.array([label_to_ind.get(x, n_labels + 1) for x in y_pred])
y_true = np.array([label_to_ind.get(x, n_labels + 1) for x in y_true])
# intersect y_pred, y_true with labels, eliminate items not in labels
ind = np.logical_and(y_pred < n_labels, y_true < n_labels)
+ if not np.all(ind):
y_pred = y_pred[ind]
y_true = y_true[ind]
# also eliminate weights of eliminated items
