BUG: Fixes axes limit for PredictionError plot #1193 (#1208)

Co-authored-by: Pkaf <[email protected]>
DistrictDataLabs · Feb 18, 2022 · 20acdcc · 20acdcc
1 parent 838063e
commit 20acdcc
Show file tree

Hide file tree

Showing 7 changed files with 6 additions and 5 deletions.
diff --git a/...s/baseline_images/test_regressor/test_prediction_error/test_peplot_no_lines.png b/...s/baseline_images/test_regressor/test_prediction_error/test_peplot_no_lines.png
diff --git a/...ne_images/test_regressor/test_prediction_error/test_peplot_no_shared_limits.png b/...ne_images/test_regressor/test_prediction_error/test_peplot_no_shared_limits.png
diff --git a/.../baseline_images/test_regressor/test_prediction_error/test_prediction_error.png b/.../baseline_images/test_regressor/test_prediction_error/test_prediction_error.png
diff --git a/...ine_images/test_regressor/test_prediction_error/test_prediction_error_numpy.png b/...ine_images/test_regressor/test_prediction_error/test_prediction_error_numpy.png
diff --git a/...ne_images/test_regressor/test_prediction_error/test_prediction_error_pandas.png b/...ne_images/test_regressor/test_prediction_error/test_prediction_error_pandas.png
diff --git a/...ges/test_regressor/test_prediction_error/test_prediction_error_quick_method.png b/...ges/test_regressor/test_prediction_error/test_prediction_error_quick_method.png
diff --git a/yellowbrick/regressor/prediction_error.py b/yellowbrick/regressor/prediction_error.py
@@ -216,12 +216,13 @@ def draw(self, y, y_pred):
                 label="best fit",
             )
 
-        # Set the axes limits based on the range of X and Y data
+        # Set the axes limits based on the overall max/min values of
+        # concatenated X and Y data
         # NOTE: shared_limits will be accounted for in finalize()
-        # TODO: do better than add one for really small residuals
-        self.ax.set_xlim(y.min() - 1, y.max() + 1)
-        self.ax.set_ylim(y_pred.min() - 1, y_pred.max() + 1)
-
+        if self.shared_limits is True:
+            self.ax.set_xlim(min(min(y), min(y_pred)), max(max(y), max(y_pred)))
+            self.ax.set_ylim(self.ax.get_xlim())
+            
         return self.ax
 
     def finalize(self, **kwargs):