uns_glindex_select.py 8.04 KB
Newer Older
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
#!/usr/bin/env python
# -----------------------------------------------------------------------
# For more information about how to use UNSIO, visit:
# http://projets.lam.fr/projects/unsio/
# -----------------------------------------------------------------------
#  Copyright Jean-Charles LAMBERT (CeSAM)- 2008-2017
#
#  e-mail:   Jean-Charles.Lambert@lam.fr                                      
#  address:  Centre de donneeS Astrophysique de Marseille (CeSAM)         
#            Laboratoire d'Astrophysique de Marseille                          
#            Pole de l'Etoile, site de Chateau-Gombert                         
#            38, rue Frederic Joliot-Curie                                     
#            13388 Marseille cedex 13 France                                   
#            CNRS U.M.R 6110
# -----------------------------------------------------------------------
from __future__ import print_function
# unsio module loading
# ==> do not forget to update PYTHONPATH environment variable with
#     py_unsio location path
from py_unsio import *

import numpy as np
# cmd line
import sys, argparse, os

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# load indexes from file
def loadIds(args):
    if not os.path.isfile(args.id_file):
        print("File <%s> does not exist, aborting...\n"%(args.id_file),file=sys.stderr)
        sys.exit()
    else:
        data=np.genfromtxt(args.id_file,dtype=np.int,skip_header=1)
        print("Size [%d]\n"%(data.size),data)
        return data
        
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# compute
def process(args):

    # load ids
    id_list=loadIds(args)

    if args.select=="all":
        args.select="gas,halo,disk,bulge,stars,bndry"

    uns=CunsIn(args.input,args.select,"all",args.verbose)

    if not uns.isValid():
        print("Unknown UNS date from file <%s>, aborting..."%(args.input),file=sys.stderr)
        sys.exit()

LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
53
54
55
56
57
58
    cpt=0
    while uns.nextFrame(args.bits):
        outfile=args.output
        if cpt==0 and args.tfirst or cpt>0:
            outfile="%s_%05d"%(outfile,cpt)
        unso=CunsOut(outfile,"gadget2")
59
60
61
        ok,tsnap=uns.getValueF("time") # return snasphot time
        #print "Snapshot time : ","%.03f"%tsnap

LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
62
        nsave=0
