Skip to content

Commit

Permalink
debug statement for remote troubleshooting
Browse files Browse the repository at this point in the history
  • Loading branch information
@mirochaj
mirochaj committed Aug 28, 2023
1 parent c6c3344 commit 068c0b6
Showing 1 changed file with 29 additions and 180 deletions.
209 changes: 29 additions & 180 deletions ares/solvers/UniformBackground.py
View file
Expand Up @@ -9,22 +9,27 @@
Description:
"""
import gc
import os
from math import ceil
import re
import types

import numpy as np
from math import ceil

from ..data import ARES
import os, re, types, gc
from ..core import GlobalVolume
from ..util import ParameterFile
from ..static import GlobalVolume
from ..util.Misc import num_freq_bins
from ..util.Math import interp1d
from .OpticalDepth import OpticalDepth
from ..util.Warnings import no_tau_table
from ..physics import Hydrogen, Cosmology
from ..populations.Composite import CompositePopulation
from ..populations.GalaxyAggregate import GalaxyAggregate
from scipy.integrate import quad, romberg, romb, trapz, simps
from ..physics.Constants import ev_per_hz, erg_per_ev, c, E_LyA, E_LL, dnu, h_p
from scipy.integrate import quad, romberg, romb, trapz, simps, cumtrapz
from ..physics.Constants import ev_per_hz, erg_per_ev, c, E_LyA, E_LL, dnu, \
h_p, cm_per_mpc
#from ..util.ReadData import flatten_energies, flatten_flux, split_flux, \
# flatten_emissivities
try:
Expand Down Expand Up @@ -150,7 +155,7 @@ def solve_rte(self):
# As long as within 0.1 eV, call it a match
if not lo_close:
lo_close = np.allclose(band[0],
pop.pf['pop_solve_rte'][0], atol=0.1)
pop.pf['pop_solve_rte'][0], atol=0.3)
if not hi_close:
hi_close = np.allclose(band[1],
pop.pf['pop_solve_rte'][1], atol=0.1)
Expand All @@ -161,7 +166,7 @@ def solve_rte(self):
tmp.append(True)
# Round (close enough?)
elif np.allclose(band, pop.pf['pop_solve_rte'],
atol=1e-1, rtol=0.):
atol=1e-3, rtol=0.):
tmp.append(True)
elif lo_close and hi_close:
tmp.append(True)
Expand Down Expand Up @@ -461,7 +466,8 @@ def _set_grid(self, pop, bands, zi=None, zf=None, nz=None,
tau = [np.zeros([len(z), len(Earr)]) for Earr in E]

if compute_emissivities:
ehat = [self.TabulateEmissivity(z, Earr, pop) for Earr in E]
ehat = [pop.get_tab_emissivity(z, Earr) \
for Earr in E]
else:
ehat = None

Expand All @@ -486,7 +492,7 @@ def _set_grid(self, pop, bands, zi=None, zf=None, nz=None,

# Tabulate emissivity
if compute_emissivities:
ehat = self.TabulateEmissivity(z, E, pop)
ehat = pop.get_tab_emissivity(z, E)
else:
ehat = None

Expand Down Expand Up @@ -586,7 +592,7 @@ def _set_tau(self, z, E, pop):
tau_solver.ionization_history = lambda z: 0.0
elif self.pf['tau_approx'] == 'post_EoR':
tau_solver.ionization_history = lambda z: 1.0
elif type(self.pf['tau_approx']) is types.FunctionType:
elif type(self.pf['tau_approx']) == types.FunctionType:
tau_solver.ionization_history = self.pf['tau_approx']
else:
raise NotImplemented('Unrecognized approx_tau option.')
Expand Down Expand Up @@ -778,7 +784,7 @@ def LymanWernerFlux(self, z, E=None, popid=0, **kwargs):
if not np.any(self.solve_rte[popid]):
norm = c * self.cosm.dtdz(z) / four_pi

rhoLW = pop.PhotonLuminosityDensity(z, Emin=E_LyA, Emax=E_LL) \
rhoLW = pop.get_photon_luminosity_density(z, Emin=E_LyA, Emax=E_LL) \
* (E_LL - 11.18) / (E_LL - E_LyA)

# Crude mean photon energy
Expand Down Expand Up @@ -872,7 +878,8 @@ def LymanAlphaFlux(self, z=None, fluxes=None, popid=0, **kwargs):
if not np.any(self.solve_rte[popid]):
norm = c * self.cosm.dtdz(z) / four_pi

rhoLW = pop.PhotonLuminosityDensity(z, Emin=E_LyA, Emax=E_LL)
rhoLW = pop.get_photon_emissivity(z, band=(E_LyA, E_LL), units='eV')
rhoLW /= cm_per_mpc**3

return norm * (1. + z)**3 * (1. + pop.pf['pop_frec_bar']) * \
rhoLW / dnu
Expand Down Expand Up @@ -906,15 +913,16 @@ def load_sed(self, prefix=None):
print("No $ARES environment variable.")
return None

input_dirs = ['{!s}/input/seds'.format(ARES)]
input_dirs = [os.path.join(ARES, "input", "seds")]

else:
if isinstance(prefix, basestring):
if isinstance(prefix, str):
input_dirs = [prefix]
else:
input_dirs = prefix

guess = '{0!s}/{1!s}.txt'.format(input_dirs[0], fn)
guess_fn = f"{fn}.txt"
guess = os.path.join(input_dirs[0], guess_fn)
self.tabname = guess
if os.path.exists(guess):
return guess
Expand All @@ -928,177 +936,16 @@ def load_sed(self, prefix=None):

# If source properties are right
if re.search(pre, fn1):
good_tab = '{0!s}/{1!s}'.format(input_dir, fn1)
good_tab = os.path.join(input_dir, fn1)

# If number of redshift bins and energy range right...
if re.search(pre, fn1) and re.search(post, fn1):
good_tab = '{0!s}/{1!s}'.format(input_dir, fn1)
good_tab = os.path.join(input_dir, fn1)
break

self.tabname = good_tab
return good_tab

def TabulateEmissivity(self, z, E, pop):
"""
Tabulate emissivity over photon energy and redshift.
For a scalable emissivity, the tabulation is done for the emissivity
in the (EminNorm, EmaxNorm) band because conversion to other bands
can simply be applied after-the-fact. However, if the emissivity is
NOT scalable, then it is tabulated separately in the (10.2, 13.6),
(13.6, 24.6), and X-ray band.
Parameters
----------
z : np.ndarray
Array of redshifts
E : np.ndarray
Array of photon energies [eV]
pop : object
Better be some kind of Galaxy population object.
Returns
-------
A 2-D array, first axis corresponding to redshift, second axis for
photon energy.
Units of emissivity are: erg / s / Hz / cMpc
"""

Nz, Nf = len(z), len(E)

Inu = np.zeros(Nf)

# Special case: delta function SED! Can't normalize a-priori without
# knowing binning, so we do it here.
if pop.src.is_delta:
# This is a little weird. Trapezoidal integration doesn't make
# sense for a delta function, but it's what happens later, so
# insert a factor of a half now so we recover all the flux we
# should.
Inu[-1] = 1.
else:
for i in range(Nf):
Inu[i] = pop.src.Spectrum(E[i])

# Convert to photon *number* (well, something proportional to it)
Inu_hat = Inu / E

# Now, redshift dependent parts
epsilon = np.zeros([Nz, Nf])

#if Nf == 1:
# return epsilon

scalable = pop.is_emissivity_scalable
separable = pop.is_emissivity_separable

H = np.array(list(map(self.cosm.HubbleParameter, z)))

if scalable:
Lbol = pop.Emissivity(z)
for ll in range(Nz):
epsilon[ll,:] = Inu_hat * Lbol[ll] * ev_per_hz / H[ll] \
/ erg_per_ev

else:

# There is only a distinction here for computational
# convenience, really. The LWB gets solved in much more detail
# than the LyC or X-ray backgrounds, so it makes sense
# to keep the different emissivity chunks separate.
ct = 0
for band in [(10.2, 13.6), (13.6, 24.6), None]:

if band is not None:

if pop.src.Emin > band[1]:
continue

if pop.src.Emax < band[0]:
continue

# Remind me of this distinction?
if band is None:
b = pop.full_band
fix = 1.

# Means we already generated the emissivity.
if ct > 0:
continue

else:
b = band

# If aging population, is handled within the pop object.
if not pop.is_aging:
fix = 1. / pop._convert_band(*band)
else:
fix = 1.

in_band = np.logical_and(E >= b[0], E <= b[1])

# Shouldn't be any filled elements yet
if np.any(epsilon[:,in_band==1] > 0):
raise ValueError("Non-zero elements already!")

if not np.any(in_band):
continue

###
# No need for spectral correction in this case, at least
# in Lyman continuum. Treat LWB more carefully.
if pop.is_aging and band == (13.6, 24.6):
fix = 1. / Inu_hat[in_band==1]

elif pop.is_aging and band == (10.2, 13.6):

if pop.pf['pop_synth_lwb_method'] == 0:
# No approximation: loop over energy below
raise NotImplemented('sorry dude')
elif pop.pf['pop_synth_lwb_method'] == 1:
# Assume SED of continuousy-star-forming source.
Inu_hat_p = pop._src_csfr.Spectrum(E[in_band==1]) \
/ E[in_band==1]
fix = Inu_hat_p / Inu_hat[in_band==1][0]
else:
raise NotImplemented('sorry dude')
###

# By definition, rho_L integrates to unity in (b[0], b[1]) band
# BUT, Inu_hat is normalized in (EminNorm, EmaxNorm) band,
# hence the 'fix'.

for ll, redshift in enumerate(z):

if (redshift < self.pf['final_redshift']):
continue
if (redshift < pop.zdead):
continue
if (redshift > pop.zform):
continue
if redshift < self.pf['kill_redshift']:
continue
if redshift > self.pf['first_light_redshift']:
continue

# Use Emissivity here rather than rho_L because only
# GalaxyCohort objects will have a rho_L attribute.
epsilon[ll,in_band==1] = fix \
* pop.Emissivity(redshift, Emin=b[0], Emax=b[1]) \
* ev_per_hz * Inu_hat[in_band==1] / H[ll] / erg_per_ev

#ehat = pop.Emissivity(redshift, Emin=b[0], Emax=b[1])

#if ll == 1:
# print("Set emissivity for pop {} band #{}".format(pop.id_num, band))
## print(f'fix={fix}, raw={ehat} z={redshift}')

ct += 1

return epsilon

def _flux_generator_generic(self, energies, redshifts, ehat, tau=None,
flux0=None, my_id=None, accept_photons=False):
"""
Expand All @@ -1120,6 +967,8 @@ def _flux_generator_generic(self, energies, redshifts, ehat, tau=None,
"""

print('hello', my_id, energies.shape, redshifts.shape, ehat.shape, tau.shape)

# Some stuff we need
x = 1. + redshifts
xsq = x**2
Expand Down Expand Up @@ -1195,9 +1044,9 @@ def _flux_generator_generic(self, energies, redshifts, ehat, tau=None,
# Equivalent to Eq. 25 in Mirocha (2014)
# Less readable, but faster!
flux = c_over_four_pi \
* ((xsq[ll] * trapz_base) * ehat[ll]) \
* ((xsq[ll] * trapz_base) * ehat[ll] / cm_per_mpc**3) \
+ exp_term * (c_over_four_pi * xsq[ll+1] \
* trapz_base * ehat_r[ll] \
* trapz_base * (ehat_r[ll] / cm_per_mpc**3) \
+ np.hstack((flux[1:], [0])) / Rsq)
#+ np.roll(flux, -1) / Rsq)

Expand Down

0 comments on commit 068c0b6

Please sign in to comment.