FIX: hexbin to return grid info

astromancer · astromancer · commit b179b29e2844 · 2024-06-13T13:55:16.000+02:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -37,7 +37,6 @@
 from matplotlib.axes._secondary_axes import SecondaryAxis
 from matplotlib.container import BarContainer, ErrorbarContainer, StemContainer
 from matplotlib.hexbin import hexbin
-
 _log = logging.getLogger(__name__)
 
 
@@ -5137,12 +5136,22 @@ def reduce_C_function(C: array) -> float
         self._process_unit_info([("x", x), ("y", y)], kwargs, convert=False)
 
         x, y, C = cbook.delete_masked_points(x, y, C)
-
-        offsets, accum, (sx, sy) = hexbin(x, y, C, gridsize,
-                                          xscale, yscale,  extent,
-                                          reduce_C_function, mincnt)
-
-        polygon = [sx, sy] * np.array(
+        
+        # Count the number of data in each hexagon
+        x = np.asarray(x, float)
+        y = np.asarray(y, float)
+        
+        (*offsets, accum), (xmin, xmax), (ymin, ymax), (nx, ny) = hexbin(
+            x, y, C, gridsize, xscale, yscale, extent, reduce_C_function, mincnt
+        )
+        offsets = np.transpose(offsets)
+        sx = (xmax - xmin) / nx
+        sy = (ymax - ymin) / ny
+        
+        if C is None:
+            C = np.ones(len(x))
+    
+        polygon = [sx, sy / 3] * np.array(
             [[.5, -.5], [.5, .5], [0., 1.], [-.5, .5], [-.5, -.5], [0., -1.]])
 
         if linewidths is None:
@@ -8658,4 +8667,4 @@ def _get_aspect_ratio(self):
         figure_size = self.get_figure().get_size_inches()
         ll, ur = self.get_position() * figure_size
         width, height = ur - ll
-        return height / (width * self.get_data_ratio())
+        return height / (width * self.get_data_ratio())
diff --git a/lib/matplotlib/hexbin.py b/lib/matplotlib/hexbin.py
@@ -19,10 +19,6 @@ def hexbin(x, y, C=None, gridsize=100,
         nx = gridsize
         ny = int(nx / math.sqrt(3))
 
-    # Count the number of data in each hexagon
-    x = np.asarray(x, float)
-    y = np.asarray(y, float)
-
     # Will be log()'d if necessary, and then rescaled.
     tx = x
     ty = y
@@ -83,7 +79,7 @@ def hexbin(x, y, C=None, gridsize=100,
         accum = np.concatenate([counts1, counts2]).astype(float)
         if mincnt is not None:
             accum[accum < mincnt] = np.nan
-        C = np.ones(len(x))
+
     else:
         # store the C values in a list per hexagon index
         Cs_at_i1 = [[] for _ in range(1 + nx1 * ny1)]
@@ -96,9 +92,9 @@ def hexbin(x, y, C=None, gridsize=100,
         if mincnt is None:
             mincnt = 0
         accum = np.array(
-            [reduce_C_function(acc) if len(acc) > mincnt else np.nan
-             for Cs_at_i in [Cs_at_i1, Cs_at_i2]
-             for acc in Cs_at_i[1:]],  # [1:] drops out-of-range points.
+            [reduce_C_function(acc) if len(acc) >= mincnt else np.nan
+                for Cs_at_i in [Cs_at_i1, Cs_at_i2]
+                for acc in Cs_at_i[1:]],  # [1:] drops out-of-range points.
             float)
 
     good_idxs = ~np.isnan(accum)
@@ -116,4 +112,4 @@ def hexbin(x, y, C=None, gridsize=100,
     offsets = offsets[good_idxs, :]
     accum = accum[good_idxs]
 
-    return (*offsets.T, accum), (sx, sy / 3)
+    return (*offsets.T, accum), (xmin, xmax), (ymin, ymax), (nx, ny)
diff --git a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py
@@ -59,8 +59,9 @@ def get_gaussian_hexs(mu=(0, 0),
                       seed=123):
     np.random.seed(seed)
     rv = multivariate_normal(mu, np.array(sigma))
-    xyz, dxy = hexbin(*rv.rvs(n).T, gridsize=res)
-    return *xyz, np.array(dxy) * 0.9
+    xyz, (xmin, xmax), (ymin, ymax), (nx, ny) = hexbin(*rv.rvs(n).T, gridsize=res)
+    dxy = np.array([(xmax - xmin) / nx, (ymax - ymin) / ny]) * 0.9
+    return *xyz, dxy
 
 
 bar3d_data_generators = {