63
64
65
66
67
68
69
        # loop on all components stored in select variable
        for onecomp in (args.select.split(",")):
            nsave += selectByIds(args,onecomp,id_list,uns,unso)

        if nsave > 0:
            unso.setValueF("time",tsnap) # save snapshot time
            unso.save()
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
70
            cpt += 1
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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
        else:
            print("Nothing to save for file <%s>."%(args.input),file=sys.stderr)
    
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# select particles from their id
# and save them in gadget2 file
def selectByIds(args,comp,id_list,uns,unso):

    info="""\
    ----------------------------------------------
    Component : [%s]
    ----------------------------------------------
    """%(comp)
    print("%s"%(info),file=sys.stderr)

    status=0

    # return a 1D numy data array with id
    ok,indexes=uns.getArrayI(comp,"id")
    if ok:
        good_id=np.in1d(indexes,id_list)
        if indexes[good_id].size == 0:
            print("\n No ids match this component [%s], skipping...\n"%(comp),file=sys.stderr)
            return 0
    else:
        print("\nThere are no IDS in input file....\n",file=sys.stderr)
        return status

    print(good_id)

    # return a 1D numpy data array with mass
    ok,mass=uns.getArrayF(comp,"mass")
    if ok:
        status=1
        print("mass =",mass.size,mass,file=sys.stderr)
        mass=mass[good_id]
        print ("mass =",mass.size,mass,file=sys.stderr)
        unso.setArrayF(comp,"mass",mass) # save mass
        
    # return a 1D numpy data array with pos
    ok,pos=uns.getArrayF(comp,"pos")
    if ok:
        status=1
        print ("pos =",pos,file=sys.stderr)
        pos = np.reshape(pos,(-1,3))
        pos = np.reshape(pos[good_id],pos[good_id].size)
        unso.setArrayF(comp,"pos",pos) # save pos
        
    # return a 1D numpy data array with vel
    ok,vel=uns.getArrayF(comp,"vel")
    if ok:
        status=1
        print ("vel =",vel,file=sys.stderr)
        vel = np.reshape(vel,(-1,3))
        vel = np.reshape(vel[good_id],vel[good_id].size)        
        unso.setArrayF(comp,"vel",vel) # save vel

    # return a 1D numpy data array with acc 
    ok,acc=uns.getArrayF(comp,"acc")
    if ok:
        status=1
        print ("acc =",acc,file=sys.stderr)
        acc = np.reshape(acc,(-1,3))
        acc = np.reshape(acc[good_id],acc[good_id].size)
        unso.setArrayF(comp,"acc",acc) # save vel

   # return a 1D numpy data array with rho
    ok,rho=uns.getArrayF(comp,"rho")
    if ok:
        status=1
        print ("rho =",rho,file=sys.stderr)
        rho=rho[good_id]
        unso.setArrayF(comp,"rho",rho) # save rho

    # return a 1D numpy data array with temperature
    ok,temp=uns.getArrayF(comp,"temp")
    if ok:
        status=1
        print ("temp =",temp,file=sys.stderr)
        temp=temp[good_id]
        unso.setArrayF(comp,"temp",temp) # save temperature

    # return a 1D numpy data array with internal energy (U)
    ok,u=uns.getArrayF(comp,"u")
    if ok:
        status=1
        print ("u =",u.size,u.dtype,u,file=sys.stderr)
        u=u[good_id]
        print ("u =",u.size,u.dtype,u,file=sys.stderr)
        unso.setArrayF(comp,"u",u) # save internal energy

    # return a 1D numpy data array with hsml
    ok,hsml=uns.getArrayF(comp,"hsml")
    if ok:
        status=1
        print ("hsml =",hsml,file=sys.stderr)
        hsml=hsml[good_id]
        unso.setArrayF(comp,"hsml",hsml) # save hsml

    # return a 1D numpy data array with particles age
    ok,age=uns.getArrayF(comp,"age")
    if ok:
        status=1
        print ("age =",age,file=sys.stderr)
        age=age[good_id]
        unso.setArrayF(comp,"age",age) # save age

    # return a 1D numpy data array with mettalicity
    ok,metal=uns.getArrayF(comp,"metal")
    if ok:
        status=1
        print ("metal =",metal,file=sys.stderr)
        metal=metal[good_id]
        unso.setArrayF(comp,"metal",metal) # save mettalicity

    # return a 1D numpy data array with potential 
    ok,pot=uns.getArrayF(comp,"pot")
    if ok:
        status=1
        print ("pot =",pot,file=sys.stderr)
        pot=pot[good_id]
        unso.setArrayF(comp,"pot",pot) # save mettalicity

    # return a 1D numy data array with id
    ok,indexes=uns.getArrayI(comp,"id")
    if ok:
        status=1
        print ("indexes =", indexes,file=sys.stderr)
        indexes=indexes[good_id]
        unso.setArrayI(comp,"id",indexes) # save id

    return status
       
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# commandLine, parse the command line 
def commandLine():
    dbname=None

    # help
    parser = argparse.ArgumentParser(description="Select particles by their ids given from file",
                                    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    # options
    
    parser.add_argument('input',  help='UNS input file')
    parser.add_argument('output', help='UNS output file')
    parser.add_argument('select', help='selected components separated by coma')
    parser.add_argument('id_file', help='input file with particles\'s id (aka glnemo2 index file)')
    parser.add_argument('--bits', help='specify which array you want to proces\nexample:\nbits=\"mxvIU\" (pos,mass,vel,id,internal energy)\nbits=\"\" (every array, default)',default="")
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
219
    parser.add_argument('--tfirst ',help='add trailing number to the first snapshot', dest="tfirst", action="store_true", default=False)
220
221
222
223
224
    parser.add_argument('--dbname',help='UNS database file name', default=dbname)
    parser.add_argument('--verbose',help='verbose mode on', dest="verbose", action="store_true", default=False)

    # parse
    args = parser.parse_args()   
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
225
    print ("tfirst=",args.tfirst)
226
227
228
229
230
231
232
    # start main funciton
    process(args)

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# main
if __name__ == "__main__":
    commandLine()