filter_galaxy.py

"""
Script for creating subsets of particles from gizmo sims

Authors: 
- Sydney Lower
- Chris Lovell
"""

import h5py
import caesar
import numpy as np
import json

verbose=1

ds = '/orange/narayanan/desika.narayanan/gizmo_runs/simba/m100n1024/snap_m100n1024_078.hdf5'
cs_file = '/orange/narayanan/desika.narayanan/gizmo_runs/simba/m100n1024/m100n1024_078.hdf5'
output_dir = '/orange/narayanan/c.lovell/m100n1024/s50/'


with open('selection_078.json', 'r') as fp:
    gal_idx = list(json.load(fp).keys())
    gal_idx = np.array([int(g[6:]) for g in gal_idx])


cs = caesar.quick_load(cs_file)


## ---- write header info and create groups
with h5py.File(ds, 'r') as input_file:

    header = input_file['Header']

    for galaxy in gal_idx:
        output = '{}/gal_subset_{:05.0f}.h5'.format(output_dir,galaxy)
        with h5py.File(output, 'a') as output_file:
            if 'Header' not in output_file:
                output_file.copy(header, 'Header')

            for group in ['PartType0','PartType1','PartType4','PartType5']:
                if group not in output_file:
                    output_file.create_group(group)



ignore_fields = []
# 'GrackleHI','GrackleHII','GrackleHM','GrackleHeI',
# 'GrackleHeII','GrackleHeIII','Dust_Metallicity','Metallicity']


## ---- write datasets
with h5py.File(ds, 'r') as input_file:

    for ptype in ['PartType0','PartType1','PartType4','PartType5']:

        pidx = int(ptype[8:])   # get particle type index

        for k in input_file[ptype]:   # loop through fields

            if k in ignore_fields:
                if verbose > 1: print(k,'skipped...')
                continue

            if verbose > 0: print(ptype,k)

            # load a given field (the bottleneck)
            temp_dset = input_file[ptype][k][:]
            if verbose > 1: print(temp_dset.shape)

            for galaxy in gal_idx:

                if ptype == 'PartType0':
                    # plist = cs.galaxies[galaxy].glist
                    plist = cs.halos[cs.galaxies[galaxy].parent_halo_index].glist
                elif ptype == 'PartType4':
                    # plist = cs.galaxies[galaxy].slist
                    plist = cs.halos[cs.galaxies[galaxy].parent_halo_index].slist
                elif ptype == 'PartType1':
                    # plist = cs.galaxies[galaxy].slist
                    plist = cs.halos[cs.galaxies[galaxy].parent_halo_index].dmlist
                elif ptype == 'PartType5':
                    # plist = cs.galaxies[galaxy].slist
                    plist = cs.halos[cs.galaxies[galaxy].parent_halo_index].bhlist
                else:
                    if verbose > 0: print('No compatible particle type specified')
                    continue

                output = '{}/gal_subset_{:05.0f}.h5'.format(output_dir,galaxy)
                
                # write to file
                with h5py.File(output, 'a') as output_file:

                    if '%s/%s'%(ptype,k) in output_file:
                        if verbose > 1: print("dataset already exists. replacing...")
                        del output_file[ptype][k]

                    output_file[ptype][k] = temp_dset[plist]

                    temp = output_file['Header'].attrs['NumPart_ThisFile']
                    temp[pidx] = len(plist)
                    output_file['Header'].attrs['NumPart_ThisFile'] = temp

                    temp = output_file['Header'].attrs['NumPart_Total']
                    temp[pidx] = len(plist)
                    output_file['Header'].attrs['NumPart_Total'] = temp

select_gals.py

import numpy as np
import caesar

snap = '078'
fname = '/orange/narayanan/desika.narayanan/gizmo_runs/simba/m100n1024/m100n1024_%s.hdf5'%snap

cs = caesar.quick_load(fname)

sfr = np.array([g.sfr.value for g in cs.galaxies])
centrals = np.array([g.central for g in cs.galaxies])
coods = np.array([g.pos.in_units('code_length').value for g in cs.galaxies])

idx_arr = np.where((sfr > 150) & centrals)[0]

print(idx_arr.shape)

dat = {}
for idx in idx_arr:
    dat['galaxy'+str(int(idx))] = {}
    dat['galaxy'+str(idx)]['snap'+snap] = list(coods[idx])


import json
with open('selection_%s.json'%snap, 'w') as fp:
    json.dump(dat, fp)