decrease_nbody.py 8.04 KB
Newer Older
1 2 3 4 5
#!/usr/bin/python
# -----------------------------------------------------------------------
# For more information about how to use UNSIO, visit:
# http://projets.lam.fr/projects/unsio/
# -----------------------------------------------------------------------
LAMBERT Jean-charles's avatar
updates  
LAMBERT Jean-charles committed
6
#  Copyright Jean-Charles LAMBERT (CeSAM)- 2008-2016
7 8 9 10 11 12 13 14 15 16 17 18
#
#  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
# -----------------------------------------------------------------------

# unsio module loading
# ==> do not forget to update PYTHONPATH environment variable with
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
19 20
#     unsio location path
from unsio import *
21 22 23 24 25 26 27 28

import numpy as np
# cmd line
import sys, getopt


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# compute
(no author)'s avatar
(no author) committed
29
def compute(file,out,comp,take,bits,mmult):
30
    compchk="gas,halo,disk,bulge,stars,bndry"
(no author)'s avatar
up  
(no author) committed
31
    uns=CunsIn(file,comp,"all")
32 33 34 35 36 37 38 39 40

    if uns.nextFrame(bits):
        unso=CunsOut(out,"gadget2")
        ok,tsnap=uns.getValueF("time") # return snasphot time
        print "Snapshot time : ","%.03f"%tsnap
        unso.setValueF("time",tsnap) # save snapshot time

        # loop on all components stored in comp variable
        for onecomp in (compchk.split(",")):
(no author)'s avatar
(no author) committed
41
            printAndSaveProp(uns,unso,onecomp,take,mmult) # print properties for the component
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
42
        unso.save()
43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# keepEveryTake:
# from an input numpy array keep one particle every take
# input array MUST be a 1d array
# dim match to the array stride
def keepEveryTake(array,dim,take):
    if (dim==1):             # 1D array arrangement
        array=array[::take]  # keep one particle every take
        array=array*1 # we *1 to have a contiguous array
    else:                    # > 1D array arrangment
        # first reshape in n X "dim", then keep every take
        array=np.reshape(array,(-1,dim))[::take]
        # reshape in 1D array
        array=np.reshape(array,array.size)
    return array
    
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# printAndSaveProp
# print properties for the component comp given as argument
# keep one particle over tale
# and save them in gadget2 file
(no author)'s avatar
(no author) committed
65
def printAndSaveProp(uns,unso,comp,take,mmult):
66 67 68 69 70 71 72 73 74 75 76
    info="""\
    ----------------------------------------------
    Component : [%s]
    ----------------------------------------------
    """
    print info % (comp)
    # return a 1D numpy data array with mass
    ok,mass=uns.getArrayF(comp,"mass")
    if ok:
        print "mass =",mass.size,mass
        mass=keepEveryTake(mass,1,take)
(no author)'s avatar
(no author) committed
77 78
	if mmult==1:
            mass=mass*take # multiply mass by take
79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
        print "mass =",mass.size,mass
        unso.setArrayF(comp,"mass",mass) # save mass
        
    # return a 1D numpy data array with pos
    ok,pos=uns.getArrayF(comp,"pos")
    if ok:
        print "pos =",pos
        pos=keepEveryTake(pos,3,take)
        unso.setArrayF(comp,"pos",pos) # save pos
        
    # return a 1D numpy data array with vel
    ok,vel=uns.getArrayF(comp,"vel")
    if ok:
        print "vel =",vel
        vel=keepEveryTake(vel,3,take)
        unso.setArrayF(comp,"vel",vel) # save vel

LAMBERT Jean-charles's avatar
updates  
LAMBERT Jean-charles committed
96 97 98 99
    # return a 1D numpy data array with acc 
    ok,acc=uns.getArrayF(comp,"acc")
    if ok:
        print "acc =",acc
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
100
        acc=keepEveryTake(acc,3,take)
