Inverse transform sampling from a PyROOT TH2 histogram

import random, ROOT
ROOT.gROOT.SetBatch(True)
ROOT.gROOT.SetStyle("Pub")
#ROOT.gROOT.SetStyle("Plain")
#ROOT.gStyle.SetOptStat(False);


def get_sampling_vars(h):
    """
    Get the following from a TH2 histogram h, since the ROOT API sucks:
    * list of global bin IDs (not even contiguous, gee thanks ROOT)
    * dict mapping global bin IDs to a tuple of axis bin IDs
    * list of nbins+1 cumulative bin values, in the same order as globalbins
    """
    globalbin_to_axisbin = {} # for reverse axis bin lookup to get edges
    globalbins = [] # because they aren't easily predicted, nor contiguous
    cheights = [0] # cumulative "histogram" from which to uniformly sample
    for ix in xrange(1, h.GetNbinsX()+1):
        for iy in xrange(1, h.GetNbinsY()+1):
            iglobal = h.GetBin(ix, iy)
            globalbins.append(iglobal)
            globalbin_to_axisbin[iglobal] = (ix, iy)
            cheights.append(cheights[-1] + h.GetBinContent(iglobal))
    return globalbins, globalbin_to_axisbin, cheights

def get_random_bin(globalbins, cheights):
    """
    Choose a random bin from the cumulative distribution list of nbins+1 entries.

    TODO: Search more efficiently (lin and log guesses, then lin search or
    binary split depending on vector size).
    """
    assert len(cheights) == len(globalbins)+1
    randomheight = random.uniform(0, cheights[-1])
    randombin = None
    for i, iglobal in enumerate(globalbins):
        if randomheight >= cheights[i] and randomheight < cheights[i+1]:
            return iglobal
    raise Exception("Sample fell outside range of cumulative distribution?!?!")

def get_random_xy(globalbins, cheights, globalbin_to_axisbin):
    """
    Choose a random bin via get_random_bin, then pick a uniform random x,y
    point in that bin (without any attempt at estimating the in-bin distribution).
    """
    irand = get_random_bin(globalbins, cheights)
    axisids = globalbin_to_axisbin.get(irand)
    #print irand, axisids
    assert axisids is not None
    xrand = random.uniform(h.GetXaxis().GetBinLowEdge(axisids[0]), h.GetXaxis().GetBinUpEdge(axisids[0]))
    yrand = random.uniform(h.GetYaxis().GetBinLowEdge(axisids[1]), h.GetYaxis().GetBinUpEdge(axisids[1]))
    #print "{:d} -> {:f}, {:f} ".format(irand, xrand, yrand)
    return xrand, yrand


## Create and populate a histo to sample from
h = ROOT.TH2D("/h", "", 3,0,1, 3,0,1)
for ix in xrange(1, h.GetNbinsX()+1):
    for iy in xrange(1, h.GetNbinsY()+1):
        i = h.GetBin(ix, iy)
        h.SetBinContent(i, i**2)

## Get some necessary properties for sampling (ROOT's API is very unhelpful)
globalbins, globalbin_to_axisbin, cheights = get_sampling_vars(h)

## Sample into a new histo with different binning and wider range
NUMFILLS = 2000
h2 = ROOT.TH2D("/h2", "", 12,-0.1,1.1, 12,-0.1,1.1)
for i in xrange(NUMFILLS):
    if i % 500 == 0:
        print "Fill #{:d}/{:d}".format(i, NUMFILLS)
    xrand, yrand = get_random_xy(globalbins, cheights, globalbin_to_axisbin)
    h2.Fill(xrand, yrand)

## Plot sampling and sampled histos overlayed on one canvas
c = ROOT.TCanvas()
h2.Draw("col")
sf = h2.GetBinContent(h2.GetMaximumBin()) / h.GetBinContent(h.GetMaximumBin())
h.Scale(sf, "")
h.SetLineColor(ROOT.kBlack)
h.Draw("box,same")
c.SaveAs("sampled.pdf")
#c.SaveAs("sampled.png")