Commit 09c42995 authored by Niklas Uwe Langner's avatar Niklas Uwe Langner
Browse files

[cosmic_rays] Shuffle crs after simulation using crs.shuffle_events()

parent 807e5c6d
Pipeline #224855 failed with stages
in 5 minutes and 34 seconds
......@@ -575,3 +575,16 @@ class CosmicRaysSets(CosmicRaysBase):
# noinspection PyTypeChecker
crs = self.get(setid)
return plot_energy_spectrum(crs, **kwargs)
def shuffle_events(self):
"""
Independently shuffle the cosmic rays of each set.
"""
# This function can be simplified in the future using np.take_along_axis()
shuffle_ids = np.random.permutation(np.prod(self.shape)).reshape(self.shape)
shuffle_ids = np.argsort(shuffle_ids, axis=1)
sets_ids = np.repeat(np.arange(self.nsets), self.ncrs).reshape(self.shape)
for _key in self.shape_array.dtype.names:
self.__setitem__(_key, self.__getitem__(_key)[sets_ids, shuffle_ids])
sets_ids_3d = np.repeat(np.arange(3), np.prod(self.shape)).reshape((3,) + self.shape)
self.__setitem__('vecs', self.__getitem__('vecs')[sets_ids_3d, sets_ids, np.stack([shuffle_ids] * 3)])
......@@ -988,6 +988,15 @@ class TestCosmicRaysSets(unittest.TestCase):
self.assertTrue(np.shape(pvals) == (shape[0], 60))
self.assertTrue(((pvals >= 0) & (pvals <= 1)).all())
def test_32b_shuffle(self):
crs = CosmicRaysSets(self.shape)
crs["log10e"] = np.random.random(self.shape)
crs["vecs"] = coord.rand_vec(self.shape)
test1 = np.unique(crs['vecs'] * crs['log10e'])
crs.shuffle_events()
test2 = np.unique(crs['vecs'] * crs['log10e'])
self.assertTrue(np.all(test1 == test2))
if __name__ == '__main__':
unittest.main()
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment