How to sample from a step function in python/scipy/numpy etc?

"If I have bin edges and counts for each bin, is there a nice succinct way to sample from the probability density function this implies?"

This is exactly the case for scipy.stats.rv_histogram.

Here's an example.

First, generate some histogram data that we can use for the example.

In [150]: sample = np.random.gamma(6, size=2000)                                                                                                                

In [151]: hist, edges = np.histogram(sample, bins=5)                                                                                                            

In [152]: hist                                                                                                                                                  
Out[152]: array([490, 949, 438, 100,  23])

In [153]: edges                                                                                                                                                 
Out[153]: 
array([ 1.23006474,  4.19769156,  7.16531838, 10.13294519, 13.10057201,
       16.06819883])

Create an instance of rv_histogram with that data.

In [154]: from scipy.stats import rv_histogram                                                                                                                  

In [155]: rv = rv_histogram((hist, edges))                                                                                                                      

Generate a random sample from rv, and plot its histogram.

In [156]: rv_sample = rv.rvs(size=100000)                                                                                                                       

In [157]: plt.hist(rv_sample, bins=50, alpha=0.5, edgecolor='k')

You have already given the answer:

If you fall in a bin/step the value you sample is uniformly sampled from the x-values of the step. So if you fall in the first step it is uniform between 0 and 2.1.

You generate the sample in a two-step fashion:

1. choose one of the bins based on their probability mass function
2. generate a sample uniformly distributed over the range of values of the selected bin

The function below does exactly this

def random_sample(bin_edges, counts):
    n_bins = len(counts) # number of bins
    p_bin = counts/np.sum(counts) # probability mass function (pmf) for the bin
    sample_bin = np.random.choice(np.arange(n_bins), p = p_bin) # sample a bin according to the bin pmf

    # now, generate a random variable uniformly distributed within the sampled bin edges
    sample = np.random.uniform(bin_edges[sample_bin], bin_edges[sample_bin + 1]) 
    return sample

To test the function, lets generate multiple independent samples and plot their histogram

bin_edges = [0,2.1,6.3,23.5]
counts = [5,2,10]
samples = [random_sample(bin_edges, counts) for _ in range(100000)]

from seaborn import distplot

distplot(samples,bins = 100,kde=False,norm_hist = True, hist_kws=dict(edgecolor="k", linewidth=1))

As expected, the histogram corresponds to the distribution function of the random variable