def _is_valid(cls, filename, *args, **kwargs):
"""Checks whether the supplied file can be read by this frontend."""
warn_h5py(filename)
try:
with h5py.File(filename, mode="r") as f:
attrs = list(f["/"].attrs.keys())
for i in opmd_required_attributes:
if i not in attrs:
return False
if (StrictVersion(f.attrs["openPMD"].decode())
not in ompd_known_versions):
return False
if f.attrs["iterationEncoding"].decode() == "fileBased":
return True
return False
except (OSError, ImportError):
return False
def _read_particle_coords(self, chunks, ptf):
# This will read chunks and yield the results.
chunks = list(chunks)
data_files = set()
for chunk in chunks:
for obj in chunk.objs:
data_files.update(obj.data_files)
for data_file in sorted(data_files, key=lambda x: (x.filename, x.start)):
with h5py.File(data_file.filename, mode="r") as f:
for ptype in sorted(ptf):
pcount = data_file.total_particles[ptype]
if pcount == 0:
continue
x = _get_position_array(ptype, f, "px")
y = _get_position_array(ptype, f, "py")
z = (
np.zeros(x.size, dtype="float64")
+ self.ds.domain_left_edge[2].to("code_length").d
)
yield ptype, (x, y, z)
def _read_particle_positions(self, ptype, f=None):
"""
Read all particle positions in this file.
"""
if f is None:
close = True
f = h5py.File(self.filename, mode="r")
else:
close = False
pcount = self.header["num_halos"]
pos = np.empty((pcount, 3), dtype="float64")
for i, ax in enumerate("xyz"):
pos[:, i] = f[f"particle_position_{ax}"][()]
if close:
f.close()
return pos
def _read_particle_fields(self, chunks, ptf, selector):
# Now we have all the sizes, and we can allocate
chunks = list(chunks)
data_files = set([])
for chunk in chunks:
for obj in chunk.objs:
data_files.update(obj.data_files)
for data_file in sorted(data_files):
with h5py.File(data_file.filename, "r") as f:
for ptype, field_list in sorted(ptf.items()):
units = _get_position_array_units(ptype, f, "x")
x, y, z = \
(self.ds.arr(_get_position_array(ptype, f, ax), units)
for ax in "xyz")
mask = selector.select_points(x, y, z, 0.0)
del x, y, z
if mask is None: continue
for field in field_list:
data = f[ptype][field][mask].astype("float64")
yield (ptype, field), data
def _initialize_index(self, data_file, regions):
all_count = self._count_particles(data_file)
pcount = sum(all_count.values())
morton = np.empty(pcount, dtype="uint64")
mylog.debug(
"Initializing index % 5i (% 7i particles)", data_file.file_id, pcount
)
ind = 0
with h5py.File(data_file.filename, mode="r") as f:
for ptype in all_count:
if ptype not in f or all_count[ptype] == 0:
continue
pos = np.empty((all_count[ptype], 3), dtype="float64")
units = _get_position_array_units(ptype, f, "x")
if ptype == "grid":
dx = f["grid"]["dx"][()].min()
dx = self.ds.quan(dx, parse_h5_attr(f["grid"]["dx"], "units")).to(
"code_length"
)
else:
dx = 2.0 * np.finfo(f[ptype]["particle_position_x"].dtype).eps
dx = self.ds.quan(dx, units).to("code_length")
pos[:, 0] = _get_position_array(ptype, f, "x")
pos[:, 1] = _get_position_array(ptype, f, "y")
pos[:, 2] = _get_position_array(ptype, f, "z")
pos = self.ds.arr(pos, units).to("code_length")
dle = self.ds.domain_left_edge.to("code_length")
dre = self.ds.domain_right_edge.to("code_length")
# These are 32 bit numbers, so we give a little lee-way.
# Otherwise, for big sets of particles, we often will bump into the
# domain edges. This helps alleviate that.
np.clip(pos, dle + dx, dre - dx, pos)
if np.any(pos.min(axis=0) < dle) or np.any(pos.max(axis=0) > dre):
raise YTDomainOverflow(pos.min(axis=0), pos.max(axis=0), dle, dre)
regions.add_data_file(pos, data_file.file_id)
morton[ind : ind + pos.shape[0]] = compute_morton(
pos[:, 0], pos[:, 1], pos[:, 2], dle, dre
)
ind += pos.shape[0]
return morton
def _read_particle_fields(self, chunks, ptf, selector):
# Now we have all the sizes, and we can allocate
chunks = list(chunks)
data_files = set([])
for chunk in chunks:
for obj in chunk.objs:
data_files.update(obj.data_files)
for data_file in sorted(data_files, key=lambda f: f.filename):
with h5py.File(data_file.filename, "r") as f:
for ptype, field_list in sorted(ptf.items()):
pcount = data_file.total_particles[ptype]
if pcount == 0: continue
coords = f[ptype]["CenterOfMass"][()].astype("float64")
coords = np.resize(coords, (pcount, 3))
x = coords[:, 0]
y = coords[:, 1]
z = coords[:, 2]
mask = selector.select_points(x, y, z, 0.0)
del x, y, z
if mask is None: continue
for field in field_list:
if field in self.offset_fields:
field_data = \
self._read_offset_particle_field(field, data_file, f)
else:
if field == "particle_identifier":
field_data = \
np.arange(data_file.total_particles[ptype]) + \
data_file.index_start[ptype]
elif field in f[ptype]:
field_data = f[ptype][field][()].astype("float64")
else:
fname = field[:field.rfind("_")]
field_data = f[ptype][fname][()].astype("float64")
my_div = field_data.size / pcount
if my_div > 1:
field_data = np.resize(field_data, (int(pcount), int(my_div)))
findex = int(field[field.rfind("_") + 1:])
field_data = field_data[:, findex]
data = field_data[mask]
yield (ptype, field), data
def _parse_parameter_file(self):
handle = h5py.File(self.parameter_filename, mode="r")
hvals = {}
hvals.update((str(k), v) for k, v in handle["/Header"].attrs.items())
hvals["NumFiles"] = hvals["NumFilesPerSnapshot"]
hvals["Massarr"] = hvals["MassTable"]
self.dimensionality = 3
self.refine_by = 2
# Set standard values
self.current_time = self.quan(hvals["Time_GYR"], "Gyr")
self.domain_left_edge = np.zeros(3, "float64")
self.domain_right_edge = np.ones(3, "float64") * hvals["BoxSize"]
self.domain_dimensions = np.ones(3, "int32")
self.cosmological_simulation = 1
self.periodicity = (True, True, True)
self.current_redshift = hvals["Redshift"]
self.omega_lambda = hvals["OmegaLambda"]
self.omega_matter = hvals["Omega0"]
self.hubble_constant = hvals["HubbleParam"]
self.parameters = hvals
prefix = os.path.abspath(
os.path.join(
os.path.dirname(self.parameter_filename),
os.path.basename(self.parameter_filename).split(".", 1)[0],
)
)
suffix = self.parameter_filename.rsplit(".", 1)[-1]
self.filename_template = f"{prefix}.%(num)i.{suffix}"
self.file_count = len(glob.glob(prefix + "*" + self._suffix))
if self.file_count == 0:
raise YTException(message="No data files found.", ds=self)
self.particle_types = ("FOF", "SUBFIND")
self.particle_types_raw = ("FOF", "SUBFIND")
# To avoid having to open files twice
self._unit_base = {}
self._unit_base.update((str(k), v) for k, v in handle["/Units"].attrs.items())
handle.close()
def _initialize_data_storage(self):
if not ytcfg.getboolean('yt', 'serialize'): return
fn = self.ds.storage_filename
if fn is None:
if os.path.isfile(
os.path.join(self.directory,
"%s.yt" % self.ds.unique_identifier)):
fn = os.path.join(self.directory,
"%s.yt" % self.ds.unique_identifier)
else:
fn = os.path.join(self.directory,
"%s.yt" % self.dataset.basename)
dir_to_check = os.path.dirname(fn)
if dir_to_check == '':
dir_to_check = '.'
# We have four options:
# Writeable, does not exist : create, open as append
# Writeable, does exist : open as append
# Not writeable, does not exist : do not attempt to open
# Not writeable, does exist : open as read-only
exists = os.path.isfile(fn)
if not exists:
writeable = os.access(dir_to_check, os.W_OK)
else:
writeable = os.access(fn, os.W_OK)
writeable = writeable and not ytcfg.getboolean('yt', 'onlydeserialize')
# We now have our conditional stuff
self.comm.barrier()
if not writeable and not exists: return
if writeable:
try:
if not exists: self.__create_data_file(fn)
self._data_mode = 'a'
except IOError:
self._data_mode = None
return
else:
self._data_mode = 'r'
self.__data_filename = fn
self._data_file = h5py.File(fn, self._data_mode)
def _read_particle_data_file(self, sub_file, ptf, selector=None):
# note: this frontend uses the variable name and terminology sub_file.
# other frontends use data_file with the understanding that it may
# actually be a sub_file, hence the super()._read_datafile is called
# ._read_datafile instead of ._read_subfile
return_data = {}
si, ei = sub_file.start, sub_file.end
f = h5py.File(sub_file.filename, mode="r")
for ptype, field_list in sorted(ptf.items()):
if sub_file.total_particles[ptype] == 0:
continue
g = f[f"/{ptype}"]
# this should load as float64
coords = g["Coordinates"][si:ei]
if ptype == "PartType0":
hsmls = self._get_smoothing_length(sub_file)
else:
hsmls = 0.0
if selector:
mask = selector.select_points(coords[:, 0], coords[:, 1],
coords[:, 2], hsmls)
del coords
if selector and mask is None:
continue
for field in field_list:
if field in ("Mass", "Masses"):
data = g[self.ds._particle_mass_name][si:ei]
else:
data = g[field][si:ei]
if selector:
data = data[mask, ...]
data.astype("float64", copy=False)
return_data[(ptype, field)] = data
f.close()
return return_data
def test_write_gdf():
"""Main test suite for write_gdf"""
tmpdir = tempfile.mkdtemp()
tmpfile = os.path.join(tmpdir, "test_gdf.h5")
try:
test_ds = fake_random_ds(64)
write_to_gdf(
test_ds, tmpfile, data_author=TEST_AUTHOR, data_comment=TEST_COMMENT
)
del test_ds
assert isinstance(load(tmpfile), GDFDataset)
h5f = h5py.File(tmpfile, mode="r")
gdf = h5f["gridded_data_format"].attrs
assert_equal(gdf["data_author"], TEST_AUTHOR)
assert_equal(gdf["data_comment"], TEST_COMMENT)
h5f.close()
finally:
shutil.rmtree(tmpdir)
def store_kd_bricks(self, fn=None):
if not self._initialized:
self.initialize_source()
if fn is None:
fn = '%s_kd_bricks.h5'%self.ds
if self.comm.rank != 0:
self.comm.recv_array(self.comm.rank-1, tag=self.comm.rank-1)
f = h5py.File(fn,'w')
for node in self.tree.depth_traverse():
i = node.node_id
if node.data is not None:
for fi,field in enumerate(self.fields):
try:
f.create_dataset("/brick_%s_%s" % (hex(i),field),
data = node.data.my_data[fi].astype('float64'))
except:
pass
f.close()
del f
if self.comm.rank != (self.comm.size-1):
self.comm.send_array([0],self.comm.rank+1, tag=self.comm.rank)
def _is_valid(self, *args, **kwargs):
warn_h5py(args[0])
try:
with h5.File(args[0], "r") as f:
attrs = list(f["/"].attrs.keys())
for i in opmd_required_attributes:
if i not in attrs:
return False
if (
StrictVersion(f.attrs["openPMD"].decode())
not in ompd_known_versions
):
return False
if f.attrs["iterationEncoding"].decode() == "groupBased":
return True
return False
except (OSError, ImportError):
return False
def _read_particle_fields(self, chunks, ptf, selector):
# Now we have all the sizes, and we can allocate
chunks = list(chunks)
data_files = set([])
for chunk in chunks:
for obj in chunk.objs:
data_files.update(obj.data_files)
for data_file in sorted(data_files):
all_count = self._count_particles(data_file)
with h5py.File(data_file.filename, "r") as f:
for ptype, field_list in sorted(ptf.items()):
x = _get_position_array(ptype, f, "px")
y = _get_position_array(ptype, f, "py")
z = np.zeros(all_count[ptype], dtype="float64") + \
self.ds.domain_left_edge[2].in_cgs().d
mask = selector.select_points(x, y, z, 0.0)
del x, y, z
if mask is None: continue
for field in field_list:
data = f[ptype][field][mask].astype("float64")
yield (ptype, field), data
def store(self, storage):
if hasattr(self, "_ds_mrep"):
self._ds_mrep.store(storage)
metadata, (final_name, chunks) = self._generate_post()
metadata["obj_type"] = self.type
with h5.File(storage) as h5f:
dset = str(uuid4())[:8]
h5f.create_group(dset)
_serialize_to_h5(h5f[dset], metadata)
if len(chunks) > 0:
g = h5f[dset].create_group("chunks")
g.attrs["final_name"] = final_name
for fname, fdata in chunks:
if isinstance(fname, (tuple, list)):
fname = "*".join(fname)
if isinstance(fdata, (YTQuantity, YTArray)):
g.create_dataset(fname, data=fdata.d, compression="lzf")
g[fname].attrs["units"] = str(fdata.units)
else:
g.create_dataset(fname, data=fdata, compression="lzf")
def _read_particle_coords(self, chunks, ptf):
pn = "particle_position_%s"
chunks = list(chunks)
for chunk in chunks:
f = None
for g in chunk.objs:
if g.filename is None: continue
if f is None:
f = h5py.File(g.filename, "r")
if g.NumberOfParticles == 0:
continue
for ptype, field_list in sorted(ptf.items()):
units = parse_h5_attr(f[ptype][pn % "x"], "units")
x, y, z = \
(self.ds.arr(f[ptype][pn % ax][()].astype("float64"), units)
for ax in "xyz")
for field in field_list:
if np.asarray(f[ptype][field]).ndim > 1:
self._array_fields[field] = f[ptype][field].shape[1:]
yield ptype, (x, y, z)
if f: f.close()
def _initialize_index(self, data_file, regions):
if self.index_ptype == "all":
ptypes = self.ds.particle_types_raw
pcount = sum(data_file.total_particles.values())
else:
ptypes = [self.index_ptype]
pcount = data_file.total_particles[self.index_ptype]
morton = np.empty(pcount, dtype='uint64')
if pcount == 0: return morton
mylog.debug("Initializing index % 5i (% 7i particles)",
data_file.file_id, pcount)
ind = 0
with h5py.File(data_file.filename, "r") as f:
if not f.keys(): return None
dx = np.finfo(f["Group"]["GroupPos"].dtype).eps
dx = 2.0 * self.ds.quan(dx, "code_length")
for ptype in ptypes:
if data_file.total_particles[ptype] == 0: continue
pos = f[ptype]["%sPos" % ptype].value.astype("float64")
pos = np.resize(pos, (data_file.total_particles[ptype], 3))
pos = data_file.ds.arr(pos, "code_length")
# These are 32 bit numbers, so we give a little lee-way.
# Otherwise, for big sets of particles, we often will bump into the
# domain edges. This helps alleviate that.
np.clip(pos, self.ds.domain_left_edge + dx,
self.ds.domain_right_edge - dx, pos)
if np.any(pos.min(axis=0) < self.ds.domain_left_edge) or \
np.any(pos.max(axis=0) > self.ds.domain_right_edge):
raise YTDomainOverflow(pos.min(axis=0), pos.max(axis=0),
self.ds.domain_left_edge,
self.ds.domain_right_edge)
regions.add_data_file(pos, data_file.file_id)
morton[ind:ind + pos.shape[0]] = compute_morton(
pos[:, 0], pos[:, 1], pos[:, 2],
data_file.ds.domain_left_edge,
data_file.ds.domain_right_edge)
ind += pos.shape[0]
return morton
def _is_valid(self, *args, **kwargs):
need_groups = ["Constants", "Header", "Parameters", "Units"]
veto_groups = [
"SUBFIND",
"FOF",
"PartType0/ChemistryAbundances",
"PartType0/ChemicalAbundances",
"RuntimePars",
"HashTable",
]
valid = True
valid_fname = args[0]
# If passed arg is a directory, look for the .0 file in that dir
if os.path.isdir(args[0]):
valid_files = []
for f in os.listdir(args[0]):
fname = os.path.join(args[0], f)
fext = os.path.splitext(fname)[-1]
if ((".0" in f) and (fext not in {".ewah", ".kdtree"})
and os.path.isfile(fname)):
valid_files.append(fname)
if len(valid_files) == 0:
valid = False
elif len(valid_files) > 1:
valid = False
else:
valid_fname = valid_files[0]
try:
fileh = h5py.File(valid_fname, mode="r")
for ng in need_groups:
if ng not in fileh["/"]:
valid = False
for vg in veto_groups:
if vg in fileh["/"]:
valid = False
fileh.close()
except Exception:
valid = False
pass
return valid
def _read_particle_fields(self, chunks, ptf, selector):
# Now we have all the sizes, and we can allocate
data_files = set([])
for chunk in chunks:
for obj in chunk.objs:
data_files.update(obj.data_files)
for data_file in sorted(data_files, key=lambda x: x.filename):
f = h5py.File(data_file.filename, "r")
for ptype, field_list in sorted(ptf.items()):
if data_file.total_particles[ptype] == 0:
continue
g = f["/%s" % ptype]
coords = g["Coordinates"][:].astype("float64")
mask = selector.select_points(
coords[:, 0], coords[:, 1], coords[:, 2], 0.0)
del coords
if mask is None:
continue
for field in field_list:
if field in ("Mass", "Masses") and \
ptype not in self.var_mass:
data = np.empty(mask.sum(), dtype="float64")
ind = self._known_ptypes.index(ptype)
data[:] = self.ds["Massarr"][ind]
elif field in self._element_names:
rfield = 'ElementAbundance/' + field
data = g[rfield][:][mask, ...]
elif field.startswith("Metallicity_"):
col = int(field.rsplit("_", 1)[-1])
data = g["Metallicity"][:, col][mask]
elif field.startswith("Chemistry_"):
col = int(field.rsplit("_", 1)[-1])
data = g["ChemistryAbundances"][:, col][mask]
else:
data = g[field][:][mask, ...]
yield (ptype, field), data
f.close()
def _read_particle_fields(self, chunks, ptf, selector):
chunks = list(chunks)
for chunk in chunks: # These should be organized by grid filename
f = None
for g in chunk.objs:
if g.filename is None:
continue
if f is None:
# print("Opening (read) %s" % g.filename)
f = h5py.File(g.filename, mode="r")
nap = sum(g.NumberOfActiveParticles.values())
if g.NumberOfParticles == 0 and nap == 0:
continue
ds = f.get("/Grid%08i" % g.id)
for ptype, field_list in sorted(ptf.items()):
if ptype != "io":
if g.NumberOfActiveParticles[ptype] == 0:
continue
pds = ds.get("Particles/%s" % ptype)
else:
pds = ds
pn = _particle_position_names.get(ptype,
r"particle_position_%s")
x, y, z = (np.asarray(pds.get(pn % ax)[()], dtype="=f8")
for ax in "xyz")
if selector is None:
# This only ever happens if the call is made from
# _read_particle_coords.
yield ptype, (x, y, z)
continue
mask = selector.select_points(x, y, z, 0.0)
if mask is None:
continue
for field in field_list:
data = np.asarray(pds.get(field)[()], "=f8")
if field in _convert_mass:
data *= g.dds.prod(dtype="f8")
yield (ptype, field), data[mask]
if f:
f.close()
def _read_particle_fields(self, chunks, ptf, selector):
pn = "particle_position_%s"
chunks = list(chunks)
for chunk in chunks: # These should be organized by grid filename
f = None
for g in chunk.objs:
if g.filename is None: continue
if f is None:
f = h5py.File(g.filename, "r")
if g.NumberOfParticles == 0:
continue
for ptype, field_list in sorted(ptf.items()):
units = parse_h5_attr(f[ptype][pn % "x"], "units")
x, y, z = \
(self.ds.arr(f[ptype][pn % ax][()].astype("float64"), units)
for ax in "xyz")
mask = selector.select_points(x, y, z, 0.0)
if mask is None: continue
for field in field_list:
data = np.asarray(f[ptype][field][()], "=f8")
yield (ptype, field), data[mask]
if f: f.close()
def _read_member_fields(self, dobj, member_fields):
all_data = defaultdict(
lambda: np.empty(dobj.particle_number, dtype=np.float64))
if not member_fields:
return all_data
field_start = 0
for i, data_file in enumerate(dobj.field_data_files):
start_index = dobj.field_data_start[i]
end_index = dobj.field_data_end[i]
pcount = end_index - start_index
if pcount == 0:
continue
field_end = field_start + end_index - start_index
with h5py.File(data_file.filename, mode="r") as f:
for ptype, field_list in sorted(member_fields.items()):
for field in field_list:
field_data = all_data[(ptype, field)]
my_data = f["particles"][field][
start_index:end_index].astype("float64")
field_data[field_start:field_end] = my_data
field_start = field_end
return all_data
def _read_field_names(self, grid):
if grid.filename is None:
return []
f = h5py.File(grid.filename, mode="r")
try:
group = f["/Grid%08i" % grid.id]
except KeyError:
group = f
fields = []
dtypes = set()
add_io = "io" in grid.ds.particle_types
add_dm = "DarkMatter" in grid.ds.particle_types
for name, v in group.items():
# NOTE: This won't work with 1D datasets or references.
# For all versions of Enzo I know about, we can assume all floats
# are of the same size. So, let's grab one.
if not hasattr(v, "shape") or v.dtype == "O":
continue
elif len(v.dims) == 1:
if grid.ds.dimensionality == 1:
fields.append(("enzo", str(name)))
elif add_io:
fields.append(("io", str(name)))
elif add_dm:
fields.append(("DarkMatter", str(name)))
else:
fields.append(("enzo", str(name)))
dtypes.add(v.dtype)
if len(dtypes) == 1:
# Now, if everything we saw was the same dtype, we can go ahead and
# set it here. We do this because it is a HUGE savings for 32 bit
# floats, since our numpy copying/casting is way faster than
# h5py's, for some reason I don't understand. This does *not* need
# to be correct -- it will get fixed later -- it just needs to be
# okay for now.
self._field_dtype = list(dtypes)[0]
f.close()
return fields
def light_cone_projection(parameter_file, simulation_type):
lc = LightCone(parameter_file,
simulation_type,
0.,
0.1,
observer_redshift=0.0,
time_data=False)
lc.calculate_light_cone_solution(seed=123456789,
filename="LC/solution.txt")
lc.project_light_cone((600.0, "arcmin"), (60.0, "arcsec"),
"density",
weight_field=None,
save_stack=True)
fh = h5py.File("LC/LightCone.h5")
data = fh["density_None"].value
units = fh["density_None"].attrs["units"]
assert units == "g/cm**2"
fh.close()
mean = data.mean()
mi = data[data.nonzero()].min()
ma = data.max()
return np.array([mean, mi, ma])
def _read_particle_fields(self, chunks, ptf, selector):
# Now we have all the sizes, and we can allocate
chunks = list(chunks)
data_files = set([])
# Only support halo reading for now.
assert (len(ptf) == 1)
assert (list(ptf.keys())[0] == "halos")
for chunk in chunks:
for obj in chunk.objs:
data_files.update(obj.data_files)
for data_file in sorted(data_files):
with h5py.File(data_file.filename, "r") as f:
for ptype, field_list in sorted(ptf.items()):
x = f['particle_position_x'].value.astype("float64")
y = f['particle_position_y'].value.astype("float64")
z = f['particle_position_z'].value.astype("float64")
mask = selector.select_points(x, y, z, 0.0)
del x, y, z
if mask is None: continue
for field in field_list:
data = f[field][mask].astype("float64")
yield (ptype, field), data
def _parse_parameter_file(self):
with h5py.File(self.parameter_filename, mode="r") as f:
self.parameters = {
str(field): val
for field, val in f["Header"].attrs.items()
}
self.dimensionality = 3
self.refine_by = 2
# Set standard values
self.domain_left_edge = np.zeros(3, "float64")
self.domain_right_edge = np.ones(
3, "float64") * self.parameters["BoxSize"]
self.domain_dimensions = np.ones(3, "int32")
self.cosmological_simulation = 1
self._periodicity = (True, True, True)
self.current_redshift = self.parameters["Redshift"]
self.omega_lambda = self.parameters["OmegaLambda"]
self.omega_matter = self.parameters["Omega0"]
self.hubble_constant = self.parameters["HubbleParam"]
cosmology = Cosmology(
hubble_constant=self.hubble_constant,
omega_matter=self.omega_matter,
omega_lambda=self.omega_lambda,
)
self.current_time = cosmology.t_from_z(self.current_redshift)
prefix = os.path.abspath(
os.path.join(
os.path.dirname(self.parameter_filename),
os.path.basename(self.parameter_filename).split(".", 1)[0],
))
suffix = self.parameter_filename.rsplit(".", 1)[-1]
self.filename_template = f"{prefix}.%(num)i.{suffix}"
self.file_count = self.parameters["NumFiles"]
self.particle_types = ("Group", "Subhalo")
self.particle_types_raw = ("Group", "Subhalo")
def _initialize_index(self, data_file, regions):
all_count = self._count_particles(data_file)
pcount = sum(all_count.values())
morton = np.empty(pcount, dtype='uint64')
mylog.debug("Initializing index % 5i (% 7i particles)",
data_file.file_id, pcount)
ind = 0
with h5py.File(data_file.filename, "r") as f:
for ptype in all_count:
if ptype not in f or all_count[ptype] == 0: continue
pos = np.empty((all_count[ptype], 3), dtype="float64")
pos = data_file.ds.arr(pos, "code_length")
if ptype == "grid":
dx = f["grid"]["pdx"].value.min()
else:
raise NotImplementedError
dx = self.ds.quan(dx, "code_length")
pos[:, 0] = _get_position_array(ptype, f, "px")
pos[:, 1] = _get_position_array(ptype, f, "py")
pos[:,2] = np.zeros(all_count[ptype], dtype="float64") + \
self.ds.domain_left_edge[2].in_cgs().d
# These are 32 bit numbers, so we give a little lee-way.
# Otherwise, for big sets of particles, we often will bump into the
# domain edges. This helps alleviate that.
np.clip(pos, self.ds.domain_left_edge + dx,
self.ds.domain_right_edge - dx, pos)
if np.any(pos.min(axis=0) < self.ds.domain_left_edge) or \
np.any(pos.max(axis=0) > self.ds.domain_right_edge):
raise YTDomainOverflow(pos.min(axis=0), pos.max(axis=0),
self.ds.domain_left_edge,
self.ds.domain_right_edge)
regions.add_data_file(pos, data_file.file_id)
morton[ind:ind + pos.shape[0]] = compute_morton(
pos[:, 0], pos[:, 1], pos[:, 2],
data_file.ds.domain_left_edge,
data_file.ds.domain_right_edge)
ind += pos.shape[0]
return morton
def _save_raw_arrays(arrays, answer_file, func_name):
r"""
Saves the raw arrays produced from answer tests to a file.
The structure of `answer_file` is: each test function (e.g.,
test_toro1d[0-None-None-0]) forms a group. Within each group is a
hdf5 dataset named after the test (e.g., field_values). The value
stored in each dataset is the raw array corresponding to that
test and function.
Parameters
----------
arrays : dict
Keys are the test name (e.g. field_values) and the values are
the actual answer arrays produced by the test.
answer_file : str
The name of the file to save the answers to, in hdf5 format.
func_name : str
The name of the function (possibly augmented by pytest with
test parameters) that called the test functions
(e.g, test_toro1d).
"""
with h5py.File(answer_file, "a") as f:
grp = f.create_group(func_name)
for test_name, test_data in arrays.items():
# Some answer tests (e.g., grid_values, projection_values)
# return a dictionary, which cannot be handled by h5py
if isinstance(test_data, dict):
sub_grp = grp.create_group(test_name)
_parse_raw_answer_dict(test_data, sub_grp)
else:
# Some tests return None, which hdf5 can't handle, and there is
# no proxy, so we have to make one ourselves. Using -1
if test_data is None:
test_data = -1
grp.create_dataset(test_name, data=test_data)
def run(self):
# Set up in a temp dir
tmpdir = tempfile.mkdtemp()
curdir = os.getcwd()
os.chdir(tmpdir)
lc = LightCone(
self.parameter_file, self.simulation_type, 0., 0.1,
observer_redshift=0.0, time_data=False)
lc.calculate_light_cone_solution(
seed=123456789, filename="LC/solution.txt")
lc.project_light_cone(
(600.0, "arcmin"), (60.0, "arcsec"), self.field,
weight_field=self.weight_field, save_stack=True)
dname = "%s_%s" % (self.field, self.weight_field)
fh = h5py.File("LC/LightCone.h5", mode="r")
data = fh[dname][()]
units = fh[dname].attrs["units"]
if self.weight_field is None:
punits = _funits[self.field] * _funits['length']
else:
punits = _funits[self.field] * _funits[self.weight_field] * \
_funits['length']
wunits = fh['weight_field_%s' % self.weight_field].attrs['units']
pwunits = _funits[self.weight_field] * _funits['length']
assert wunits == str(pwunits.units)
assert units == str(punits.units)
fh.close()
# clean up
os.chdir(curdir)
shutil.rmtree(tmpdir)
mean = data.mean()
mi = data[data.nonzero()].min()
ma = data.max()
return np.array([mean, mi, ma])
def _read_scalar_fields(self, dobj, scalar_fields):
all_data = {}
if not scalar_fields: return all_data
pcount = 1
with h5py.File(dobj.scalar_data_file.filename, "r") as f:
for ptype, field_list in sorted(scalar_fields.items()):
for field in field_list:
if field == "particle_identifier":
field_data = \
np.arange(dobj.scalar_data_file.total_particles[ptype]) + \
dobj.scalar_data_file.index_start[ptype]
elif field in f[ptype]:
field_data = f[ptype][field][()].astype("float64")
else:
fname = field[:field.rfind("_")]
field_data = f[ptype][fname][()].astype("float64")
my_div = field_data.size / pcount
if my_div > 1:
findex = int(field[field.rfind("_") + 1:])
field_data = field_data[:, findex]
data = np.array([field_data[dobj.scalar_index]])
all_data[(ptype, field)] = data
return all_data
def _detect_active_particle_fields(self):
ap_list = self.dataset["AppendActiveParticleType"]
_fields = dict((ap, []) for ap in ap_list)
fields = []
for ptype in self.dataset["AppendActiveParticleType"]:
select_grids = self.grid_active_particle_count[ptype].flat
if not np.any(select_grids):
current_ptypes = self.dataset.particle_types
new_ptypes = [p for p in current_ptypes if p != ptype]
self.dataset.particle_types = new_ptypes
self.dataset.particle_types_raw = new_ptypes
continue
gs = self.grids[select_grids > 0]
g = gs[0]
handle = h5py.File(g.filename, "r")
node = handle["/Grid%08i/Particles/" % g.id]
for ptype in (str(p) for p in node):
if ptype not in _fields: continue
for field in (str(f) for f in node[ptype]):
_fields[ptype].append(field)
fields += [(ptype, field) for field in _fields.pop(ptype)]
handle.close()
return set(fields)
