From ceea0e1c41c4edddc93561491552cfd41ec91394 Mon Sep 17 00:00:00 2001
From: Maoz Gelbart <13831112+MaozGelbart@users.noreply.github.com>
Date: Wed, 26 Feb 2020 14:44:47 +0200
Subject: [PATCH] FIX: Ignore masked cells when finding heatmap data limits
 (#1956)

* respect mask when setting heatmap limits

* respect mask when setting heatmap limits

* improve code style

(cherry picked from commit 47570655ccc3780cc3332360b7478993190ce467)
---
 seaborn/matrix.py            | 14 +++++++++++---
 seaborn/tests/test_matrix.py | 19 +++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/seaborn/matrix.py b/seaborn/matrix.py
index a5fc7fac6a..b76e42eb4b 100644
--- a/seaborn/matrix.py
+++ b/seaborn/matrix.py
@@ -194,11 +194,19 @@ def __init__(self, data, vmin, vmax, cmap, center, robust, annot, fmt,
     def _determine_cmap_params(self, plot_data, vmin, vmax,
                                cmap, center, robust):
         """Use some heuristics to set good defaults for colorbar and range."""
-        calc_data = plot_data.data[~np.isnan(plot_data.data)]
+
+        # plot_data is a np.ma.array instance
+        calc_data = plot_data.filled(np.nan)
         if vmin is None:
-            vmin = np.percentile(calc_data, 2) if robust else calc_data.min()
+            if robust:
+                vmin = np.nanpercentile(calc_data, 2)
+            else:
+                vmin = np.nanmin(calc_data)
         if vmax is None:
-            vmax = np.percentile(calc_data, 98) if robust else calc_data.max()
+            if robust:
+                vmax = np.nanpercentile(calc_data, 98)
+            else:
+                vmax = np.nanmax(calc_data)
         self.vmin, self.vmax = vmin, vmax
 
         # Choose default colormaps if not provided
diff --git a/seaborn/tests/test_matrix.py b/seaborn/tests/test_matrix.py
index e75084991f..4b9b58a0ed 100644
--- a/seaborn/tests/test_matrix.py
+++ b/seaborn/tests/test_matrix.py
@@ -97,6 +97,25 @@ def test_mask_input(self):
 
         npt.assert_array_equal(p.plot_data, plot_data)
 
+    def test_mask_limits(self):
+        """Make sure masked cells are not used to calculate extremes"""
+
+        kws = self.default_kws.copy()
+
+        mask = self.x_norm > 0
+        kws['mask'] = mask
+        p = mat._HeatMapper(self.x_norm, **kws)
+
+        assert p.vmax == np.ma.array(self.x_norm, mask=mask).max()
+        assert p.vmin == np.ma.array(self.x_norm, mask=mask).min()
+
+        mask = self.x_norm < 0
+        kws['mask'] = mask
+        p = mat._HeatMapper(self.x_norm, **kws)
+
+        assert p.vmin == np.ma.array(self.x_norm, mask=mask).min()
+        assert p.vmax == np.ma.array(self.x_norm, mask=mask).max()
+
     def test_default_vlims(self):
 
         p = mat._HeatMapper(self.df_unif, **self.default_kws)