def test_re_evaluate(self):
a = array([1., 2., 3.])
b = array([4., 5., 6.])
c = array([7., 8., 9.])
x = evaluate("2*a + 3*b*c")
x = re_evaluate()
assert_array_equal(x, array([86., 124., 168.]))
def test_re_evaluate_dict(self):
a = array([1., 2., 3.])
b = array([4., 5., 6.])
c = array([7., 8., 9.])
x = evaluate("2*a + 3*b*c", local_dict={'a': a, 'b': b, 'c': c})
x = re_evaluate()
assert_array_equal(x, array([86., 124., 168.]))
def action(inputs, **ignore):
if all(
isinstance(x, ak._v2.contents.NumpyArray)
or not isinstance(x, ak._v2.contents.Content) for x in inputs):
return (ak._v2.contents.NumpyArray(
numexpr.re_evaluate(dict(zip(names, inputs)))), )
else:
return None
def getfunction(inputs):
if all(
isinstance(x, ak.layout.NumpyArray)
or not isinstance(x, ak.layout.Content) for x in inputs):
return lambda: (ak.layout.NumpyArray(
numexpr.re_evaluate(dict(zip(names, inputs)))), )
else:
return None
def process(self, **kwargs):
"""Process module."""
lum_key = self.key('luminosities')
kwargs = self.prepare_input(lum_key, **kwargs)
self._luminosities = kwargs[lum_key]
self._bands = kwargs['all_bands']
self._band_indices = kwargs['all_band_indices']
self._frequencies = kwargs['all_frequencies']
self._radius_phot = kwargs[self.key('radiusphot')]
self._temperature_phot = kwargs[self.key('temperaturephot')]
xc = self.X_CONST # noqa: F841
fc = self.FLUX_CONST # noqa: F841
cc = self.C_CONST
temperature_phot = self._temperature_phot
# Some temp vars for speed.
zp1 = 1.0 + kwargs[self.key('redshift')]
Azp1 = u.Angstrom.cgs.scale / zp1
czp1 = cc / zp1
seds = []
rest_wavs_dict = {}
evaled = False
for li, lum in enumerate(self._luminosities):
bi = self._band_indices[li]
if lum == 0.0:
seds.append(
np.zeros(
len(self._sample_wavelengths[bi]) if bi >= 0 else 1))
continue
if bi >= 0:
rest_wavs = rest_wavs_dict.setdefault(
bi, self._sample_wavelengths[bi] * Azp1)
else:
rest_wavs = np.array( # noqa: F841
[czp1 / self._frequencies[li]])
radius_phot = self._radius_phot[li] # noqa: F841
temperature_phot = self._temperature_phot[li] # noqa: F841
if not evaled:
seds.append(
ne.evaluate('fc * radius_phot**2 / rest_wavs**5 / '
'expm1(xc / rest_wavs / temperature_phot)'))
evaled = True
else:
seds.append(ne.re_evaluate())
seds[-1][np.isnan(seds[-1])] = 0.0
seds = self.add_to_existing_seds(seds, **kwargs)
# Units of `seds` is ergs / s / Angstrom.
return {'sample_wavelengths': self._sample_wavelengths, 'seds': seds}
def process(self, **kwargs):
"""Process module."""
Transform.process(self, **kwargs)
self._kappa = kwargs[self.key('kappa')]
self._mass = kwargs[self.key('mcsm')] * M_SUN_CGS
self._R0 = kwargs[self.key('r0')] * AU_CGS # AU to cm
self._s = kwargs[self.key('s')]
self._rho = kwargs[self.key('rho')]
# scaling constant for CSM density profile
self._q = self._rho * self._R0 ** self._s
# outer radius of CSM shell
self._Rcsm = (
(3.0 - self._s) / (4.0 * np.pi * self._q) * self._mass +
self._R0 ** (3.0 - self._s)) ** (1.0 / (3.0 - self._s))
# radius of photosphere (should be within CSM)
self._Rph = abs(
(-2.0 * (1.0 - self._s) / (3.0 * self._kappa * self._q) +
self._Rcsm ** (1.0 - self._s)) ** (1.0 / (1.0 - self._s)))
self._tau_diff = (
self._kappa * self._mass) / (13.8 * C_CGS * self._Rph) / DAY_CGS
tbarg = self.MIN_EXP_ARG * self._tau_diff ** 2
new_lum = np.zeros_like(self._times_to_process)
evaled = False
lum_cache = OrderedDict()
min_te = min(self._dense_times_since_exp)
for ti, te in enumerate(self._times_to_process):
if te <= 0.0:
continue
if te in lum_cache:
new_lum[ti] = lum_cache[te]
continue
te2 = te ** 2
tb = max(np.sqrt(max(te2 - tbarg, 0.0)), min_te)
int_times = np.linspace(tb, te, self.N_INT_TIMES)
dt = int_times[1] - int_times[0]
td = self._tau_diff # noqa: F841
int_lums = np.interp( # noqa: F841
int_times, self._dense_times_since_exp,
self._dense_luminosities)
if not evaled:
int_arg = ne.evaluate('int_lums * int_times / td**2 * '
'exp((int_times - te) / td)')
evaled = True
else:
int_arg = ne.re_evaluate()
int_arg[np.isnan(int_arg)] = 0.0
lum_val = np.trapz(int_arg, dx=dt)
lum_cache[te] = lum_val
new_lum[ti] = lum_val
return {self.dense_key('luminosities'): new_lum}
def process(self, **kwargs):
"""Process module."""
lum_key = self.key('luminosities')
kwargs = self.prepare_input(lum_key, **kwargs)
self._luminosities = kwargs[lum_key]
self._bands = kwargs['all_bands']
self._band_indices = kwargs['all_band_indices']
self._frequencies = kwargs['all_frequencies']
self._radius_phot = kwargs[self.key('radiusphot')]
self._temperature_phot = kwargs[self.key('temperaturephot')]
xc = self.X_CONST # noqa: F841
fc = self.FLUX_CONST # noqa: F841
cc = self.C_CONST
temperature_phot = self._temperature_phot
zp1 = 1.0 + kwargs[self.key('redshift')]
seds = []
evaled = False
for li, lum in enumerate(self._luminosities):
radius_phot = self._radius_phot[li] # noqa: F841
temperature_phot = self._temperature_phot[li] # noqa: F841
bi = self._band_indices[li]
if lum == 0.0:
if bi >= 0:
seds.append(np.zeros_like(self._sample_wavelengths[bi]))
else:
seds.append([0.0])
continue
if bi >= 0:
rest_wavs = (self._sample_wavelengths[bi] *
u.Angstrom.cgs.scale / zp1)
else:
rest_wavs = [cc / (self._frequencies[li] * zp1)] # noqa: F841
if not evaled:
sed = ne.evaluate('fc * radius_phot**2 / rest_wavs**5 / '
'expm1(xc / rest_wavs / temperature_phot)')
evaled = True
else:
sed = ne.re_evaluate()
sed[np.isnan(sed)] = 0.0
seds.append(sed)
seds = self.add_to_existing_seds(seds, **kwargs)
return {'sample_wavelengths': self._sample_wavelengths, 'seds': seds}
def process(self, **kwargs):
"""Process module."""
raise NotImplementedError('`MultiBlackbody` is not yet functional.')
kwargs = self.prepare_input(self.key('luminosities'), **kwargs)
self._luminosities = kwargs[self.key('luminosities')]
self._bands = kwargs['all_bands']
self._band_indices = kwargs['all_band_indices']
self._areas = kwargs[self.key('areas')]
self._radius_phots = kwargs[self.key('radiusphots')]
self._temperature_phots = kwargs[self.key('temperaturephots')]
xc = self.X_CONST # noqa: F841
fc = self.FLUX_CONST # noqa: F841
zp1 = 1.0 + kwargs[self.key('redshift')]
seds = []
for li, lum in enumerate(self._luminosities):
cur_band = self._bands[li] # noqa: F841
bi = self._band_indices[li]
rest_freqs = [
x * zp1 # noqa: F841
for x in self._sample_frequencies[bi]
]
wav_arr = np.array(self._sample_wavelengths[bi]) # noqa: F841
radius_phot = self._radius_phots[li] # noqa: F841
temperature_phot = self._temperature_phots[li] # noqa: F841
if li == 0:
sed = ne.evaluate(
'fc * radius_phot**2 * rest_freqs**3 / '
'(exp(xc * rest_freqs / temperature_phot) - 1.0)')
else:
sed = ne.re_evaluate()
sed = np.nan_to_num(sed)
seds.append(list(sed))
seds = self.add_to_existing_seds(seds, **kwargs)
# Units of `seds` is ergs / s / Angstrom.
return {
'sample_wavelengths': self._sample_wavelengths,
self.key('seds'): seds
}
def process(self, **kwargs):
"""Process module."""
kwargs = self.prepare_input(self.key('luminosities'), **kwargs)
prt = self._printer
self._rest_t_explosion = kwargs[self.key('resttexplosion')]
self._times = kwargs[self.key('rest_times')]
self._seds = kwargs.get(self.key('seds'))
self._bands = kwargs['all_bands']
self._band_indices = kwargs['all_band_indices']
self._sample_wavelengths = kwargs['sample_wavelengths']
self._frequencies = kwargs['all_frequencies']
self._luminosities = kwargs[self.key('luminosities')]
self._line_wavelength = kwargs[self.key('line_wavelength')]
self._line_width = kwargs[self.key('line_width')]
self._line_time = kwargs[self.key('line_time')]
self._line_duration = kwargs[self.key('line_duration')]
self._line_amplitude = kwargs[self.key('line_amplitude')]
lw = self._line_wavelength # noqa: F841
ls = self._line_width
cc = self.C_CONST
# Some temp vars for speed.
zp1 = 1.0 + kwargs[self.key('redshift')]
czp1A = cc / (zp1 * u.Angstrom.cgs.scale)
amps = self._line_amplitude * np.array([
np.exp(-0.5 * ((x - self._rest_t_explosion - self._line_time) /
self._line_duration)**2) for x in self._times
])
if self._seds is None:
raise ValueError(prt.message('line_sed'))
seds = [x * (1.0 - amps[xi]) for xi, x in enumerate(self._seds)]
amps *= self._luminosities / (ls * SQRT_2_PI)
amps_dict = {}
evaled = False
for li, lum in enumerate(self._luminosities):
bi = self._band_indices[li]
if lum == 0.0:
continue
bind = czp1A / self._frequencies[li] if bi < 0 else bi
if bind not in amps_dict:
# Leave `rest_wavs` in Angstroms.
if bi >= 0:
rest_wavs = self._sample_wavelengths[bi] / zp1
else:
rest_wavs = np.array([bind]) # noqa: F841
if not evaled:
amps_dict[bind] = ne.evaluate(
'exp(-0.5 * ((rest_wavs - lw) / ls) ** 2)')
evaled = True
else:
amps_dict[bind] = ne.re_evaluate()
seds[li] += amps[li] * amps_dict[bind]
# seds[li][np.isnan(seds[li])] = 0.0
# Units of `seds` is ergs / s / Angstrom.
return {
'sample_wavelengths': self._sample_wavelengths,
self.key('seds'): seds
}
import numexpr as ne
import numpy as np
nrange = (2 ** np.arange(6, 24)).astype(int)
t_numpy = []
t_numexpr = []
for n in nrange:
a = np.random.random(n)
b = np.arange(n, dtype=np.double)
c = np.random.random(n)
c1 = ne.evaluate("a ** 2 + b ** 2 + 2 * a * b * c ", optimization='aggressive')
t1 = %timeit -oq -n 10 a ** 2 + b ** 2 + 2 * a * b * c
t2 = %timeit -oq -n 10 ne.re_evaluate()
t_numpy.append(t1.best)
t_numexpr.append(t2.best)
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
import matplotlib.pyplot as plt
import seaborn; seaborn.set()
plt.loglog(nrange, t_numpy, label='numpy')
plt.loglog(nrange, t_numexpr, label='numexpr')
plt.legend(loc='lower right')
plt.xlabel('Vectors size')
plt.ylabel('Execution Time (s)');
def hydrodynamic_model(Input):
""" Iterates the friction coefficient and yields flow velocities.
Computes the flow solution in the inlets and basin, and iterates the velocity scale
until it matches that found in the hydrodynamic solution (given a certain tolerance).
Solves the system :math:`\mathbf{A} \mathbf{\hat{u}} = \mathbf{f}`
for the flow velocity amplitude :math:`\mathbf{\hat{u}}`. The forcing therm :math:`\mathbf{f}`
consists of a tidal forcing, and matrix :math:`\mathbf{A} = \mathbf{A}_1 + \mathbf{A}_2 - \mathbf{A}_3`.
The first term is associated with friction in the channel.
The second with the sea impedance.
The third term is associated with basin impedance.
Args:
Input (Class): Class containing all data for the various geograpical elements.
Consists of the Basin, OpenInlets, Ocean, and Pars classes.
l2 (np.ndarray): 3d-array containing the L2-norm for the eigenfunctinos :math:`\phi_{m, n}(x, y)`.
phij (np.ndarray): 3d-array containing the integrated eigenfunctions over all j inlets.
Returns:
returnObject (dict):
uj (np.ndarray): flow velocity in the inlets.
ub (float): flow velocity in the basin.
mui2 (np.ndarray): frictional correction factor for the inlets.
mub2 (complex): frictional correction factor.
"""
Basin = Input.Basin
Inlets = Input.OpenInlets
Ocean = Input.Ocean
Pars = Input.Pars
l2 = Basin.EigenFunctions.l2
phij = Basin.EigenFunctions.phij
A1 = np.zeros((Inlets.numinlets, Inlets.numinlets), dtype=complex)
A2 = np.zeros((Inlets.numinlets, Inlets.numinlets), dtype=complex)
A3 = np.zeros((Inlets.numinlets, Inlets.numinlets), dtype=complex)
A = np.zeros((Inlets.numinlets, Inlets.numinlets), dtype=complex)
A1c = 1j * Ocean.tidefreq * Inlets.lengths / g * np.eye(Inlets.numinlets)
a_width = Inlets.widths[0, :]
beta = Inlets.widths / a_width[:, np.newaxis]
betap = (beta + 1) / 2 * (1 - np.eye(np.sum(Inlets.numinlets)))
betam = (beta - 1) / 2
alpha = (
abs(np.transpose(Inlets.locations) - Inlets.locations) / a_width[:, np.newaxis]
)
alpha[alpha == 0] = 1e-7 # Ensure that alpha**2 != 0
if np.sum(Inlets.numinlets) > 1:
A2 = (
(Inlets.depths * Ocean.tidefreq * a_width[:, np.newaxis])
/ (2 * g * Ocean.depth)
* (
beta
+ (2j / np.pi)
* (
(3 / 2 - eg) * beta
+ betam ** 2
* np.log(
Ocean.ko
* a_width[:, np.newaxis]
/ 2
* np.sqrt(alpha ** 2 - betam ** 2)
)
- betap ** 2
* np.log(
Ocean.ko
* a_width[:, np.newaxis]
/ 2
* np.sqrt(alpha ** 2 - betap ** 2)
)
+ alpha
* (
betam * np.log((alpha + betam) / (alpha - betam))
- betap * np.log((alpha + betap) / (alpha - betap))
)
+ alpha ** 2
* np.log(
np.sqrt((alpha ** 2 - betam ** 2) / (alpha ** 2 - betap ** 2))
)
)
)
)
A2self = (
(Inlets.depths[0, :] * Ocean.tidefreq * Inlets.widths[0, :])
/ (2 * g * Ocean.depth)
* (
1
+ 2j / np.pi * (3 / 2 - eg - np.log(Ocean.ko * Inlets.widths[0, :] / 2))
)
)
A2[np.eye(Inlets.numinlets) == 1] = A2self
else:
A2 = (
(Inlets.depths[0, :] * Ocean.tidefreq * Inlets.widths[0, :])
/ (2 * g * Ocean.depth)
* (
1
+ 2j / np.pi * (3 / 2 - eg - np.log(Ocean.ko * Inlets.widths[0, :] / 2))
)
)
A3c = (
Ocean.tidefreq
/ (g * 1j)
* np.reshape(
(
Inlets.widths[np.newaxis, np.newaxis, :, :]
* Inlets.depths[np.newaxis, np.newaxis, :, :]
* (phij[:, :, :, np.newaxis] * phij[:, :, np.newaxis, :])
/ (l2[:, :, :, np.newaxis] * Basin.depth)
),
(-1, np.sum(Inlets.numinlets), np.sum(Inlets.numinlets)),
)
)
cmnqc = (
Inlets.widths[np.newaxis, :, :]
* Inlets.depths[np.newaxis, :, :]
* (phij[:, :, :])
/ (Basin.depth)
)
ubc = cmnqc / np.sqrt(l2)
kmn2 = np.reshape(
Basin.EigenFunctions.kmn2[:, :, np.newaxis, np.newaxis], (-1, 1, 1)
)
kb2 = Basin.kb ** 2
uj = Inlets.uj
ub = Basin.ub
res = 42
r = 1
while res > 1e-10 and r < 50:
rb = 8 / (3 * np.pi) * Basin.cd * ub
rj = 8 / (3 * np.pi) * Inlets.cd * uj
mub2 = 1 - 1j * rb / (Ocean.tidefreq * Basin.depth)
mui2 = 1 - 1j * rj / (Ocean.tidefreq * Inlets.depths)
A1 = A1c * mui2
if r == 1:
A3 = mub2 * ne.evaluate("sum(A3c/(kmn2 - mub2*kb2), axis=0)")
else:
A3 = mub2 * ne.re_evaluate()
A = A1 + A2 - A3
B = -Ocean.tideamp * np.exp(1j * Ocean.wavenumber * Inlets.locations[0, :])
sol = np.linalg.solve(A, B)
if np.isnan(sol).any():
print(alpha, beta, betap, betam)
print(Inlets.widths)
print(Inlets.locations)
raise NameError("aiaiai")
ubn = np.sqrt(
1
/ (Basin.area)
* np.sum(
Basin.EigenFunctions.kmn2
* np.abs(
np.sum(ubc * sol[np.newaxis, np.newaxis, :], axis=2)
/ (Basin.EigenFunctions.kmn2 - mub2 * Basin.kb ** 2)
)
** 2
)
)
ires = np.abs(uj[0, :] - abs(sol))
bres = np.abs(ub - ubn)
res = np.max(np.r_[ires, bres])
uj = (abs(abs(sol[np.newaxis, :])) + uj) / 2
ub = (ub + abs(ubn)) / 2
r = r + 1
return uj, ub, mui2, mub2, sol[np.newaxis, :]
def process(self, **kwargs):
"""Process module."""
kwargs = self.prepare_input(self.key('luminosities'), **kwargs)
self._luminosities = kwargs[self.key('luminosities')]
self._bands = kwargs['all_bands']
self._band_indices = kwargs['all_band_indices']
self._frequencies = kwargs['all_frequencies']
self._radius_phot = np.array(kwargs[self.key('radiusphot')])
self._temperature_phot = np.array(kwargs[self.key('temperaturephot')])
self._cutoff_wavelength = kwargs[self.key('cutoff_wavelength')]
self._times = np.array(kwargs['rest_times'])
xc = self.X_CONST # noqa: F841
fc = self.FLUX_CONST
cc = self.C_CONST
ac = ANG_CGS
cwave_ac = self._cutoff_wavelength * ac
cwave_ac2 = cwave_ac * cwave_ac
cwave_ac3 = cwave_ac2 * cwave_ac # noqa: F841
zp1 = 1.0 + kwargs[self.key('redshift')]
lt = len(self._times)
seds = np.empty(lt, dtype=object)
rp2 = self._radius_phot ** 2
tp = self._temperature_phot
evaled = False
for li, lum in enumerate(self._luminosities):
bi = self._band_indices[li]
# tpi = tp[li]
# rp2i = rp2[li]
if lum == 0.0:
seds[li] = np.zeros(
len(self._sample_wavelengths[bi]) if bi >= 0 else 1)
continue
if bi >= 0:
rest_wavs = self._sample_wavelengths[bi] * ac / zp1
else:
rest_wavs = np.array([cc / (self._frequencies[li] * zp1)])
# Apply absorption to SED only bluewards of cutoff wavelength
ab = rest_wavs < cwave_ac # noqa: F841
tpi = tp[li] # noqa: F841
rp2i = rp2[li] # noqa: F841
if not evaled:
# Absorbed blackbody: 0% transmission at 0 Angstroms 100% at
# >3000 Angstroms.
sed = ne.evaluate(
"where(ab, fc * (rp2i / cwave_ac / "
"rest_wavs ** 4) / expm1(xc / rest_wavs / tpi), "
"fc * (rp2i / rest_wavs ** 5) / "
"expm1(xc / rest_wavs / tpi))"
)
evaled = True
else:
sed = ne.re_evaluate()
sed[np.isnan(sed)] = 0.0
seds[li] = sed
uniq_times = np.unique(self._times)
tsort = np.argsort(self._times)
uniq_is = np.searchsorted(self._times, uniq_times, sorter=tsort)
lu = len(uniq_times)
norms = self._luminosities[
uniq_is] / (fc / ac * rp2[uniq_is] * tp[uniq_is])
rp2 = rp2[uniq_is].reshape(lu, 1)
tp = tp[uniq_is].reshape(lu, 1)
tp2 = tp * tp
tp3 = tp2 * tp # noqa: F841
nxcs = self._nxcs # noqa: F841
f_blue_reds = ne.evaluate(
"sum((exp(-nxcs / (cwave_ac * tp)) * ("
"nxcs ** 2 + 2 * ("
"nxcs * cwave_ac * tp + cwave_ac2 * tp2)) / ("
"nxcs ** 3 * cwave_ac3)) + "
"(6 * tp3 - exp(-nxcs / (cwave_ac * tp)) * ("
"nxcs ** 3 + 3 * nxcs ** 2 * cwave_ac * tp + 6 * ("
"nxcs * cwave_ac2 * tp2 + cwave_ac3 *"
"tp3)) / cwave_ac3) / (nxcs ** 4), 1)"
)
norms /= f_blue_reds
# Apply renormalisation
seds *= norms[np.searchsorted(uniq_times, self._times)]
seds = self.add_to_existing_seds(seds, **kwargs)
# Units of `seds` is ergs / s / Angstrom.
return {'sample_wavelengths': self._sample_wavelengths,
self.key('seds'): seds}
def process(self, **kwargs):
"""Process module."""
kwargs = self.prepare_input(self.key('luminosities'), **kwargs)
self._rest_t_explosion = kwargs[self.key('resttexplosion')]
self._times = kwargs[self.key('rest_times')]
self._seds = kwargs[self.key('seds')]
self._bands = kwargs['all_bands']
self._band_indices = kwargs['all_band_indices']
self._sample_wavelengths = kwargs['sample_wavelengths']
self._frequencies = kwargs['all_frequencies']
self._luminosities = kwargs[self.key('luminosities')]
self._line_wavelength = kwargs[self.key('line_wavelength')]
self._line_width = kwargs[self.key('line_width')]
self._line_time = kwargs[self.key('line_time')]
self._line_duration = kwargs[self.key('line_duration')]
self._line_amplitude = kwargs[self.key('line_amplitude')]
lw = self._line_wavelength
ls = self._line_width
cc = self.C_CONST
zp1 = 1.0 + kwargs[self.key('redshift')]
amps = [
self._line_amplitude *
np.exp(-0.5 * ((x - self._rest_t_explosion - self._line_time) /
self._line_duration)**2) for x in self._times
]
seds = self._seds
evaled = False
for li, lum in enumerate(self._luminosities):
bi = self._band_indices[li]
if lum == 0.0:
if bi >= 0:
seds.append(np.zeros_like(self._sample_wavelengths[bi]))
else:
seds.append([0.0])
continue
if bi >= 0:
rest_wavs = (self._sample_wavelengths[bi] *
u.Angstrom.cgs.scale / zp1)
else:
rest_wavs = [cc / (self._frequencies[li] * zp1)] # noqa: F841
amp = lum * amps[li]
if not evaled:
sed = ne.evaluate(
'amp * exp(-0.5 * ((rest_wavs - lw) / ls) ** 2)')
evaled = True
else:
sed = ne.re_evaluate()
sed = np.nan_to_num(sed)
norm = (lum + amp / zp1 * np.sqrt(np.pi / 2.0) *
(1.0 + erf(lw / (np.sqrt(2.0) * ls)))) / lum
seds[li] += sed
seds[li] /= norm
return {
'sample_wavelengths': self._sample_wavelengths,
self.key('seds'): seds
}
h_resolution = 1000
iterations = 500
sub_pixel_sample = 2
record_history_length = 500
h_resolution *= sub_pixel_sample
v_resolution = int(h_resolution * (im_max - im_min) / (re_max - re_min))
initial_zs = (np.expand_dims(np.linspace(re_min, re_max, h_resolution), 0) +
1j * np.expand_dims(np.linspace(im_min, im_max, v_resolution), 1)
).flatten()
# Iterate raising to the power
zs = initial_zs
zs = ne.evaluate('initial_zs**zs')
for i in tqdm(range(iterations)):
zs = ne.re_evaluate()
final_zs = np.copy(zs)
# Record periodicities
period = np.zeros_like(final_zs, dtype=np.int32)
for i in tqdm(range(record_history_length)):
zs = ne.evaluate('initial_zs**zs')
close = is_close(final_zs, zs, atol=1e-6)
# If we've found a repeat for a point we've not previously
# found a repeat for, record the period
period[np.logical_and(period == 0, close)] = i + 1
# Image array, initially white
rgb = np.ones((h_resolution * v_resolution, 3)) * 255
def process(self, **kwargs):
"""Process module."""
kwargs = self.prepare_input(self.key('luminosities'), **kwargs)
self._luminosities = kwargs[self.key('luminosities')]
self._bands = kwargs['all_bands']
self._band_indices = kwargs['all_band_indices']
self._frequencies = kwargs['all_frequencies']
self._radius_phot = np.array(kwargs[self.key('radiusphot')])
self._temperature_phot = np.array(kwargs[self.key('temperaturephot')])
self._cutoff_wavelength = kwargs[self.key('cutoff_wavelength')]
self._alpha = kwargs[self.key('alpha')]
self._times = np.array(kwargs['rest_times'])
xc = self.X_CONST # noqa: F841
fc = self.FLUX_CONST
cc = self.C_CONST
ac = ANG_CGS
zp1 = 1.0 + kwargs[self.key('redshift')]
lt = len(self._times)
seds = np.empty(lt, dtype=object)
rp2 = self._radius_phot**2
tp = self._temperature_phot
evaled = False
for li, lum in enumerate(self._luminosities):
bi = self._band_indices[li]
if lum == 0.0:
seds[li] = np.zeros(
len(self._sample_wavelengths[bi]) if bi >= 0 else 1)
continue
if bi >= 0:
rest_wavs = self._sample_wavelengths[bi] * ac / zp1
else:
rest_wavs = np.array([cc / (self._frequencies[li] * zp1)])
#if float(tp[li]) <= float(9e9):
cwave_ac = float(max(1, self._cutoff_wavelength)) * ac
#else:
# cwave_ac = float(1) * ac
# Apply absorption to SED only bluewards of cutoff wavelength
ab = rest_wavs < cwave_ac # noqa: F841
tpi = tp[li] # noqa: F841
rp2i = rp2[li] # noqa: F841
# Exponent of suppresion and rest wavelength should sum to 5
sup_power = float(max(0, self._alpha))
wavs_power = (5 - sup_power)
if not evaled:
# Absorbed blackbody: 0% transmission at 0 Angstroms 100% at
# >3000 Angstroms.
sed = ne.evaluate(
"where(ab, fc * (rp2i / cwave_ac ** sup_power/ "
"rest_wavs ** wavs_power) / expm1(xc / rest_wavs / tpi), "
"fc * (rp2i / rest_wavs ** 5) / "
"expm1(xc / rest_wavs / tpi))")
evaled = True
else:
sed = ne.re_evaluate()
sed[np.isnan(sed)] = 0.0
seds[li] = sed
uniq_times = np.unique(self._times)
tsort = np.argsort(self._times)
uniq_is = np.searchsorted(self._times, uniq_times, sorter=tsort)
sample_wavelengths = np.linspace(100, 100000, 1000)
STEF_CONST = (4.0 * pi * c.sigma_sb).cgs.value
wavelength_list = np.ones(len(self._times)) * np.array(
self._cutoff_wavelength).astype(float) # * ac
power_list = np.ones(len(self._times)) * np.array(
self._alpha).astype(float)
wavelength_list[wavelength_list < 0] = 0
power_list[power_list < 0] = 0
norms = np.array([(R2 * STEF_CONST * T**4) / np.trapz(
bbody(sample_wavelengths, T, R2, wave, power), sample_wavelengths)
for T, R2, wave, power in zip(
tp[uniq_is], rp2[uniq_is],
wavelength_list[uniq_is], power_list[uniq_is])])
# Apply renormalisation
seds *= norms[np.searchsorted(uniq_times, self._times)]
seds = self.add_to_existing_seds(seds, **kwargs)
# Units of `seds` is ergs / s / Angstrom.
return {
'sample_wavelengths': self._sample_wavelengths,
self.key('seds'): seds,
'luminosities_out': self._luminosities,
'power_list': power_list,
'wavelength_list': wavelength_list,
'times_out': self._times
}