prevent nan values for zero count bins

This commit is contained in:
ignis 2021-03-11 05:11:48 +09:00
parent 2a3fad7b55
commit 0692237c7e

View file

@ -42,13 +42,15 @@ def cmean_binning (c, v, nbins=100, boundary=0.1, quiet=False):
'''
Compute conditional mean of v given condition c by data binning
'''
eps = np.finfo(np.float).eps
c_bins = np.hstack(
(np.linspace(0, boundary, nbins),
np.linspace(boundary, 1-boundary, nbins)[1:],
np.linspace(1-boundary, 1, nbins)[1:],
)
)
c_bins[0] = -np.finfo(np.float).eps
c_bins[0] = -eps # to include c = 0 samples to the first bin
cstar = np.zeros(c_bins.size + 1)
cstar[1:-1] = (c_bins[1:]+c_bins[:-1])/2.
@ -75,6 +77,8 @@ def cmean_binning (c, v, nbins=100, boundary=0.1, quiet=False):
v_plane = v[tidx]
bin_indices = np.digitize(cplane, c_bins, right=True)
# assign bin indices. bin intervals are right inclusive
# i-th bin: c_bin[i-1] < c <= c_bin[i]
for i in range(1, c_bins.size):
binned = v_plane[bin_indices == i]
@ -82,7 +86,9 @@ def cmean_binning (c, v, nbins=100, boundary=0.1, quiet=False):
v_counts[i] += binned.size
v_avg = np.zeros(v_sums.size+1)
v_avg[:-1] = v_sums / v_counts
v_avg[:-1] = v_sums / (v_counts + eps) # add eps to prevent nan values from divide by zero
# Boundary treatment
return cstar, v_avg, c_bins, v_counts