#!/usr/bin/env python3 ############################################################################### # This file is part of SWIFT. # Copyright (c) 2021 Mladen Ivkovic (mladen.ivkovic@hotmail.com) # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Lesser General Public License as published # by the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . # ############################################################################## # -------------------------------------------------- # 'Module' containing RT I/O routines for the RT # debugging scheme # -------------------------------------------------- import os import numpy as np import swiftsimio import unyt class RTGasData(object): """ Object to store RT gas particle data of a snapshot """ def __init__(self): self.IDs = None self.coords = None self.h = None self.volumes = None self.PhotonEnergies = None self.PhotonFluxes = None return class RTStarData(object): """ Object to store RT star particle data of a snapshot """ def __init__(self): self.IDs = None self.coords = None self.h = None self.InjectedPhotonEnergy = None return class RTSnapData(object): """ Object to store RT snapshot data """ def __init__(self): self.snapnr = None self.ncells = None self.boxsize = None self.time = None self.has_stars = False self.has_star_debug_data = False self.nstars = None self.npart = None self.stars = RTStarData() self.gas = RTGasData() return class Rundata(object): """ Object to store some global swift data """ def __init__(self): self.units = None self.use_const_emission_rate = False self.has_stars = False # assume we don't have stars, check while reading in self.has_star_debug_data = ( False ) # assume we don't have stars, check while reading in self.ngroups = 0 # photon frequency groups self.const_emission_rates = None self.reduced_speed_of_light = -1.0 self.with_mpi = False return def get_snap_data(prefix="output", skip_snap_zero=False, skip_last_snap=False): """ Finds all prefix_XXXX.hdf5 files and reads the RT data in. Parameters ---------- prefix: str file name prefix of snapshots. Don't include the underscore! skip_snap_zero: bool whether to skip the snapshot prefix_0000.hdf5 skip_last_snap: bool whether to skip the last available snapshot Returns ------- snapdata: list list of RTSnapData objects filled out with actual snapshot data rundata: Rundata object object containing data of the entire run/simulation/compilation """ snapdata = [] ls = os.listdir() hdf5files = [] for f in ls: if f.startswith(prefix + "_") and f.endswith(".hdf5"): hdf5files.append(f) if len(hdf5files) == 0: raise IOError("No " + prefix + "_XXXX.hdf5 files found in this directory") hdf5files.sort() # -------------------------------------- # Get global data from first snapshot # -------------------------------------- rundata = Rundata() firstfile = swiftsimio.load(hdf5files[0]) try: scheme = str(firstfile.metadata.subgrid_scheme["RT Scheme"]) except KeyError: raise ValueError( "These tests only work for the GEAR RT scheme.", "Compile swift --with-rt=GEAR_N", ) if "GEAR" not in scheme: raise ValueError( "These tests only work for the GEAR RT scheme.", "Compile swift --with-rt=GEAR_N", ) ngroups = int(firstfile.metadata.subgrid_scheme["PhotonGroupNumber"][0]) rundata.ngroups = ngroups luminosity_model = firstfile.metadata.parameters["GEARRT:stellar_luminosity_model"] rundata.use_const_emission_rate = luminosity_model.decode("utf-8") == "const" rundata.units = firstfile.units if rundata.use_const_emission_rate: # read emission rate parameter as string emissionstr = firstfile.metadata.parameters[ "GEARRT:const_stellar_luminosities_LSol" ].decode("utf-8") # clean string up if emissionstr.startswith("["): emissionstr = emissionstr[1:] if emissionstr.endswith("]"): emissionstr = emissionstr[:-1] # transform string values to floats with unyts emissions = emissionstr.split(",") emlist = [] for er in emissions: emlist.append(float(er)) rundata.const_emission_rates = unyt.unyt_array(emlist, unyt.L_Sun) if len(rundata.const_emission_rates) != rundata.ngroups: print("Got number of emission rates different from number of groups?") print( "paramfile:", len(rundata.const_emission_rates), "vs", rundata.ngroups, "groups", ) if len(rundata.const_emission_rates) > rundata.ngroups: rundata.const_emission_rates = rundata.const_emission_rates[ : rundata.ngroups ] print( "Only using first", rundata.ngroups, "emission rates:", rundata.const_emission_rates, ) else: quit() else: print( "Didn't detect use of constant stellar emission rates. Proceeding without." ) rundata.reduced_speed_of_light = firstfile.metadata.reduced_lightspeed with_mpi = False if firstfile.metadata.code["MPI library"] != b"Non-MPI version of SWIFT": with_mpi = True rundata.with_mpi = with_mpi # ------------------- # Read in all files # ------------------- for f in hdf5files: snapnrstr = f[len(prefix) + 1 : len(prefix) + 5] snapnr = int(snapnrstr) if skip_snap_zero and snapnr == 0: continue if skip_last_snap and f == hdf5files[-1]: continue # Get snap general data newsnap = RTSnapData() newsnap.snapnr = snapnr data = swiftsimio.load(f) newsnap.boxsize = data.metadata.boxsize newsnap.time = data.metadata.time # Get gas data Gas = RTGasData() Gas.IDs = data.gas.particle_ids Gas.coords = data.gas.coordinates Gas.h = data.gas.smoothing_lengths masses = data.gas.masses densities = data.gas.densities # skip potential div by zero mask = densities == 0.0 masses[mask] = 0.0 densities[mask] = 1.0 Gas.volumes = masses / densities Gas.PhotonEnergies = [ getattr(data.gas.photon_energies, "group" + str(g + 1)) for g in range(rundata.ngroups) ] Gas.PhotonFluxes = [ unyt.uvstack( [ getattr(data.gas.photon_fluxes, "Group" + str(g + 1) + d) for d in ("X", "Y", "Z") ] ).T for g in range(rundata.ngroups) ] newsnap.gas = Gas newsnap.npart = Gas.IDs.shape[0] # Get star data Stars = RTStarData() nstars = 0 has_stars = False has_star_debug_data = False try: Stars.IDs = data.stars.particle_ids Stars.coords = data.stars.coordinates Stars.h = data.stars.smoothing_lengths nstars = Stars.IDs.shape[0] has_stars = True except AttributeError: pass if has_stars: try: inj = np.atleast_2d(data.stars.rtdebug_injected_photon_energy) Stars.InjectedPhotonEnergy = np.reshape( inj, (Stars.IDs.shape[0], ngroups) ) has_star_debug_data = True except AttributeError: pass newsnap.stars = Stars newsnap.nstars = nstars newsnap.has_stars = has_stars newsnap.has_star_debug_data = has_star_debug_data snapdata.append(newsnap) for snap in snapdata: rundata.has_stars = rundata.has_stars or snap.has_stars rundata.has_star_debug_data = ( rundata.has_star_debug_data or snap.has_star_debug_data ) return snapdata, rundata