def __init__(self, nH, emid, sigma):
self.nH = YTQuantity(nH * 1.0e22, "cm**-2")
self.emid = YTArray(emid, "keV")
self.sigma = YTArray(sigma, "cm**2")
def __missing__(self, item):
if not isinstance(item, tuple):
field = ("unknown", item)
else:
field = item
finfo = self.ds._get_field_info(*field)
params, permute_params = finfo._get_needed_parameters(self)
self.field_parameters.update(params)
# For those cases where we are guessing the field type, we will
# need to re-update -- otherwise, our item will always not have the
# field type. This can lead to, for instance, "unknown" particle
# types not getting correctly identified.
# Note that the *only* way this works is if we also fix our field
# dependencies during checking. Bug #627 talks about this.
item = self.ds._last_freq
if finfo is not None and finfo._function.__name__ != "NullFunc":
try:
for param, param_v in permute_params.items():
for v in param_v:
self.field_parameters[param] = v
vv = finfo(self)
if not permute_params:
vv = finfo(self)
except NeedsGridType as exc:
ngz = exc.ghost_zones
nfd = FieldDetector(
self.nd + ngz * 2,
ds=self.ds,
field_parameters=self.field_parameters.copy(),
)
nfd._num_ghost_zones = ngz
vv = finfo(nfd)
if ngz > 0:
vv = vv[ngz:-ngz, ngz:-ngz, ngz:-ngz]
for i in nfd.requested:
if i not in self.requested:
self.requested.append(i)
for i in nfd.requested_parameters:
if i not in self.requested_parameters:
self.requested_parameters.append(i)
if vv is not None:
if not self.flat:
self[item] = vv
else:
self[item] = vv.ravel()
return self[item]
elif finfo is not None and finfo.sampling_type == "particle":
io = io_registry[self.ds.dataset_type](self.ds)
if hasattr(io, "_vector_fields") and (
item in io._vector_fields or item[1] in io._vector_fields
):
try:
cols = io._vector_fields[item]
except KeyError:
cols = io._vector_fields[item[1]]
# A vector
self[item] = YTArray(
np.ones((self.NumberOfParticles, cols)),
finfo.units,
registry=self.ds.unit_registry,
)
else:
# Not a vector
self[item] = YTArray(
np.ones(self.NumberOfParticles),
finfo.units,
registry=self.ds.unit_registry,
)
if item == ("STAR", "BIRTH_TIME"):
# hack for the artio frontend so we pass valid times to
# the artio functions for calculating physical times
# from internal times
self[item] *= -0.1
self.requested.append(item)
return self[item]
self.requested.append(item)
if item not in self:
self[item] = self._read_data(item)
return self[item]
def __init__(self,
simulation_ds=None,
halos_ds=None,
make_analytic=True,
omega_matter0=0.2726,
omega_lambda0=0.7274,
omega_baryon0=0.0456,
hubble0=0.704,
sigma8=0.86,
primordial_index=1.0,
this_redshift=0,
log_mass_min=None,
log_mass_max=None,
num_sigma_bins=360,
fitting_function=4):
self.simulation_ds = simulation_ds
self.halos_ds = halos_ds
self.omega_matter0 = omega_matter0
self.omega_lambda0 = omega_lambda0
self.omega_baryon0 = omega_baryon0
self.hubble0 = hubble0
self.sigma8 = sigma8
self.primordial_index = primordial_index
self.this_redshift = this_redshift
self.log_mass_min = log_mass_min
self.log_mass_max = log_mass_max
self.num_sigma_bins = num_sigma_bins
self.fitting_function = fitting_function
self.make_analytic = make_analytic
self.make_simulated = False
"""
If we want to make an analytic mass function, grab what we can from either the
halo file or the data set, and make sure that the user supplied everything else
that is needed.
"""
# If we don't have any datasets, make the analytic function with user values
if simulation_ds is None and halos_ds is None:
# Set a reasonable mass min and max if none were provided
if log_mass_min is None:
self.log_mass_min = 5
if log_mass_max is None:
self.log_mass_max = 16
# If we're making the analytic function...
if self.make_analytic is True:
# Try to set cosmological parameters from the simulation dataset
if simulation_ds is not None:
self.omega_matter0 = self.simulation_ds.omega_matter
self.omega_lambda0 = self.simulation_ds.omega_lambda
self.hubble0 = self.simulation_ds.hubble_constant
self.this_redshift = self.simulation_ds.current_redshift
# Set a reasonable mass min and max if none were provided
if log_mass_min is None:
self.log_mass_min = 5
if log_mass_max is None:
self.log_mass_max = 16
# If we have a halo dataset but not a simulation dataset, use that instead
if simulation_ds is None and halos_ds is not None:
self.omega_matter0 = self.halos_ds.omega_matter
self.omega_lambda0 = self.halos_ds.omega_lambda
self.hubble0 = self.halos_ds.hubble_constant
self.this_redshift = self.halos_ds.current_redshift
# If the user didn't specify mass min and max, set them from the halos
if log_mass_min is None:
self.set_mass_from_halos("min_mass")
if log_mass_max is None:
self.set_mass_from_halos("max_mass")
# Do the calculations.
self.sigmaM()
self.dndm()
# Return the mass array in M_solar rather than M_solar/h
self.masses_analytic = YTArray(self.masses_analytic / self.hubble0,
"Msun")
# The halo arrays will already have yt units, but the analytic forms do
# not. If a dataset has been provided, use that to give them units. At the
# same time, convert to comoving (Mpc)^-3
if simulation_ds is not None:
self.n_cumulative_analytic = simulation_ds.arr(
self.n_cumulative_analytic, "(Mpccm)**(-3)")
elif halos_ds is not None:
self.n_cumulative_analytic = halos_ds.arr(
self.n_cumulative_analytic, "(Mpccm)**(-3)")
else:
from yt.units.unit_registry import UnitRegistry
from yt.units.dimensions import length
hmf_registry = UnitRegistry()
for my_unit in ["m", "pc", "AU", "au"]:
new_unit = "%scm" % my_unit
hmf_registry.add(
new_unit, hmf_registry.lut[my_unit][0] /
(1 + self.this_redshift), length,
"\\rm{%s}/(1+z)" % my_unit)
self.n_cumulative_analytic = YTArray(
self.n_cumulative_analytic,
"(Mpccm)**(-3)",
registry=hmf_registry)
"""
If a halo file has been supplied, make a mass function for the simulated halos.
"""
if halos_ds is not None:
# Used to check if a simulated halo mass funciton exists to write out
self.make_simulated = True
# Calculate the simulated halo mass function
self.create_sim_hmf()
def get_data(self, field):
"""
Return the data array of the image corresponding to *field*
with units attached.
"""
return YTArray(self.hdulist[field].data, self.field_units[field])
def interpolate_area(self, energy):
"""
Interpolate the effective area to the energies provided by the supplied *energy* array.
"""
earea = np.interp(energy, self.emid, self.eff_area, left=0.0, right=0.0)
return YTArray(earea, "cm**2")
def unary_ufunc_comparison(ufunc, a):
out = a.copy()
a_array = a.to_ndarray()
if ufunc in (
np.isreal,
np.iscomplex,
):
# According to the numpy docs, these two explicitly do not do
# in-place copies.
ret = ufunc(a)
assert_true(not hasattr(ret, 'units'))
assert_array_equal(ret, ufunc(a))
elif ufunc in yield_np_ufuncs([
'exp', 'exp2', 'log', 'log2', 'log10', 'expm1', 'log1p', 'sin',
'cos', 'tan', 'arcsin', 'arccos', 'arctan', 'sinh', 'cosh', 'tanh',
'arccosh', 'arcsinh', 'arctanh', 'deg2rad', 'rad2deg', 'isfinite',
'isinf', 'isnan', 'signbit', 'sign', 'rint', 'logical_not'
]):
# These operations should return identical results compared to numpy.
with np.errstate(invalid='ignore'):
try:
ret = ufunc(a, out=out)
except YTUnitOperationError:
assert_true(ufunc in (np.deg2rad, np.rad2deg))
ret = ufunc(YTArray(a, '1'))
assert_array_equal(ret, out)
assert_array_equal(ret, ufunc(a_array))
# In-place copies do not drop units.
assert_true(hasattr(out, 'units'))
assert_true(not hasattr(ret, 'units'))
elif ufunc in yield_np_ufuncs([
'absolute', 'fabs', 'conjugate', 'floor', 'ceil', 'trunc',
'negative', 'spacing', 'positive'
]):
ret = ufunc(a, out=out)
assert_array_equal(ret, out)
assert_array_equal(ret.to_ndarray(), ufunc(a_array))
assert_true(ret.units == out.units)
elif ufunc in yield_np_ufuncs(
['ones_like', 'square', 'sqrt', 'reciprocal']):
if ufunc is np.ones_like:
ret = ufunc(a)
else:
with np.errstate(invalid='ignore'):
ret = ufunc(a, out=out)
assert_array_equal(ret, out)
with np.errstate(invalid='ignore'):
assert_array_equal(ret.to_ndarray(), ufunc(a_array))
if ufunc is np.square:
assert_true(out.units == a.units**2)
assert_true(ret.units == a.units**2)
elif ufunc is np.sqrt:
assert_true(out.units == a.units**0.5)
assert_true(ret.units == a.units**0.5)
elif ufunc is np.reciprocal:
assert_true(out.units == a.units**-1)
assert_true(ret.units == a.units**-1)
elif ufunc is np.modf:
ret1, ret2 = ufunc(a)
npret1, npret2 = ufunc(a_array)
assert_array_equal(ret1.to_ndarray(), npret1)
assert_array_equal(ret2.to_ndarray(), npret2)
elif ufunc is np.frexp:
ret1, ret2 = ufunc(a)
npret1, npret2 = ufunc(a_array)
assert_array_equal(ret1, npret1)
assert_array_equal(ret2, npret2)
elif ufunc is np.invert:
assert_raises(TypeError, ufunc, a)
elif hasattr(np, 'isnat') and ufunc is np.isnat:
# numpy 1.13 raises ValueError, numpy 1.14 and newer raise TypeError
assert_raises((TypeError, ValueError), ufunc, a)
else:
# There shouldn't be any untested ufuncs.
assert_true(False)
def test_reductions():
arr = YTArray([[1, 2, 3], [4, 5, 6]], 'cm')
ev_result = arr.dot(YTArray([1, 2, 3], 'cm'))
res = YTArray([14., 32.], 'cm**2')
assert_equal(ev_result, res)
assert_equal(ev_result.units, res.units)
assert_isinstance(ev_result, YTArray)
answers = {
'prod': (YTQuantity(720, 'cm**6'), YTArray([4, 10, 18], 'cm**2'),
YTArray([6, 120], 'cm**3')),
'sum': (
YTQuantity(21, 'cm'),
YTArray([5., 7., 9.], 'cm'),
YTArray([6, 15], 'cm'),
),
'mean': (YTQuantity(3.5, 'cm'), YTArray([2.5, 3.5, 4.5],
'cm'), YTArray([2, 5], 'cm')),
'std': (YTQuantity(1.707825127659933,
'cm'), YTArray([1.5, 1.5, 1.5], 'cm'),
YTArray([0.81649658, 0.81649658], 'cm')),
}
for op, (result1, result2, result3) in answers.items():
ev_result = getattr(arr, op)()
assert_almost_equal(ev_result, result1)
assert_almost_equal(ev_result.units, result1.units)
assert_isinstance(ev_result, YTQuantity)
for axis, result in [(0, result2), (1, result3), (-1, result3)]:
ev_result = getattr(arr, op)(axis=axis)
assert_almost_equal(ev_result, result)
assert_almost_equal(ev_result.units, result.units)
assert_isinstance(ev_result, YTArray)
units={'radius': 'code_length'})
# Find the radius of the I-front (f_HI=0.5)
res = np.abs(prof1d["Neutral_Fraction"] - 0.5)
ir = np.where(res == res.min())[0]
r = np.interp(0.5, prof1d["Neutral_Fraction"], prof1d.x.value)
r = pf.quan(r, 'code_length')
time.append(pf.current_time.to('s'))
radius.append(r.to('cm'))
#print("radius = ", r)
del pf
del prof1d
#print("time1 = ", time)
time = np.array(time)
radius = np.array(radius)
radius = YTArray(radius, 'cm')
#print("radius = ", radius.to('kpc'))
#print("time = ", time)
# Calculate analytic Stroemgren radius
alpha = 2.59e-13 # Recombination rate at T=1e4 K
nH = 1e-3 # Hydrogen density
lum = 5e48 # Ionizing photon luminosity
trec = 1.0 / (alpha * nH) # Recombination time
rs = ((3 * lum) / (4 * np.pi * alpha * nH**2))**(1.0 / 3)
anyl_radius = rs * (1.0 - np.exp(-time / trec))**(1.0 / 3)
#print("anyl_radius = ", anyl_radius)
anyl_radius = YTArray(anyl_radius, 'cm')
#print("anyl_radius = ", anyl_radius.to('kpc'))
ratio = radius / anyl_radius
error = np.abs(1 - ratio)
def test_subtraction():
"""
Test subtraction of two YTArrays
"""
# Same units
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'cm')
a3 = [4 * cm, 5 * cm, 6 * cm]
answer1 = YTArray([-3, -3, -3], 'cm')
answer2 = YTArray([3, 3, 3], 'cm')
yield operate_and_compare, a1, a2, operator.sub, answer1
yield operate_and_compare, a2, a1, operator.sub, answer2
yield operate_and_compare, a1, a3, operator.sub, answer1
yield operate_and_compare, a3, a1, operator.sub, answer2
yield operate_and_compare, a1, a2, np.subtract, answer1
yield operate_and_compare, a2, a1, np.subtract, answer2
yield operate_and_compare, a1, a3, np.subtract, answer1
yield operate_and_compare, a3, a1, np.subtract, answer2
# different units
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'm')
a3 = [4 * m, 5 * m, 6 * m]
answer1 = YTArray([-399, -498, -597], 'cm')
answer2 = YTArray([3.99, 4.98, 5.97], 'm')
answer3 = YTArray([399, 498, 597], 'cm')
yield operate_and_compare, a1, a2, operator.sub, answer1
yield operate_and_compare, a2, a1, operator.sub, answer2
yield operate_and_compare, a1, a3, operator.sub, answer1
yield operate_and_compare, a3, a1, operator.sub, answer3
yield assert_raises, YTUfuncUnitError, np.subtract, a1, a2
yield assert_raises, YTUfuncUnitError, np.subtract, a1, a3
# Test dimensionless quantities
a1 = YTArray([1, 2, 3])
a2 = array([4, 5, 6])
a3 = [4, 5, 6]
answer1 = YTArray([-3, -3, -3])
answer2 = YTArray([3, 3, 3])
yield operate_and_compare, a1, a2, operator.sub, answer1
yield operate_and_compare, a2, a1, operator.sub, answer2
yield operate_and_compare, a1, a3, operator.sub, answer1
yield operate_and_compare, a3, a1, operator.sub, answer2
yield operate_and_compare, a1, a2, np.subtract, answer1
yield operate_and_compare, a2, a1, np.subtract, answer2
yield operate_and_compare, a1, a3, np.subtract, answer1
yield operate_and_compare, a3, a1, np.subtract, answer2
# Catch the different dimensions error
a1 = YTArray([1, 2, 3], 'm')
a2 = YTArray([4, 5, 6], 'kg')
yield assert_raises, YTUnitOperationError, operator.sub, a1, a2
yield assert_raises, YTUnitOperationError, operator.isub, a1, a2
def pixelize_line(self, field, start_point, end_point, npoints):
"""
Method for sampling datasets along a line in preparation for
one-dimensional line plots. For UnstructuredMesh, relies on a
sampling routine written in cython
"""
if npoints < 2:
raise ValueError(
"Must have at least two sample points in order to draw a line plot."
)
index = self.ds.index
if hasattr(index, "meshes") and not isinstance(
index.meshes[0], SemiStructuredMesh
):
ftype, fname = field
if ftype == "all":
mesh_id = 0
indices = np.concatenate(
[mesh.connectivity_indices for mesh in index.mesh_union]
)
else:
mesh_id = int(ftype[-1]) - 1
indices = index.meshes[mesh_id].connectivity_indices
coords = index.meshes[mesh_id].connectivity_coords
if coords.shape[1] != end_point.size != start_point.size:
raise ValueError(
"The coordinate dimension doesn't match the "
"start and end point dimensions."
)
offset = index.meshes[mesh_id]._index_offset
ad = self.ds.all_data()
field_data = ad[field]
# if this is an element field, promote to 2D here
if len(field_data.shape) == 1:
field_data = np.expand_dims(field_data, 1)
# if this is a higher-order element, we demote to 1st order
# here, for now.
elif field_data.shape[1] == 27:
# hexahedral
mylog.warning(
"High order elements not yet supported, dropping to 1st order."
)
field_data = field_data[:, 0:8]
indices = indices[:, 0:8]
arc_length, plot_values = pixelize_element_mesh_line(
coords,
indices,
start_point,
end_point,
npoints,
field_data,
index_offset=offset,
)
arc_length = YTArray(arc_length, start_point.units)
plot_values = YTArray(plot_values, field_data.units)
else:
ray = self.ds.ray(start_point, end_point)
arc_length, plot_values = _sample_ray(ray, npoints, field)
return arc_length, plot_values
def __missing__(self, item):
if hasattr(self.ds, "field_info"):
if not isinstance(item, tuple):
field = ("unknown", item)
finfo = self.ds._get_field_info(*field)
#mylog.debug("Guessing field %s is %s", item, finfo.name)
else:
field = item
finfo = self.ds._get_field_info(*field)
# For those cases where we are guessing the field type, we will
# need to re-update -- otherwise, our item will always not have the
# field type. This can lead to, for instance, "unknown" particle
# types not getting correctly identified.
# Note that the *only* way this works is if we also fix our field
# dependencies during checking. Bug #627 talks about this.
item = self.ds._last_freq
else:
FI = getattr(self.ds, "field_info", FieldInfo)
if item in FI:
finfo = FI[item]
else:
finfo = None
if finfo is not None and finfo._function.func_name != 'NullFunc':
try:
vv = finfo(self)
except NeedsGridType as exc:
ngz = exc.ghost_zones
nfd = FieldDetector(self.nd + ngz * 2, ds=self.ds)
nfd._num_ghost_zones = ngz
vv = finfo(nfd)
if ngz > 0: vv = vv[ngz:-ngz, ngz:-ngz, ngz:-ngz]
for i in nfd.requested:
if i not in self.requested: self.requested.append(i)
for i in nfd.requested_parameters:
if i not in self.requested_parameters:
self.requested_parameters.append(i)
if vv is not None:
if not self.flat: self[item] = vv
else: self[item] = vv.ravel()
return self[item]
elif finfo is not None and finfo.particle_type:
if "particle_position" in (item, item[1]) or \
"particle_velocity" in (item, item[1]) or \
"Velocity" in (item, item[1]) or \
"Velocities" in (item, item[1]) or \
"Coordinates" in (item, item[1]):
# A vector
self[item] = \
YTArray(np.ones((self.NumberOfParticles, 3)),
finfo.units, registry=self.ds.unit_registry)
else:
# Not a vector
self[item] = \
YTArray(np.ones(self.NumberOfParticles),
finfo.units, registry=self.ds.unit_registry)
self.requested.append(item)
return self[item]
self.requested.append(item)
if item not in self:
self[item] = self._read_data(item)
return self[item]
def test_display_ytarray_too_large():
arr = YTArray([1,2,3,4], 'cm')
assert_raises(YTArrayTooLargeToDisplay, display_ytarray, arr)
def project_photons(self,
normal,
sky_center,
absorb_model=None,
nH=None,
no_shifting=False,
north_vector=None,
sigma_pos=None,
kernel="top_hat",
prng=None,
**kwargs):
r"""
Projects photons onto an image plane given a line of sight.
Returns a new :class:`~pyxsim.event_list.EventList`.
Parameters
----------
normal : character or array-like
Normal vector to the plane of projection. If "x", "y", or "z", will
assume to be along that axis (and will probably be faster). Otherwise,
should be an off-axis normal vector, e.g [1.0, 2.0, -3.0]
sky_center : array-like
Center RA, Dec of the events in degrees.
absorb_model : string or :class:`~pyxsim.spectral_models.AbsorptionModel`
A model for foreground galactic absorption, to simulate the
absorption of events before being detected. This cannot be applied
here if you already did this step previously in the creation of the
:class:`~pyxsim.photon_list.PhotonList` instance. Known options for
strings are "wabs" and "tbabs".
nH : float, optional
The foreground column density in units of 10^22 cm^{-2}. Only used
if absorption is applied.
no_shifting : boolean, optional
If set, the photon energies will not be Doppler shifted.
north_vector : a sequence of floats
A vector defining the "up" direction. This option sets the
orientation of the plane of projection. If not set, an arbitrary
grid-aligned north_vector is chosen. Ignored in the case where a
particular axis (e.g., "x", "y", or "z") is explicitly specified.
sigma_pos : float, optional
Apply a gaussian smoothing operation to the sky positions of the
events. This may be useful when the binned events appear blocky due
to their uniform distribution within simulation cells. However, this
will move the events away from their originating position on the
sky, and so may distort surface brightness profiles and/or spectra.
Should probably only be used for visualization purposes. Supply a
float here to smooth with a standard deviation with this fraction
of the cell size. Default: None
kernel : string, optional
The kernel used when smoothing positions of X-rays originating from
SPH particles, "gaussian" or "top_hat". Default: "top_hat".
prng : integer or :class:`~numpy.random.RandomState` object
A pseudo-random number generator. Typically will only be specified
if you have a reason to generate the same set of random numbers,
such as for a test. Default is to use the :mod:`numpy.random`
module.
Examples
--------
>>> L = np.array([0.1,-0.2,0.3])
>>> events = my_photons.project_photons(L, [30., 45.])
"""
prng = parse_prng(prng)
scale_shift = -1.0 / clight.to("km/s")
if "smooth_positions" in kwargs:
issue_deprecation_warning(
"'smooth_positions' has been renamed to "
"'sigma_pos' and the former is deprecated!")
sigma_pos = kwargs["smooth_positions"]
if "redshift_new" in kwargs or "area_new" in kwargs or \
"exp_time_new" in kwargs or "dist_new" in kwargs:
issue_deprecation_warning(
"Changing the redshift, distance, area, or "
"exposure time has been deprecated in "
"project_photons!")
if sigma_pos is not None and self.parameters[
"data_type"] == "particles":
raise RuntimeError(
"The 'smooth_positions' argument should not be used with "
"particle-based datasets!")
if isinstance(absorb_model, string_types):
if absorb_model not in absorb_models:
raise KeyError("%s is not a known absorption model!" %
absorb_model)
absorb_model = absorb_models[absorb_model]
if absorb_model is not None:
if nH is None:
raise RuntimeError(
"You specified an absorption model, but didn't "
"specify a value for nH!")
absorb_model = absorb_model(nH)
sky_center = YTArray(sky_center, "degree")
n_ph = self.photons["num_photons"]
if not isinstance(normal, string_types):
L = np.array(normal)
orient = Orientation(L, north_vector=north_vector)
x_hat = orient.unit_vectors[0]
y_hat = orient.unit_vectors[1]
z_hat = orient.unit_vectors[2]
else:
x_hat = np.zeros(3)
y_hat = np.zeros(3)
z_hat = np.zeros(3)
parameters = {}
D_A = self.parameters["fid_d_a"]
events = {}
eobs = self.photons["energy"].v
if not no_shifting:
if comm.rank == 0:
mylog.info("Doppler-shifting photon energies.")
if isinstance(normal, string_types):
shift = self.photons["vel"][:,
"xyz".index(normal)] * scale_shift
else:
shift = np.dot(self.photons["vel"], z_hat) * scale_shift
doppler_shift(shift, n_ph, eobs)
if absorb_model is None:
det = np.ones(eobs.size, dtype='bool')
num_det = eobs.size
else:
if comm.rank == 0:
mylog.info("Foreground galactic absorption: using "
"the %s model and nH = %g." %
(absorb_model._name, nH))
det = absorb_model.absorb_photons(eobs, prng=prng)
num_det = det.sum()
events["eobs"] = YTArray(eobs[det], "keV")
num_events = comm.mpi_allreduce(num_det)
if comm.rank == 0:
mylog.info("%d events have been detected." % num_events)
if num_det > 0:
if comm.rank == 0:
mylog.info("Assigning positions to events.")
if isinstance(normal, string_types):
norm = "xyz".index(normal)
else:
norm = normal
xsky, ysky = scatter_events(norm, prng, kernel,
self.parameters["data_type"], num_det,
det, self.photons["num_photons"],
self.photons["pos"].d,
self.photons["dx"].d, x_hat, y_hat)
if self.parameters[
"data_type"] == "cells" and sigma_pos is not None:
if comm.rank == 0:
mylog.info("Optionally smoothing sky positions.")
sigma = sigma_pos * np.repeat(self.photons["dx"].d, n_ph)[det]
xsky += sigma * prng.normal(loc=0.0, scale=1.0, size=num_det)
ysky += sigma * prng.normal(loc=0.0, scale=1.0, size=num_det)
d_a = D_A.to("kpc").v
xsky /= d_a
ysky /= d_a
if comm.rank == 0:
mylog.info("Converting pixel to sky coordinates.")
pixel_to_cel(xsky, ysky, sky_center)
else:
xsky = []
ysky = []
events["xsky"] = YTArray(xsky, "degree")
events["ysky"] = YTArray(ysky, "degree")
parameters["exp_time"] = self.parameters["fid_exp_time"]
parameters["area"] = self.parameters["fid_area"]
parameters["sky_center"] = sky_center
return EventList(events, parameters)
def from_file(cls, filename):
r"""
Initialize a :class:`~pyxsim.photon_list.PhotonList` from
the HDF5 file *filename*.
"""
mylog.info("Reading photons from %s." % filename)
photons = {}
parameters = {}
f = h5py.File(filename, "r")
p = f["/parameters"]
parameters["fid_exp_time"] = YTQuantity(p["fid_exp_time"].value, "s")
parameters["fid_area"] = YTQuantity(p["fid_area"].value, "cm**2")
parameters["fid_redshift"] = p["fid_redshift"].value
parameters["fid_d_a"] = YTQuantity(p["fid_d_a"].value, "Mpc")
parameters["hubble"] = p["hubble"].value
parameters["omega_matter"] = p["omega_matter"].value
parameters["omega_lambda"] = p["omega_lambda"].value
if "data_type" in p:
parameters["data_type"] = force_unicode(p["data_type"].value)
else:
parameters["data_type"] = "cells"
d = f["/data"]
num_cells = d["x"].size
start_c = comm.rank * num_cells // comm.size
end_c = (comm.rank + 1) * num_cells // comm.size
photons["pos"] = YTArray(np.zeros((num_cells, 3)), "kpc")
photons["vel"] = YTArray(np.zeros((num_cells, 3)), "km/s")
photons["pos"][:, 0] = d["x"][start_c:end_c]
photons["pos"][:, 1] = d["y"][start_c:end_c]
photons["pos"][:, 2] = d["z"][start_c:end_c]
photons["vel"][:, 0] = d["vx"][start_c:end_c]
photons["vel"][:, 1] = d["vy"][start_c:end_c]
photons["vel"][:, 2] = d["vz"][start_c:end_c]
photons["dx"] = YTArray(d["dx"][start_c:end_c], "kpc")
n_ph = d["num_photons"][:]
if comm.rank == 0:
start_e = np.int64(0)
else:
start_e = n_ph[:start_c].sum()
end_e = start_e + np.int64(n_ph[start_c:end_c].sum())
photons["num_photons"] = n_ph[start_c:end_c]
photons["energy"] = YTArray(d["energy"][start_e:end_e], "keV")
f.close()
cosmo = Cosmology(hubble_constant=parameters["hubble"],
omega_matter=parameters["omega_matter"],
omega_lambda=parameters["omega_lambda"])
mylog.info("Read %d photons from %d %s." %
(n_ph.sum(), num_cells, parameters["data_type"]))
return cls(photons, parameters, cosmo)
def test_division():
"""
Test multiplication of two YTArrays
"""
# Same units
a1 = YTArray([1., 2., 3.], 'cm')
a2 = YTArray([4., 5., 6.], 'cm')
a3 = [4 * cm, 5 * cm, 6 * cm]
answer1 = YTArray([0.25, 0.4, 0.5])
answer2 = YTArray([4, 2.5, 2])
if "div" in dir(operator):
op = operator.div
else:
op = operator.truediv
operate_and_compare(a1, a2, op, answer1)
operate_and_compare(a2, a1, op, answer2)
operate_and_compare(a1, a3, op, answer1)
operate_and_compare(a3, a1, op, answer2)
operate_and_compare(a1, a2, np.divide, answer1)
operate_and_compare(a2, a1, np.divide, answer2)
operate_and_compare(a1, a3, np.divide, answer1)
operate_and_compare(a3, a1, np.divide, answer2)
# different units, same dimension
a1 = YTArray([1., 2., 3.], 'cm')
a2 = YTArray([4., 5., 6.], 'm')
a3 = [4 * m, 5 * m, 6 * m]
answer1 = YTArray([.0025, .004, .005])
answer2 = YTArray([400, 250, 200])
answer3 = YTArray([0.25, 0.4, 0.5], 'cm/m')
answer4 = YTArray([4.0, 2.5, 2.0], 'm/cm')
operate_and_compare(a1, a2, op, answer1)
operate_and_compare(a2, a1, op, answer2)
operate_and_compare(a1, a3, op, answer1)
operate_and_compare(a3, a1, op, answer2)
if LooseVersion(np.__version__) < LooseVersion('1.13.0'):
operate_and_compare(a1, a2, np.divide, answer3)
operate_and_compare(a2, a1, np.divide, answer4)
operate_and_compare(a1, a3, np.divide, answer3)
operate_and_compare(a3, a1, np.divide, answer4)
else:
operate_and_compare(a1, a2, np.divide, answer1)
operate_and_compare(a2, a1, np.divide, answer2)
operate_and_compare(a1, a3, np.divide, answer1)
operate_and_compare(a3, a1, np.divide, answer2)
# different dimensions
a1 = YTArray([1., 2., 3.], 'cm')
a2 = YTArray([4., 5., 6.], 'g')
a3 = [4 * g, 5 * g, 6 * g]
answer1 = YTArray([0.25, 0.4, 0.5], 'cm/g')
answer2 = YTArray([4, 2.5, 2], 'g/cm')
operate_and_compare(a1, a2, op, answer1)
operate_and_compare(a2, a1, op, answer2)
operate_and_compare(a1, a3, op, answer1)
operate_and_compare(a3, a1, op, answer2)
operate_and_compare(a1, a2, np.divide, answer1)
operate_and_compare(a2, a1, np.divide, answer2)
operate_and_compare(a1, a3, np.divide, answer1)
operate_and_compare(a3, a1, np.divide, answer2)
# One dimensionless, one unitful
a1 = YTArray([1., 2., 3.], 'cm')
a2 = array([4., 5., 6.])
a3 = [4, 5, 6]
answer1 = YTArray([0.25, 0.4, 0.5], 'cm')
answer2 = YTArray([4, 2.5, 2], '1/cm')
operate_and_compare(a1, a2, op, answer1)
operate_and_compare(a2, a1, op, answer2)
operate_and_compare(a1, a3, op, answer1)
operate_and_compare(a3, a1, op, answer2)
operate_and_compare(a1, a2, np.divide, answer1)
operate_and_compare(a2, a1, np.divide, answer2)
operate_and_compare(a1, a3, np.divide, answer1)
operate_and_compare(a3, a1, np.divide, answer2)
# Both dimensionless quantities
a1 = YTArray([1., 2., 3.])
a2 = array([4., 5., 6.])
a3 = [4, 5, 6]
answer1 = YTArray([0.25, 0.4, 0.5])
answer2 = YTArray([4, 2.5, 2])
operate_and_compare(a1, a2, op, answer1)
operate_and_compare(a2, a1, op, answer2)
operate_and_compare(a1, a3, op, answer1)
operate_and_compare(a3, a1, op, answer2)
operate_and_compare(a1, a3, np.divide, answer1)
operate_and_compare(a3, a1, np.divide, answer2)
operate_and_compare(a1, a3, np.divide, answer1)
operate_and_compare(a3, a1, np.divide, answer2)
def test_division():
"""
Test multiplication of two YTArrays
"""
# Same units
a1 = YTArray([1., 2., 3.], 'cm')
a2 = YTArray([4., 5., 6.], 'cm')
a3 = [4 * cm, 5 * cm, 6 * cm]
answer1 = YTArray([0.25, 0.4, 0.5])
answer2 = YTArray([4, 2.5, 2])
yield operate_and_compare, a1, a2, operator.div, answer1
yield operate_and_compare, a2, a1, operator.div, answer2
yield operate_and_compare, a1, a3, operator.div, answer1
yield operate_and_compare, a3, a1, operator.div, answer2
yield operate_and_compare, a1, a2, np.divide, answer1
yield operate_and_compare, a2, a1, np.divide, answer2
yield operate_and_compare, a1, a3, np.divide, answer1
yield operate_and_compare, a3, a1, np.divide, answer2
# different units, same dimension
a1 = YTArray([1., 2., 3.], 'cm')
a2 = YTArray([4., 5., 6.], 'm')
a3 = [4 * m, 5 * m, 6 * m]
answer1 = YTArray([.0025, .004, .005])
answer2 = YTArray([400, 250, 200])
answer3 = YTArray([0.25, 0.4, 0.5], 'cm/m')
answer4 = YTArray([4.0, 2.5, 2.0], 'm/cm')
yield operate_and_compare, a1, a2, operator.div, answer1
yield operate_and_compare, a2, a1, operator.div, answer2
yield operate_and_compare, a1, a3, operator.div, answer1
yield operate_and_compare, a3, a1, operator.div, answer2
yield operate_and_compare, a1, a2, np.divide, answer3
yield operate_and_compare, a2, a1, np.divide, answer4
yield operate_and_compare, a1, a3, np.divide, answer3
yield operate_and_compare, a3, a1, np.divide, answer4
# different dimensions
a1 = YTArray([1., 2., 3.], 'cm')
a2 = YTArray([4., 5., 6.], 'g')
a3 = [4 * g, 5 * g, 6 * g]
answer1 = YTArray([0.25, 0.4, 0.5], 'cm/g')
answer2 = YTArray([4, 2.5, 2], 'g/cm')
yield operate_and_compare, a1, a2, operator.div, answer1
yield operate_and_compare, a2, a1, operator.div, answer2
yield operate_and_compare, a1, a3, operator.div, answer1
yield operate_and_compare, a3, a1, operator.div, answer2
yield operate_and_compare, a1, a2, np.divide, answer1
yield operate_and_compare, a2, a1, np.divide, answer2
yield operate_and_compare, a1, a3, np.divide, answer1
yield operate_and_compare, a3, a1, np.divide, answer2
# One dimensionless, one unitful
a1 = YTArray([1., 2., 3.], 'cm')
a2 = array([4., 5., 6.])
a3 = [4, 5, 6]
answer1 = YTArray([0.25, 0.4, 0.5], 'cm')
answer2 = YTArray([4, 2.5, 2], '1/cm')
yield operate_and_compare, a1, a2, operator.div, answer1
yield operate_and_compare, a2, a1, operator.div, answer2
yield operate_and_compare, a1, a3, operator.div, answer1
yield operate_and_compare, a3, a1, operator.div, answer2
yield operate_and_compare, a1, a2, np.divide, answer1
yield operate_and_compare, a2, a1, np.divide, answer2
yield operate_and_compare, a1, a3, np.divide, answer1
yield operate_and_compare, a3, a1, np.divide, answer2
# Both dimensionless quantities
a1 = YTArray([1., 2., 3.])
a2 = array([4., 5., 6.])
a3 = [4, 5, 6]
answer1 = YTArray([0.25, 0.4, 0.5])
answer2 = YTArray([4, 2.5, 2])
yield operate_and_compare, a1, a2, operator.div, answer1
yield operate_and_compare, a2, a1, operator.div, answer2
yield operate_and_compare, a1, a3, operator.div, answer1
yield operate_and_compare, a3, a1, operator.div, answer2
yield operate_and_compare, a1, a3, np.divide, answer1
yield operate_and_compare, a3, a1, np.divide, answer2
yield operate_and_compare, a1, a3, np.divide, answer1
yield operate_and_compare, a3, a1, np.divide, answer2
def test_addition():
"""
Test addition of two YTArrays
"""
# Same units
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'cm')
a3 = [4 * cm, 5 * cm, 6 * cm]
answer = YTArray([5, 7, 9], 'cm')
operate_and_compare(a1, a2, operator.add, answer)
operate_and_compare(a2, a1, operator.add, answer)
operate_and_compare(a1, a3, operator.add, answer)
operate_and_compare(a3, a1, operator.add, answer)
operate_and_compare(a2, a1, np.add, answer)
operate_and_compare(a1, a2, np.add, answer)
operate_and_compare(a1, a3, np.add, answer)
operate_and_compare(a3, a1, np.add, answer)
# different units
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'm')
a3 = [4 * m, 5 * m, 6 * m]
answer1 = YTArray([401, 502, 603], 'cm')
answer2 = YTArray([4.01, 5.02, 6.03], 'm')
operate_and_compare(a1, a2, operator.add, answer1)
operate_and_compare(a2, a1, operator.add, answer2)
operate_and_compare(a1, a3, operator.add, answer1)
if LooseVersion(np.__version__) < LooseVersion('1.13.0'):
operate_and_compare(a3, a1, operator.add, answer1)
assert_raises(YTUfuncUnitError, np.add, a1, a2)
assert_raises(YTUfuncUnitError, np.add, a1, a3)
else:
operate_and_compare(a3, a1, operator.add, answer2)
operate_and_compare(a1, a2, np.add, answer1)
operate_and_compare(a2, a1, np.add, answer2)
operate_and_compare(a1, a3, np.add, answer1)
operate_and_compare(a3, a1, np.add, answer2)
# Test dimensionless quantities
a1 = YTArray([1, 2, 3])
a2 = array([4, 5, 6])
a3 = [4, 5, 6]
answer = YTArray([5, 7, 9])
operate_and_compare(a1, a2, operator.add, answer)
operate_and_compare(a2, a1, operator.add, answer)
operate_and_compare(a1, a3, operator.add, answer)
operate_and_compare(a3, a1, operator.add, answer)
operate_and_compare(a1, a2, np.add, answer)
operate_and_compare(a2, a1, np.add, answer)
operate_and_compare(a1, a3, np.add, answer)
operate_and_compare(a3, a1, np.add, answer)
# Catch the different dimensions error
a1 = YTArray([1, 2, 3], 'm')
a2 = YTArray([4, 5, 6], 'kg')
a3 = [7, 8, 9]
a4 = YTArray([10, 11, 12], '')
assert_raises(YTUnitOperationError, operator.add, a1, a2)
assert_raises(YTUnitOperationError, operator.iadd, a1, a2)
assert_raises(YTUnitOperationError, operator.add, a1, a3)
assert_raises(YTUnitOperationError, operator.iadd, a1, a3)
assert_raises(YTUnitOperationError, operator.add, a3, a1)
assert_raises(YTUnitOperationError, operator.iadd, a3, a1)
assert_raises(YTUnitOperationError, operator.add, a1, a4)
assert_raises(YTUnitOperationError, operator.iadd, a1, a4)
assert_raises(YTUnitOperationError, operator.add, a4, a1)
assert_raises(YTUnitOperationError, operator.iadd, a4, a1)
# adding with zero is allowed irrespective of the units
zeros = np.zeros(3)
zeros_yta_dimless = YTArray(zeros, 'dimensionless')
zeros_yta_length = YTArray(zeros, 'm')
zeros_yta_mass = YTArray(zeros, 'kg')
operands = [
0,
YTQuantity(0),
YTQuantity(0, 'kg'), zeros, zeros_yta_dimless, zeros_yta_length,
zeros_yta_mass
]
for op in [operator.add, np.add]:
for operand in operands:
operate_and_compare(a1, operand, op, a1)
operate_and_compare(operand, a1, op, a1)
operate_and_compare(4 * m, operand, op, 4 * m)
operate_and_compare(operand, 4 * m, op, 4 * m)
def test_addition():
"""
Test addition of two YTArrays
"""
# Same units
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'cm')
a3 = [4 * cm, 5 * cm, 6 * cm]
answer = YTArray([5, 7, 9], 'cm')
yield operate_and_compare, a1, a2, operator.add, answer
yield operate_and_compare, a2, a1, operator.add, answer
yield operate_and_compare, a1, a3, operator.add, answer
yield operate_and_compare, a3, a1, operator.add, answer
yield operate_and_compare, a2, a1, np.add, answer
yield operate_and_compare, a1, a2, np.add, answer
yield operate_and_compare, a1, a3, np.add, answer
yield operate_and_compare, a3, a1, np.add, answer
# different units
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'm')
a3 = [4 * m, 5 * m, 6 * m]
answer1 = YTArray([401, 502, 603], 'cm')
answer2 = YTArray([4.01, 5.02, 6.03], 'm')
yield operate_and_compare, a1, a2, operator.add, answer1
yield operate_and_compare, a2, a1, operator.add, answer2
yield operate_and_compare, a1, a3, operator.add, answer1
yield operate_and_compare, a3, a1, operator.add, answer1
yield assert_raises, YTUfuncUnitError, np.add, a1, a2
yield assert_raises, YTUfuncUnitError, np.add, a1, a3
# Test dimensionless quantities
a1 = YTArray([1, 2, 3])
a2 = array([4, 5, 6])
a3 = [4, 5, 6]
answer = YTArray([5, 7, 9])
yield operate_and_compare, a1, a2, operator.add, answer
yield operate_and_compare, a2, a1, operator.add, answer
yield operate_and_compare, a1, a3, operator.add, answer
yield operate_and_compare, a3, a1, operator.add, answer
yield operate_and_compare, a1, a2, np.add, answer
yield operate_and_compare, a2, a1, np.add, answer
yield operate_and_compare, a1, a3, np.add, answer
yield operate_and_compare, a3, a1, np.add, answer
# Catch the different dimensions error
a1 = YTArray([1, 2, 3], 'm')
a2 = YTArray([4, 5, 6], 'kg')
yield assert_raises, YTUnitOperationError, operator.add, a1, a2
yield assert_raises, YTUnitOperationError, operator.iadd, a1, a2
def test_ufuncs():
for ufunc in unary_operators:
if ufunc is None:
continue
unary_ufunc_comparison(ufunc, YTArray([.3, .4, .5], 'cm'))
unary_ufunc_comparison(ufunc, YTArray([12, 23, 47], 'g'))
unary_ufunc_comparison(ufunc, YTArray([2, 4, -6], 'erg/m**3'))
for ufunc in binary_operators:
if ufunc is None:
continue
# arr**arr is undefined for arrays with units because
# each element of the result would have different units.
if ufunc is np.power:
a = YTArray([.3, .4, .5], 'cm')
b = YTArray([.1, .2, .3], 'dimensionless')
c = np.array(b)
d = YTArray([1., 2., 3.], 'g')
binary_ufunc_comparison(ufunc, a, b)
binary_ufunc_comparison(ufunc, a, c)
assert_raises(YTUnitOperationError, ufunc, a, d)
continue
a = YTArray([.3, .4, .5], 'cm')
b = YTArray([.1, .2, .3], 'cm')
c = YTArray([.1, .2, .3], 'm')
d = YTArray([.1, .2, .3], 'g')
e = YTArray([.1, .2, .3], 'erg/m**3')
for pair in itertools.product([a, b, c, d, e], repeat=2):
binary_ufunc_comparison(ufunc, pair[0], pair[1])
def unary_ufunc_comparison(ufunc, a):
out = a.copy()
a_array = a.to_ndarray()
if ufunc in (
np.isreal,
np.iscomplex,
):
# According to the numpy docs, these two explicitly do not do
# in-place copies.
ret = ufunc(a)
assert_true(not hasattr(ret, 'units'))
assert_array_equal(ret, ufunc(a))
elif ufunc in (np.exp, np.exp2, np.log, np.log2, np.log10, np.expm1,
np.log1p, np.sin, np.cos, np.tan, np.arcsin, np.arccos,
np.arctan, np.sinh, np.cosh, np.tanh, np.arccosh,
np.arcsinh, np.arctanh, np.deg2rad, np.rad2deg, np.isfinite,
np.isinf, np.isnan, np.signbit, np.sign, np.rint,
np.logical_not):
# These operations should return identical results compared to numpy.
try:
ret = ufunc(a, out=out)
except YTUnitOperationError:
assert_true(ufunc in (np.deg2rad, np.rad2deg))
ret = ufunc(YTArray(a, '1'))
assert_array_equal(ret, out)
assert_array_equal(ret, ufunc(a_array))
# In-place copies do not drop units.
assert_true(hasattr(out, 'units'))
assert_true(not hasattr(ret, 'units'))
elif ufunc in (np.absolute, np.conjugate, np.floor, np.ceil, np.trunc,
np.negative):
ret = ufunc(a, out=out)
assert_array_equal(ret, out)
assert_array_equal(ret.to_ndarray(), ufunc(a_array))
assert_true(ret.units == out.units)
elif ufunc in (np.ones_like, np.square, np.sqrt, np.reciprocal):
if ufunc is np.ones_like:
ret = ufunc(a)
else:
ret = ufunc(a, out=out)
assert_array_equal(ret, out)
assert_array_equal(ret.to_ndarray(), ufunc(a_array))
if ufunc is np.square:
assert_true(out.units == a.units**2)
assert_true(ret.units == a.units**2)
elif ufunc is np.sqrt:
assert_true(out.units == a.units**0.5)
assert_true(ret.units == a.units**0.5)
elif ufunc is np.reciprocal:
assert_true(out.units == a.units**-1)
assert_true(ret.units == a.units**-1)
elif ufunc is np.modf:
ret1, ret2 = ufunc(a)
npret1, npret2 = ufunc(a_array)
assert_array_equal(ret1.to_ndarray(), npret1)
assert_array_equal(ret2.to_ndarray(), npret2)
elif ufunc is np.frexp:
ret1, ret2 = ufunc(a)
npret1, npret2 = ufunc(a_array)
assert_array_equal(ret1, npret1)
assert_array_equal(ret2, npret2)
else:
# There shouldn't be any untested ufuncs.
assert_true(False)
def calculate_spectrum(self,
data_source=None,
star_mass=None,
star_creation_time=None,
star_metallicity_fraction=None,
star_metallicity_constant=None,
min_age=YTQuantity(0.0, 'yr')):
r"""For the set of stars, calculate the collective spectrum.
Attached to the output are several useful objects:
Attributes
----------
final_spec: array
The collective spectrum in units of flux binned in wavelength.
wavelength: array
The wavelength for the spectrum bins, in Angstroms.
total_mass: float
Total mass of all the stars.
avg_mass: float
Average mass of all the stars.
avg_metal: float
Average metallicity of all the stars.
Parameters
----------
data_source : AMRRegion object, optional
The region from which stars are extracted for analysis. If this is
not specified, the next three parameters must be supplied.
star_mass : Array or list of floats
An array of star masses in Msun units.
star_creation_time : Array or list of floats
An array of star creation times in code units.
star_metallicity_fraction : Array or list of floats
An array of star metallicity fractions, in code
units (which is not Z/Zsun, rather just Z).
star_metallicity_constant : Float
If desired, override the star
metallicity fraction of all the stars to the given value.
min_age : Float
Removes young stars younger than this number (in years)
from the spectrum. Default: 0 (all stars).
Examples
--------
>>> import yt
>>> from yt.analysis_modules.star_analysis.api import SpectrumBuilder
>>> ds = yt.load("Enzo_64/RD0006/RedshiftOutput0006")
>>> spec = SpectrumBuilder(ds, "bc", model="salpeter")
>>> sp = ds.sphere([0.5, 0.5, 0.5], 0.1)
>>> spec.calculate_spectrum(data_source=sp, min_age=1.e6)
"""
# Initialize values
self.final_spec = np.zeros(self.wavelength.size, dtype='float64')
self._data_source = data_source
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(star_metallicity_fraction, YTArray):
assert \
star_metallicity_fraction.units.same_dimensions_as(Zsun.units)
elif star_metallicity_fraction is not None:
star_metallicity_fraction = self._ds.arr(star_metallicity_fraction,
'code_metallicity')
self.star_metallicity_fraction = star_metallicity_fraction
if isinstance(min_age, YTQuantity):
assert min_age.units.same_dimensions_as(s.units)
elif min_age is not None:
min_age = YTQuantity(min_age, 'yr')
self.min_age = min_age
# Check to make sure we have the right set of data.
if data_source is None:
if self.star_mass is None or self.star_creation_time is None or \
(star_metallicity_fraction is None and
star_metallicity_constant is None):
mylog.error("""
If data_source is not provided, all of these paramters
need to be set:
star_mass (array, Msun),
star_creation_time (array, code units),
And one of:
star_metallicity_fraction (array, code units).
--OR--
star_metallicity_constant (float, code units).
""")
return None
if star_metallicity_fraction is not None:
self.star_metal = star_metallicity_fraction
else:
self.star_metal = \
self._ds.arr(np.ones_like(self.star_mass) *
star_metallicity_constant, 'Zsun')
else:
# Get the data we need.
if self.filter_provided:
ct = self._filter['creation_time']
# mass_stars = self._data_source[self._filter, "particle_mass"]
if star_metallicity_constant is None:
self.star_metal = self._data_source[
self._filter, "metallicity_fraction"].in_units('Zsun')
else:
self.star_metal = \
self._ds.arr(np.ones_like(
self._data_source[self._filter,
"metallicity_fraction"]) *
star_metallicity_constant, "Zsun")
else:
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']
self.star_creation_time = ct[mask]
self.star_mass = self._data_source['particle_mass'][
mask].in_units('Msun')
if star_metallicity_constant is not None:
self.star_metal = self._ds.arr(
np.ones_like(self.star_mass) *
star_metallicity_constant, 'Zsun')
else:
self.star_metal = self._data_source[
"metallicity_fraction"][mask].in_units('Zsun')
# Age of star in years.
dt = (self.time_now - self.star_creation_time).in_units('yr')
dt[dt < 0.0] = 0.0
# Remove young stars
sub = dt >= self.min_age
if len(sub) == 0:
return
self.star_metal = self.star_metal[sub]
dt = dt[sub]
self.star_creation_time = self.star_creation_time[sub]
# Figure out which METALS bin the star goes into.
Mindex = np.digitize(self.star_metal.in_units('Zsun'), METALS)
# Replace the indices with strings.
Mname = MtoD[Mindex]
# Figure out which age bin this star goes into.
Aindex = np.digitize(dt, self.age)
# Ratios used for the interpolation.
ratio1 = (dt - self.age[Aindex - 1]) / \
(self.age[Aindex] - self.age[Aindex - 1])
ratio2 = (self.age[Aindex] - dt) / \
(self.age[Aindex] - self.age[Aindex - 1])
# Sort the stars by metallicity and then by age, which should reduce
# memory access time by a little bit in the loop.
indexes = np.arange(self.star_metal.size)
sort = np.asarray(
[indexes[i] for i in np.lexsort([indexes, Aindex, Mname])])
Mname = Mname[sort]
Aindex = Aindex[sort]
ratio1 = ratio1[sort]
ratio2 = ratio2[sort]
self.star_mass = self.star_mass[sort]
self.star_creation_time = self.star_creation_time[sort]
self.star_metal = self.star_metal[sort]
# Interpolate the flux for each star, adding to the total by weight.
pbar = get_pbar("Calculating fluxes", len(self.star_mass))
for i, star in enumerate(
izip(Mname, Aindex, ratio1, ratio2, self.star_mass)):
# Pick the right age bin for the right flux array.
flux = self.flux[star[0]][star[1], :]
# Get the one just before the one above.
flux_1 = self.flux[star[0]][star[1] - 1, :]
# interpolate in log(flux), linear in time.
int_flux = star[3] * np.log10(flux_1) + star[2] * np.log10(flux)
# Add this flux to the total, weighted by mass.
self.final_spec += np.power(10., int_flux) * star[4]
pbar.update(i)
pbar.finish()
# Normalize.
self.total_mass = self.star_mass.sum()
self.avg_mass = self.star_mass.mean()
tot_metal = (self.star_metal * self.star_mass).sum()
if tot_metal > 0:
self.avg_metal = math.log10(
(tot_metal / self.total_mass).in_units('Zsun'))
else:
self.avg_metal = -99
def __missing__(self, item):
if not isinstance(item, tuple):
field = ("unknown", item)
else:
field = item
finfo = self.ds._get_field_info(*field)
params, permute_params = finfo._get_needed_parameters(self)
self.field_parameters.update(params)
# For those cases where we are guessing the field type, we will
# need to re-update -- otherwise, our item will always not have the
# field type. This can lead to, for instance, "unknown" particle
# types not getting correctly identified.
# Note that the *only* way this works is if we also fix our field
# dependencies during checking. Bug #627 talks about this.
item = self.ds._last_freq
if finfo is not None and finfo._function.__name__ != 'NullFunc':
try:
for param, param_v in permute_params.items():
for v in param_v:
self.field_parameters[param] = v
vv = finfo(self)
if not permute_params:
vv = finfo(self)
except NeedsGridType as exc:
ngz = exc.ghost_zones
nfd = FieldDetector(
self.nd + ngz * 2,
ds=self.ds,
field_parameters=self.field_parameters.copy())
nfd._num_ghost_zones = ngz
vv = finfo(nfd)
if ngz > 0: vv = vv[ngz:-ngz, ngz:-ngz, ngz:-ngz]
for i in nfd.requested:
if i not in self.requested: self.requested.append(i)
for i in nfd.requested_parameters:
if i not in self.requested_parameters:
self.requested_parameters.append(i)
if vv is not None:
if not self.flat: self[item] = vv
else: self[item] = vv.ravel()
return self[item]
elif finfo is not None and finfo.particle_type:
if "particle_position" in (item, item[1]) or \
"particle_velocity" in (item, item[1]) or \
"particle_magnetic_field" in (item, item[1]) or \
"Velocity" in (item, item[1]) or \
"Velocities" in (item, item[1]) or \
"Coordinates" in (item, item[1]) or \
"MagneticField" in (item, item[1]):
# A vector
self[item] = \
YTArray(np.ones((self.NumberOfParticles, 3)),
finfo.units, registry=self.ds.unit_registry)
elif "FourMetalFractions" in (item, item[1]):
self[item] = \
YTArray(np.ones((self.NumberOfParticles, 4)),
finfo.units, registry=self.ds.unit_registry)
else:
# Not a vector
self[item] = \
YTArray(np.ones(self.NumberOfParticles),
finfo.units, registry=self.ds.unit_registry)
if item == ('STAR', 'BIRTH_TIME'):
# hack for the artio frontend so we pass valid times to
# the artio functions for calculating physical times
# from internal times
self[item] *= -0.1
self.requested.append(item)
return self[item]
self.requested.append(item)
if item not in self:
self[item] = self._read_data(item)
return self[item]
def create_spectral_slabs(filename, slab_centers, slab_width, **kwargs):
r"""
Given a dictionary of spectral slab centers and a width in
spectral units, extract data from a spectral cube at these slab
centers and return a `FITSDataset` instance containing the different
slabs as separate yt fields. Useful for extracting individual
lines from a spectral cube and separating them out as different fields.
Requires the SpectralCube (https://spectral-cube.readthedocs.io/en/latest/)
library.
All keyword arguments will be passed on to the `FITSDataset` constructor.
Parameters
----------
filename : string
The spectral cube FITS file to extract the data from.
slab_centers : dict of (float, string) tuples or YTQuantities
The centers of the slabs, where the keys are the names
of the new fields and the values are (float, string) tuples or
YTQuantities, specifying a value for each center and its unit.
slab_width : YTQuantity or (float, string) tuple
The width of the slab along the spectral axis.
Examples
--------
>>> slab_centers = {'13CN': (218.03117, 'GHz'),
... 'CH3CH2CHO': (218.284256, 'GHz'),
... 'CH3NH2': (218.40956, 'GHz')}
>>> slab_width = (0.05, "GHz")
>>> ds = create_spectral_slabs("intensity_cube.fits",
... slab_centers, slab_width,
... nan_mask=0.0)
"""
from spectral_cube import SpectralCube
from yt.visualization.fits_image import FITSImageData
from yt.frontends.fits.api import FITSDataset
cube = SpectralCube.read(filename)
if not isinstance(slab_width, YTQuantity):
slab_width = YTQuantity(slab_width[0], slab_width[1])
slab_data = {}
field_units = cube.header.get("bunit", "dimensionless")
for k, v in slab_centers.items():
if not isinstance(v, YTQuantity):
slab_center = YTQuantity(v[0], v[1])
else:
slab_center = v
mylog.info("Adding slab field %s at %g %s" %
(k, slab_center.v, slab_center.units))
slab_lo = (slab_center - 0.5 * slab_width).to_astropy()
slab_hi = (slab_center + 0.5 * slab_width).to_astropy()
subcube = cube.spectral_slab(slab_lo, slab_hi)
slab_data[k] = YTArray(subcube.filled_data[:, :, :], field_units)
width = subcube.header["naxis3"] * cube.header["cdelt3"]
w = subcube.wcs.copy()
w.wcs.crpix[-1] = 0.5
w.wcs.crval[-1] = -0.5 * width
fid = FITSImageData(slab_data, wcs=w)
for hdu in fid:
hdu.header.pop("RESTFREQ", None)
hdu.header.pop("RESTFRQ", None)
ds = FITSDataset(fid, **kwargs)
return ds
from formulas.physics import \
maxwellian_speed, \
maxwellian_velocity
from numpy.testing import assert_allclose
import numpy as np
import yt.units as u
from yt.units.yt_array import YTArray
v_th = 1000 * u.km / u.s
n_0 = 1.0 * u.cm**-3 * u.km**-3 * u.s**3
v = YTArray(np.linspace(-1000., 1000., 100), "km/s")
A = 1.0
def test_maxwellian_velocity():
p = maxwellian_velocity(sigma="v_th")
p.set_param_values(A=A, v_th=v_th)
assert_allclose(
p(v).v,
(np.exp(-0.5 * (v / v_th)**2) / (np.sqrt(2. * np.pi) * v_th)).v)
assert str(p(v).units) == str((1 / v_th).units)
def test_maxwellian_speed():
p = maxwellian_speed(sigma="v_th", A="n_0")
p.set_param_values(n_0=n_0, v_th=v_th)
assert_allclose(
p(v).v, (4. * np.pi * v * v * n_0 * np.exp(-0.5 * (v / v_th)**2) /
(np.sqrt(2. * np.pi) * v_th)**3).v)
assert str(p(v).units) == str((n_0 * v_th**-1).units)
"set_cmap": [(('density', 'RdBu'), {}), (('density', 'kamae'), {})],
"set_font": [((OrderedDict(sorted(FPROPS.items(),
key=lambda t: t[0])), ), {})],
"set_log": [(('density', False), {})],
"set_window_size": [((7.0, ), {})],
"set_zlim": [(('density', 1e-25, 1e-23), {}),
(('density', 1e-25, None), {
'dynamic_range': 4
})],
"zoom": [((10, ), {})],
"toggle_right_handed": [((), {})]
}
CENTER_SPECS = ("m", "M", "max", "Max", "c", "C", "center", "Center",
[0.5, 0.5, 0.5], [[0.2, 0.3, 0.4],
"cm"], YTArray([0.3, 0.4, 0.7], "cm"))
WIDTH_SPECS = {
# Width choices map to xlim, ylim, width, axes_unit_name 4-tuples
None: (
((0, 'code_length'), (1, 'code_length')),
((0, 'code_length'), (1, 'code_length')),
((1, 'code_length'), (1, 'code_length')),
None,
),
0.2: (
((0.4, 'code_length'), (0.6, 'code_length')),
((0.4, 'code_length'), (0.6, 'code_length')),
((0.2, 'code_length'), (0.2, 'code_length')),
None,
),
def test_subtraction():
"""
Test subtraction of two YTArrays
"""
# Same units
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'cm')
a3 = [4 * cm, 5 * cm, 6 * cm]
answer1 = YTArray([-3, -3, -3], 'cm')
answer2 = YTArray([3, 3, 3], 'cm')
operate_and_compare(a1, a2, operator.sub, answer1)
operate_and_compare(a2, a1, operator.sub, answer2)
operate_and_compare(a1, a3, operator.sub, answer1)
operate_and_compare(a3, a1, operator.sub, answer2)
operate_and_compare(a1, a2, np.subtract, answer1)
operate_and_compare(a2, a1, np.subtract, answer2)
operate_and_compare(a1, a3, np.subtract, answer1)
operate_and_compare(a3, a1, np.subtract, answer2)
# different units
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'm')
a3 = [4 * m, 5 * m, 6 * m]
answer1 = YTArray([-399, -498, -597], 'cm')
answer2 = YTArray([3.99, 4.98, 5.97], 'm')
answer3 = YTArray([399, 498, 597], 'cm')
operate_and_compare(a1, a2, operator.sub, answer1)
operate_and_compare(a2, a1, operator.sub, answer2)
operate_and_compare(a1, a3, operator.sub, answer1)
operate_and_compare(a3, a1, operator.sub, answer3)
if LooseVersion(np.__version__) < LooseVersion('1.13.0'):
assert_raises(YTUfuncUnitError, np.subtract, a1, a2)
assert_raises(YTUfuncUnitError, np.subtract, a1, a3)
else:
operate_and_compare(a1, a2, np.subtract, answer1)
operate_and_compare(a2, a1, np.subtract, answer2)
operate_and_compare(a1, a3, np.subtract, answer1)
operate_and_compare(a3, a1, np.subtract, answer3)
# Test dimensionless quantities
a1 = YTArray([1, 2, 3])
a2 = array([4, 5, 6])
a3 = [4, 5, 6]
answer1 = YTArray([-3, -3, -3])
answer2 = YTArray([3, 3, 3])
operate_and_compare(a1, a2, operator.sub, answer1)
operate_and_compare(a2, a1, operator.sub, answer2)
operate_and_compare(a1, a3, operator.sub, answer1)
operate_and_compare(a3, a1, operator.sub, answer2)
operate_and_compare(a1, a2, np.subtract, answer1)
operate_and_compare(a2, a1, np.subtract, answer2)
operate_and_compare(a1, a3, np.subtract, answer1)
operate_and_compare(a3, a1, np.subtract, answer2)
# Catch the different dimensions error
a1 = YTArray([1, 2, 3], 'm')
a2 = YTArray([4, 5, 6], 'kg')
a3 = [7, 8, 9]
a4 = YTArray([10, 11, 12], '')
assert_raises(YTUnitOperationError, operator.sub, a1, a2)
assert_raises(YTUnitOperationError, operator.isub, a1, a2)
assert_raises(YTUnitOperationError, operator.sub, a1, a3)
assert_raises(YTUnitOperationError, operator.isub, a1, a3)
assert_raises(YTUnitOperationError, operator.sub, a3, a1)
assert_raises(YTUnitOperationError, operator.isub, a3, a1)
assert_raises(YTUnitOperationError, operator.sub, a1, a4)
assert_raises(YTUnitOperationError, operator.isub, a1, a4)
assert_raises(YTUnitOperationError, operator.sub, a4, a1)
assert_raises(YTUnitOperationError, operator.isub, a4, a1)
# subtracting with zero is allowed irrespective of the units
zeros = np.zeros(3)
zeros_yta_dimless = YTArray(zeros, 'dimensionless')
zeros_yta_length = YTArray(zeros, 'm')
zeros_yta_mass = YTArray(zeros, 'kg')
operands = [
0,
YTQuantity(0),
YTQuantity(0, 'kg'), zeros, zeros_yta_dimless, zeros_yta_length,
zeros_yta_mass
]
for op in [operator.sub, np.subtract]:
for operand in operands:
operate_and_compare(a1, operand, op, a1)
operate_and_compare(operand, a1, op, -a1)
operate_and_compare(4 * m, operand, op, 4 * m)
operate_and_compare(operand, 4 * m, op, -4 * m)
def z_from_t_analytic(my_time, hubble_constant=0.7,
omega_matter=0.3, omega_lambda=0.7):
"""
Compute the redshift from time after the big bang. This is based on
Enzo's CosmologyComputeExpansionFactor.C, but altered to use physical
units.
"""
hubble_constant = YTQuantity(hubble_constant, "100*km/s/Mpc")
omega_curvature = 1.0 - omega_matter - omega_lambda
OMEGA_TOLERANCE = 1e-5
ETA_TOLERANCE = 1.0e-10
# Convert the time to Time * H0.
if not isinstance(my_time, YTArray):
my_time = YTArray(my_time, "s")
t0 = (my_time.in_units("s") *
hubble_constant.in_units("1/s")).to_ndarray()
# For a flat universe with omega_matter = 1, it's easy.
if np.fabs(omega_matter-1) < OMEGA_TOLERANCE and \
omega_lambda < OMEGA_TOLERANCE:
a = np.power(1.5 * t0, 2.0/3.0)
# For omega_matter < 1 and omega_lambda == 0 see
# Peebles 1993, eq. 13-3, 13-10.
# Actually, this is a little tricky since we must solve an equation
# of the form eta - np.sinh(eta) + x = 0..
elif omega_matter < 1 and omega_lambda < OMEGA_TOLERANCE:
x = 2*t0*np.power(1.0 - omega_matter, 1.5) / omega_matter
# Compute eta in a three step process, first from a third-order
# Taylor expansion of the formula above, then use that in a fifth-order
# approximation. Then finally, iterate on the formula itself, solving for
# eta. This works well because parts 1 & 2 are an excellent approximation
# when x is small and part 3 converges quickly when x is large.
eta = np.power(6*x, 1.0/3.0) # part 1
eta = np.power(120*x/(20+eta*eta), 1.0/3.0) # part 2
mask = np.ones(eta.size, dtype=bool)
max_iter = 1000
for i in range(max_iter): # part 3
eta_old = eta[mask]
eta[mask] = np.arcsinh(eta[mask] + x[mask])
mask[mask] = np.fabs(eta[mask]-eta_old) >= ETA_TOLERANCE
if not mask.any():
break
if (i == max_iter-1):
raise RuntimeError(
"No convergence after %d iterations." % i)
# Now use eta to compute the expansion factor (eq. 13-10, part 2).
a = omega_matter/(2.0*(1.0 - omega_matter))*\
(np.cosh(eta) - 1.0)
# For flat universe, with non-zero omega_lambda, see eq. 13-20.
elif np.fabs(omega_curvature) < OMEGA_TOLERANCE and \
omega_lambda > OMEGA_TOLERANCE:
a = np.power(omega_matter / (1 - omega_matter), 1.0/3.0) * \
np.power(np.sinh(1.5 * np.sqrt(1.0 - omega_matter) * \
t0), 2.0/3.0)
else:
raise NotImplementedError
redshift = (1.0/a) - 1.0
return redshift
def test_multiplication():
"""
Test multiplication of two YTArrays
"""
# Same units
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'cm')
a3 = [4 * cm, 5 * cm, 6 * cm]
answer = YTArray([4, 10, 18], 'cm**2')
operate_and_compare(a1, a2, operator.mul, answer)
operate_and_compare(a2, a1, operator.mul, answer)
operate_and_compare(a1, a3, operator.mul, answer)
operate_and_compare(a3, a1, operator.mul, answer)
operate_and_compare(a1, a2, np.multiply, answer)
operate_and_compare(a2, a1, np.multiply, answer)
operate_and_compare(a1, a3, np.multiply, answer)
operate_and_compare(a3, a1, np.multiply, answer)
# different units, same dimension
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'm')
a3 = [4 * m, 5 * m, 6 * m]
answer1 = YTArray([400, 1000, 1800], 'cm**2')
answer2 = YTArray([.04, .10, .18], 'm**2')
answer3 = YTArray([4, 10, 18], 'cm*m')
operate_and_compare(a1, a2, operator.mul, answer1)
operate_and_compare(a2, a1, operator.mul, answer2)
operate_and_compare(a1, a3, operator.mul, answer1)
operate_and_compare(a3, a1, operator.mul, answer2)
operate_and_compare(a1, a2, np.multiply, answer3)
operate_and_compare(a2, a1, np.multiply, answer3)
operate_and_compare(a1, a3, np.multiply, answer3)
operate_and_compare(a3, a1, np.multiply, answer3)
# different dimensions
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([4, 5, 6], 'g')
a3 = [4 * g, 5 * g, 6 * g]
answer = YTArray([4, 10, 18], 'cm*g')
operate_and_compare(a1, a2, operator.mul, answer)
operate_and_compare(a2, a1, operator.mul, answer)
operate_and_compare(a1, a3, operator.mul, answer)
operate_and_compare(a3, a1, operator.mul, answer)
operate_and_compare(a1, a2, np.multiply, answer)
operate_and_compare(a2, a1, np.multiply, answer)
operate_and_compare(a1, a3, np.multiply, answer)
operate_and_compare(a3, a1, np.multiply, answer)
# One dimensionless, one unitful
a1 = YTArray([1, 2, 3], 'cm')
a2 = array([4, 5, 6])
a3 = [4, 5, 6]
answer = YTArray([4, 10, 18], 'cm')
operate_and_compare(a1, a2, operator.mul, answer)
operate_and_compare(a2, a1, operator.mul, answer)
operate_and_compare(a1, a3, operator.mul, answer)
operate_and_compare(a3, a1, operator.mul, answer)
operate_and_compare(a1, a2, np.multiply, answer)
operate_and_compare(a2, a1, np.multiply, answer)
operate_and_compare(a1, a3, np.multiply, answer)
operate_and_compare(a3, a1, np.multiply, answer)
# Both dimensionless quantities
a1 = YTArray([1, 2, 3])
a2 = array([4, 5, 6])
a3 = [4, 5, 6]
answer = YTArray([4, 10, 18])
operate_and_compare(a1, a2, operator.mul, answer)
operate_and_compare(a2, a1, operator.mul, answer)
operate_and_compare(a1, a3, operator.mul, answer)
operate_and_compare(a3, a1, operator.mul, answer)
operate_and_compare(a1, a2, np.multiply, answer)
operate_and_compare(a2, a1, np.multiply, answer)
operate_and_compare(a1, a3, np.multiply, answer)
operate_and_compare(a3, a1, np.multiply, answer)
def data(self):
return YTArray(self.hdu.data, self.units)
def _creation_time(field, data):
return YTArray(
data.ds._handle.tphys_from_tcode_array(data["STAR",
"BIRTH_TIME"]),
"yr",
)
