SciPy's binned_statistic_2d returning 3d statistic

I am using SciPy's binned_statistic_2d() function to determine the binned averages of data in an XArray Dataset. My usage is shown below:

LTSbins = list(cloudcorrLTSw.LTS.values)[::5]
OMEGAbins = list(cloudcorrLTSw.OMEGA_700.values)[::5]
teststat = binned_statistic_2d(cloudcorrLTSw.LTS, cloudcorrLTSw.OMEGA_700, cloudcorrLTSw.CLD_RHO,
                               statistic=np.nanmean,bins=[LTSbins, OMEGAbins])

where cloudcorrLTSw.LTS, cloudcorrLTSw.OMEGA_700, and cloudcorrLTSw.CLD_RHO are 138 entry vectors, the former two being dimensions of the xarray dataset and the latter being a variable in the dataset on LTS-OMEGA_700 space. The LTS, OMEGA_700, and CLD_RHO data are themselves from a timeseries developed from spatial data, and have been transformed into their current state, so there are, as a result, as many CLD_RHO values as there are of the coordinate values, meaning the data is sparse with many NaNs per row and column. This is the reasoning behind using np.nanmean as the statistic instead of the inbuilt mean function.

Per the SciPy documentation page for this function, it is intended to return (among other things) an ndarray of shape (nx, ny), where x and y are determined by the bins kwarg in the binned_statistic_2d function call. What I'm getting, however, is a 138x27x27 (27x27 because the described binning results in 28 values), with the first dimension being completely full of NaNs. As a result, I must pass teststat.statistic through np.nanmean again to remove that surplus dimension, an operation that doesn't take much time but makes me nervous about interfering with the data. It's probably wise of me to interpolate the sparse data so that the plot doesn't look like this,

but that's another question in and of itself. So, before I can solve that, is this the intended output of binned_statistic_2d()?

Updated image to follow some3128's suggestion:

Updated image following some3128's answer using statistic = 'count':

The data I was using, I had 'forced' to be in LTS-omega space by creating an xarray object using LTS and omega as the coordinates and CLD_RHO as the variable on those axes. By using the original data (pre- making a new dataset), I was able to get this image, which seems much more reasonable. I think I can count this as solved.

It looks like you might be facing an issue either in the way you're utilizing the function or due to the characteristics of your input data. One possible culprit could be related to how you've defined the bins using the LTSbins and OMEGAbins lists. It's worth confirming that these lists accurately represent the bin edges you intend to work with.

Furthermore, it's a good idea to verify that the arrays cloudcorrLTSw.LTS, cloudcorrLTSw.OMEGA_700, and cloudcorrLTSw.CLD_RHO possess the correct shapes and are appropriately aligned with the binning strategy you're attempting to implement. If any of these arrays have unexpected shapes or contain NaN values, they could potentially be causing the unexpected behavior you're seeing in the binned_statistic_2d function. Issues like these can sometimes lead to outcomes that don't quite match your initial expectations.

I synthesised some test data (with 50% NaNs in rho) and ran it through the binning function. Results and code below.

Data:

Output from binning function visualised using a heatmap:

When I binned the data, I set the statistic to "count". This provides a sanity check on the number of observations in each bracket. The binning function also returns the edge values of the bins, which I've overlaid onto the plot. The colouring reflects the density of the raw data on the scatter plot, so it looks like it's making sense.

import numpy as np
import matplotlib.pyplot as plt
np.random.seed(0)

#Mock data
n_pts = 138
lts1 = np.random.uniform(-0.4, 1, n_pts // 2)
lts2 = np.random.randn(n_pts // 2) * 0.1 - 0.1
lts = np.empty(n_pts)
lts[::2] = lts1
lts[1::2] = lts2

hpa = np.random.uniform(-0.001, 0.0015, n_pts)
#rho has 50% NaNs
rho = np.random.randn(n_pts) + np.where(np.random.uniform(size=n_pts) > 0.5, np.nan, 0)

#View the data
plt.scatter(lts, hpa, c=rho, label='rho')
plt.ylabel('hpa')
plt.xlabel('lts')
plt.legend()
plt.show()
import seaborn as sns
import scipy

#2D binning
#hpa is the "first dimension" (row index), which the function refers to as "x"
#  on the plots, this is actually the y axis
statistic, hpa_edge, lts_edge, binnumber = scipy.stats.binned_statistic_2d(
    hpa, lts, rho, statistic='count', bins=9
)

#Plot
sns.heatmap(statistic, annot=True,
            xticklabels=np.round(lts_edge, 1),
            yticklabels=np.round(hpa_edge * 1e3, 1),
            cmap='plasma')
plt.gca().set_ylabel('hpa * 1e3 bin')
plt.gca().set_xlabel('lts bin')
plt.gca().invert_yaxis()

I think it'd be worth seeing if you can run your data through the steps above.