import os
import weakref
from functools import cached_property
import numpy as np
from yt.data_objects.index_subobjects.unstructured_mesh import SemiStructuredMesh
from yt.data_objects.static_output import Dataset
from yt.funcs import setdefaultattr
from yt.geometry.unstructured_mesh_handler import UnstructuredIndex
from yt.utilities.file_handler import HDF5FileHandler
from yt.utilities.on_demand_imports import _h5py as h5py
from .fields import MoabFieldInfo, PyneFieldInfo
[docs]
class MoabHex8Mesh(SemiStructuredMesh):
_connectivity_length = 8
_index_offset = 1
[docs]
class MoabHex8Hierarchy(UnstructuredIndex):
def __init__(self, ds, dataset_type="h5m"):
self.dataset = weakref.proxy(ds)
self.dataset_type = dataset_type
self.index_filename = self.dataset.parameter_filename
self.directory = os.path.dirname(self.index_filename)
self._fhandle = h5py.File(self.index_filename, mode="r")
UnstructuredIndex.__init__(self, ds, dataset_type)
self._fhandle.close()
def _initialize_mesh(self):
con = self._fhandle["/tstt/elements/Hex8/connectivity"][:]
con = np.asarray(con, dtype="int64")
coords = self._fhandle["/tstt/nodes/coordinates"][:]
coords = np.asarray(coords, dtype="float64")
self.meshes = [MoabHex8Mesh(0, self.index_filename, con, coords, self)]
def _detect_output_fields(self):
self.field_list = [
("moab", f) for f in self._fhandle["/tstt/elements/Hex8/tags"].keys()
]
def _count_grids(self):
self.num_grids = 1
[docs]
class MoabHex8Dataset(Dataset):
_load_requirements = ["h5py"]
_index_class = MoabHex8Hierarchy
_field_info_class = MoabFieldInfo
periodicity = (False, False, False)
def __init__(
self,
filename,
dataset_type="moab_hex8",
storage_filename=None,
units_override=None,
unit_system="cgs",
):
self.fluid_types += ("moab",)
Dataset.__init__(
self,
filename,
dataset_type,
units_override=units_override,
unit_system=unit_system,
)
self.storage_filename = storage_filename
self._handle = HDF5FileHandler(filename)
def _set_code_unit_attributes(self):
# Almost everything is regarded as dimensionless in MOAB, so these will
# not be used very much or at all.
setdefaultattr(self, "length_unit", self.quan(1.0, "cm"))
setdefaultattr(self, "time_unit", self.quan(1.0, "s"))
setdefaultattr(self, "mass_unit", self.quan(1.0, "g"))
def _parse_parameter_file(self):
self._handle = h5py.File(self.parameter_filename, mode="r")
coords = self._handle["/tstt/nodes/coordinates"]
self.domain_left_edge = coords[0]
self.domain_right_edge = coords[-1]
self.domain_dimensions = self.domain_right_edge - self.domain_left_edge
self.refine_by = 2
self.dimensionality = len(self.domain_dimensions)
self.current_time = 0.0
self.cosmological_simulation = False
self.num_ghost_zones = 0
self.current_redshift = 0.0
self.omega_lambda = 0.0
self.omega_matter = 0.0
self.hubble_constant = 0.0
self.cosmological_simulation = 0
@classmethod
def _is_valid(cls, filename: str, *args, **kwargs) -> bool:
return filename.endswith(".h5m") and not cls._missing_load_requirements()
def __str__(self):
return self.basename.rsplit(".", 1)[0]
[docs]
class PyneHex8Mesh(SemiStructuredMesh):
_connectivity_length = 8
_index_offset = 0
[docs]
class PyneMeshHex8Hierarchy(UnstructuredIndex):
def __init__(self, ds, dataset_type="moab_hex8_pyne"):
self.dataset = weakref.proxy(ds)
self.dataset_type = dataset_type
self.index_filename = self.dataset.parameter_filename
self.directory = os.getcwd()
self.pyne_mesh = ds.pyne_mesh
super().__init__(ds, dataset_type)
def _initialize_mesh(self):
from pymoab import types
ents = list(self.pyne_mesh.structured_iterate_vertex())
coords = self.pyne_mesh.mesh.get_coords(ents).astype("float64")
coords = coords.reshape(len(coords) // 3, 3)
hexes = self.pyne_mesh.mesh.get_entities_by_type(0, types.MBHEX)
vind = []
for h in hexes:
vind.append(
self.pyne_mesh.mesh.get_adjacencies(
h, 0, create_if_missing=True, op_type=types.UNION
)
)
vind = np.asarray(vind, dtype=np.int64)
if vind.ndim == 1:
vind = vind.reshape(len(vind) // 8, 8)
assert vind.ndim == 2 and vind.shape[1] == 8
self.meshes = [PyneHex8Mesh(0, self.index_filename, vind, coords, self)]
def _detect_output_fields(self):
self.field_list = [("pyne", f) for f in self.pyne_mesh.tags.keys()]
def _count_grids(self):
self.num_grids = 1
[docs]
class PyneMoabHex8Dataset(Dataset):
_index_class = PyneMeshHex8Hierarchy
_fieldinfo_fallback = MoabFieldInfo
_field_info_class = PyneFieldInfo
periodicity = (False, False, False)
def __init__(
self,
pyne_mesh,
dataset_type="moab_hex8_pyne",
storage_filename=None,
units_override=None,
unit_system="cgs",
):
self.fluid_types += ("pyne",)
filename = f"pyne_mesh_{id(pyne_mesh)}"
self.pyne_mesh = pyne_mesh
Dataset.__init__(
self,
str(filename),
dataset_type,
units_override=units_override,
unit_system=unit_system,
)
self.storage_filename = storage_filename
@property
def filename(self) -> str:
return self._input_filename
@cached_property
def unique_identifier(self) -> str:
return self.filename
def _set_code_unit_attributes(self):
# Almost everything is regarded as dimensionless in MOAB, so these will
# not be used very much or at all.
setdefaultattr(self, "length_unit", self.quan(1.0, "cm"))
setdefaultattr(self, "time_unit", self.quan(1.0, "s"))
setdefaultattr(self, "mass_unit", self.quan(1.0, "g"))
def _parse_parameter_file(self):
ents = list(self.pyne_mesh.structured_iterate_vertex())
coords = self.pyne_mesh.mesh.get_coords(ents)
self.domain_left_edge = coords[0:3]
self.domain_right_edge = coords[-3:]
self.domain_dimensions = self.domain_right_edge - self.domain_left_edge
self.refine_by = 2
self.dimensionality = len(self.domain_dimensions)
self.current_time = 0.0
self.cosmological_simulation = False
self.num_ghost_zones = 0
self.current_redshift = 0.0
self.omega_lambda = 0.0
self.omega_matter = 0.0
self.hubble_constant = 0.0
self.cosmological_simulation = 0
@classmethod
def _is_valid(cls, filename: str, *args, **kwargs) -> bool:
return False
def __str__(self):
return self.basename.rsplit(".", 1)[0]