import numpy as np
from yt.funcs import mylog
from yt.utilities.io_handler import BaseParticleIOHandler
[docs]
class IOHandlerSDF(BaseParticleIOHandler):
_dataset_type = "sdf_particles"
@property
def _handle(self):
return self.ds.sdf_container
def _read_fluid_selection(self, chunks, selector, fields, size):
raise NotImplementedError
def _read_particle_coords(self, chunks, ptf):
assert len(ptf) == 1
assert ptf.keys()[0] == "dark_matter"
data_files = self._get_data_files(chunks)
assert len(data_files) == 1
for _data_file in sorted(data_files, key=lambda x: (x.filename, x.start)):
yield (
"dark_matter",
(
self._handle["x"],
self._handle["y"],
self._handle["z"],
),
0.0,
)
def _read_particle_fields(self, chunks, ptf, selector):
assert len(ptf) == 1
assert ptf.keys()[0] == "dark_matter"
data_files = self._get_data_files(chunks)
assert len(data_files) == 1
for _data_file in sorted(data_files, key=lambda x: (x.filename, x.start)):
for ptype, field_list in sorted(ptf.items()):
x = self._handle["x"]
y = self._handle["y"]
z = self._handle["z"]
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 == "mass":
data = np.ones(mask.sum(), dtype="float64")
data *= self.ds.parameters["particle_mass"]
else:
data = self._handle[field][mask]
yield (ptype, field), data
def _identify_fields(self, data_file):
fields = [("dark_matter", v) for v in self._handle.keys()]
fields.append(("dark_matter", "mass"))
return fields, {}
def _count_particles(self, data_file):
pcount = self._handle["x"].size
if pcount > 1e9:
mylog.warning(
"About to load %i particles into memory. "
"You may want to consider a midx-enabled load",
pcount,
)
return {"dark_matter": pcount}
[docs]
class IOHandlerHTTPSDF(IOHandlerSDF):
_dataset_type = "http_sdf_particles"
def _read_particle_coords(self, chunks, ptf):
chunks = list(chunks)
data_files = set()
assert len(ptf) == 1
assert ptf.keys()[0] == "dark_matter"
for chunk in chunks:
for obj in chunk.objs:
data_files.update(obj.data_files)
assert len(data_files) == 1
for _data_file in data_files:
pcount = self._handle["x"].size
yield (
"dark_matter",
(
self._handle["x"][:pcount],
self._handle["y"][:pcount],
self._handle["z"][:pcount],
),
0.0,
)
def _read_particle_fields(self, chunks, ptf, selector):
chunks = list(chunks)
data_files = set()
assert len(ptf) == 1
assert ptf.keys()[0] == "dark_matter"
for chunk in chunks:
for obj in chunk.objs:
data_files.update(obj.data_files)
assert len(data_files) == 1
for _data_file in data_files:
pcount = self._handle["x"].size
for ptype, field_list in sorted(ptf.items()):
x = self._handle["x"][:pcount]
y = self._handle["y"][:pcount]
z = self._handle["z"][:pcount]
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 == "mass":
if self.ds.field_info._mass_field is None:
pm = 1.0
if "particle_mass" in self.ds.parameters:
pm = self.ds.parameters["particle_mass"]
else:
raise RuntimeError
data = pm * np.ones(mask.sum(), dtype="float64")
else:
data = self._handle[self.ds.field_info._mass_field][:][mask]
else:
data = self._handle[field][:][mask]
yield (ptype, field), data
def _count_particles(self, data_file):
return {"dark_matter": self._handle["x"].http_array.shape}
[docs]
class IOHandlerSIndexSDF(IOHandlerSDF):
_dataset_type = "midx_sdf_particles"
def _read_particle_coords(self, chunks, ptf):
dle = self.ds.domain_left_edge.in_units("code_length").d
dre = self.ds.domain_right_edge.in_units("code_length").d
for dd in self.ds.midx.iter_bbox_data(dle, dre, ["x", "y", "z"]):
yield "dark_matter", (dd["x"], dd["y"], dd["z"]), 0.0
def _read_particle_fields(self, chunks, ptf, selector):
dle = self.ds.domain_left_edge.in_units("code_length").d
dre = self.ds.domain_right_edge.in_units("code_length").d
required_fields = []
for field_list in sorted(ptf.values()):
for field in field_list:
if field == "mass":
continue
required_fields.append(field)
for dd in self.ds.midx.iter_bbox_data(dle, dre, required_fields):
for ptype, field_list in sorted(ptf.items()):
x = dd["x"]
y = dd["y"]
z = dd["z"]
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 == "mass":
data = np.ones(mask.sum(), dtype="float64")
data *= self.ds.parameters["particle_mass"]
else:
data = dd[field][mask]
yield (ptype, field), data
def _count_particles(self, data_file):
dle = self.ds.domain_left_edge.in_units("code_length").d
dre = self.ds.domain_right_edge.in_units("code_length").d
pcount_estimate = self.ds.midx.get_nparticles_bbox(dle, dre)
if pcount_estimate > 1e9:
mylog.warning(
"Filtering %i particles to find total. "
"You may want to reconsider your bounding box.",
pcount_estimate,
)
pcount = 0
for dd in self.ds.midx.iter_bbox_data(dle, dre, ["x"]):
pcount += dd["x"].size
return {"dark_matter": pcount}
def _identify_fields(self, data_file):
fields = [("dark_matter", v) for v in self._handle.keys()]
fields.append(("dark_matter", "mass"))
return fields, {}
[docs]
class IOHandlerSIndexHTTPSDF(IOHandlerSIndexSDF):
_dataset_type = "midx_http_sdf_particles"