def __init__(
self,
hubble_constant,
omega_matter,
omega_lambda,
omega_radiation,
omega_curvature,
initial_redshift,
unit_registry=None,
):
Cosmology.__init__(
self,
hubble_constant=hubble_constant,
omega_matter=omega_matter,
omega_lambda=omega_lambda,
omega_radiation=omega_radiation,
omega_curvature=omega_curvature,
unit_registry=unit_registry,
)
self.initial_redshift = initial_redshift
self.initial_time = self.t_from_z(self.initial_redshift)
# time units = 1 / sqrt(4 * pi * G rho_0 * (1 + z_i)**3),
# rho_0 = (3 * Omega_m * h**2) / (8 * pi * G)
self.time_unit = (
(
1.5
* self.omega_matter
* self.hubble_constant**2
* (1 + self.initial_redshift) ** 3
)
** -0.5
).in_units("s")
self.time_unit.units.registry = self.unit_registry
def _set_units(self):
self.unit_registry = UnitRegistry()
self.time_unit = self.quan(1.0, "s")
if self.cosmological_simulation:
# Instantiate Cosmology object for units and time conversions.
self.cosmology = \
Cosmology(hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda,
unit_registry=self.unit_registry)
self.unit_registry.modify("h", self.hubble_constant)
# Comoving lengths
for my_unit in ["m", "pc", "AU", "au"]:
new_unit = "%scm" % my_unit
# technically not true, but should be ok
self.unit_registry.add(new_unit,
self.unit_registry.lut[my_unit][0],
dimensions.length,
"\\rm{%s}/(1+z)" % my_unit)
self.length_unit = self.quan(self.unit_base["UnitLength_in_cm"],
"cmcm / h",
registry=self.unit_registry)
self.box_size *= self.length_unit.in_units("Mpccm / h")
else:
# Read time from file for non-cosmological sim
self.time_unit = self.quan(
self.unit_base["UnitLength_in_cm"]/ \
self.unit_base["UnitVelocity_in_cm_per_s"], "s")
self.unit_registry.add("code_time", 1.0, dimensions.time)
self.unit_registry.modify("code_time", self.time_unit)
# Length
self.length_unit = self.quan(self.unit_base["UnitLength_in_cm"],
"cm")
self.unit_registry.add("code_length", 1.0, dimensions.length)
self.unit_registry.modify("code_length", self.length_unit)
def find_radius_mass(m_r, delta, z=0.0, cosmo=None):
"""
Given a mass profile and an overdensity, find the radius
and mass (e.g. M200, r200)
Parameters
----------
m_r : RadialProfile
The mass profile.
delta : float
The overdensity to compute the mass and radius for.
z : float, optional
The redshift of the halo formation. Default: 0.0
cosmo : yt ``Cosmology`` object
The cosmology to be used when computing the critical
density. If not supplied, a default one from yt will
be used.
"""
from yt.utilities.cosmology import Cosmology
from scipy.optimize import bisect
if cosmo is None:
cosmo = Cosmology()
rho_crit = cosmo.critical_density(z).to_value("Msun/kpc**3")
f = lambda r: 3.0*m_r(r)/(4.*np.pi*r**3) - delta*rho_crit
r_delta = bisect(f, 0.01, 10000.0)
return r_delta, m_r(r_delta)
def _parse_parameter_file(self):
with open(self.parameter_filename, "rb") as f:
hvals = fpu.read_cattrs(f, header_dt)
hvals.pop("unused")
self.dimensionality = 3
self.refine_by = 2
prefix = ".".join(self.parameter_filename.rsplit(".", 2)[:-2])
self.filename_template = "%s.%%(num)s%s" % (prefix, self._suffix)
self.file_count = len(glob.glob(prefix + ".*" + self._suffix))
# Now we can set up things we already know.
self.cosmological_simulation = 1
self.current_redshift = (1.0 / hvals["scale"]) - 1.0
self.hubble_constant = hvals["h0"]
self.omega_lambda = hvals["Ol"]
self.omega_matter = hvals["Om"]
cosmo = Cosmology(
hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda,
)
self.current_time = cosmo.lookback_time(self.current_redshift,
1e6).in_units("s")
self.periodicity = (True, True, True)
self.particle_types = "halos"
self.particle_types_raw = "halos"
self.domain_left_edge = np.array([0.0, 0.0, 0.0])
self.domain_right_edge = np.array([hvals["box_size"]] * 3)
self.domain_dimensions = np.ones(3, "int32")
self.parameters.update(hvals)
def _parse_parameter_file(self):
# Read all parameters.
simu = self._read_log_simu()
param = self._read_parameter()
# Set up general information.
self.filename_template = self.parameter_filename
self.file_count = 1
self.parameters.update(param)
self.particle_types = ('halos')
self.particle_types_raw = ('halos')
# Set up geometrical information.
self.refine_by = 2
self.dimensionality = 3
nz = 1 << self.over_refine_factor
self.domain_dimensions = np.ones(self.dimensionality, "int32") * nz
self.domain_left_edge = np.array([0.0, 0.0, 0.0])
# Note that boxsize is in Mpc but particle positions are in kpc.
self.domain_right_edge = np.array([simu['boxsize']] * 3) * 1000
self.periodicity = (True, True, True)
# Set up cosmological information.
self.cosmological_simulation = 1
self.current_redshift = param['z']
self.omega_lambda = simu['lambda0']
self.omega_matter = simu['omega0']
cosmo = Cosmology(hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda)
self.current_time = cosmo.lookback_time(param['z'], 1e6).in_units('s')
def _load_star_data(self, select='all'):
"""If star is present load Metallicity if present"""
if self.obj.simulation.nstar == 0:
return
if select is 'all':
flag = [True]*self.obj.simulation.nstar
else:
flag = (select>=0)
if has_property(self.obj, 'star', 'metallicity'):
self.sZ = get_property(self.obj, 'metallicity', 'star')[flag]
elif has_property(self.obj, 'star', 'met_tng'): # try Illustris/TNG alias
self.sZ = get_property(self.obj, 'met_tng', 'star')[flag]
#self.sZ = np.sum(self.sZ.T[2:],axis=0) # first two are H,He; the rest sum to give metallicity
#self.sZ[self.sZ<0] = 0. # some (very small) negative values, set to 0
else:
mylog.warning('Metallicity not found: setting all stars to solar=0.0134')
self.sZ = 0.0134*np.ones(self.obj.simulation.nstar,dtype=MY_DTYPE)
ds = self.obj.yt_dataset
if has_property(self.obj, 'star', 'aform'):
self.age = get_property(self.obj, 'aform', 'star')[flag] # a_exp at time of formation
elif has_property(self.obj, 'star', 'aform_tng'): # try Illustris/TNG alias
self.age = get_property(self.obj, 'aform_tng', 'star')[flag]
self.age = abs(self.age) # some negative values here too; not sure what to do?
else:
self.age = np.zeros(self.obj.simulation.nstar,dtype=MY_DTYPE)
mylog.warning('Stellar age not found -- photometry will be incorrect!')
if ds.cosmological_simulation:
from yt.utilities.cosmology import Cosmology
co = Cosmology(hubble_constant=ds.hubble_constant, omega_matter=ds.omega_matter, omega_lambda=ds.omega_lambda)
self.age = (ds.current_time - co.t_from_z(1./self.age-1.)).in_units('Gyr').astype(MY_DTYPE) # age at time of snapshot
def _parse_parameter_file(self):
hvals = self._get_hvals()
self.dimensionality = 3
self.refine_by = 2
self.parameters["HydroMethod"] = "sph"
self.unique_identifier = \
int(os.stat(self.parameter_filename)[stat.ST_CTIME])
# Set standard values
# We may have an overridden bounding box.
if self.domain_left_edge is None:
self.domain_left_edge = np.zeros(3, "float64")
self.domain_right_edge = np.ones(3, "float64") * hvals["BoxSize"]
nz = 1 << self.over_refine_factor
self.domain_dimensions = np.ones(3, "int32") * nz
self.periodicity = (True, True, True)
self.cosmological_simulation = 1
self.current_redshift = hvals["Redshift"]
self.omega_lambda = hvals["OmegaLambda"]
self.omega_matter = hvals["Omega0"]
self.hubble_constant = hvals["HubbleParam"]
# According to the Gadget manual, OmegaLambda will be zero for
# non-cosmological datasets. However, it may be the case that
# individuals are running cosmological simulations *without* Lambda, in
# which case we may be doing something incorrect here.
# It may be possible to deduce whether ComovingIntegration is on
# somehow, but opinions on this vary.
if self.omega_lambda == 0.0:
mylog.info("Omega Lambda is 0.0, so we are turning off Cosmology.")
self.hubble_constant = 1.0 # So that scaling comes out correct
self.cosmological_simulation = 0
self.current_redshift = 0.0
# This may not be correct.
self.current_time = hvals["Time"]
else:
# Now we calculate our time based on the cosmology, because in
# ComovingIntegration hvals["Time"] will in fact be the expansion
# factor, not the actual integration time, so we re-calculate
# global time from our Cosmology.
cosmo = Cosmology(self.hubble_constant,
self.omega_matter, self.omega_lambda)
self.current_time = cosmo.hubble_time(self.current_redshift)
mylog.info("Calculating time from %0.3e to be %0.3e seconds",
hvals["Time"], self.current_time)
self.parameters = hvals
prefix = os.path.abspath(
os.path.join(os.path.dirname(self.parameter_filename),
os.path.basename(self.parameter_filename).split(".", 1)[0]))
if hvals["NumFiles"] > 1:
self.filename_template = "%s.%%(num)s%s" % (prefix, self._suffix)
else:
self.filename_template = self.parameter_filename
self.file_count = hvals["NumFiles"]
def _parse_parameter_file(self):
with h5py.File(self.parameter_filename, mode="r") as f:
self.parameters = \
dict((str(field), val)
for field, val in f["Header"].attrs.items())
self.dimensionality = 3
self.refine_by = 2
# Set standard values
self.domain_left_edge = np.zeros(3, "float64")
self.domain_right_edge = np.ones(3, "float64") * \
self.parameters["BoxSize"]
self.domain_dimensions = np.ones(3, "int32")
self.cosmological_simulation = 1
self.periodicity = (True, True, True)
self.current_redshift = self.parameters["Redshift"]
self.omega_lambda = self.parameters["OmegaLambda"]
self.omega_matter = self.parameters["Omega0"]
self.hubble_constant = self.parameters["HubbleParam"]
cosmology = Cosmology(hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda)
self.current_time = cosmology.t_from_z(self.current_redshift)
prefix = os.path.abspath(
os.path.join(os.path.dirname(self.parameter_filename),
os.path.basename(self.parameter_filename).split(".", 1)[0]))
suffix = self.parameter_filename.rsplit(".", 1)[-1]
self.filename_template = "%s.%%(num)i.%s" % (prefix, suffix)
self.file_count = self.parameters["NumFiles"]
self.particle_types = ("Group", "Subhalo")
self.particle_types_raw = ("Group", "Subhalo")
def find_virial_radius(ds,center):
z = ds.current_redshift
co = Cosmology(hubble_constant=ds.hubble_constant, omega_matter=ds.omega_matter,
omega_lambda=ds.omega_lambda, omega_curvature=0.0)
rho_c = co.critical_density(z=z)
r = 0
rs = []
densities = []
density = np.inf
while density > 200*rho_c:
r+=5
rs.append(r)
sp = ds.sphere(center,(r,'kpc'))
cell_mass,particle_mass = sp.quantities.total_mass()
volume = 4/3*np.pi*ds.arr(r,'kpc')**3
new_density = (cell_mass+particle_mass)/volume
if new_density > density:
print('density not decreasing!')
break
density = new_density.in_units('g/cm**3')
densities.append(density)
toprint = (r,cell_mass+particle_mass,(cell_mass+particle_mass).units,particle_mass,particle_mass.units,cell_mass,cell_mass.units)
rs = np.array(rs)
densities = np.array(densities)
Rvir = np.interp(200*rho_c,np.flip(densities),np.flip(rs))
Mvir = cell_mass+particle_mass
return ds.arr(Rvir,'kpc'),ds.arr(Mvir,'g')
def _parse_parameter_file(self):
with open(self.parameter_filename, "rb") as f:
hvals = fpu.read_cattrs(f, header_dt)
hvals.pop("unused")
self.dimensionality = 3
self.refine_by = 2
self.unique_identifier = int(os.stat(self.parameter_filename)[stat.ST_CTIME])
prefix = ".".join(self.parameter_filename.rsplit(".", 2)[:-2])
self.filename_template = "%s.%%(num)s%s" % (prefix, self._suffix)
self.file_count = len(glob.glob(prefix + ".*" + self._suffix))
# Now we can set up things we already know.
self.cosmological_simulation = 1
self.current_redshift = (1.0 / hvals["scale"]) - 1.0
self.hubble_constant = hvals["h0"]
self.omega_lambda = hvals["Ol"]
self.omega_matter = hvals["Om"]
cosmo = Cosmology(self.hubble_constant, self.omega_matter, self.omega_lambda)
self.current_time = cosmo.hubble_time(self.current_redshift).in_units("s")
self.periodicity = (True, True, True)
self.particle_types = "halos"
self.particle_types_raw = "halos"
self.domain_left_edge = np.array([0.0, 0.0, 0.0])
self.domain_right_edge = np.array([hvals["box_size"]] * 3)
nz = 1 << self.over_refine_factor
self.domain_dimensions = np.ones(3, "int32") * nz
self.parameters.update(hvals)
def nfw_scale_density(conc, z=0.0, delta=200.0, cosmo=None):
"""
Compute a scale density parameter for an NFW profile
given a concentration parameter, and optionally
a redshift, overdensity, and cosmology.
Parameters
----------
conc : float
The concentration parameter for the halo, which should
correspond the selected overdensity (which has a default
of 200).
z : float, optional
The redshift of the halo formation. Default: 0.0
delta : float, optional
The overdensity parameter for which the concentration
is defined. Default: 200.0
cosmo : yt Cosmology object
The cosmology to be used when computing the critical
density. If not supplied, a default one from yt will
be used.
"""
from yt.utilities.cosmology import Cosmology
if cosmo is None:
cosmo = Cosmology()
rho_crit = cosmo.critical_density(z).to_value("Msun/kpc**3")
rho_s = delta*rho_crit*conc**3*_nfw_factor(conc)/3.
return rho_s
def read_caesar(c_file, selected_gals):
print('Reading caesar file...')
sim = caesar.load(c_file, LoadHalo=False)
redshift = np.round(sim.simulation.redshift, decimals=2)
h = sim.simulation.hubble_constant
cosmo = Cosmology(hubble_constant=sim.simulation.hubble_constant,
omega_matter=sim.simulation.omega_matter,
omega_lambda=sim.simulation.omega_lambda,
omega_curvature=0)
thubble = cosmo.hubble_time(redshift).in_units("Gyr")
H0 = (100 * sim.simulation.hubble_constant) * km / s / Mpc
rhocrit = 3. * H0.to('1/s')**2 / (8 * np.pi * sim.simulation.G)
mlim = 32 * rhocrit.to('Msun/kpc**3') * sim.simulation.boxsize.to(
'kpc'
)**3 * sim.simulation.omega_baryon / sim.simulation.effective_resolution**3 / sim.simulation.scale_factor**3 # galaxy mass resolution limit: 32 gas particle masses
gals = np.array([])
# read galaxy particle data for the selected galaxies
if isinstance(selected_gals, np.ndarray):
gals = np.asarray([
i for i in sim.galaxies
if i.GroupID in selected_gals and i.masses['stellar'] > mlim
]) # select resolved galaxies
print('Galaxy data from caesar file extracted and saved.')
return sim, redshift, h, cosmo, thubble, H0, rhocrit, mlim, gals
def _parse_parameter_file(self):
self.dimensionality = 3
self._periodicity = (True, True, True)
self.refine_by = 2
self.unique_identifier = \
int(os.stat(self.parameter_filename)[stat.ST_CTIME])
self.file_count = 1
ct_file = CTHLMiniFile(self, self.parameter_filename)
self.parameters.update(ct_file.parameters)
for par in [
'hubble_constant', 'omega_matter', 'omega_lambda',
'current_redshift'
]:
setattr(self, par, self.parameters.get(par))
self.cosmological_simulation = 1
co = Cosmology(hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda)
self.current_time = co.t_from_z(self.current_redshift)
self.domain_left_edge = np.zeros(self.dimensionality)
self.domain_right_edge = float(self.parameters['box_size'][0]) * \
np.ones(self.dimensionality)
self.domain_dimensions = np.ones(self.dimensionality, "int32")
def _parse_parameter_file(self):
with open(self.parameter_filename, "rb") as f:
hvals = fpu.read_cattrs(f, header_dt)
hvals.pop("unused")
self.dimensionality = 3
self.refine_by = 2
self.unique_identifier = \
int(os.stat(self.parameter_filename)[stat.ST_CTIME])
prefix = ".".join(self.parameter_filename.rsplit(".", 2)[:-2])
self.filename_template = "%s.%%(num)s%s" % (prefix, self._suffix)
self.file_count = len(glob.glob(prefix + ".*" + self._suffix))
# Now we can set up things we already know.
self.cosmological_simulation = 1
self.current_redshift = (1.0 / hvals['scale']) - 1.0
self.hubble_constant = hvals['h0']
self.omega_lambda = hvals['Ol']
self.omega_matter = hvals['Om']
cosmo = Cosmology(self.hubble_constant, self.omega_matter,
self.omega_lambda)
self.current_time = cosmo.hubble_time(
self.current_redshift).in_units("s")
self.periodicity = (True, True, True)
self.particle_types = ("halos")
self.particle_types_raw = ("halos")
self.domain_left_edge = np.array([0.0, 0.0, 0.0])
self.domain_right_edge = np.array([hvals['box_size']] * 3)
nz = 1 << self.over_refine_factor
self.domain_dimensions = np.ones(3, "int32") * nz
self.parameters.update(hvals)
def _set_code_unit_attributes(self):
# First try to set units based on parameter file
if self.cosmological_simulation:
mu = self.parameters.get("dMsolUnit", 1.0)
self.mass_unit = self.quan(mu, "Msun")
lu = self.parameters.get("dKpcUnit", 1000.0)
# In cosmological runs, lengths are stored as length*scale_factor
self.length_unit = self.quan(lu, "kpc") * self.scale_factor
density_unit = self.mass_unit / (self.length_unit / self.scale_factor) ** 3
if "dHubble0" in self.parameters:
# Gasoline's internal hubble constant, dHubble0, is stored in
# units of proper code time
self.hubble_constant *= np.sqrt(G * density_unit)
# Finally, we scale the hubble constant by 100 km/s/Mpc
self.hubble_constant /= self.quan(100, "km/s/Mpc")
# If we leave it as a YTQuantity, the cosmology object
# used below will add units back on.
self.hubble_constant = self.hubble_constant.to_value("")
else:
mu = self.parameters.get("dMsolUnit", 1.0)
self.mass_unit = self.quan(mu, "Msun")
lu = self.parameters.get("dKpcUnit", 1.0)
self.length_unit = self.quan(lu, "kpc")
# If unit base is defined by the user, override all relevant units
if self._unit_base is not None:
for my_unit in ["length", "mass", "time"]:
if my_unit in self._unit_base:
my_val = self._unit_base[my_unit]
my_val = (
self.quan(*my_val)
if isinstance(my_val, tuple)
else self.quan(my_val)
)
setattr(self, f"{my_unit}_unit", my_val)
# Finally, set the dependent units
if self.cosmological_simulation:
cosmo = Cosmology(
hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda,
)
self.current_time = cosmo.lookback_time(self.current_redshift, 1e6)
# mass units are rho_crit(z=0) * domain volume
mu = (
cosmo.critical_density(0.0)
* (1 + self.current_redshift) ** 3
* self.length_unit ** 3
)
self.mass_unit = self.quan(mu.in_units("Msun"), "Msun")
density_unit = self.mass_unit / (self.length_unit / self.scale_factor) ** 3
# need to do this again because we've modified the hubble constant
self.unit_registry.modify("h", self.hubble_constant)
else:
density_unit = self.mass_unit / self.length_unit ** 3
if not hasattr(self, "time_unit"):
self.time_unit = 1.0 / np.sqrt(density_unit * G)
def __init__(self,
ds,
data_source=None,
star_mass=None,
star_creation_time=None,
bins=300,
volume=None,
star_filter=None):
self._ds = ds
self._data_source = data_source
self._filter = star_filter
self.ds_provided = self._data_source is not None
self.filter_provided = self._filter is not None
self.bin_count = bins
# Set up for time conversion.
self.cosm = Cosmology(hubble_constant=self._ds.hubble_constant,
omega_matter=self._ds.omega_matter,
omega_lambda=self._ds.omega_lambda)
# Find the time right now.
self.time_now = self._ds.current_time
if not self.ds_provided:
# Check to make sure we have the right set of informations.
if star_mass is None or star_creation_time is None \
or volume is None:
mylog.error("""
If data_source is not provided, all of these parameters
need to be set:
star_mass (array, Msun),
star_creation_time (array, code units),
volume (float, cMpc**3).""")
return None
if isinstance(star_mass, YTArray):
assert star_mass.units.same_dimensions_as(g.units)
elif star_mass is not None:
star_mass = YTArray(star_mass, 'Msun')
self.star_mass = star_mass
if isinstance(star_creation_time, YTArray):
assert star_creation_time.units.same_dimensions_as(s.units)
elif star_creation_time is not None:
star_creation_time = self._ds.arr(star_creation_time,
'code_time')
self.star_creation_time = star_creation_time
if isinstance(volume, YTQuantity):
assert volume.units.same_dimensions_as(
self._ds.quan(1.0, 'Mpccm**3').units)
elif volume is not None:
volume = self._ds.quan(volume, 'Mpccm**3')
self.volume = volume
# Build the distribution.
self.build_dist()
# Attach some convenience arrays.
self.attach_arrays()
def test_z_t_analytic():
"""
Test z/t conversions against analytic solutions.
"""
cosmos = (
{
"hubble_constant": 0.7,
"omega_matter": 0.3,
"omega_lambda": 0.7
},
{
"hubble_constant": 0.7,
"omega_matter": 1.0,
"omega_lambda": 0.0
},
{
"hubble_constant": 0.7,
"omega_matter": 0.3,
"omega_lambda": 0.0
},
)
for cosmo in cosmos:
omega_curvature = 1 - cosmo["omega_matter"] - cosmo["omega_lambda"]
co = Cosmology(omega_curvature=omega_curvature, **cosmo)
# random sample in log(a) from -6 to 6
my_random = np.random.RandomState(10132324)
la = 12 * my_random.random_sample(1000) - 6
z = 1 / np.power(10, la) - 1
t_an = t_from_z_analytic(z, **cosmo).to("Gyr")
t_co = co.t_from_z(z).to("Gyr")
assert_rel_equal(
t_an,
t_co,
4,
err_msg=
f"t_from_z does not match analytic version for cosmology {cosmo}.",
)
# random sample in log(t/t0) from -3 to 1
t0 = np.power(10, 4 * my_random.random_sample(1000) - 3)
t = (t0 / co.hubble_constant).to("Gyr")
z_an = z_from_t_analytic(t, **cosmo)
z_co = co.z_from_t(t)
# compare scale factors since z approaches 0
assert_rel_equal(
1 / (1 + z_an),
1 / (1 + z_co),
5,
err_msg=
f"z_from_t does not match analytic version for cosmology {cosmo}.",
)
def _set_units(self):
self.unit_registry = UnitRegistry()
self.time_unit = self.quan(1.0, "s")
if self.cosmological_simulation:
# Instantiate Cosmology object for units and time conversions.
self.cosmology = Cosmology(
hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda,
unit_registry=self.unit_registry,
)
if "h" in self.unit_registry:
self.unit_registry.modify("h", self.hubble_constant)
else:
self.unit_registry.add("h", self.hubble_constant,
dimensions.dimensionless)
# Comoving lengths
for my_unit in ["m", "pc", "AU"]:
new_unit = f"{my_unit}cm"
# technically not true, but should be ok
self.unit_registry.add(
new_unit,
self.unit_registry.lut[my_unit][0],
dimensions.length,
"\\rm{%s}/(1+z)" % my_unit,
prefixable=True,
)
self.length_unit = self.quan(
self.unit_base["UnitLength_in_cm"],
"cmcm / h",
registry=self.unit_registry,
)
self.mass_unit = self.quan(self.unit_base["UnitMass_in_g"],
"g / h",
registry=self.unit_registry)
self.box_size = self.box_size * self.length_unit
self.domain_left_edge = self.domain_left_edge * self.length_unit
self.domain_right_edge = self.domain_right_edge * self.length_unit
self.unit_registry.add(
"unitary",
float(self.box_size.in_base()),
self.length_unit.units.dimensions,
)
else:
# Read time from file for non-cosmological sim
self.time_unit = self.quan(
self.unit_base["UnitLength_in_cm"] /
self.unit_base["UnitVelocity_in_cm_per_s"],
"s",
)
self.unit_registry.add("code_time", 1.0, dimensions.time)
self.unit_registry.modify("code_time", self.time_unit)
# Length
self.length_unit = self.quan(self.unit_base["UnitLength_in_cm"],
"cm")
self.unit_registry.add("code_length", 1.0, dimensions.length)
self.unit_registry.modify("code_length", self.length_unit)
def __init__(self, parameter_filename, simulation_type, find_outputs=False):
self.parameter_filename = parameter_filename
self.simulation_type = simulation_type
self.simulation = simulation(parameter_filename, simulation_type,
find_outputs=find_outputs)
self.cosmology = Cosmology(
hubble_constant=(self.simulation.hubble_constant),
omega_matter=self.simulation.omega_matter,
omega_lambda=self.simulation.omega_lambda)
def test_hubble_time():
"""
Make sure hubble_time and t_from_z functions agree.
"""
for i in range(10):
co = Cosmology()
# random sample over interval (-1,100]
z = -101 * np.random.random() + 100
assert_rel_equal(co.hubble_time(z), co.t_from_z(z), 5)
def test_dark_factor():
"""
Test that dark factor returns same value for when not
being used and when w_0 = -1 and w_z = 0.
"""
co = Cosmology(w_0=-1, w_a=0, use_dark_factor=False)
assert_equal(co.get_dark_factor(0), 1.0)
co.use_dark_factor = True
assert_equal(co.get_dark_factor(0), 1.0)
def test_hubble_time():
"""
Make sure hubble_time and t_from_z functions agree.
"""
for i in range(10):
co = Cosmology()
# random sample over interval (-1,100]
z = -101 * np.random.random() + 100
yield assert_rel_equal, co.hubble_time(z), co.t_from_z(z), 5
def test_z_t_conversion():
"""
Make sure t_from_z and z_from_t are consistent.
"""
for i in range(10):
co = Cosmology()
# random sample over interval (-1,100]
z1 = -101 * np.random.random() + 100
t = co.t_from_z(z1)
z2 = co.z_from_t(t)
yield assert_rel_equal, z1, z2, 10
def test_z_t_conversion():
"""
Make sure t_from_z and z_from_t are consistent.
"""
for i in range(10):
co = Cosmology()
# random sample over interval (-1,100]
z1 = -101 * np.random.random() + 100
t = co.t_from_z(z1)
z2 = co.z_from_t(t)
assert_rel_equal(z1, z2, 10)
def __init__(self, hubble_constant, omega_matter, omega_lambda,
omega_curvature, initial_redshift, unit_registry = None):
Cosmology.__init__(self,
hubble_constant=hubble_constant,
omega_matter=omega_matter,
omega_lambda=omega_lambda,
omega_curvature=omega_curvature,
unit_registry=unit_registry)
self.initial_redshift = initial_redshift
# time units = 1 / sqrt(4 * pi * G rho_0 * (1 + z_i)**3),
# rho_0 = (3 * Omega_m * h**2) / (8 * pi * G)
self.time_unit = ((1.5 * self.omega_matter * self.hubble_constant**2 *
(1 + self.initial_redshift)**3)**-0.5).in_units("s")
self.time_unit.units.registry = self.unit_registry
def test_dataset_cosmology_calculator():
"""
Test datasets's cosmology calculator against standalone.
"""
ds = data_dir_load(enzotiny)
co = Cosmology(hubble_constant=ds.hubble_constant,
omega_matter=ds.omega_matter,
omega_lambda=ds.omega_lambda)
v1 = ds.cosmology.comoving_radial_distance(1, 5).to('Mpccm').v
v2 = co.comoving_radial_distance(1, 5).to('Mpccm').v
assert_equal(v1, v2)
def test_z_t_roundtrip():
"""
Make sure t_from_z and z_from_t are consistent.
"""
co = Cosmology()
# random sample in log(a) from -6 to 6
my_random = np.random.RandomState(6132305)
la = 12 * my_random.random_sample(10000) - 6
z1 = 1 / np.power(10, la) - 1
t = co.t_from_z(z1)
z2 = co.z_from_t(t)
assert_rel_equal(z1, z2, 4)
def z_to_Time(z,cosmo=None,verbose=False):
if cosmo is None:
#use Planck 2015 from last column (TT, TE, EE+lowP+lensing+ext) of Table 4 from http://arxiv.org/pdf/1502.01589v2.pdf
from yt.utilities.cosmology import Cosmology
h = 0.6774
om = 0.3089
ol = 0.6911
#behroozi parameters
# h=.7
# om=.27
# ol=1-om
if verbose: print "Assuming a Planck 2015 cosmology (H0 = {0}, Om0 = {1}, OL = {2})".format(h*100,om,ol)
cosmo = Cosmology(hubble_constant=h,omega_matter=om,omega_lambda=ol)
return cosmo.t_from_z(z).in_units('yr')
def _set_units(self):
"""
Set "cm" units for explicitly comoving.
Note, we are using comoving units all the time since
we are dealing with data at multiple redshifts.
"""
for my_unit in ["m", "pc", "AU"]:
new_unit = f"{my_unit}cm"
self.unit_registry.add(
new_unit, self.unit_registry.lut[my_unit][0],
length, self.unit_registry.lut[my_unit][3])
setup = True
for attr in ["hubble_constant",
"omega_matter",
"omega_lambda"]:
if getattr(self, attr) is None:
setup = False
ytreeLogger.warning(
f"{attr} missing from data. "
"Arbor will have no cosmology calculator.")
if setup:
self.cosmology = Cosmology(
hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda,
omega_radiation=self.omega_radiation,
unit_registry=self.unit_registry)
def _set_units(self):
self.unit_registry = UnitRegistry()
self.time_unit = self.quan(1.0, "s")
if self.cosmological_simulation:
# Instantiate Cosmology object for units and time conversions.
self.cosmology = \
Cosmology(hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda,
unit_registry=self.unit_registry)
self.unit_registry.modify("h", self.hubble_constant)
# Comoving lengths
for my_unit in ["m", "pc", "AU", "au"]:
new_unit = "%scm" % my_unit
# technically not true, but should be ok
self.unit_registry.add(
new_unit, self.unit_registry.lut[my_unit][0],
dimensions.length, "\\rm{%s}/(1+z)" % my_unit)
self.length_unit = self.quan(self.unit_base["UnitLength_in_cm"],
"cmcm / h", registry=self.unit_registry)
self.box_size *= self.length_unit.in_units("Mpccm / h")
else:
# Read time from file for non-cosmological sim
self.time_unit = self.quan(
self.unit_base["UnitLength_in_cm"]/ \
self.unit_base["UnitVelocity_in_cm_per_s"], "s")
self.unit_registry.add("code_time", 1.0, dimensions.time)
self.unit_registry.modify("code_time", self.time_unit)
# Length
self.length_unit = self.quan(
self.unit_base["UnitLength_in_cm"],"cm")
self.unit_registry.add("code_length", 1.0, dimensions.length)
self.unit_registry.modify("code_length", self.length_unit)
def realistic_ds(fields, particle_fields, nprocs):
np.random.seed(int(0x4d3d3d3))
units = [base_ds._get_field_info(*f).units for f in fields]
punits = [base_ds._get_field_info('io', f).units for f in particle_fields]
fields = [_strip_ftype(f) for f in fields]
ds = fake_random_ds(16,
fields=fields,
units=units,
nprocs=nprocs,
particle_fields=particle_fields,
particle_field_units=punits,
particles=base_ds.stream_handler.particle_count[0][0])
ds.parameters["HydroMethod"] = "streaming"
ds.parameters["EOSType"] = 1.0
ds.parameters["EOSSoundSpeed"] = 1.0
ds.conversion_factors["Time"] = 1.0
ds.conversion_factors.update(dict((f, 1.0) for f in fields))
ds.gamma = 5.0 / 3.0
ds.current_redshift = 0.0001
ds.cosmological_simulation = 1
ds.hubble_constant = 0.7
ds.omega_matter = 0.27
ds.omega_lambda = 0.73
ds.cosmology = Cosmology(hubble_constant=ds.hubble_constant,
omega_matter=ds.omega_matter,
omega_lambda=ds.omega_lambda,
unit_registry=ds.unit_registry)
return ds
def __init__(self,
ds,
bcdir="",
model="chabrier",
time_now=None,
star_filter=None):
self._ds = ds
if not os.path.isdir(bcdir):
bcdir = os.path.join(ytcfg.get("yt", "test_data_dir"), bcdir)
if not os.path.isdir(bcdir):
raise RuntimeError("Failed to locate %s" % bcdir)
self.bcdir = bcdir
self._filter = star_filter
self.filter_provided = self._filter is not None
if model == "chabrier":
self.model = CHABRIER
elif model == "salpeter":
self.model = SALPETER
# Set up for time conversion.
self.cosm = Cosmology(hubble_constant=self._ds.hubble_constant,
omega_matter=self._ds.omega_matter,
omega_lambda=self._ds.omega_lambda)
# Find the time right now.
if time_now is None:
self.time_now = self._ds.current_time
else:
self.time_now = time_now
# Read the tables.
self.read_bclib()
def test_cosmology_calculator_answers():
"""
Test cosmology calculator functions against previously calculated values.
"""
fn = os.path.join(local_dir, 'cosmology_answers.yml')
data = yaml.load(open(fn, 'r'), Loader=yaml.FullLoader)
cosmologies = data['cosmologies']
functions = data['functions']
for cname, copars in cosmologies.items():
omega_curvature = 1 - copars['omega_matter'] - \
copars['omega_lambda'] - copars['omega_radiation']
cosmology = Cosmology(omega_curvature=omega_curvature, **copars)
for fname, finfo in functions.items():
func = getattr(cosmology, fname)
args = finfo.get('args', [])
val = func(*args)
units = finfo.get('units')
if units is not None:
val.convert_to_units(units)
val = float(val)
err_msg = '%s answer has changed for %s cosmology, old: %f, new: %f.' % \
(fname, cname, finfo['answers'][cname], val)
assert_almost_equal(
val, finfo['answers'][cname], 10,
err_msg=err_msg)
def get_base_ds(nprocs):
fields, units = [], []
for fname, (code_units, aliases, dn) in StreamFieldInfo.known_other_fields:
fields.append(("gas", fname))
units.append(code_units)
ds = fake_random_ds(4,
fields=fields,
units=units,
particles=20,
nprocs=nprocs)
ds.parameters["HydroMethod"] = "streaming"
ds.parameters["EOSType"] = 1.0
ds.parameters["EOSSoundSpeed"] = 1.0
ds.conversion_factors["Time"] = 1.0
ds.conversion_factors.update(dict((f, 1.0) for f in fields))
ds.gamma = 5.0 / 3.0
ds.current_redshift = 0.0001
ds.cosmological_simulation = 1
ds.hubble_constant = 0.7
ds.omega_matter = 0.27
ds.omega_lambda = 0.73
ds.cosmology = Cosmology(hubble_constant=ds.hubble_constant,
omega_matter=ds.omega_matter,
omega_lambda=ds.omega_lambda,
unit_registry=ds.unit_registry)
# ensures field errors are raised during testing
# see FieldInfoContainer.check_derived_fields
ds._field_test_dataset = True
ds.index
return ds
def set_units(self):
"""
Creates the unit registry for this dataset.
"""
from yt.units.dimensions import length
if hasattr(self, "cosmological_simulation") \
and getattr(self, "cosmological_simulation"):
# this dataset is cosmological, so add cosmological units.
self.unit_registry.modify("h", self.hubble_constant)
# Comoving lengths
for my_unit in ["m", "pc", "AU", "au"]:
new_unit = "%scm" % my_unit
self.unit_registry.add(
new_unit, self.unit_registry.lut[my_unit][0] /
(1 + self.current_redshift), length,
"\\rm{%s}/(1+z)" % my_unit)
self.set_code_units()
if hasattr(self, "cosmological_simulation") \
and getattr(self, "cosmological_simulation"):
# this dataset is cosmological, add a cosmology object
setattr(
self, "cosmology",
Cosmology(hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda,
unit_registry=self.unit_registry))
setattr(self, "critical_density",
self.cosmology.critical_density(self.current_redshift))
def z_to_Time(z, cosmo=None, verbose=False):
if cosmo is None:
#use Planck 2015 from last column (TT, TE, EE+lowP+lensing+ext) of Table 4 from http://arxiv.org/pdf/1502.01589v2.pdf
from yt.utilities.cosmology import Cosmology
h = 0.6774
om = 0.3089
ol = 0.6911
#behroozi parameters
# h=.7
# om=.27
# ol=1-om
if verbose:
print "Assuming a Planck 2015 cosmology (H0 = {0}, Om0 = {1}, OL = {2})".format(
h * 100, om, ol)
cosmo = Cosmology(hubble_constant=h, omega_matter=om, omega_lambda=ol)
return cosmo.t_from_z(z).in_units('yr')
def __init__(self, parameter_filename, simulation_type, find_outputs=False):
self.parameter_filename = parameter_filename
self.simulation_type = simulation_type
self.simulation = simulation(parameter_filename, simulation_type,
find_outputs=find_outputs)
self.cosmology = Cosmology(
hubble_constant=(self.simulation.hubble_constant),
omega_matter=self.simulation.omega_matter,
omega_lambda=self.simulation.omega_lambda)
def __init__(self, ds, data_source=None, star_mass=None,
star_creation_time=None, bins=300, volume=None,
star_filter=None):
self._ds = ds
self._data_source = data_source
self._filter = star_filter
self.ds_provided = self._data_source is not None
self.filter_provided = self._filter is not None
self.bin_count = bins
# Set up for time conversion.
self.cosm = Cosmology(
hubble_constant=self._ds.hubble_constant,
omega_matter=self._ds.omega_matter,
omega_lambda=self._ds.omega_lambda)
# Find the time right now.
self.time_now = self._ds.current_time
if not self.ds_provided:
# Check to make sure we have the right set of informations.
if star_mass is None or star_creation_time is None \
or volume is None:
mylog.error("""
If data_source is not provided, all of these parameters
need to be set:
star_mass (array, Msun),
star_creation_time (array, code units),
volume (float, cMpc**3).""")
return None
if isinstance(star_mass, YTArray):
assert star_mass.units.same_dimensions_as(g.units)
elif star_mass is not None:
star_mass = YTArray(star_mass, 'Msun')
self.star_mass = star_mass
if isinstance(star_creation_time, YTArray):
assert star_creation_time.units.same_dimensions_as(s.units)
elif star_creation_time is not None:
star_creation_time = self._ds.arr(star_creation_time,
'code_time')
self.star_creation_time = star_creation_time
if isinstance(volume, YTQuantity):
assert volume.units.same_dimensions_as(
self._ds.quan(1.0, 'Mpccm**3').units
)
elif volume is not None:
volume = self._ds.quan(volume, 'Mpccm**3')
self.volume = volume
# Build the distribution.
self.build_dist()
# Attach some convenience arrays.
self.attach_arrays()
def _parse_parameter_file(self):
handle = h5py.File(self.parameter_filename, mode="r")
hvals = {}
hvals.update((str(k), v) for k, v in handle["/Header"].attrs.items())
hvals["NumFiles"] = hvals["NumFiles"]
self.dimensionality = 3
self.refine_by = 2
self.unique_identifier = \
int(os.stat(self.parameter_filename)[stat.ST_CTIME])
# Set standard values
self.domain_left_edge = np.zeros(3, "float64")
self.domain_right_edge = np.ones(3, "float64") * hvals["BoxSize"]
nz = 1 << self.over_refine_factor
self.domain_dimensions = np.ones(3, "int32") * nz
self.cosmological_simulation = 1
self.periodicity = (True, True, True)
self.current_redshift = hvals["Redshift"]
self.omega_lambda = hvals["OmegaLambda"]
self.omega_matter = hvals["Omega0"]
self.hubble_constant = hvals["HubbleParam"]
cosmology = Cosmology(hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda)
self.current_time = cosmology.t_from_z(self.current_redshift)
self.parameters = hvals
prefix = os.path.abspath(
os.path.join(os.path.dirname(self.parameter_filename),
os.path.basename(self.parameter_filename).split(".", 1)[0]))
suffix = self.parameter_filename.rsplit(".", 1)[-1]
self.filename_template = "%s.%%(num)i.%s" % (prefix, suffix)
self.file_count = len(glob.glob(prefix + "*" + self._suffix))
if self.file_count == 0:
raise YTException(message="No data files found.", ds=self)
self.particle_types = ("Group", "Subhalo")
self.particle_types_raw = ("Group", "Subhalo")
handle.close()
def _set_code_unit_attributes(self):
if self.cosmological_simulation:
mu = self.parameters.get('dMsolUnit', 1.)
lu = self.parameters.get('dKpcUnit', 1000.)
# In cosmological runs, lengths are stored as length*scale_factor
self.length_unit = self.quan(lu, 'kpc')*self.scale_factor
self.mass_unit = self.quan(mu, 'Msun')
density_unit = self.mass_unit/ (self.length_unit/self.scale_factor)**3
# Gasoline's hubble constant, dHubble0, is stored units of proper code time.
self.hubble_constant *= np.sqrt(G.in_units('kpc**3*Msun**-1*s**-2')*density_unit).value/(3.2407793e-18)
cosmo = Cosmology(self.hubble_constant,
self.omega_matter, self.omega_lambda)
self.current_time = cosmo.hubble_time(self.current_redshift)
else:
mu = self.parameters.get('dMsolUnit', 1.0)
self.mass_unit = self.quan(mu, 'Msun')
lu = self.parameters.get('dKpcUnit', 1.0)
self.length_unit = self.quan(lu, 'kpc')
density_unit = self.mass_unit / self.length_unit**3
self.time_unit = 1.0 / np.sqrt(G * density_unit)
def time_to_z(t,cosmo=None,verbose=False): # H0=YTQuantity(70.2,'km/s/Mpc')):
from yt import YTQuantity,YTArray
if cosmo is None:
#use Planck 2015 from last column (TT, TE, EE+lowP+lensing+ext) of Table 4 from http://arxiv.org/pdf/1502.01589v2.pdf
from yt.utilities.cosmology import Cosmology
h = 0.6774
om = 0.3089
ol = 0.6911
#behroozi parameters
# h=.7
# om=.27
# ol=1-om
if verbose: print "Assuming a Planck 2015 cosmology (H0 = {0}, Om0 = {1}, OL = {2})".format(h*100,om,ol)
cosmo = Cosmology(hubble_constant=h,omega_matter=om,omega_lambda=ol)
if type(t) != type(YTQuantity(1,'Gyr')) and type(t) != type(YTArray([1,2,3],'Gyr')):
#then I need to figure out units and wrap in a yt object
if type(t) == type(1.23): #single float
if t < 15: #assume Gyr
t = YTArray(t,'Gyr')
if verbose: print "Assuming time in Gyr"
elif t < 1e11: #assume yr
t = YTArray(t,'yr')
if verbose: print "Assuming time in yr"
else: #then it's probably in seconds
t = YTArray(t,'s')
if verbose: print "Assuming time in seconds"
else:
from numpy import array
t = array(t)
if (t < 15).all():
t = YTArray(t,'Gyr')
if verbose: print "Assuming time in Gyr"
elif (t < 1e11).all(): #assume yr
t = YTArray(t,'yr')
if verbose: print "Assuming time in yr"
else: #then it's probably in seconds
t = YTArray(t,'s')
if verbose: print "Assuming time in seconds"
return cosmo.z_from_t(t)
def __init__(self, parameter_filename, simulation_type=None,
near_redshift=None, far_redshift=None,
use_minimum_datasets=True, deltaz_min=0.0,
minimum_coherent_box_fraction=0.0,
time_data=True, redshift_data=True,
find_outputs=False, load_kwargs=None):
self.near_redshift = near_redshift
self.far_redshift = far_redshift
self.use_minimum_datasets = use_minimum_datasets
self.deltaz_min = deltaz_min
self.minimum_coherent_box_fraction = minimum_coherent_box_fraction
self.parameter_filename = parameter_filename
if load_kwargs is None:
self.load_kwargs = {}
else:
self.load_kwargs = load_kwargs
self.light_ray_solution = []
self._data = {}
# Make a light ray from a single, given dataset.
if simulation_type is None:
ds = load(parameter_filename, **self.load_kwargs)
if ds.cosmological_simulation:
redshift = ds.current_redshift
self.cosmology = Cosmology(
hubble_constant=ds.hubble_constant,
omega_matter=ds.omega_matter,
omega_lambda=ds.omega_lambda,
unit_registry=ds.unit_registry)
else:
redshift = 0.
self.light_ray_solution.append({"filename": parameter_filename,
"redshift": redshift})
# Make a light ray from a simulation time-series.
else:
# Get list of datasets for light ray solution.
CosmologySplice.__init__(self, parameter_filename, simulation_type,
find_outputs=find_outputs)
self.light_ray_solution = \
self.create_cosmology_splice(self.near_redshift, self.far_redshift,
minimal=self.use_minimum_datasets,
deltaz_min=self.deltaz_min,
time_data=time_data,
redshift_data=redshift_data)
class CosmologySplice(object):
"""
Class for splicing together datasets to extend over a
cosmological distance.
"""
def __init__(self, parameter_filename, simulation_type, find_outputs=False):
self.parameter_filename = parameter_filename
self.simulation_type = simulation_type
self.simulation = simulation(parameter_filename, simulation_type,
find_outputs=find_outputs)
self.cosmology = Cosmology(
hubble_constant=(self.simulation.hubble_constant),
omega_matter=self.simulation.omega_matter,
omega_lambda=self.simulation.omega_lambda)
def create_cosmology_splice(self, near_redshift, far_redshift,
minimal=True, deltaz_min=0.0,
time_data=True, redshift_data=True):
r"""Create list of datasets capable of spanning a redshift
interval.
For cosmological simulations, the physical width of the simulation
box corresponds to some \Delta z, which varies with redshift.
Using this logic, one can stitch together a series of datasets to
create a continuous volume or length element from one redshift to
another. This method will return such a list
Parameters
----------
near_redshift : float
The nearest (lowest) redshift in the cosmology splice list.
far_redshift : float
The furthest (highest) redshift in the cosmology splice list.
minimal : bool
If True, the minimum number of datasets is used to connect the
initial and final redshift. If false, the list will contain as
many entries as possible within the redshift
interval.
Default: True.
deltaz_min : float
Specifies the minimum delta z between consecutive datasets
in the returned
list.
Default: 0.0.
time_data : bool
Whether or not to include time outputs when gathering
datasets for time series.
Default: True.
redshift_data : bool
Whether or not to include redshift outputs when gathering
datasets for time series.
Default: True.
Examples
--------
>>> co = CosmologySplice("enzo_tiny_cosmology/32Mpc_32.enzo", "Enzo")
>>> cosmo = co.create_cosmology_splice(1.0, 0.0)
"""
if time_data and redshift_data:
self.splice_outputs = self.simulation.all_outputs
elif time_data:
self.splice_outputs = self.simulation.all_time_outputs
elif redshift_data:
self.splice_outputs = self.simulation.all_redshift_outputs
else:
mylog.error('Both time_data and redshift_data are False.')
return
# Link datasets in list with pointers.
# This is used for connecting datasets together.
for i, output in enumerate(self.splice_outputs):
if i == 0:
output['previous'] = None
output['next'] = self.splice_outputs[i + 1]
elif i == len(self.splice_outputs) - 1:
output['previous'] = self.splice_outputs[i - 1]
output['next'] = None
else:
output['previous'] = self.splice_outputs[i - 1]
output['next'] = self.splice_outputs[i + 1]
# Calculate maximum delta z for each data dump.
self._calculate_deltaz_max()
# Calculate minimum delta z for each data dump.
self._calculate_deltaz_min(deltaz_min=deltaz_min)
cosmology_splice = []
if near_redshift == far_redshift:
self.simulation.get_time_series(redshifts=[near_redshift])
cosmology_splice.append({'time': self.simulation[0].current_time,
'redshift': self.simulation[0].current_redshift,
'filename': os.path.join(self.simulation[0].fullpath,
self.simulation[0].basename),
'next': None})
mylog.info("create_cosmology_splice: Using %s for z = %f ." %
(cosmology_splice[0]['filename'], near_redshift))
return cosmology_splice
# Use minimum number of datasets to go from z_i to z_f.
if minimal:
z_Tolerance = 1e-3
z = far_redshift
# fill redshift space with datasets
while ((z > near_redshift) and
(np.abs(z - near_redshift) > z_Tolerance)):
# For first data dump, choose closest to desired redshift.
if (len(cosmology_splice) == 0):
# Sort data outputs by proximity to current redshift.
self.splice_outputs.sort(key=lambda obj:np.fabs(z - \
obj['redshift']))
cosmology_splice.append(self.splice_outputs[0])
# Move forward from last slice in stack until z > z_max.
else:
current_slice = cosmology_splice[-1]
while current_slice['next'] is not None and \
(z < current_slice['next']['redshift'] or \
np.abs(z - current_slice['next']['redshift']) <
z_Tolerance):
current_slice = current_slice['next']
if current_slice is cosmology_splice[-1]:
near_redshift = cosmology_splice[-1]['redshift'] - \
cosmology_splice[-1]['dz_max']
mylog.error("Cosmology splice incomplete due to insufficient data outputs.")
break
else:
cosmology_splice.append(current_slice)
z = cosmology_splice[-1]['redshift'] - \
cosmology_splice[-1]['dz_max']
# Make light ray using maximum number of datasets (minimum spacing).
else:
# Sort data outputs by proximity to current redsfhit.
self.splice_outputs.sort(key=lambda obj:np.abs(far_redshift -
obj['redshift']))
# For first data dump, choose closest to desired redshift.
cosmology_splice.append(self.splice_outputs[0])
nextOutput = cosmology_splice[-1]['next']
while (nextOutput is not None):
if (nextOutput['redshift'] <= near_redshift):
break
if ((cosmology_splice[-1]['redshift'] - nextOutput['redshift']) >
cosmology_splice[-1]['dz_min']):
cosmology_splice.append(nextOutput)
nextOutput = nextOutput['next']
if (cosmology_splice[-1]['redshift'] -
cosmology_splice[-1]['dz_max']) > near_redshift:
mylog.error("Cosmology splice incomplete due to insufficient data outputs.")
near_redshift = cosmology_splice[-1]['redshift'] - \
cosmology_splice[-1]['dz_max']
mylog.info("create_cosmology_splice: Used %d data dumps to get from z = %f to %f." %
(len(cosmology_splice), far_redshift, near_redshift))
# change the 'next' and 'previous' pointers to point to the correct outputs for the created
# splice
for i, output in enumerate(cosmology_splice):
if len(cosmology_splice) == 1:
output['previous'] = None
output['next'] = None
elif i == 0:
output['previous'] = None
output['next'] = cosmology_splice[i + 1]
elif i == len(cosmology_splice) - 1:
output['previous'] = cosmology_splice[i - 1]
output['next'] = None
else:
output['previous'] = cosmology_splice[i - 1]
output['next'] = cosmology_splice[i + 1]
self.splice_outputs.sort(key=lambda obj: obj['time'])
return cosmology_splice
def plan_cosmology_splice(self, near_redshift, far_redshift,
decimals=3, filename=None,
start_index=0):
r"""Create imaginary list of redshift outputs to maximally
span a redshift interval.
If you want to run a cosmological simulation that will have just
enough data outputs to create a cosmology splice,
this method will calculate a list of redshifts outputs that will
minimally connect a redshift interval.
Parameters
----------
near_redshift : float
The nearest (lowest) redshift in the cosmology splice list.
far_redshift : float
The furthest (highest) redshift in the cosmology splice list.
decimals : int
The decimal place to which the output redshift will be rounded.
If the decimal place in question is nonzero, the redshift will
be rounded up to
ensure continuity of the splice. Default: 3.
filename : string
If provided, a file will be written with the redshift outputs in
the form in which they should be given in the enzo dataset.
Default: None.
start_index : int
The index of the first redshift output. Default: 0.
Examples
--------
>>> from yt.analysis_modules.api import CosmologySplice
>>> my_splice = CosmologySplice('enzo_tiny_cosmology/32Mpc_32.enzo', 'Enzo')
>>> my_splice.plan_cosmology_splice(0.0, 0.1, filename='redshifts.out')
"""
z = far_redshift
outputs = []
while z > near_redshift:
rounded = np.round(z, decimals=decimals)
if rounded - z < 0:
rounded += np.power(10.0, (-1.0*decimals))
z = rounded
deltaz_max = self._deltaz_forward(z, self.simulation.box_size)
outputs.append({'redshift': z, 'dz_max': deltaz_max})
z -= deltaz_max
mylog.info("%d data dumps will be needed to get from z = %f to %f." %
(len(outputs), near_redshift, far_redshift))
if filename is not None:
self.simulation._write_cosmology_outputs(filename, outputs,
start_index,
decimals=decimals)
return outputs
def _calculate_deltaz_max(self):
r"""Calculate delta z that corresponds to full box length going
from z to (z - delta z).
"""
d_Tolerance = 1e-4
max_Iterations = 100
target_distance = self.simulation.box_size
for output in self.splice_outputs:
z = output['redshift']
# Calculate delta z that corresponds to the length of the box
# at a given redshift using Newton's method.
z1 = z
z2 = z1 - 0.1 # just an initial guess
distance1 = self.simulation.quan(0.0, "Mpccm / h")
distance2 = self.cosmology.comoving_radial_distance(z2, z)
iteration = 1
while ((np.abs(distance2-target_distance)/distance2) > d_Tolerance):
m = (distance2 - distance1) / (z2 - z1)
z1 = z2
distance1 = distance2
z2 = ((target_distance - distance2) / m.in_units("Mpccm / h")) + z2
distance2 = self.cosmology.comoving_radial_distance(z2, z)
iteration += 1
if (iteration > max_Iterations):
mylog.error("calculate_deltaz_max: Warning - max iterations " +
"exceeded for z = %f (delta z = %f)." %
(z, np.abs(z2 - z)))
break
output['dz_max'] = np.abs(z2 - z)
def _calculate_deltaz_min(self, deltaz_min=0.0):
r"""Calculate delta z that corresponds to a single top grid pixel
going from z to (z - delta z).
"""
d_Tolerance = 1e-4
max_Iterations = 100
target_distance = self.simulation.box_size / \
self.simulation.domain_dimensions[0]
for output in self.splice_outputs:
z = output['redshift']
# Calculate delta z that corresponds to the length of a
# top grid pixel at a given redshift using Newton's method.
z1 = z
z2 = z1 - 0.01 # just an initial guess
distance1 = self.simulation.quan(0.0, "Mpccm / h")
distance2 = self.cosmology.comoving_radial_distance(z2, z)
iteration = 1
while ((np.abs(distance2 - target_distance) / distance2) > d_Tolerance):
m = (distance2 - distance1) / (z2 - z1)
z1 = z2
distance1 = distance2
z2 = ((target_distance - distance2) / m.in_units("Mpccm / h")) + z2
distance2 = self.cosmology.comoving_radial_distance(z2, z)
iteration += 1
if (iteration > max_Iterations):
mylog.error("calculate_deltaz_max: Warning - max iterations " +
"exceeded for z = %f (delta z = %f)." %
(z, np.abs(z2 - z)))
break
# Use this calculation or the absolute minimum specified by the user.
output['dz_min'] = max(np.abs(z2 - z), deltaz_min)
def _deltaz_forward(self, z, target_distance):
r"""Calculate deltaz corresponding to moving a comoving distance
starting from some redshift.
"""
d_Tolerance = 1e-4
max_Iterations = 100
# Calculate delta z that corresponds to the length of the
# box at a given redshift.
z1 = z
z2 = z1 - 0.1 # just an initial guess
distance1 = self.cosmology.quan(0.0, "Mpccm / h")
distance2 = self.cosmology.comoving_radial_distance(z2, z)
iteration = 1
while ((np.abs(distance2 - target_distance)/distance2) > d_Tolerance):
m = (distance2 - distance1) / (z2 - z1)
z1 = z2
distance1 = distance2
z2 = ((target_distance - distance2) / m.in_units("Mpccm / h")) + z2
distance2 = self.cosmology.comoving_radial_distance(z2, z)
iteration += 1
if (iteration > max_Iterations):
mylog.error("deltaz_forward: Warning - max iterations " +
"exceeded for z = %f (delta z = %f)." %
(z, np.abs(z2 - z)))
break
return np.abs(z2 - z)
class LightRay(CosmologySplice):
"""
LightRay(parameter_filename, simulation_type=None,
near_redshift=None, far_redshift=None,
use_minimum_datasets=True, deltaz_min=0.0,
minimum_coherent_box_fraction=0.0,
time_data=True, redshift_data=True,
find_outputs=False, load_kwargs=None):
Create a LightRay object. A light ray is much like a light cone,
in that it stacks together multiple datasets in order to extend a
redshift interval. Unlike a light cone, which does randomly
oriented projections for each dataset, a light ray consists of
randomly oriented single rays. The purpose of these is to create
synthetic QSO lines of sight.
Light rays can also be made from single datasets.
Once the LightRay object is set up, use LightRay.make_light_ray to
begin making rays. Different randomizations can be created with a
single object by providing different random seeds to make_light_ray.
Parameters
----------
parameter_filename : string
The path to the simulation parameter file or dataset.
simulation_type : optional, string
The simulation type. If None, the first argument is assumed to
refer to a single dataset.
Default: None
near_redshift : optional, float
The near (lowest) redshift for a light ray containing multiple
datasets. Do not use is making a light ray from a single
dataset.
Default: None
far_redshift : optional, float
The far (highest) redshift for a light ray containing multiple
datasets. Do not use is making a light ray from a single
dataset.
Default: None
use_minimum_datasets : optional, bool
If True, the minimum number of datasets is used to connect the
initial and final redshift. If false, the light ray solution
will contain as many entries as possible within the redshift
interval.
Default: True.
deltaz_min : optional, float
Specifies the minimum :math:`\Delta z` between consecutive
datasets in the returned list.
Default: 0.0.
minimum_coherent_box_fraction : optional, float
Used with use_minimum_datasets set to False, this parameter
specifies the fraction of the total box size to be traversed
before rerandomizing the projection axis and center. This
was invented to allow light rays with thin slices to sample
coherent large scale structure, but in practice does not work
so well. Try setting this parameter to 1 and see what happens.
Default: 0.0.
time_data : optional, bool
Whether or not to include time outputs when gathering
datasets for time series.
Default: True.
redshift_data : optional, bool
Whether or not to include redshift outputs when gathering
datasets for time series.
Default: True.
find_outputs : optional, bool
Whether or not to search for datasets in the current
directory.
Default: False.
load_kwargs : optional, dict
Optional dictionary of kwargs to be passed to the "load"
function, appropriate for use of certain frontends. E.g.
Tipsy using "bounding_box"
Gadget using "unit_base", etc.
Default : None
"""
def __init__(self, parameter_filename, simulation_type=None,
near_redshift=None, far_redshift=None,
use_minimum_datasets=True, deltaz_min=0.0,
minimum_coherent_box_fraction=0.0,
time_data=True, redshift_data=True,
find_outputs=False, load_kwargs=None):
self.near_redshift = near_redshift
self.far_redshift = far_redshift
self.use_minimum_datasets = use_minimum_datasets
self.deltaz_min = deltaz_min
self.minimum_coherent_box_fraction = minimum_coherent_box_fraction
self.parameter_filename = parameter_filename
if load_kwargs is None:
self.load_kwargs = {}
else:
self.load_kwargs = load_kwargs
self.light_ray_solution = []
self._data = {}
# Make a light ray from a single, given dataset.
if simulation_type is None:
ds = load(parameter_filename, **self.load_kwargs)
if ds.cosmological_simulation:
redshift = ds.current_redshift
self.cosmology = Cosmology(
hubble_constant=ds.hubble_constant,
omega_matter=ds.omega_matter,
omega_lambda=ds.omega_lambda,
unit_registry=ds.unit_registry)
else:
redshift = 0.
self.light_ray_solution.append({"filename": parameter_filename,
"redshift": redshift})
# Make a light ray from a simulation time-series.
else:
# Get list of datasets for light ray solution.
CosmologySplice.__init__(self, parameter_filename, simulation_type,
find_outputs=find_outputs)
self.light_ray_solution = \
self.create_cosmology_splice(self.near_redshift, self.far_redshift,
minimal=self.use_minimum_datasets,
deltaz_min=self.deltaz_min,
time_data=time_data,
redshift_data=redshift_data)
def _calculate_light_ray_solution(self, seed=None,
start_position=None, end_position=None,
trajectory=None, filename=None):
"Create list of datasets to be added together to make the light ray."
# Calculate dataset sizes, and get random dataset axes and centers.
np.random.seed(seed)
# If using only one dataset, set start and stop manually.
if start_position is not None:
if len(self.light_ray_solution) > 1:
raise RuntimeError("LightRay Error: cannot specify start_position " + \
"if light ray uses more than one dataset.")
if not ((end_position is None) ^ (trajectory is None)):
raise RuntimeError("LightRay Error: must specify either end_position " + \
"or trajectory, but not both.")
self.light_ray_solution[0]['start'] = np.array(start_position)
if end_position is not None:
self.light_ray_solution[0]['end'] = np.array(end_position)
else:
# assume trajectory given as r, theta, phi
if len(trajectory) != 3:
raise RuntimeError("LightRay Error: trajectory must have length 3.")
r, theta, phi = trajectory
self.light_ray_solution[0]['end'] = self.light_ray_solution[0]['start'] + \
r * np.array([np.cos(phi) * np.sin(theta),
np.sin(phi) * np.sin(theta),
np.cos(theta)])
self.light_ray_solution[0]['traversal_box_fraction'] = \
vector_length(self.light_ray_solution[0]['start'],
self.light_ray_solution[0]['end'])
# the normal way (random start positions and trajectories for each dataset)
else:
# For box coherence, keep track of effective depth travelled.
box_fraction_used = 0.0
for q in range(len(self.light_ray_solution)):
if (q == len(self.light_ray_solution) - 1):
z_next = self.near_redshift
else:
z_next = self.light_ray_solution[q+1]['redshift']
# Calculate fraction of box required for a depth of delta z
self.light_ray_solution[q]['traversal_box_fraction'] = \
self.cosmology.comoving_radial_distance(z_next, \
self.light_ray_solution[q]['redshift']).in_units("Mpccm / h") / \
self.simulation.box_size
# Simple error check to make sure more than 100% of box depth
# is never required.
if (self.light_ray_solution[q]['traversal_box_fraction'] > 1.0):
mylog.error("Warning: box fraction required to go from z = %f to %f is %f" %
(self.light_ray_solution[q]['redshift'], z_next,
self.light_ray_solution[q]['traversal_box_fraction']))
mylog.error("Full box delta z is %f, but it is %f to the next data dump." %
(self.light_ray_solution[q]['dz_max'],
self.light_ray_solution[q]['redshift']-z_next))
# Get dataset axis and center.
# If using box coherence, only get start point and vector if
# enough of the box has been used,
# or if box_fraction_used will be greater than 1 after this slice.
if (q == 0) or (self.minimum_coherent_box_fraction == 0) or \
(box_fraction_used >
self.minimum_coherent_box_fraction) or \
(box_fraction_used +
self.light_ray_solution[q]['traversal_box_fraction'] > 1.0):
# Random start point
self.light_ray_solution[q]['start'] = np.random.random(3)
theta = np.pi * np.random.random()
phi = 2 * np.pi * np.random.random()
box_fraction_used = 0.0
else:
# Use end point of previous segment and same theta and phi.
self.light_ray_solution[q]['start'] = \
self.light_ray_solution[q-1]['end'][:]
self.light_ray_solution[q]['end'] = \
self.light_ray_solution[q]['start'] + \
self.light_ray_solution[q]['traversal_box_fraction'] * \
np.array([np.cos(phi) * np.sin(theta),
np.sin(phi) * np.sin(theta),
np.cos(theta)])
box_fraction_used += \
self.light_ray_solution[q]['traversal_box_fraction']
if filename is not None:
self._write_light_ray_solution(filename,
extra_info={'parameter_filename':self.parameter_filename,
'random_seed':seed,
'far_redshift':self.far_redshift,
'near_redshift':self.near_redshift})
def make_light_ray(self, seed=None,
start_position=None, end_position=None,
trajectory=None,
fields=None, setup_function=None,
solution_filename=None, data_filename=None,
get_los_velocity=True, redshift=None,
njobs=-1):
"""
make_light_ray(seed=None, start_position=None, end_position=None,
trajectory=None, fields=None, setup_function=None,
solution_filename=None, data_filename=None,
get_los_velocity=True, redshift=None,
njobs=-1)
Create a light ray and get field values for each lixel. A light
ray consists of a list of field values for cells intersected by
the ray and the path length of the ray through those cells.
Light ray data can be written out to an hdf5 file.
Parameters
----------
seed : optional, int
Seed for the random number generator.
Default: None.
start_position : optional, list of floats
Used only if creating a light ray from a single dataset.
The coordinates of the starting position of the ray.
Default: None.
end_position : optional, list of floats
Used only if creating a light ray from a single dataset.
The coordinates of the ending position of the ray.
Default: None.
trajectory : optional, list of floats
Used only if creating a light ray from a single dataset.
The (r, theta, phi) direction of the light ray. Use either
end_position or trajectory, not both.
Default: None.
fields : optional, list
A list of fields for which to get data.
Default: None.
setup_function : optional, callable, accepts a ds
This function will be called on each dataset that is loaded
to create the light ray. For, example, this can be used to
add new derived fields.
Default: None.
solution_filename : optional, string
Path to a text file where the trajectories of each
subray is written out.
Default: None.
data_filename : optional, string
Path to output file for ray data.
Default: None.
get_los_velocity : optional, bool
If True, the line of sight velocity is calculated for
each point in the ray.
Default: True.
redshift : optional, float
Used with light rays made from single datasets to specify a
starting redshift for the ray. If not used, the starting
redshift will be 0 for a non-cosmological dataset and
the dataset redshift for a cosmological dataset.
Default: None.
njobs : optional, int
The number of parallel jobs over which the segments will
be split. Choose -1 for one processor per segment.
Default: -1.
Examples
--------
Make a light ray from multiple datasets:
>>> import yt
>>> from yt.analysis_modules.cosmological_observation.light_ray.api import \
... LightRay
>>> my_ray = LightRay("enzo_tiny_cosmology/32Mpc_32.enzo", "Enzo",
... 0., 0.1, time_data=False)
...
>>> my_ray.make_light_ray(seed=12345,
... solution_filename="solution.txt",
... data_filename="my_ray.h5",
... fields=["temperature", "density"],
... get_los_velocity=True)
Make a light ray from a single dataset:
>>> import yt
>>> from yt.analysis_modules.cosmological_observation.light_ray.api import \
... LightRay
>>> my_ray = LightRay("IsolatedGalaxy/galaxy0030/galaxy0030")
...
>>> my_ray.make_light_ray(start_position=[0., 0., 0.],
... end_position=[1., 1., 1.],
... solution_filename="solution.txt",
... data_filename="my_ray.h5",
... fields=["temperature", "density"],
... get_los_velocity=True)
"""
# Calculate solution.
self._calculate_light_ray_solution(seed=seed,
start_position=start_position,
end_position=end_position,
trajectory=trajectory,
filename=solution_filename)
# Initialize data structures.
self._data = {}
if fields is None: fields = []
data_fields = fields[:]
all_fields = fields[:]
all_fields.extend(['dl', 'dredshift', 'redshift'])
if get_los_velocity:
all_fields.extend(['velocity_x', 'velocity_y',
'velocity_z', 'velocity_los'])
data_fields.extend(['velocity_x', 'velocity_y', 'velocity_z'])
all_ray_storage = {}
for my_storage, my_segment in parallel_objects(self.light_ray_solution,
storage=all_ray_storage,
njobs=njobs):
# Load dataset for segment.
ds = load(my_segment['filename'], **self.load_kwargs)
my_segment['unique_identifier'] = ds.unique_identifier
if redshift is not None:
if ds.cosmological_simulation and redshift != ds.current_redshift:
mylog.warn("Generating light ray with different redshift than " +
"the dataset itself.")
my_segment["redshift"] = redshift
if setup_function is not None:
setup_function(ds)
if start_position is not None:
my_segment["start"] = ds.arr(my_segment["start"], "code_length")
my_segment["end"] = ds.arr(my_segment["end"], "code_length")
else:
my_segment["start"] = ds.domain_width * my_segment["start"] + \
ds.domain_left_edge
my_segment["end"] = ds.domain_width * my_segment["end"] + \
ds.domain_left_edge
if not ds.cosmological_simulation:
next_redshift = my_segment["redshift"]
elif self.near_redshift == self.far_redshift:
next_redshift = my_segment["redshift"] - \
self._deltaz_forward(my_segment["redshift"],
ds.domain_width[0].in_units("Mpccm / h") *
my_segment["traversal_box_fraction"])
elif my_segment.get("next", None) is None:
next_redshift = self.near_redshift
else:
next_redshift = my_segment['next']['redshift']
mylog.info("Getting segment at z = %s: %s to %s." %
(my_segment['redshift'], my_segment['start'],
my_segment['end']))
# Break periodic ray into non-periodic segments.
sub_segments = periodic_ray(my_segment['start'], my_segment['end'],
left=ds.domain_left_edge,
right=ds.domain_right_edge)
# Prepare data structure for subsegment.
sub_data = {}
sub_data['segment_redshift'] = my_segment['redshift']
for field in all_fields:
sub_data[field] = []
# Get data for all subsegments in segment.
for sub_segment in sub_segments:
mylog.info("Getting subsegment: %s to %s." %
(list(sub_segment[0]), list(sub_segment[1])))
sub_ray = ds.ray(sub_segment[0], sub_segment[1])
asort = np.argsort(sub_ray["t"])
sub_data['dl'].extend(sub_ray['dts'][asort] *
vector_length(sub_ray.start_point,
sub_ray.end_point))
for field in data_fields:
sub_data[field].extend(sub_ray[field][asort])
if get_los_velocity:
line_of_sight = sub_segment[1] - sub_segment[0]
line_of_sight /= ((line_of_sight**2).sum())**0.5
sub_vel = ds.arr([sub_ray['velocity_x'],
sub_ray['velocity_y'],
sub_ray['velocity_z']])
sub_data['velocity_los'].extend((np.rollaxis(sub_vel, 1) *
line_of_sight).sum(axis=1)[asort])
del sub_vel
sub_ray.clear_data()
del sub_ray, asort
for key in sub_data:
sub_data[key] = ds.arr(sub_data[key]).in_cgs()
# Get redshift for each lixel. Assume linear relation between l and z.
sub_data['dredshift'] = (my_segment['redshift'] - next_redshift) * \
(sub_data['dl'] / vector_length(my_segment['start'],
my_segment['end']).in_cgs())
sub_data['redshift'] = my_segment['redshift'] - \
sub_data['dredshift'].cumsum() + sub_data['dredshift']
# Remove empty lixels.
sub_dl_nonzero = sub_data['dl'].nonzero()
for field in all_fields:
sub_data[field] = sub_data[field][sub_dl_nonzero]
del sub_dl_nonzero
# Add to storage.
my_storage.result = sub_data
del ds
# Reconstruct ray data from parallel_objects storage.
all_data = [my_data for my_data in all_ray_storage.values()]
# This is now a list of segments where each one is a dictionary
# with all the fields.
all_data.sort(key=lambda a:a['segment_redshift'], reverse=True)
# Flatten the list into a single dictionary containing fields
# for the whole ray.
all_data = _flatten_dict_list(all_data, exceptions=['segment_redshift'])
if data_filename is not None:
self._write_light_ray(data_filename, all_data)
self._data = all_data
return all_data
@parallel_root_only
def _write_light_ray(self, filename, data):
"""
_write_light_ray(filename, data)
Write light ray data to hdf5 file.
"""
mylog.info("Saving light ray data to %s." % filename)
output = h5py.File(filename, 'w')
for field in data.keys():
# if the field is a tuple, only use the second part of the tuple
# in the hdf5 output (i.e. ('gas', 'density') -> 'density')
if isinstance(field, tuple):
fieldname = field[1]
else:
fieldname = field
output.create_dataset(fieldname, data=data[field])
output[fieldname].attrs["units"] = str(data[field].units)
output.close()
@parallel_root_only
def _write_light_ray_solution(self, filename, extra_info=None):
"""
_write_light_ray_solution(filename, extra_info=None)
Write light ray solution to a file.
"""
mylog.info("Writing light ray solution to %s." % filename)
f = open(filename, 'w')
if extra_info is not None:
for par, val in extra_info.items():
f.write("%s = %s\n" % (par, val))
f.write("\nSegment Redshift dl/box Start x y " + \
"z End x y z Dataset\n")
for q, my_segment in enumerate(self.light_ray_solution):
f.write("%04d %.6f %.6f % .10f % .10f % .10f % .10f % .10f % .10f %s\n" % \
(q, my_segment['redshift'], my_segment['traversal_box_fraction'],
my_segment['start'][0], my_segment['start'][1], my_segment['start'][2],
my_segment['end'][0], my_segment['end'][1], my_segment['end'][2],
my_segment['filename']))
f.close()
class GadgetSimulation(SimulationTimeSeries):
r"""
Initialize an Gadget Simulation object.
Upon creation, the parameter file is parsed and the time and redshift
are calculated and stored in all_outputs. A time units dictionary is
instantiated to allow for time outputs to be requested with physical
time units. The get_time_series can be used to generate a
DatasetSeries object.
parameter_filename : str
The simulation parameter file.
find_outputs : bool
If True, the OutputDir directory is searched for datasets.
Time and redshift information are gathered by temporarily
instantiating each dataset. This can be used when simulation
data was created in a non-standard way, making it difficult
to guess the corresponding time and redshift information.
Default: False.
Examples
--------
>>> import yt
>>> gs = yt.simulation("my_simulation.par", "Gadget")
>>> gs.get_time_series()
>>> for ds in gs:
... print ds.current_time
"""
def __init__(self, parameter_filename, find_outputs=False):
self.simulation_type = "particle"
self.dimensionality = 3
SimulationTimeSeries.__init__(self, parameter_filename,
find_outputs=find_outputs)
def _set_units(self):
self.unit_registry = UnitRegistry()
self.time_unit = self.quan(1.0, "s")
if self.cosmological_simulation:
# Instantiate Cosmology object for units and time conversions.
self.cosmology = \
Cosmology(hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda,
unit_registry=self.unit_registry)
self.unit_registry.modify("h", self.hubble_constant)
# Comoving lengths
for my_unit in ["m", "pc", "AU", "au"]:
new_unit = "%scm" % my_unit
# technically not true, but should be ok
self.unit_registry.add(
new_unit, self.unit_registry.lut[my_unit][0],
dimensions.length, "\\rm{%s}/(1+z)" % my_unit)
self.length_unit = self.quan(self.unit_base["UnitLength_in_cm"],
"cmcm / h", registry=self.unit_registry)
self.box_size *= self.length_unit.in_units("Mpccm / h")
else:
# Read time from file for non-cosmological sim
self.time_unit = self.quan(
self.unit_base["UnitLength_in_cm"]/ \
self.unit_base["UnitVelocity_in_cm_per_s"], "s")
self.unit_registry.add("code_time", 1.0, dimensions.time)
self.unit_registry.modify("code_time", self.time_unit)
# Length
self.length_unit = self.quan(
self.unit_base["UnitLength_in_cm"],"cm")
self.unit_registry.add("code_length", 1.0, dimensions.length)
self.unit_registry.modify("code_length", self.length_unit)
def get_time_series(self, initial_time=None, final_time=None,
initial_redshift=None, final_redshift=None,
times=None, redshifts=None, tolerance=None,
parallel=True, setup_function=None):
"""
Instantiate a DatasetSeries object for a set of outputs.
If no additional keywords given, a DatasetSeries object will be
created with all potential datasets created by the simulation.
Outputs can be gather by specifying a time or redshift range
(or combination of time and redshift), with a specific list of
times or redshifts), or by simply searching all subdirectories
within the simulation directory.
initial_time : tuple of type (float, str)
The earliest time for outputs to be included. This should be
given as the value and the string representation of the units.
For example, (5.0, "Gyr"). If None, the initial time of the
simulation is used. This can be used in combination with
either final_time or final_redshift.
Default: None.
final_time : tuple of type (float, str)
The latest time for outputs to be included. This should be
given as the value and the string representation of the units.
For example, (13.7, "Gyr"). If None, the final time of the
simulation is used. This can be used in combination with either
initial_time or initial_redshift.
Default: None.
times : tuple of type (float array, str)
A list of times for which outputs will be found and the units
of those values. For example, ([0, 1, 2, 3], "s").
Default: None.
initial_redshift : float
The earliest redshift for outputs to be included. If None,
the initial redshift of the simulation is used. This can be
used in combination with either final_time or
final_redshift.
Default: None.
final_redshift : float
The latest redshift for outputs to be included. If None,
the final redshift of the simulation is used. This can be
used in combination with either initial_time or
initial_redshift.
Default: None.
redshifts : array_like
A list of redshifts for which outputs will be found.
Default: None.
tolerance : float
Used in combination with "times" or "redshifts" keywords,
this is the tolerance within which outputs are accepted
given the requested times or redshifts. If None, the
nearest output is always taken.
Default: None.
parallel : bool/int
If True, the generated DatasetSeries will divide the work
such that a single processor works on each dataset. If an
integer is supplied, the work will be divided into that
number of jobs.
Default: True.
setup_function : callable, accepts a ds
This function will be called whenever a dataset is loaded.
Examples
--------
>>> import yt
>>> gs = yt.simulation("my_simulation.par", "Gadget")
>>> gs.get_time_series(initial_redshift=10, final_time=(13.7, "Gyr"))
>>> gs.get_time_series(redshifts=[3, 2, 1, 0])
>>> # after calling get_time_series
>>> for ds in gs.piter():
... p = ProjectionPlot(ds, "x", "density")
... p.save()
>>> # An example using the setup_function keyword
>>> def print_time(ds):
... print ds.current_time
>>> gs.get_time_series(setup_function=print_time)
>>> for ds in gs:
... SlicePlot(ds, "x", "Density").save()
"""
if (initial_redshift is not None or \
final_redshift is not None) and \
not self.cosmological_simulation:
raise InvalidSimulationTimeSeries(
"An initial or final redshift has been given for a " +
"noncosmological simulation.")
my_all_outputs = self.all_outputs
if not my_all_outputs:
DatasetSeries.__init__(self, outputs=[], parallel=parallel,
unit_base=self.unit_base)
mylog.info("0 outputs loaded into time series.")
return
# Apply selection criteria to the set.
if times is not None:
my_outputs = self._get_outputs_by_key("time", times,
tolerance=tolerance,
outputs=my_all_outputs)
elif redshifts is not None:
my_outputs = self._get_outputs_by_key("redshift",
redshifts, tolerance=tolerance,
outputs=my_all_outputs)
else:
if initial_time is not None:
if isinstance(initial_time, float):
initial_time = self.quan(initial_time, "code_time")
elif isinstance(initial_time, tuple) and len(initial_time) == 2:
initial_time = self.quan(*initial_time)
elif not isinstance(initial_time, YTArray):
raise RuntimeError(
"Error: initial_time must be given as a float or " +
"tuple of (value, units).")
elif initial_redshift is not None:
my_initial_time = self.cosmology.t_from_z(initial_redshift)
else:
my_initial_time = self.initial_time
if final_time is not None:
if isinstance(final_time, float):
final_time = self.quan(final_time, "code_time")
elif isinstance(final_time, tuple) and len(final_time) == 2:
final_time = self.quan(*final_time)
elif not isinstance(final_time, YTArray):
raise RuntimeError(
"Error: final_time must be given as a float or " +
"tuple of (value, units).")
my_final_time = final_time.in_units("s")
elif final_redshift is not None:
my_final_time = self.cosmology.t_from_z(final_redshift)
else:
my_final_time = self.final_time
my_initial_time.convert_to_units("s")
my_final_time.convert_to_units("s")
my_times = np.array([a["time"] for a in my_all_outputs])
my_indices = np.digitize([my_initial_time, my_final_time], my_times)
if my_initial_time == my_times[my_indices[0] - 1]: my_indices[0] -= 1
my_outputs = my_all_outputs[my_indices[0]:my_indices[1]]
init_outputs = []
for output in my_outputs:
if os.path.exists(output["filename"]):
init_outputs.append(output["filename"])
if len(init_outputs) == 0 and len(my_outputs) > 0:
mylog.warn("Could not find any datasets. " +
"Check the value of OutputDir in your parameter file.")
DatasetSeries.__init__(self, outputs=init_outputs, parallel=parallel,
setup_function=setup_function,
unit_base=self.unit_base)
mylog.info("%d outputs loaded into time series.", len(init_outputs))
def _parse_parameter_file(self):
"""
Parses the parameter file and establishes the various
dictionaries.
"""
self.unit_base = {}
# Let's read the file
lines = open(self.parameter_filename).readlines()
comments = ["%", ";"]
for line in (l.strip() for l in lines):
for comment in comments:
if comment in line: line = line[0:line.find(comment)]
if len(line) < 2: continue
param, vals = (i.strip() for i in line.split(None, 1))
# First we try to decipher what type of value it is.
vals = vals.split()
# Special case approaching.
if "(do" in vals: vals = vals[:1]
if len(vals) == 0:
pcast = str # Assume NULL output
else:
v = vals[0]
# Figure out if it's castable to floating point:
try:
float(v)
except ValueError:
pcast = str
else:
if any("." in v or "e" in v for v in vals):
pcast = float
elif v == "inf":
pcast = str
else:
pcast = int
# Now we figure out what to do with it.
if param.startswith("Unit"):
self.unit_base[param] = float(vals[0])
if len(vals) == 0:
vals = ""
elif len(vals) == 1:
vals = pcast(vals[0])
else:
vals = np.array([pcast(i) for i in vals])
self.parameters[param] = vals
if self.parameters["ComovingIntegrationOn"]:
cosmo_attr = {"box_size": "BoxSize",
"omega_lambda": "OmegaLambda",
"omega_matter": "Omega0",
"hubble_constant": "HubbleParam"}
self.initial_redshift = 1.0 / self.parameters["TimeBegin"] - 1.0
self.final_redshift = 1.0 / self.parameters["TimeMax"] - 1.0
self.cosmological_simulation = 1
for a, v in cosmo_attr.items():
if not v in self.parameters:
raise MissingParameter(self.parameter_filename, v)
setattr(self, a, self.parameters[v])
else:
self.cosmological_simulation = 0
self.omega_lambda = self.omega_matter = \
self.hubble_constant = 0.0
def _snapshot_format(self, index=None):
"""
The snapshot filename for a given index. Modify this for different
naming conventions.
"""
if self.parameters["OutputDir"].startswith("/"):
data_dir = self.parameters["OutputDir"]
else:
data_dir = os.path.join(self.directory,
self.parameters["OutputDir"])
if self.parameters["NumFilesPerSnapshot"] > 1:
suffix = ".0"
else:
suffix = ""
if self.parameters["SnapFormat"] == 3:
suffix += ".hdf5"
if index is None:
count = "*"
else:
count = "%03d" % index
filename = "%s_%s%s" % (self.parameters["SnapshotFileBase"],
count, suffix)
return os.path.join(data_dir, filename)
def _get_all_outputs(self, find_outputs=False):
"""
Get all potential datasets and combine into a time-sorted list.
"""
# Create the set of outputs from which further selection will be done.
if find_outputs:
self._find_outputs()
else:
if self.parameters["OutputListOn"]:
a_values = [float(a) for a in
file(self.parameters["OutputListFilename"], "r").readlines()]
else:
a_values = [float(self.parameters["TimeOfFirstSnapshot"])]
time_max = float(self.parameters["TimeMax"])
while a_values[-1] < time_max:
if self.cosmological_simulation:
a_values.append(
a_values[-1] * self.parameters["TimeBetSnapshot"])
else:
a_values.append(
a_values[-1] + self.parameters["TimeBetSnapshot"])
if a_values[-1] > time_max:
a_values[-1] = time_max
if self.cosmological_simulation:
self.all_outputs = \
[{"filename": self._snapshot_format(i),
"redshift": (1. / a - 1)}
for i, a in enumerate(a_values)]
# Calculate times for redshift outputs.
for output in self.all_outputs:
output["time"] = self.cosmology.t_from_z(output["redshift"])
else:
self.all_outputs = \
[{"filename": self._snapshot_format(i),
"time": self.quan(a, "code_time")}
for i, a in enumerate(a_values)]
self.all_outputs.sort(key=lambda obj:obj["time"].to_ndarray())
def _calculate_simulation_bounds(self):
"""
Figure out the starting and stopping time and redshift for the simulation.
"""
# Convert initial/final redshifts to times.
if self.cosmological_simulation:
self.initial_time = self.cosmology.t_from_z(self.initial_redshift)
self.initial_time.units.registry = self.unit_registry
self.final_time = self.cosmology.t_from_z(self.final_redshift)
self.final_time.units.registry = self.unit_registry
# If not a cosmology simulation, figure out the stopping criteria.
else:
if "TimeBegin" in self.parameters:
self.initial_time = self.quan(self.parameters["TimeBegin"], "code_time")
else:
self.initial_time = self.quan(0., "code_time")
if "TimeMax" in self.parameters:
self.final_time = self.quan(self.parameters["TimeMax"], "code_time")
else:
self.final_time = None
if not "TimeMax" in self.parameters:
raise NoStoppingCondition(self.parameter_filename)
def _find_outputs(self):
"""
Search for directories matching the data dump keywords.
If found, get dataset times py opening the ds.
"""
potential_outputs = glob.glob(self._snapshot_format())
self.all_outputs = self._check_for_outputs(potential_outputs)
self.all_outputs.sort(key=lambda obj: obj["time"])
only_on_root(mylog.info, "Located %d total outputs.", len(self.all_outputs))
# manually set final time and redshift with last output
if self.all_outputs:
self.final_time = self.all_outputs[-1]["time"]
if self.cosmological_simulation:
self.final_redshift = self.all_outputs[-1]["redshift"]
def _check_for_outputs(self, potential_outputs):
r"""
Check a list of files to see if they are valid datasets.
"""
only_on_root(mylog.info, "Checking %d potential outputs.",
len(potential_outputs))
my_outputs = {}
for my_storage, output in parallel_objects(potential_outputs,
storage=my_outputs):
if os.path.exists(output):
try:
ds = load(output)
if ds is not None:
my_storage.result = {"filename": output,
"time": ds.current_time.in_units("s")}
if ds.cosmological_simulation:
my_storage.result["redshift"] = ds.current_redshift
except YTOutputNotIdentified:
mylog.error("Failed to load %s", output)
my_outputs = [my_output for my_output in my_outputs.values() \
if my_output is not None]
return my_outputs
def _write_cosmology_outputs(self, filename, outputs, start_index,
decimals=3):
r"""
Write cosmology output parameters for a cosmology splice.
"""
mylog.info("Writing redshift output list to %s.", filename)
f = open(filename, "w")
for output in outputs:
f.write("%f\n" % (1. / (1. + output["redshift"])))
f.close()
class StarFormationRate(object):
r"""Calculates the star formation rate for a given population of
star particles.
Parameters
----------
ds : EnzoDataset object
data_source : AMRRegion object, optional
The region from which stars are extracted for analysis. If this
is not supplied, the next three must be, otherwise the next
three do not need to be specified.
star_mass : Ordered array or list of floats
The mass of the stars to be analyzed in units of Msun.
star_creation_time : Ordered array or list of floats
The creation time for the stars in code units.
volume : Float
The comoving volume of the region for the specified list of stars.
bins : Integer
The number of time bins used for binning the stars. Default = 300.
star_filter : A user-defined filtering rule for stars.
See: http://yt-project.org/docs/dev/analyzing/filtering.html
Default: ct>0
Examples
--------
>>> import yt
>>> from yt.analysis_modules.star_analysis.api import StarFormationRate
>>> ds = yt.load("Enzo_64/RD0006/RedshiftOutput0006")
>>> sp = ds.sphere([0.5, 0.5, 0.5], 0.1)
>>> sfr = StarFormationRate(ds, sp)
"""
def __init__(self, ds, data_source=None, star_mass=None,
star_creation_time=None, bins=300, volume=None,
star_filter=None):
self._ds = ds
self._data_source = data_source
self._filter = star_filter
self.ds_provided = self._data_source is not None
self.filter_provided = self._filter is not None
self.bin_count = bins
# Set up for time conversion.
self.cosm = Cosmology(
hubble_constant=self._ds.hubble_constant,
omega_matter=self._ds.omega_matter,
omega_lambda=self._ds.omega_lambda)
# Find the time right now.
self.time_now = self._ds.current_time
if not self.ds_provided:
# Check to make sure we have the right set of informations.
if star_mass is None or star_creation_time is None \
or volume is None:
mylog.error("""
If data_source is not provided, all of these parameters
need to be set:
star_mass (array, Msun),
star_creation_time (array, code units),
volume (float, cMpc**3).""")
return None
if isinstance(star_mass, YTArray):
assert star_mass.units.same_dimensions_as(g.units)
elif star_mass is not None:
star_mass = YTArray(star_mass, 'Msun')
self.star_mass = star_mass
if isinstance(star_creation_time, YTArray):
assert star_creation_time.units.same_dimensions_as(s.units)
elif star_creation_time is not None:
star_creation_time = self._ds.arr(star_creation_time,
'code_time')
self.star_creation_time = star_creation_time
if isinstance(volume, YTQuantity):
assert volume.units.same_dimensions_as(
self._ds.quan(1.0, 'Mpccm**3').units
)
elif volume is not None:
volume = self._ds.quan(volume, 'Mpccm**3')
self.volume = volume
# Build the distribution.
self.build_dist()
# Attach some convenience arrays.
self.attach_arrays()
def build_dist(self):
"""
Build the data for plotting.
"""
# Pick out the stars.
if self.filter_provided:
ct = self._filter['creation_time']
mass_stars = self._data_source[self._filter, "particle_mass"]
else:
if self.ds_provided:
ct = self._data_source['creation_time']
if ct is None:
errmsg = 'data source must have particle_age!'
mylog.error(errmsg)
raise RuntimeError(errmsg)
mask = ct > 0
if not any(mask):
errmsg = 'all particles have age < 0'
mylog.error(errmsg)
raise RuntimeError(errmsg)
# type = self._data_source['particle_type']
ct_stars = ct[mask]
mass_stars = self._data_source[
'particle_mass'][mask].in_units('Msun')
del mask
else:
ct_stars = self.star_creation_time
mass_stars = self.star_mass
# Find the oldest stars in units of code time.
tmin = ct_stars.min().in_units("s")
# Multiply the end to prevent numerical issues.
self.time_bins = np.linspace(
tmin * 1.01, self._ds.current_time.in_units("s"),
num=self.bin_count + 1)
# Figure out which bins the stars go into.
inds = np.digitize(ct_stars.in_units("s"), self.time_bins) - 1
# Sum up the stars created in each time bin.
self.mass_bins = YTArray(
np.zeros(self.bin_count + 1, dtype='float64'), "Msun"
)
for index in np.unique(inds):
self.mass_bins[index] += (mass_stars[inds == index]).sum()
# We will want the time taken between bins.
self.time_bins_dt = self.time_bins[1:] - self.time_bins[:-1]
def attach_arrays(self):
"""
Attach convenience arrays to the class for easy access.
"""
if self.ds_provided:
try:
vol = self._data_source[
'cell_volume'].in_units('Mpccm ** 3').sum()
except AttributeError:
# If we're here, this is probably a HOPHalo object, and we
# can get the volume this way.
ds = self._data_source.get_sphere()
vol = ds['cell_volume'].in_units('Mpccm ** 3').sum()
else:
vol = self.volume.in_units('Mpccm ** 3')
# Use the center of the time_bin, not the left edge.
self.time = 0.5 * \
(self.time_bins[1:] + self.time_bins[:-1]).in_units('yr')
self.lookback_time = self.time_now - self.time # now in code_time...
self.redshift = self.cosm.z_from_t(self.time)
self.Msol_yr = (
self.mass_bins[:-1] / self.time_bins_dt[:]).in_units('Msun/yr')
# changed vol from mpc to mpccm used in literature
self.Msol_yr_vol = self.Msol_yr / vol
self.Msol = self.mass_bins[:-1].in_units("Msun")
self.Msol_cumulative = self.Msol.cumsum()
def write_out(self, name="StarFormationRate.out"):
r"""Write out the star analysis to a text file *name*. The columns are in
order.
The columns in the output file are:
1. Time (yrs)
2. Look-back time (yrs)
3. Redshift
4. Star formation rate in this bin per year (Msol/yr)
5. Star formation rate in this bin per year
per Mpc**3 (Msol/yr/Mpc**3)
6. Stars formed in this time bin (Msol)
7. Cumulative stars formed up to this time bin (Msol)
Parameters
----------
name : String
The name of the file to write to. Default = StarFormationRate.out.
Examples
--------
>>> import yt
>>> from yt.analysis_modules.star_analysis.api import StarFormationRate
>>> ds = yt.load("Enzo_64/RD0006/RedshiftOutput0006")
>>> sp = ds.sphere([0.5, 0.5, 0.5], 0.1)
>>> sfr = StarFormationRate(ds, sp)
>>> sfr.write_out("stars-SFR.out")
"""
fp = open(name, "w")
fp.write(
"#time\tlookback\tredshift\tMsol/yr\tMsol/yr/Mpc3\tMsol\tcumMsol\t\n")
for i, time in enumerate(self.time):
line = "%1.5e %1.5e %1.5e %1.5e %1.5e %1.5e %1.5e\n" % \
(time.in_units("yr"), # Time
self.lookback_time[i].in_units('yr'), # Lookback time
self.redshift[i], # Redshift
self.Msol_yr[i].in_units("Msun/yr"),
self.Msol_yr_vol[i],
self.Msol[i].in_units("Msun"), # Msol in bin
self.Msol_cumulative[i].in_units("Msun")) # cumulative
fp.write(line)
fp.close()
def from_data_source(cls, data_source, redshift, area,
exp_time, source_model, point_sources=False,
parameters=None, center=None, dist=None,
cosmology=None, velocity_fields=None):
r"""
Initialize a :class:`~pyxsim.photon_list.PhotonList` from a yt data
source. The redshift, collecting area, exposure time, and cosmology
are stored in the *parameters* dictionary which is passed to the
*source_model* function.
Parameters
----------
data_source : :class:`~yt.data_objects.data_containers.YTSelectionContainer`
The data source from which the photons will be generated.
redshift : float
The cosmological redshift for the photons.
area : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
The collecting area to determine the number of photons. If units are
not specified, it is assumed to be in cm^2.
exp_time : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
The exposure time to determine the number of photons. If units are
not specified, it is assumed to be in seconds.
source_model : :class:`~pyxsim.source_models.SourceModel`
A source model used to generate the photons.
point_sources : boolean, optional
If True, the photons will be assumed to be generated from the exact
positions of the cells or particles and not smeared around within
a volume. Default: False
parameters : dict, optional
A dictionary of parameters to be passed for the source model to use,
if necessary.
center : string or array_like, optional
The origin of the photon spatial coordinates. Accepts "c", "max", or
a coordinate. If not specified, pyxsim attempts to use the "center"
field parameter of the data_source.
dist : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
The angular diameter distance, used for nearby sources. This may be
optionally supplied instead of it being determined from the
*redshift* and given *cosmology*. If units are not specified, it is
assumed to be in kpc. To use this, the redshift must be set to zero.
cosmology : :class:`~yt.utilities.cosmology.Cosmology`, optional
Cosmological information. If not supplied, we try to get the
cosmology from the dataset. Otherwise, LCDM with the default yt
parameters is assumed.
velocity_fields : list of fields
The yt fields to use for the velocity. If not specified, the
following will be assumed:
['velocity_x', 'velocity_y', 'velocity_z'] for grid datasets
['particle_velocity_x', 'particle_velocity_y', 'particle_velocity_z'] for particle datasets
Examples
--------
>>> thermal_model = ThermalSourceModel(apec_model, Zmet=0.3)
>>> redshift = 0.05
>>> area = 6000.0 # assumed here in cm**2
>>> time = 2.0e5 # assumed here in seconds
>>> sp = ds.sphere("c", (500., "kpc"))
>>> my_photons = PhotonList.from_data_source(sp, redshift, area,
... time, thermal_model)
"""
ds = data_source.ds
if parameters is None:
parameters = {}
if cosmology is None:
if hasattr(ds, 'cosmology'):
cosmo = ds.cosmology
else:
cosmo = Cosmology()
else:
cosmo = cosmology
if dist is None:
if redshift <= 0.0:
msg = "If redshift <= 0.0, you must specify a distance to the " \
"source using the 'dist' argument!"
mylog.error(msg)
raise ValueError(msg)
D_A = cosmo.angular_diameter_distance(0.0, redshift).in_units("Mpc")
else:
D_A = parse_value(dist, "kpc")
if redshift > 0.0:
mylog.warning("Redshift must be zero for nearby sources. "
"Resetting redshift to 0.0.")
redshift = 0.0
if isinstance(center, string_types):
if center == "center" or center == "c":
parameters["center"] = ds.domain_center
elif center == "max" or center == "m":
parameters["center"] = ds.find_max("density")[-1]
elif iterable(center):
if isinstance(center, YTArray):
parameters["center"] = center.in_units("code_length")
elif isinstance(center, tuple):
if center[0] == "min":
parameters["center"] = ds.find_min(center[1])[-1]
elif center[0] == "max":
parameters["center"] = ds.find_max(center[1])[-1]
else:
raise RuntimeError
else:
parameters["center"] = ds.arr(center, "code_length")
elif center is None:
if hasattr(data_source, "left_edge"):
parameters["center"] = 0.5*(data_source.left_edge+data_source.right_edge)
else:
parameters["center"] = data_source.get_field_parameter("center")
parameters["fid_exp_time"] = parse_value(exp_time, "s")
parameters["fid_area"] = parse_value(area, "cm**2")
parameters["fid_redshift"] = redshift
parameters["fid_d_a"] = D_A
parameters["hubble"] = cosmo.hubble_constant
parameters["omega_matter"] = cosmo.omega_matter
parameters["omega_lambda"] = cosmo.omega_lambda
if redshift > 0.0:
mylog.info("Cosmology: h = %g, omega_matter = %g, omega_lambda = %g" %
(cosmo.hubble_constant, cosmo.omega_matter, cosmo.omega_lambda))
else:
mylog.info("Observing local source at distance %s." % D_A)
D_A = parameters["fid_d_a"].in_cgs()
dist_fac = 1.0/(4.*np.pi*D_A.value*D_A.value*(1.+redshift)**2)
spectral_norm = parameters["fid_area"].v*parameters["fid_exp_time"].v*dist_fac
source_model.setup_model(data_source, redshift, spectral_norm)
p_fields, v_fields, w_field = determine_fields(ds,
source_model.source_type,
point_sources)
if velocity_fields is not None:
v_fields = velocity_fields
if p_fields[0] == ("index", "x"):
parameters["data_type"] = "cells"
else:
parameters["data_type"] = "particles"
citer = data_source.chunks([], "io")
photons = defaultdict(list)
for chunk in parallel_objects(citer):
chunk_data = source_model(chunk)
if chunk_data is not None:
ncells, number_of_photons, idxs, energies = chunk_data
photons["num_photons"].append(number_of_photons)
photons["energy"].append(energies)
photons["pos"].append(np.array([chunk[p_fields[0]].d[idxs],
chunk[p_fields[1]].d[idxs],
chunk[p_fields[2]].d[idxs]]))
photons["vel"].append(np.array([chunk[v_fields[0]].d[idxs],
chunk[v_fields[1]].d[idxs],
chunk[v_fields[2]].d[idxs]]))
if w_field is None:
photons["dx"].append(np.zeros(ncells))
else:
photons["dx"].append(chunk[w_field].d[idxs])
source_model.cleanup_model()
photon_units = {"pos": ds.field_info[p_fields[0]].units,
"vel": ds.field_info[v_fields[0]].units,
"energy": "keV"}
if w_field is None:
photon_units["dx"] = "kpc"
else:
photon_units["dx"] = ds.field_info[w_field].units
concatenate_photons(ds, photons, photon_units)
c = parameters["center"].to("kpc")
if sum(ds.periodicity) > 0:
# Fix photon coordinates for regions crossing a periodic boundary
dw = ds.domain_width.to("kpc")
le, re = find_object_bounds(data_source)
for i in range(3):
if ds.periodicity[i] and photons["pos"].shape[0] > 0:
tfl = photons["pos"][:,i] < le[i]
tfr = photons["pos"][:,i] > re[i]
photons["pos"][tfl,i] += dw[i]
photons["pos"][tfr,i] -= dw[i]
# Re-center all coordinates
if photons["pos"].shape[0] > 0:
photons["pos"] -= c
mylog.info("Finished generating photons.")
mylog.info("Number of photons generated: %d" % int(np.sum(photons["num_photons"])))
mylog.info("Number of cells with photons: %d" % photons["dx"].size)
return cls(photons, parameters, cosmo)
