1import os 2import re 3import string 4import time 5import weakref 6from collections import defaultdict 7 8import numpy as np 9from more_itertools import always_iterable 10 11from yt.data_objects.index_subobjects.grid_patch import AMRGridPatch 12from yt.data_objects.static_output import Dataset 13from yt.fields.field_info_container import NullFunc 14from yt.frontends.enzo.misc import cosmology_get_units 15from yt.funcs import get_pbar, iter_fields, setdefaultattr 16from yt.geometry.geometry_handler import YTDataChunk 17from yt.geometry.grid_geometry_handler import GridIndex 18from yt.utilities.logger import ytLogger as mylog 19from yt.utilities.on_demand_imports import _h5py as h5py, _libconf as libconf 20 21from .fields import EnzoFieldInfo 22 23 24class EnzoGrid(AMRGridPatch): 25 """ 26 Class representing a single Enzo Grid instance. 27 """ 28 29 def __init__(self, id, index): 30 """ 31 Returns an instance of EnzoGrid with *id*, associated with 32 *filename* and *index*. 33 """ 34 # All of the field parameters will be passed to us as needed. 35 AMRGridPatch.__init__(self, id, filename=None, index=index) 36 self._children_ids = [] 37 self._parent_id = -1 38 self.Level = -1 39 40 def set_filename(self, filename): 41 """ 42 Intelligently set the filename. 43 """ 44 if filename is None: 45 self.filename = filename 46 return 47 if self.index._strip_path: 48 self.filename = os.path.join( 49 self.index.directory, os.path.basename(filename) 50 ) 51 elif filename[0] == os.path.sep: 52 self.filename = filename 53 else: 54 self.filename = os.path.join(self.index.directory, filename) 55 return 56 57 def __repr__(self): 58 return "EnzoGrid_%04i" % (self.id) 59 60 @property 61 def Parent(self): 62 if self._parent_id == -1: 63 return None 64 return self.index.grids[self._parent_id - self._id_offset] 65 66 @property 67 def Children(self): 68 return [self.index.grids[cid - self._id_offset] for cid in self._children_ids] 69 70 @property 71 def NumberOfActiveParticles(self): 72 if not hasattr(self.index, "grid_active_particle_count"): 73 return {} 74 id = self.id - self._id_offset 75 nap = { 76 ptype: self.index.grid_active_particle_count[ptype][id] 77 for ptype in self.index.grid_active_particle_count 78 } 79 return nap 80 81 82class EnzoGridInMemory(EnzoGrid): 83 __slots__ = ["proc_num"] 84 85 def set_filename(self, filename): 86 pass 87 88 89class EnzoGridGZ(EnzoGrid): 90 91 __slots__ = () 92 93 def retrieve_ghost_zones(self, n_zones, fields, all_levels=False, smoothed=False): 94 NGZ = self.ds.parameters.get("NumberOfGhostZones", 3) 95 if n_zones > NGZ: 96 return EnzoGrid.retrieve_ghost_zones( 97 self, n_zones, fields, all_levels, smoothed 98 ) 99 100 # ----- Below is mostly the original code, except we remove the field 101 # ----- access section 102 # We will attempt this by creating a datacube that is exactly bigger 103 # than the grid by nZones*dx in each direction 104 nl = self.get_global_startindex() - n_zones 105 new_left_edge = nl * self.dds + self.ds.domain_left_edge 106 # Something different needs to be done for the root grid, though 107 level = self.Level 108 kwargs = { 109 "dims": self.ActiveDimensions + 2 * n_zones, 110 "num_ghost_zones": n_zones, 111 "use_pbar": False, 112 } 113 # This should update the arguments to set the field parameters to be 114 # those of this grid. 115 kwargs.update(self.field_parameters) 116 if smoothed: 117 # cube = self.index.smoothed_covering_grid( 118 # level, new_left_edge, new_right_edge, **kwargs) 119 cube = self.index.smoothed_covering_grid(level, new_left_edge, **kwargs) 120 else: 121 cube = self.index.covering_grid(level, new_left_edge, **kwargs) 122 # ----- This is EnzoGrid.get_data, duplicated here mostly for 123 # ---- efficiency's sake. 124 start_zone = NGZ - n_zones 125 if start_zone == 0: 126 end_zone = None 127 else: 128 end_zone = -(NGZ - n_zones) 129 sl = tuple(slice(start_zone, end_zone) for i in range(3)) 130 if fields is None: 131 return cube 132 for field in iter_fields(fields): 133 if field in self.field_list: 134 conv_factor = 1.0 135 if field in self.ds.field_info: 136 conv_factor = self.ds.field_info[field]._convert_function(self) 137 if self.ds.field_info[field].sampling_type == "particle": 138 continue 139 temp = self.index.io._read_raw_data_set(self, field) 140 temp = temp.swapaxes(0, 2) 141 cube.field_data[field] = np.multiply(temp, conv_factor, temp)[sl] 142 return cube 143 144 145class EnzoHierarchy(GridIndex): 146 147 _strip_path = False 148 grid = EnzoGrid 149 _preload_implemented = True 150 151 def __init__(self, ds, dataset_type): 152 153 self.dataset_type = dataset_type 154 if ds.file_style is not None: 155 self._bn = ds.file_style 156 else: 157 self._bn = "%s.cpu%%04i" 158 self.index_filename = os.path.abspath(f"{ds.parameter_filename}.hierarchy") 159 if os.path.getsize(self.index_filename) == 0: 160 raise OSError(-1, "File empty", self.index_filename) 161 self.directory = os.path.dirname(self.index_filename) 162 163 # For some reason, r8 seems to want Float64 164 if "CompilerPrecision" in ds and ds["CompilerPrecision"] == "r4": 165 self.float_type = "float32" 166 else: 167 self.float_type = "float64" 168 169 GridIndex.__init__(self, ds, dataset_type) 170 # sync it back 171 self.dataset.dataset_type = self.dataset_type 172 173 def _count_grids(self): 174 self.num_grids = None 175 test_grid = test_grid_id = None 176 self.num_stars = 0 177 for line in rlines(open(self.index_filename, "rb")): 178 if ( 179 line.startswith("BaryonFileName") 180 or line.startswith("ParticleFileName") 181 or line.startswith("FileName ") 182 ): 183 test_grid = line.split("=")[-1].strip().rstrip() 184 if line.startswith("NumberOfStarParticles"): 185 self.num_stars = int(line.split("=")[-1]) 186 if line.startswith("Grid "): 187 if self.num_grids is None: 188 self.num_grids = int(line.split("=")[-1]) 189 test_grid_id = int(line.split("=")[-1]) 190 if test_grid is not None: 191 break 192 self._guess_dataset_type(self.ds.dimensionality, test_grid, test_grid_id) 193 194 def _guess_dataset_type(self, rank, test_grid, test_grid_id): 195 if test_grid[0] != os.path.sep: 196 test_grid = os.path.join(self.directory, test_grid) 197 if not os.path.exists(test_grid): 198 test_grid = os.path.join(self.directory, os.path.basename(test_grid)) 199 mylog.debug("Your data uses the annoying hardcoded path.") 200 self._strip_path = True 201 if self.dataset_type is not None: 202 return 203 if rank == 3: 204 mylog.debug("Detected packed HDF5") 205 if self.parameters.get("WriteGhostZones", 0) == 1: 206 self.dataset_type = "enzo_packed_3d_gz" 207 self.grid = EnzoGridGZ 208 else: 209 self.dataset_type = "enzo_packed_3d" 210 elif rank == 2: 211 mylog.debug("Detect packed 2D") 212 self.dataset_type = "enzo_packed_2d" 213 elif rank == 1: 214 mylog.debug("Detect packed 1D") 215 self.dataset_type = "enzo_packed_1d" 216 else: 217 raise NotImplementedError 218 219 # Sets are sorted, so that won't work! 220 def _parse_index(self): 221 def _next_token_line(token, f): 222 for line in f: 223 if line.startswith(token): 224 return line.split()[2:] 225 226 pattern = r"Pointer: Grid\[(\d*)\]->NextGrid(Next|This)Level = (\d*)\s+$" 227 patt = re.compile(pattern) 228 f = open(self.index_filename) 229 self.grids = [self.grid(1, self)] 230 self.grids[0].Level = 0 231 si, ei, LE, RE, fn, npart = [], [], [], [], [], [] 232 pbar = get_pbar("Parsing Hierarchy ", self.num_grids) 233 version = self.dataset.parameters.get("VersionNumber", None) 234 params = self.dataset.parameters 235 if version is None and "Internal" in params: 236 version = float(params["Internal"]["Provenance"]["VersionNumber"]) 237 if version >= 3.0: 238 active_particles = True 239 nap = { 240 ap_type: [] 241 for ap_type in params["Physics"]["ActiveParticles"][ 242 "ActiveParticlesEnabled" 243 ] 244 } 245 else: 246 if "AppendActiveParticleType" in self.parameters: 247 nap = {} 248 active_particles = True 249 for type in self.parameters.get("AppendActiveParticleType", []): 250 nap[type] = [] 251 else: 252 nap = None 253 active_particles = False 254 for grid_id in range(self.num_grids): 255 pbar.update(grid_id + 1) 256 # We will unroll this list 257 si.append(_next_token_line("GridStartIndex", f)) 258 ei.append(_next_token_line("GridEndIndex", f)) 259 LE.append(_next_token_line("GridLeftEdge", f)) 260 RE.append(_next_token_line("GridRightEdge", f)) 261 nb = int(_next_token_line("NumberOfBaryonFields", f)[0]) 262 fn.append([None]) 263 if nb > 0: 264 fn[-1] = _next_token_line("BaryonFileName", f) 265 npart.append(int(_next_token_line("NumberOfParticles", f)[0])) 266 # Below we find out what active particles exist in this grid, 267 # and add their counts individually. 268 if active_particles: 269 ptypes = _next_token_line("PresentParticleTypes", f) 270 counts = [int(c) for c in _next_token_line("ParticleTypeCounts", f)] 271 for ptype in self.parameters.get("AppendActiveParticleType", []): 272 if ptype in ptypes: 273 nap[ptype].append(counts[ptypes.index(ptype)]) 274 else: 275 nap[ptype].append(0) 276 if nb == 0 and npart[-1] > 0: 277 fn[-1] = _next_token_line("ParticleFileName", f) 278 for line in f: 279 if len(line) < 2: 280 break 281 if line.startswith("Pointer:"): 282 vv = patt.findall(line)[0] 283 self.__pointer_handler(vv) 284 pbar.finish() 285 self._fill_arrays(ei, si, LE, RE, npart, nap) 286 temp_grids = np.empty(self.num_grids, dtype="object") 287 temp_grids[:] = self.grids 288 self.grids = temp_grids 289 self.filenames = fn 290 291 def _initialize_grid_arrays(self): 292 super()._initialize_grid_arrays() 293 if "AppendActiveParticleType" in self.parameters.keys() and len( 294 self.parameters["AppendActiveParticleType"] 295 ): 296 gac = { 297 ptype: np.zeros(self.num_grids, dtype="i4") 298 for ptype in self.parameters["AppendActiveParticleType"] 299 } 300 self.grid_active_particle_count = gac 301 302 def _fill_arrays(self, ei, si, LE, RE, npart, nap): 303 self.grid_dimensions.flat[:] = ei 304 self.grid_dimensions -= np.array(si, dtype="i4") 305 self.grid_dimensions += 1 306 self.grid_left_edge.flat[:] = LE 307 self.grid_right_edge.flat[:] = RE 308 self.grid_particle_count.flat[:] = npart 309 if nap is not None: 310 for ptype in nap: 311 self.grid_active_particle_count[ptype].flat[:] = nap[ptype] 312 313 def __pointer_handler(self, m): 314 sgi = int(m[2]) - 1 315 if sgi == -1: 316 return # if it's 0, then we're done with that lineage 317 # Okay, so, we have a pointer. We make a new grid, with an id of the length+1 318 # (recall, Enzo grids are 1-indexed) 319 self.grids.append(self.grid(len(self.grids) + 1, self)) 320 # We'll just go ahead and make a weakref to cache 321 second_grid = self.grids[sgi] # zero-indexed already 322 first_grid = self.grids[int(m[0]) - 1] 323 if m[1] == "Next": 324 first_grid._children_ids.append(second_grid.id) 325 second_grid._parent_id = first_grid.id 326 second_grid.Level = first_grid.Level + 1 327 elif m[1] == "This": 328 if first_grid.Parent is not None: 329 first_grid.Parent._children_ids.append(second_grid.id) 330 second_grid._parent_id = first_grid._parent_id 331 second_grid.Level = first_grid.Level 332 self.grid_levels[sgi] = second_grid.Level 333 334 def _rebuild_top_grids(self, level=0): 335 mylog.info("Rebuilding grids on level %s", level) 336 cmask = self.grid_levels.flat == (level + 1) 337 cmsum = cmask.sum() 338 mask = np.zeros(self.num_grids, dtype="bool") 339 for grid in self.select_grids(level): 340 mask[:] = 0 341 LE = self.grid_left_edge[grid.id - grid._id_offset] 342 RE = self.grid_right_edge[grid.id - grid._id_offset] 343 grids, grid_i = self.get_box_grids(LE, RE) 344 mask[grid_i] = 1 345 grid._children_ids = [] 346 cgrids = self.grids[(mask * cmask).astype("bool")] 347 mylog.info("%s: %s / %s", grid, len(cgrids), cmsum) 348 for cgrid in cgrids: 349 grid._children_ids.append(cgrid.id) 350 cgrid._parent_id = grid.id 351 mylog.info("Finished rebuilding") 352 353 def _populate_grid_objects(self): 354 for g, f in zip(self.grids, self.filenames): 355 g._prepare_grid() 356 g._setup_dx() 357 g.set_filename(f[0]) 358 del self.filenames # No longer needed. 359 self.max_level = self.grid_levels.max() 360 361 def _detect_active_particle_fields(self): 362 ap_list = self.dataset["AppendActiveParticleType"] 363 _fields = {ap: [] for ap in ap_list} 364 fields = [] 365 for ptype in self.dataset["AppendActiveParticleType"]: 366 select_grids = self.grid_active_particle_count[ptype].flat 367 if not np.any(select_grids): 368 current_ptypes = self.dataset.particle_types 369 new_ptypes = [p for p in current_ptypes if p != ptype] 370 self.dataset.particle_types = new_ptypes 371 self.dataset.particle_types_raw = new_ptypes 372 continue 373 if ptype != "DarkMatter": 374 gs = self.grids[select_grids > 0] 375 g = gs[0] 376 handle = h5py.File(g.filename, "r") 377 grid_group = handle[f"/Grid{g.id:08d}"] 378 for pname in ["Active Particles", "Particles"]: 379 if pname in grid_group: 380 break 381 else: 382 raise RuntimeError("Could not find active particle group in data.") 383 node = grid_group[pname] 384 for ptype in (str(p) for p in node): 385 if ptype not in _fields: 386 continue 387 for field in (str(f) for f in node[ptype]): 388 _fields[ptype].append(field) 389 if node[ptype][field].ndim > 1: 390 self.io._array_fields[field] = ( 391 node[ptype][field].shape[1:], 392 ) 393 fields += [(ptype, field) for field in _fields.pop(ptype)] 394 handle.close() 395 return set(fields) 396 397 def _setup_derived_fields(self): 398 super()._setup_derived_fields() 399 aps = self.dataset.parameters.get("AppendActiveParticleType", []) 400 for fname, field in self.ds.field_info.items(): 401 if not field.sampling_type == "particle": 402 continue 403 if isinstance(fname, tuple): 404 continue 405 if field._function is NullFunc: 406 continue 407 for apt in aps: 408 dd = field._copy_def() 409 dd.pop("name") 410 self.ds.field_info.add_field((apt, fname), sampling_type="cell", **dd) 411 412 def _detect_output_fields(self): 413 self.field_list = [] 414 # Do this only on the root processor to save disk work. 415 if self.comm.rank in (0, None): 416 mylog.info("Gathering a field list (this may take a moment.)") 417 field_list = set() 418 random_sample = self._generate_random_grids() 419 for grid in random_sample: 420 if not hasattr(grid, "filename"): 421 continue 422 try: 423 gf = self.io._read_field_names(grid) 424 except self.io._read_exception as e: 425 raise OSError("Grid %s is a bit funky?", grid.id) from e 426 mylog.debug("Grid %s has: %s", grid.id, gf) 427 field_list = field_list.union(gf) 428 if "AppendActiveParticleType" in self.dataset.parameters: 429 ap_fields = self._detect_active_particle_fields() 430 field_list = list(set(field_list).union(ap_fields)) 431 if not any(f[0] == "io" for f in field_list): 432 if "io" in self.dataset.particle_types_raw: 433 ptypes_raw = list(self.dataset.particle_types_raw) 434 ptypes_raw.remove("io") 435 self.dataset.particle_types_raw = tuple(ptypes_raw) 436 437 if "io" in self.dataset.particle_types: 438 ptypes = list(self.dataset.particle_types) 439 ptypes.remove("io") 440 self.dataset.particle_types = tuple(ptypes) 441 ptypes = self.dataset.particle_types 442 ptypes_raw = self.dataset.particle_types_raw 443 else: 444 field_list = None 445 ptypes = None 446 ptypes_raw = None 447 self.field_list = list(self.comm.mpi_bcast(field_list)) 448 self.dataset.particle_types = list(self.comm.mpi_bcast(ptypes)) 449 self.dataset.particle_types_raw = list(self.comm.mpi_bcast(ptypes_raw)) 450 451 def _generate_random_grids(self): 452 if self.num_grids > 40: 453 starter = np.random.randint(0, 20) 454 random_sample = np.mgrid[starter : len(self.grids) - 1 : 20j].astype( 455 "int32" 456 ) 457 # We also add in a bit to make sure that some of the grids have 458 # particles 459 gwp = self.grid_particle_count > 0 460 if np.any(gwp) and not np.any(gwp[(random_sample,)]): 461 # We just add one grid. This is not terribly efficient. 462 first_grid = np.where(gwp)[0][0] 463 random_sample.resize((21,)) 464 random_sample[-1] = first_grid 465 mylog.debug("Added additional grid %s", first_grid) 466 mylog.debug("Checking grids: %s", random_sample.tolist()) 467 else: 468 random_sample = np.mgrid[0 : max(len(self.grids), 1)].astype("int32") 469 return self.grids[(random_sample,)] 470 471 def _get_particle_type_counts(self): 472 try: 473 ret = {} 474 for ptype in self.grid_active_particle_count: 475 ret[ptype] = self.grid_active_particle_count[ptype].sum() 476 return ret 477 except AttributeError: 478 return super()._get_particle_type_counts() 479 480 def find_particles_by_type(self, ptype, max_num=None, additional_fields=None): 481 """ 482 Returns a structure of arrays with all of the particles' 483 positions, velocities, masses, types, IDs, and attributes for 484 a particle type **ptype** for a maximum of **max_num** 485 particles. If non-default particle fields are used, provide 486 them in **additional_fields**. 487 """ 488 # Not sure whether this routine should be in the general HierarchyType. 489 if self.grid_particle_count.sum() == 0: 490 mylog.info("Data contains no particles.") 491 return None 492 if additional_fields is None: 493 additional_fields = [ 494 "metallicity_fraction", 495 "creation_time", 496 "dynamical_time", 497 ] 498 pfields = [f for f in self.field_list if f.startswith("particle_")] 499 nattr = self.dataset["NumberOfParticleAttributes"] 500 if nattr > 0: 501 pfields += additional_fields[:nattr] 502 # Find where the particles reside and count them 503 if max_num is None: 504 max_num = 1e100 505 total = 0 506 pstore = [] 507 for level in range(self.max_level, -1, -1): 508 for grid in self.select_grids(level): 509 index = np.where(grid["particle_type"] == ptype)[0] 510 total += len(index) 511 pstore.append(index) 512 if total >= max_num: 513 break 514 if total >= max_num: 515 break 516 result = None 517 if total > 0: 518 result = {} 519 for p in pfields: 520 result[p] = np.zeros(total, "float64") 521 # Now we retrieve data for each field 522 ig = count = 0 523 for level in range(self.max_level, -1, -1): 524 for grid in self.select_grids(level): 525 nidx = len(pstore[ig]) 526 if nidx > 0: 527 for p in pfields: 528 result[p][count : count + nidx] = grid[p][pstore[ig]] 529 count += nidx 530 ig += 1 531 if count >= total: 532 break 533 if count >= total: 534 break 535 # Crop data if retrieved more than max_num 536 if count > max_num: 537 for p in pfields: 538 result[p] = result[p][0:max_num] 539 return result 540 541 542class EnzoHierarchyInMemory(EnzoHierarchy): 543 544 grid = EnzoGridInMemory 545 _enzo = None 546 547 @property 548 def enzo(self): 549 if self._enzo is None: 550 import enzo 551 552 self._enzo = enzo 553 return self._enzo 554 555 def __init__(self, ds, dataset_type=None): 556 self.dataset_type = dataset_type 557 self.float_type = "float64" 558 self.dataset = weakref.proxy(ds) # for _obtain_enzo 559 self.float_type = self.enzo.hierarchy_information["GridLeftEdge"].dtype 560 self.directory = os.getcwd() 561 GridIndex.__init__(self, ds, dataset_type) 562 563 def _initialize_data_storage(self): 564 pass 565 566 def _count_grids(self): 567 self.num_grids = self.enzo.hierarchy_information["GridDimensions"].shape[0] 568 569 def _parse_index(self): 570 self._copy_index_structure() 571 mylog.debug("Copying reverse tree") 572 reverse_tree = self.enzo.hierarchy_information["GridParentIDs"].ravel().tolist() 573 # Initial setup: 574 mylog.debug("Reconstructing parent-child relationships") 575 grids = [] 576 # We enumerate, so it's 0-indexed id and 1-indexed pid 577 self.filenames = ["-1"] * self.num_grids 578 for id, pid in enumerate(reverse_tree): 579 grids.append(self.grid(id + 1, self)) 580 grids[-1].Level = self.grid_levels[id, 0] 581 if pid > 0: 582 grids[-1]._parent_id = pid 583 grids[pid - 1]._children_ids.append(grids[-1].id) 584 self.max_level = self.grid_levels.max() 585 mylog.debug("Preparing grids") 586 self.grids = np.empty(len(grids), dtype="object") 587 for i, grid in enumerate(grids): 588 if (i % 1e4) == 0: 589 mylog.debug("Prepared % 7i / % 7i grids", i, self.num_grids) 590 grid.filename = "Inline_processor_%07i" % (self.grid_procs[i, 0]) 591 grid._prepare_grid() 592 grid._setup_dx() 593 grid.proc_num = self.grid_procs[i, 0] 594 self.grids[i] = grid 595 mylog.debug("Prepared") 596 597 def _initialize_grid_arrays(self): 598 EnzoHierarchy._initialize_grid_arrays(self) 599 self.grid_procs = np.zeros((self.num_grids, 1), "int32") 600 601 def _copy_index_structure(self): 602 # Dimensions are important! 603 self.grid_dimensions[:] = self.enzo.hierarchy_information["GridEndIndices"][:] 604 self.grid_dimensions -= self.enzo.hierarchy_information["GridStartIndices"][:] 605 self.grid_dimensions += 1 606 self.grid_left_edge[:] = self.enzo.hierarchy_information["GridLeftEdge"][:] 607 self.grid_right_edge[:] = self.enzo.hierarchy_information["GridRightEdge"][:] 608 self.grid_levels[:] = self.enzo.hierarchy_information["GridLevels"][:] 609 self.grid_procs = self.enzo.hierarchy_information["GridProcs"].copy() 610 self.grid_particle_count[:] = self.enzo.hierarchy_information[ 611 "GridNumberOfParticles" 612 ][:] 613 614 def save_data(self, *args, **kwargs): 615 pass 616 617 _cached_field_list = None 618 _cached_derived_field_list = None 619 620 def _generate_random_grids(self): 621 my_rank = self.comm.rank 622 my_grids = self.grids[self.grid_procs.ravel() == my_rank] 623 if len(my_grids) > 40: 624 starter = np.random.randint(0, 20) 625 random_sample = np.mgrid[starter : len(my_grids) - 1 : 20j].astype("int32") 626 mylog.debug("Checking grids: %s", random_sample.tolist()) 627 else: 628 random_sample = np.mgrid[0 : max(len(my_grids) - 1, 1)].astype("int32") 629 return my_grids[(random_sample,)] 630 631 def _chunk_io(self, dobj, cache=True, local_only=False): 632 gfiles = defaultdict(list) 633 gobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) 634 for g in gobjs: 635 gfiles[g.filename].append(g) 636 for fn in sorted(gfiles): 637 if local_only: 638 gobjs = [g for g in gfiles[fn] if g.proc_num == self.comm.rank] 639 gfiles[fn] = gobjs 640 gs = gfiles[fn] 641 count = self._count_selection(dobj, gs) 642 yield YTDataChunk(dobj, "io", gs, count, cache=cache) 643 644 645class EnzoHierarchy1D(EnzoHierarchy): 646 def _fill_arrays(self, ei, si, LE, RE, npart, nap): 647 self.grid_dimensions[:, :1] = ei 648 self.grid_dimensions[:, :1] -= np.array(si, dtype="i4") 649 self.grid_dimensions += 1 650 self.grid_left_edge[:, :1] = LE 651 self.grid_right_edge[:, :1] = RE 652 self.grid_particle_count.flat[:] = npart 653 self.grid_left_edge[:, 1:] = 0.0 654 self.grid_right_edge[:, 1:] = 1.0 655 self.grid_dimensions[:, 1:] = 1 656 if nap is not None: 657 raise NotImplementedError 658 659 660class EnzoHierarchy2D(EnzoHierarchy): 661 def _fill_arrays(self, ei, si, LE, RE, npart, nap): 662 self.grid_dimensions[:, :2] = ei 663 self.grid_dimensions[:, :2] -= np.array(si, dtype="i4") 664 self.grid_dimensions += 1 665 self.grid_left_edge[:, :2] = LE 666 self.grid_right_edge[:, :2] = RE 667 self.grid_particle_count.flat[:] = npart 668 self.grid_left_edge[:, 2] = 0.0 669 self.grid_right_edge[:, 2] = 1.0 670 self.grid_dimensions[:, 2] = 1 671 if nap is not None: 672 raise NotImplementedError 673 674 675class EnzoDataset(Dataset): 676 """ 677 Enzo-specific output, set at a fixed time. 678 """ 679 680 _index_class = EnzoHierarchy 681 _field_info_class = EnzoFieldInfo 682 683 def __init__( 684 self, 685 filename, 686 dataset_type=None, 687 file_style=None, 688 parameter_override=None, 689 conversion_override=None, 690 storage_filename=None, 691 units_override=None, 692 unit_system="cgs", 693 default_species_fields=None, 694 ): 695 """ 696 This class is a stripped down class that simply reads and parses 697 *filename* without looking at the index. *dataset_type* gets passed 698 to the index to pre-determine the style of data-output. However, 699 it is not strictly necessary. Optionally you may specify a 700 *parameter_override* dictionary that will override anything in the 701 parameter file and a *conversion_override* dictionary that consists 702 of {fieldname : conversion_to_cgs} that will override the #DataCGS. 703 """ 704 self.fluid_types += ("enzo",) 705 if filename.endswith(".hierarchy"): 706 filename = filename[:-10] 707 if parameter_override is None: 708 parameter_override = {} 709 self._parameter_override = parameter_override 710 if conversion_override is None: 711 conversion_override = {} 712 self._conversion_override = conversion_override 713 self.storage_filename = storage_filename 714 Dataset.__init__( 715 self, 716 filename, 717 dataset_type, 718 file_style=file_style, 719 units_override=units_override, 720 unit_system=unit_system, 721 default_species_fields=default_species_fields, 722 ) 723 724 def _setup_1d(self): 725 self._index_class = EnzoHierarchy1D 726 self.domain_left_edge = np.concatenate([[self.domain_left_edge], [0.0, 0.0]]) 727 self.domain_right_edge = np.concatenate([[self.domain_right_edge], [1.0, 1.0]]) 728 729 def _setup_2d(self): 730 self._index_class = EnzoHierarchy2D 731 self.domain_left_edge = np.concatenate([self.domain_left_edge, [0.0]]) 732 self.domain_right_edge = np.concatenate([self.domain_right_edge, [1.0]]) 733 734 def get_parameter(self, parameter, type=None): 735 """ 736 Gets a parameter not in the parameterDict. 737 """ 738 if parameter in self.parameters: 739 return self.parameters[parameter] 740 for line in open(self.parameter_filename): 741 if line.find("#") >= 1: # Keep the commented lines 742 line = line[: line.find("#")] 743 line = line.strip().rstrip() 744 if len(line) < 2: 745 continue 746 try: 747 param, vals = map(string.strip, map(string.rstrip, line.split("="))) 748 except ValueError: 749 mylog.error("ValueError: '%s'", line) 750 if parameter == param: 751 if type is None: 752 t = vals.split() 753 else: 754 t = map(type, vals.split()) 755 if len(t) == 1: 756 self.parameters[param] = t[0] 757 else: 758 self.parameters[param] = t 759 if param.endswith("Units") and not param.startswith("Temperature"): 760 dataType = param[:-5] 761 self.conversion_factors[dataType] = self.parameters[param] 762 return self.parameters[parameter] 763 764 return "" 765 766 def _parse_parameter_file(self): 767 """ 768 Parses the parameter file and establishes the various 769 dictionaries. 770 """ 771 # Let's read the file 772 with open(self.parameter_filename) as f: 773 line = f.readline().strip() 774 f.seek(0) 775 if line == "Internal:": 776 self._parse_enzo3_parameter_file(f) 777 else: 778 self._parse_enzo2_parameter_file(f) 779 780 def _parse_enzo3_parameter_file(self, f): 781 self.parameters = p = libconf.load(f) 782 sim = p["SimulationControl"] 783 internal = p["Internal"] 784 phys = p["Physics"] 785 self.refine_by = sim["AMR"]["RefineBy"] 786 self._periodicity = tuple( 787 a == 3 for a in sim["Domain"]["LeftFaceBoundaryCondition"] 788 ) 789 self.dimensionality = sim["Domain"]["TopGridRank"] 790 self.domain_dimensions = np.array( 791 sim["Domain"]["TopGridDimensions"], dtype="int64" 792 ) 793 self.domain_left_edge = np.array( 794 sim["Domain"]["DomainLeftEdge"], dtype="float64" 795 ) 796 self.domain_right_edge = np.array( 797 sim["Domain"]["DomainRightEdge"], dtype="float64" 798 ) 799 self.gamma = phys["Hydro"]["Gamma"] 800 self.unique_identifier = internal["Provenance"]["CurrentTimeIdentifier"] 801 self.current_time = internal["InitialTime"] 802 self.cosmological_simulation = phys["Cosmology"]["ComovingCoordinates"] 803 if self.cosmological_simulation == 1: 804 cosmo = phys["Cosmology"] 805 self.current_redshift = internal["CosmologyCurrentRedshift"] 806 self.omega_lambda = cosmo["OmegaLambdaNow"] 807 self.omega_matter = cosmo["OmegaMatterNow"] 808 self.hubble_constant = cosmo["HubbleConstantNow"] 809 else: 810 self.current_redshift = 0.0 811 self.omega_lambda = 0.0 812 self.omega_matter = 0.0 813 self.hubble_constant = 0.0 814 self.cosmological_simulation = 0 815 self.particle_types = ["DarkMatter"] + phys["ActiveParticles"][ 816 "ActiveParticlesEnabled" 817 ] 818 self.particle_types = tuple(self.particle_types) 819 self.particle_types_raw = self.particle_types 820 if self.dimensionality == 1: 821 self._setup_1d() 822 elif self.dimensionality == 2: 823 self._setup_2d() 824 825 def _parse_enzo2_parameter_file(self, f): 826 for line in (l.strip() for l in f): 827 if (len(line) < 2) or ("=" not in line): 828 continue 829 param, vals = (i.strip() for i in line.split("=", 1)) 830 # First we try to decipher what type of value it is. 831 vals = vals.split() 832 # Special case approaching. 833 if "(do" in vals: 834 vals = vals[:1] 835 if len(vals) == 0: 836 pcast = str # Assume NULL output 837 else: 838 v = vals[0] 839 # Figure out if it's castable to floating point: 840 try: 841 float(v) 842 except ValueError: 843 pcast = str 844 else: 845 if any("." in v or "e+" in v or "e-" in v for v in vals): 846 pcast = float 847 elif v == "inf": 848 pcast = str 849 else: 850 pcast = int 851 # Now we figure out what to do with it. 852 if len(vals) == 0: 853 vals = "" 854 elif len(vals) == 1: 855 vals = pcast(vals[0]) 856 else: 857 vals = np.array([pcast(i) for i in vals if i != "-99999"]) 858 if param.startswith("Append"): 859 if param not in self.parameters: 860 self.parameters[param] = [] 861 self.parameters[param].append(vals) 862 else: 863 self.parameters[param] = vals 864 self.refine_by = self.parameters["RefineBy"] 865 _periodicity = tuple( 866 always_iterable(self.parameters["LeftFaceBoundaryCondition"] == 3) 867 ) 868 self.dimensionality = self.parameters["TopGridRank"] 869 if "MetaDataDatasetUUID" in self.parameters: 870 self.unique_identifier = self.parameters["MetaDataDatasetUUID"] 871 elif "CurrentTimeIdentifier" in self.parameters: 872 self.unique_identifier = self.parameters["CurrentTimeIdentifier"] 873 if self.dimensionality > 1: 874 self.domain_dimensions = self.parameters["TopGridDimensions"] 875 if len(self.domain_dimensions) < 3: 876 tmp = self.domain_dimensions.tolist() 877 tmp.append(1) 878 self.domain_dimensions = np.array(tmp) 879 _periodicity += (False,) 880 self.domain_left_edge = np.array( 881 self.parameters["DomainLeftEdge"], "float64" 882 ).copy() 883 self.domain_right_edge = np.array( 884 self.parameters["DomainRightEdge"], "float64" 885 ).copy() 886 else: 887 self.domain_left_edge = np.array( 888 self.parameters["DomainLeftEdge"], "float64" 889 ) 890 self.domain_right_edge = np.array( 891 self.parameters["DomainRightEdge"], "float64" 892 ) 893 self.domain_dimensions = np.array( 894 [self.parameters["TopGridDimensions"], 1, 1] 895 ) 896 _periodicity += (False, False) 897 assert len(_periodicity) == 3 898 self._periodicity = _periodicity 899 900 self.gamma = self.parameters["Gamma"] 901 if self.parameters["ComovingCoordinates"]: 902 self.cosmological_simulation = 1 903 self.current_redshift = self.parameters["CosmologyCurrentRedshift"] 904 self.omega_lambda = self.parameters["CosmologyOmegaLambdaNow"] 905 self.omega_matter = self.parameters["CosmologyOmegaMatterNow"] 906 self.omega_radiation = self.parameters.get( 907 "CosmologyOmegaRadiationNow", 0.0 908 ) 909 self.hubble_constant = self.parameters["CosmologyHubbleConstantNow"] 910 else: 911 self.current_redshift = 0.0 912 self.omega_lambda = 0.0 913 self.omega_matter = 0.0 914 self.hubble_constant = 0.0 915 self.cosmological_simulation = 0 916 self.particle_types = [] 917 self.current_time = self.parameters["InitialTime"] 918 if ( 919 self.parameters["NumberOfParticles"] > 0 920 and "AppendActiveParticleType" in self.parameters.keys() 921 ): 922 # If this is the case, then we know we should have a DarkMatter 923 # particle type, and we don't need the "io" type. 924 self.parameters["AppendActiveParticleType"].append("DarkMatter") 925 else: 926 # We do not have an "io" type for Enzo particles if the 927 # ActiveParticle machinery is on, as we simply will ignore any of 928 # the non-DarkMatter particles in that case. However, for older 929 # datasets, we call this particle type "io". 930 self.particle_types = ["io"] 931 for ptype in self.parameters.get("AppendActiveParticleType", []): 932 self.particle_types.append(ptype) 933 self.particle_types = tuple(self.particle_types) 934 self.particle_types_raw = self.particle_types 935 936 if self.dimensionality == 1: 937 self._setup_1d() 938 elif self.dimensionality == 2: 939 self._setup_2d() 940 941 def _set_code_unit_attributes(self): 942 if self.cosmological_simulation: 943 k = cosmology_get_units( 944 self.hubble_constant, 945 self.omega_matter, 946 self.parameters["CosmologyComovingBoxSize"], 947 self.parameters["CosmologyInitialRedshift"], 948 self.current_redshift, 949 ) 950 # Now some CGS values 951 box_size = self.parameters["CosmologyComovingBoxSize"] 952 setdefaultattr(self, "length_unit", self.quan(box_size, "Mpccm/h")) 953 setdefaultattr( 954 self, 955 "mass_unit", 956 self.quan(k["urho"], "g/cm**3") * (self.length_unit.in_cgs()) ** 3, 957 ) 958 setdefaultattr(self, "time_unit", self.quan(k["utim"], "s")) 959 setdefaultattr(self, "velocity_unit", self.quan(k["uvel"], "cm/s")) 960 else: 961 if "LengthUnits" in self.parameters: 962 length_unit = self.parameters["LengthUnits"] 963 mass_unit = self.parameters["DensityUnits"] * length_unit ** 3 964 time_unit = self.parameters["TimeUnits"] 965 elif "SimulationControl" in self.parameters: 966 units = self.parameters["SimulationControl"]["Units"] 967 length_unit = units["Length"] 968 mass_unit = units["Density"] * length_unit ** 3 969 time_unit = units["Time"] 970 else: 971 mylog.warning("Setting 1.0 in code units to be 1.0 cm") 972 mylog.warning("Setting 1.0 in code units to be 1.0 s") 973 length_unit = mass_unit = time_unit = 1.0 974 975 setdefaultattr(self, "length_unit", self.quan(length_unit, "cm")) 976 setdefaultattr(self, "mass_unit", self.quan(mass_unit, "g")) 977 setdefaultattr(self, "time_unit", self.quan(time_unit, "s")) 978 setdefaultattr(self, "velocity_unit", self.length_unit / self.time_unit) 979 980 density_unit = self.mass_unit / self.length_unit ** 3 981 magnetic_unit = np.sqrt(4 * np.pi * density_unit) * self.velocity_unit 982 magnetic_unit = np.float64(magnetic_unit.in_cgs()) 983 setdefaultattr(self, "magnetic_unit", self.quan(magnetic_unit, "gauss")) 984 985 @classmethod 986 def _is_valid(cls, filename, *args, **kwargs): 987 if str(filename).endswith(".hierarchy"): 988 return True 989 return os.path.exists(f"{filename}.hierarchy") 990 991 @classmethod 992 def _guess_candidates(cls, base, directories, files): 993 candidates = [ 994 _ 995 for _ in files 996 if _.endswith(".hierarchy") 997 and os.path.exists(os.path.join(base, _.rsplit(".", 1)[0])) 998 ] 999 # Typically, Enzo won't have nested outputs. 1000 return candidates, (len(candidates) == 0) 1001 1002 1003class EnzoDatasetInMemory(EnzoDataset): 1004 _index_class = EnzoHierarchyInMemory 1005 _dataset_type = "enzo_inline" 1006 1007 def __new__(cls, *args, **kwargs): 1008 obj = object.__new__(cls) 1009 obj.__init__(*args, **kwargs) 1010 return obj 1011 1012 def __init__(self, parameter_override=None, conversion_override=None): 1013 self.fluid_types += ("enzo",) 1014 if parameter_override is None: 1015 parameter_override = {} 1016 self._parameter_override = parameter_override 1017 if conversion_override is None: 1018 conversion_override = {} 1019 self._conversion_override = conversion_override 1020 1021 Dataset.__init__(self, "InMemoryParameterFile", self._dataset_type) 1022 1023 def _parse_parameter_file(self): 1024 enzo = self._obtain_enzo() 1025 self.basename = "cycle%08i" % (enzo.yt_parameter_file["NumberOfPythonCalls"]) 1026 self.parameters["CurrentTimeIdentifier"] = time.time() 1027 self.parameters.update(enzo.yt_parameter_file) 1028 self.conversion_factors.update(enzo.conversion_factors) 1029 for i in self.parameters: 1030 if isinstance(self.parameters[i], tuple): 1031 self.parameters[i] = np.array(self.parameters[i]) 1032 if i.endswith("Units") and not i.startswith("Temperature"): 1033 dataType = i[:-5] 1034 self.conversion_factors[dataType] = self.parameters[i] 1035 self.domain_left_edge = self.parameters["DomainLeftEdge"].copy() 1036 self.domain_right_edge = self.parameters["DomainRightEdge"].copy() 1037 for i in self.conversion_factors: 1038 if isinstance(self.conversion_factors[i], tuple): 1039 self.conversion_factors[i] = np.array(self.conversion_factors[i]) 1040 for p, v in self._parameter_override.items(): 1041 self.parameters[p] = v 1042 for p, v in self._conversion_override.items(): 1043 self.conversion_factors[p] = v 1044 self.refine_by = self.parameters["RefineBy"] 1045 self._periodicity = tuple( 1046 always_iterable(self.parameters["LeftFaceBoundaryCondition"] == 3) 1047 ) 1048 self.dimensionality = self.parameters["TopGridRank"] 1049 self.domain_dimensions = self.parameters["TopGridDimensions"] 1050 self.current_time = self.parameters["InitialTime"] 1051 if "CurrentTimeIdentifier" in self.parameters: 1052 self.unique_identifier = self.parameters["CurrentTimeIdentifier"] 1053 if self.parameters["ComovingCoordinates"]: 1054 self.cosmological_simulation = 1 1055 self.current_redshift = self.parameters["CosmologyCurrentRedshift"] 1056 self.omega_lambda = self.parameters["CosmologyOmegaLambdaNow"] 1057 self.omega_matter = self.parameters["CosmologyOmegaMatterNow"] 1058 self.hubble_constant = self.parameters["CosmologyHubbleConstantNow"] 1059 else: 1060 self.current_redshift = 0.0 1061 self.omega_lambda = 0.0 1062 self.omega_matter = 0.0 1063 self.hubble_constant = 0.0 1064 self.cosmological_simulation = 0 1065 1066 def _obtain_enzo(self): 1067 import enzo 1068 1069 return enzo 1070 1071 @classmethod 1072 def _is_valid(cls, filename, *args, **kwargs): 1073 return False 1074 1075 1076# These next two functions are taken from 1077# http://www.reddit.com/r/Python/comments/6hj75/reverse_file_iterator/c03vms4 1078# Credit goes to "Brian" on Reddit 1079 1080 1081def rblocks(f, blocksize=4096): 1082 """Read file as series of blocks from end of file to start. 1083 1084 The data itself is in normal order, only the order of the blocks is reversed. 1085 ie. "hello world" -> ["ld","wor", "lo ", "hel"] 1086 Note that the file must be opened in binary mode. 1087 """ 1088 if "b" not in f.mode.lower(): 1089 raise Exception("File must be opened using binary mode.") 1090 size = os.stat(f.name).st_size 1091 fullblocks, lastblock = divmod(size, blocksize) 1092 1093 # The first(end of file) block will be short, since this leaves 1094 # the rest aligned on a blocksize boundary. This may be more 1095 # efficient than having the last (first in file) block be short 1096 f.seek(-lastblock, 2) 1097 yield f.read(lastblock).decode("ascii") 1098 1099 for i in range(fullblocks - 1, -1, -1): 1100 f.seek(i * blocksize) 1101 yield f.read(blocksize).decode("ascii") 1102 1103 1104def rlines(f, keepends=False): 1105 """Iterate through the lines of a file in reverse order. 1106 1107 If keepends is true, line endings are kept as part of the line. 1108 """ 1109 buf = "" 1110 for block in rblocks(f): 1111 buf = block + buf 1112 lines = buf.splitlines(keepends) 1113 # Return all lines except the first (since may be partial) 1114 if lines: 1115 lines.reverse() 1116 buf = lines.pop() # Last line becomes end of new first line. 1117 yield from lines 1118 yield buf # First line. 1119