""" 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
from astrotools.statistics import random_choice_multi
PATH = os.path.split(__file__)[0]
[docs]def set_fisher_smeared_sources(nside, sources, delta, source_fluxes=None):
"""
Smears the source positions (optional fluxes) with a fisher distribution of width delta.
:param nside: nside of the HEALPix pixelization (default: 64)
:type nside: int
:param sources: array of shape (3, n_sources) that point towards the center of the sources
:param delta: float or array with same length as sources: width of the fisher distribution (in radians)
:param source_fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
:return: healpy map (with npix entries) for the smeared sources normalized to sum 1
"""
npix = hpt.nside2npix(nside)
nsrc = np.shape(sources)[1]
eg_map = np.zeros(npix)
if isinstance(delta, (int, float)):
delta = np.repeat(delta, nsrc)
if len(delta) != nsrc:
raise ValueError("Number of deltas must be 1 or equal to number of sources")
for i, v_src in enumerate(sources.T):
eg_map_add = hpt.fisher_pdf(nside, *v_src, k=1. / delta[i] ** 2)
if source_fluxes is not None:
eg_map_add *= source_fluxes[i]
eg_map += eg_map_add
return eg_map / eg_map.sum()
[docs]def sample_from_m_distributions(weight_matrix, n):
"""
Sample n events from m distributions each having k bins and defined in weight_matrix.
:param weight_matrix: shape (m, k), hosting m different distributions each with k bins.
:param n: Number of events drawn for each of the m distributions.
:return indices: Indice array with values 0..(k-1), in shape (m, n)
"""
weight_matrix /= weight_matrix.sum(axis=1, keepdims=True)
s, r = weight_matrix.cumsum(axis=1), np.random.random(n)
return (r[np.newaxis, np.newaxis] > s[:, :, np.newaxis]).sum(axis=1)
[docs]class BaseSimulation:
"""
Base class where the classes ObservedBound() and SourceBound() inherit from.
"""
def __init__(self, nsets, ncrs):
nsets = int(nsets)
ncrs = int(ncrs)
self.shape = (nsets, ncrs)
self.crs = cosmic_rays.CosmicRaysSets((nsets, ncrs))
self.nsets = nsets
self.ncrs = ncrs
self.signal_label = None
[docs] def set_xmax(self, z2a='double', model='EPOS-LHC'):
"""
Calculate Xmax bei gumbel distribution for the simulated energies and charges.
:param z2a: How the charge is converted to mass number ['double', 'empiric', 'stable', 'abundance']
:param model: Hadronic interaction for gumbel distribution
:return: no return
"""
assert 'xmax' not in self.crs.keys(), "Try to re-assign xmax values!"
if (not hasattr(self.crs, 'log10e')) or (not hasattr(self.crs, 'charge')):
raise Exception("""Use function set_energy() and set_charges() before using function set_xmax.
If you Use SourceBound simulation attenuate() also has to be used additionally""")
mass = getattr(nt.Charge2Mass(self.crs['charge']), z2a)()
mass = np.hstack(mass) if isinstance(mass, np.ndarray) else mass
xmax = auger.rand_xmax(np.hstack(self.crs['log10e']), mass, model=model)
self.crs['xmax'] = np.reshape(xmax, self.shape)
return self.crs['xmax']
[docs] def apply_uncertainties(self, err_e=None, err_a=None, err_xmax=None):
"""
Apply measurement uncertainties.
:param err_e: relative uncertainty on the energy (typical 0.14)
:param err_a: angular uncertainty in degree on the arrival directions (typical 0.5 degree)
:param err_xmax: absolute uncertainty on the shower depth xmax (typical 15 g/cm^2)
"""
if err_e is not None:
self.crs['log10e'] += np.log10(1 + np.random.normal(scale=err_e, size=self.shape))
if err_a is not None:
vecs = coord.rand_fisher_vec(self.crs['vecs'], 1./np.deg2rad(err_a)**2)
lon, lat = coord.vec2ang(vecs)
self.crs['lon'] = lon.reshape(self.shape)
self.crs['lat'] = lat.reshape(self.shape)
if err_xmax is not None:
self.crs['xmax'] += np.random.normal(err_xmax)
[docs] 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.crs.shape_array.dtype.names:
self.crs[_key] = self.crs[_key][sets_ids, shuffle_ids]
self.signal_label = self.signal_label[sets_ids, shuffle_ids]
if 'vecs' in self.crs.keys():
sets_ids_3d = np.repeat(np.arange(3), np.prod(self.shape)).reshape((3,) + self.shape)
self.crs['vecs'] = self.crs['vecs'][sets_ids_3d, sets_ids, np.stack([shuffle_ids] * 3)]
[docs]class ObservedBound(BaseSimulation):
"""
Class to simulate cosmic ray arrival scenario by sources located at the sky, including energies, charges, smearing
and galactic magnetic field effects.
This is an observed bound simulation, thus energies and composition is set on Earth and might differ at sources.
"""
def __init__(self, nside, nsets, ncrs):
"""
Initialization of ObservedBound simulation.
:param nside: nside of the HEALPix pixelization (default: 64)
:param nsets: number of simulated cosmic ray sets
:param ncrs: number of cosmic rays per set
"""
BaseSimulation.__init__(self, nsets, ncrs)
self.nside = nside
self.npix = hpt.nside2npix(nside)
self.crs['nside'] = nside
self.sources = None
self.source_fluxes = None
self.rigidities = None
self.rig_bins = None
self.cr_map = None
self.lensed = None
self.exposure = {'map': None, 'a0': None, 'zmax': None}
self.signal_idx = None
[docs] def set_energy(self, log10e_min, log10e_max=None, energy_spectrum=None, **kwargs):
"""
Setting the energies of the simulated cosmic ray set.
:param log10e_min: Either minimum energy (in log10e) for AUGER setup or power-law
or numpy.array of energies in shape (nsets, ncrs)
:type log10e_min: Union[np.ndarray, float]
:param log10e_max: Maximum energy for AUGER setup
:param energy_spectrum: model that is defined in below class EnergySpectrum
:param kwargs: additional named keywords which will be passed to class EnergySpectrum
:return: energies in log10e
"""
assert 'log10e' not in self.crs.keys(), "Try to re-assign energies!"
if isinstance(log10e_min, np.ndarray):
if log10e_min.shape == self.shape:
self.crs['log10e'] = log10e_min
elif log10e_min.size == self.ncrs:
print("Warning: the same energies have been used for all simulated sets (nsets).")
self.crs['log10e'] = np.tile(log10e_min, self.nsets).reshape(self.shape)
else:
raise Exception("Shape of input energies not in format (nsets, ncrs).")
elif isinstance(log10e_min, (float, np.float, int, np.int)):
if energy_spectrum is not None:
log10e = getattr(EnergySpectrum(self.shape, log10e_min, log10e_max), energy_spectrum)(**kwargs)
else:
if (log10e_min < 17.) or (log10e_min > 21.):
print("Warning: Specified parameter log10e_min below 17 or above 20.5.")
log10e = auger.rand_energy_from_auger(self.shape, log10e_min=log10e_min, log10e_max=log10e_max)
self.crs['log10e'] = log10e
else:
raise Exception("Input of emin could not be understood.")
return self.crs['log10e']
[docs] def set_charges(self, charge, **kwargs):
"""
Setting the charges of the simulated cosmic ray set.
:param charge: Either charge number that is used or numpy.array of charges in shape (nsets, ncrs) or keyword
:type: charge: Union[np.ndarray, str, float]
:return: charges
"""
assert 'charge' not in self.crs.keys(), "Try to re-assign charges!"
if isinstance(charge, np.ndarray):
if charge.shape == self.shape:
self.crs['charge'] = charge
elif charge.size == self.ncrs:
self.crs['charge'] = np.tile(charge, self.nsets).reshape(self.shape)
else:
raise Exception("Shape of input charges not in format (nsets, ncrs).")
elif isinstance(charge, (float, np.float, int, np.int)):
self.crs['charge'] = charge
elif isinstance(charge, str):
if not hasattr(self.crs, 'log10e'):
raise Exception("Use function set_energy() before accessing a composition model.")
self.crs['charge'] = getattr(CompositionModel(self.shape, self.crs['log10e']), charge.lower())(**kwargs)
elif isinstance(charge, dict):
z = np.array([nt.ELEMENT_CHARGE[key] for key in charge])
self.crs['charge'] = np.random.choice(z, self.shape, p=[charge[key] for key in charge])
else:
raise Exception("Input of charge could not be understood.")
return self.crs['charge']
[docs] def set_xmax(self, z2a='double', model='EPOS-LHC'):
"""
Calculate Xmax bei gumbel distribution for the simulated energies and charges.
For more information refer to BaseSimulation.set_xmax().
"""
return BaseSimulation.set_xmax(self, z2a, model)
[docs] def set_sources(self, sources, fluxes=None):
"""
Define source position and optional weights (cosmic ray luminosity).
:param sources: array of shape (3, n_sources) that point towards the center of the sources or integer for
number of random sources or keyword ['sbg']
:param fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
:return: no return
"""
if isinstance(sources, np.ndarray):
if (len(np.shape(sources)) == 1) and len(sources) == 3:
sources = np.reshape(sources, (3, 1))
assert len(np.shape(sources)) == 2
assert np.shape(sources)[0] == 3
self.sources = sources
if fluxes is not None:
assert fluxes.size == len(sources.T)
self.source_fluxes = fluxes
elif isinstance(sources, (int, np.int)):
src_pix = np.random.randint(0, self.npix, sources)
self.sources = np.array(hpt.pix2vec(self.nside, src_pix))
if fluxes is not None:
assert fluxes.size == sources
self.source_fluxes = fluxes
elif isinstance(sources, str):
self.sources, self.source_fluxes = getattr(SourceScenario(), sources.lower())()[:2]
else:
raise Exception("Source scenario not understood.")
[docs] def set_rigidity_bins(self, lens_or_bins, cover_rigidity=True):
"""
Defines the bin centers of the rigidities.
:param lens_or_bins: Either the binning of the lens (left bin borders) or the lens itself
:return: no return
"""
if self.rig_bins is None:
if 'log10e' not in self.crs.keys():
raise Exception("Cannot define rigidity bins without energies specified.")
if 'charge' not in self.crs.keys():
print("Warning: Energy dependent deflection instead of rigidity dependent (set_charges to avoid)")
if isinstance(lens_or_bins, np.ndarray):
bins = lens_or_bins # type: np.array
dbins = bins[1] - bins[0]
else:
bins = np.array(lens_or_bins.log10r_mins)
dbins = lens_or_bins.dlog10r
rigidities = self.crs['log10e']
if 'charge' in self.crs.keys():
rigidities = rigidities - np.log10(self.crs['charge'])
if cover_rigidity:
assert np.all(np.min(rigidities) >= np.min(bins - dbins / 2.)), "Rigidities not covered by lens!"
idx = np.digitize(rigidities, bins) - 1
rigs = (bins + dbins / 2.)[idx]
rigs = rigs.reshape(self.shape)
self.rigidities = rigs
self.rig_bins = np.unique(rigs)
return self.rig_bins
[docs] def smear_sources(self, delta, dynamic=None):
"""
Smears the source positions with a fisher distribution of width delta (optional dynamic smearing).
:param delta: either the constant width of the fisher distribution or dynamic (delta / R[10 EV]), in radians
:param dynamic: if True, function applies dynamic smearing (delta / R[EV]) with value delta at 10 EV rigidity
:return: no return
"""
if self.sources is None:
raise Exception("Cannot smear sources without positions.")
if (dynamic is None) or (dynamic is False):
shape = (1, self.npix)
eg_map = np.reshape(set_fisher_smeared_sources(self.nside, self.sources, delta, self.source_fluxes), shape)
else:
if self.rig_bins is None:
raise Exception("Cannot dynamically smear sources without rigidity bins (use set_rigidity_bins()).")
eg_map = np.zeros((self.rig_bins.size, self.npix))
for i, rig in enumerate(self.rig_bins):
delta_temp = delta / 10 ** (rig - 19.)
eg_map[i] = set_fisher_smeared_sources(self.nside, self.sources, delta_temp, self.source_fluxes)
self.cr_map = eg_map
[docs] def lensing_map(self, lens, cache=None):
"""
Apply a galactic magnetic field to the extragalactic map.
:param lens: Instance of astrotools.gamale.Lens class, for the galactic magnetic field
:param cache: Caches all the loaded lens parts (increases speed, but may consume a lot of memory!)
:return: no return
"""
if self.lensed:
print("Warning: Cosmic Ray maps were already lensed before.")
if self.rig_bins is None:
self.set_rigidity_bins(lens)
if self.cr_map is None:
self._fix_point_source()
arrival_map = np.zeros((self.rig_bins.size, self.npix))
for i, rig in enumerate(self.rig_bins):
lp = lens.get_lens_part(rig, cache=cache)
eg_map_bin = self.cr_map[0] if self.cr_map.size == self.npix else self.cr_map[i]
lensed_map = lp.dot(eg_map_bin)
if not cache:
del lp.data, lp.indptr, lp.indices
arrival_map[i] = lensed_map / np.sum(lensed_map) if np.sum(lensed_map) > 0 else 1. / self.npix
self.lensed = True
self.cr_map = arrival_map
[docs] def apply_exposure(self, a0=-35.25, zmax=60):
"""
Apply the exposure (coverage) of any experiment (default: AUGER) to the observed maps.
:param a0: equatorial declination [deg] of the experiment (default: AUGER, a0=-35.25 deg)
:type a0: float, int
:param zmax: maximum zenith angle [deg] for the events
:return: no return
"""
self.exposure.update({'map': hpt.exposure_pdf(self.nside, a0, zmax), 'a0': a0, 'zmax': zmax})
if self.cr_map is None:
self.cr_map = np.reshape(self.exposure['map'], (1, self.npix))
else:
self.cr_map = self.cr_map * self.exposure['map']
self.cr_map /= np.sum(self.cr_map, axis=-1)[:, np.newaxis]
[docs] def arrival_setup(self, fsig, seed=1):
"""
Creates the realizations of the arrival maps.
:param fsig: signal fraction of cosmic rays per set (signal = originate from sources)
:type fsig: float
:type seed: random seed
:return: no return
"""
np.random.seed(seed)
dtype = np.uint16 if self.nside <= 64 else np.uint32
pixel = np.zeros(self.shape).astype(dtype)
# Setup the signal part
n_sig = int(fsig * self.ncrs)
signal_label = np.in1d(range(self.ncrs), np.arange(0, n_sig, 1))
if n_sig == 0:
pass
elif (self.cr_map is None) and (self.sources is None):
pixel[:, signal_label] = np.random.choice(self.npix, (self.nsets, n_sig))
elif (self.cr_map is None):
self._fix_point_source()
if self.cr_map.size == self.npix:
pixel[:, signal_label] = np.random.choice(self.npix, (self.nsets, n_sig),
p=np.hstack(self.cr_map)/np.sum(np.hstack(self.cr_map)))
else:
for i, rig in enumerate(self.rig_bins):
mask_rig = (rig == self.rigidities) * signal_label # type: np.ndarray
n = np.sum(mask_rig)
if n == 0:
continue
pixel[mask_rig] = np.random.choice(self.npix, n, p=self.cr_map[i])
elif np.sum(self.cr_map) > 0:
if self.cr_map.size == self.npix:
pixel[:, signal_label] = np.random.choice(self.npix, (self.nsets, n_sig), p=np.hstack(self.cr_map))
else:
for i, rig in enumerate(self.rig_bins):
mask_rig = (rig == self.rigidities) * signal_label # type: np.ndarray
n = np.sum(mask_rig)
if n == 0:
continue
pixel[mask_rig] = np.random.choice(self.npix, n, p=self.cr_map[i])
else:
raise Exception("No signal probability to sample signal from!")
# Setup the background part
n_back = self.ncrs - n_sig
bpdf = self.exposure['map'] if self.exposure['map'] is not None else np.ones(self.npix) / float(self.npix)
pixel[:, np.invert(signal_label)] = np.random.choice(self.npix, (self.nsets, n_back), p=bpdf)
self.signal_label = np.repeat(signal_label[np.newaxis], self.nsets, axis=0)
self.crs['pixel'] = pixel
[docs] def apply_uncertainties(self, err_e=None, err_a=None, err_xmax=None):
""" Apply measurement uncertainties (see BaseSimulation.apply_uncertainties()). """
BaseSimulation.apply_uncertainties(self, err_e, err_a, err_xmax)
[docs] def get_data(self, convert_all=False, shuffle=False):
"""
Returns the data in the form of the cosmic_rays.CosmicRaysSets() container.
:param convert_all: if True, also vectors and lon/lat of the cosmic ray sets are saved (more memory expensive)
:type convert_all: bool
:return: Instance of cosmic_rays.CosmicRaysSets() classes
Example:
sim = ObservedBound()
...
crs = sim.get_data(convert_all=True)
pixel = crs['pixel']
lon = crs['lon']
lat = crs['lat']
log10e = crs['log10e']
charge = crs['charge']
"""
if convert_all is True:
if not hasattr(self.crs, 'lon') or not hasattr(self.crs, 'lat'):
self.convert_pixel(convert_all=True)
if shuffle:
self.shuffle_events()
self.crs['signal_label'] = self.signal_label
return self.crs
[docs] def convert_pixel(self, keyword='vecs', convert_all=False):
"""
Converts pixelized information stored under key 'pixel' to vectors (keyword='vecs')
or angles (keyword='angles'), accessible via 'lon', 'lat'. When convert_all is True, both are saved.
This can be used at a later stage, if convert_all was set to False originally.
"""
shape = (-1, self.shape[0], self.shape[1])
a0, zmax = self.exposure['a0'], self.exposure['zmax']
if (self.exposure['map'] is not None) and (a0 is not None) and (zmax is not None):
vecs = hpt.rand_exposure_vec_in_pix(self.nside, np.hstack(self.crs['pixel']), a0, zmax).reshape(shape)
else:
vecs = hpt.rand_vec_in_pix(self.nside, np.hstack(self.crs['pixel'])).reshape(shape)
if keyword == 'vecs' or convert_all is True:
if hasattr(self.crs, 'lon') and hasattr(self.crs, 'lat') and not all:
raise Exception('Not useful to convert pixels to vecs, information already there!')
self.crs['vecs'] = vecs
if keyword == 'angles' or convert_all is True:
if keyword == 'angles' and not convert_all:
if hasattr(self.crs, 'vecs') and not convert_all:
raise Exception('Not useful to convert pixels to angles, information already there!')
lon, lat = coord.vec2ang(vecs)
self.crs['lon'] = lon.reshape(self.shape)
self.crs['lat'] = lat.reshape(self.shape)
else:
raise Exception('keyword not understood! Use angles or vecs!')
def _fix_point_source(self):
""" Internal function to set non-smeared sources automatically """
print("Warning: No extragalactic smearing of the sources performed before lensing (smear_sources). "
"Sources are considered point-like.")
eg_map = np.zeros((1, self.npix))
weights = self.source_fluxes if self.source_fluxes is not None else 1.
eg_map[:, hpt.vec2pix(self.nside, *self.sources)] = weights
self.cr_map = eg_map
[docs]class SourceBound(BaseSimulation):
"""
Class to simulate cosmic ray arrival scenario by sources located at the sky, including energies, charges, smearing.
This is a source bound simulation, thus energies and composition is set at the sources and might differ at Earth.
"""
def __init__(self, nsets, ncrs):
"""
Initialization of SourceBound simulation.
:param nsets: number of simulated cosmic ray sets
:param ncrs: number of cosmic rays per set
"""
BaseSimulation.__init__(self, nsets, ncrs)
self.arrival_matrix, self.source_matrix = None, None
self.dis_bins, self.log10e_bins = None, None
self.energy_setting = None
self.charge_weights = None
self.universe = SourceGeometry(nsets, cosmo=True)
self.redshift_bins = 0
[docs] def set_energy(self, log10e_min, gamma=-2, log10_cut=20., rig_cut=True):
"""
Define energy spectrum and cut off energy of sources.
:param log10e_min: Minimum threshold energy for observation
:param gamma: Spectral index of energy spectrum at sources
:param log10_cut: Maximum cut-off energy or rigidity for sources
:param rig_cut: if True, log10_cut refers to a rigidity cut
"""
if not isinstance(gamma, np.ndarray):
gamma = np.array([gamma])
if not isinstance(log10_cut, np.ndarray):
log10_cut = np.array([log10_cut])
self.energy_setting = {'log10e_min': log10e_min, 'gamma': gamma, 'log10_cut': log10_cut, 'rig_cut': rig_cut}
[docs] 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 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)
or string ('first_minimum'/'second_minimum')
"""
if isinstance(charges, str):
# (CTG, CTD) according to naming in Auger's combined fit paper (arXiv:1612.07155)
if charges in ('first_minimum_CTG', 'first_minimum'):
self.energy_setting['gamma'], self.energy_setting['log10_cut'] = np.array([1.03]), np.array([18.21])
charges = {'h': 0.6794, 'he': 0.31, 'n': 0.01, 'si': 0.0006}
elif charges == 'second_minimum_CTG':
self.energy_setting['gamma'], self.energy_setting['log10_cut'] = np.array([-0.87]), np.array([18.62])
charges = {'n': 0.88, 'si': 0.12}
elif charges == 'first_minimum_CTD':
self.energy_setting['gamma'], self.energy_setting['log10_cut'] = np.array([1.47]), np.array([18.15])
charges = {'h': 0.4494, 'he': 0.52, 'n': 0.03, 'si': 0.0006}
elif charges == 'first_minimum_walz':
self.energy_setting['gamma'], self.energy_setting['log10_cut'] = np.array([-0.62]), np.array([18.56])
charges = {'h': 0.001, 'he': 0.001, 'n': 0.985, 'fe': 0.012}
elif charges == 'second_minimum_walz':
self.energy_setting['gamma'], self.energy_setting['log10_cut'] = np.array([-2.03]), np.array([19.88])
charges = {'he': 0.003, 'n': 0.92, 'fe': 0.077}
else:
raise Exception('Charge keyword not understood (first_minimum/second_minimum)')
self.energy_setting['rig_cut'] = True
if isinstance(charges, dict):
if not isinstance(charges[list(charges.keys())[0]], np.ndarray):
charges = {key: np.array([charges[key]]) for key in charges}
fraction = np.sum([charges[key] for key in charges], axis=0)
assert np.all(np.abs(fraction - 1) < 1e-4), "Fractions of charges dictionary must be normalized!"
self.charge_weights = charges
else:
raise Exception('Charge type not understood (dictionary or string)')
[docs] 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).
: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 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
"""
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
[docs] def attenuate(self, library_path=None, inside_fraction=None, evolution=0):
"""
Apply attenuation for far away sources based on CRPropa simulations
:param library_path: Input library file to use.
:param evolution: m for source evolution of the homogeneous background
(right now: not for foreground sources)
:param cosmological_expansion: if the cosmological expansion of the universe
is taken into account. If this is True, all
distances (e.g. rmax) have to be given in comoving distance.
"""
# Prepare the arrival and source matrix by reweighting
self._prepare_arrival_matrix(library_path)
# Assign source allocation of cosmic rays
self._allocate_sources(inside_fraction)
# Assign the arrival directions of the cosmic rays
self._set_arrival_directions()
# Assign charges and energies of the cosmic rays
self._set_charges_energies(evolution=evolution)
[docs] def set_xmax(self, z2a='double', model='EPOS-LHC'):
"""
Calculate Xmax by Gumbel distribution for the simulated energies and charges.
For more information refer to BaseSimulation.set_xmax().
"""
return BaseSimulation.set_xmax(self, z2a, model)
[docs] def smear_sources(self, delta, dynamic=None):
"""
Smears the source positions with a fisher distribution of width delta (optional dynamic smearing).
:param delta: either the constant width of the fisher distribution or dynamic (delta / R[10 EV]), in radians
:param dynamic: if True, function applies dynamic smearing (delta / R[EV]) with value delta at 10 EV rigidity
:return: no return
"""
assert self.universe.sources is not None, "Cannot smear sources without positions!"
assert self.crs is not None, "Cannot smear sources without calling attenuate() before!"
mask_close = self.crs['source_labels'] >= 0
d = delta if dynamic is None else delta / \
10 ** (self.crs['log10e'] - np.log10(self.crs['charge']) - 19.)[mask_close]
self.crs['vecs'][:, mask_close] = coord.rand_fisher_vec(self.crs['vecs'][:, mask_close], kappa=1/d**2)
[docs] def get_data(self, shuffle=False):
"""
Returns the data in the form of the cosmic_rays.CosmicRaysSets() container.
:return: Instance of cosmic_rays.CosmicRaysSets() classes
"""
if shuffle:
self.shuffle_events()
self.crs['signal_label'] = self.signal_label
return self.crs
def _prepare_arrival_matrix(self, library_path):
# pylint: disable = unexpected-keyword-arg
"""
Prepare the arrival and source matrix by reweighting
:param library_path: Input library file to use.
"""
if (self.energy_setting is None) or (self.charge_weights is None):
raise Exception("You have to define energies and charges before (set_energy() and set_charges())!")
if library_path is None:
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']
try:
self.redshift_bins = data['redshifts']
except KeyError:
self.redshift_bins = None
if self.universe.has_sources():
# Assign distance indices for all simulated sources, 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))
self.arrival_matrix = np.zeros((self.nsets, self.dis_bins.size, 5, len(self.log10e_bins) - 1))
for key in self.charge_weights:
f = self.charge_weights[key]
if np.all(f == 0):
continue
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: (nsets, distances, charges_out, energy_out)
self.arrival_matrix += coord.atleast_kd(f, fractions.ndim) * fractions
if self.universe.has_sources():
self.source_matrix += coord.atleast_kd(f, fractions.ndim - 1) * np.take_along_axis(np.sum(fractions, axis=-1), dis_bin_idx[..., np.newaxis], axis=-2)
# 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):
""" Internal function to reweight to desired energy spectrum and rigidity cut """
assert fractions.ndim == 4, "Element arrival matrix fraction must have 4 dimensions!"
nlog10cut = self.energy_setting['log10_cut'].shape[0]
bin_center = (self.log10e_bins[:-1] + self.log10e_bins[1:]) / 2.
# reweight spectrum (simulated is gamma=-1 as resulting from equally binning in log space)
fractions = fractions[np.newaxis, :]
fractions = fractions * 10**((coord.atleast_kd(self.energy_setting['gamma'], fractions.ndim) + 1) * (coord.atleast_kd(bin_center, fractions.ndim - 1) - 19)[np.newaxis, :])
# Apply the rigidity cut
log10_cut = self.energy_setting['log10_cut']
if self.energy_setting['rig_cut']:
# convert rigidity cut (log10_cut) to corresponding energy cut (log10e_cut)
log10e_cut = log10_cut + np.log10(c)
fcut = np.ones((nlog10cut, bin_center.size))
# cut function from combined fit paper: https://arxiv.org/pdf/1612.07155.pdf
fcut[bin_center > log10e_cut[..., np.newaxis]] = np.exp(1 - 10**((np.vstack([bin_center] * nlog10cut) - log10e_cut[..., np.newaxis])[bin_center > log10e_cut[..., np.newaxis]]))
fractions = fractions * coord.atleast_kd(fcut, fractions.ndim)
else:
if log10_cut <= self.energy_setting['log10e_min']:
print("Warning: element with charge Z=%i is not injected with rigidity cut of" % c)
print("%s and log10e > %s" % (self.energy_setting['log10_cut'], self.energy_setting['log10e_min']))
fractions[:, bin_center > log10_cut, :, :, :] = 0
fractions[:, :, :, :, bin_center < self.energy_setting['log10e_min']] = 0. # energy threshold (measured)
fractions = np.sum(fractions, axis=1) # sum over injected energies to reweight the spectrum
return fractions
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)
dist_frac = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3) # dim self.arrival_matrix: (nsets, distances, charges_out, energy_out)
rmax = np.atleast_1d(self.universe.rmax) * np.ones((self.nsets))
rmax_mask = coord.atleast_kd(rmax[:, np.newaxis] >= self.dis_bins, dist_frac.ndim)
inside_fraction = np.sum(dist_frac * rmax_mask, axis=(1, 2, 3)) / np.sum(dist_frac, axis=(1, 2, 3))
return inside_fraction
def _allocate_sources(self, inside_fraction=None):
""" Internal function to assign the source allocation of the signal (inside rmax) cosmic rays """
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 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(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 < 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
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):
print("Warning: not enough sources. Set keyword 'n_src' (currently %s) higher?" % self.universe.n_src)
self.crs['source_labels'] = source_labels
occ = np.apply_along_axis(lambda x: np.bincount(x+1)[x+1], axis=1, arr=source_labels)
self.signal_label = (occ >= 2) & (source_labels >= 0)
return source_labels
def _set_arrival_directions(self):
""" 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]]
distances[mask] = self.universe.distances[np.where(mask)[0], self.crs['source_labels'][mask]]
self.crs['vecs'] = vecs
self.crs['distances'] = distances
return vecs
def _set_charges_energies(self, evolution=0):
""" Internal function to assign charges and energies of all cosmic rays """
log10e = np.zeros(self.shape)
charge = np.zeros(self.shape)
c = [1, 2, 7, 14, 26]
d_dis, d_log10e = np.diff(np.log10(self.dis_bins))[0], np.around(np.diff(self.log10e_bins), decimals=2)[0]
if np.diff(self.log10e_bins)[0] < 0.009:
d_dis, d_log10e = np.diff(np.log10(self.dis_bins))[0], np.diff(self.log10e_bins)[0]
print('Warning, your log10ebins are very fine and we cannot guarantee that emin will \
be set completely correclty (floating point accuracy)')
# Assign distances, charges and energies of background cosmic rays
mask_close = self.crs['source_labels'] >= 0
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)
if evolution != 0:
try:
dis_bins_e = self.dis_bins * (1 + self.redshift_bins)**evolution
except TypeError:
import crpropa as crp
print('have to import crpropa: your attenuation file does not include redshift bins!')
self.redshift_bins = np.array([crp.comovingDistance2Redshift(dis * crp.Mpc) for dis in self.dis_bins])
dis_bins_e = self.dis_bins * (1 + self.redshift_bins)**evolution
arrival_mat_far = self.arrival_matrix * coord.atleast_kd(dis_bins_e, 3)[np.newaxis, :] # dim: (nsets, distances, charges_out, energy_out)
else:
arrival_mat_far = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3)[np.newaxis, :] # dim: (nsets, distances, charges_out, energy_out)
# determine for each set the distance bin of the horizon rmax
dis_idx = np.argmin(np.abs(np.atleast_1d(self.universe.rmax)[None] - self.dis_bins[:, None]), axis=0) * np.ones(self.nsets, dtype=int)
# loop over all occuring distance bins and fill events of the respective sets
for _dis in np.sort(np.unique(dis_idx)):
current_sets = (dis_idx == _dis)
arrival_mat_far[current_sets, 0:_dis] = 0 # set probabilities inside rmax to zero
arrival_mat_far /= coord.atleast_kd(np.sum(arrival_mat_far, axis=(1, 2, 3)), arrival_mat_far.ndim) # normalize arrival matrix to then draw random samples
mask_pick = ~mask_close & (dis_idx == _dis)[:, np.newaxis] # select background events of respective set
n_pick = np.sum(mask_pick, axis=-1)
max_num = np.max(n_pick)
arrival_idx = np.zeros((self.nsets, max_num), dtype=int)
for k in np.where(current_sets)[0]:
amf = arrival_mat_far[k]
arrival_idx[k][:n_pick[k]] = np.random.choice(amf.size, size=n_pick[k], p=amf.flatten()) # get the indices in terms of the original shape (distances, charges, energies)
idx = np.unravel_index(arrival_idx, arrival_mat_far[0].shape) # shape (distances, charges_out, energy_out)
perm_idx = np.random.rand(arrival_idx.shape[0], arrival_idx.shape[1]).argsort(axis=-1)
del arrival_idx
_d = 10**(np.log10(self.dis_bins)[idx[0]] + (np.random.random(idx[0].shape) - 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
del idx
samplenum_mask = np.take_along_axis(np.array([[True] * _s + [False] * (max_num - _s) for _s in n_pick]), perm_idx, axis=-1)
log10e[mask_pick] = np.take_along_axis(_lge, perm_idx, axis=-1)[samplenum_mask]
charge[mask_pick] = np.take_along_axis(_c, perm_idx, axis=-1)[samplenum_mask]
self.crs['distances'][mask_pick] = np.take_along_axis(_d, perm_idx, axis=-1)[samplenum_mask]
# 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)
w_mats = self.arrival_matrix / coord.atleast_kd(self.arrival_matrix.sum(axis=(-1, -2)), self.arrival_matrix.ndim)
cr_idx = dis_bin_idx[np.where(mask_close)[0], self.crs['source_labels'][mask_close]]
for i in range(1 + np.max(dis_bin_idx)):
w_mat = w_mats[:, i, ...] # shape charges, energies
# assign distance indices to CRs
mask = cr_idx == i
if np.sum(mask) == 0:
continue
_mask = np.copy(mask_close)
_mask[_mask] = mask
# procedure in analogy to code block above
samplenum = np.sum(_mask, axis=-1)
max_num = samplenum.max()
samplenum_mask = np.array([[True] * _s + [False] * (max_num - _s) for _s in samplenum])
arrival_idx = random_choice_multi(a=w_mat[0].size, k=max_num, p=w_mat.reshape(w_mat.shape[0], -1))
idx = np.unravel_index(arrival_idx, w_mat[0].shape)
perm_idx = np.random.rand(w_mat.shape[0], max_num).argsort(axis=-1)
charge[_mask] = np.take_along_axis(np.array(c)[idx[0]], perm_idx, axis=-1)[samplenum_mask]
log10e[_mask] = np.take_along_axis(self.log10e_bins[:-1][idx[1]], perm_idx, axis=-1)[samplenum_mask]
log10e += d_log10e * np.random.random(self.shape)
self.crs['log10e'] = log10e
self.crs['charge'] = charge
return log10e, charge
def _get_charge_id(self):
""" Return charge id of universe """
chargegroups = ['h', 'he', 'n', 'si', 'fe']
return ''.join(['__%s_%s' % (key, self.charge_weights[key]) for key in chargegroups
if key in self.charge_weights])
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')
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 """
ns = np.array([np.sum(self.crs['source_labels'][idx] == k) for k in range(self.universe.n_src)])
n_t = self.ncrs
n_thres = ns >= n_t
while (np.sum(n_thres) < min_number_src) & (n_t >= 3):
n_t = int(0.8 * n_t)
n_thres = ns >= n_t
try:
src_idx = np.squeeze(np.argwhere(n_thres))[np.argsort(ns[n_thres])]
except IndexError:
src_idx = []
return src_idx, ns
[docs] def plot_spectrum(self, opath=None): # pragma: no cover
""" Plot energy spectrum """
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
from astrotools import statistics as st
log10e = np.array(self.crs['log10e'])
plt.figure(figsize=(6, 4))
dspectrum = auger.SPECTRA_DICT[17]
log10e_center = dspectrum['logE']
# log10e_center = log10e_center # [log10e_center + 0.05 >= self.energy_setting['log10e_min']]
flux = (10 ** log10e_center) ** 3 * dspectrum['mean']
flux_high = (10 ** log10e_center) ** 3 * dspectrum['stathi']
flux_low = (10 ** log10e_center) ** 3 * dspectrum['statlo']
plt.errorbar(log10e_center, flux, yerr=[flux_low, flux_high], color='red',
fmt='.', linewidth=1, markersize=8, capsize=0, label='Auger 2017')
log10e_bins = np.arange(np.round(np.min(log10e), 1), np.max(log10e) + 0.1, 0.1)
n, bins = np.histogram(log10e.flatten(), bins=log10e_bins)
nall = np.apply_along_axis(lambda x: np.histogram(x, bins=log10e_bins)[0], 1, log10e)
flux_sim = (10 ** st.mid(bins)) ** 2 * n / self.nsets
flux_unc = (10 ** st.mid(bins)) ** 2 * nall
norm = flux[np.argmin(np.abs(log10e_center - st.mid(bins)[0]))] / flux_sim[0]
plt.errorbar(st.mid(bins), norm * flux_sim, marker='.', ls='none', color='k', xerr=0.05,
yerr=np.std(norm*flux_unc, axis=0), zorder=-1)
# plot arriving composition (spline approximation)
colors = ['firebrick', 'darkorange', 'forestgreen', 'deepskyblue', 'darkblue']
e, c = ['h', 'he', 'n', 'si', 'fe'], [1, 2, 7, 14, 26]
for i, ei in enumerate(e):
mask = self.crs['charge'] == c[i]
if np.sum(mask) == 0:
continue
_n, _bins = np.histogram(log10e[mask].flatten(), bins=log10e_bins)
_flux_sim = (10 ** st.mid(_bins)) ** 2 * _n / self.nsets
smooth_spline = interp1d(st.mid(_bins), norm * _flux_sim, kind='cubic', bounds_error=False)
x = np.linspace(log10e_bins[0], log10e_bins[-1], 100)
plt.plot(x, smooth_spline(x), color=colors[i], label=ei)
plt.yscale('log')
plt.xlim([self.energy_setting['log10e_min'] - 0.01, np.max(log10e) + 0.07])
plt.ylim([1e36, 1e38])
plt.legend(loc='lower left', fontsize=12)
yl = r'$E^{3} \, J(E)$ [km$^{-2}$ yr$^{-1}$ sr$^{-1}$ eV$^{2}$]'
plt.ylabel(yl, fontsize=16)
plt.xlabel(r'$\log_{10}$($E$/eV)', fontsize=16)
if opath is None:
opath = '/tmp/spectrum%s__emin_%s__ecut_%s.pdf' % (self._get_charge_id(), self.energy_setting['log10e_min'],
self.energy_setting['log10_cut'])
plt.savefig(opath, bbox_inches='tight')
plt.close()
[docs] def plot_composition(self, idx=None, opath=None): # pragma: no cover
""" Plots composition (comparison measurements/simulation) """
import matplotlib.pyplot as plt
from scipy.stats import binned_statistic
if opath is None:
e_id = 'emin_%s__ecut_%s' % (self.energy_setting['log10e_min'], self.energy_setting['log10_cut'])
opath = '/tmp/composition%s__%s.pdf' % (self._get_charge_id(), e_id)
crs = self.crs
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]
std_xmax = binned_statistic(crs['log10e'][idx, :], values=crs['xmax'][idx, :], statistic='std', bins=bins)[0]
n_xmax = binned_statistic(crs['log10e'][idx, :], values=crs['xmax'][idx, :], statistic='count', bins=bins)[0]
mids = binned_statistic(crs['log10e'][idx, :], values=crs['log10e'][idx, :], statistic='mean', bins=bins)[0]
fig = plt.figure(figsize=(6, 4))
ax = fig.add_subplot(111)
auger.plot_mean_xmax(ax=ax, models=['EPOS-LHC', 'QGSJetII-04'])
xmax_err = std_xmax/np.sqrt(n_xmax)
ax.errorbar(mids, mean_xmax, yerr=xmax_err, marker='o', color='darkgray', label='simulation', fmt='')
ax.set_xlim(right=20.6)
fig.savefig(opath, bbox_inches='tight')
fig = plt.figure(figsize=(6, 4))
ax = fig.add_subplot(111)
auger.plot_std_xmax(ax=ax, models=['EPOS-LHC', 'QGSJetII-04'])
ax.scatter(mids, std_xmax, marker='o', color='darkgray', label='simulation')
ax.set_xlim(left=17.5, right=20.6)
fig.savefig(opath.replace('.', '_std.'), bbox_inches='tight')
[docs] def plot_arrivals(self, idx=None, opath=None, emin=None): # pragma: no cover
""" Plot arrival map """
import matplotlib.pyplot as plt
from astrotools import skymap
if idx is None:
idx = self._select_representative_set()
src_idx, ns = self._get_strongest_sources(idx)
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)):
labels_p[labels_p == idxj] = j
cmap = plt.get_cmap('jet', len(src_idx))
cmap.set_over('k')
cmap.set_under('gray')
mask = np.ones(self.crs.ncrs).astype(bool)
if emin is not None:
mask = self.crs['log10e'][idx, :] > np.log10(emin)+18.
if len(src_idx) > 0:
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)
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'])
plt.savefig(opath, bbox_inches='tight')
plt.close()
[docs] def plot_composition_skymap(self, idx=None, opath=None, emin=None): # pragma: no cover
""" Plot arrival map """
import matplotlib.pyplot as plt
from astrotools import skymap
idx = self._select_representative_set() if idx is None else idx
charges = self.crs['charge'][idx, :]
cmap = plt.get_cmap('jet_r', 26)
src_idx, ns = self._get_strongest_sources(idx)
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(src_idx):
labels_p[labels_p == idxj] = j
mask = np.ones(self.crs.ncrs).astype(bool)
if emin is not None:
mask = self.crs['log10e'][idx, :] > np.log10(emin)+18.
skymap.eventmap(self.crs['vecs'][:, idx, labels_p >= 0][:, mask[labels_p >= 0]],
c=charges[labels_p >= 0][mask[labels_p >= 0]],
cmap=cmap, cblabel='$Z$',
cticks=np.array([1, 2, 6, 12, 20, 26]), vmin=0.5, vmax=26.5,
s=25, alpha=0.6, marker='v')
lons, lats = coord.vec2ang(self.crs['vecs'][:, idx, labels_p < 0][:, mask[labels_p < 0]])
plt.scatter(-lons, lats, c=charges[labels_p < 0][mask[labels_p < 0]],
cmap=cmap, s=15, alpha=0.4, marker='o', vmin=0.5, vmax=26.5,)
lon_src, lat_src = coord.vec2ang(self.universe.sources[:, idx])
plt.scatter(-lon_src, lat_src, c='0.5', linewidth=0.5, edgecolor='k', marker='*', s=7*ns)
ns = np.sort(ns)[::-1]
plt.title('Strongest sources: (%i, %i, %i)' % (ns[0], ns[1], ns[2]), fontsize=15)
if opath is None:
opath = '/tmp/arrival_charge%s__emin_%s__ecut_%s.pdf' % (self._get_charge_id(),
self.energy_setting['log10e_min'],
self.energy_setting['log10_cut'])
plt.savefig(opath, bbox_inches='tight')
plt.close()
[docs] def plot_distance(self, sets='all', opath=None, emin=None, sig='both'): # pragma: no cover
# pylint: disable=too-many-statements,invalid-unary-operand-type
""" Plot histogram of travel distances (either mean of all sets or one specific) """
import matplotlib.pyplot as plt
dists = self.crs['distances']
amax = np.amax(dists)
bin_width = 5 if amax <= 300 else int(amax/50)
bins = np.arange(-bin_width, np.amax(dists), bin_width)
bin_centers = bins[1:-1] + bin_width/2. # without first artificial bin
if sig == 'both':
mask_sig = np.ones((self.nsets, self.ncrs)).astype(bool)
elif sig == 'sig':
mask_sig = self.signal_label
elif sig == 'background':
mask_sig = ~self.signal_label
else:
raise Exception('Signal keyword not understood, use sig / background / both')
if isinstance(sets, int):
mask = np.ones(self.crs.ncrs).astype(bool)
if emin is not None:
mask = self.crs['log10e'][sets] > np.log10(emin)+18.
mask = mask & mask_sig[sets]
hist_p = np.histogram(dists[sets][(self.crs['charge'][sets] == 1) & mask], bins)[0][1::]
hist_he = np.histogram(dists[sets][(self.crs['charge'][sets] == 2) & mask], bins)[0][1::]
hist_cno = np.histogram(dists[sets][(self.crs['charge'][sets] > 2) &
(self.crs['charge'][sets] <= 11) & mask], bins)[0][1::]
hist_medium = np.histogram(dists[sets][(self.crs['charge'][sets] >= 12) &
(self.crs['charge'][sets] < 16) & mask], bins)[0][1::]
hist_heavy = np.histogram(dists[sets][(self.crs['charge'][sets] >= 17) & mask], bins)[0][1::]
yerr_p, yerr_he, yerr_cno, yerr_medium, yerr_heavy = 0, 0, 0, 0, 0
amax = np.amax(dists[sets])
elif sets == 'all':
# median and std histogram over nsets
mask = np.ones((self.crs.nsets, self.crs.ncrs)).astype(bool)
if emin is not None:
mask = self.crs['log10e'] > np.log10(emin)+18.
mask = mask & mask_sig
hists_p = np.apply_along_axis(lambda a: np.histogram(a, bins)[0],
1, np.where((self.crs['charge'] == 1) & mask, dists,
np.ones_like(dists)*(-1)))
hists_he = np.apply_along_axis(lambda a: np.histogram(a, bins)[0],
1, np.where((self.crs['charge'] == 2) & mask, dists,
np.ones_like(dists)*(-1)))
hists_cno = np.apply_along_axis(lambda a: np.histogram(a, bins)[0],
1, np.where((self.crs['charge'] > 2) &
(self.crs['charge'] <= 11) & mask,
dists, np.ones_like(dists)*(-1)))
hists_medium = np.apply_along_axis(lambda a: np.histogram(a, bins)[0], 1,
np.where((self.crs['charge'] >= 12) &
(self.crs['charge'] < 16) & mask,
dists, np.ones_like(dists)*(-1)))
hists_heavy = np.apply_along_axis(lambda a: np.histogram(a, bins)[0], 1,
np.where((self.crs['charge'] >= 17) & mask,
dists, np.ones_like(dists)*(-1)))
hist_p = np.median(hists_p, axis=0)[1::]
hist_he = np.median(hists_he, axis=0)[1::]
hist_cno = np.median(hists_cno, axis=0)[1::]
hist_medium = np.median(hists_medium, axis=0)[1::]
hist_heavy = np.median(hists_heavy, axis=0)[1::]
yerr_p = np.std(hists_p, axis=0)[1::]
yerr_he = np.std(hists_he, axis=0)[1::]
yerr_cno = np.std(hists_cno, axis=0)[1::]
yerr_medium = np.std(hists_medium, axis=0)[1::]
yerr_heavy = np.std(hists_heavy, axis=0)[1::]
else:
raise Exception("Sets not understood, either give set id number or keyword all!")
plt.figure(figsize=(12, 9))
plt.bar(bin_centers, hist_p,
color='firebrick', label=r'$Z = 1$', width=bin_width*0.8,
yerr=yerr_p)
plt.bar(bin_centers, hist_he,
color='darkorange', label=r'$Z = 2$', width=bin_width*0.8,
yerr=yerr_he, bottom=hist_p)
plt.bar(bin_centers, hist_cno,
color='forestgreen', label=r'$3 \leq Z \leq 11$', width=bin_width*0.8,
yerr=yerr_cno, bottom=hist_p+hist_he)
plt.bar(bin_centers, hist_medium,
color='deepskyblue', label=r'$12 \leq Z \leq 16$', width=bin_width*0.8,
bottom=hist_p+hist_he+hist_cno, yerr=yerr_medium)
plt.bar(bin_centers, hist_heavy,
bottom=hist_p+hist_he+hist_cno+hist_medium,
color='darkblue', label=r'$Z \geq 17$', width=bin_width*0.8,
yerr=yerr_heavy)
plt.axvline(x=self.universe.rmax, color='0.5', linestyle='dashed', label='Source shell')
plt.ylabel(r'flux [a.u.] / %s Mpc' % bin_width, fontsize=22)
plt.xlabel('distance / Mpc', fontsize=22)
plt.xticks(fontsize=22)
plt.yticks(fontsize=22)
plt.legend(fontsize=20)
plt.xlim(-0.1, amax)
plt.ylim(bottom=0)
plt.grid()
if opath is None:
opath = '/tmp/distance%s__emin_%s__ecut_%s__set%s.png' % (self._get_charge_id(),
self.energy_setting['log10e_min'],
self.energy_setting['log10_cut'], sets)
plt.savefig(opath, bbox_inches='tight')
plt.close()
plt.clf()
[docs]class SourceGeometry:
"""
Class to set up a 3d Universe out of isotropically distributed sources.
"""
def __init__(self, nsets, cosmo=True):
self.nsets = nsets
self.cosmo = cosmo
self.rmax = None
self.n_src = None
self.sources = None
self.source_fluxes = None
self.distances = None
[docs] 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_density: sets source density (in 1 / Mpc^3) as float or array with shape (self.nsets, )
can be keyword of source catalog like 'sbg'
:param sources: explicitly set source coordinates of shape (3, None)
Here, distances are always lookback distances, never comoving!
: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
"""
if ((source_density is None) and (sources is None)) or (n_src == 0):
print("No explicit sources have been set!")
self.n_src = 0
self.rmax = 0 if rmax is None else rmax
return
if isinstance(source_density, str) or isinstance(sources, str):
self._set_source_catalog(source_density, sources)
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 density!"
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)
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
uses uniform distribution in comoving volume, assigns self.distances as lookback distance like in the reweighting database
:param cosmo: use luminosity distance to assign fluxes
"""
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, self.n_src))
if self.cosmo:
comoving_distances = coord.atleast_kd(self.rmax, 2) * u**(1/3.) # shape (nsets, n_src), uniform in comoving volume
# use parameters obtained by fitting the redshift - comoving distance relation from 1 to 5000 Mpc with a polynomial order 6
# leads to differences between z(d)-of < 10^-5 and improves speed of calculation by factor 1000
# the chosen model is the LambdaCDM model with H0=67.3, Om0=0.315, Ode0=0.685
# similarly for the lookback/light-travel distance dt: polynomial of order 12, difference dt(z)-dt < 1e-5 Mpc << dt(0.001) - dc (0.001) = 2e-3 Mpc
fitted_params = [4.80166012e-24, -2.41614004e-20, 2.85709487e-16, 1.50785163e-12, 1.21672083e-08, 2.24387564e-04, 9.34693823e-06]
comoving2redshift = np.poly1d(fitted_params)
redshifts = comoving2redshift(comoving_distances)
fitted_params_lookback = [-6.07682754e-01, 7.72159574e+00, -4.45567484e+01, 1.54948994e+02
-3.63854131e+02, 6.15857662e+02, -8.01697479e+02, 9.19036748e+02
-1.20144989e+03, 1.98091151e+03, -3.27958548e+03, 4.45456591e+03, 1.78292637e-06]
redshift2lookback = np.poly1d(fitted_params_lookback)
self.distances = redshift2lookback(redshifts)
luminosity_distances = (1+redshifts)*comoving_distances
self.source_fluxes = 1 / luminosity_distances**2
else:
self.distances = coord.atleast_kd(self.rmax, 2) * u**(1/3.) # shape (nsets, n_src) # without cosmo
self.source_fluxes = 1 / self.distances**2
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)
[docs] def distance_indices(self, distance_bins):
"""
Get indices of given distance bins in the shape of the sources (nsets, n_src).
:param distance_bins: Distance bins which refer to indices
:return indices: indices of distance_bins in shape (nsets, n_src)
"""
diff = np.abs(np.log(self.distances[np.newaxis] / distance_bins[:, np.newaxis, np.newaxis]))
return np.argmin(diff, axis=0)
[docs] def has_sources(self):
""" Check if geometry features sources """
return not (np.all(self.n_src == 0) or np.all(self.rmax == 0))
[docs]class SourceScenario:
"""Predefined source scenarios"""
def __init__(self):
pass
[docs] @staticmethod
def sbg():
"""Starburst Galaxy Scenario used in GAP note 2017_007"""
# Position, fluxes, distances, names of starburst galaxies proposed as UHECRs sources
# by J. Biteau & O. Deligny (2017)
# Internal Auger publication: GAP note 2017_007
lon = np.array([97.4, 141.4, 305.3, 314.6, 138.2, 95.7, 208.7, 106, 240.9, 242, 142.8, 104.9, 140.4, 148.3,
141.6, 135.7, 157.8, 172.1, 238, 141.9, 36.6, 20.7, 121.6])
lat = np.array([-88, 40.6, 13.3, 32, 10.6, 11.7, 44.5, 74.3, 64.8, 64.4, 84.2, 68.6, -17.4, 56.3, -47.4, 24.9,
48.4, -51.9, -54.6, 55.4, 53, 27.3, 60.2])
vecs = coord.ang2vec(np.deg2rad(lon), np.deg2rad(lat))
distance = np.array([2.7, 3.6, 4, 4, 4, 5.9, 6.6, 7.8, 8.1, 8.1, 8.7, 10.3, 11, 11.4, 15, 16.3, 17.4, 17.9,
22.3, 46, 80, 105, 183])
flux = np.array([13.6, 18.6, 16., 6.3, 5.5, 3.4, 1.1, 0.9, 1.3, 1.1, 2.9, 3.6, 1.7, 0.7, 0.9, 2.6, 2.1, 12.1,
1.3, 1.6, 0.8, 1., 0.8])
names = np.array(['NGC 253', 'M82', 'NGC 4945', 'M83', 'IC 342', 'NGC 6946', 'NGC 2903', 'NGC 5055', 'NGC 3628',
'NGC 3627', 'NGC 4631', 'M51', 'NGC 891', 'NGC 3556', 'NGC 660', 'NGC 2146', 'NGC 3079',
'NGC 1068', 'NGC 1365', 'Arp 299', 'Arp 220', 'NGC 6240', 'Mkn 231'])
return vecs, flux, distance, names
[docs] @staticmethod
def sbg_lunardini():
"""Lunardini SBG catalog + Circinus (44 sources), from J.Biteau, received May 2019 per email"""
opath = os.path.split(__file__)[0] + '/data/Lunardini_SBG_catalog.dat'
data = np.genfromtxt(opath, skip_header=1, usecols=(2, 4, 5, 8, 10),
dtype={'names': ('names', 'ra', 'dec', 'distance', 'flux'),
'formats': ('|S11', float, float, float, float)})
vecs = coord.eq2gal(coord.ang2vec(np.deg2rad(data['ra']), np.deg2rad(data['dec'])))
flux = data['flux'] * 100 / np.sum(data['flux'])
names = np.array([data['names'][i].decode('UTF-8') for i in range(len(data['names']))]).flatten()
return vecs, flux, data['distance'], names
[docs] @staticmethod
def sbg_NGC253():
"""Lunardini SBG catalog + Circinus (44 sources), without NGC 253 for cross checks"""
opath = os.path.split(__file__)[0] + '/data/Lunardini_SBG_catalog.dat'
data = np.genfromtxt(opath, skip_header=1, usecols=(2, 4, 5, 8, 10),
dtype={'names': ('names', 'ra', 'dec', 'distance', 'flux'),
'formats': ('|S11', float, float, float, float)})
data = np.delete(data, 6)
vecs = coord.eq2gal(coord.ang2vec(np.deg2rad(data['ra']), np.deg2rad(data['dec'])))
flux = data['flux'] * 100 / np.sum(data['flux'])
names = np.array([data['names'][i].decode('UTF-8') for i in range(len(data['names']))]).flatten()
return vecs, flux, data['distance'], names
[docs] @staticmethod
def sbg_NGC1068():
"""Lunardini SBG catalog + Circinus (44 sources), without NGC 253 for cross checks"""
opath = os.path.split(__file__)[0] + '/data/Lunardini_SBG_catalog.dat'
data = np.genfromtxt(opath, skip_header=1, usecols=(2, 4, 5, 8, 10),
dtype={'names': ('names', 'ra', 'dec', 'distance', 'flux'),
'formats': ('|S11', float, float, float, float)})
data = np.delete(data, 16)
vecs = coord.eq2gal(coord.ang2vec(np.deg2rad(data['ra']), np.deg2rad(data['dec'])))
flux = data['flux'] * 100 / np.sum(data['flux'])
names = np.array([data['names'][i].decode('UTF-8') for i in range(len(data['names']))]).flatten()
return vecs, flux, data['distance'], names
[docs] @staticmethod
def agn_3fhl():
"""catalog of gamma-AGN emitters from 3FHL catalog, downloaded 15.12.2020 from
https://www.auger.unam.mx/AugerWiki/ArrivalDirections/HEfd"""
opath = os.path.split(__file__)[0] + '/data/Fermi3FHL_AGN_sample.dat'
data = np.genfromtxt(opath, skip_header=1, usecols=(2, 4, 5, 8, 10),
dtype={'names': ('names', 'ra', 'dec', 'distance', 'flux'),
'formats': ('|S11', float, float, float, float)})
vecs = coord.eq2gal(coord.ang2vec(np.deg2rad(data['ra']), np.deg2rad(data['dec'])))
flux = data['flux'] * 100 / np.sum(data['flux'])
names = np.array([data['names'][i].decode('UTF-8') for i in range(len(data['names']))]).flatten()
return vecs, flux, data['distance'], names
[docs] @staticmethod
def gamma_agn():
"""AGN scenario used in GAP note 2017_007"""
# Position, fluxes, distances, names of gamma_AGNs proposed as UHECRs sources by J. Biteau & O. Deligny (2017)
# Internal Auger publication: GAP note 2017_007
lon = np.array([309.6, 283.7, 150.6, 150.2, 235.8, 127.9, 179.8, 280.2, 63.6, 112.9, 131.9, 98, 340.7, 135.8,
160, 243.4, 77.1])
lat = np.array([19.4, 74.5, -13.3, -13.7, 73, 9, 65, -54.6, 38.9, -9.9, 45.6, 17.7, 27.6, -9, 14.6, -20, 33.5])
vecs = coord.ang2vec(np.deg2rad(lon), np.deg2rad(lat))
distance = np.array([3.7, 18.5, 76, 83, 95, 96, 136, 140, 148, 195, 199, 209, 213, 218, 232, 245, 247])
flux = np.array([0.8, 1, 2.2, 1, 0.5, 0.5, 54, 0.5, 20.8, 3.3, 1.9, 6.8, 1.7, 0.9, 0.4, 1.3, 2.3])
names = np.array(['Cen A Core', 'M 87', 'NGC 1275', 'IC 310', '3C 264', 'TXS 0149+710', 'Mkn 421',
'PKS 0229-581', 'Mkn 501', '1ES 2344+514', 'Mkn 180', '1ES 1959+650', 'AP Librae',
'TXS 0210+515', 'GB6 J0601+5315', 'PKS 0625-35', 'I Zw 187'])
return vecs, flux, distance, names
[docs] @staticmethod
def cena():
"""Only Cen A as source https://www.auger.unam.mx/AugerWiki/ArrivalDirections/HEfd"""
opath = os.path.split(__file__)[0] + '/data/Fermi3FHL_AGN_sample.dat'
data = np.genfromtxt(opath, skip_header=1, usecols=(2, 4, 5, 8, 10),
dtype={'names': ('names', 'ra', 'dec', 'distance', 'flux'),
'formats': ('|S11', float, float, float, float)})
vecs = coord.eq2gal(coord.ang2vec(np.deg2rad(data['ra']), np.deg2rad(data['dec'])))
names = np.array([data['names'][i].decode('UTF-8') for i in range(len(data['names']))]).flatten()
return vecs[:, 0].reshape((3, 1)), np.array([100]), np.array([data['distance'][0]]), np.array([names[0]])
[docs]class CompositionModel:
"""Predefined composition models"""
def __init__(self, shape, log10e=None):
self.shape = shape
self.log10e = log10e
[docs] def mixed(self):
"""Simple estimate of the composition above ~20 EeV by M. Erdmann (2017)"""
z = {'z': [1, 2, 6, 7, 8], 'p': [0.15, 0.45, 0.4 / 3., 0.4 / 3., 0.4 / 3.]}
charges = np.random.choice(z['z'], self.shape, p=z['p'])
return charges
[docs] def mixed_clipped(self):
"""mixed from above, but CNO group only Z=6 because of no lenses at low rigidities"""
z = {'z': [1, 2, 6], 'p': [0.15, 0.45, 0.4]}
charges = np.random.choice(z['z'], self.shape, p=z['p'])
return charges
[docs] def equal(self):
"""Assumes a equal distribution in (H, He, N, Fe) groups."""
z = {'z': [1, 2, 7, 26], 'p': [0.25, 0.25, 0.25, 0.25]}
charges = np.random.choice(z['z'], self.shape, p=z['p'])
return charges
[docs] def auger(self, smoothed=True, model='EPOS-LHC'):
"""Simple estimate from AUGER Xmax measurements"""
log10e = self.log10e
charges = auger.rand_charge_from_auger(np.hstack(log10e), model=model,
smoothed=smoothed).reshape(self.shape)
return charges
[docs] def auger_exponential(self):
"""Simple exponential estimate from AUGER Xmax measurements"""
log10e = self.log10e
charges = auger.rand_charge_from_exponential(log10e)
return charges
[docs]class EnergySpectrum:
"""Predefined energy spectra"""
def __init__(self, shape, log10e_min, log10e_max=20.5):
self.shape = shape
self.log10e_min = log10e_min
self.log10e_max = log10e_max
[docs] def power_law(self, gamma=-3):
"""
Power law spectrum, with spectral index corresponding to non differential spectrum,
where gamma=-3.7 corresponds to the AUGER fit at intermediate energies.
:param gamma: non-differential spectral index (E ~ E^(gamma))
:return: energies in shape self.shape
"""
emin = 10**(self.log10e_min - 18.)
emax = 10**(self.log10e_max - 18.)
u = np.random.random(self.shape)
if np.abs(1 + gamma) < 1e-3:
e = np.exp((np.log(emax) - np.log(emin)) * u + np.log(emin))
else:
exp = 1. / (1 + gamma)
e = ((emax**(1+gamma) - emin**(1+gamma)) * u + emin**(1+gamma))**exp
return 18. + np.log10(e)
[docs] def auger_fit(self):
""" Energies following the AUGER spectrum above log10e_min 17.5. """
return auger.rand_energy_from_auger(self.shape, self.log10e_min, self.log10e_max)