LAMBERT Jean-charles's avatar
updates  
LAMBERT Jean-charles committed
101 102 103
        unso.setArrayF(comp,"acc",acc) # save vel

   # return a 1D numpy data array with rho
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
    ok,rho=uns.getArrayF(comp,"rho")
    if ok:
        print "rho =",rho
        rho=keepEveryTake(rho,1,take)
        unso.setArrayF(comp,"rho",rho) # save rho

    # return a 1D numpy data array with temperature
    ok,temp=uns.getArrayF(comp,"temp")
    if ok:
        print "temp =",temp
        temp=keepEveryTake(temp,1,take)
        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:
        print "u =",u.size,u.dtype,u
        u=keepEveryTake(u,1,take)
        print "u =",u.size,u.dtype,u
        unso.setArrayF(comp,"u",u) # save internal energy

    # return a 1D numpy data array with hsml
    ok,hsml=uns.getArrayF(comp,"hsml")
    if ok:
        print "hsml =",hsml
        hsml=keepEveryTake(hsml,1,take)
        unso.setArrayF(comp,"hsml",hsml) # save hsml

    # return a 1D numpy data array with particles age
    ok,age=uns.getArrayF(comp,"age")
    if ok:
        print "age =",age
        age=keepEveryTake(age,1,take)
        unso.setArrayF(comp,"age",age) # save age

    # return a 1D numpy data array with mettalicity
    ok,metal=uns.getArrayF(comp,"metal")
    if ok:
        print "metal =",metal
        metal=keepEveryTake(metal,1,take)
        unso.setArrayF(comp,"metal",metal) # save mettalicity

LAMBERT Jean-charles's avatar
updates  
LAMBERT Jean-charles committed
146 147 148 149
    # return a 1D numpy data array with potential 
    ok,pot=uns.getArrayF(comp,"pot")
    if ok:
        print "pot =",pot
LAMBERT Jean-charles's avatar
updates  
LAMBERT Jean-charles committed
150
        pot=keepEveryTake(pot,1,take)
LAMBERT Jean-charles's avatar
updates  
LAMBERT Jean-charles committed
151 152 153
        unso.setArrayF(comp,"pot",pot) # save mettalicity


154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171

    # return a 1D numy data array with id
    ok,indexes=uns.getArrayI(comp,"id")
    if ok:
        print "indexes =", indexes
        indexes=keepEveryTake(indexes,1,take)
        unso.setArrayI(comp,"id",indexes) # save id

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# main
def main(argv):
    prog = argv.pop(0) # program name
    infile=''
    out=''
    take=0
    comp='all'
    times='all'
    bits=''
(no author)'s avatar
(no author) committed
172
    mmult=1
173 174

    try:
(no author)'s avatar
(no author) committed
175
        opts,args=getopt.getopt(argv,"hi:o:k:c:b:m:",["in=","out=","take=","comp=","bits=","mmult="])
176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194

    except getopt.GetoptError:
        print "\nUnknown parameters, please check ....\n\n"
        printHelp(prog)
        sys.exit()
    for opt, arg in opts:
        if opt == '-h':
            printHelp(prog)
            sys.exit()
        elif opt in ("-i", "--in"):
            infile = arg
        elif opt in ("-o", "--out"):
            out = arg
        elif opt in ("-c", "--comp"):
            comp = arg
        elif opt in ("-k", "--take"):
            take = int(arg)
        elif opt in ("-b", "--bits"):
            bits = arg
(no author)'s avatar
(no author) committed
195 196
	elif opt in ("-m", "--mmult"):
	            mmult = int(arg)
197 198 199


    if (infile != '' and out != '' and take !=0):
(no author)'s avatar
(no author) committed
200
        compute(infile,out,comp,take,bits,mmult)
201 202 203 204 205 206 207 208 209 210 211 212 213 214
    else:
        print "\nYou must provide input, output files and number of particles to take.....\n\n"
        printHelp(prog)
        
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# printHelp
def printHelp(prog):
    help= """\
    --------------------------------------------------
    From an input UNS input file:
    - Keep one particles every "take" particles, and
    - multiply masses by "take" and save into a gadget2 file
    --------------------------------------------------
    
(no author)'s avatar
(no author) committed
215
    Syntaxe : %s  -i <inputfile> -o <outfile> -c <components> -k <taKe> -b <bits> -m <mmult>
216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231
    Example : %s  -i gtr118_1912 -o myfile.g2 -c all -k 10
    
    Notes :
        inputfile  : UNS input snapshot

        outfile    : gadget2 out filename
        
        components : specify just one or a coma separeted list of components
                     among => disk,gas,stars,halo,bulge,bndry or all
                     exemple : -c disk,stars
                     
        take       : number of particles to take

        bits       : you specify which array you want to process
                     example:
                       bits="mxvIU" (pos,mass,vel,id,internal energy)
(no author)'s avatar
(no author) committed
232 233
                       bits="" (every array, default)    
        mmult      : multiply mass by take (default 1 - yes,  else  0 -  no)
234 235 236 237 238 239 240 241

    """
    print help % (prog,prog)
        
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# main
if __name__ == "__main__":
    main(sys.argv[0:])