def test_latex_repr():
registry = UnitRegistry()
# create a fake comoving unit
registry.add('pccm',
registry.lut['pc'][0] / (1 + 2),
length,
"\\rm{pc}/(1+z)",
prefixable=True)
test_unit = Unit('Mpccm', registry=registry)
assert_almost_equal(test_unit.base_value, m_per_mpc / 3)
assert_equal(test_unit.latex_repr, r'\rm{Mpc}/(1+z)')
test_unit = Unit('cm**-3', base_value=1.0, registry=registry)
assert_equal(test_unit.latex_repr, '\\frac{1}{\\rm{cm}^{3}}')
test_unit = Unit('m_geom/l_geom**3')
assert_equal(test_unit.latex_repr, '\\frac{1}{M_\\odot^{2}}')
test_unit = Unit('1e9*cm')
assert_equal(test_unit.latex_repr, '1.0 \\times 10^{9}\\ \\rm{cm}')
test_unit = Unit('1.0*cm')
assert_equal(test_unit.latex_repr, '\\rm{cm}')
def test_latex_repr():
registry = UnitRegistry()
# create a fake comoving unit
registry.add(
"pccm",
registry.lut["pc"][0] / (1 + 2),
length,
"\\rm{pc}/(1+z)",
prefixable=True,
)
test_unit = Unit("Mpccm", registry=registry)
assert_almost_equal(test_unit.base_value, m_per_mpc / 3)
assert_equal(test_unit.latex_repr, r"\rm{Mpc}/(1+z)")
test_unit = Unit("cm**-3", base_value=1.0, registry=registry)
assert_equal(test_unit.latex_repr, "\\frac{1}{\\rm{cm}^{3}}")
test_unit = Unit("m_geom/l_geom**3")
assert_equal(test_unit.latex_repr, "\\frac{1}{\\rm{M}_\\odot^{2}}")
test_unit = Unit("1e9*cm")
assert_equal(test_unit.latex_repr, "1.0 \\times 10^{9}\\ \\rm{cm}")
test_unit = Unit("1.0*cm")
assert_equal(test_unit.latex_repr, "\\rm{cm}")
def test_registry_json():
reg = UnitRegistry()
json_reg = reg.to_json()
unserialized_reg = UnitRegistry.from_json(json_reg)
assert reg.lut == unserialized_reg.lut
assert reg.lut["m"][1] is length
assert reg.lut["erg"][1] is energy
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 test_bad_unit_system():
with pytest.raises(IllDefinedUnitSystem):
UnitSystem("atomic", "nm", "fs", "nK", "rad")
with pytest.raises(IllDefinedUnitSystem):
UnitSystem("atomic", "nm", "fs", "nK", "rad", registry=UnitRegistry())
with pytest.raises(IllDefinedUnitSystem):
UnitSystem("atomic", us.nm, us.fs, us.nK, us.rad)
with pytest.raises(IllDefinedUnitSystem):
UnitSystem("atomic", us.nm, us.fs, us.nK, us.rad, registry=UnitRegistry())
def test_old_registry_multiple_load():
# See Issue #157 for details
reg1 = UnitRegistry()
reg1.add("code_length", 1.0, length)
reg1.add("code_mass", 1.0, mass)
reg1.add("code_time", 1.0, time)
reg1.add("code_temperature", 1.0, temperature)
UnitSystem(
reg1.unit_system_id,
"code_length",
"code_mass",
"code_time",
"code_temperature",
registry=reg1,
)
cm = Unit("code_mass", registry=reg1)
cl = Unit("code_length", registry=reg1)
(cm / cl).latex_representation()
with open(OLD_JSON_PATH) as f:
json_data = f.read()
reg2 = UnitRegistry.from_json(json_data)
UnitSystem(
reg2.unit_system_id,
"code_length",
"code_mass",
"code_time",
"code_temperature",
registry=reg2,
)
def test_prefixable_units():
ureg = UnitRegistry()
pu = ureg.prefixable_units
assert 'm' in pu
assert 'pc' in pu
assert 'mol' in pu
ureg.add('foobar', 1.0, length, prefixable=True)
assert 'foobar' in ureg.prefixable_units
mfoobar = Unit('mfoobar', registry=ureg)
foobar = Unit('foobar', registry=ureg)
assert (1 * foobar) / (1 * mfoobar) == 1000
def test_prefixable_units():
ureg = UnitRegistry()
pu = ureg.prefixable_units
assert "m" in pu
assert "pc" in pu
assert "mol" in pu
ureg.add("foobar", 1.0, length, prefixable=True)
assert "foobar" in ureg.prefixable_units
mfoobar = Unit("mfoobar", registry=ureg)
foobar = Unit("foobar", registry=ureg)
assert (1 * foobar) / (1 * mfoobar) == 1000
def test_mixed_registry_operations():
reg = UnitRegistry(unit_system="cgs")
reg.add("fake_length", 0.001, length)
a = unyt_quantity(1, units="fake_length", registry=reg)
b = unyt_quantity(1, "cm")
assert_almost_equal(a + b, b + a)
assert_almost_equal(a - b, -(b - a))
assert_almost_equal(a * b, b * a)
assert_almost_equal(b / a, b / a.in_units("km"))
assert_almost_equal(a / b, a / b.in_units("km"))
def test_old_registry_json():
with open(OLD_JSON_PATH) as f:
json_text = f.read()
reg = UnitRegistry.from_json(json_text)
default_reg = UnitRegistry()
loaded_keys = reg.keys()
for k in default_reg.keys():
assert k in loaded_keys
loaded_val = reg[k]
val = default_reg[k]
assert_allclose(loaded_val[0], val[0])
assert loaded_val[1:] == val[1:]
def test_old_registry_json():
path = os.sep.join([os.path.dirname(__file__), "old_json_registry.txt"])
with open(path) as f:
json_text = f.read()
reg = UnitRegistry.from_json(json_text)
default_reg = UnitRegistry()
loaded_keys = reg.keys()
for k in default_reg.keys():
assert k in loaded_keys
loaded_val = reg[k]
val = default_reg[k]
assert_allclose(loaded_val[0], val[0])
assert loaded_val[1:] == val[1:]
def _set_units(self):
self.unit_registry = UnitRegistry()
self.unit_registry.add("code_time", 1.0, dimensions.time)
self.unit_registry.add("code_length", 1.0, dimensions.length)
if self.cosmological_simulation:
# Instantiate EnzoCosmology object for units and time conversions.
self.cosmology = EnzoCosmology(
self.parameters["CosmologyHubbleConstantNow"],
self.parameters["CosmologyOmegaMatterNow"],
self.parameters["CosmologyOmegaLambdaNow"],
self.parameters.get("CosmologyOmegaRadiationNow", 0.0),
0.0,
self.parameters["CosmologyInitialRedshift"],
unit_registry=self.unit_registry,
)
self.time_unit = self.cosmology.time_unit.in_units("s")
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.box_size, "Mpccm / h", registry=self.unit_registry
)
else:
self.time_unit = self.quan(self.parameters["TimeUnits"], "s")
self.length_unit = self.quan(self.parameters["LengthUnits"], "cm")
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.modify("code_time", self.time_unit)
self.unit_registry.modify("code_length", self.length_unit)
self.unit_registry.add(
"unitary", float(self.box_size.in_base()), self.length_unit.units.dimensions
)
def test_define_unit():
define_unit("mph", (1.0, "mile/hr"))
a = unyt_quantity(2.0, "mph")
b = unyt_quantity(1.0, "mile")
c = unyt_quantity(1.0, "hr")
assert a == 2.0 * b / c
d = unyt_quantity(1000.0, "cm**3")
define_unit("Baz", d, prefixable=True)
e = unyt_quantity(1.0, "mBaz")
f = unyt_quantity(1.0, "cm**3")
assert e == f
define_unit("Foo", (1.0, "V/sqrt(s)"))
g = unyt_quantity(1.0, "Foo")
volt = unyt_quantity(1.0, "V")
second = unyt_quantity(1.0, "s")
assert g == volt / second**(0.5)
# Test custom registry
reg = UnitRegistry()
define_unit("Foo", (1, "m"), registry=reg)
define_unit("Baz", (1, "Foo**2"), registry=reg)
h = unyt_quantity(1, "Baz", registry=reg)
i = unyt_quantity(1, "m**2", registry=reg)
assert h == i
def unit_registry(self):
"""
Unit system registry.
"""
if self._unit_registry is None:
self._unit_registry = UnitRegistry()
return self._unit_registry
def _parse_parameter_file(self):
self._prefix = \
self.filename[:self.filename.rfind(self._suffix)]
fh = h5py.File(self.filename, mode="r")
for attr in ["hubble_constant", "omega_matter", "omega_lambda"]:
setattr(self, attr, fh.attrs.get(attr, None))
if "unit_registry_json" in fh.attrs:
self.unit_registry = \
UnitRegistry.from_json(
parse_h5_attr(fh, "unit_registry_json"))
if "box_size" in fh.attrs:
self.box_size = _hdf5_yt_attr(fh,
"box_size",
unit_registry=self.unit_registry)
self.field_info.update(json.loads(parse_h5_attr(fh, "field_info")))
self._size = fh.attrs["total_trees"]
fh.close()
# analysis fields in sidecar files
analysis_filename = f"{self._prefix}-analysis{self._suffix}"
if os.path.exists(analysis_filename):
self.analysis_filename = analysis_filename
fh = h5py.File(analysis_filename, mode="r")
analysis_fi = json.loads(parse_h5_attr(fh, "field_info"))
fh.close()
for field in analysis_fi:
analysis_fi[field]["type"] = "analysis_saved"
self.field_info.update(analysis_fi)
else:
self.analysis_filename = None
self.field_list = list(self.field_info.keys())
def test_registry_contains():
ureg = UnitRegistry()
assert 'm' in ureg
assert 'cm' in ureg
assert 'erg' in ureg
assert 'Merg' in ureg
assert 'foobar' not in ureg
assert Unit('m', registry=ureg) in ureg
def __deepcopy__(self, memodict=None):
expr = str(self.expr)
base_value = copy.deepcopy(self.base_value)
base_offset = copy.deepcopy(self.base_offset)
dimensions = copy.deepcopy(self.dimensions)
lut = copy.deepcopy(self.registry.lut)
registry = UnitRegistry(lut=lut)
return Unit(expr, base_value, base_offset, dimensions, registry)
def test_registry_contains():
ureg = UnitRegistry()
assert "m" in ureg
assert "cm" in ureg
assert "erg" in ureg
assert "Merg" in ureg
assert "foobar" not in ureg
assert Unit("m", registry=ureg) in ureg
def hubble_constant(self, value):
self._hubble_constant = value
# reset the unit registry lut while preserving other changes
self.unit_registry = UnitRegistry.from_json(
self.unit_registry.to_json())
if 'h' in self.unit_registry:
self.unit_registry.modify("h", self.hubble_constant)
else:
self.unit_registry.add(
'h', self.hubble_constant, dimensionless)
def _parse_parameter_file(self):
fh = h5py.File(self.filename, "r")
for attr in ["hubble_constant", "omega_matter", "omega_lambda"]:
setattr(self, attr, fh.attrs[attr])
my_ur = UnitRegistry.from_json(parse_h5_attr(fh, "unit_registry_json"))
right = _hdf5_yt_attr(fh, "domain_right_edge", unit_registry=my_ur)
left = _hdf5_yt_attr(fh, "domain_left_edge", unit_registry=my_ur)
# Drop the "cm" suffix because all lengths will
# be in comoving units.
self.box_size = self.quan((right - left)[0].to("Mpccm/h"), "Mpc/h")
fh.close()
def __init__(self,
filename,
log_filename=None,
hubble_constant=1.0,
box_size=None,
omega_matter=None,
omega_lambda=None):
self.unit_registry = UnitRegistry()
self.log_filename = log_filename
self.hubble_constant = hubble_constant
self.omega_matter = omega_matter
self.omega_lambda = omega_lambda
self._box_size_user = box_size
super().__init__(filename)
def test_code_unit():
from unyt import UnitRegistry
ureg = UnitRegistry()
ureg.add('code_length', 10., length)
u = Unit('code_length', registry=ureg)
assert u.is_code_unit is True
assert u.get_base_equivalent() == Unit('m')
u = Unit('cm')
assert u.is_code_unit is False
UnitSystem(ureg.unit_system_id, 'code_length', 'kg', 's', registry=ureg)
u = Unit('cm', registry=ureg)
ue = u.get_base_equivalent('code')
assert str(ue) == 'code_length'
assert ue.base_value == 10
assert ue.dimensions is length
class FakeDataset(object):
unit_registry = ureg
ds = FakeDataset()
UnitSystem(ds, 'code_length', 'kg', 's', registry=ureg)
u = Unit('cm', registry=ureg)
ue = u.get_base_equivalent(ds)
assert str(ue) == 'code_length'
assert ue.base_value == 10
assert ue.dimensions is length
with pytest.raises(UnitParseError):
Unit('code_length')
def test_base_equivalent():
"""
Check base equivalent of a unit.
"""
Msun_mks = mass_sun_kg
Mpc_mks = m_per_mpc
u1 = Unit("Msun * Mpc**-3")
u2 = Unit("kg * m**-3")
u3 = u1.get_base_equivalent()
assert u2.expr == u3.expr
assert u2 == u3
assert_allclose_units(u1.base_value, Msun_mks / Mpc_mks ** 3, 1e-12)
assert u2.base_value == 1
assert u3.base_value == 1
mass_density = mass / length ** 3
assert u1.dimensions == mass_density
assert u2.dimensions == mass_density
assert u3.dimensions == mass_density
assert_allclose_units(
u1.get_conversion_factor(u3)[0], Msun_mks / Mpc_mks ** 3, 1e-12
)
with pytest.raises(UnitConversionError):
u1.get_conversion_factor(Unit("m"))
with pytest.raises(UnitConversionError):
u1.get_conversion_factor(Unit("degF"))
reg = UnitRegistry(unit_system=cgs_unit_system)
u = Unit("kg", registry=reg)
assert u.get_base_equivalent() == Unit("g")
u = Unit("kg")
assert u.get_base_equivalent() == Unit("kg")
u = Unit("A")
assert u.get_base_equivalent(unit_system="mks") == Unit("A")
def test_code_unit():
from unyt import UnitRegistry
ureg = UnitRegistry()
ureg.add("code_length", 10.0, length)
ureg.add("code_magnetic_field", 2.0, magnetic_field_mks)
u = Unit("code_length", registry=ureg)
assert u.is_code_unit is True
assert u.get_base_equivalent() == Unit("m")
u = Unit("cm")
assert u.is_code_unit is False
u = Unit("code_magnetic_field", registry=ureg)
assert u.get_base_equivalent("mks") == Unit("T")
with pytest.raises(UnitsNotReducible):
assert u.get_base_equivalent("cgs")
# see issue #60
u = Unit("s/m")
assert u.get_mks_equivalent() == Unit("s/m")
assert u.get_mks_equivalent() != Unit("ohm")
assert u.get_cgs_equivalent() == Unit("s/cm")
u = Unit("kC")
assert u.get_cgs_equivalent() == Unit("kesu")
assert u.get_cgs_equivalent().get_mks_equivalent() == u
UnitSystem(ureg.unit_system_id, "code_length", "kg", "s", registry=ureg)
u = Unit("cm", registry=ureg)
ue = u.get_base_equivalent("code")
assert str(ue) == "code_length"
assert ue.base_value == 10
assert ue.dimensions is length
class FakeDataset(object):
unit_registry = ureg
ds = FakeDataset()
UnitSystem(ds, "code_length", "kg", "s", registry=ureg)
u = Unit("cm", registry=ureg)
ue = u.get_base_equivalent(ds)
assert str(ue) == "code_length"
assert ue.base_value == 10
assert ue.dimensions is length
with pytest.raises(UnitParseError):
Unit("code_length")
def _parse_parameter_file(self):
self._prefix = \
self.filename[:self.filename.rfind(self._suffix)]
fh = h5py.File(self.filename, "r")
for attr in ["hubble_constant",
"omega_matter",
"omega_lambda"]:
setattr(self, attr, fh.attrs[attr])
if "unit_registry_json" in fh.attrs:
self.unit_registry = \
UnitRegistry.from_json(
parse_h5_attr(fh, "unit_registry_json"))
self.box_size = _hdf5_yt_attr(
fh, "box_size", unit_registry=self.unit_registry)
self.field_info.update(
json.loads(parse_h5_attr(fh, "field_info")))
self.field_list = list(self.field_info.keys())
self._size = fh.attrs["total_trees"]
fh.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)
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 = "%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,
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 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, unyt_array):
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, unyt_array):
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.warning("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
# Domain dimensions for Gadget datasets are always 2x2x2 for octree
self.domain_dimensions = np.array([2,2,2])
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 v not in self.parameters:
raise MissingParameter(self.parameter_filename, v)
setattr(self, a, self.parameters[v])
self.domain_left_edge = np.array([0., 0., 0.])
self.domain_right_edge = np.array([1., 1., 1.]) * self.parameters['BoxSize']
else:
self.cosmological_simulation = 0
self.omega_lambda = self.omega_matter = \
self.hubble_constant = 0.0
def _find_data_dir(self):
"""
Find proper location for datasets. First look where parameter file
points, but if this doesn't exist then default to the current
directory.
"""
if self.parameters["OutputDir"].startswith("/"):
data_dir = self.parameters["OutputDir"]
else:
data_dir = os.path.join(self.directory,
self.parameters["OutputDir"])
if not os.path.exists(data_dir):
mylog.info("OutputDir not found at %s, instead using %s." %
(data_dir, self.directory))
data_dir = self.directory
self.data_dir = data_dir
def _snapshot_format(self, index=None):
"""
The snapshot filename for a given index. Modify this for different
naming conventions.
"""
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(self.data_dir, filename)
def _get_all_outputs(self, find_outputs=False):
"""
Get all potential datasets and combine into a time-sorted list.
"""
# Find the data directory where the outputs are
self._find_data_dir()
# 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
open(os.path.join(self.data_dir,
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 "TimeMax" not 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()
def test_registry_json():
reg = UnitRegistry()
json_reg = reg.to_json()
unserialized_reg = UnitRegistry.from_json(json_reg)
assert_equal(reg.lut, unserialized_reg.lut)
def test_keys():
ureg = UnitRegistry()
assert sorted(ureg.keys()) == sorted(ureg.lut.keys())
def __init__(self, filename, hubble_constant=1.0):
self.unit_registry = UnitRegistry()
self.hubble_constant = hubble_constant
super(AHFArbor, self).__init__(filename)
def test_add_modify_error():
from unyt import m
ureg = UnitRegistry()
with pytest.raises(UnitParseError):
ureg.add('tayne', 1, length)
with pytest.raises(UnitParseError):
ureg.add('tayne', 'blah', length)
with pytest.raises(UnitParseError):
ureg.add('tayne', 1.0, length, offset=1)
with pytest.raises(UnitParseError):
ureg.add('tayne', 1.0, length, offset='blah')
ureg.add('tayne', 1.1, length)
with pytest.raises(SymbolNotFoundError):
ureg.remove('tayn')
with pytest.raises(SymbolNotFoundError):
ureg.modify('tayn', 1.2)
ureg.modify('tayne', 1.0 * m)
assert ureg['tayne'][:3] == ureg['m'][:3]
