select particles from ids stored in a file

py/mains/uns_glindex_select.py

+#!/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()
+
+    nsave=0
+    if uns.nextFrame(args.bits):
+        unso=CunsOut(args.output,"gadget2")
+        ok,tsnap=uns.getValueF("time") # return snasphot time
+        #print "Snapshot time : ","%.03f"%tsnap
+
+        # 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()
+        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="")
+    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()   
+    
+    # start main funciton
+    process(args)
+
+# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+# main
+if __name__ == "__main__":
+    commandLine()