class FakeDS:
domain_left_edge = None
domain_right_edge = None
domain_width = None
unit_registry = UnitRegistry()
unit_registry.add('code_length', 1.0, dimensions.length)
periodicity = (False, False, False)
def _set_units(self):
self.unit_registry = UnitRegistry()
self.time_unit = self.quan(1.0, "s")
if self.cosmological_simulation:
# Instantiate Cosmology object for units and time conversions.
self.cosmology = \
Cosmology(hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda,
unit_registry=self.unit_registry)
self.unit_registry.modify("h", self.hubble_constant)
# Comoving lengths
for my_unit in ["m", "pc", "AU", "au"]:
new_unit = "%scm" % my_unit
# technically not true, but should be ok
self.unit_registry.add(new_unit,
self.unit_registry.lut[my_unit][0],
dimensions.length,
"\\rm{%s}/(1+z)" % my_unit)
self.length_unit = self.quan(self.unit_base["UnitLength_in_cm"],
"cmcm / h",
registry=self.unit_registry)
self.box_size *= self.length_unit.in_units("Mpccm / h")
else:
# Read time from file for non-cosmological sim
self.time_unit = self.quan(
self.unit_base["UnitLength_in_cm"]/ \
self.unit_base["UnitVelocity_in_cm_per_s"], "s")
self.unit_registry.add("code_time", 1.0, dimensions.time)
self.unit_registry.modify("code_time", self.time_unit)
# Length
self.length_unit = self.quan(self.unit_base["UnitLength_in_cm"],
"cm")
self.unit_registry.add("code_length", 1.0, dimensions.length)
self.unit_registry.modify("code_length", self.length_unit)
def from_hdf5(cls, filename, dataset_name=None):
r"""Attempts read in and convert a dataset in an hdf5 file into a YTArray.
Parameters
----------
filename: string
The filename to of the hdf5 file.
dataset_name: string
The name of the dataset to read from. If the dataset has a units
attribute, attempt to infer units as well.
"""
import h5py
from yt.extern.six.moves import cPickle as pickle
if dataset_name is None:
dataset_name = 'array_data'
f = h5py.File(filename)
dataset = f[dataset_name]
data = dataset[:]
units = dataset.attrs.get('units', '')
if 'unit_registry' in dataset.attrs.keys():
unit_lut = pickle.loads(dataset.attrs['unit_registry'].tostring())
else:
unit_lut = None
f.close()
registry = UnitRegistry(lut=unit_lut, add_default_symbols=False)
return cls(data, units, registry=registry)
def mpi_bcast(self, data, root=0):
# The second check below makes sure that we know how to communicate
# this type of array. Otherwise, we'll pickle it.
if isinstance(data, np.ndarray) and \
get_mpi_type(data.dtype) is not None:
if self.comm.rank == root:
if isinstance(data, YTArray):
info = (data.shape, data.dtype, str(data.units),
data.units.registry.lut)
else:
info = (data.shape, data.dtype)
else:
info = ()
info = self.comm.bcast(info, root=root)
if self.comm.rank != root:
if len(info) == 4:
registry = UnitRegistry(lut=info[3],
add_default_symbols=False)
data = YTArray(np.empty(info[0], dtype=info[1]),
info[2],
registry=registry)
else:
data = np.empty(info[0], dtype=info[1])
mpi_type = get_mpi_type(info[1])
self.comm.Bcast([data, mpi_type], root=root)
return data
else:
# Use pickled methods.
data = self.comm.bcast(data, root=root)
return data
def _set_units(self):
self.unit_registry = UnitRegistry()
self.unit_registry.lut["code_time"] = (1.0, dimensions.time)
if self.cosmological_simulation:
# Instantiate EnzoCosmology object for units and time conversions.
self.cosmology = \
EnzoCosmology(self.parameters['CosmologyHubbleConstantNow'],
self.parameters['CosmologyOmegaMatterNow'],
self.parameters['CosmologyOmegaLambdaNow'],
0.0, self.parameters['CosmologyInitialRedshift'],
unit_registry=self.unit_registry)
self.time_unit = self.cosmology.time_unit.in_units("s")
self.unit_registry.modify("h", self.hubble_constant)
# Comoving lengths
for my_unit in ["m", "pc", "AU", "au"]:
new_unit = "%scm" % my_unit
# technically not true, but should be ok
self.unit_registry.add(new_unit,
self.unit_registry.lut[my_unit][0],
dimensions.length,
"\\rm{%s}/(1+z)" % my_unit)
self.length_unit = self.quan(self.box_size,
"Mpccm / h",
registry=self.unit_registry)
self.box_size = self.length_unit
else:
self.time_unit = self.quan(self.parameters["TimeUnits"], "s")
self.unit_registry.modify("code_time", self.time_unit)
def __init__(self, hubble_constant = 0.71,
omega_matter = 0.27,
omega_lambda = 0.73,
omega_curvature = 0.0,
unit_registry = None,
unit_system = "cgs",
use_dark_factor = False,
w_0 = -1.0,
w_a = 0.0):
self.omega_matter = float(omega_matter)
self.omega_lambda = float(omega_lambda)
self.omega_curvature = float(omega_curvature)
if unit_registry is None:
unit_registry = UnitRegistry()
unit_registry.modify("h", hubble_constant)
for my_unit in ["m", "pc", "AU", "au"]:
new_unit = "%scm" % my_unit
# technically not true, but distances here are actually comoving
unit_registry.add(new_unit, unit_registry.lut[my_unit][0],
dimensions.length, "\\rm{%s}/(1+z)" % my_unit)
self.unit_registry = unit_registry
self.hubble_constant = self.quan(hubble_constant, "100*km/s/Mpc")
self.unit_system = unit_system
# For non-standard dark energy. If false, use default cosmological constant
# This only affects the expansion_factor function.
self.use_dark_factor = use_dark_factor
self.w_0 = w_0
self.w_a = w_a
def unit_registry(self):
"""
Unit system registry.
"""
if self._unit_registry is None:
self._unit_registry = UnitRegistry()
return self._unit_registry
def fget(self):
ur = self._unit_registry
if ur is None:
ur = UnitRegistry()
# This will be updated when we add a volume source
ur.add("unitary", 1.0, length)
self._unit_registry = ur
return self._unit_registry
def __setstate__(self, state):
"""Pickle setstate method
This is called inside pickle.read() and restores the unit data from the
metadata extracted in __reduce__ and then serialized by pickle.
"""
super(YTArray, self).__setstate__(state[1:])
unit, lut = state[0]
registry = UnitRegistry(lut=lut, add_default_symbols=False)
self.units = Unit(unit, registry=registry)
def __deepcopy__(self, memodict=None):
if memodict is None:
memodict = {}
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 _create_unit_registry(self):
self.unit_registry = UnitRegistry()
import yt.units.dimensions as dimensions
self.unit_registry.add("code_length", 1.0, dimensions.length)
self.unit_registry.add("code_mass", 1.0, dimensions.mass)
self.unit_registry.add("code_time", 1.0, dimensions.time)
self.unit_registry.add("code_magnetic", 1.0, dimensions.magnetic_field)
self.unit_registry.add("code_temperature", 1.0, dimensions.temperature)
self.unit_registry.add("code_velocity", 1.0, dimensions.velocity)
self.unit_registry.add("code_metallicity", 1.0,
dimensions.dimensionless)
def recv_array(self, source, tag=0):
metadata = self.comm.recv(source=source, tag=tag)
dt, ne = metadata[:2]
if ne is None and dt is None:
return self.comm.recv(source=source, tag=tag)
arr = np.empty(ne, dtype=dt)
if len(metadata) == 4:
registry = UnitRegistry(lut=metadata[3], add_default_symbols=False)
arr = YTArray(arr, metadata[2], registry=registry)
tmp = arr.view(self.__tocast)
self.comm.Recv([tmp, MPI.CHAR], source=source, tag=tag)
return arr
def __init__(self, filename):
"""
Initialize an Arbor given an input file.
"""
self.filename = filename
self.basename = os.path.basename(filename)
self.unit_registry = UnitRegistry()
self._parse_parameter_file()
self._set_units()
self._root_field_data = FieldContainer(self)
self._setup_fields()
self._set_default_selector()
self._node_io = self._tree_field_io_class(self)
self._root_io = self._root_field_io_class(self)
def _set_new_unit_registry(self, input_registry):
self.unit_registry = UnitRegistry(add_default_symbols=False,
lut=input_registry.lut)
# Validate that the new unit registry makes sense
current_scaling = self.unit_registry['unitary'][0]
if current_scaling != input_registry['unitary'][0]:
for source in self.sources.items():
data_source = getattr(source, 'data_source', None)
if data_source is None:
continue
scaling = data_source.ds.unit_registry['unitary'][0]
if scaling != current_scaling:
raise NotImplementedError(
"Simultaneously rendering data from datasets with "
"different units is not supported")
def __init__(self, hubble_constant = 0.71,
omega_matter = 0.27,
omega_lambda = 0.73,
omega_curvature = 0.0,
unit_registry = None):
self.omega_matter = omega_matter
self.omega_lambda = omega_lambda
self.omega_curvature = omega_curvature
if unit_registry is None:
unit_registry = UnitRegistry()
unit_registry.modify("h", hubble_constant)
for my_unit in ["m", "pc", "AU", "au"]:
new_unit = "%scm" % my_unit
# technically not true, but distances here are actually comoving
unit_registry.add(new_unit, unit_registry.lut[my_unit][0],
dimensions.length, "\\rm{%s}/(1+z)" % my_unit)
self.unit_registry = unit_registry
self.hubble_constant = self.quan(hubble_constant, "100*km/s/Mpc")
def __init__(
self,
hubble_constant=0.71,
omega_matter=0.27,
omega_lambda=0.73,
omega_radiation=0.0,
omega_curvature=0.0,
unit_registry=None,
unit_system="cgs",
use_dark_factor=False,
w_0=-1.0,
w_a=0.0,
):
self.omega_matter = float(omega_matter)
self.omega_radiation = float(omega_radiation)
self.omega_lambda = float(omega_lambda)
self.omega_curvature = float(omega_curvature)
hubble_constant = float(hubble_constant)
if unit_registry is None:
unit_registry = UnitRegistry(unit_system=unit_system)
unit_registry.add("h", hubble_constant, dimensions.dimensionless,
r"h")
for my_unit in ["m", "pc", "AU", "au"]:
new_unit = f"{my_unit}cm"
my_u = Unit(my_unit, registry=unit_registry)
# technically not true, but distances here are actually comoving
unit_registry.add(
new_unit,
my_u.base_value,
dimensions.length,
"\\rm{%s}/(1+z)" % my_unit,
prefixable=True,
)
self.unit_registry = unit_registry
self.hubble_constant = self.quan(hubble_constant, "100*km/s/Mpc")
self.unit_system = unit_system
# For non-standard dark energy. If false, use default cosmological constant
# This only affects the expansion_factor function.
self.use_dark_factor = use_dark_factor
self.w_0 = w_0
self.w_a = w_a
equivalence_registry
from yt.units.unit_lookup_table import \
unit_prefixes, prefixable_units, latex_prefixes, \
default_base_units
from yt.units.unit_registry import \
UnitRegistry, \
UnitParseError
from yt.utilities.exceptions import YTUnitsNotReducible
import copy
import token
class InvalidUnitOperation(Exception):
pass
default_unit_registry = UnitRegistry()
sympy_one = sympify(1)
global_dict = {
'Symbol': Symbol,
'Integer': Integer,
'Float': Float,
'Rational': Rational,
'sqrt': sqrt
}
unit_system_registry = {}
def auto_positive_symbol(tokens, local_dict, global_dict):
"""
def __init__(self, filename, hubble_constant=1.0):
self.unit_registry = UnitRegistry()
self.hubble_constant = hubble_constant
super(AHFArbor, self).__init__(filename)
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 function exists to write out
self.make_simulated = True
# Calculate the simulated halo mass function
self.create_sim_hmf()
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)
import functools
import numpy as np
import optparse
from yt.units.unit_registry import UnitRegistry
from yt.units.yt_array import YTQuantity
from yt.utilities.physical_ratios import \
rho_crit_g_cm3_h2
refine_by = 2
unit_registry = UnitRegistry()
def quan(v, u):
return YTQuantity(v, u, registry=unit_registry)
def get_smallest_appropriate_unit(v,
quantity='distance',
return_quantity=False):
"""
Returns the largest whole unit smaller than the YTQuantity passed to
it as a string.
The quantity keyword can be equal to `distance` or `time`. In the
case of distance, the units are: 'Mpc', 'kpc', 'pc', 'au', 'rsun',
'km', etc. For time, the units are: 'Myr', 'kyr', 'yr', 'day', 'hr',
's', 'ms', etc.
If return_quantity is set to True, it finds the largest YTQuantity
object with a whole unit and a power of ten as the coefficient, and it
returns this YTQuantity.
from functools import partial
from yt.data_objects.static_output import Dataset
from yt.testing import assert_raises
from yt.units import YTQuantity
from yt.units.unit_registry import UnitRegistry
mock_quan = partial(YTQuantity, registry=UnitRegistry())
def test_schema_validation():
valid_schemas = [
{
"length_unit": 1.0
},
{
"length_unit": [1.0]
},
{
"length_unit": (1.0, )
},
{
"length_unit": int(1.0)
},
{
"length_unit": (1.0, "m")
},
{
"length_unit": [1.0, "m"]
},
