1import glob 2import os 3 4import numpy as np 5from unyt import dimensions, unyt_array 6from unyt.unit_registry import UnitRegistry 7 8from yt.data_objects.time_series import DatasetSeries, SimulationTimeSeries 9from yt.funcs import only_on_root 10from yt.loaders import load 11from yt.utilities.cosmology import Cosmology 12from yt.utilities.exceptions import ( 13 InvalidSimulationTimeSeries, 14 MissingParameter, 15 NoStoppingCondition, 16 YTUnidentifiedDataType, 17) 18from yt.utilities.logger import ytLogger as mylog 19from yt.utilities.parallel_tools.parallel_analysis_interface import parallel_objects 20 21 22class EnzoSimulation(SimulationTimeSeries): 23 r""" 24 Initialize an Enzo Simulation object. 25 26 Upon creation, the parameter file is parsed and the time and redshift 27 are calculated and stored in all_outputs. A time units dictionary is 28 instantiated to allow for time outputs to be requested with physical 29 time units. The get_time_series can be used to generate a 30 DatasetSeries object. 31 32 parameter_filename : str 33 The simulation parameter file. 34 find_outputs : bool 35 If True, subdirectories within the GlobalDir directory are 36 searched one by one for datasets. Time and redshift 37 information are gathered by temporarily instantiating each 38 dataset. This can be used when simulation data was created 39 in a non-standard way, making it difficult to guess the 40 corresponding time and redshift information. 41 Default: False. 42 43 Examples 44 -------- 45 >>> import yt 46 >>> es = yt.load_simulation("enzo_tiny_cosmology/32Mpc_32.enzo", "Enzo") 47 >>> es.get_time_series() 48 >>> for ds in es: 49 ... print(ds.current_time) 50 51 """ 52 53 def __init__(self, parameter_filename, find_outputs=False): 54 self.simulation_type = "grid" 55 self.key_parameters = ["stop_cycle"] 56 SimulationTimeSeries.__init__( 57 self, parameter_filename, find_outputs=find_outputs 58 ) 59 60 def _set_units(self): 61 self.unit_registry = UnitRegistry() 62 self.unit_registry.add("code_time", 1.0, dimensions.time) 63 self.unit_registry.add("code_length", 1.0, dimensions.length) 64 if self.cosmological_simulation: 65 # Instantiate EnzoCosmology object for units and time conversions. 66 self.cosmology = EnzoCosmology( 67 self.parameters["CosmologyHubbleConstantNow"], 68 self.parameters["CosmologyOmegaMatterNow"], 69 self.parameters["CosmologyOmegaLambdaNow"], 70 self.parameters.get("CosmologyOmegaRadiationNow", 0.0), 71 0.0, 72 self.parameters["CosmologyInitialRedshift"], 73 unit_registry=self.unit_registry, 74 ) 75 76 self.time_unit = self.cosmology.time_unit.in_units("s") 77 if "h" in self.unit_registry: 78 self.unit_registry.modify("h", self.hubble_constant) 79 else: 80 self.unit_registry.add( 81 "h", self.hubble_constant, dimensions.dimensionless 82 ) 83 # Comoving lengths 84 for my_unit in ["m", "pc", "AU"]: 85 new_unit = f"{my_unit}cm" 86 # technically not true, but should be ok 87 self.unit_registry.add( 88 new_unit, 89 self.unit_registry.lut[my_unit][0], 90 dimensions.length, 91 "\\rm{%s}/(1+z)" % my_unit, 92 prefixable=True, 93 ) 94 self.length_unit = self.quan( 95 self.box_size, "Mpccm / h", registry=self.unit_registry 96 ) 97 else: 98 self.time_unit = self.quan(self.parameters["TimeUnits"], "s") 99 self.length_unit = self.quan(self.parameters["LengthUnits"], "cm") 100 self.box_size = self.length_unit 101 self.domain_left_edge = self.domain_left_edge * self.length_unit 102 self.domain_right_edge = self.domain_right_edge * self.length_unit 103 self.unit_registry.modify("code_time", self.time_unit) 104 self.unit_registry.modify("code_length", self.length_unit) 105 self.unit_registry.add( 106 "unitary", float(self.box_size.in_base()), self.length_unit.units.dimensions 107 ) 108 109 def get_time_series( 110 self, 111 time_data=True, 112 redshift_data=True, 113 initial_time=None, 114 final_time=None, 115 initial_redshift=None, 116 final_redshift=None, 117 initial_cycle=None, 118 final_cycle=None, 119 times=None, 120 redshifts=None, 121 tolerance=None, 122 parallel=True, 123 setup_function=None, 124 ): 125 126 """ 127 Instantiate a DatasetSeries object for a set of outputs. 128 129 If no additional keywords given, a DatasetSeries object will be 130 created with all potential datasets created by the simulation. 131 132 Outputs can be gather by specifying a time or redshift range 133 (or combination of time and redshift), with a specific list of 134 times or redshifts, a range of cycle numbers (for cycle based 135 output), or by simply searching all subdirectories within the 136 simulation directory. 137 138 time_data : bool 139 Whether or not to include time outputs when gathering 140 datasets for time series. 141 Default: True. 142 redshift_data : bool 143 Whether or not to include redshift outputs when gathering 144 datasets for time series. 145 Default: True. 146 initial_time : tuple of type (float, str) 147 The earliest time for outputs to be included. This should be 148 given as the value and the string representation of the units. 149 For example, (5.0, "Gyr"). If None, the initial time of the 150 simulation is used. This can be used in combination with 151 either final_time or final_redshift. 152 Default: None. 153 final_time : tuple of type (float, str) 154 The latest time for outputs to be included. This should be 155 given as the value and the string representation of the units. 156 For example, (13.7, "Gyr"). If None, the final time of the 157 simulation is used. This can be used in combination with either 158 initial_time or initial_redshift. 159 Default: None. 160 times : tuple of type (float array, str) 161 A list of times for which outputs will be found and the units 162 of those values. For example, ([0, 1, 2, 3], "s"). 163 Default: None. 164 initial_redshift : float 165 The earliest redshift for outputs to be included. If None, 166 the initial redshift of the simulation is used. This can be 167 used in combination with either final_time or 168 final_redshift. 169 Default: None. 170 final_redshift : float 171 The latest redshift for outputs to be included. If None, 172 the final redshift of the simulation is used. This can be 173 used in combination with either initial_time or 174 initial_redshift. 175 Default: None. 176 redshifts : array_like 177 A list of redshifts for which outputs will be found. 178 Default: None. 179 initial_cycle : float 180 The earliest cycle for outputs to be included. If None, 181 the initial cycle of the simulation is used. This can 182 only be used with final_cycle. 183 Default: None. 184 final_cycle : float 185 The latest cycle for outputs to be included. If None, 186 the final cycle of the simulation is used. This can 187 only be used in combination with initial_cycle. 188 Default: None. 189 tolerance : float 190 Used in combination with "times" or "redshifts" keywords, 191 this is the tolerance within which outputs are accepted 192 given the requested times or redshifts. If None, the 193 nearest output is always taken. 194 Default: None. 195 parallel : bool/int 196 If True, the generated DatasetSeries will divide the work 197 such that a single processor works on each dataset. If an 198 integer is supplied, the work will be divided into that 199 number of jobs. 200 Default: True. 201 setup_function : callable, accepts a ds 202 This function will be called whenever a dataset is loaded. 203 204 Examples 205 -------- 206 207 >>> import yt 208 >>> es = yt.load_simulation("enzo_tiny_cosmology/32Mpc_32.enzo", "Enzo") 209 >>> es.get_time_series( 210 ... initial_redshift=10, final_time=(13.7, "Gyr"), redshift_data=False 211 ... ) 212 >>> for ds in es: 213 ... print(ds.current_time) 214 >>> es.get_time_series(redshifts=[3, 2, 1, 0]) 215 >>> for ds in es: 216 ... print(ds.current_time) 217 218 """ 219 220 if ( 221 initial_redshift is not None or final_redshift is not None 222 ) and not self.cosmological_simulation: 223 raise InvalidSimulationTimeSeries( 224 "An initial or final redshift has been given for a " 225 + "noncosmological simulation." 226 ) 227 228 if time_data and redshift_data: 229 my_all_outputs = self.all_outputs 230 elif time_data: 231 my_all_outputs = self.all_time_outputs 232 elif redshift_data: 233 my_all_outputs = self.all_redshift_outputs 234 else: 235 raise InvalidSimulationTimeSeries( 236 "Both time_data and redshift_data are False." 237 ) 238 239 if not my_all_outputs: 240 DatasetSeries.__init__(self, outputs=[], parallel=parallel) 241 mylog.info("0 outputs loaded into time series.") 242 return 243 244 # Apply selection criteria to the set. 245 if times is not None: 246 my_outputs = self._get_outputs_by_key( 247 "time", times, tolerance=tolerance, outputs=my_all_outputs 248 ) 249 250 elif redshifts is not None: 251 my_outputs = self._get_outputs_by_key( 252 "redshift", redshifts, tolerance=tolerance, outputs=my_all_outputs 253 ) 254 255 elif initial_cycle is not None or final_cycle is not None: 256 if initial_cycle is None: 257 initial_cycle = 0 258 else: 259 initial_cycle = max(initial_cycle, 0) 260 if final_cycle is None: 261 final_cycle = self.parameters["StopCycle"] 262 else: 263 final_cycle = min(final_cycle, self.parameters["StopCycle"]) 264 265 my_outputs = my_all_outputs[ 266 int( 267 np.ceil(float(initial_cycle) / self.parameters["CycleSkipDataDump"]) 268 ) : (final_cycle / self.parameters["CycleSkipDataDump"]) 269 + 1 270 ] 271 272 else: 273 if initial_time is not None: 274 if isinstance(initial_time, float): 275 my_initial_time = self.quan(initial_time, "code_time") 276 elif isinstance(initial_time, tuple) and len(initial_time) == 2: 277 my_initial_time = self.quan(*initial_time) 278 elif not isinstance(initial_time, unyt_array): 279 raise RuntimeError( 280 "Error: initial_time must be given as a float or " 281 + "tuple of (value, units)." 282 ) 283 elif initial_redshift is not None: 284 my_initial_time = self.cosmology.t_from_z(initial_redshift) 285 else: 286 my_initial_time = self.initial_time 287 288 if final_time is not None: 289 if isinstance(final_time, float): 290 my_final_time = self.quan(final_time, "code_time") 291 elif isinstance(final_time, tuple) and len(final_time) == 2: 292 my_final_time = self.quan(*final_time) 293 elif not isinstance(final_time, unyt_array): 294 raise RuntimeError( 295 "Error: final_time must be given as a float or " 296 + "tuple of (value, units)." 297 ) 298 elif final_redshift is not None: 299 my_final_time = self.cosmology.t_from_z(final_redshift) 300 else: 301 my_final_time = self.final_time 302 303 my_initial_time.convert_to_units("s") 304 my_final_time.convert_to_units("s") 305 my_times = np.array([a["time"] for a in my_all_outputs]) 306 my_indices = np.digitize([my_initial_time, my_final_time], my_times) 307 if my_initial_time == my_times[my_indices[0] - 1]: 308 my_indices[0] -= 1 309 my_outputs = my_all_outputs[my_indices[0] : my_indices[1]] 310 311 init_outputs = [] 312 for output in my_outputs: 313 if os.path.exists(output["filename"]): 314 init_outputs.append(output["filename"]) 315 316 DatasetSeries.__init__( 317 self, outputs=init_outputs, parallel=parallel, setup_function=setup_function 318 ) 319 mylog.info("%d outputs loaded into time series.", len(init_outputs)) 320 321 def _parse_parameter_file(self): 322 """ 323 Parses the parameter file and establishes the various 324 dictionaries. 325 """ 326 327 self.conversion_factors = {} 328 redshift_outputs = [] 329 330 # Let's read the file 331 lines = open(self.parameter_filename).readlines() 332 for line in (l.strip() for l in lines): 333 if "#" in line: 334 line = line[0 : line.find("#")] 335 if "//" in line: 336 line = line[0 : line.find("//")] 337 if len(line) < 2: 338 continue 339 param, vals = (i.strip() for i in line.split("=", 1)) 340 # First we try to decipher what type of value it is. 341 vals = vals.split() 342 # Special case approaching. 343 if "(do" in vals: 344 vals = vals[:1] 345 if len(vals) == 0: 346 pcast = str # Assume NULL output 347 else: 348 v = vals[0] 349 # Figure out if it's castable to floating point: 350 try: 351 float(v) 352 except ValueError: 353 pcast = str 354 else: 355 if any("." in v or "e" in v for v in vals): 356 pcast = float 357 elif v == "inf": 358 pcast = str 359 else: 360 pcast = int 361 # Now we figure out what to do with it. 362 if param.endswith("Units") and not param.startswith("Temperature"): 363 dataType = param[:-5] 364 # This one better be a float. 365 self.conversion_factors[dataType] = float(vals[0]) 366 if param.startswith("CosmologyOutputRedshift["): 367 index = param[param.find("[") + 1 : param.find("]")] 368 redshift_outputs.append( 369 {"index": int(index), "redshift": float(vals[0])} 370 ) 371 elif len(vals) == 0: 372 vals = "" 373 elif len(vals) == 1: 374 vals = pcast(vals[0]) 375 else: 376 vals = np.array([pcast(i) for i in vals if i != "-99999"]) 377 self.parameters[param] = vals 378 self.refine_by = self.parameters["RefineBy"] 379 self.dimensionality = self.parameters["TopGridRank"] 380 if self.dimensionality > 1: 381 self.domain_dimensions = self.parameters["TopGridDimensions"] 382 if len(self.domain_dimensions) < 3: 383 tmp = self.domain_dimensions.tolist() 384 tmp.append(1) 385 self.domain_dimensions = np.array(tmp) 386 self.domain_left_edge = np.array( 387 self.parameters["DomainLeftEdge"], "float64" 388 ).copy() 389 self.domain_right_edge = np.array( 390 self.parameters["DomainRightEdge"], "float64" 391 ).copy() 392 else: 393 self.domain_left_edge = np.array( 394 self.parameters["DomainLeftEdge"], "float64" 395 ) 396 self.domain_right_edge = np.array( 397 self.parameters["DomainRightEdge"], "float64" 398 ) 399 self.domain_dimensions = np.array( 400 [self.parameters["TopGridDimensions"], 1, 1] 401 ) 402 403 if self.parameters["ComovingCoordinates"]: 404 cosmo_attr = { 405 "box_size": "CosmologyComovingBoxSize", 406 "omega_lambda": "CosmologyOmegaLambdaNow", 407 "omega_matter": "CosmologyOmegaMatterNow", 408 "omega_radiation": "CosmologyOmegaRadiationNow", 409 "hubble_constant": "CosmologyHubbleConstantNow", 410 "initial_redshift": "CosmologyInitialRedshift", 411 "final_redshift": "CosmologyFinalRedshift", 412 } 413 self.cosmological_simulation = 1 414 for a, v in cosmo_attr.items(): 415 if v not in self.parameters: 416 raise MissingParameter(self.parameter_filename, v) 417 setattr(self, a, self.parameters[v]) 418 else: 419 self.cosmological_simulation = 0 420 self.omega_lambda = self.omega_matter = self.hubble_constant = 0.0 421 422 # make list of redshift outputs 423 self.all_redshift_outputs = [] 424 if not self.cosmological_simulation: 425 return 426 for output in redshift_outputs: 427 output["filename"] = os.path.join( 428 self.parameters["GlobalDir"], 429 "%s%04d" % (self.parameters["RedshiftDumpDir"], output["index"]), 430 "%s%04d" % (self.parameters["RedshiftDumpName"], output["index"]), 431 ) 432 del output["index"] 433 self.all_redshift_outputs = redshift_outputs 434 435 def _calculate_time_outputs(self): 436 """ 437 Calculate time outputs and their redshifts if cosmological. 438 """ 439 440 self.all_time_outputs = [] 441 if ( 442 self.final_time is None 443 or "dtDataDump" not in self.parameters 444 or self.parameters["dtDataDump"] <= 0.0 445 ): 446 return [] 447 448 index = 0 449 current_time = self.initial_time.copy() 450 dt_datadump = self.quan(self.parameters["dtDataDump"], "code_time") 451 while current_time <= self.final_time + dt_datadump: 452 filename = os.path.join( 453 self.parameters["GlobalDir"], 454 "%s%04d" % (self.parameters["DataDumpDir"], index), 455 "%s%04d" % (self.parameters["DataDumpName"], index), 456 ) 457 458 output = {"index": index, "filename": filename, "time": current_time.copy()} 459 output["time"] = min(output["time"], self.final_time) 460 if self.cosmological_simulation: 461 output["redshift"] = self.cosmology.z_from_t(current_time) 462 463 self.all_time_outputs.append(output) 464 if np.abs(self.final_time - current_time) / self.final_time < 1e-4: 465 break 466 current_time += dt_datadump 467 index += 1 468 469 def _calculate_cycle_outputs(self): 470 """ 471 Calculate cycle outputs. 472 """ 473 474 mylog.warning("Calculating cycle outputs. Dataset times will be unavailable.") 475 476 if ( 477 self.stop_cycle is None 478 or "CycleSkipDataDump" not in self.parameters 479 or self.parameters["CycleSkipDataDump"] <= 0.0 480 ): 481 return [] 482 483 self.all_time_outputs = [] 484 index = 0 485 for cycle in range( 486 0, self.stop_cycle + 1, self.parameters["CycleSkipDataDump"] 487 ): 488 filename = os.path.join( 489 self.parameters["GlobalDir"], 490 "%s%04d" % (self.parameters["DataDumpDir"], index), 491 "%s%04d" % (self.parameters["DataDumpName"], index), 492 ) 493 494 output = {"index": index, "filename": filename, "cycle": cycle} 495 self.all_time_outputs.append(output) 496 index += 1 497 498 def _get_all_outputs(self, find_outputs=False): 499 """ 500 Get all potential datasets and combine into a time-sorted list. 501 """ 502 503 # Create the set of outputs from which further selection will be done. 504 if find_outputs: 505 self._find_outputs() 506 507 elif ( 508 self.parameters["dtDataDump"] > 0 509 and self.parameters["CycleSkipDataDump"] > 0 510 ): 511 mylog.info( 512 "Simulation %s has both dtDataDump and CycleSkipDataDump set.", 513 self.parameter_filename, 514 ) 515 mylog.info( 516 " Unable to calculate datasets. " 517 "Attempting to search in the current directory" 518 ) 519 self._find_outputs() 520 521 else: 522 # Get all time or cycle outputs. 523 if self.parameters["CycleSkipDataDump"] > 0: 524 self._calculate_cycle_outputs() 525 else: 526 self._calculate_time_outputs() 527 528 # Calculate times for redshift outputs. 529 if self.cosmological_simulation: 530 for output in self.all_redshift_outputs: 531 output["time"] = self.cosmology.t_from_z(output["redshift"]) 532 self.all_redshift_outputs.sort(key=lambda obj: obj["time"]) 533 534 self.all_outputs = self.all_time_outputs + self.all_redshift_outputs 535 if self.parameters["CycleSkipDataDump"] <= 0: 536 self.all_outputs.sort(key=lambda obj: obj["time"].to_ndarray()) 537 538 def _calculate_simulation_bounds(self): 539 """ 540 Figure out the starting and stopping time and redshift for the simulation. 541 """ 542 543 if "StopCycle" in self.parameters: 544 self.stop_cycle = self.parameters["StopCycle"] 545 546 # Convert initial/final redshifts to times. 547 if self.cosmological_simulation: 548 self.initial_time = self.cosmology.t_from_z(self.initial_redshift) 549 self.initial_time.units.registry = self.unit_registry 550 self.final_time = self.cosmology.t_from_z(self.final_redshift) 551 self.final_time.units.registry = self.unit_registry 552 553 # If not a cosmology simulation, figure out the stopping criteria. 554 else: 555 if "InitialTime" in self.parameters: 556 self.initial_time = self.quan( 557 self.parameters["InitialTime"], "code_time" 558 ) 559 else: 560 self.initial_time = self.quan(0.0, "code_time") 561 562 if "StopTime" in self.parameters: 563 self.final_time = self.quan(self.parameters["StopTime"], "code_time") 564 else: 565 self.final_time = None 566 if not ("StopTime" in self.parameters or "StopCycle" in self.parameters): 567 raise NoStoppingCondition(self.parameter_filename) 568 if self.final_time is None: 569 mylog.warning( 570 "Simulation %s has no stop time set, stopping condition " 571 "will be based only on cycles.", 572 self.parameter_filename, 573 ) 574 575 def _set_parameter_defaults(self): 576 """ 577 Set some default parameters to avoid problems 578 if they are not in the parameter file. 579 """ 580 581 self.parameters["GlobalDir"] = self.directory 582 self.parameters["DataDumpName"] = "data" 583 self.parameters["DataDumpDir"] = "DD" 584 self.parameters["RedshiftDumpName"] = "RedshiftOutput" 585 self.parameters["RedshiftDumpDir"] = "RD" 586 self.parameters["ComovingCoordinates"] = 0 587 self.parameters["TopGridRank"] = 3 588 self.parameters["DomainLeftEdge"] = np.zeros(self.parameters["TopGridRank"]) 589 self.parameters["DomainRightEdge"] = np.ones(self.parameters["TopGridRank"]) 590 self.parameters["RefineBy"] = 2 # technically not the enzo default 591 self.parameters["StopCycle"] = 100000 592 self.parameters["dtDataDump"] = 0.0 593 self.parameters["CycleSkipDataDump"] = 0.0 594 self.parameters["LengthUnits"] = 1.0 595 self.parameters["TimeUnits"] = 1.0 596 self.parameters["CosmologyOmegaRadiationNow"] = 0.0 597 598 def _find_outputs(self): 599 """ 600 Search for directories matching the data dump keywords. 601 If found, get dataset times py opening the ds. 602 """ 603 604 # look for time outputs. 605 potential_time_outputs = glob.glob( 606 os.path.join( 607 self.parameters["GlobalDir"], f"{self.parameters['DataDumpDir']}*" 608 ) 609 ) 610 self.all_time_outputs = self._check_for_outputs(potential_time_outputs) 611 self.all_time_outputs.sort(key=lambda obj: obj["time"]) 612 613 # look for redshift outputs. 614 potential_redshift_outputs = glob.glob( 615 os.path.join( 616 self.parameters["GlobalDir"], f"{self.parameters['RedshiftDumpDir']}*" 617 ) 618 ) 619 self.all_redshift_outputs = self._check_for_outputs(potential_redshift_outputs) 620 self.all_redshift_outputs.sort(key=lambda obj: obj["time"]) 621 622 self.all_outputs = self.all_time_outputs + self.all_redshift_outputs 623 self.all_outputs.sort(key=lambda obj: obj["time"]) 624 only_on_root(mylog.info, "Located %d total outputs.", len(self.all_outputs)) 625 626 # manually set final time and redshift with last output 627 if self.all_outputs: 628 self.final_time = self.all_outputs[-1]["time"] 629 if self.cosmological_simulation: 630 self.final_redshift = self.all_outputs[-1]["redshift"] 631 632 def _check_for_outputs(self, potential_outputs): 633 """ 634 Check a list of files to see if they are valid datasets. 635 """ 636 637 only_on_root( 638 mylog.info, "Checking %d potential outputs.", len(potential_outputs) 639 ) 640 641 my_outputs = {} 642 llevel = mylog.level 643 # suppress logging as we load every dataset, unless set to debug 644 if llevel > 10 and llevel < 40: 645 mylog.setLevel(40) 646 for my_storage, output in parallel_objects( 647 potential_outputs, storage=my_outputs 648 ): 649 if self.parameters["DataDumpDir"] in output: 650 dir_key = self.parameters["DataDumpDir"] 651 output_key = self.parameters["DataDumpName"] 652 else: 653 dir_key = self.parameters["RedshiftDumpDir"] 654 output_key = self.parameters["RedshiftDumpName"] 655 index = output[output.find(dir_key) + len(dir_key) :] 656 filename = os.path.join( 657 self.parameters["GlobalDir"], 658 f"{dir_key}{index}", 659 f"{output_key}{index}", 660 ) 661 try: 662 ds = load(filename) 663 except (FileNotFoundError, YTUnidentifiedDataType): 664 mylog.error("Failed to load %s", filename) 665 continue 666 my_storage.result = { 667 "filename": filename, 668 "time": ds.current_time.in_units("s"), 669 } 670 if ds.cosmological_simulation: 671 my_storage.result["redshift"] = ds.current_redshift 672 mylog.setLevel(llevel) 673 my_outputs = [ 674 my_output for my_output in my_outputs.values() if my_output is not None 675 ] 676 677 return my_outputs 678 679 def _write_cosmology_outputs(self, filename, outputs, start_index, decimals=3): 680 """ 681 Write cosmology output parameters for a cosmology splice. 682 """ 683 684 mylog.info("Writing redshift output list to %s.", filename) 685 f = open(filename, "w") 686 for q, output in enumerate(outputs): 687 f.write( 688 (f"CosmologyOutputRedshift[%d] = %.{decimals}f\n") 689 % ((q + start_index), output["redshift"]) 690 ) 691 f.close() 692 693 694class EnzoCosmology(Cosmology): 695 def __init__( 696 self, 697 hubble_constant, 698 omega_matter, 699 omega_lambda, 700 omega_radiation, 701 omega_curvature, 702 initial_redshift, 703 unit_registry=None, 704 ): 705 Cosmology.__init__( 706 self, 707 hubble_constant=hubble_constant, 708 omega_matter=omega_matter, 709 omega_lambda=omega_lambda, 710 omega_radiation=omega_radiation, 711 omega_curvature=omega_curvature, 712 unit_registry=unit_registry, 713 ) 714 self.initial_redshift = initial_redshift 715 self.initial_time = self.t_from_z(self.initial_redshift) 716 # time units = 1 / sqrt(4 * pi * G rho_0 * (1 + z_i)**3), 717 # rho_0 = (3 * Omega_m * h**2) / (8 * pi * G) 718 self.time_unit = ( 719 ( 720 1.5 721 * self.omega_matter 722 * self.hubble_constant ** 2 723 * (1 + self.initial_redshift) ** 3 724 ) 725 ** -0.5 726 ).in_units("s") 727 self.time_unit.units.registry = self.unit_registry 728