-
Notifications
You must be signed in to change notification settings - Fork 2
/
filter_eagle.py
118 lines (94 loc) · 3.13 KB
/
filter_eagle.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
import h5py
import numpy as np
import pandas as pd
import sys, os
import tqdm
import glob
try:
import read_eagle
except:
print('Need to install read_eagle routine from https://github.com/jchelly/read_eagle')
sys.exit()
try:
from mpi4py import MPI
except:
print('Need to instal mpi4py')
sys.exit()
comm = MPI.COMM_WORLD
comm_rank = comm.Get_rank()
comm_size = comm.Get_size()
#############
# Line arguments
#############
path_to_snapshots = sys.argv[1]
galaxy_list_path = sys.argv[2]
output_path = sys.argv[3]
galaxy_number = sys.argv[4]
####################
IDs = pd.read_csv(galaxy_list_path)
group = IDs['GroupNumber'][int(galaxy_number)]
subgroup = IDs['SubGroupNumber'][int(galaxy_number)]
print('Read in galaxy ID')
snap_dir = path_to_snapshots
snapshot_0 = glob.glob(snap_dir+'/snap_*_z*.0.hdf5')
snap = read_eagle.EagleSnapshot(snapshot_0)
#find box size
try:
start = snap_dir.find('RefL') + 4
end = snap_dir.find('N', start)
box_size = float(snap_dir[start:end])
except:
print('Need to input boxsize')
sys.exit()
snap.select_region(0, box_size * 0.6777, 0, box_size * 0.6777, 0, box_size * 0.6777)
snap.split_selection(comm_rank, comm_size)
print('Selected region')
output_file = h5py.File(output_path+'/galaxy_'+str(galaxy)+'.hdf5', 'w')
input1 = h5py.File(snap_dir+'/'+fname, 'r')
attrs = ['Config',
'Constants',
'HashTable',
'Header',
'Parameters',
'RuntimePars',
'Units']
for key in attrs:
output_file.copy(input1[key], key)
input1.close()
print('Copied Header and others')
gas_attr = snap.datasets(0)
print('Found gas attributes')
star_attr = snap.datasets(4)
print('Found star attributes')
gas_groups = snap.read_dataset(0, 'GroupNumber')
print('got gas groups')
gas_subgroups = snap.read_dataset(0, 'SubGroupNumber')
print('got gas subgroups')
star_groups = snap.read_dataset(4, 'GroupNumber')
print('got star groups')
star_subgroups = snap.read_dataset(4, 'SubGroupNumber')
print('got star subgroups')
mask_g = np.logical_and(gas_groups == group, gas_subgroups == subgroup)
print('got gas mask')
mask_s = np.logical_and(star_groups == group, star_subgroups == subgroup)
print('got star mask')
output_file.create_group('PartType0')
output_file.create_group('PartType4')
print('going into gas attributes')
for attr in tqdm.tqdm(gas_attr):
output_file['PartType0'][attr[1:]] = snap.read_dataset(0, attr)[mask_g]
if attr[1:] == 'Mass':
output_file['PartType0']['Masses'] = snap.read_dataset(0, attr)[mask_g]
print('going into star attributes')
for attr in tqdm.tqdm(star_attr):
output_file['PartType4'][attr[1:]] = snap.read_dataset(4, attr)[mask_s]
if attr[1:] == 'Mass':
output_file['PartType4']['Masses'] = snap.read_dataset(4, attr)[mask_s]
output_file.close()
re_out = h5py.File(output_path+'/galaxy_'+str(galaxy)+'.hdf5', 'r+')
slist = len(re_out['PartType4']['Coordinates'])
glist = len(re_out['PartType0']['Coordinates'])
re_out['Header'].attrs.modify('NumPart_ThisFile', np.array([glist, 0, 0, 0, slist, 0]))
re_out['Header'].attrs.modify('NumPart_Total', np.array([glist, 0, 0, 0, slist, 0]))
re_out['Header'].attrs.modify('NumFilesPerSnapshot', 1)
re_out.close()