Source code for prince_cr.data

"""Module inteded to contain some prince-specific data structures."""
import pickle as pickle
import os.path as path
import numpy as np
import scipy.constants as spc
import h5py


from prince_cr.util import convert_to_namedtuple, info
import prince_cr.config as config

#: Dictionary containing particle properties, like mass, charge
#: lifetime or branching ratios
try:
    spec_data = pickle.load(
        open(path.join(config.data_dir, "particle_data.ppo"), "rb"))
except UnicodeDecodeError:
    spec_data = pickle.load(
        open(path.join(config.data_dir, "particle_data.ppo"), "rb"), encoding='latin1')
except FileNotFoundError:
    info(0, 'Warning, particle database "particle_data.ppo" file not found.')


# Default units in Prince are ***cm, s, GeV***
# Define here all constants and unit conversions and use
# throughout the code. Don't write c=2.99.. whatever.
# Write clearly which units a function returns.
# Convert them if not standard unit
# Accept only arguments in the units above

UNITS_AND_CONVERSIONS_DEF = dict(
    c=1e2 * spc.c,
    cm2Mpc=1. / (spc.parsec * spc.mega * 1e2),
    Mpc2cm=spc.mega * spc.parsec * 1e2,
    m_proton=spc.physical_constants['proton mass energy equivalent in MeV'][0]
    * 1e-3,
    m_electron=spc.physical_constants['electron mass energy equivalent in MeV']
    [0] * 1e-3,
    r_electron=spc.physical_constants['classical electron radius'][0] * 1e2,
    fine_structure=spc.fine_structure,
    GeV2erg=1. / 624.15,
    erg2GeV=624.15,
    km2cm=1e5,
    yr2sec=spc.year,
    Gyr2sec=spc.giga * spc.year,
    cm2sec=1e-2 / spc.c,
    sec2cm=spc.c * 1e2)

# This is the immutable unit object to be imported throughout the code
PRINCE_UNITS = convert_to_namedtuple(UNITS_AND_CONVERSIONS_DEF, "PriNCeUnits")


