def setUp(self):
self._n_events = 1e5
self._signal = spectra.Spectra("signal", self._n_events)
self._backgrounds = [
spectra.Spectra("bkg_1", self._n_events),
spectra.Spectra("bkg_2", self._n_events)
]
def test_fill(self):
""" Test the fill method.
Basically tests bin positioning makes sense.
"""
test_decays = 10
test_spectra = spectra.Spectra("Test", test_decays)
for x in range(0, test_decays):
energy = random.uniform(0, test_spectra._energy_high)
radius = random.uniform(0, test_spectra._radial_high)
time = random.uniform(0, test_spectra._time_high)
test_spectra.fill(energy, radius, time)
x_bin = energy / test_spectra._energy_high * test_spectra._energy_bins
y_bin = radius / test_spectra._radial_high * test_spectra._radial_bins
z_bin = time / test_spectra._time_high * test_spectra._time_bins
self.assertTrue(test_spectra._data[x_bin, y_bin, z_bin] > 0)
# Also test the sum method at the same time
self.assertTrue(test_spectra.sum() == test_decays)
self.assertRaises(ValueError, test_spectra.fill, -1, 0, 0)
self.assertRaises(ValueError, test_spectra.fill, 0, -1, 0)
self.assertRaises(ValueError, test_spectra.fill, 0, 0, -1)
self.assertRaises(ValueError, test_spectra.fill,
test_spectra._energy_high + 1, 0, 0)
self.assertRaises(ValueError, test_spectra.fill, 0,
test_spectra._radial_high + 1, 0)
self.assertRaises(ValueError, test_spectra.fill, 0, 0,
test_spectra._time_high + 1)
def test_slicing(self):
""" Test the slicing shirnks the spectra in the correct way.
"""
test_decays = 10
test_spectra = spectra.Spectra("Test", test_decays)
self.assertRaises(ValueError, test_spectra.shrink,
test_spectra._energy_low,
2 * test_spectra._energy_high,
test_spectra._radial_low, test_spectra._radial_high,
test_spectra._time_low, test_spectra._time_high)
self.assertRaises(ValueError, test_spectra.shrink,
test_spectra._energy_low, test_spectra._energy_high,
test_spectra._radial_low,
2 * test_spectra._radial_high,
test_spectra._time_low, test_spectra._time_high)
self.assertRaises(ValueError, test_spectra.shrink,
test_spectra._energy_low, test_spectra._energy_high,
test_spectra._radial_low, test_spectra._radial_high,
test_spectra._time_low, 2 * test_spectra._time_high)
test_spectra.shrink(test_spectra._energy_low,
test_spectra._energy_high / 2,
test_spectra._radial_low,
test_spectra._radial_high / 2,
test_spectra._time_low,
test_spectra._time_high / 2)
self.assertTrue(test_spectra._data.shape == (test_spectra._energy_bins,
test_spectra._radial_bins,
test_spectra._time_bins))
def test_project(self):
""" Test the projection method of the spectra.
This creates projected spectra alongside the actual spectra.
"""
test_decays = 10
test_spectra = spectra.Spectra("Test", test_decays)
energy_projection = numpy.ndarray(shape=(test_spectra._energy_bins),
dtype=float)
energy_projection.fill(0)
radial_projection = numpy.ndarray(shape=(test_spectra._radial_bins),
dtype=float)
radial_projection.fill(0)
time_projection = numpy.ndarray(shape=(test_spectra._time_bins),
dtype=float)
time_projection.fill(0)
for x in range(0, test_decays):
energy = random.uniform(0, 10.0)
radius = random.uniform(0, 6000.0)
time = random.uniform(0, 10.0)
test_spectra.fill(energy, radius, time)
x_bin = energy / test_spectra._energy_high * test_spectra._energy_bins
y_bin = radius / test_spectra._radial_high * test_spectra._radial_bins
z_bin = time / test_spectra._time_high * test_spectra._time_bins
energy_projection[x_bin] += 1.0
radial_projection[y_bin] += 1.0
time_projection[z_bin] += 1.0
self.assertTrue(
numpy.array_equal(energy_projection, test_spectra.project(0)))
self.assertTrue(
numpy.array_equal(radial_projection, test_spectra.project(1)))
self.assertTrue(
numpy.array_equal(time_projection, test_spectra.project(2)))
def test_rebin(self):
""" Tests that the spectra are being rebinned correctly.
"""
test_decays = 10
test_spectra = spectra.Spectra("Test", test_decays)
test_spectra._energy_bins = 1000
test_spectra._radial_bins = 1000
test_spectra._time_bins = 10
test_spectra.calc_widths()
old_energy_width = test_spectra._energy_width
old_radial_width = test_spectra._radial_width
old_time_width = test_spectra._time_width
for decay in range(test_decays):
energy = random.uniform(test_spectra._energy_low,
test_spectra._energy_high)
radius = random.uniform(test_spectra._radial_low,
test_spectra._radial_high)
time = random.uniform(test_spectra._time_low,
test_spectra._time_high)
test_spectra.fill(energy, radius, time)
new_bins = (1, 2, 3, 4)
self.assertRaises(ValueError, test_spectra.rebin, new_bins)
new_bins = (99999, 99999, 99999)
self.assertRaises(ValueError, test_spectra.rebin, new_bins)
old_sum = test_spectra.sum()
new_bins = (500, 250, 2)
test_spectra.rebin(new_bins)
self.assertTrue(old_sum == test_spectra.sum())
self.assertTrue(test_spectra._data.shape == new_bins)
self.assertTrue(test_spectra._energy_width == old_energy_width * 2.)
self.assertTrue(test_spectra._radial_width == old_radial_width * 4.)
self.assertTrue(test_spectra._time_width == old_time_width * 5.)
def test_copy(self):
""" Tests that the spectra are being copied correctly.
"""
test_decays = 10
test_spectra = spectra.Spectra("Test", test_decays)
# Modify spectra to non default values
test_spectra._energy_bins = 250
test_spectra._radial_bins = 500
test_spectra._time_bins = 2
test_spectra._energy_low = 1.
test_spectra._energy_high = 11.
test_spectra._radial_low = 100.
test_spectra._radial_high = 10100.
test_spectra._time_low = 1.
test_spectra._time_high = 11.
test_spectra._raw_events = 9.
test_spectra.calc_widths()
for decay in range(test_decays):
energy = random.uniform(test_spectra._energy_low,
test_spectra._energy_high)
radius = random.uniform(test_spectra._radial_low,
test_spectra._radial_high)
time = random.uniform(test_spectra._time_low,
test_spectra._time_high)
test_spectra.fill(energy, radius, time)
new_spectra = test_spectra.copy()
self.assertTrue(new_spectra._data.all() == test_spectra._data.all())
self.assertTrue(new_spectra._name == test_spectra._name)
self.assertTrue(new_spectra._energy_low == test_spectra._energy_low)
self.assertTrue(new_spectra._energy_high == test_spectra._energy_high)
self.assertTrue(new_spectra._energy_bins == test_spectra._energy_bins)
self.assertTrue(
new_spectra._energy_width == test_spectra._energy_width)
self.assertTrue(new_spectra._radial_low == test_spectra._radial_low)
self.assertTrue(new_spectra._radial_high == test_spectra._radial_high)
self.assertTrue(new_spectra._radial_bins == test_spectra._radial_bins)
self.assertTrue(
new_spectra._radial_width == test_spectra._radial_width)
self.assertTrue(new_spectra._time_low == test_spectra._time_low)
self.assertTrue(new_spectra._time_high == test_spectra._time_high)
self.assertTrue(new_spectra._time_bins == test_spectra._time_bins)
self.assertTrue(new_spectra._time_width == test_spectra._time_width)
self.assertTrue(new_spectra._num_decays == test_spectra._num_decays)
self.assertTrue(new_spectra._raw_events == test_spectra._raw_events)
new_spectra2 = test_spectra.copy(name="Copy")
self.assertTrue(new_spectra2._name != test_spectra._name)
self.assertTrue(new_spectra2._name == "Copy")
def random_gaussian_radius_spectra(self, true_spectrum):
""" Smears the radius of a spectra object by generating a
number of random points from a Gaussian pdf generated for
that bin. The number of points generated is equivalent to the
number of entries in that bin.
Args:
true_spectrum (spectra): spectrum to be smeared
Returns:
A smeared spectra object.
"""
raw_events = true_spectrum._raw_events
energy_step = (true_spectrum._energy_high -
true_spectrum._energy_low) / true_spectrum._energy_bins
time_step = (true_spectrum._time_high -
true_spectrum._time_low) / true_spectrum._time_bins
radial_step = (true_spectrum._radial_high -
true_spectrum._radial_low) / true_spectrum._radial_bins
smeared_spectrum = spectra.Spectra(
true_spectrum._name + str(self._position_resolution) +
"_position_resolution", true_spectrum._num_decays)
for time_bin in range(true_spectrum._time_bins):
mean_time = time_bin * time_step + 0.5 * time_step
for energy_bin in range(true_spectrum._energy_bins):
mean_energy = energy_bin * energy_step + 0.5 * energy_step
for radial_bin in range(true_spectrum._radial_bins):
mean_radius = radial_bin * radial_step + 0.5 * radial_step
entries = true_spectrum._data[energy_bin, radial_bin,
time_bin]
for i in range(int(entries)):
try:
smeared_spectrum.fill(
mean_energy,
np.fabs(
np.random.normal(
mean_radius,
self._position_resolution)), mean_time)
except ValueError:
# Occurs when smeared radius is outside bin range
print "Warning: Smeared radius out of bounds. Skipping."
continue
smeared_spectrum._raw_events = raw_events
return smeared_spectrum
def test_scale(self):
""" Test the scale method of the spectra.
This creates a spectra and then scales it.
"""
test_decays = 10
test_spectra = spectra.Spectra("Test", test_decays)
for x in range(0, test_decays):
energy = random.uniform(test_spectra._energy_low,
test_spectra._energy_high)
radius = random.uniform(test_spectra._radial_low,
test_spectra._radial_high)
time = random.uniform(test_spectra._time_low,
test_spectra._time_high)
test_spectra.fill(energy, radius, time)
self.assertTrue(test_spectra.sum() == test_decays)
count = 150
test_spectra.scale(count)
self.assertTrue(test_spectra.sum() == count)
def test_random_radius(self):
""" Tests the random gaussian smearing method for radius.
Creates a delta function and fits the mean and sigma after smearing.
Mean and sigma are checked against set values within 1 %.
Integral of smeared spectrum is checked against original number of
entries.
"""
num_entries = 100000
test_spectra = spectra.Spectra("Test", num_entries)
radius = 1000. # mm
for i in range(num_entries):
test_spectra.fill(0, radius, 0)
smearing = smear.Smear()
smearing._position_resolution = 100. # mm
smeared_spectra = smearing.random_gaussian_radius_spectra(test_spectra)
mean, sigma, integral = self.fit_gaussian_radius(smeared_spectra)
self.assertTrue(radius < 1.01*mean and radius > 0.99*mean)
self.assertTrue(smearing._position_resolution < 1.01*sigma and smearing._position_resolution > 0.99*sigma)
self.assertAlmostEqual(integral/float(num_entries), 1.0)
def test_random_energy(self):
""" Tests the random gaussian smearing method for energy.
Creates a delta function and fits the mean and sigma after smearing.
Mean and sigma are checked against set values within 1 %.
Integral of smeared spectrum is checked against original number of
entries.
"""
num_entries = 100000
test_spectra = spectra.Spectra("Test",num_entries)
energy = 2.5 # MeV
for i in range(num_entries):
test_spectra.fill(energy, 0, 0)
smearing = smear.Smear()
smearing._light_yield = 200. # NHit per MeV
smeared_spectra = smearing.random_gaussian_energy_spectra(test_spectra)
expected_sigma = numpy.sqrt(energy/200.)
mean, sigma, integral = self.fit_gaussian_energy(smeared_spectra)
self.assertTrue(energy < 1.01*mean and energy > 0.99*mean)
self.assertTrue(expected_sigma < 1.01*sigma and expected_sigma > 0.99*sigma)
self.assertAlmostEqual(integral/float(num_entries), 1.0)
def test_serialisation(self):
""" Test saving and then reloading a test spectra.
"""
test_decays = 10
test_spectra = spectra.Spectra("Test", test_decays)
for x in range(0, test_decays):
energy = random.uniform(0, test_spectra._energy_high)
radius = random.uniform(0, test_spectra._radial_high)
time = random.uniform(0, test_spectra._time_high)
test_spectra.fill(energy, radius, time)
store.dump("test.hdf5", test_spectra)
loaded_spectra = store.load("test.hdf5")
self.assertTrue(loaded_spectra.sum() == test_decays)
self.assertTrue(
numpy.array_equal(test_spectra._data, loaded_spectra._data))
self.assertTrue(test_spectra._energy_low == loaded_spectra._energy_low)
self.assertTrue(
test_spectra._energy_high == loaded_spectra._energy_high)
self.assertTrue(
test_spectra._energy_bins == loaded_spectra._energy_bins)
self.assertTrue(
test_spectra._energy_width == loaded_spectra._energy_width)
self.assertTrue(test_spectra._radial_low == loaded_spectra._radial_low)
self.assertTrue(
test_spectra._radial_high == loaded_spectra._radial_high)
self.assertTrue(
test_spectra._radial_bins == loaded_spectra._radial_bins)
self.assertTrue(
test_spectra._radial_width == loaded_spectra._radial_width)
self.assertTrue(test_spectra._time_low == loaded_spectra._time_low)
self.assertTrue(test_spectra._time_high == loaded_spectra._time_high)
self.assertTrue(test_spectra._time_bins == loaded_spectra._time_bins)
self.assertTrue(test_spectra._time_width == loaded_spectra._time_width)
self.assertTrue(test_spectra._num_decays == loaded_spectra._num_decays)
def weight_gaussian_radius_spectra(self, true_spectrum, num_sigma=5.):
""" Smears the radius of a spectra object by calculating a
Gaussian pdf for each bin and applies a weight to the bin and
corresponding bins a default 5 sigma apart.
Args:
true_spectrum (spectra): spectrum to be smeared
num_sigma (float): Width of window to apply the weight method.
Default is 5.
Returns:
A smeared spectra object.
"""
raw_events = true_spectrum._raw_events
energy_step = (true_spectrum._energy_high -
true_spectrum._energy_low) / true_spectrum._energy_bins
time_step = (true_spectrum._time_high -
true_spectrum._time_low) / true_spectrum._time_bins
radial_step = (true_spectrum._radial_high -
true_spectrum._radial_low) / true_spectrum._radial_bins
smeared_spectrum = spectra.Spectra(
true_spectrum._name + str(self._position_resolution) +
"_position_resolution", true_spectrum._num_decays)
for time_bin in range(true_spectrum._time_bins):
mean_time = time_bin * time_step + 0.5 * time_step
for energy_bin in range(true_spectrum._energy_bins):
mean_energy = energy_bin * energy_step + 0.5 * energy_step
for radial_bin in range(true_spectrum._radial_bins):
mean_radius = radial_bin * radial_step + 0.5 * radial_step
entries = float(true_spectrum._data[energy_bin, radial_bin,
time_bin])
if entries == 0:
continue # Bin Empty
lower_bin = self.floor_to_bin(
mean_radius - num_sigma * self._position_resolution,
radial_step) + 0.5 * radial_step
upper_bin = self.ceil_to_bin(
mean_radius + num_sigma * self._position_resolution,
radial_step) - 0.5 * radial_step
if upper_bin > true_spectrum._radial_high:
upper_bin = true_spectrum._radial_high - 0.5 * energy_step
if lower_bin < true_spectrum._radial_low:
lower_bin = true_spectrum._radial_low + 0.5 * energy_step
weights = []
for radius in np.arange(lower_bin, upper_bin, radial_step):
weights.append(
self.calc_gaussian(radius, mean_radius,
self._position_resolution))
weight_tot = np.array(weights).sum()
i = 0
for radius in np.arange(lower_bin, upper_bin, radial_step):
try:
smeared_spectrum.fill(
mean_energy, radius, mean_time,
entries * weights[i] / weight_tot)
except ValueError:
# Occurs when smeared radius is outside bin range
print "Warning: Smeared radius out of bounds. Skipping."
continue
i += 1
smeared_spectrum._raw_events = raw_events
return smeared_spectrum
def fill_reco_spectrum(filename, T, spectrumname="", spectrum=None):
"""**Weights have been disabled.**
This function fills in the ndarray of energies, radii, times
and weights. It takes the reconstructed energies and positions
of the events from the root file. In order to keep the statistics,
the time dependence is performed via adding weights to every event
depending on the time period. Both, studied time and Half-life must
be written in the same units.
Args:
filename (str): A root file to study
T (float): The Half-life of a studied background
spectrumname (str, optional): A name of future spectrum. Not
required when appending a spectrum.
spectrum (:class:`echidna.core.spectra.Spectra`, optional):
Spectrum you wish to append. Not required when creating a
new spectrum.
Raises:
ValueError: If spectrumname is not set when creating a new
spectrum.
Returns:
spectrum (:class:`echidna.core.spectra.Spectra`)
"""
print filename
print spectrumname
dsreader = RAT.DU.DSReader(filename)
if spectrum is None:
if spectrumname == "":
raise ValueError("Name not set when creating new spectra.")
spectrum = spectra.Spectra(str(spectrumname),
10. * dsreader.GetEntryCount())
else:
spectrum._num_decays += 10. * dsreader.GetEntryCount()
spectrumname = spectrum._name
print spectrumname
times = [0]
for time_step in range(0, spectrum._time_bins):
time = time_step * spectrum._time_width + spectrum._time_low
times.append(time)
if 'AV' in spectrumname:
print "AV WEIGHTS ARE CURRENTLY UNAVAILABLE"
weights = _av_weights(times, T)
else:
weights = _scint_weights(times, T)
for ievent in range(0, dsreader.GetEntryCount()):
ds = dsreader.GetEntry(ievent)
for ievent in range(0, ds.GetEVCount()):
ev = ds.GetEV(ievent)
if not ev.DefaultFitVertexExists() or not ev.GetDefaultFitVertex(
).ContainsEnergy() or not ev.GetDefaultFitVertex().ValidEnergy():
continue
energy = ev.GetDefaultFitVertex().GetEnergy()
position = ev.GetDefaultFitVertex().GetPosition().Mag()
spectrum._raw_events += 1
for time, weight in zip(times, weights):
try:
spectrum.fill(energy, position, time, 1.)
except ValueError:
pass
return spectrum
def fill_mc_ntuple_spectrum(filename, T, spectrumname="", spectrum=None):
"""**Weights have been disabled.**
This function fills in the ndarray of energies, radii, times
and weights. It takes the reconstructed energies and positions
of the events from ntuple. In order to keep the statistics,
the time dependence is performed via adding weights to every event
depending on the time period. Both, studied time and Half-life must
be written in the same units.
Args:
filename (str): The ntuple to study
T (float): The Half-life of a studied background
spectrumname (str, optional): A name of future spectrum. Not
required when appending a spectrum.
spectrum (:class:`echidna.core.spectra.Spectra`, optional):
Spectrum you wish to append. Not required when creating a
new spectrum.
Raises:
ValueError: If spectrumname is not set when creating a new
spectrum.
Returns:
spectrum (:class:`echidna.core.spectra.Spectra`)
"""
print filename
chain = TChain("output")
chain.Add(filename)
if spectrum is None:
if spectrumname == "":
raise ValueError("Name not set when creating new spectra.")
spectrum = spectra.Spectra(str(spectrumname), 10. * chain.GetEntries())
else:
spectrum._num_decays += 10. * chain.GetEntries()
spectrumname = spectrum._name
print spectrumname
times = []
for time_step in range(0, spectrum._time_bins):
time = time_step * spectrum._time_width + spectrum._time_low
times.append(time)
if 'AV' in spectrumname:
print "AV WEIGHTS ARE CURRENTLY UNAVAILABLE"
weights = _av_weights(times, T)
else:
weights = _scint_weights(times, T)
for event in chain:
energy = event.mcEdepQuenched
position = math.fabs(
math.sqrt((event.mcPosx)**2 + (event.mcPosy)**2 +
(event.mcPosz)**2))
spectrum._raw_events += 1
for time, weight in zip(times, weights):
try:
spectrum.fill(energy, position, time, 1.)
except ValueError:
pass
return spectrum
