def shift_eng(self, new_eng):
""" Shifts the abscissa while conserving number.
This function can be used to subtract or add some amount of energy from each bin in the spectrum. The dN/dE is adjusted to conserve number in each bin.
Parameters
----------
new_eng : ndarray
The new energy abscissa.
Returns
-------
None
"""
if new_eng.size != self.eng.size:
raise TypeError(
"The new abscissa must have the same length as the old one.")
if not all(np.diff(new_eng) > 0):
raise TypeError("abscissa must be ordered in increasing energy.")
new_log_bin_width = get_log_bin_width(new_eng)
if self._spec_type == 'dNdE':
new_dNdE = self.totN('bin') / (new_eng * new_log_bin_width)
self.eng = new_eng
self._data = new_dNdE
elif self._spec_type == 'N':
self.eng = new_eng
def switch_spec_type(self, target=None):
"""Switches between data being stored as N or dN/dE.
Parameters
----------
target : {'N', 'dNdE'}, optional
The target type to switch to. If not specified, performs a switch regardless.
Notes
------
Although both N and dN/dE can be accessed regardless of which values
are stored, performing a switch before repeated computations can
speed up the computation.
"""
if target is not None:
if target != 'N' and target != 'dNdE':
raise ValueError('Invalid target specified.')
log_bin_width = get_log_bin_width(self.eng)
if self._spec_type == 'N' and not target == 'N':
self._data = self._data / (self.eng * log_bin_width)
self._spec_type = 'dNdE'
elif self._spec_type == 'dNdE' and not target == 'dNdE':
self._data = self._data * self.eng * log_bin_width
self._spec_type = 'N'
def switch_spec_type(self, target=None):
"""Switches between the type of values to be stored.
Parameters
----------
target : {'N', 'dNdE'}
The target type to switch to.
"""
log_bin_width = get_log_bin_width(self.eng)
if self.spec_type == 'N' and not target == 'N':
self._grid_vals = self.grid_vals / (self.eng * log_bin_width)
self._spec_type = 'dNdE'
elif self.spec_type == 'dNdE' and not target == 'dNdE':
self._grid_vals = self.grid_vals * self.eng * log_bin_width
self._spec_type = 'N'
def get_optical_depth(rs_vec, xe_vec):
"""Computes the optical depth given an ionization history.
Parameters
----------
rs_vec : ndarray
Redshift (1+z).
xe_vec : ndarray
Free electron fraction xe = ne/nH.
Returns
-------
float
The optical depth.
"""
from darkhistory.spec.spectools import get_log_bin_width
rs_log_bin_width = get_log_bin_width(rs_vec)
abs_dtdz_vec = -dtdz(rs_vec)
return np.dot(xe_vec * nH * thomson_xsec * c * abs_dtdz_vec,
rs_vec**4 * rs_log_bin_width)
def toteng(self, bound_type=None, bound_arr=None):
"""Returns the total energy of particles in part of the spectra.
The part of the `Spectrum` objects to find the total energy of particles can be specified in two ways, and is specified by `bound_type`. Multiple totals can be obtained through `bound_arr`.
Parameters
----------
bound_type : {'bin', 'eng', None}
The type of bounds to use. Bound values do not have to be within the [0:eng.size] for `'bin'` or within the abscissa for `'eng'`. `None` should only be used when computing the total particle number in the spectrum.
Specifying ``bound_type=='bin'`` without bound_arr gives the total energy in each bin.
bound_arr : ndarray of length N, optional
An array of boundaries (bin or energy), between which the total number of particles will be computed. If unspecified, the total number of particles in the whole spectrum is computed.
For 'bin', bounds are specified as the bin *boundary*, with 0 being the left most boundary, 1 the right-hand of the first bin and so on. This is equivalent to integrating over a histogram. For 'eng', bounds are specified by energy values.
These boundaries need not be integer values for 'bin': specifying np.array([0.5, 1.5]) for example will include half of the first bin and half of the second.
Returns
-------
ndarray of shape (self.rs.size, N-1) or length N-1
Total number of particles in the spectra or between the specified boundaries.
Examples
--------
>>> from darkhistory.spec.spectrum import Spectrum
>>> eng = np.array([1, 10, 100, 1000])
>>> spec_arr = [Spectrum(eng, np.arange(4) + 4*i, rs=100, spec_type='N') for i in np.arange(4)]
>>> test_spectra = Spectra(spec_arr)
>>> test_spectra.toteng()
array([ 3210., 7654., 12098., 16542.])
>>> test_spectra.toteng('bin', np.array([1, 3]))
array([[ 210., 650., 1090., 1530.]])
>>> test_spectra.toteng('eng', np.array([10, 1e4]))
array([[ 3205., 7625., 12045., 16465.]])
See Also
---------
:meth:`Spectra.totN`
"""
log_bin_width = get_log_bin_width(self.eng)
# Using the broadcasting rules here.
if self.spec_type == 'dNdE':
dNdlogE = self.grid_vals * self.eng
elif self.spec_type == 'N':
dNdlogE = self.grid_vals / log_bin_width
if bound_type is not None:
if bound_arr is None:
return dNdlogE * self.eng * log_bin_width
if bound_type == 'bin':
if not all(np.diff(bound_arr) >= 0):
raise TypeError('bound_arr must have increasing entries.')
# Size is number of totals requested x number of Spectrums.
eng_in_bin = np.zeros((bound_arr.size - 1, self.in_eng.size))
if bound_arr[0] > self.eng.size or bound_arr[-1] < 0:
return eng_in_bin
for i, (low,
upp) in enumerate(zip(bound_arr[:-1], bound_arr[1:])):
# Set the lower and upper bounds, including case where
# low and upp are outside of the bins.
if low > self.eng.size or upp < 0:
continue
low_ceil = int(np.ceil(low))
low_floor = int(np.floor(low))
upp_ceil = int(np.ceil(upp))
upp_floor = int(np.floor(upp))
# Sum the bins that are completely between the bounds.
eng_full_bins = np.dot(
dNdlogE[:, low_ceil:upp_floor] *
self.eng[low_ceil:upp_floor],
log_bin_width[low_ceil:upp_floor])
eng_part_bins = np.zeros_like(self.in_eng)
if low_floor == upp_floor or low_ceil == upp_ceil:
# Bin indices are within the same bin. The second
# requirement covers the case where upp_ceil is
# eng.size.
eng_part_bins += (dNdlogE[:, low_floor] * (upp - low) *
self.eng[low_floor] *
log_bin_width[low_floor])
else:
# Add up part of the bin for the low partial bin
# and the high partial bin.
eng_part_bins += (dNdlogE[:, low_floor] *
(low_ceil - low) *
self.eng[low_floor] *
log_bin_width[low_floor])
if upp_floor < self.eng.size:
# If upp_floor is eng.size, then there is no
# partial bin for the upper index.
eng_part_bins += (dNdlogE[:, upp_floor] *
(upp - upp_floor) *
self.eng[upp_floor] *
log_bin_width[upp_floor])
eng_in_bin[i] = eng_full_bins + eng_part_bins
return eng_in_bin
if bound_type == 'eng':
bin_boundary = get_bin_bound(self.eng)
eng_bin_ind = np.interp(np.log(bound_arr),
np.log(bin_boundary),
np.arange(bin_boundary.size),
left=0,
right=self.eng.size + 1)
return self.toteng('bin', eng_bin_ind)
else:
return (np.dot(dNdlogE, self.eng * log_bin_width) +
self.eng_underflow)
def rebin(self, out_eng):
""" Re-bins all `Spectrum` objects according to a new abscissa.
Rebinning conserves total number and total energy.
Parameters
----------
out_eng : ndarray
The new abscissa to bin into. If `self.eng` has values that are smaller than `out_eng[0]`, then the new underflow will be filled. If `self.eng` has values that exceed `out_eng[-1]`, then an error is returned.
rebin_type : {'1D', '2D'}, optional
Whether to rebin each `Spectrum` separately (`'1D'`), or the whole `Spectra` object at once (`'2D'`). Default is `'2D'`.
See Also
--------
spec.spectools.rebin_N_2D_arr
"""
if not np.all(np.diff(out_eng) > 0):
raise TypeError(
'new abscissa must be ordered in increasing energy.')
# Get the bin indices that the current abscissa (self.eng)
# corresponds to in the new abscissa (new_eng). Bin indices are
# with respect to bin centers.
# Add an additional bin at the lower end of out_eng so that
# underflow can be treated easily.
first_bin_eng = np.exp(
np.log(out_eng[0]) - (np.log(out_eng[1]) - np.log(out_eng[0])))
new_eng = np.insert(out_eng, 0, first_bin_eng)
# Find the relative bin indices for self.eng. The first bin in
# new_eng has bin index -1. Underflow has index -2, overflow
# corresponds to new_eng.size
bin_ind = np.interp(self.eng,
new_eng,
np.arange(new_eng.size) - 1,
left=-2,
right=new_eng.size)
# Locate where bin_ind is below 0, above self.length-1
# or in between.
ind_low = np.where(bin_ind < 0)[0]
ind_high = np.where(bin_ind == new_eng.size)[0]
ind_reg = np.where((bin_ind >= 0) & (bin_ind <= new_eng.size - 1))[0]
if ind_high.size > 0:
warnings.warn(
"The new abscissa lies below the old one: only bins that lie within the new abscissa will be rebinned, bins above the abscissa will be discarded.",
RuntimeWarning)
# These arrays are of size in_eng x eng.
N_arr = self.totN('bin')
toteng_arr = self.toteng('bin')
N_arr_low = N_arr[:, ind_low]
N_arr_high = N_arr[:, ind_high]
N_arr_reg = N_arr[:, ind_reg]
toteng_arr_low = toteng_arr[:, ind_low]
# Factor depends on the spec_type.
if self.spec_type == 'dNdE':
# E dlog E of the new array.
fac = new_eng * get_log_bin_width(new_eng)
elif self.spec_type == 'N':
fac = np.ones_like(new_eng)
# Regular bins first.
# reg_bin_low is the array of the lower bins to be allocated the
# particles in N_arr_reg, similarly reg_bin_upp. This should also
# take care of the case where bin_ind is an integer.
reg_bin_low = np.floor(bin_ind[ind_reg]).astype(int)
reg_bin_upp = reg_bin_low + 1
# Takes care of the case where eng[-1] = new_eng[-1], which falls
# under regular indices. Remember the extra bin on the left.
reg_bin_low[reg_bin_low == new_eng.size - 2] = new_eng.size - 3
reg_bin_upp[reg_bin_upp == new_eng.size - 1] = new_eng.size - 2
# Split the particles up into the lower bin and upper bin.
# Remember there's an extra bin on the left when indexing into
# new_E_dlogE.
reg_data_low = ((reg_bin_upp - bin_ind[ind_reg]) * N_arr_reg /
fac[reg_bin_low + 1])
reg_data_upp = ((bin_ind[ind_reg] - reg_bin_low) * N_arr_reg /
fac[reg_bin_upp + 1])
# Handle low bins.
low_bin_low = np.floor(bin_ind[ind_low]).astype(int)
N_above_underflow = np.sum(
(bin_ind[ind_low] - low_bin_low) * N_arr_low, axis=1)
eng_above_underflow = N_above_underflow * new_eng[1]
N_underflow = np.sum(N_arr_low, axis=1) - N_above_underflow
eng_underflow = (np.sum(toteng_arr_low, axis=1) - eng_above_underflow)
low_data = N_above_underflow / fac[1]
# Add up, obtain the new dN/dE.
new_data = np.zeros((self.in_eng.size, new_eng.size))
new_data[:, 1] += low_data
# np.add.at(new_data, (slice(None), reg_bin_low+1), reg_data_low)
# np.add.at(new_data, (slice(None), reg_bin_upp+1), reg_data_upp)
# Replace with agg.aggregate
from darkhistory.numpy_groupies import aggregate as agg
# low_data = agg.aggregate(
# reg_bin_low+1, reg_data_low, func='sum', fill_value=0, axis=1, size=new_eng.size
# )
# upp_data = agg.aggregate(
# reg_bin_upp+1, reg_data_upp, func='sum', fill_value=0, axis=1, size=new_eng.size
# )
# new_data = (low_data + upp_data)
low_data = agg.aggregate(reg_bin_low + 1,
np.transpose(reg_data_low),
func='sum',
fill_value=0,
axis=0,
size=new_eng.size)
upp_data = agg.aggregate(reg_bin_upp + 1,
np.transpose(reg_data_upp),
func='sum',
fill_value=0,
axis=0,
size=new_eng.size)
new_data += np.transpose(low_data + upp_data)
# new_data[:,reg_bin_low+1] += reg_data_low
# new_data[:,reg_bin_upp+1] += reg_data_upp
self._eng = new_eng[1:]
self._grid_vals = new_data[:, 1:]
self._N_underflow += N_underflow
self._eng_underflow += eng_underflow
def N(self, value):
if self._spec_type == 'dNdE':
self._data = value * self.eng * get_log_bin_width(self.eng)
elif self._spec_type == 'N':
self._data = value
def N(self):
if self._spec_type == 'dNdE':
return self._data * self.eng * get_log_bin_width(self.eng)
elif self._spec_type == 'N':
return self._data
def rebin(self, out_eng):
""" Rebins according to a new abscissa.
The total number and total energy is conserved.
If a bin in the old abscissa self.eng is below the lowest bin of the new abscissa out_eng, then the total number and energy not assigned to the lowest bin are assigned to the underflow.
If a bin in self.eng is above the highest bin in out_eng, a warning is thrown, the values are simply discarded, and the total number and energy can no longer be conserved.
Parameters
----------
out_eng : ndarray
The new abscissa to bin into.
Returns
-------
None
Notes
-----
Total number and energy are conserved by assigning the number of particles :math:`N` in a bin of energy :math:`E` to two adjacent bins in the new abscissa out_eng, with energies :math:`E_\\text{low}` and :math:`E_\\text{upp}` such that :math:`E_\\text{low} < E < E_\\text{upp}`\ . The number of particles :math:`N_\\text{low}` and :math:`N_\\text{upp}` assigned to these two bins are given by
.. math::
N_\\text{low} &= \\frac{E_\\text{upp} - E}{E_\\text{upp} - E_\\text{low}} N \\,, \\\\
N_\\text{upp} &= \\frac{E - E_\\text{low}}{E_\\text{upp} - E_\\text{low}} N
Rebinning works best when going from a finer binning to a coarser binning. Going the other way can result in spiky features, since the coarser binning simply does not contain enough information to reconstruct the finer binning in this way.
See Also
--------
:func:`.spectools.rebin_N_arr`
"""
if not np.all(np.diff(out_eng) > 0):
raise TypeError(
"new abscissa must be ordered in increasing energy.")
# if out_eng[-1] < self.eng[-1]:
# raise OverflowError("the new abscissa lies below the old one: this function cannot handle overflow (yet?).")
# Get the bin indices that the current abscissa (self.eng) corresponds to in the new abscissa (new_eng). Can be any number between 0 and self.length-1. Bin indices are wrt the bin centers.
# Add an additional bin at the lower end of out_eng so that underflow can be treated easily.
# Forces out_eng to be float, avoids strange problems with np.insert
# below if out_eng is of type int.
out_eng = out_eng.astype(float)
first_bin_eng = np.exp(
np.log(out_eng[0]) - (np.log(out_eng[1]) - np.log(out_eng[0])))
new_eng = np.insert(out_eng, 0, first_bin_eng)
# Find the relative bin indices for self.eng wrt new_eng. The first bin in new_eng has bin index -1.
bin_ind_interp = interp1d(new_eng,
np.arange(new_eng.size) - 1,
bounds_error=False,
fill_value=(-2, new_eng.size))
bin_ind = bin_ind_interp(self.eng)
# bin_ind = np.interp(self.eng, new_eng,
# np.arange(new_eng.size)-1, left = -2, right = new_eng.size)
# Locate where bin_ind is below 0, above self.length-1 and in between.
ind_low = np.where(bin_ind < 0)
ind_high = np.where(bin_ind == new_eng.size)
ind_reg = np.where((bin_ind >= 0) & (bin_ind <= new_eng.size - 1))
if ind_high[0].size > 0:
warnings.warn(
"The new abscissa lies below the old one: only bins that lie within the new abscissa will be rebinned, bins above the abscissa will be discarded.",
RuntimeWarning)
# raise OverflowError("the new abscissa lies below the old one: this function cannot handle overflow (yet?).")
# Get the total N and toteng in each bin of self._data
if self._spec_type == 'dNdE':
N_arr = self.totN('bin')
toteng_arr = self.toteng('bin')
elif self._spec_type == 'N':
N_arr = self.N
toteng_arr = self.N * self.eng
N_arr_low = N_arr[ind_low]
N_arr_high = N_arr[ind_high]
N_arr_reg = N_arr[ind_reg]
toteng_arr_low = toteng_arr[ind_low]
# Bin width of the new array. Use only the log bin width, so that dN/dE = N/(E d log E)
if self._spec_type == 'dNdE':
new_E_dlogE = new_eng * get_log_bin_width(new_eng)
# Regular bins first, done in a completely vectorized fashion.
# reg_bin_low is the array of the lower bins to be allocated the particles in N_arr_reg, similarly reg_bin_upp. This should also take care of the fact that bin_ind is an integer.
reg_bin_low = np.floor(bin_ind[ind_reg]).astype(int)
reg_bin_upp = reg_bin_low + 1
# Takes care of the case where eng[-1] = new_eng[-1]
reg_bin_low[reg_bin_low == new_eng.size - 2] = new_eng.size - 3
reg_bin_upp[reg_bin_upp == new_eng.size - 1] = new_eng.size - 2
if self._spec_type == 'dNdE':
reg_dNdE_low = ((reg_bin_upp - bin_ind[ind_reg]) * N_arr_reg /
new_E_dlogE[reg_bin_low + 1])
reg_dNdE_upp = ((bin_ind[ind_reg] - reg_bin_low) * N_arr_reg /
new_E_dlogE[reg_bin_upp + 1])
elif self._spec_type == 'N':
reg_N_low = (reg_bin_upp - bin_ind[ind_reg]) * N_arr_reg
reg_N_upp = (bin_ind[ind_reg] - reg_bin_low) * N_arr_reg
# Low bins.
low_bin_low = np.floor(bin_ind[ind_low]).astype(int)
N_above_underflow = np.sum(
(bin_ind[ind_low] - low_bin_low) * N_arr_low)
eng_above_underflow = N_above_underflow * new_eng[1]
N_underflow = np.sum(N_arr_low) - N_above_underflow
eng_underflow = np.sum(toteng_arr_low) - eng_above_underflow
if self._spec_type == 'dNdE':
low_dNdE = N_above_underflow / new_E_dlogE[1]
# Add up, obtain the new data.
new_data = np.zeros(new_eng.size)
if self._spec_type == 'dNdE':
new_data[1] += low_dNdE
# reg_dNdE_low = -1 refers to new_eng[0]
np.add.at(new_data, reg_bin_low + 1, reg_dNdE_low)
np.add.at(new_data, reg_bin_upp + 1, reg_dNdE_upp)
# print(new_data[reg_bin_low+1])
# new_data[reg_bin_low+1] += reg_dNdE_low
# new_data[reg_bin_upp+1] += reg_dNdE_upp
elif self._spec_type == 'N':
new_data[1] += N_above_underflow
np.add.at(new_data, reg_bin_low + 1, reg_N_low)
np.add.at(new_data, reg_bin_upp + 1, reg_N_upp)
# new_data[reg_bin_low+1] += reg_N_low
# new_data[reg_bin_upp+1] += reg_N_upp
# Implement changes.
self.eng = new_eng[1:]
self._data = new_data[1:]
self.length = self.eng.size
self.underflow['N'] += N_underflow
self.underflow['eng'] += eng_underflow
def dNdE(self):
if self._spec_type == 'dNdE':
return self._data
elif self._spec_type == 'N':
return self._data / (self.eng * get_log_bin_width(self.eng))
def toteng(self, bound_type=None, bound_arr=None):
"""Returns the total energy of particles in part of the spectrum.
The part of the spectrum can be specified in two ways, and is specified by bound_type. Multiple totals can be obtained through bound_arr.
Parameters
----------
bound_type : {'bin', 'eng', None}
The type of bounds to use. Bound values do not have to be within the [0:length] for 'bin' or within the abscissa for 'eng'. None should only be used to obtain the total energy.
Specifying ``bound_type='bin'`` without bound_arr gives the total energy in each bin.
bound_arr : ndarray of length N, optional
An array of boundaries (bin or energy), between which the total number of particles will be computed. If unspecified, the total number of particles in the whole spectrum is computed.
For 'bin', bounds are specified as the bin *boundary*, with 0 being the left most boundary, 1 the right-hand of the first bin and so on. This is equivalent to integrating over a histogram. For 'eng', bounds are specified by energy values.
These boundaries need not be integer values for 'bin': specifying np.array([0.5, 1.5]) for example will include half of the first bin and half of the second.
Returns
-------
ndarray of length N-1, or float
Total energy in the spectrum or between the specified boundaries.
Examples
---------
>>> eng = np.array([1, 10, 100, 1000])
>>> N = np.array([1, 2, 3, 4])
>>> spec = Spectrum(eng, N, spec_type='N')
>>> spec.toteng()
4321.0
>>> spec.toteng('bin', np.array([1, 3]))
array([320.])
>>> spec.toteng('eng', np.array([10, 1e4]))
array([4310.])
See Also
---------
:meth:`.Spectrum.totN`
"""
eng = self.eng
length = self.length
log_bin_width = get_log_bin_width(self.eng)
if self._spec_type == 'dNdE':
dNdlogE = self.eng * self.dNdE
elif self._spec_type == 'N':
dNdlogE = self.N / log_bin_width
if bound_type is not None:
if bound_arr is None:
return dNdlogE * eng * log_bin_width
if bound_type == 'bin':
if not all(np.diff(bound_arr) >= 0):
raise TypeError("bound_arr must have increasing entries.")
eng_in_bin = np.zeros(bound_arr.size - 1)
if bound_arr[0] > length or bound_arr[-1] < 0:
return eng_in_bin
for low, upp, i in zip(bound_arr[:-1], bound_arr[1:],
np.arange(eng_in_bin.size)):
if low > length or upp < 0:
eng_in_bin[i] = 0
continue
low_ceil = int(np.ceil(low))
low_floor = int(np.floor(low))
upp_ceil = int(np.ceil(upp))
upp_floor = int(np.floor(upp))
# Sum the bins that are completely between the bounds.
eng_full_bins = np.dot(
eng[low_ceil:upp_floor] * dNdlogE[low_ceil:upp_floor],
log_bin_width[low_ceil:upp_floor])
eng_part_bins = 0
if low_floor == upp_floor or low_ceil == upp_ceil:
# Bin indices are within the same bin. The second requirement covers the case where upp_ceil is length.
eng_part_bins += (eng[low_floor] * dNdlogE[low_floor] *
(upp - low) *
log_bin_width[low_floor])
else:
# Add up part of the bin for the low partial bin and the high partial bin.
eng_part_bins += (eng[low_floor] * dNdlogE[low_floor] *
(low_ceil - low) *
log_bin_width[low_floor])
if upp_floor < length:
# If upp_floor is length, then there is no partial bin for the upper index.
eng_part_bins += (eng[upp_floor] *
dNdlogE[upp_floor] *
(upp - upp_floor) *
log_bin_width[upp_floor])
eng_in_bin[i] = eng_full_bins + eng_part_bins
return eng_in_bin
if bound_type == 'eng':
bin_boundary = get_bin_bound(self.eng)
eng_bin_ind = np.interp(np.log(bound_arr),
np.log(bin_boundary),
np.arange(bin_boundary.size),
left=0,
right=length + 1)
return self.toteng('bin', eng_bin_ind)
else:
return (np.dot(dNdlogE, eng * log_bin_width) +
self.underflow['eng'])
