Fix hexbin marginals

dstansby · dstansby · commit 7c800f96d494 · 2021-09-10T14:03:47.000+01:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -4646,6 +4646,11 @@ def reduce_C_function(C: array) -> float
         # Count the number of data in each hexagon
         x = np.array(x, float)
         y = np.array(y, float)
+
+        if marginals:
+            xorig = x.copy()
+            yorig = y.copy()
+
         if xscale == 'log':
             if np.any(x <= 0.0):
                 raise ValueError("x contains non-positive values, so can not"
@@ -4674,10 +4679,6 @@ def reduce_C_function(C: array) -> float
         sx = (xmax - xmin) / nx
         sy = (ymax - ymin) / ny
 
-        if marginals:
-            xorig = x.copy()
-            yorig = y.copy()
-
         x = (x - xmin) / sx
         y = (y - ymin) / sy
         ix1 = np.round(x).astype(int)
@@ -4833,40 +4834,42 @@ def reduce_C_function(C: array) -> float
         if not marginals:
             return collection
 
+        # Process marginals
         if C is None:
             C = np.ones(len(x))
 
-        def coarse_bin(x, y, coarse):
-            ind = coarse.searchsorted(x).clip(0, len(coarse) - 1)
-            mus = np.zeros(len(coarse))
-            for i in range(len(coarse)):
-                yi = y[ind == i]
+        def coarse_bin(x, y, bin_edges):
+            """
+            Sort x-values into bins defined by *bin_edges*, then for all the
+            corresponding y-values in each bin use *reduce_c_function* to
+            compute the bin value.
+            """
+            nbins = len(bin_edges) - 1
+            # Sort x-values into bins
+            bin_idxs = np.searchsorted(bin_edges, x) - 1
+            mus = np.zeros(nbins) * np.nan
+            for i in range(nbins):
+                # Get y-values for each bin
+                yi = y[bin_idxs == i]
                 if len(yi) > 0:
-                    mu = reduce_C_function(yi)
-                else:
-                    mu = np.nan
-                mus[i] = mu
+                    mus[i] = reduce_C_function(yi)
             return mus
 
-        coarse = np.linspace(xmin, xmax, gridsize)
+        if xscale == 'log':
+            bin_edges = np.geomspace(xmin, xmax, nx + 1)
+        else:
+            bin_edges = np.linspace(xmin, xmax, nx + 1)
+        xcoarse = coarse_bin(xorig, C, bin_edges)
 
-        xcoarse = coarse_bin(xorig, C, coarse)
-        valid = ~np.isnan(xcoarse)
         verts, values = [], []
-        for i, val in enumerate(xcoarse):
-            thismin = coarse[i]
-            if i < len(coarse) - 1:
-                thismax = coarse[i + 1]
-            else:
-                thismax = thismin + np.diff(coarse)[-1]
-
-            if not valid[i]:
+        for bin_left, bin_right, val in zip(
+                bin_edges[:-1], bin_edges[1:], xcoarse):
+            if np.isnan(val):
                 continue
-
-            verts.append([(thismin, 0),
-                          (thismin, 0.05),
-                          (thismax, 0.05),
-                          (thismax, 0)])
+            verts.append([(bin_left, 0),
+                          (bin_left, 0.05),
+                          (bin_right, 0.05),
+                          (bin_right, 0)])
             values.append(val)
 
         values = np.array(values)
@@ -4881,20 +4884,21 @@ def coarse_bin(x, y, coarse):
         hbar.update(kwargs)
         self.add_collection(hbar, autolim=False)
 
-        coarse = np.linspace(ymin, ymax, gridsize)
-        ycoarse = coarse_bin(yorig, C, coarse)
-        valid = ~np.isnan(ycoarse)
+        if yscale == 'log':
+            bin_edges = np.geomspace(ymin, ymax, 2 * ny + 1)
+        else:
+            bin_edges = np.linspace(ymin, ymax, 2 * ny + 1)
+        ycoarse = coarse_bin(yorig, C, bin_edges)
+
         verts, values = [], []
-        for i, val in enumerate(ycoarse):
-            thismin = coarse[i]
-            if i < len(coarse) - 1:
-                thismax = coarse[i + 1]
-            else:
-                thismax = thismin + np.diff(coarse)[-1]
-            if not valid[i]:
+        for bin_bottom, bin_top, val in zip(
+                bin_edges[:-1], bin_edges[1:], ycoarse):
+            if np.isnan(val):
                 continue
-            verts.append([(0, thismin), (0.0, thismax),
-                          (0.05, thismax), (0.05, thismin)])
+            verts.append([(0, bin_bottom),
+                          (0, bin_top),
+                          (0.05, bin_top),
+                          (0.05, bin_bottom)])
             values.append(val)
 
         values = np.array(values)
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/hexbin_log.png b/lib/matplotlib/tests/baseline_images/test_axes/hexbin_log.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -781,7 +781,8 @@ def test_hexbin_log():
     y = np.power(2, y * 0.5)
 
     fig, ax = plt.subplots()
-    h = ax.hexbin(x, y, yscale='log', bins='log')
+    h = ax.hexbin(x, y, yscale='log', bins='log',
+                  marginals=True, reduce_C_function=np.sum)
     plt.colorbar(h)
 
 
