Commit acad8431 authored by jclamber's avatar jclamber
Browse files

add uns_select_byid.py

git-svn-id: http://svn.oamp.fr/repos/uns_projects/trunk@109 f264a43e-d52d-4b82-913b-c2bd5215a18a
parent 68f7f3ef
#!/usr/bin/python
from __future__ import print_function
import sys
from py_unsio import * # import py_unsio package (UNSIO)
import numpy as np
import argparse
class snap:
time = None
nbody = None
mass = None
pos = None
vel = None
ids = None
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# commandLine, parse the command line
def commandLine():
time="all"
verbose=False
component="halo"
offset=0
# help
parser = argparse.ArgumentParser(description="Select components by their id",
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
# options
parser.add_argument('simname', help='Simulation name or snapshot')
parser.add_argument('outname', help='Output name snapshot')
parser.add_argument('npoints', help='#points to keep',type=int, default=component)
parser.add_argument('--component', help='Simulation component', default=component)
parser.add_argument('--offset',help='index offset of the isolated component', type=int,default=offset)
#parser.add_argument('--time',help='simulation time', default=time)
parser.add_argument('--verbose',help='verbosity', default=verbose)
# parse
args = parser.parse_args()
# start main funciton
process(args.simname,args.outname,args.npoints,args.component,args.offset,time,args.verbose)
# -----------------------------------------------------
#
def saveSnapshot(snap,idnew,outname):
snap.pos=np.reshape(snap.pos,(-1,3)) # pos reshaped in a 2D array [nbody,3]
snap.vel=np.reshape(snap.vel,(-1,3)) # pos reshaped in a 2D array [nbody,3]
print ("pos size:",snap.pos.size,snap.pos.shape)
# rescale pos
snap.pos = snap.pos[idnew] # keep only good ids
snap.pos = np.reshape(snap.pos,snap.pos.size) # flatten the array (mandatory for unsio)
#print ("pos size:",snap.pos.size,snap.pos.shape)
# rescale vel
snap.vel = snap.vel[idnew] # keep only good ids
snap.vel = np.reshape(snap.vel,snap.vel.size) # flatten the array (mandatory for unsio)
# rescale mass
snap.mass = snap.mass[idnew]
# save data
unso=CunsOut(outname,"nemo") # instantiate object
unso.setValueF("time",snap.time) # save snapshot time
unso.setArrayF("all","pos",snap.pos) # save pos
unso.setArrayF("all","vel",snap.vel) # save vel
unso.setArrayF("all","mass",snap.mass) # save mass
unso.save()
# -----------------------------------------------------
#
def getComponentByIds(snap,offset,npoints):
print("ID min %d / max %d"%(snap.ids.min(),snap.ids.max()),file=sys.stderr)
istart = snap.ids.min()+offset
iend = istart+npoints-1
if (iend-istart > snap.nbody) :
print("Error, iend %d computed > %d bodies"%(iend-istart,snap.nbody),file=sys.stderr)
sys.exit()
idnew = (snap.ids>=istart) & (snap.ids<=iend) # selection criteria
indexes = np.array((idnew.nonzero())).reshape(-1,) # retreive indexes array
print(idnew,indexes)
return True,indexes
# -----------------------------------------------------
# process, is the core function
def process(simname,outname,npoints,component,offset,time,verbose):
#print (simname, " " , time,file=sys.stderr)
# Create a UNSIO object
uns = CunsIn(simname,component,time,verbose)
find=False
bits="mxvI"
while (uns.nextFrame(bits)): # loop while there is something to read
ok,snap.time = uns.getValueF("time")
ok,snap.nbody = uns.getValueI("nsel")
ok,snap.mass = uns.getArrayF(component,"mass")
ok,snap.pos = uns.getArrayF(component,"pos")
ok,snap.vel = uns.getArrayF(component,"vel")
ok,snap.ids = uns.getArrayI(component,"id")
if ok:
find=True
print ( simname, component, snap.time, snap.nbody,snap.mass[0], snap.nbody*snap.mass[0])
break;
if not find:
print ("none")
else:
ok,idnew=getComponentByIds(snap,offset,npoints)
if ok:
saveSnapshot(snap,idnew,outname)
# -----------------------------------------------------
# main program
commandLine() # parse command line
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment