add_ic_units.py

from swiftsimio import Writer
from swiftsimio.units import cosmo_units

import unyt
import numpy as np
import h5py
import glob
  
files = glob.glob('ics*.hdf5')

h = 0.6711

for _f in files:
    with h5py.File(_f, 'r') as hf:
        coods_gas = hf['PartType0/Coordinates'][:]
        pids_gas = hf['PartType0/ParticleIDs'][:]
        vels_gas = hf['PartType0/Velocities'][:]

        coods = hf['PartType1/Coordinates'][:]
        pids = hf['PartType1/ParticleIDs'][:]
        vels = hf['PartType1/Velocities'][:]

        dm_mass = hf['Header'].attrs['MassTable'][1]
        gas_mass = hf['Header'].attrs['MassTable'][0]
        scalefactor = hf['Header'].attrs['Time']
        numpart_total = hf['Header'].attrs['NumPart_Total']
        numfiles_snap = hf['Header'].attrs['NumFilesPerSnapshot']


    boxsize = (25 / h) * unyt.Mpc
    
    # Generate object. cosmo_units corresponds to default Gadget-oid units
    # of 10^10 Msun, Mpc, and km/s
    x = Writer(cosmo_units, boxsize, scale_factor=scalefactor, 
               extra_header={'Time': scalefactor})
    
    x.dark_matter.coordinates = (coods / h) * unyt.Mpc
    x.dark_matter.velocities = vels * (unyt.km / unyt.s)
    x.dark_matter.masses = (dm_mass * np.ones(len(pids)) / h) * (1e10 * unyt.msun)
    x.dark_matter.particle_ids = pids
    
    x.gas.coordinates = (coods_gas / h) * unyt.Mpc
    x.gas.velocities = vels_gas * (unyt.km / unyt.s)
    x.gas.masses = (gas_mass * np.ones(len(pids_gas)) / h) * (1e10 * unyt.msun)
    x.gas.particle_ids = pids_gas
    
    # Generate internal energy corresponding to 273 K
    x.gas.internal_energy = (
        np.ones(len(pids_gas), dtype=float) * (273 * unyt.kb * unyt.K) / (1e6 * unyt.msun)
    )
    
    # # Generate initial guess for smoothing lengths based on MIPS
    x.gas.generate_smoothing_lengths(boxsize=boxsize, dimension=3)
    
    # If IDs are not present, this automatically generates
    x.write("new_%s"%_f)

    with h5py.File("new_%s"%_f, 'a') as hf:
        hf['Header'].attrs['NumPart_Total'] = numpart_total
        hf['Header'].attrs['NumFilesPerSnapshot'] = numfiles_snap