import h5py
import glob
import numpy as np

sim_label = { 'aqc4':'C', 'aqd4':'D', 'aqe4':'E' }
max_id = { 'aqc4' : (53967599, 103), 'aqd4' : (41535584 61), 'aqe4' : (9223281412051116774, 80) }
last_snapnum = { 'aqc4' : 127, 'aqd4' : 511, 'aqe4' : 511 }

GAS = 0
DARK_MATTER = DM = 1
STARS = 4

class Struct:
    def __init__(self, *args, **kwds):
        for arg in args:
            self.__dict__.update(arg)
        self.__dict__.update(kwds)
        
    def __repr__(self):
        return str(self.__dict__)

def get_dir(sim_name):
    return "/gpfs/data/Aquila/TO/Aq-%s/400/data" % sim_label[sim_name]
        
def get_data_dir(sim_name):
    return "/gpfs/data/Aquila/halo/data/Aq-%s-4" % sim_label[sim_name]

def get_files(sim_name, snap_num):
    directory = get_dir(sim_name) + "/snapshot_%03d" % snap_num    
    return glob.glob(directory + "/*.hdf5")

def read_header(filename):
    file = h5py.File(filename, "r")
    group = file['/Header']
    header = Struct(dict(group.attrs))
    return header
    
def read_data(sim_name):
    filename = get_data_dir(sim_name) + "/Aq-%s-4_halo_data.hdf5" % sim_label[sim_name]
    file = h5py.File(filename, 'r')
    data = { name : np.array(file['/stars'][name]) for name in file['/stars'].dtype.names }
    return Struct(data)
    
def read_attr(sim_name, snap_num, part_type, attr_name):
    files = get_files(sim_name, snap_num)

    data = []
    for file in files:
        f = h5py.File(file, "r")
        dataset = np.array(f['/PartType%d/%s' % (part_type, attr_name)])
        data.append(dataset)
        
    return np.concatenate(data)
    
def make_unique_star_ids(sim_name, snap_num, ids):
    child_ids = read_attr(sim_name, snap_num, 4,'ChildIDforStars')
    unique_ids = max_id[sim_name][0] + (max_id[sim_name][1] + 1) * (1 + ids) + child_ids
    return unique_ids


if __name__=="__main__":
    #ids = read_attr("aqd4", 127, 4, 'ParticleIDs')
    #print np.where(ids==588263426)

    data = read_data("aqd4")
    idxs = np.where((data.origin==0) & (data.snap_last_in_sub==127))
    ids_snapshot = data.ID[idxs]

    print ids_snapshot

@@ -143,16 +143,17 @@
 41535584
+,
  61), 'a
@@ -1832,34 +1832,587 @@
 ids%0A
-%0A%0Aif __name__==%22__main__%22:
+    %0Adef get_halo_centre(sim_name, snap_num):%0A    fname = get_dir(sim_name) + %22/groups_%2503d/subhalo_tab_%2503d.0%22 %25 (snap_num, snap_num)%0A    file = open(fname, %22rb%22)%0A    %0A    #Load number of groups and subgroups%0A    num_groups = np.fromfile(file, np.int32, 1)%5B0%5D%0A    f.seek(20, 1)%0A    num_subgroups = np.fromfile(file, np.int32, 1)%5B0%5D%0A    f.seek(4 + 76*num_groups + 16*num_subgroups, 1)%0A    %0A    centre = np.fromfile(file, np.float32, 3)%0A    file.close()%0A    return centre%0A%0Aif __name__==%22__main__%22:%0A    sim_name = %22aqd4%22%0A    make_unique_star_ids(sim_name, last_snapnum%5Bsim_name%5D, %0A
 %0A   
@@ -2497,24 +2497,25 @@
 63426)%0A%0A    
+#
 data = read_
@@ -2531,16 +2531,17 @@
 4%22)%0A    
+#
 idxs = n
@@ -2601,16 +2601,17 @@
 7))%0A    
+#
 ids_snap
@@ -2636,16 +2636,17 @@
 s%5D%0A%0A    
+#
 print id

Read halo position from Aquila files