import glob
import os
import struct
import weakref
import numpy as np
import yt.utilities.fortran_utils as fpu
from yt.data_objects.index_subobjects.octree_subset import OctreeSubset
from yt.data_objects.static_output import Dataset, ParticleFile
from yt.data_objects.unions import ParticleUnion
from yt.frontends.art.definitions import (
amr_header_struct,
constants,
dmparticle_header_struct,
filename_pattern,
fluid_fields,
particle_fields,
particle_header_struct,
seek_extras,
)
from yt.frontends.art.fields import ARTFieldInfo
from yt.frontends.art.io import (
_read_art_level_info,
_read_child_level,
_read_root_level,
a2b,
b2t,
)
from yt.funcs import mylog, setdefaultattr
from yt.geometry.geometry_handler import Index, YTDataChunk
from yt.geometry.oct_container import ARTOctreeContainer
from yt.geometry.oct_geometry_handler import OctreeIndex
from yt.geometry.particle_geometry_handler import ParticleIndex
[docs]
class ARTIndex(OctreeIndex):
def __init__(self, ds, dataset_type="art"):
self.fluid_field_list = fluid_fields
self.dataset_type = dataset_type
self.dataset = weakref.proxy(ds)
self.index_filename = self.dataset.parameter_filename
self.directory = os.path.dirname(self.index_filename)
self.max_level = ds.max_level
self.float_type = np.float64
super().__init__(ds, dataset_type)
[docs]
def get_smallest_dx(self):
"""
Returns (in code units) the smallest cell size in the simulation.
"""
# Overloaded
ds = self.dataset
return (ds.domain_width / ds.domain_dimensions / (2**self.max_level)).min()
def _initialize_oct_handler(self):
"""
Just count the number of octs per domain and
allocate the requisite memory in the oct tree
"""
nv = len(self.fluid_field_list)
self.oct_handler = ARTOctreeContainer(
self.dataset.domain_dimensions / 2, # dd is # of root cells
self.dataset.domain_left_edge,
self.dataset.domain_right_edge,
1,
)
# The 1 here refers to domain_id == 1 always for ARTIO.
self.domains = [ARTDomainFile(self.dataset, nv, self.oct_handler, 1)]
self.octs_per_domain = [dom.level_count.sum() for dom in self.domains]
self.total_octs = sum(self.octs_per_domain)
mylog.debug("Allocating %s octs", self.total_octs)
self.oct_handler.allocate_domains(self.octs_per_domain)
domain = self.domains[0]
domain._read_amr_root(self.oct_handler)
domain._read_amr_level(self.oct_handler)
self.oct_handler.finalize()
def _detect_output_fields(self):
self.particle_field_list = list(particle_fields)
self.field_list = [("art", f) for f in fluid_fields]
# now generate all of the possible particle fields
for ptype in self.dataset.particle_types_raw:
for pfield in self.particle_field_list:
pfn = (ptype, pfield)
self.field_list.append(pfn)
def _identify_base_chunk(self, dobj):
"""
Take the passed in data source dobj, and use its embedded selector
to calculate the domain mask, build the reduced domain
subsets and oct counts. Attach this information to dobj.
"""
if getattr(dobj, "_chunk_info", None) is None:
# Get all octs within this oct handler
domains = [dom for dom in self.domains if dom.included(dobj.selector)]
base_region = getattr(dobj, "base_region", dobj)
if len(domains) > 1:
mylog.debug("Identified %s intersecting domains", len(domains))
subsets = [
ARTDomainSubset(base_region, domain, self.dataset) for domain in domains
]
dobj._chunk_info = subsets
dobj._current_chunk = list(self._chunk_all(dobj))[0]
def _chunk_all(self, dobj):
oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info)
# We pass the chunk both the current chunk and list of chunks,
# as well as the referring data source
yield YTDataChunk(dobj, "all", oobjs, None)
def _chunk_spatial(self, dobj, ngz, sort=None, preload_fields=None):
sobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info)
for og in sobjs:
if ngz > 0:
g = og.retrieve_ghost_zones(ngz, [], smoothed=True)
else:
g = og
yield YTDataChunk(dobj, "spatial", [g], None)
def _chunk_io(self, dobj, cache=True, local_only=False):
"""
Since subsets are calculated per domain,
i.e. per file, yield each domain at a time to
organize by IO. We will eventually chunk out NMSU ART
to be level-by-level.
"""
oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info)
for subset in oobjs:
yield YTDataChunk(dobj, "io", [subset], None, cache=cache)
[docs]
class ARTDataset(Dataset):
_index_class: type[Index] = ARTIndex
_field_info_class = ARTFieldInfo
def __init__(
self,
filename,
dataset_type="art",
fields=None,
storage_filename=None,
skip_particles=False,
skip_stars=False,
limit_level=None,
spread_age=True,
force_max_level=None,
file_particle_header=None,
file_particle_data=None,
file_particle_stars=None,
units_override=None,
unit_system="cgs",
default_species_fields=None,
):
self.fluid_types += ("art",)
if fields is None:
fields = fluid_fields
filename = os.path.abspath(filename)
self._fields_in_file = fields
self._file_amr = filename
self._file_particle_header = file_particle_header
self._file_particle_data = file_particle_data
self._file_particle_stars = file_particle_stars
self._find_files(filename)
self.skip_particles = skip_particles
self.skip_stars = skip_stars
self.limit_level = limit_level
self.max_level = limit_level
self.force_max_level = force_max_level
self.spread_age = spread_age
Dataset.__init__(
self,
filename,
dataset_type,
units_override=units_override,
unit_system=unit_system,
default_species_fields=default_species_fields,
)
self.storage_filename = storage_filename
def _find_files(self, file_amr):
"""
Given the AMR base filename, attempt to find the
particle header, star files, etc.
"""
base_prefix, base_suffix = filename_pattern["amr"]
numericstr = file_amr.rsplit("_", 1)[1].replace(base_suffix, "")
possibles = glob.glob(
os.path.join(os.path.dirname(os.path.abspath(file_amr)), "*")
)
for filetype, (prefix, suffix) in filename_pattern.items():
# if this attribute is already set skip it
if getattr(self, "_file_" + filetype, None) is not None:
continue
match = None
for possible in possibles:
if possible.endswith(numericstr + suffix):
if os.path.basename(possible).startswith(prefix):
match = possible
if match is not None:
mylog.info("discovered %s:%s", filetype, match)
setattr(self, "_file_" + filetype, match)
else:
setattr(self, "_file_" + filetype, None)
def __str__(self):
return self._file_amr.split("/")[-1]
def _set_code_unit_attributes(self):
"""
Generates the conversion to various physical units based
on the parameters from the header
"""
# spatial units
z = self.current_redshift
h = self.hubble_constant
boxcm_cal = self.parameters["boxh"]
boxcm_uncal = boxcm_cal / h
box_proper = boxcm_uncal / (1 + z)
aexpn = self.parameters["aexpn"]
# all other units
Om0 = self.parameters["Om0"]
ng = self.parameters["ng"]
boxh = self.parameters["boxh"]
aexpn = self.parameters["aexpn"]
hubble = self.parameters["hubble"]
r0 = boxh / ng
v0 = 50.0 * r0 * np.sqrt(Om0)
rho0 = 2.776e11 * hubble**2.0 * Om0
aM0 = rho0 * (boxh / hubble) ** 3.0 / ng**3.0
velocity = v0 / aexpn * 1.0e5 # proper cm/s
mass = aM0 * 1.98892e33
self.cosmological_simulation = True
setdefaultattr(self, "mass_unit", self.quan(mass, f"g*{ng ** 3}"))
setdefaultattr(self, "length_unit", self.quan(box_proper, "Mpc"))
setdefaultattr(self, "velocity_unit", self.quan(velocity, "cm/s"))
setdefaultattr(self, "time_unit", self.length_unit / self.velocity_unit)
def _parse_parameter_file(self):
"""
Get the various simulation parameters & constants.
"""
self.domain_left_edge = np.zeros(3, dtype="float")
self.domain_right_edge = np.zeros(3, dtype="float") + 1.0
self.dimensionality = 3
self.refine_by = 2
self._periodicity = (True, True, True)
self.cosmological_simulation = True
self.parameters = {}
self.parameters.update(constants)
self.parameters["Time"] = 1.0
# read the amr header
with open(self._file_amr, "rb") as f:
amr_header_vals = fpu.read_attrs(f, amr_header_struct, ">")
n_to_skip = len(("tl", "dtl", "tlold", "dtlold", "iSO"))
fpu.skip(f, n_to_skip, endian=">")
(self.ncell) = fpu.read_vector(f, "i", ">")[0]
# Try to figure out the root grid dimensions
est = int(np.rint(self.ncell ** (1.0 / 3.0)))
# Note here: this is the number of *cells* on the root grid.
# This is not the same as the number of Octs.
# domain dimensions is the number of root *cells*
self.domain_dimensions = np.ones(3, dtype="int64") * est
self.root_grid_mask_offset = f.tell()
self.root_nocts = self.domain_dimensions.prod() // 8
self.root_ncells = self.root_nocts * 8
mylog.debug(
"Estimating %i cells on a root grid side, %i root octs",
est,
self.root_nocts,
)
self.root_iOctCh = fpu.read_vector(f, "i", ">")[: self.root_ncells]
self.root_iOctCh = self.root_iOctCh.reshape(
self.domain_dimensions, order="F"
)
self.root_grid_offset = f.tell()
self.root_nhvar = fpu.skip(f, endian=">")
self.root_nvar = fpu.skip(f, endian=">")
# make sure that the number of root variables is a multiple of
# rootcells
assert self.root_nhvar % self.root_ncells == 0
assert self.root_nvar % self.root_ncells == 0
self.nhydro_variables = (
self.root_nhvar + self.root_nvar
) / self.root_ncells
self.iOctFree, self.nOct = fpu.read_vector(f, "i", ">")
self.child_grid_offset = f.tell()
# lextra needs to be loaded as a string, but it's actually
# array values. So pop it off here, and then re-insert.
lextra = amr_header_vals.pop("lextra")
amr_header_vals["lextra"] = np.frombuffer(lextra, ">f4")
self.parameters.update(amr_header_vals)
amr_header_vals = None
# estimate the root level
float_center, fl, iocts, nocts, root_level = _read_art_level_info(
f, [0, self.child_grid_offset], 1, coarse_grid=self.domain_dimensions[0]
)
del float_center, fl, iocts, nocts
self.root_level = root_level
mylog.info("Using root level of %02i", self.root_level)
# read the particle header
self.particle_types = []
self.particle_types_raw = ()
if not self.skip_particles and self._file_particle_header:
with open(self._file_particle_header, "rb") as fh:
particle_header_vals = fpu.read_attrs(fh, particle_header_struct, ">")
fh.seek(seek_extras)
n = particle_header_vals["Nspecies"]
wspecies = np.fromfile(fh, dtype=">f", count=10)
lspecies = np.fromfile(fh, dtype=">i", count=10)
# extras needs to be loaded as a string, but it's actually
# array values. So pop it off here, and then re-insert.
extras = particle_header_vals.pop("extras")
particle_header_vals["extras"] = np.frombuffer(extras, ">f4")
self.parameters["wspecies"] = wspecies[:n]
self.parameters["lspecies"] = lspecies[:n]
for specie in range(n):
self.particle_types.append("specie%i" % specie)
self.particle_types_raw = tuple(self.particle_types)
ls_nonzero = np.diff(lspecies)[: n - 1]
ls_nonzero = np.append(lspecies[0], ls_nonzero)
self.star_type = len(ls_nonzero)
mylog.info("Discovered %i species of particles", len(ls_nonzero))
info_str = "Particle populations: " + "%9i " * len(ls_nonzero)
mylog.info(info_str, *ls_nonzero)
self._particle_type_counts = dict(zip(self.particle_types_raw, ls_nonzero))
for k, v in particle_header_vals.items():
if k in self.parameters.keys():
if not self.parameters[k] == v:
mylog.info(
"Inconsistent parameter %s %1.1e %1.1e",
k,
v,
self.parameters[k],
)
else:
self.parameters[k] = v
self.parameters_particles = particle_header_vals
self.parameters.update(particle_header_vals)
self.parameters["ng"] = self.parameters["Ngridc"]
self.parameters["ncell0"] = self.parameters["ng"] ** 3
# setup standard simulation params yt expects to see
self.current_redshift = self.parameters["aexpn"] ** -1.0 - 1.0
self.omega_lambda = self.parameters["Oml0"]
self.omega_matter = self.parameters["Om0"]
self.hubble_constant = self.parameters["hubble"]
self.min_level = self.parameters["min_level"]
self.max_level = self.parameters["max_level"]
if self.limit_level is not None:
self.max_level = min(self.limit_level, self.parameters["max_level"])
if self.force_max_level is not None:
self.max_level = self.force_max_level
self.hubble_time = 1.0 / (self.hubble_constant * 100 / 3.08568025e19)
self.current_time = self.quan(b2t(self.parameters["t"]), "Gyr")
self.gamma = self.parameters["gamma"]
mylog.info("Max level is %02i", self.max_level)
[docs]
def create_field_info(self):
super().create_field_info()
if "wspecies" in self.parameters:
# We create dark_matter and stars unions.
ptr = self.particle_types_raw
pu = ParticleUnion("darkmatter", list(ptr[:-1]))
self.add_particle_union(pu)
pu = ParticleUnion("stars", list(ptr[-1:]))
self.add_particle_union(pu)
@classmethod
def _is_valid(cls, filename: str, *args, **kwargs) -> bool:
"""
Defined for the NMSU file naming scheme.
This could differ for other formats.
"""
f = str(filename)
prefix, suffix = filename_pattern["amr"]
if not os.path.isfile(f):
return False
if not f.endswith(suffix):
return False
with open(f, "rb") as fh:
try:
fpu.read_attrs(fh, amr_header_struct, ">")
return True
except Exception:
return False
[docs]
class ARTParticleFile(ParticleFile):
def __init__(self, ds, io, filename, file_id):
super().__init__(ds, io, filename, file_id, range=None)
self.total_particles = {}
for ptype, count in zip(
ds.particle_types_raw, ds.parameters["total_particles"]
):
self.total_particles[ptype] = count
with open(filename, "rb") as f:
f.seek(0, os.SEEK_END)
self._file_size = f.tell()
[docs]
class ARTParticleIndex(ParticleIndex):
def _setup_filenames(self):
# no need for template, all data in one file
template = self.dataset.filename_template
ndoms = self.dataset.file_count
cls = self.dataset._file_class
self.data_files = []
fi = 0
for i in range(int(ndoms)):
df = cls(self.dataset, self.io, template % {"num": i}, fi)
fi += 1
self.data_files.append(df)
[docs]
class DarkMatterARTDataset(ARTDataset):
_index_class = ARTParticleIndex
_file_class = ARTParticleFile
filter_bbox = False
def __init__(
self,
filename,
dataset_type="dm_art",
fields=None,
storage_filename=None,
skip_particles=False,
skip_stars=False,
limit_level=None,
spread_age=True,
force_max_level=None,
file_particle_header=None,
file_particle_stars=None,
units_override=None,
unit_system="cgs",
):
self.num_zones = 2
self.n_ref = 64
self.particle_types += ("all",)
if fields is None:
fields = particle_fields
filename = os.path.abspath(filename)
self._fields_in_file = fields
self._file_particle = filename
self._file_particle_header = file_particle_header
self._find_files(filename)
self.skip_stars = skip_stars
self.spread_age = spread_age
Dataset.__init__(
self,
filename,
dataset_type,
units_override=units_override,
unit_system=unit_system,
)
self.storage_filename = storage_filename
def _find_files(self, file_particle):
"""
Given the particle base filename, attempt to find the
particle header and star files.
"""
base_prefix, base_suffix = filename_pattern["particle_data"]
aexpstr = file_particle.rsplit("s0", 1)[1].replace(base_suffix, "")
possibles = glob.glob(
os.path.join(os.path.dirname(os.path.abspath(file_particle)), "*")
)
for filetype, (prefix, suffix) in filename_pattern.items():
# if this attribute is already set skip it
if getattr(self, "_file_" + filetype, None) is not None:
continue
match = None
for possible in possibles:
if possible.endswith(aexpstr + suffix):
if os.path.basename(possible).startswith(prefix):
match = possible
if match is not None:
mylog.info("discovered %s:%s", filetype, match)
setattr(self, "_file_" + filetype, match)
else:
setattr(self, "_file_" + filetype, None)
def __str__(self):
return self._file_particle.split("/")[-1]
def _set_code_unit_attributes(self):
"""
Generates the conversion to various physical units based
on the parameters from the header
"""
# spatial units
z = self.current_redshift
h = self.hubble_constant
boxcm_cal = self.parameters["boxh"]
boxcm_uncal = boxcm_cal / h
box_proper = boxcm_uncal / (1 + z)
aexpn = self.parameters["aexpn"]
# all other units
Om0 = self.parameters["Om0"]
ng = self.parameters["ng"]
boxh = self.parameters["boxh"]
aexpn = self.parameters["aexpn"]
hubble = self.parameters["hubble"]
r0 = boxh / ng
rho0 = 2.776e11 * hubble**2.0 * Om0
aM0 = rho0 * (boxh / hubble) ** 3.0 / ng**3.0
velocity = 100.0 * r0 / aexpn * 1.0e5 # proper cm/s
mass = aM0 * 1.98892e33
self.cosmological_simulation = True
self.mass_unit = self.quan(mass, f"g*{ng ** 3}")
self.length_unit = self.quan(box_proper, "Mpc")
self.velocity_unit = self.quan(velocity, "cm/s")
self.time_unit = self.length_unit / self.velocity_unit
def _parse_parameter_file(self):
"""
Get the various simulation parameters & constants.
"""
self.domain_left_edge = np.zeros(3, dtype="float")
self.domain_right_edge = np.zeros(3, dtype="float") + 1.0
self.dimensionality = 3
self.refine_by = 2
self._periodicity = (True, True, True)
self.cosmological_simulation = True
self.parameters = {}
self.parameters.update(constants)
self.parameters["Time"] = 1.0
self.file_count = 1
self.filename_template = self.parameter_filename
# read the particle header
self.particle_types = []
self.particle_types_raw = ()
assert self._file_particle_header
with open(self._file_particle_header, "rb") as fh:
seek = 4
fh.seek(seek)
headerstr = fh.read(45).decode("ascii")
aexpn = np.fromfile(fh, count=1, dtype=">f4")
aexp0 = np.fromfile(fh, count=1, dtype=">f4")
amplt = np.fromfile(fh, count=1, dtype=">f4")
astep = np.fromfile(fh, count=1, dtype=">f4")
istep = np.fromfile(fh, count=1, dtype=">i4")
partw = np.fromfile(fh, count=1, dtype=">f4")
tintg = np.fromfile(fh, count=1, dtype=">f4")
ekin = np.fromfile(fh, count=1, dtype=">f4")
ekin1 = np.fromfile(fh, count=1, dtype=">f4")
ekin2 = np.fromfile(fh, count=1, dtype=">f4")
au0 = np.fromfile(fh, count=1, dtype=">f4")
aeu0 = np.fromfile(fh, count=1, dtype=">f4")
nrowc = np.fromfile(fh, count=1, dtype=">i4")
ngridc = np.fromfile(fh, count=1, dtype=">i4")
nspecs = np.fromfile(fh, count=1, dtype=">i4")
nseed = np.fromfile(fh, count=1, dtype=">i4")
Om0 = np.fromfile(fh, count=1, dtype=">f4")
Oml0 = np.fromfile(fh, count=1, dtype=">f4")
hubble = np.fromfile(fh, count=1, dtype=">f4")
Wp5 = np.fromfile(fh, count=1, dtype=">f4")
Ocurv = np.fromfile(fh, count=1, dtype=">f4")
wspecies = np.fromfile(fh, count=10, dtype=">f4")
lspecies = np.fromfile(fh, count=10, dtype=">i4")
extras = np.fromfile(fh, count=79, dtype=">f4")
boxsize = np.fromfile(fh, count=1, dtype=">f4")
n = nspecs[0]
particle_header_vals = {}
tmp = [
headerstr,
aexpn,
aexp0,
amplt,
astep,
istep,
partw,
tintg,
ekin,
ekin1,
ekin2,
au0,
aeu0,
nrowc,
ngridc,
nspecs,
nseed,
Om0,
Oml0,
hubble,
Wp5,
Ocurv,
wspecies,
lspecies,
extras,
boxsize,
]
for i, arr in enumerate(tmp):
a1 = dmparticle_header_struct[0][i]
a2 = dmparticle_header_struct[1][i]
if a2 == 1:
particle_header_vals[a1] = arr[0]
else:
particle_header_vals[a1] = arr[:a2]
for specie in range(n):
self.particle_types.append("specie%i" % specie)
self.particle_types_raw = tuple(self.particle_types)
ls_nonzero = np.diff(lspecies)[: n - 1]
ls_nonzero = np.append(lspecies[0], ls_nonzero)
self.star_type = len(ls_nonzero)
mylog.info("Discovered %i species of particles", len(ls_nonzero))
info_str = "Particle populations: " + "%9i " * len(ls_nonzero)
mylog.info(info_str, *ls_nonzero)
for k, v in particle_header_vals.items():
if k in self.parameters.keys():
if not self.parameters[k] == v:
mylog.info(
"Inconsistent parameter %s %1.1e %1.1e",
k,
v,
self.parameters[k],
)
else:
self.parameters[k] = v
self.parameters_particles = particle_header_vals
self.parameters.update(particle_header_vals)
self.parameters["wspecies"] = wspecies[:n]
self.parameters["lspecies"] = lspecies[:n]
self.parameters["ng"] = self.parameters["Ngridc"]
self.parameters["ncell0"] = self.parameters["ng"] ** 3
self.parameters["boxh"] = self.parameters["boxsize"]
self.parameters["total_particles"] = ls_nonzero
self.domain_dimensions = np.ones(3, dtype="int64") * 2 # NOT ng
# setup standard simulation params yt expects to see
# Convert to float to please unyt
self.current_redshift = float(self.parameters["aexpn"] ** -1.0 - 1.0)
self.omega_lambda = float(particle_header_vals["Oml0"])
self.omega_matter = float(particle_header_vals["Om0"])
self.hubble_constant = float(particle_header_vals["hubble"])
self.min_level = 0
self.max_level = 0
# self.min_level = particle_header_vals['min_level']
# self.max_level = particle_header_vals['max_level']
# if self.limit_level is not None:
# self.max_level = min(
# self.limit_level, particle_header_vals['max_level'])
# if self.force_max_level is not None:
# self.max_level = self.force_max_level
self.hubble_time = 1.0 / (self.hubble_constant * 100 / 3.08568025e19)
self.parameters["t"] = a2b(self.parameters["aexpn"])
self.current_time = self.quan(b2t(self.parameters["t"]), "Gyr")
self.gamma = self.parameters["gamma"]
mylog.info("Max level is %02i", self.max_level)
[docs]
def create_field_info(self):
super(ARTDataset, self).create_field_info()
ptr = self.particle_types_raw
pu = ParticleUnion("darkmatter", list(ptr))
self.add_particle_union(pu)
pass
@classmethod
def _is_valid(cls, filename: str, *args, **kwargs) -> bool:
"""
Defined for the NMSU file naming scheme.
This could differ for other formats.
"""
f = str(filename)
prefix, suffix = filename_pattern["particle_data"]
if not os.path.isfile(f):
return False
if not f.endswith(suffix):
return False
if "s0" not in f:
# ATOMIC.DAT, for instance, passes the other tests, but then dies
# during _find_files because it can't be split.
return False
with open(f, "rb") as fh:
try:
amr_prefix, amr_suffix = filename_pattern["amr"]
possibles = glob.glob(
os.path.join(os.path.dirname(os.path.abspath(f)), "*")
)
for possible in possibles:
if possible.endswith(amr_suffix):
if os.path.basename(possible).startswith(amr_prefix):
return False
except Exception:
pass
try:
seek = 4
fh.seek(seek)
headerstr = np.fromfile(fh, count=1, dtype=(str, 45)) # NOQA
aexpn = np.fromfile(fh, count=1, dtype=">f4") # NOQA
aexp0 = np.fromfile(fh, count=1, dtype=">f4") # NOQA
amplt = np.fromfile(fh, count=1, dtype=">f4") # NOQA
astep = np.fromfile(fh, count=1, dtype=">f4") # NOQA
istep = np.fromfile(fh, count=1, dtype=">i4") # NOQA
partw = np.fromfile(fh, count=1, dtype=">f4") # NOQA
tintg = np.fromfile(fh, count=1, dtype=">f4") # NOQA
ekin = np.fromfile(fh, count=1, dtype=">f4") # NOQA
ekin1 = np.fromfile(fh, count=1, dtype=">f4") # NOQA
ekin2 = np.fromfile(fh, count=1, dtype=">f4") # NOQA
au0 = np.fromfile(fh, count=1, dtype=">f4") # NOQA
aeu0 = np.fromfile(fh, count=1, dtype=">f4") # NOQA
nrowc = np.fromfile(fh, count=1, dtype=">i4") # NOQA
ngridc = np.fromfile(fh, count=1, dtype=">i4") # NOQA
nspecs = np.fromfile(fh, count=1, dtype=">i4") # NOQA
nseed = np.fromfile(fh, count=1, dtype=">i4") # NOQA
Om0 = np.fromfile(fh, count=1, dtype=">f4") # NOQA
Oml0 = np.fromfile(fh, count=1, dtype=">f4") # NOQA
hubble = np.fromfile(fh, count=1, dtype=">f4") # NOQA
Wp5 = np.fromfile(fh, count=1, dtype=">f4") # NOQA
Ocurv = np.fromfile(fh, count=1, dtype=">f4") # NOQA
wspecies = np.fromfile(fh, count=10, dtype=">f4") # NOQA
lspecies = np.fromfile(fh, count=10, dtype=">i4") # NOQA
extras = np.fromfile(fh, count=79, dtype=">f4") # NOQA
boxsize = np.fromfile(fh, count=1, dtype=">f4") # NOQA
return True
except Exception:
return False
[docs]
class ARTDomainSubset(OctreeSubset):
[docs]
def fill(self, content, ftfields, selector):
"""
This is called from IOHandler. It takes content
which is a binary stream, reads the requested field
over this while domain. It then uses oct_handler fill
to reorganize values from IO read index order to
the order they are in in the octhandler.
"""
oct_handler = self.oct_handler
all_fields = self.domain.ds.index.fluid_field_list
fields = [f for ft, f in ftfields]
field_idxs = [all_fields.index(f) for f in fields]
source, tr = {}, {}
cell_count = selector.count_oct_cells(self.oct_handler, self.domain_id)
levels, cell_inds, file_inds = self.oct_handler.file_index_octs(
selector, self.domain_id, cell_count
)
for field in fields:
tr[field] = np.zeros(cell_count, "float64")
data = _read_root_level(
content, self.domain.level_child_offsets, self.domain.level_count
)
ns = (self.domain.ds.domain_dimensions.prod() // 8, 8)
for field, fi in zip(fields, field_idxs):
source[field] = np.empty(ns, dtype="float64", order="C")
dt = data[fi, :].reshape(self.domain.ds.domain_dimensions, order="F")
for i in range(2):
for j in range(2):
for k in range(2):
ii = ((k * 2) + j) * 2 + i
# Note: C order because our index converts C to F.
source[field][:, ii] = dt[i::2, j::2, k::2].ravel(order="C")
oct_handler.fill_level(0, levels, cell_inds, file_inds, tr, source)
del source
# Now we continue with the additional levels.
for level in range(1, self.ds.index.max_level + 1):
no = self.domain.level_count[level]
noct_range = [0, no]
source = _read_child_level(
content,
self.domain.level_child_offsets,
self.domain.level_offsets,
self.domain.level_count,
level,
fields,
self.domain.ds.domain_dimensions,
self.domain.ds.parameters["ncell0"],
noct_range=noct_range,
)
oct_handler.fill_level(level, levels, cell_inds, file_inds, tr, source)
return tr
[docs]
class ARTDomainFile:
"""
Read in the AMR, left/right edges, fill out the octhandler
"""
# We already read in the header in static output,
# and since these headers are defined in only a single file it's
# best to leave them in the static output
_last_mask = None
_last_selector_id = None
def __init__(self, ds, nvar, oct_handler, domain_id):
self.nvar = nvar
self.ds = ds
self.domain_id = domain_id
self._level_count = None
self._level_oct_offsets = None
self._level_child_offsets = None
self.oct_handler = oct_handler
@property
def level_count(self):
# this is number of *octs*
if self._level_count is not None:
return self._level_count
self.level_offsets
return self._level_count
@property
def level_child_offsets(self):
if self._level_count is not None:
return self._level_child_offsets
self.level_offsets
return self._level_child_offsets
@property
def level_offsets(self):
# this is used by the IO operations to find the file offset,
# and then start reading to fill values
# note that this is called hydro_offset in ramses
if self._level_oct_offsets is not None:
return self._level_oct_offsets
# We now have to open the file and calculate it
f = open(self.ds._file_amr, "rb")
(
nhydrovars,
inoll,
_level_oct_offsets,
_level_child_offsets,
) = self._count_art_octs(
f, self.ds.child_grid_offset, self.ds.min_level, self.ds.max_level
)
# remember that the root grid is by itself; manually add it back in
inoll[0] = self.ds.domain_dimensions.prod() // 8
_level_child_offsets[0] = self.ds.root_grid_offset
self.nhydrovars = nhydrovars
self.inoll = inoll # number of octs
self._level_oct_offsets = _level_oct_offsets
self._level_child_offsets = _level_child_offsets
self._level_count = inoll
return self._level_oct_offsets
def _count_art_octs(self, f, offset, MinLev, MaxLevelNow):
level_oct_offsets = [
0,
]
level_child_offsets = [
0,
]
f.seek(offset)
nchild, ntot = 8, 0
Level = np.zeros(MaxLevelNow + 1 - MinLev, dtype="int64")
iNOLL = np.zeros(MaxLevelNow + 1 - MinLev, dtype="int64")
iHOLL = np.zeros(MaxLevelNow + 1 - MinLev, dtype="int64")
for Lev in range(MinLev + 1, MaxLevelNow + 1):
level_oct_offsets.append(f.tell())
# Get the info for this level, skip the rest
# print("Reading oct tree data for level", Lev)
# print('offset:',f.tell())
Level[Lev], iNOLL[Lev], iHOLL[Lev] = fpu.read_vector(f, "i", ">")
# print('Level %i : '%Lev, iNOLL)
# print('offset after level record:',f.tell())
nLevel = iNOLL[Lev]
ntot = ntot + nLevel
# Skip all the oct hierarchy data
ns = fpu.peek_record_size(f, endian=">")
size = struct.calcsize(">i") + ns + struct.calcsize(">i")
f.seek(f.tell() + size * nLevel)
level_child_offsets.append(f.tell())
# Skip the child vars data
ns = fpu.peek_record_size(f, endian=">")
size = struct.calcsize(">i") + ns + struct.calcsize(">i")
f.seek(f.tell() + size * nLevel * nchild)
# find nhydrovars
nhydrovars = 8 + 2
f.seek(offset)
return nhydrovars, iNOLL, level_oct_offsets, level_child_offsets
def _read_amr_level(self, oct_handler):
"""Open the oct file, read in octs level-by-level.
For each oct, only the position, index, level and domain
are needed - its position in the octree is found automatically.
The most important is finding all the information to feed
oct_handler.add
"""
self.level_offsets
f = open(self.ds._file_amr, "rb")
for level in range(1, self.ds.max_level + 1):
unitary_center, fl, iocts, nocts, root_level = _read_art_level_info(
f,
self._level_oct_offsets,
level,
coarse_grid=self.ds.domain_dimensions[0],
root_level=self.ds.root_level,
)
nocts_check = oct_handler.add(self.domain_id, level, unitary_center)
assert nocts_check == nocts
mylog.debug(
"Added %07i octs on level %02i, cumulative is %07i",
nocts,
level,
oct_handler.nocts,
)
def _read_amr_root(self, oct_handler):
self.level_offsets
# add the root *cell* not *oct* mesh
root_octs_side = self.ds.domain_dimensions[0] / 2
NX = np.ones(3) * root_octs_side
LE = np.array([0.0, 0.0, 0.0], dtype="float64")
RE = np.array([1.0, 1.0, 1.0], dtype="float64")
root_dx = (RE - LE) / NX
LL = LE + root_dx / 2.0
RL = RE - root_dx / 2.0
# compute floating point centers of root octs
root_fc = np.mgrid[
LL[0] : RL[0] : NX[0] * 1j,
LL[1] : RL[1] : NX[1] * 1j,
LL[2] : RL[2] : NX[2] * 1j,
]
root_fc = np.vstack([p.ravel() for p in root_fc]).T
oct_handler.add(self.domain_id, 0, root_fc)
assert oct_handler.nocts == root_fc.shape[0]
mylog.debug(
"Added %07i octs on level %02i, cumulative is %07i",
root_octs_side**3,
0,
oct_handler.nocts,
)
[docs]
def included(self, selector):
return True