Commit 5dc27e04 authored by Marcus Wirtz's avatar Marcus Wirtz
Browse files

[simulations] Implement parallelized simulation of multiple source densities

parent 3840e157
Pipeline #352248 passed with stages
in 6 minutes and 1 second
""" Module to setup a parametrized simulation, that work on probability distributions """
import os
import numpy as np
from scipy.stats import mode
from astrotools import auger, coord, cosmic_rays, healpytools as hpt, nucleitools as nt
PATH = os.path.split(__file__)[0]
......@@ -499,7 +500,7 @@ class SourceBound(BaseSimulation):
def set_charges(self, charges):
"""
Define fraction of charge groups in form of dictionary (e.g. {'h':0.5, 'fe':0.5}) at source
or as keyword 'first_minimum'/'second_minimum' from Auger's best fit paper (arXiv:1612.07155)
or as keyword 'first_minimum'/'second_minimum' from Auger's combined fit paper (arXiv:1612.07155)
If string is given, gamma and Rcut are also set to the respective best fit values.
:param charges: dictionary hosting the fractions of injected elements (H, He, N, Si, Fe possible)
......@@ -511,6 +512,7 @@ class SourceBound(BaseSimulation):
self.charge_weights = charges
elif isinstance(charges, str):
# (CTG, CTD) according to naming in Auger's combined fit paper (arXiv:1612.07155)
if charges == 'first_minimum_CTG':
self.energy_setting['gamma'], self.energy_setting['log10_cut'] = 1.03, 18.21
self.charge_weights = {'h': 0.6794, 'he': 0.31, 'n': 0.01, 'si': 0.0006}
......@@ -610,7 +612,8 @@ class SourceBound(BaseSimulation):
self.dis_bins, self.log10e_bins = data['distances'], data['log10e_bins']
if np.median(np.min(self.universe.distances, axis=1)) < np.min(self.dis_bins):
print("Warning: Distance binning of simulation starts too far outside (%s Mpc)! " % np.min(self.dis_bins))
print("Required would be: %.2fMpc." % np.median(np.min(self.universe.distances, axis=1)))
print("\tRequired would be: %.2fMpc." % np.median(np.min(self.universe.distances, axis=1)))
print("\tThis is only relevant if there is substantial attenuation below %s Mpc." % np.min(self.dis_bins))
# Assign distance indices for all simulated soures, shape: (self.nsets, self.universe.n_src)
dis_bin_idx = self.universe.distance_indices(self.dis_bins)
......@@ -656,8 +659,9 @@ class SourceBound(BaseSimulation):
def _get_inside_fraction(self):
""" Internal function to determine fraction of cosmic rays which come from inside rmax (signal cosmic rays) """
# as log-space binning the width of the distance bin is increasing with distance (d log(x) = 1/x * dx)
distance_fractions = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3)
inside_fraction = np.sum(distance_fractions[self.dis_bins <= self.universe.rmax]) / np.sum(distance_fractions)
dist_frac = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3)
rmax = np.atleast_1d(self.universe.rmax)
inside_fraction = np.array([np.sum(dist_frac[self.dis_bins <= r]) for r in rmax]) / np.sum(dist_frac)
return inside_fraction
def _allocate_sources(self, inside_fraction=None):
......@@ -665,17 +669,19 @@ class SourceBound(BaseSimulation):
if inside_fraction is None:
inside_fraction = self._get_inside_fraction()
self.crs['inside_fraction'] = inside_fraction
# Sample for each set the number of CRs coming from inside and outside rmax
nsplit = np.random.multinomial(self.ncrs, [inside_fraction, 1-inside_fraction], size=self.nsets).T
# Assign the CRs from inside rmax to their separate sources (by index label)
# Sample for each set the number of CRs coming from inside the horizon rmax
std = np.sqrt(self.ncrs * inside_fraction * (1 - inside_fraction)) # approximate fluctutation from
n_fluc = np.random.normal(scale=std, size=np.size(inside_fraction)) # multinomial distribution
n_inside = np.rint(np.clip(self.ncrs * inside_fraction + n_fluc, 0, self.ncrs)).astype(int)
# Assign the CRs from inside rmax to their separate sources by index label (-1 for outside rmax / no source)
source_labels = -np.ones(self.shape).astype(int)
n_max = np.max(nsplit[0]) # max events from inside over all sets
n_max = np.max(n_inside) # max events from inside over all sets
source_labels[:, :n_max] = sample_from_m_distributions(self.source_matrix.sum(axis=-1), n_max)
nrange = np.tile(np.arange(self.ncrs), self.nsets).reshape(self.shape)
mask_close = nrange < nsplit[0][:, np.newaxis] # Create mask for CRs inside rmax
source_labels[~mask_close] = -1 # correct the ones resulting by max(nsplit[0])
mask_close = nrange < n_inside[:, np.newaxis] # Create mask for CRs inside rmax
source_labels[~mask_close] = -1 # correct the ones resulting by max(n_inside)
# Checks if for all sets there is at least one explicitly simulated source WITHOUT cosmic-ray contribution
# Checks if for ALL sets there is at least one explicitly simulated source WITHOUT cosmic-ray contribution
range_src = np.arange(self.universe.n_src)
check = np.apply_along_axis(lambda x: np.in1d(range_src, x, invert=True).any(), axis=1, arr=source_labels).all()
if (not check):
......@@ -711,21 +717,26 @@ class SourceBound(BaseSimulation):
if np.sum(~mask_close) > 0:
# as log-space binning the width of the distance bin is increasing with distance (d log(x) = 1/x * dx)
arrival_mat_far = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3)
if self.universe.background_horizon is not None:
mask_out = self.dis_bins >= self.universe.background_horizon
else:
mask_out = self.dis_bins >= self.universe.rmax
arrival_mat_far[~mask_out] = 0
arrival_mat_far /= arrival_mat_far.sum()
arrival_idx = np.random.choice(arrival_mat_far.size, size=np.sum(~mask_close), p=arrival_mat_far.flatten())
idx = np.unravel_index(arrival_idx, arrival_mat_far.shape)
_d = 10**(np.log10(self.dis_bins)[idx[0]] + (np.random.random(idx[0].size) - 0.5) * d_dis)
_c, _lge = np.array(c)[idx[1]], self.log10e_bins[:-1][idx[2]]
perm = np.arange(np.sum(~mask_close))
np.random.shuffle(perm)
log10e[~mask_close] = _lge[perm]
charge[~mask_close] = _c[perm]
self.crs['distances'][~mask_close] = _d[perm]
rmax = self.universe.rmax if self.universe.background_horizon is None else self.universe.background_horizon
# determine for each set the distance bin of the horizon rmax
dis_idx = np.argmin(np.abs(np.atleast_1d(rmax)[np.newaxis] - self.dis_bins[:, np.newaxis]), axis=0)
# loop over all occuring distance bins and fill events of the respective sets
for _dis in np.sort(np.unique(dis_idx)):
arrival_mat_far[self.dis_bins < _dis] = 0 # set probabilities inside rmax to zero
mask_pick = ~mask_close & (dis_idx == _dis)[:, np.newaxis] # select background events of respective set
n_pick = np.sum(mask_pick)
arrival_mat_far /= arrival_mat_far.sum() # normalize arrival matrix to then draw random samples
arrival_idx = np.random.choice(arrival_mat_far.size, size=n_pick, p=arrival_mat_far.flatten())
# get the indices in terms of the original shape (distances, charges, energies)
idx = np.unravel_index(arrival_idx, arrival_mat_far.shape)
_d = 10**(np.log10(self.dis_bins)[idx[0]] + (np.random.random(idx[0].size) - 0.5) * d_dis)
_c, _lge = np.array(c)[idx[1]], self.log10e_bins[:-1][idx[2]]
# fill in the distances, charges, and energies with a random permutation
perm = np.arange(n_pick)
np.random.shuffle(perm)
log10e[mask_pick] = _lge[perm]
charge[mask_pick] = _c[perm]
self.crs['distances'][mask_pick] = _d[perm]
# Assign charges and energies of close-by cosmic rays
dis_bin_idx = self.universe.distance_indices(self.dis_bins)
......@@ -735,6 +746,7 @@ class SourceBound(BaseSimulation):
mask = cr_idx == i
if np.sum(mask) == 0:
continue
# procedure in analogy to code block above
w_mat = self.arrival_matrix[i] / self.arrival_matrix[i].sum()
arrival_idx = np.random.choice(w_mat.size, size=np.sum(mask), p=w_mat.flatten())
idx = np.unravel_index(arrival_idx, w_mat.shape)
......@@ -754,13 +766,14 @@ class SourceBound(BaseSimulation):
return ''.join(['__%s_%s' % (key, self.charge_weights[key]) for key in chargegroups
if key in self.charge_weights])
def _select_representative_set(self): # pragma: no cover
""" Select a representative set in terms of anisotropies """
from scipy.stats import mode
src_labels = np.copy(self.crs['source_labels']).astype(float)
def _select_representative_set(self, mask=None): # pragma: no cover
""" Select a representative set in terms of anisotropies. Returns index """
labels = self.crs['source_labels'] if mask is None else self.crs['source_labels'][mask]
src_labels = np.copy(labels).astype(float)
src_labels[src_labels < 0] = np.nan
_, counts = mode(src_labels, axis=1, nan_policy='omit')
return np.argsort(np.squeeze(counts))[int(self.nsets/2.)]
idx_select = np.argsort(np.squeeze(counts))[int(len(labels)/2.)]
return idx_select if mask is None else np.arange(self.nsets)[mask][idx_select]
def _get_strongest_sources(self, idx=None, min_number_src=5):
""" Select strongest sources """
......@@ -1055,24 +1068,28 @@ class SourceGeometry:
def set_sources(self, source_density, fluxes, n_src, background_horizon=1.):
"""
Define source density or directly positions and optional weights (cosmic ray luminosity).
Define source density or directly positions of sources and optional weights (cosmic ray luminosity).
:param source_density: source density (in 1 / Mpc^3) or array of shape (3, n_sources)
that point towards the center of the sources or keyword 'sbg'
:param source: either source density (in 1 / Mpc^3) as float or array or source coordinates of shape
(3, n_sources) or keyword of source catalog like 'sbg'
:param fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
:param n_src: Number of point sources to be considered.
:return: no return
"""
self.n_src = n_src
self.background_horizon = background_horizon
if isinstance(source_density, (int, float, np.int, np.float)):
float_like = isinstance(source_density, (int, float, np.int, np.float))
array_like = isinstance(source_density, np.ndarray)
string_like = isinstance(source_density, str)
if float_like or (array_like and np.shape(source_density) == (self.nsets, )):
# maximum radius for one source per cosmic ray (isotropy condition)
self.rmax = (3*n_src/4/np.pi/source_density)**(1/3.)
self.sources = coord.rand_vec((self.nsets, n_src)) # shape (3, nsets, n_src)
# random radius in volume (r^2 distribution)
self.distances = self.rmax * (np.random.random((self.nsets, n_src)))**(1/3.) # shape (nsets, n_src)
u = np.random.random((self.nsets, n_src))
self.distances = coord.atleast_kd(self.rmax, 2) * u**(1/3.) # shape (nsets, n_src)
self.source_fluxes = 1 / self.distances**2
elif isinstance(source_density, np.ndarray):
elif array_like:
source_density = np.reshape(source_density, (3, -1))
self.n_src = source_density.shape[1]
self.sources = source_density
......@@ -1080,7 +1097,7 @@ class SourceGeometry:
if fluxes is not None:
assert fluxes.shape == len(source_density[0].shape)
self.source_fluxes = fluxes
elif isinstance(source_density, str):
elif string_like:
sources, source_fluxes, distances = getattr(SourceScenario(), source_density.lower())()[:3]
self.n_src = len(source_fluxes)
self.sources = np.tile(sources, self.nsets).reshape(sources.shape[0], self.nsets, -1)
......
......@@ -335,8 +335,19 @@ class TestSourceBound(unittest.TestCase):
dmin = np.min(sim.universe.distances, axis=-1)
self.assertTrue((np.median(dmin) > 20))
densities = np.array([1e-6, 1e-3, 1])
n_each = 100
sim = SourceBound(int(n_each * len(densities)), self.ncrs)
sim.set_sources(source_density=np.repeat(densities, n_each))
dmin = np.min(sim.universe.distances[:100], axis=-1)
self.assertTrue((np.median(dmin) > 20))
dmin = np.min(sim.universe.distances[100:200], axis=-1)
self.assertTrue((np.median(dmin) > 3) & (np.median(dmin) < 10))
dmin = np.min(sim.universe.distances[200:], axis=-1)
self.assertTrue((np.median(dmin) < 1))
@unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
def test_03_fluxes(self):
def test_03a_fluxes(self):
sim = SourceBound(self.nsets, self.ncrs)
sim.set_energy(gamma=-2, log10e_min=19.6, log10_cut=20.5, rig_cut=False)
sim.set_charges(charges={'h': 1.})
......@@ -404,6 +415,36 @@ class TestSourceBound(unittest.TestCase):
fraction_inside = np.sum(mask_inside_10) / float(sim.ncrs * sim.nsets)
self.assertTrue(fraction_inside < 0.2)
@unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
def test_03b_fluxes(self):
densities = [1e-5, 1e-3, 1e-1]
bound_source = [[120, 180], [30, 60], [3, 15]]
for i, d in enumerate(densities):
# make a separate simulation for each source density
sim = SourceBound(self.nsets, 1000)
sim.set_energy(log10e_min=19.6)
sim.set_charges('first_minimum_CTG')
sim.set_sources(source_density=d)
sim.attenuate()
# check if number of events from the strongest source of the representative set is within the expected range
idx_select = sim._select_representative_set()
ns = np.array([np.sum(sim.crs['source_labels'][idx_select] == k) for k in range(sim.universe.n_src)]).max()
self.assertTrue((ns >= bound_source[i][0]) & (ns <= bound_source[i][1]))
# same check as above but with all source densities simulated simultaneously
n_each = 100
sim = SourceBound(int(n_each * len(densities)), 1000)
sim.set_energy(log10e_min=19.6)
sim.set_charges('first_minimum_CTG')
sim.set_sources(source_density=np.repeat(densities, n_each))
sim.attenuate()
for i in range(len(densities)):
mask = (np.arange(sim.nsets) < (i+1)*n_each) & (np.arange(sim.nsets) >= i*n_each)
idx_select = sim._select_representative_set(mask=mask)
ns = np.array([np.sum(sim.crs['source_labels'][idx_select] == k) for k in range(sim.universe.n_src)]).max()
self.assertTrue((ns >= bound_source[i][0]) & (ns <= bound_source[i][1]))
@unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
def test_04a_shuffle(self):
sim = SourceBound(self.nsets, self.ncrs)
......@@ -431,8 +472,6 @@ class TestSourceBound(unittest.TestCase):
@unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
def test_05a_first_minimum(self):
sim = SourceBound(self.nsets, 1000)
# sim.set_energy(gamma=-0.96, log10e_min=19.6, log10_cut=18.68, rig_cut=True)
# sim.set_charges(charges={'he': 0.673, 'n': 0.281, 'si': 0.046})
sim.set_energy(log10e_min=19.6)
sim.set_charges('first_minimum_CTG')
sim.set_sources(source_density=1e-3)
......@@ -443,13 +482,14 @@ class TestSourceBound(unittest.TestCase):
# sim.plot_distance()
sim = SourceBound(self.nsets, 1000)
# sim.set_energy(gamma=-0.96, log10e_min=19.6, log10_cut=18.68, rig_cut=True)
# sim.set_charges(charges={'he': 0.673, 'n': 0.281, 'si': 0.046})
sim.set_energy(log10e_min=19.6)
sim.set_charges('first_minimum_CTD')
sim.set_sources(source_density=1e-3)
sim.attenuate()
sim.attenuate(library_path=PATH + '/simulation/crpropa3__emin_18.5__emax_21.0__IRB_Dominguez11.npz')
sim.smear_sources(np.deg2rad(3))
# sim.plot_arrivals()
# sim.plot_spectrum()
# sim.plot_distance()
@unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
def test_05b_second_minimum(self):
......
Markdown is supported
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