<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 sPlot
This notebook explains sPlot and how to use hep_ml.splot.
sPlot is a way to reconstruct features of mixture components based on known properties of distributions. This method is frequently used in High Energy Physics.
If you prefer explanations without code, find them here
Step1:  Simple sPlot example
First we start with a simple (and not very useful in practice) example
Step2:  Observed distributions
Picture above shows how this distibution should look like, 
but due to inaccuracies during classification we will observe a different picture.
Let's assume that with a probability of 80% particle is classified correctly (and we are not using $p_x$ during classification).
And when we look at distribution of px for particles which were classified as electrons or positrons, we see that they were distorted. 
We lost the original shapes of distributions.
Step3:  Applying sWeights
We can think of it in the following way
Step4:  Compare
let's compare reconstructed distribution for electrons with an original one
Step5:  More complex case
In the case when we have only two 'bins' is simple and straightforward. But when there are more than two bins, the solution is not unique. There are many appropriate combinations of sWeights, which one to choose?
Step6:  Things in practice are however even more complex. We have not bins, but continuos distribution (which can be treated as many bins). 
Typically this is a distribution over mass. By fitting mass we are able to split mixture into two parts
Step7:  Of course we don't have labels which events are signal and which are background beforehand
And we observe the mixture of two distributions
Step8:  We have no information about real labels
But we know a priori that background is distributed as exponential distribution and signal - as gaussian (more complex models can be met in practice, but idea is the same).
After fitting the mixture (let me skip this process), we will get the following result
Step9:  Fitting doesn't give us information about real labels
But it gives an information about probabilities, thus now we can estimate number of signal and background events within each bin. 
We won't use bins, but instead we will get for each event probability that it is signal or background
Step10:  Appying sPlot
sPlot converts probabilities to sWeights, using an implementation from hep_ml
Step11:  As you can see, there are also negative sWeights, which are needed to compensate the contributions of other class.
Step12:  Using sWeights to reconstruct initial distribution
Let's check that we achieved our goal and can reconstruct momentum distribution for signal and background
Step13:  An important requirement of sPlot
Reconstructed variable (i.e. $p$) and splotted variable (i.e. mass) shall be statistically independent for each class.
Read the line above again. Reconstructed and splotted variable are correlated
Step14:  But within each class there is no correlation, so the requirement is satisfied
Step15:  as a demonstration why this is important let's use sweights to reconstruct mass (obviously mass is correlated with mass)
<ASSISTANT_TASK:>
Python Code:
import numpy
from matplotlib import pyplot as plt
size = 10000
sig_data = numpy.random.normal(-1, 1, size=size)
bck_data = numpy.random.normal(1, 1, size=size)
Explanation: sPlot
This notebook explains sPlot and how to use hep_ml.splot.
sPlot is a way to reconstruct features of mixture components based on known properties of distributions. This method is frequently used in High Energy Physics.
If you prefer explanations without code, find them here
End of explanation
plt.hist(sig_data, color='b', alpha=0.5, bins=30, label='electron')
plt.hist(bck_data, color='r', alpha=0.5, bins=30, label='positron')
plt.xlim(-5, 5), plt.xlabel('$p_x$')
plt.legend()
Explanation: Simple sPlot example
First we start with a simple (and not very useful in practice) example:
Assume we have two types of particles (say, electrons and positrons). 
Distribution of some characteristic is different for them (let this be $p_x$ momentum projection).
End of explanation
n_sig1, n_bck1 = 8000, 2000
n_sig2, n_bck2 = 2000, 8000
first_bin = numpy.concatenate([sig_data[:n_sig1], bck_data[:n_bck1]])
second_bin = numpy.concatenate([sig_data[n_sig1:], bck_data[n_bck1:]])
plt.figure(figsize=[15, 6])
plt.subplot(121)
plt.bar([0, 2], [n_sig1, n_sig2], width=1, color='b', alpha=0.5)
plt.bar([0, 2], [n_bck1, n_bck2], width=1, bottom=[n_sig1, n_sig2], color='r', alpha=0.5)
plt.xlim(-0.5, 3.5)
plt.axis('off')
plt.xticks([0.5, 2.5], ['as electrons', 'as positrons'])
plt.text(0.5, -300, 'as electron', horizontalalignment='center', verticalalignment='top', fontsize=20)
plt.text(2.5, -300, 'as positron', horizontalalignment='center', verticalalignment='top', fontsize=20)
plt.title('Proportion of events being classified as')
plt.subplot(122)
plt.hist(first_bin, alpha=0.5, bins=30, label='as electrons', color=(0.22, 0., 0.66))
plt.hist(second_bin, alpha=0.5, bins=30, label='as positrons', color=(0.66, 0., 0.22))
plt.legend()
plt.title('Distributions')
plt.xlim(-5, 5), plt.xlabel('$p_x$')
pass
Explanation: Observed distributions
Picture above shows how this distibution should look like, 
but due to inaccuracies during classification we will observe a different picture.
Let's assume that with a probability of 80% particle is classified correctly (and we are not using $p_x$ during classification).
And when we look at distribution of px for particles which were classified as electrons or positrons, we see that they were distorted. 
We lost the original shapes of distributions.
End of explanation
def plot_with_weights(datas, weights, **kargs):
    assert len(datas) == len(weights)
    data = numpy.concatenate(datas)
    weight = numpy.concatenate([numpy.ones(len(d)) * w for d, w in zip(datas, weights) ])
    plt.hist(data, weights=weight, alpha=0.5, bins=30, **kargs)
plot_with_weights([first_bin, second_bin], [n_bck2, -n_bck1], density=True, label='reconstructed electron')
plot_with_weights([first_bin, second_bin], [-n_sig2, n_sig1], density=True, color='r', label='reconstructed positron')
plt.xlabel('px')
plt.legend();
Explanation: Applying sWeights
We can think of it in the following way: there are 2 bins. In first 80% are electrons, 20% are signal. And visa versa in second bin.
To reconstruct initial distribution, we can plot histogram, where each event from first bin has weight 0.8, 
and each event from second bin has weight -0.2. This numbers are called sWeights. 
So, if we had 8000 $e^{-}$ + 2000 $e^{+}$ in first bin and 8000 $e^{+}$ + 2000 $e^{-}$ ($ e^-, e^+$ are electron and positron). After summing with introduced sWeights:
$$
\big[ 8000 e^{-} + 2000 e^{+} \big] \times 0.8 + \big[ 2000 e^{-} + 8000 e^{+} \big] \times (- 0.2) = 
6800 e^{-}
$$
Positrons with positive and negative weights compensated each other, and we will get pure electrons.
At this moment we ignore normalization of sWeights (because it doesn't play role when we want to reconstruct shape).
End of explanation
plot_with_weights([first_bin, second_bin], [n_bck2, -n_bck1], density=True, label='reconstructed electons', edgecolor='none')
plot_with_weights([sig_data], [1], density=True, label='original electons', edgecolor='none')
plt.legend();
Explanation: Compare
let's compare reconstructed distribution for electrons with an original one:
End of explanation
plt.bar([0, 2, 4], [3, 2, 1], width=1, color='b', alpha=0.5)
plt.bar([0, 2, 4], [1, 2, 3], width=1, bottom=[3, 2, 1], color='r', alpha=0.5)
plt.xlim(-1, 6)
plt.ylim(-0.5, 5)
plt.axis('off')
plt.text(0.5, -0.5, 'Bin 1', horizontalalignment='center', verticalalignment='top', fontsize=20)
plt.text(2.5, -0.5, 'Bin 2', horizontalalignment='center', verticalalignment='top', fontsize=20)
plt.text(4.5, -0.5, 'Bin 3', horizontalalignment='center', verticalalignment='top', fontsize=20)
pass
Explanation: More complex case
In the case when we have only two 'bins' is simple and straightforward. But when there are more than two bins, the solution is not unique. There are many appropriate combinations of sWeights, which one to choose?
End of explanation
from scipy.stats import norm, expon
plt.figure(figsize=[15, 6])
size = 10000
sig_mass_distr = norm(loc=4, scale=1)
bck_mass_distr = expon(scale=4)
sig_mass = sig_mass_distr.rvs(size=size)
bck_mass = bck_mass_distr.rvs(size=size)
sig_p = numpy.random.normal(5, 1, size=size)
bck_p = numpy.random.normal(3, 1, size=size)
plt.subplot(121)
plt.hist(sig_mass, bins=20, density=True)
plt.hist(bck_mass, bins=20, density=True, range=(0, 10), alpha=0.5)
plt.xlabel('mass')
plt.subplot(122)
plt.hist(sig_p, bins=20, density=True)
plt.hist(bck_p, bins=20, density=True, range=(0, 10), alpha=0.5)
plt.xlabel('p');
Explanation: Things in practice are however even more complex. We have not bins, but continuos distribution (which can be treated as many bins). 
Typically this is a distribution over mass. By fitting mass we are able to split mixture into two parts: signal channel and everything else.
Building sPlot over mass
Let's show how this works. First we generate two fake distributions (signal and background) with 2 variables: mass and momentum $p$.
End of explanation
plt.figure(figsize=[15, 6])
mass = numpy.concatenate([sig_mass, bck_mass])
p = numpy.concatenate([sig_p, bck_p])
sorter = numpy.argsort(mass)
mass = mass[sorter]
p = p[sorter]
plt.subplot(121)
plt.hist(mass, bins=20, range=(0, 10))
plt.xlabel('mass')
plt.subplot(122)
plt.hist(p, bins=20)
plt.xlabel('p')
Explanation: Of course we don't have labels which events are signal and which are background beforehand
And we observe the mixture of two distributions:
End of explanation
x = numpy.linspace(0, 10)
plt.hist(mass, bins=30, range=[0, 10], density=True, alpha=0.4)
plt.plot(x, norm.pdf(x, loc=4, scale=1) / 2., label='signal')
plt.plot(x, expon.pdf(x, scale=4) / 2., label='bck')
plt.plot(x, 0.5 * (norm.pdf(x, loc=4, scale=1) + expon.pdf(x, scale=4)), label='sig + bck')
plt.legend(fontsize=20)
Explanation: We have no information about real labels
But we know a priori that background is distributed as exponential distribution and signal - as gaussian (more complex models can be met in practice, but idea is the same).
After fitting the mixture (let me skip this process), we will get the following result:
End of explanation
import pandas
probs = pandas.DataFrame(dict(sig=sig_mass_distr.pdf(mass), bck=bck_mass_distr.pdf(mass)))
probs = probs.div(probs.sum(axis=1), axis=0)
plt.plot(mass, probs.sig, label='sig probability')
plt.plot(mass, probs.bck, label='bck probability')
plt.xlim(0, 10), plt.legend(), plt.xlabel('mass')
Explanation: Fitting doesn't give us information about real labels
But it gives an information about probabilities, thus now we can estimate number of signal and background events within each bin. 
We won't use bins, but instead we will get for each event probability that it is signal or background:
End of explanation
from hep_ml import splot
sWeights = splot.compute_sweights(probs)
Explanation: Appying sPlot
sPlot converts probabilities to sWeights, using an implementation from hep_ml:
End of explanation
plt.plot(mass, sWeights.sig, label='sig sWeight')
plt.plot(mass, sWeights.bck, label='bck sWeight')
plt.xlim(0, 10), plt.legend(), plt.xlabel('mass');
Explanation: As you can see, there are also negative sWeights, which are needed to compensate the contributions of other class.
End of explanation
plt.figure(figsize=[15, 7])
plt.subplot(121)
hist_conf = dict(bins=30, alpha=0.4, range=[0, 10])
plt.hist(sig_p, label='original sig p', **hist_conf)
plt.hist(p, weights=sWeights.sig, label='reconstructed sig p', **hist_conf)
plt.xlabel('p'), plt.legend()
plt.subplot(122)
plt.hist(bck_p, label='original bck p', **hist_conf)
plt.hist(p, weights=sWeights.bck, label='reconstructed bck p', **hist_conf)
plt.xlabel('p'), plt.legend()
pass
Explanation: Using sWeights to reconstruct initial distribution
Let's check that we achieved our goal and can reconstruct momentum distribution for signal and background:
End of explanation
numpy.corrcoef(abs(mass - 4), p) [0, 1]
Explanation: An important requirement of sPlot
Reconstructed variable (i.e. $p$) and splotted variable (i.e. mass) shall be statistically independent for each class.
Read the line above again. Reconstructed and splotted variable are correlated:
End of explanation
print(numpy.corrcoef(abs(sig_mass - 4), sig_p)[0, 1])
print(numpy.corrcoef(abs(bck_mass - 4), bck_p)[0, 1])
Explanation: But within each class there is no correlation, so the requirement is satisfied:
End of explanation
hist_conf = dict(bins=30, alpha=0.5, range=[-1, 7])
plt.hist(sig_mass, label='original sig mass', **hist_conf)
plt.hist(mass, weights=sWeights.sig, label='reconstructed sig mass', **hist_conf)
plt.xlabel('mass'), plt.legend()
pass
Explanation: as a demonstration why this is important let's use sweights to reconstruct mass (obviously mass is correlated with mass):
End of explanation
<END_TASK>