def _init_coupling_mat(self):
"""Initialises the coupling matrix directly in sparse (csr) format.
"""
info(0, 'Initiating coupling matrix in ({:}) format'.format('CSR'))
from scipy.sparse import csr_matrix
if using_cupy:
# For GPU we initialize the csr matrix on the host and then cast to GPU
from cupyx.scipy.sparse import csr_matrix as cp_csr_matrix
self.coupling_mat_np = csr_matrix(
(self._batch_vec.astype(np.float32),
(self._batch_rows, self._batch_cols)),
copy=True)
self.coupling_mat = cp_csr_matrix(self.coupling_mat_np, copy=True)
self._batch_vec = self.coupling_mat.data
del self.coupling_mat_np
else:
self.coupling_mat = csr_matrix(
(self._batch_vec, (self._batch_rows, self._batch_cols)),
copy=True)
# create an index to sort by rows and then columns,
# which is the same ordering CSR has internally
# lexsort sorts by last argument first!!!
self.sortidx = np.lexsort((self._batch_cols, self._batch_rows))
self._batch_rows = self._batch_rows[self.sortidx]
self._batch_cols = self._batch_cols[self.sortidx]
# Reorder batch matrix according to order in coupling_mat
if using_cupy:
self._batch_matrix = cupy.array(
self._batch_matrix[self.sortidx, :], dtype=np.float32)
else:
self._batch_matrix = self._batch_matrix[self.sortidx, :]
def _update_rates(self, z, force_update=False):
"""Batch compute all nonel and inclusive rates if z changes.
The result is always stored in the same vectors, since '_init_rate_matstruc'
makes use of views to link ranges of the vector to locations in the matrix.
Args:
z (float): Redshift value at which the photon field is taken.
Returns:
(bool): True if fields we indeed updated, False if nothing happened.
"""
if self._ratemat_zcache != z or force_update:
info(5, 'Updating batch rate vectors.')
if using_cupy:
if isinstance(self._batch_matrix, np.ndarray):
self._init_coupling_mat()
cupy.dot(self._batch_matrix,
cupy.array(self.photon_vector(z), dtype=np.float32),
out=self._batch_vec)
else:
np.dot(self._batch_matrix,
self.photon_vector(z),
out=self._batch_vec)
self._ratemat_zcache = z
return True
else:
return False
def get_particle_channels(mo, mo_energy, da_energy):
"""
Loops over a all daughers for a given mother and generates
a list of redistribution matrices on the grid:
np.outer( da_energy , 1 / mo_energy )
Args:
mo (int): id of the mother particle
mo_energy (float): energy grid of the mother particle
da_energy (float): energy grid of the daughter particle (same for all daughters)
Returns:
list of np.array: list of redistribution functions on on xgrid
"""
info(10, 'Generating decay redistribution for', mo)
dbentry = spec_data[mo]
x_grid = np.outer(da_energy, (1 / mo_energy))
redist = {}
for branching, daughters in dbentry['branchings']:
for da in daughters:
# daughter is a nucleus, we have lorentz factor conservation
if da > 99:
res = np.zeros(x_grid.shape)
res[x_grid == 1.] = 1.
else:
res = get_decay_matrix(mo, da, x_grid)
redist[da] = branching * res
return x_grid, redist
def __init__(self, prince_run=None, with_dense_jac=True, *args, **kwargs):
info(3, 'creating instance')
self.with_dense_jac = with_dense_jac
#: Reference to PhotonField object
self.photon_field = prince_run.photon_field
#: Reference to CrossSection object
self.cross_sections = prince_run.cross_sections
#: Reference to species manager
self.spec_man = prince_run.spec_man
# Initialize grids
self.e_photon = prince_run.ph_grid
self.e_cosmicray = prince_run.cr_grid
# Initialize cache of redshift value
self._ph_vec_zcache = None
self._ratemat_zcache = None
# Initialize the matrices for batch computation
self._batch_rows = None
self._batch_cols = None
self._batch_matrix = None
self._batch_vec = None
self.coupling_mat = None
self.dense_coupling_mat = None
self._estimate_batch_matrix()
self._init_matrices()
self._init_coupling_mat()
def _arange_on_xgrid(self, incl_cs):
"""Returns the inclusive cross section on an xgrid at x=1."""
egr, cs = None, None
if isinstance(incl_cs, tuple):
egr, cs = incl_cs
else:
cs = incl_cs
nxbins = len(self.xbins) - 1
if len(cs.shape) > 1 and cs.shape[0] != nxbins:
raise Exception(
'One dimensional cross section expected, instead got',
cs.shape, '\n', cs)
elif len(cs.shape) == 2 and cs.shape[0] == nxbins:
info(20, 'Supplied 2D distribution seems to be distributed in x.')
if isinstance(incl_cs, tuple):
return egr, cs
return cs
csec = np.zeros((nxbins, cs.shape[0]))
# NOTE: The factor 2 in the following line is a workarround to account for the latter linear interpolation
# This is needed because linear spline integral will result in a trapz,
# which has half the area of the actual first bin
corr_factor = 2 * self.xwidths[-1] / (self.xcenters[-1] -
self.xcenters[-2])
csec[-1, :] = cs / self.xwidths[-1] * corr_factor
info(
4,
'Warning! Workaround to account for linear interpolation in x, factor 2 added!'
)
if isinstance(incl_cs, tuple):
return egr, csec
return csec
def _estimate_batch_matrix(self):
''' estimate dimension of the batch matrix'''
dcr = self.e_cosmicray.d
dph = self.e_photon.d
batch_dim = 0
for specid in self.spec_man.known_species:
if specid < 100:
continue
# Add the main diagonal self-couplings (absoption)
batch_dim += dcr
for rtup in self.cross_sections.reactions[specid]:
# Off main diagonal couplings (reinjection)
if rtup in self.cross_sections.known_bc_channels:
batch_dim += dcr
elif rtup in self.cross_sections.known_diff_channels:
# Only half of the elements can be non-zero (energy conservation)
batch_dim += int(dcr**2 / 2) + 1
info(2, 'Batch matrix dimensions are {0}x{1}'.format(batch_dim, dph))
self._batch_matrix = np.zeros((batch_dim, dph))
self._batch_rows = []
self._batch_cols = []
info(
3,
'Memory usage: {0} MB'.format(self._batch_matrix.nbytes / 1024**2))
def nu_from_beta_decay_old(x_grid, mother, daughter):
"""
Energy distribution of a neutrinos from beta-decays of mother to daughter
Args:
x (float): energy fraction transferred to the secondary
mother (int): id of mother
daughter (int): id of daughter
Returns:
float: probability density at x
"""
info(10, 'Calculating neutrino energy from beta decay', mother, daughter)
mass_el = spec_data[20]['mass']
mass_mo = spec_data[mother]['mass']
mass_da = spec_data[daughter]['mass']
Z_mo = spec_data[mother]['charge']
Z_da = spec_data[daughter]['charge']
print(mother, daughter)
if mother == 100 and daughter == 101:
# for this channel the masses are already nucleon masses
qval = mass_mo - mass_da - mass_el
elif Z_da == Z_mo + 1: # beta+ decay
qval = mass_mo - mass_da - 2 * mass_el
elif Z_da == Z_mo - 1: # beta- decay
qval = mass_mo - mass_da
else:
raise Exception('Not an allowed beta decay channel: {:} -> {:}'.format(
mother, daughter))
E0 = qval + mass_el
print('Qval', qval, 'E0', E0)
ye = mass_el / E0
y_grid = x_grid * mass_mo / 2 / E0
# norm factor, nomalizing the formula to 1
norm = 1. / 60. * (np.sqrt(1. - ye**2) * (2 - 9 * ye**2 - 8 * ye**4) +
15 * ye**4 * np.log(ye / (1 - np.sqrt(1 - ye**2))))
cond = y_grid < 1 - ye
print((1 - ye) * 2 * E0 / mass_mo)
yshort = y_grid[cond]
result = np.zeros(y_grid.shape)
# factor for substitution y -> x
subst = mass_mo / 2 / E0
# total formula
result[cond] = subst / norm * yshort**2 * (1 - yshort) * np.sqrt(
(1 - yshort)**2 - ye**2)
result[x_grid > 1] *= 0.
return result
def resp(self):
"""Return ResponseFunction corresponding to this cross section
Will only create the Response function once.
"""
if not hasattr(self, '_resp'):
info(2, 'First Call, creating instance of ResponseFunction now')
from .response import ResponseFunction
self._resp = ResponseFunction(self)
return self._resp
def __init__(self, prince_run, energy='grid', *args, **kwargs):
info(3, 'creating instance')
#: Reference to species manager
self.spec_man = prince_run.spec_man
# Initialize grids
self.e_cosmicray = prince_run.cr_grid
self.dim_states = prince_run.dim_states
self.dim_bins = prince_run.dim_bins
# Init adiabatic loss vector
self.energy_vector = self._init_energy_vec(energy)
def _update_jacobian(self, z):
info(5, 'Updating jacobian matrix at redshift', z)
# enable photohadronic losses, or use a zero matrix
if self.enable_photohad_losses:
self.jacobian = self.had_int_rates.get_hadr_jacobian(
z, self.dldz(z), force_update=True)
else:
self.jacobian = self.had_int_rates.get_hadr_jacobian(
z, self.dldz(z), force_update=True)
self.jacobian.data *= 0.
self.last_hadr_jac = None
def _precompute_interpolators(self):
"""Interpolate each response function and store interpolators.
Uses :func:`prince_cr.util.get_interp_object` as interpolator.
This might result in too many knots and can be subject to
future optimization.
"""
info(2, 'Computing interpolators for response functions')
info(5, 'Nonelastic response functions f(y)')
self.nonel_intp = {}
for mother in self.nonel_idcs:
self.nonel_intp[mother] = get_interp_object(
*self.get_channel(mother))
info(5, 'Inclusive (boost conserving) response functions g(y)')
self.incl_intp = {}
for mother, daughter in self.incl_idcs:
self.incl_intp[(mother, daughter)] = get_interp_object(
*self.get_channel(mother, daughter))
info(5, 'Inclusive (redistributed) response functions h(y)')
self.incl_diff_intp = {}
for mother, daughter in self.incl_diff_idcs:
ygr, rfunc = self.get_channel(mother, daughter)
self.incl_diff_intp[(mother, daughter)] = get_2Dinterp_object(
self.xcenters, ygr, rfunc, self.cross_section.xbins)
from scipy.integrate import cumtrapz
integral = cumtrapz(rfunc, ygr, axis=1,initial=0)
integral = cumtrapz(integral, self.xcenters, axis=0,initial=0)
self.incl_diff_intp_integral[(mother, daughter)] = get_2Dinterp_object(
self.xcenters, ygr, integral, self.cross_section.xbins)
def incl(self, mother, daughter):
r"""Returns inclusive cross section.
Inclusive cross section for daughter in photo-nuclear
interactions of `mother`.
Args:
mother (int): Mother nucleus(on)
daughter (int): Daughter nucleus(on)
Returns:
(numpy.array, numpy.array): self._egrid_tab (:math:`\epsilon_r`),
inclusive cross section in :math:`cm^{-2}`
"""
_, Z, N = get_AZN(mother)
if daughter <= 101:
# raise Exception('Boost conserving cross section called ' +
# 'for redistributed particle')
from scipy.integrate import trapz
_, cs_diff = self.incl_diff(mother, daughter)
cs_incl = trapz(cs_diff,
x=self.xcenters,
dx=bin_widths(self.xbins),
axis=0)
return self.egrid, cs_incl[self._range]
elif daughter >= 200 and daughter not in [mother - 101, mother - 100]:
info(10, 'mother, daughter', mother, daughter, 'out of range')
return self.egrid[[0, -1]], np.array([0., 0.])
if daughter in [mother - 101]:
cgrid = Z * self.cs_proton_grid
# created incl. diff. index for all particle created in p-gamma
for da in self.redist_proton:
self.incl_diff_idcs.append((mother, da))
return self.egrid, cgrid[self._range]
elif daughter in [mother - 100]:
cgrid = N * self.cs_neutron_grid
# created incl. diff. channel index for all particle created in n-gamma
for da in self.redist_neutron:
self.incl_diff_idcs.append((mother, da))
return self.egrid, cgrid[self._range]
else:
raise Exception(
'Channel {:} to {:} not allowed in this superposition model'.
format(mother, daughter))
def _join_nonel(self, mother):
"""Returns the non-elastic cross section of the joined models.
"""
info(5, 'Joining nonelastic cross sections for', mother)
egrid = []
nonel = []
for model in self.model_refs:
e, csec = model.nonel(mother)
egrid.append(e)
nonel.append(csec)
#return np.concatenate(nonel)
return np.concatenate(egrid), np.concatenate(nonel)
def _join_incl(self, mother, daughter):
"""Returns joined incl cross sections."""
info(5, 'Joining inclusive cross sections for channel',
(mother, daughter))
egrid = []
incl = []
for model in self.model_refs:
e, csec = model.incl(mother, daughter)
egrid.append(e)
incl.append(csec)
#print np.concatenate(egrid), np.concatenate(incl)
#print '---'*30
#return np.concatenate(incl)
return np.concatenate(egrid), np.concatenate(incl)
def _join_incl_diff(self, mother, daughter):
"""Returns joined incl diff cross sections.
The function assumes the same `x` bins for all models.
"""
info(5, 'Joining inclusive differential cross sections for channel',
(mother, daughter))
egrid = []
incl_diff = []
# Get an x grid from a model which supports it
for model in self.model_refs:
if model.supports_redistributions:
self.xbins = model.xbins
break
if self.xbins is None:
raise Exception('Redistributions requested but none of the ' +
'models supports it')
for model in self.model_refs:
egr, csec = None, None
if config.debug_level > 1:
if not np.allclose(self.xbins, model.xbins):
raise Exception('Unequal x bins. Aborting...',
self.xbins.shape, model.xbins)
if (mother, daughter) in model.incl_diff_idcs:
egr, csec = model.incl_diff(mother, daughter)
info(10, model.mname, mother, daughter, 'is differential.')
elif (mother, daughter) in model.incl_idcs:
# try to use incl and extend by zeros
egr, csec_1d = model.incl(mother, daughter)
print(mother, daughter, csec_1d.shape)
# no x-distribution given, so x = 1
csec = self._arange_on_xgrid(csec_1d)
info(1, model.mname, mother, daughter,
'not differential, x=1.')
else:
info(
5, 'Model', model.mname, 'does not provide cross',
'sections for channel {0}/{1}. Setting to zero.'.format(
mother, daughter))
# Tried with reduced energy grids to save memory, but
# matrix addition in decay chains becomes untrasparent
# egr = np.array((model.egrid[0], model.egrid[-1]))
# csec = np.zeros((len(self.xbins) - 1, 2))
egr = model.egrid
csec = np.zeros((len(self.xbins) - 1, model.egrid.size))
egrid.append(egr)
incl_diff.append(csec)
#return np.concatenate(incl_diff, axis=1)
return np.concatenate(egrid), np.concatenate(incl_diff, axis=1)
def _compute_injection_grid(self):
"""Precompute the injection for all species on a single grid.
Assumes that the injection is factorized in E and z"""
self.injection_grid = np.zeros(self.dim_states)
for pid in self.ncoids:
if pid in self.params:
params = self.params[pid]
else:
params = params
info(
4, 'Injecting particle {:} with parameters {:}'.format(
pid, params))
inj_spec = self.spec_man.ncoid2sref[pid]
self.injection_grid[inj_spec.sl] = self.injection_spectrum(
pid, self.cr_grid, params)
def set_range(self, e_min=None, e_max=None):
"""Set energy range within which to return tabulated data.
Args:
e_min (float): minimal energy in GeV
e_max (float): maximal energy in GeV
"""
if e_min is None:
e_min = np.min(self._egrid_tab)
if e_max is None:
e_max = np.max(self._egrid_tab)
info(5, "Setting range to {0:3.2e} - {1:3.2e}".format(e_min, e_max))
self._range = np.where((self._egrid_tab >= e_min)
& (self._egrid_tab <= e_max))[0]
info(
2, "Range set to {0:3.2e} - {1:3.2e}".format(
np.min(self._egrid_tab[self._range]),
np.max(self._egrid_tab[self._range])))
def _load(self):
from prince_cr.data import db_handler
info(2, "Load tabulated cross sections")
photo_nuclear_tables = db_handler.photo_meson_db('SOPHIA')
info(2, "Loading SOPHIA cross sections from file.")
egrid = photo_nuclear_tables["energy_grid"]
xbins = photo_nuclear_tables["xbins"]
info(2, "Egrid loading finished")
# Integer idices of mothers and inclusive channels are stored
# in first column(s)
pid_nonel = photo_nuclear_tables["inel_mothers"]
pids_incl = photo_nuclear_tables["mothers_daughters"]
# the rest of the line denotes the crosssection on the egrid in mbarn,
# which is converted here to cm^2
nonel_raw = photo_nuclear_tables["inelastic_cross_sctions"]
incl_raw = photo_nuclear_tables["fragment_yields"]
info(2, "Data file loading finished")
self._egrid_tab = egrid
self.cs_proton_grid = nonel_raw[pid_nonel == 101].flatten()
self.cs_neutron_grid = nonel_raw[pid_nonel == 100].flatten()
self.xbins = xbins
self.redist_proton = {}
self.redist_neutron = {}
for (mo, da), csgrid in zip(pids_incl, incl_raw):
if mo == 101:
self.redist_proton[da] = csgrid
elif mo == 100:
self.redist_neutron[da] = csgrid
else:
raise Exception(
f'Sophia model should only contain protons and neutrons, but has mother id {mo}'
)
# set up inclusive differential channels for protons and neutron
# The model can return both, integrated over x and redistributed.
for da in sorted(self.redist_proton):
self.incl_diff_idcs.append((101, da))
#self.incl_idcs.append((101, da))
for da in sorted(self.redist_neutron):
self.incl_diff_idcs.append((100, da))
#self.incl_idcs.append((100, da))
# For more convenient generation of trivial redistribution matrices when joining
self.redist_shape = (self.xbins.shape[0], self._egrid_tab.shape[0])
self.set_range()
def convolve_with_decay_distribution(diff_dist, mother, daughter,
branching_ratio):
r"""Computes the prompt decay xdist by convolving the x distribution
of the unstable particle with the decay product distribution.
:math:`\frac{{\rm d}N^{A\gamma \to \mu}}{{\rm d}x_j} =
\sum_{i=0}^{N_x}~\Delta x_i
\frac{{\rm d}N^{A\gamma \to \pi}}{{\rm d} x_i}~
\frac{{\rm d}N^{\pi \to \mu}}{{\rm d} x_j}`
"""
# dec_dist = int_scale * decs.get_decay_matrix(
# mother, daughter, dec_grid)
dec_dist = decay_cached(mother,daughter)
info(20, 'convolving with decay dist', mother, daughter)
# Handle the case where table entry is (energy_grid, matrix)
if not isinstance(diff_dist, tuple):
return branching_ratio * dec_dist.dot(diff_dist)
else:
return diff_dist[0], branching_ratio * dec_dist.dot(
diff_dist[1])
def is_differential(self, mother, daughter):
"""Returns true if the model supports redistributions and requested
mother/daughter combination should return non-zero redistribution matrices.
Args:
mother (bool): Neucosma ID of mother particle
daughter (bool): Neucosma ID of daughter particle
Returns:
(bool): ``True`` if the model has this particular redistribution function
"""
# info(10, mother, daughter, " asking for redist")
# if not self.supports_redistributions:
# info(10, mother, daughter, " model doesn't support redist")
# return False
if (daughter <= config.redist_threshold_ID
or (mother, daughter) in self.incl_diff_idcs):
info(60, 'Daughter requires redistribution.', mother, daughter)
return True
info(60, 'Daughter conserves boost.', mother, daughter)
return False
def __init__(self, prince_run, energy='grid', *args, **kwargs):
info(3, 'creating instance')
#: Reference to species manager
self.spec_man = prince_run.spec_man
#: Reference to PhotonField object
self.photon_field = prince_run.photon_field
# Initialize grids
self.e_cosmicray = prince_run.cr_grid
self.dim_states = prince_run.dim_states
self.dim_bins = prince_run.dim_bins
self.e_photon = prince_run.ph_grid
# xi is dimensionless (natural units) variable
xi_steps = 400 if 'xi_steps' not in kwargs else kwargs['xi_steps']
info(2, 'using', xi_steps, 'steps in xi')
self.xi = np.logspace(np.log10(2 + 1e-8), 16., xi_steps)
# weights for integration
self.phi_xi2 = self._phi(self.xi) / (self.xi**2)
# Scale vector containing the units and factors of Z**2 for nuclei
self.scale_vec = self._init_scale_vec(energy)
# Gamma factor of the cosmic ray
if energy == 'grid':
gamma = self.e_cosmicray.grid / PRINCE_UNITS.m_proton
elif energy == 'bins':
gamma = self.e_cosmicray.bins / PRINCE_UNITS.m_proton
else:
raise Exception(
'Unexpected energy keyword ({:}), use either (grid) or (bins)',
format(energy))
# Grid of photon energies for interpolation
self.photon_grid = np.outer(1 / gamma,
self.xi) * PRINCE_UNITS.m_electron / 2.
self.pg_desort = self.photon_grid.reshape(-1).argsort()
self.pg_sorted = self.photon_grid.reshape(-1)[self.pg_desort]
def _load(self, model_prefix):
from prince_cr.data import db_handler
info(2, "Load tabulated cross sections")
# The energy grid is given in MeV, so we convert to GeV
photo_nuclear_tables = db_handler.photo_nuclear_db(model_prefix)
egrid = photo_nuclear_tables["energy_grid"]
info(2, "Egrid loading finished")
# Integer idices of mothers and inclusive channels are stored
# in first column(s)
pid_nonel = photo_nuclear_tables["inel_mothers"]
pids_incl = photo_nuclear_tables["mothers_daughters"]
# the rest of the line denotes the crosssection on the egrid in mbarn,
# which is converted here to cm^2
nonel_raw = photo_nuclear_tables["inelastic_cross_sctions"]
incl_raw = photo_nuclear_tables["fragment_yields"]
info(2, "Data file loading finished")
# Now write the raw data into a dict structure
_nonel_tab = {}
for pid, csgrid in zip(pid_nonel, nonel_raw):
if get_AZN(pid)[0] > config.max_mass:
continue
_nonel_tab[pid] = csgrid
# If proton and neutron cross sections are not in contained
# in the files, set them to 0. Needed for TALYS and CRPropa2
for pid in [101, 100]:
if pid not in _nonel_tab:
_nonel_tab[pid] = np.zeros_like(egrid)
# mo = mother, da = daughter
_incl_tab = {}
for (mo, da), csgrid in zip(pids_incl, incl_raw):
if get_AZN(mo)[0] > config.max_mass:
continue
_incl_tab[mo, da] = csgrid
self._egrid_tab = egrid
self._nonel_tab = _nonel_tab
self._incl_tab = _incl_tab
# Set initial range to whole egrid
self.set_range()
info(2, "Finished initialization")
def get_full(self, mother, daughter, ygrid, xgrid=None):
"""Return the full response function :math:`f(y) + g(y) + h(x,y)`
on the grid that is provided. xgrid is ignored if `h(x,y)` not in the channel.
"""
if xgrid is not None and ygrid.shape != xgrid.shape:
raise Exception('ygrid and xgrid do not have the same shape!!')
if get_AZN(mother)[0] < get_AZN(daughter)[0]:
info(
3,
'WARNING: channel {:} -> {:} with daughter heavier than mother!'
.format(mother, daughter))
res = np.zeros(ygrid.shape)
if (mother, daughter) in self.incl_intp:
res += self.incl_intp[(mother, daughter)](ygrid)
elif (mother, daughter) in self.incl_diff_intp:
#incl_diff_res = self.incl_diff_intp[(mother, daughter)](
# xgrid, ygrid, grid=False)
#if mother == 101:
# incl_diff_res = np.where(xgrid < 0.9, incl_diff_res, 0.)
#res += incl_diff_res
#if not(mother == daughter):
res += self.incl_diff_intp[(mother, daughter)].inteval(xgrid,
ygrid,
grid=False)
if mother == daughter and mother in self.nonel_intp:
# nonel cross section leads to absorption, therefore the minus
if xgrid is None:
res -= self.nonel_intp[mother](ygrid)
else:
diagonal = xgrid == 1.
res[diagonal] -= self.nonel_intp[mother](ygrid[diagonal])
return res
def follow_chain(first_mo, da, csection, reclev):
"""Recursive function to follow decay chains until all
final state particles are stable.
The result is saved in two dictionaries; one for the boost
conserving inclusive channels and the other one collects
channels with meson or lepton decay products, which will
need special care due to energy redistributions of these
secondaries.
"""
info(10, dbg_indent(reclev), 'Entering with', first_mo, da)
if da not in spec_data:
info(
3, dbg_indent(reclev),
'daughter {0} unknown, forcing beta decay. Not Implemented yet!!'
.format(da))
return
# Daughter is stable. Add it to the new dictionary and terminate
# recursion
if spec_data[da]["lifetime"] >= threshold:
if self.is_differential(None, da):
# If the daughter is a meson or lepton, use the dictionary for
# differential channels
info(
20, dbg_indent(reclev),
'daughter {0} stable and differential. Adding to ({1}, {2})'
.format(da, first_mo, da))
new_dec_diff_tab[(first_mo, da)] = csection
else:
info(
20, dbg_indent(reclev),
'daughter {0} stable. Adding to ({1}, {2})'.format(
da, first_mo, da))
new_incl_tab[(first_mo, da)] = csection
return
# ..otherwise follow decay products of this daughter, tracking the
# original mother particle (first_mo). The cross section (csection) is
# reduced by the branching ratio (br) of this particular channel
for br, daughters in spec_data[da]["branchings"]:
info(10, dbg_indent(reclev),
("{3} -> {0:4d} -> {2:4.2f}: {1}").format(
da, ", ".join(map(str, daughters)), br, first_mo))
for chained_daughter in daughters:
# Follow each secondary and increment the recursion level by one
if self.is_differential(None, chained_daughter):
info(10, 'daughter', chained_daughter, 'of', da,
'is differential')
follow_chain(
first_mo, chained_daughter,
convolve_with_decay_distribution(
self._arange_on_xgrid(csection), da,
chained_daughter, br), reclev + 1)
else:
follow_chain(first_mo, chained_daughter, br * csection,
reclev + 1)
def _reduce_channels(self):
"""Follows decay chains until all inclusive reactions point to
stable final state particles.
The "tau_dec_threshold" parameter in the config controls the
definition of stable. Unstable nuclei for which no decay channels
are known, will be forced to beta-decay until they reach a stable
element.
"""
from prince_cr.util import AdditiveDictionary
# TODO: check routine, how to avoid empty channels and
# mothers with zero nonel cross sections
# The new dictionary that will replace _incl_tab
new_incl_tab = AdditiveDictionary()
new_dec_diff_tab = AdditiveDictionary()
threshold = config.tau_dec_threshold
# How to indent debug printout for recursion
dbg_indent = lambda lev: 4 * lev * "-" + ">" if lev else ""
info(2, "Integrating out species with lifetime smaller than",
threshold)
info(3, (
"Before optimization, the number of known primaries is {0} with " +
"in total {1} inclusive channels").format(len(self._nonel_tab),
len(self._incl_tab)))
if self.xbins is None:
info(
4,
'Model does not provide a native xbins. Assuming JH special sophia',
'binning.')
from .photo_meson import SophiaSuperposition
self.xbins = SophiaSuperposition().xbins
bc = self.xcenters
bw = bin_widths(self.xbins)
# The x_mu/x_pi grid
# dec_grid = np.fromfunction(
# lambda j, i: 10**(np.log10(bc[1] / bc[0]) * (j - i)), (len(bc),
# len(bc)))
# dec_grid = np.outer(bc, 1 / bc)
dec_bins = np.outer(self.xbins, 1 / bc)
dec_bins_lower = dec_bins[:-1]
dec_bins_upper = dec_bins[1:]
# dec_grid[dec_grid > 1.] *= 0.
# The differential element dx_mu/x_pi
int_scale = np.tile(bw / bc, (len(bc), 1))
from functools import lru_cache
@lru_cache(maxsize=512, typed=False)
def decay_cached(mother,daughter):
dec_dist = int_scale * decs.get_decay_matrix_bin_average(
mother, daughter, dec_bins_lower, dec_bins_upper)
return dec_dist
def convolve_with_decay_distribution(diff_dist, mother, daughter,
branching_ratio):
r"""Computes the prompt decay xdist by convolving the x distribution
of the unstable particle with the decay product distribution.
:math:`\frac{{\rm d}N^{A\gamma \to \mu}}{{\rm d}x_j} =
\sum_{i=0}^{N_x}~\Delta x_i
\frac{{\rm d}N^{A\gamma \to \pi}}{{\rm d} x_i}~
\frac{{\rm d}N^{\pi \to \mu}}{{\rm d} x_j}`
"""
# dec_dist = int_scale * decs.get_decay_matrix(
# mother, daughter, dec_grid)
dec_dist = decay_cached(mother,daughter)
info(20, 'convolving with decay dist', mother, daughter)
# Handle the case where table entry is (energy_grid, matrix)
if not isinstance(diff_dist, tuple):
return branching_ratio * dec_dist.dot(diff_dist)
else:
return diff_dist[0], branching_ratio * dec_dist.dot(
diff_dist[1])
def follow_chain(first_mo, da, csection, reclev):
"""Recursive function to follow decay chains until all
final state particles are stable.
The result is saved in two dictionaries; one for the boost
conserving inclusive channels and the other one collects
channels with meson or lepton decay products, which will
need special care due to energy redistributions of these
secondaries.
"""
info(10, dbg_indent(reclev), 'Entering with', first_mo, da)
if da not in spec_data:
info(
3, dbg_indent(reclev),
'daughter {0} unknown, forcing beta decay. Not Implemented yet!!'
.format(da))
return
# Daughter is stable. Add it to the new dictionary and terminate
# recursion
if spec_data[da]["lifetime"] >= threshold:
if self.is_differential(None, da):
# If the daughter is a meson or lepton, use the dictionary for
# differential channels
info(
20, dbg_indent(reclev),
'daughter {0} stable and differential. Adding to ({1}, {2})'
.format(da, first_mo, da))
new_dec_diff_tab[(first_mo, da)] = csection
else:
info(
20, dbg_indent(reclev),
'daughter {0} stable. Adding to ({1}, {2})'.format(
da, first_mo, da))
new_incl_tab[(first_mo, da)] = csection
return
# ..otherwise follow decay products of this daughter, tracking the
# original mother particle (first_mo). The cross section (csection) is
# reduced by the branching ratio (br) of this particular channel
for br, daughters in spec_data[da]["branchings"]:
info(10, dbg_indent(reclev),
("{3} -> {0:4d} -> {2:4.2f}: {1}").format(
da, ", ".join(map(str, daughters)), br, first_mo))
for chained_daughter in daughters:
# Follow each secondary and increment the recursion level by one
if self.is_differential(None, chained_daughter):
info(10, 'daughter', chained_daughter, 'of', da,
'is differential')
follow_chain(
first_mo, chained_daughter,
convolve_with_decay_distribution(
self._arange_on_xgrid(csection), da,
chained_daughter, br), reclev + 1)
else:
follow_chain(first_mo, chained_daughter, br * csection,
reclev + 1)
# Remove all unstable particles from the dictionaries
for mother in sorted(self._nonel_tab.keys()):
if mother not in spec_data or spec_data[mother][
"lifetime"] < threshold:
info(
20,
"Primary species {0} does not fulfill stability criteria.".
format(mother))
_ = self._nonel_tab.pop(mother)
# Only stable (interacting) mother particles are left
self._update_indices()
for (mother, daughter) in self.incl_idcs:
if mother not in self.nonel_idcs:
info(
30, "Removing {0}/{1} from incl, since mother not stable ".
format(mother, daughter))
_ = self._incl_tab.pop((mother, daughter))
elif self.is_differential(mother, daughter):
# Move the distributions which are expected to be differential
# to _incl_diff_tab
self._incl_diff_tab[(mother,
daughter)] = self._arange_on_xgrid(
self._incl_tab.pop(
(mother, daughter)))
self._update_indices()
for (mother, daughter) in self.incl_diff_idcs:
if mother not in self.nonel_idcs:
info(
30,
"Removing {0}/{1} from diff incl, since mother not stable "
.format(mother, daughter))
_ = self._incl_diff_tab.pop((mother, daughter))
self._update_indices()
# Launch the reduction for each inclusive channel
for (mo, da), value in list(self._incl_tab.items()):
#print mo, da, value
#print '---'*30
follow_chain(mo, da, value, 0)
for (mo, da), value in list(self._incl_diff_tab.items()):
#print mo, da, value
#print '---'*30
follow_chain(mo, da, value, 0)
# Overwrite the old incl dictionary
self._incl_tab = dict(new_incl_tab)
# Overwrite the old incl_diff dictionary
self._incl_diff_tab = dict(new_dec_diff_tab)
# Reduce also the incl_diff_tab by removing the unknown mothers. At this stage
# of the code, the particles with redistributions are
info(
3,
("After optimization, the number of known primaries is {0} with " +
"in total {1} inclusive channels").format(
len(self._nonel_tab),
len(self._incl_tab) + len(self._incl_diff_tab)))
info(2, f'Cache used for decays, {decay_cached.cache_info()}') # pylint:disable=no-value-for-parameter
def _optimize_and_generate_index(self):
"""Construct a list of mothers and (mother, daughter) indices.
Args:
just_reactions (bool): If True then fill just the reactions index.
"""
# Integrate out short lived processes and leave only stable particles
# in the databases
self._reduce_channels()
# Go through all three cross section categories
# index contents in the ..known..variable
self.reactions = {}
self._update_indices()
for mo, da in self.incl_idcs:
if da >= 100 and get_AZN(da)[0] > get_AZN(mo)[0]:
raise Exception(
'Daughter {0} heavier than mother {1}. Physics??'.format(
da, mo))
if mo not in self.reactions:
self.reactions[mo] = []
self.known_species.append(mo)
if (mo, da) not in self.reactions[mo]:
# Make sure it's a unique list
self.reactions[mo].append((mo, da))
if self.is_differential(mo, da):
# Move the distributions which are expected to be differential
# to _incl_diff_tab
self._incl_diff_tab[(mo, da)] = self._arange_on_xgrid(
self._incl_tab.pop((mo, da)))
info(10, "Channel {0} -> {1} forced to be differential.")
else:
self.known_bc_channels.append((mo, da))
self.known_species.append(da)
for mo, da in list(self._incl_diff_tab.keys()):
if da >= 100 and get_AZN(da)[0] > get_AZN(mo)[0]:
raise Exception(
'Daughter {0} heavier than mother {1}. Physics??'.format(
da, mo))
if mo not in self.reactions:
self.reactions[mo] = []
self.known_species.append(mo)
if (mo, da) not in self.reactions[mo]:
# Make sure it's a unique list to avoid unnecessary loops
self.reactions[mo].append((mo, da))
self.known_diff_channels.append((mo, da))
self.known_species.append(da)
# Remove duplicates
self.known_species = sorted(list(set(self.known_species)))
self.known_bc_channels = sorted(list(set(self.known_bc_channels)))
self.known_diff_channels = sorted(list(set(self.known_diff_channels)))
for sp in self.known_species:
if sp >= 100 and (sp, sp) not in self.known_diff_channels:
self.known_bc_channels.append((mo, mo))
if (mo, mo) not in self.reactions[mo]:
self.reactions[mo].append((mo, mo))
# Make sure the indices are up to date
self._update_indices()
def nu_from_beta_decay(x_grid, mother, daughter, Gamma=200, angle=None):
"""
Energy distribution of a neutrinos from beta-decays of mother to daughter
The res frame distrution is boosted to the observers frame and then angular averaging is done numerically
Args:
x_grid (float): energy fraction transferred to the secondary
mother (int): id of mother
daughter (int): id of daughter
Gamma (float): Lorentz factor of the parent particle, default: 200
For large Gamma this should not play a role, as the decay is scale invariant
angle (float): collision angle, if None this will be averaged over 2 pi
Returns:
float: probability density on x_grid
"""
import warnings
info(10, 'Calculating neutrino energy from beta decay', mother, daughter)
mass_el = spec_data[20]['mass']
mass_mo = spec_data[mother]['mass']
mass_da = spec_data[daughter]['mass']
Z_mo = spec_data[mother]['charge']
Z_da = spec_data[daughter]['charge']
A_mo, _, _ = get_AZN(mother)
if mother == 100 and daughter == 101:
# for this channel the masses are already nucleon masses
qval = mass_mo - mass_da - mass_el
elif Z_da == Z_mo - 1: # beta+ decay
qval = mass_mo - mass_da - 2 * mass_el
elif Z_da == Z_mo + 1: # beta- decay
qval = mass_mo - mass_da
else:
raise Exception('Not an allowed beta decay channel: {:} -> {:}'.format(
mother, daughter))
# substitute this to the energy grid
E0 = qval + mass_el
# NOTE: we subsitute into energy per nucleon here
Emo = Gamma * mass_mo / A_mo
E = x_grid * Emo
# print '------','beta decay','------'
# print mother
# print E0
# print A_mo
# print Emo
if angle is None:
# ctheta = np.linspace(-1, 1, 1000)
# we use here logspace, as high resolution is mainly needed at small energies
# otherwise the solution will oscillate at low energy
ctheta = np.unique(
np.concatenate((
np.logspace(-8, 0, 1000) - 1,
1 - np.logspace(0, -8, 1000),
)))
else:
ctheta = angle
boost = Gamma * (1 - ctheta)
Emax = E0 * boost
E_mesh, boost_mesh = np.meshgrid(E, boost, indexing='ij')
with warnings.catch_warnings():
warnings.simplefilter("ignore")
res = E_mesh**2 / boost_mesh**5 * (Emax - E_mesh) * np.sqrt(
(E_mesh - Emax)**2 - boost_mesh**2 * mass_el**2)
res[E_mesh > Emax] = 0.
res = np.nan_to_num(res)
if np.all(res == 0):
info(10, 'Differential distribution is all zeros for', mother,
daughter, 'No angle averaging performed!')
elif angle is None:
# now average over angle
res = trapz(res, x=ctheta, axis=1)
res = res / trapz(res, x=x_grid)
else:
res = res[:, 0]
res = res / trapz(res, x=x_grid)
return res
def get_decay_matrix(mo, da, x_grid):
"""
Selects the correct redistribution for the given decay channel.
If the channel is unknown a zero grid is returned instead of raising an error
Args:
mo (int): index of the mother
da (int): index of the daughter
x_grid (float): grid in x = E_da / E_mo on which to return the result
(If x is a 2D matrix only the last column is computed
and then repeated over the matrix assuming that the
main diagonal is always x = 1)
Returns:
float: redistribution on the grid mo_energy / da_energy
"""
info(10, 'Generating decay redistribution for', mo, da)
# --------------------------------
# pi+ to numu or pi- to nummubar
# --------------------------------
if mo in [2, 3] and da in [13, 14]:
return pion_to_numu(x_grid)
# --------------------------------
# pi+ to mu+ or pi- to mu-
# --------------------------------
elif mo in [2, 3] and da in [5, 6, 7, 8, 9, 10]:
# (any helicity)
if da in [7, 10]:
return pion_to_muon(x_grid)
# left handed, hel = -1
elif da in [5, 8]:
return pion_to_muon(x_grid) * prob_muon_hel(x_grid, -1.)
# right handed, hel = 1
elif da in [6, 9]:
return pion_to_muon(x_grid) * prob_muon_hel(x_grid, 1.)
else:
raise Exception(
'This should newer have happened, check if-statements above!')
# --------------------------------
# muon to neutrino
# --------------------------------
elif mo in [5, 6, 7, 8, 9, 10] and da in [11, 12, 13, 14]:
# translating muon ids to helicity
muon_hel = {
5: 1.,
6: -1.,
7: 0.,
8: 1.,
9: -1.,
10: 0.,
}
hel = muon_hel[mo]
# muon+ to electron neutrino
if mo in [5, 6, 7] and da in [11]:
return muonplus_to_nue(x_grid, hel)
# muon+ to muon anti-neutrino
elif mo in [5, 6, 7] and da in [14]:
return muonplus_to_numubar(x_grid, hel)
# muon- to elec anti-neutrino
elif mo in [8, 9, 10] and da in [12]:
return muonplus_to_nue(x_grid, -1 * hel)
# muon- to muon neutrino
elif mo in [8, 9, 10] and da in [13]:
return muonplus_to_numubar(x_grid, -1 * hel)
# --------------------------------
# neutrinos from beta decays
# --------------------------------
# beta-
elif mo > 99 and da == 11:
info(10, 'nu_e from beta- decay', mo, mo - 1, da)
return nu_from_beta_decay(x_grid, mo, mo - 1)
# beta+
elif mo > 99 and da == 12:
info(10, 'nubar_e from beta+ decay', mo, mo + 1, da)
return nu_from_beta_decay(x_grid, mo, mo + 1)
# neutron
elif mo > 99 and 99 < da < 200:
info(10, 'beta decay boost conservation', mo, da)
return boost_conservation(x_grid)
else:
info(
5,
'Called with unknown channel {:} to {:}, returning an empty redistribution'
.format(mo, da))
# no known channel, return zeros
return np.zeros(x_grid.shape)
def get_decay_matrix_bin_average(mo, da, x_lower, x_upper):
"""
Selects the correct redistribution for the given decay channel.
If the channel is unknown a zero grid is returned instead of raising an error
Args:
mo (int): index of the mother
da (int): index of the daughter
x_grid (float): grid in x = E_da / E_mo on which to return the result
Returns:
float: redistribution on the grid mo_energy / da_energy
"""
# TODO: Some of the distribution are not averaged yet.
# The error is small for smooth distributions though
info(10, 'Generating decay redistribution for', mo, da)
x_grid = (x_upper + x_lower) / 2
# remember shape, but only calculate for last column, as x repeats in each column
from scipy.integrate import trapz
shape = x_grid.shape
if len(shape) == 2:
x_grid = x_grid[:, -1]
x_upper = x_upper[:, -1]
x_lower = x_lower[:, -1]
# --------------------------------
# pi+ to numu or pi- to nummubar
# --------------------------------
if mo in [2, 3] and da in [13, 14]:
result = pion_to_numu_avg(x_lower, x_upper)
# --------------------------------
# pi+ to mu+ or pi- to mu-
# --------------------------------
# TODO: The helicity distr need to be averaged analyticaly
elif mo in [2, 3] and da in [5, 6, 7, 8, 9, 10]:
# (any helicity)
if da in [7, 10]:
result = pion_to_muon_avg(x_lower, x_upper)
# left handed, hel = -1
elif da in [5, 8]:
result = pion_to_muon_avg(x_lower, x_upper) * prob_muon_hel(
x_grid, -1.)
# right handed, hel = 1
elif da in [6, 9]:
result = pion_to_muon_avg(x_lower, x_upper) * prob_muon_hel(
x_grid, 1.)
else:
raise Exception(
'This should newer have happened, check if-statements above!')
# --------------------------------
# muon to neutrino
# --------------------------------
# TODO: The following distr need to be averaged analyticaly
elif mo in [5, 6, 7, 8, 9, 10] and da in [11, 12, 13, 14]:
# translating muon ids to helicity
muon_hel = {
5: 1.,
6: -1.,
7: 0.,
8: 1.,
9: -1.,
10: 0.,
}
hel = muon_hel[mo]
# muon+ to electron neutrino
if mo in [5, 6, 7] and da in [11]:
result = muonplus_to_nue(x_grid, hel)
# muon+ to muon anti-neutrino
elif mo in [5, 6, 7] and da in [14]:
result = muonplus_to_numubar(x_grid, hel)
# muon- to elec anti-neutrino
elif mo in [8, 9, 10] and da in [12]:
result = muonplus_to_nue(x_grid, -1 * hel)
# muon- to muon neutrino
elif mo in [8, 9, 10] and da in [13]:
result = muonplus_to_numubar(x_grid, -1 * hel)
# --------------------------------
# neutrinos from beta decays
# --------------------------------
# TODO: The following beta decay to neutrino distr need to be averaged analyticaly
# TODO: Also the angular averaging is done numerically still
# beta-
elif mo > 99 and da == 11:
info(10, 'nu_e from beta+ decay', mo, mo - 1, da)
result = nu_from_beta_decay(x_grid, mo, mo - 1)
# beta+
elif mo > 99 and da == 12:
info(10, 'nubar_e from beta- decay', mo, mo + 1, da)
result = nu_from_beta_decay(x_grid, mo, mo + 1)
# neutron
elif mo > 99 and 99 < da < 200:
info(10, 'beta decay boost conservation', mo, da)
result = boost_conservation_avg(x_lower, x_upper)
else:
info(
5,
'Called with unknown channel {:} to {:}, returning an empty redistribution'
.format(mo, da))
# no known channel, return zeros
result = np.zeros(x_grid.shape)
# now fill this into diagonals of matrix
if len(shape) == 2:
#'filling matrix'
res_mat = np.zeros(shape)
for idx, val in enumerate(result[::-1]):
np.fill_diagonal(res_mat[:, idx:], val)
result = res_mat
return result
def _init_matrices(self):
""" A new take on filling the matrices"""
# Define some short-cuts
known_species = self.spec_man.known_species[::-1]
sp_id_ref = self.spec_man.ncoid2sref
resp = self.cross_sections.resp
m_pr = PRINCE_UNITS.m_proton
# Energy variables
dcr = self.e_cosmicray.d
dph = self.e_photon.d
ecr = self.e_cosmicray.grid
bcr = self.e_cosmicray.bins
eph = self.e_photon.grid
bph = self.e_photon.bins
delta_ec = self.e_cosmicray.widths
delta_ph = self.e_photon.widths
# Edges of each CR energy bin and photon energy bin
elims = np.vstack([bcr[:-1], bcr[1:]])
plims = np.vstack([bph[:-1], bph[1:]])
# CR and photon grid indices
emo_idcs = np.arange(dcr)
eda_idcs = np.arange(dcr)
p_idcs = np.arange(dph)
# values for x and y to cut on:
x_cut = config.x_cut
y_cut = config.y_cut
x_cut_proton = config.x_cut_proton
ibatch = 0
import itertools
spec_iter = itertools.product(known_species, known_species)
for moid, daid in spec_iter:
if moid < 100:
continue
else:
info(10, f'Filling channel {moid} -> {daid}')
has_nonel = moid == daid
if has_nonel:
intp_nonel = resp.nonel_intp[moid].antiderivative()
if (((moid, daid) in self.cross_sections.known_bc_channels) or
(has_nonel and
(moid, daid) not in self.cross_sections.known_diff_channels)):
has_incl = (moid, daid) in resp.incl_intp
if has_incl:
intp_bc = resp.incl_intp[(moid, daid)].antiderivative()
else:
info(1, 'Inclusive interpolator not found for',
(moid, daid))
if not (has_nonel or has_incl):
raise Exception('Channel without interactions:',
(moid, daid))
# The cross sections need to be evaluated
# on x = E_{CR,da} / E_{CR,mo} and y = E_ph * E_{CR,mo} / m_proton
# To vectorize the evaluation, we create outer products using numpy
# broadcasting:
emo = ecr
xl = elims[0] / emo
xu = elims[1] / emo
delta_x = delta_ec / emo
yl = plims[0, None, :] * emo[:, None] / m_pr
yu = plims[1, None, :] * emo[:, None] / m_pr
delta_y = delta_ph[None, :] * emo[:, None] / m_pr
int_fac = (delta_ec[:, None] * delta_ph[None, :] /
emo[:, None])
diff_fac = 1. / delta_x[:, None] / delta_y
# This takes the average by evaluating the integral and dividing by bin
# width
if has_incl:
self._batch_matrix[ibatch:ibatch + len(emo), :] = (
intp_bc(yu) - intp_bc(yl)) * int_fac * diff_fac
if has_nonel:
self._batch_matrix[ibatch:ibatch + len(emo), :] -= (
intp_nonel(yu) - intp_nonel(yl)) * int_fac * diff_fac
# finally map this to the coupling matrix
ibatch += len(emo)
self._batch_rows.append(sp_id_ref[daid].lidx() + eda_idcs)
self._batch_cols.append(sp_id_ref[moid].lidx() + emo_idcs)
elif (moid, daid) in self.cross_sections.known_diff_channels:
has_redist = (moid, daid) in resp.incl_diff_intp
if has_redist:
intp_diff = resp.incl_diff_intp[(moid, daid)]
intp_diff_integral = resp.incl_diff_intp_integral[(moid,
daid)]
intp_nonel = resp.nonel_intp[moid]
intp_nonel_antid = resp.nonel_intp[moid].antiderivative()
ymin = np.min(intp_diff.get_knots()[1])
else:
raise Exception('This should not occur.')
ibatch_bf = ibatch
# generate outer products using broadcasting
emo = ecr[:, None, None]
eda = ecr[None, :, None]
epho = eph[None, None, :]
target_shape = np.ones_like(emo * eda * epho)
xl = elims[0, None, :, None] / emo * target_shape
xu = elims[1, None, :, None] / emo * target_shape
delta_x = delta_ec[None, :, None] / emo
yl = plims[0, None, None, :] * emo / m_pr * target_shape
yu = plims[1, None, None, :] * emo / m_pr * target_shape
delta_y = delta_ph[None, None, :] * emo / m_pr
int_fac = (delta_ec[:, None, None] * delta_ph[None, None, :] /
emo) * target_shape
diff_fac = 1. / delta_x / delta_y
# Generate boolean arrays to cut on xvalues
if daid == 101:
cuts = np.logical_and(xl >= x_cut_proton, xl <= 1)
else:
# or (yu < ymin) or (yl > y_cut)
cuts = np.logical_and(xl >= x_cut, xl <= 1)
cuts = cuts[:, :, 0]
# # NOTE JH: This is an old version, which brute force vectorizes the integral with numpy
# # I am leaving this in the comments, in case we want to go back for testing-
# integrator = np.vectorize(intp_diff.integral)
# res = integrator(xl[cuts], xu[cuts], yl[cuts], yu[cuts]) * diff_fac[cuts] * int_fac[cuts]
# This takes the average by evaluating the integral and dividing by bin width
# intp_diff_integral contains the antiderivate, to to get the integral (xl,yl,xu,yu)
# we need to substract INT = (0,0,xu,yu) - (0,0,xl,yu) - (0,0,xu,yl) +
# (0,0,xl,yl)
res = intp_diff_integral.ev(xu[cuts], yu[cuts])
res -= intp_diff_integral.ev(xl[cuts], yu[cuts])
res -= intp_diff_integral.ev(xu[cuts], yl[cuts])
res += intp_diff_integral.ev(xl[cuts], yl[cuts])
res *= diff_fac[cuts] * int_fac[cuts]
res[res < 0] = 0.
# Since we made cuts on x, we need to make the same cut on the index
# mapping
emoidx, edaidx, _ = np.meshgrid(
sp_id_ref[moid].lidx() + emo_idcs,
sp_id_ref[daid].lidx() + eda_idcs,
p_idcs,
indexing='ij')
emoidx, edaidx = emoidx[cuts], edaidx[cuts]
# Now add the nonel interactions on the main diagonal
if has_nonel:
res[emoidx == edaidx] -= (
intp_nonel_antid(yu[cuts][emoidx == edaidx]) -
intp_nonel_antid(yl[cuts][emoidx == edaidx])) * (
diff_fac[cuts][emoidx == edaidx] *
int_fac[cuts][emoidx == edaidx])
# Finally write this to the batch matrix
self._batch_matrix[ibatch:ibatch + len(emoidx), :] = res
self._batch_rows.append(edaidx[:, 0])
self._batch_cols.append(emoidx[:, 0])
ibatch += len(emoidx)
else:
info(20, 'Species combination not included in model', moid,
daid)
self._batch_matrix = self._batch_matrix[:ibatch, :]
self._batch_rows = np.concatenate(self._batch_rows, axis=None)
self._batch_cols = np.concatenate(self._batch_cols, axis=None)
self._batch_vec = np.zeros(ibatch)
info(2, f'Batch matrix shape: {self._batch_matrix.shape}')
info(2, f'Batch rows shape: {self._batch_rows.shape}')
info(2, f'Batch cols shape: {self._batch_cols.shape}')
info(2, f'Batch vector shape: {self._batch_vec.shape}')
memory = (self._batch_matrix.nbytes + self._batch_rows.nbytes +
self._batch_cols.nbytes + self._batch_vec.nbytes) / 1024**2
info(3, "Memory usage after initialization: {:} MB".format(memory))
