An utility function to read specific properties for all SAGE galaxies at a given snapshot from the SAGE hdf5 output file

read_sage_hdf5.py

def read_all_sage_galaxies_from_hdf5(h5file, snapnum=63, wanted_fields=None):
    import h5py
    import numpy as np
    if wanted_fields == None:
        wanted_fields = ['StellarMass', 'EjectedMass', 'ColdGas', 'BlackHoleMass', 'BulgeMass', 'CentralMvir', 'SfrDisk', 'Mvir', 'Rvir', 
                         'infallVmax', 'TimeOfLastMajorMerger', 'TimeOfLastMinorMerger', 'SAGETreeIndex', 'Type']

    with h5py.File(h5file, 'r') as hf:
        if hf['Header/Misc'].attrs['sage_data_version'] == b'1.00':
            ncores = hf['Header/Misc'].attrs['num_cores']
        else:
            ncores = hf['Header/Metadata'].attrs['num_cores']
        
        # Check that all the requested fields exist, and also retrieve the
        # dtype for every field. Only check on Core_0 (assuming that data corruption
        # has not occurred)
        field_dtype = []
        expected_ngal = None
        for field_key in wanted_fields:
            
            # Check that the field exists in the file
            hdf5_field = 'Core_0/Snap_{}/{}'.format(snapnum, field_key)
            try:
                field_dtype.append((field_key, hf[hdf5_field].dtype))
            except KeyError:
                msg = 'Error: Could not find field = {} in the galaxy data'\
                    .format(field_key)
                raise KeyError(msg)

            # Ensure that all the requested fields are exactly 1D
            try:
                hf[hdf5_field].shape[1]
            except IndexError:
                # Okay the array is 1D
                pass

            # Check that all the fields have the exact same number of galaxies
            if expected_ngal is None:
                expected_ngal = hf[hdf5_field].shape[0]
            else:
                assert hf[hdf5_field].shape[0] == expected_ngal

        totngals = 0
        # Retrieve the number of galaxies per file (each core writes out its own file)
        for icore in range(ncores):
            base_field_at_snap = 'Core_{}/Snap_{}'.format(icore, snapnum)
            ngals = hf[base_field_at_snap][wanted_fields[0]].shape[0]
            totngals += ngals

        galaxy_dtype = np.dtype(field_dtype)
        galaxy_data = np.empty(totngals, dtype=galaxy_dtype)
        
        for field_key in wanted_fields:
            offset = 0
            for icore in range(ncores):
                hdf5_field = 'Core_{}/Snap_{}/{}'.format(icore, snapnum, field_key)
                ngals = hf[hdf5_field].shape[0]
                galaxy_data[field_key][offset:offset + ngals] = hf[hdf5_field][:]
                offset += ngals
            
            assert offset == totngals
    
    return galaxy_data