Commit 28ac0b19 authored by Marcus Wirtz's avatar Marcus Wirtz
Browse files

[simulations] Allow to set source density and catalog the same time (sampling)...

[simulations] Allow to set source density and catalog the same time (sampling) + allow n_src=0 for only background simulation
parent 5ebf081b
Pipeline #361199 failed with stages
in 5 minutes and 21 seconds
......@@ -534,7 +534,7 @@ class SourceBound(BaseSimulation):
else:
raise Exception('Charge type not understood (dictionary or string)')
def set_sources(self, source_density, fluxes=None, n_src=100, background_horizon=None):
def set_sources(self, source_density=None, sources=None, fluxes=None, n_src=None, rmax=None):
"""
Define source density or directly positions and optional weights (cosmic ray luminosity).
......@@ -544,7 +544,12 @@ class SourceBound(BaseSimulation):
:param n_src: Number of point sources to be considered.
:return: no return
"""
self.universe.set_sources(source_density, fluxes, n_src, background_horizon)
if (n_src is not None) and (n_src > self.ncrs):
print("Number of sources should not be larger than number of cosmic rays. Automatically set n_src = ncrs.")
n_src = self.ncrs
elif n_src is None:
n_src = max(30, min(500, self.ncrs // 10))
self.universe.set_sources(source_density, sources, fluxes, n_src, rmax)
self.crs['sources'] = self.universe.sources
self.crs['source_density'] = source_density
......@@ -610,12 +615,14 @@ class SourceBound(BaseSimulation):
library_path = PATH + '/simulation/crpropa3__emin_18.5__emax_21.0__IRB_Gilmore12.npz'
data = np.load(library_path, allow_pickle=True)
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("\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))
if self.universe.has_sources():
# Assign distance indices for all simulated soures, shape: (self.nsets, self.universe.n_src)
dis_bin_idx = self.universe.distance_indices(self.dis_bins)
min_dis = np.min(self.dis_bins)
if np.median(np.min(self.universe.distances, axis=1)) < min_dis:
print("Warning: Distance binning of simulation starts too far outside (%s Mpc)! " % min_dis)
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." % min_dis)
charge = {'h': 1, 'he': 2, 'n': 7, 'si': 14, 'fe': 26}
self.source_matrix = np.zeros((self.nsets, self.universe.n_src, 5))
......@@ -627,9 +634,12 @@ class SourceBound(BaseSimulation):
fractions = data['fractions'].item()[key] # dimensions: (energy_in, distances, charges_out, energy_out)
# reweight to spectral index (simulated gamma=-1) and apply energy / rigidity cut
fractions = self._reweight_spectrum(fractions, charge[key]) # dim: (distances, charges_out, energy_out)
self.source_matrix += f * np.sum(fractions, axis=-1)[dis_bin_idx]
self.arrival_matrix += f * fractions
if self.universe.has_sources():
self.source_matrix += f * np.sum(fractions, axis=-1)[dis_bin_idx]
# Account for optional luminosity weights of sources
if self.universe.has_sources():
self.source_matrix *= coord.atleast_kd(self.universe.source_fluxes, k=self.source_matrix.ndim)
def _reweight_spectrum(self, fractions, c):
......@@ -696,12 +706,13 @@ class SourceBound(BaseSimulation):
""" Internal function to sample the arrival directions """
# First, set random direction of all events
vecs = coord.rand_vec(self.shape)
distances = np.zeros(self.shape)
if self.universe.has_sources():
# Set source directions of simulated sources
mask = self.crs['source_labels'] >= 0 # mask all cosmic rays which originate within rmax
vecs[:, mask] = self.universe.sources[:, np.argwhere(mask)[:, 0], self.crs['source_labels'][mask]]
self.crs['vecs'] = vecs
distances = np.zeros(self.shape)
distances[mask] = self.universe.distances[np.where(mask)[0], self.crs['source_labels'][mask]]
self.crs['vecs'] = vecs
self.crs['distances'] = distances
return vecs
......@@ -717,9 +728,8 @@ 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)
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)
dis_idx = np.argmin(np.abs(np.atleast_1d(self.universe.rmax)[None] - self.dis_bins[:, None]), 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
......@@ -739,6 +749,7 @@ class SourceBound(BaseSimulation):
self.crs['distances'][mask_pick] = _d[perm]
# Assign charges and energies of close-by cosmic rays
if np.sum(mask_close) > 0:
dis_bin_idx = self.universe.distance_indices(self.dis_bins)
for i in range(1 + np.max(dis_bin_idx)):
# assign distance indices to CRs
......@@ -756,6 +767,7 @@ class SourceBound(BaseSimulation):
_mask[_mask] = mask
charge[_mask], log10e[_mask] = np.array(c)[idx[0]][perm], self.log10e_bins[:-1][idx[1]][perm]
log10e += d_log10e * np.random.random(self.shape)
self.crs['log10e'] = log10e
self.crs['charge'] = charge
return log10e, charge
......@@ -849,7 +861,8 @@ class SourceBound(BaseSimulation):
opath = '/tmp/composition%s__%s.pdf' % (self._get_charge_id(), e_id)
crs = self.crs
idx = self._select_representative_set() if idx is None else idx
if idx is None:
idx = self._select_representative_set()
bins = np.array(np.concatenate((np.arange(18.5, 20.1, 0.1), np.array([20.5]))))
mean_xmax = binned_statistic(crs['log10e'][idx, :], values=crs['xmax'][idx, :], statistic='mean', bins=bins)[0]
......@@ -876,10 +889,12 @@ class SourceBound(BaseSimulation):
""" Plot arrival map """
import matplotlib.pyplot as plt
from astrotools import skymap
idx = self._select_representative_set() if idx is None else idx
if idx is None:
idx = self._select_representative_set()
src_idx, ns = self._get_strongest_sources(idx)
if self.crs['source_density'] == 'sbg_lunardini':
if isinstance(self.crs['source_density'], str) and (self.crs['source_density'] == 'sbg_lunardini'):
src_idx = np.arange(0, 44, 1)
if len(src_idx) > 0:
labels_p = np.copy(self.crs['source_labels'][idx])
labels_p[~np.in1d(labels_p, src_idx) & (labels_p >= 0)] = 10*self.universe.n_src
for j, idxj in enumerate(np.sort(src_idx)):
......@@ -894,12 +909,13 @@ class SourceBound(BaseSimulation):
skymap.eventmap(self.crs['vecs'][:, idx][:, mask], c=labels_p[mask], cmap=cmap, cblabel='Source ID',
cticks=np.arange(0, len(src_idx)+1, 1), clabels=np.sort(src_idx),
vmin=-0.5, vmax=len(src_idx)-0.5)
else:
skymap.eventmap(self.crs['vecs'][:, idx][:, mask], c='0.5')
lon_src, lat_src = coord.vec2ang(self.universe.sources[:, idx])
plt.scatter(-lon_src, lat_src, c='k', marker='*', s=2*ns)
ns = np.sort(ns)[::-1]
plt.title('Strongest sources: (%i, %i, %i)' % (ns[0], ns[1], ns[2]), fontsize=15)
else:
skymap.eventmap(self.crs['vecs'][:, idx][:, mask], c='0.5')
if opath is None:
opath = '/tmp/arrival%s__emin_%s__ecut_%s.pdf' % (self._get_charge_id(), self.energy_setting['log10e_min'],
self.energy_setting['log10_cut'])
......@@ -1064,48 +1080,94 @@ class SourceGeometry:
self.sources = None
self.source_fluxes = None
self.distances = None
self.background_horizon = None
def set_sources(self, source_density, fluxes, n_src, background_horizon=1.):
def set_sources(self, source_density=None, sources=None, fluxes=None, n_src=None, rmax=None):
"""
Define source density or directly positions of sources and optional weights (cosmic ray luminosity).
: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 source_density: sets source density (in 1 / Mpc^3) as float or array with shape (self.nsets, )
:param sources: explicitly set source coordinates of shape (3, None) 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.
:param rmax: distance where background sources start.
:return: no return
"""
self.n_src = n_src
self.background_horizon = background_horizon
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)
if ((source_density is None) and (sources is None)) or (n_src == 0):
print("No explicit sources have been set!")
self.n_src, self.rmax = 0, 0
return
self.rmax, self.n_src = rmax, n_src
if (source_density is not None):
# Set number of sources or maximum radius depending on source density
if isinstance(source_density, np.ndarray):
assert np.shape(source_density) == (self.nsets, ), "If source density is array: shape == (nsets, )!"
assert (n_src is not None) or (rmax is not None), "Can not interpret n_src and rmax simultaneously!"
if rmax is None:
if sources is not None:
self.n_src = min(self.n_src, sources.shape[-1])
self.rmax = (3*self.n_src/4/np.pi/source_density)**(1/3.)
else:
assert not isinstance(source_density, np.ndarray), "Dont set 'rmax' when using a varying source denity!"
self.n_src = int(round(4/3*np.pi*self.rmax**3 * source_density))
if sources is not None:
assert np.all(sources.shape[-1] >= self.n_src), "Too few 'sources' for given keywords \
'source_density' and 'rmax'!"
if (sources is not None):
assert len(sources) == 3, "First dimension of source array needs to be (x, y, z), thus len(sources) == 3 !"
assert len(sources.shape) >= 2, "Shape of array sources requires at least two dimensions!"
if (source_density is not None) and (sources is None):
self._set_source_density()
elif (sources is not None):
self._set_source_positions(source_density, sources, fluxes)
elif isinstance(source_density, str) or isinstance(sources, str):
self._set_source_catalog(source_density, sources)
else:
raise Exception("Source scenario not understood.")
def _set_source_density(self):
""" Internal function to set source density of n_src sources between [0, rmax] Mpc"""
self.sources = coord.rand_vec((self.nsets, self.n_src)) # shape (3, nsets, n_src)
# random radius in volume (r^2 distribution)
u = np.random.random((self.nsets, n_src))
u = np.random.random((self.nsets, self.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 array_like:
source_density = np.reshape(source_density, (3, -1))
self.n_src = source_density.shape[1]
self.sources = source_density
self.distances = np.sqrt(self.sources**2, axis=0)
if fluxes is not None:
assert fluxes.shape == len(source_density[0].shape)
self.source_fluxes = fluxes
elif string_like:
sources, source_fluxes, distances = getattr(SourceScenario(), source_density.lower())()[:3]
def _set_source_positions(self, source_density, sources, fluxes):
""" Internal function to set sources by position """
if (source_density is None):
self.n_src = sources.shape[-1]
self.sources = sources
else:
sources = sources[:, np.sqrt(np.sum(sources**2, axis=0)) <= np.max(np.atleast_1d(self.rmax))]
n_tot = sources.shape[-1]
n_tot_min = np.sum(np.sqrt(np.sum(sources**2, axis=0)) <= np.min(np.atleast_1d(self.rmax)))
assert n_tot_min > self.n_src, "Not enough sources provided for given source density!"
if isinstance(source_density, np.ndarray):
d = np.sqrt(np.sum(sources**2, axis=0))
src_idx = [np.random.choice(n_tot, self.n_src, False, p=(d <= r)/(d <= r).sum()) for r in self.rmax]
else:
a = np.full((self.nsets, 1), self.n_src)
src_idx = np.apply_along_axis(lambda x: np.random.choice(n_tot, x, replace=False), 1, a)
self.sources = sources[:, np.array(src_idx)]
self.distances = np.sqrt(np.sum(self.sources**2, axis=0))
sources_normed = np.allclose(self.distances, 1.)
if sources_normed:
print("Warning: keyword 'sources' seems to be normalized to 1, thus all distances are set to 1 Mpc!")
self.source_fluxes = fluxes if fluxes is not None else 1 / self.distances**2
assert self.source_fluxes.shape == self.sources.shape[1:], "Sources and fluxes have different shape!"
def _set_source_catalog(self, source_density, sources):
""" Internal function to set sources by position """
catalog = source_density if isinstance(source_density, str) else sources
sources, source_fluxes, distances = getattr(SourceScenario(), catalog.lower())()[:3]
self.n_src = len(source_fluxes)
self.sources = np.tile(sources, self.nsets).reshape(sources.shape[0], self.nsets, -1)
self.source_fluxes = np.tile(source_fluxes, self.nsets).reshape(-1, source_fluxes.shape[0])
self.distances = np.tile(distances, self.nsets).reshape(-1, distances.shape[0])
self.rmax = np.amax(distances)
else:
raise Exception("Source scenario not understood.")
def distance_indices(self, distance_bins):
"""
......@@ -1117,6 +1179,10 @@ class SourceGeometry:
diff = np.abs(np.log(self.distances[np.newaxis] / distance_bins[:, np.newaxis, np.newaxis]))
return np.argmin(diff, axis=0)
def has_sources(self):
""" Check if geometry features sources """
return not (np.all(self.n_src == 0) or np.all(self.rmax == 0))
class SourceScenario:
"""Predefined source scenarios"""
......
......@@ -4,7 +4,7 @@ import numpy as np
import unittest
from astrotools import coord, gamale, healpytools as hpt
from astrotools.simulations import PATH, ObservedBound, SourceBound
from astrotools.simulations import PATH, ObservedBound, SourceBound, SourceGeometry
nside = 64
ncrs = 1000
......@@ -416,7 +416,7 @@ class TestSourceBound(unittest.TestCase):
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):
def test_03b_fluxes_multiple_densities(self):
densities = [1e-5, 1e-3, 1e-1]
bound_source = [[120, 180], [30, 60], [3, 15]]
......@@ -445,6 +445,17 @@ class TestSourceBound(unittest.TestCase):
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_03c_fluxes_no_sources(self):
sim = SourceBound(self.nsets, self.ncrs)
sim.set_energy(log10e_min=19.6)
sim.set_charges('first_minimum')
sim.set_sources()
sim.attenuate()
# 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_04a_shuffle(self):
sim = SourceBound(self.nsets, self.ncrs)
......@@ -505,6 +516,104 @@ class TestSourceBound(unittest.TestCase):
# sim.plot_spectrum()
# sim.plot_distance()
@unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
def test_6a_source_density_from_samples(self):
n_src = 10000
distances = 50 * np.random.random((1, n_src))**(1/3)
sources = coord.rand_fisher_vec(np.array([1, 0, 0]), kappa=1/np.deg2rad(30)**2, n=n_src) * distances
sim = SourceBound(self.nsets, 1000)
sim.set_energy(log10e_min=19.6)
sim.set_charges('first_minimum')
sim.set_sources(source_density=1e-4, sources=sources, n_src=300)
sim.attenuate()
sim.smear_sources(np.deg2rad(3))
sim.plot_arrivals(opath='source_hybrid_1e-4.png')
# sim.plot_spectrum()
# sim.plot_distance()
vecs_sum = np.mean(sim.crs['vecs'], axis=(-2, -1))
self.assertTrue(coord.angle(vecs_sum, np.array([1, 0, 0]))[0] < 0.5)
amplitude = 3 * np.sqrt(np.sum(vecs_sum**2))
self.assertTrue(amplitude > 1.)
@unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
def test_6b_source_density_from_samples(self):
densities = [1e-4, 1e-3, 1e-2]
bound_amplitude = [1., 0.5, 0.1]
# same check as above but with all source densities simulated simultaneously
n_each = 100
ncrs = 1000
sim = SourceBound(int(n_each * len(densities)), ncrs)
sim.set_energy(log10e_min=19.6)
sim.set_charges('first_minimum')
n_src = 10000
distances = 50 * np.random.random((1, n_src))**(1/3)
sources = coord.rand_fisher_vec(np.array([1, 0, 0]), kappa=1/np.deg2rad(30)**2, n=n_src) * distances
sim.set_sources(source_density=np.repeat(densities, n_each), sources=sources, n_src=300)
sim.attenuate()
sim.smear_sources(np.deg2rad(3))
for i in range(len(densities)):
mask = (np.arange(sim.nsets) < (i+1)*n_each) & (np.arange(sim.nsets) >= i*n_each)
sim.plot_arrivals(idx=sim._select_representative_set(mask), opath='source_hybrid_sd_%s.png' % densities[i])
vecs = sim.crs['vecs'][:, mask]
self.assertTrue(vecs.shape == (3, n_each, ncrs))
vecs_sum = np.mean(vecs, axis=(-2, -1))
self.assertTrue(coord.angle(vecs_sum, np.array([1, 0, 0]))[0] < 0.5)
amplitude = 3 * np.sqrt(np.sum(vecs_sum**2))
self.assertTrue(amplitude > bound_amplitude[i])
class TestSourceGeometry(unittest.TestCase):
def setUp(self):
self.nsets = 50
self.n_src = 100
def test_01_no_soures(self):
universe = SourceGeometry(self.nsets)
universe.set_sources()
self.assertTrue((universe.n_src == 0) and (universe.rmax == 0))
universe2 = SourceGeometry(self.nsets)
universe2.set_sources(source_density=1e-3, n_src=0)
self.assertTrue((universe2.n_src == 0) and (universe2.rmax == 0))
def test_02_source_density(self):
universe = SourceGeometry(self.nsets)
universe.set_sources(source_density=1e-3, n_src=self.n_src)
self.assertTrue(np.abs(universe.rmax - 28.8) < 0.1)
universe2 = SourceGeometry(self.nsets)
universe2.set_sources(source_density=1e-3, rmax=universe.rmax)
self.assertTrue((universe2.n_src == self.n_src) and (universe.rmax == universe2.rmax))
sources = coord.rand_vec(1000) * 30 * np.random.random((1, 1000))**(1/3)
distances = np.sqrt(np.sum(sources**2, axis=0))
universe = SourceGeometry(self.nsets)
universe.set_sources(source_density=1e-3, sources=sources, n_src=self.n_src)
self.assertTrue(np.abs(universe.rmax - 28.8) < 0.1)
self.assertTrue(np.all(universe.distances <= universe.rmax))
self.assertTrue(np.shape(universe.distances) == (self.nsets, self.n_src))
self.assertTrue(not np.allclose(universe.distances[0], universe.distances[1]))
for dis in universe.distances:
self.assertTrue(np.isin(dis, distances).all())
universe2 = SourceGeometry(self.nsets)
universe2.set_sources(source_density=1e-3, rmax=universe.rmax)
self.assertTrue((universe2.n_src == self.n_src) and (universe.rmax == universe2.rmax))
universe = SourceGeometry(self.nsets)
with self.assertRaises(AssertionError):
universe.set_sources(source_density=1, sources=sources)
with self.assertRaises(AssertionError):
universe.set_sources(source_density=1, sources=sources, n_src=100)
def test_03_sources(self):
sources = coord.rand_vec(1000) * 30 * np.random.random((1, 1000))**(1/3)
universe = SourceGeometry(self.nsets)
universe.set_sources(sources=sources)
self.assertTrue(np.allclose(universe.sources, sources))
class TestReweight(unittest.TestCase):
......
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