Search code examples
python-3.xmatplotlibroot-framework

How to mimic the Draw('same') from ROOT with matplotlib

I have a use case from ROOT that I have not been able to reproduce with matplotlib. Here is a minimal example

import numpy as np
import matplotlib.pyplot as plt

fig, ax = plt.subplots()
dist1x = np.random.normal(5, 0.05, 10_000)  
dist1y = np.random.normal(5, 0.05, 10_000)
dist2x = np.random.normal(15, 0.05, 10_000)  
dist2y = np.random.normal(15, 0.05, 10_000)

ax.hist2d(dist1x, dist1y, bins=100, cmap='viridis')
ax.hist2d(dist2x, dist2y, bins=100, cmap='viridis')

plt.show()

and the output is

enter image description here

With ROOT one can do:

TCanvas *c1 = new TCanvas("c1","c1");

TH1D *h1 = new TH1D("h1","h1",500,-5,5);
h1->FillRandom("gaus");

TH1D *h2 = new TH1D("h2","h2",500,-5,5);
h2->FillRandom("gaus");

h1->Draw();
h2->Draw("SAME");

and the two histograms will share the canvas, axes, etc. Why plotting the two histograms in the same figure only shows the last one? How can I reproduce the ROOT behavior?

Solution

  • I think the intended behavior is to draw the sum of both histograms. You can do this by concatenating the arrays before plotting:

    ax.hist2d(np.concatenate([dist1x, dist2x]),
          np.concatenate([dist1y, dist2y]),
          bins=100, cmap='viridis')

    Matplotlib sum of histograms

    (I've modified the number a bit, to make sure the two blobs overlap.)

    The default behavior in ROOT for SAME with TH2F is probably not desirable.

    ROOT SAME

    The second histogram is drawn over the other, overwriting the fill color of the bins. The information from the first histogram is discarded in every cell if there is at least one event from the second histogram.

    To reproduce this behavior, I'd suggest to use numpy.histogram2d. Set the bins of the first histogram to zero if there are entries in the second one, and then plot the sum of both.

    bins = np.linspace(0, 20, 100), np.linspace(0, 20, 100)
hist1, _, _ = np.histogram2d(dist1x, dist1y, bins=bins)
hist2, _, _ = np.histogram2d(dist2x, dist2y, bins=bins)
hist1[hist2 > 0] = 0
sum_hist = hist1 + hist2
plt.pcolormesh(*bins, sum_hist)

    Matplotlib "SAME"

    If the two histograms don't have any populated bin in common, the two behaviors are identical.