[docs]class PrinceDB(object): """Provides access to data stored in an HDF5 file. The file contains all tables for runnin PriNCe. Currently the only still required file is the particle database. The tools to generate this database are publicly available in `PriNCe-data-utils <https://github.com/joheinze/PriNCe-data-utils>`_. """ def __init__(self): info(2, 'Opening HDF5 file', config.db_fname) self.prince_db_fname = path.join(config.data_dir, config.db_fname) if not path.isfile(self.prince_db_fname): raise Exception( 'Prince DB file {0} not found in "data" directory.'.format( config.db_fname)) with h5py.File(self.prince_db_fname, 'r') as prince_db: self.version = (prince_db.attrs['version']) def _check_subgroup_exists(self, subgroup, mname): available_models = list(subgroup) if mname not in available_models: info(0, 'Invalid choice/model', mname) info(0, 'Choose from:\n', '\n'.join(available_models)) raise Exception('Unknown selections.') def photo_nuclear_db(self, model_tag): info(10, 'Reading photo-nuclear db. tag={0}'.format(model_tag)) db_entry = {} with h5py.File(self.prince_db_fname, 'r') as prince_db: self._check_subgroup_exists(prince_db['photo_nuclear'], model_tag) for entry in ['energy_grid', 'fragment_yields', 'inel_mothers', 'inelastic_cross_sctions', 'mothers_daughters']: info(10, 'Reading entry {0} from db.'.format(entry)) db_entry[entry] = prince_db['photo_nuclear'][model_tag][entry][:] return db_entry def photo_meson_db(self, model_tag): info(10, 'Reading photo-nuclear db. tag={0}'.format(model_tag)) db_entry = {} with h5py.File(self.prince_db_fname, 'r') as prince_db: self._check_subgroup_exists(prince_db['photo_nuclear'], model_tag) for entry in ['energy_grid', 'xbins', 'fragment_yields', 'inel_mothers', 'inelastic_cross_sctions', 'mothers_daughters']: info(10, 'Reading entry {0} from db.'.format(entry)) db_entry[entry] = prince_db['photo_nuclear'][model_tag][entry][:] return db_entry def ebl_spline(self, model_tag, subset='base'): from scipy.interpolate import interp2d info(10, 'Reading EBL field splines. tag={0}'.format(model_tag)) with h5py.File(self.prince_db_fname, 'r') as prince_db: self._check_subgroup_exists(prince_db['EBL_models'], model_tag) self._check_subgroup_exists(prince_db['EBL_models'][model_tag], subset) spl_gr = prince_db['EBL_models'][model_tag][subset] return interp2d(spl_gr['x'], spl_gr['y'], spl_gr['z'], fill_value=0., kind='linear')
#: db_handler is the HDF file interface db_handler = PrinceDB()
[docs]class EnergyGrid(object): """Class for constructing a grid for discrete distributions. Since we discretize everything in energy, the name seems appropriate. All grids are log spaced. Args: lower (float): log10 of low edge of the lowest bin upper (float): log10 of upper edge of the highest bin bins_dec (int): bins per decade of energy """ def __init__(self, lower, upper, bins_dec): self.bins = np.logspace(lower, upper, int((upper - lower) * bins_dec + 1)) self.grid = 0.5 * (self.bins[1:] + self.bins[:-1]) self.widths = self.bins[1:] - self.bins[:-1] self.d = self.grid.size info( 5, 'Energy grid initialized {0:3.1e} - {1:3.1e}, {2} bins'.format( self.bins[0], self.bins[-1], self.grid.size))
[docs]class PrinceSpecies(object): """Bundles different particle properties for simplified availability of particle properties in :class:`prince_cr.core.PriNCeRun`. Args: pdgid (int): PDG ID of the particle particle_db (object): a dictionary with particle properties d (int): dimension of the energy grid """
[docs] @staticmethod def calc_AZN(nco_id): """Returns mass number :math:`A`, charge :math:`Z` and neutron number :math:`N` of ``nco_id``.""" Z, A = 1, 1 if nco_id >= 100: Z = nco_id % 100 A = (nco_id - Z) // 100 else: Z, A = 0, 0 return A, Z, A - Z
def __init__(self, ncoid, princeidx, d): info(5, 'Initializing new species', ncoid) #: Neucosma ID of particle self.ncoid = ncoid #: (bool) particle is a hadron (meson or baryon) self.is_hadron = False #: (bool) particle is a meson self.is_meson = False #: (bool) particle is a baryon self.is_baryon = False #: (bool) particle is a lepton self.is_lepton = False #: (bool) if it's an electromagnetic particle self.is_em = False #: (bool) particle is a lepton self.is_charged = False #: (bool) particle is a nucleus self.is_nucleus = False #: (bool) particle has an energy redistribution self.has_redist = False #: (bool) particle is stable self.is_stable = True #: (float) lifetime self.lifetime = np.inf #: (bool) particle is an alias (PDG ID encodes special scoring behavior) self.is_alias = False #: (str) species name in string representation self.sname = None #: decay channels if any self.decay_channels = {} #: Mass, charge, neutron number self.A, self.Z, self.N = 1, None, None #: Mass in atomic units or GeV self.mass = None #: (int) Prince index (in state vector) self.princeidx = princeidx # (dict) Dimension of energy grids (for idx calculations) self.grid_dims = {'default': d} # Obtain values for the attributes self._init_species() def _init_species(self): """Fill all class attributes with values from :var:`spec_data`, depending on ncoid.""" ncoid = self.ncoid dbentry = spec_data[ncoid] if ncoid < 200: self.is_nucleus = False if ncoid == 0: self.is_em = True elif ncoid in [100, 101]: self.is_hadron = True self.is_baryon = True self.is_nucleus = True self.A, self.Z, self.N = self.calc_AZN(ncoid) elif ncoid not in [2, 3, 4, 50]: self.is_hadron = True self.is_meson = True else: self.is_lepton = True if ncoid in [20, 21]: self.is_em = True elif ncoid in [7, 10]: self.is_alias = True else: self.is_nucleus = True self.A, self.Z, self.N = self.calc_AZN(ncoid) self.AZN = self.A, self.Z, self.N if ncoid <= config.redist_threshold_ID: self.has_redist = True if "name" not in dbentry: info(5, "Name for species", ncoid, "not defined") self.sname = "nucleus_{0}".format(ncoid) else: self.sname = dbentry["name"] self.charge = dbentry["charge"] self.is_charged = self.charge != 0 self.is_stable = dbentry["stable"] self.lifetime = dbentry["lifetime"] self.mass = dbentry["mass"] self.decay_channels = dbentry["branchings"] @property def sl(self): """Return the slice for this species on the grid can be used as spec[s.sl] Returns: (slice): a slice object pointing to the species in the state vecgtor """ idx = self.princeidx dim = self.grid_dims['default'] return slice(idx * dim, (idx + 1) * dim)
[docs] def lidx(self, grid_tag='default'): """Returns lower index of particle range in state vector. Returns: (int): lower index in state vector :attr:`PrinceRun.phi` """ return self.princeidx * self.grid_dims[grid_tag]
[docs] def uidx(self, grid_tag='default'): """Returns upper index of particle range in state vector. Returns: (int): upper index in state vector :attr:`PrinceRun.phi` """ return (self.princeidx + 1) * self.grid_dims[grid_tag]
[docs] def lbin(self, grid_tag='default'): """Returns lower bin of particle range in state vector. Returns: (int): lower bin in state vector :attr:`PrinceRun.phi` """ return self.princeidx * (self.grid_dims[grid_tag] + 1)
[docs] def ubin(self, grid_tag='default'): """Returns upper bin of particle range in state vector. Returns: (int): upper bin in state vector :attr:`PrinceRun.phi` """ return (self.princeidx + 1) * (self.grid_dims[grid_tag] + 1)
[docs] def indices(self, grid_tag='default'): """Returns a list of all indices in the state vector. Returns: (numpy.array): array of indices in state vector :attr:`PrinceRun.phi` """ idx = self.princeidx dim = self.grid_dims[grid_tag] return np.arange(idx * dim, (idx + 1) * dim)
[docs]class SpeciesManager(object): """Provides a database with particle and species.""" def __init__(self, ncoid_list, ed): # (dict) Dimension of primary grid self.grid_dims = {'default': ed} # Particle index shortcuts #: (dict) Converts Neucosma ID to index in state vector self.ncoid2princeidx = {} #: (dict) Converts particle name to index in state vector self.sname2princeidx = {} #: (dict) Converts Neucosma ID to reference of # :class:`data.PrinceSpecies` self.ncoid2sref = {} #: (dict) Converts particle name to reference of #:class:`data.PrinceSpecies` self.sname2sref = {} #: (dict) Converts prince index to reference of #:class:`data.PrinceSpecies` self.princeidx2sref = {} #: (dict) Converts index in state vector to Neucosma ID self.princeidx2ncoid = {} #: (dict) Converts index in state vector to reference # of :class:`data.PrinceSpecies` self.princeidx2pname = {} #: (int) Total number of species self.nspec = 0 self._gen_species(ncoid_list) self._init_species_tables() def _gen_species(self, ncoid_list): info(4, "Generating list of species.") # ncoid_list += spec_data["non_nuclear_species"] # Make sure list is unique and sorted ncoid_list = sorted(list(set(ncoid_list))) self.species_refs = [] # Define position in state vector (princeidx) by simply # incrementing it with the (sorted) list of Neucosma IDs for princeidx, ncoid in enumerate(ncoid_list): info( 4, "Appending species {0} at position {1}".format( ncoid, princeidx)) self.species_refs.append( PrinceSpecies(ncoid, princeidx, self.grid_dims['default'])) self.known_species = [s.ncoid for s in self.species_refs] self.redist_species = [ s.ncoid for s in self.species_refs if s.has_redist ] self.boost_conserv_species = [ s.ncoid for s in self.species_refs if not s.has_redist ] def _init_species_tables(self): for s in self.species_refs: self.ncoid2princeidx[s.ncoid] = s.princeidx self.sname2princeidx[s.sname] = s.princeidx self.princeidx2ncoid[s.princeidx] = s.ncoid self.princeidx2pname[s.princeidx] = s.sname self.ncoid2sref[s.ncoid] = s self.princeidx2sref[s.princeidx] = s self.sname2sref[s.sname] = s self.nspec = len(self.species_refs)
[docs] def add_grid(self, grid_tag, dimension): """Defines additional grid dimensions under a certain tag. Propagates changes to this variable to all known species. """ info(2, 'New grid_tag', grid_tag, 'with dimension', dimension) self.grid_dims[grid_tag] = dimension for s in self.species_refs: s.grid_dims = self.grid_dims
def __repr__(self): str_out = "" ident = 3 * ' ' for s in self.species_refs: str_out += s.sname + '\n' + ident str_out += 'NCO id : ' + str(s.ncoid) + '\n' + ident str_out += 'PriNCe idx : ' + str(s.princeidx) + '\n\n' return str_out