decrease_nbody.py 8.14 KB
Newer Older
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
1
#!/usr/bin/env python
2 3 4 5
# -----------------------------------------------------------------------
# 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
#     unsio location path
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
20
from __future__ import print_function
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
21
from unsio import *
22 23 24 25 26 27 28 29

import numpy as np
# cmd line
import sys, getopt


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

    if uns.nextFrame(bits):
        unso=CunsOut(out,"gadget2")
        ok,tsnap=uns.getValueF("time") # return snasphot time
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
37
        print ("Snapshot time : ","%.03f"%tsnap)
38 39 40 41
        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
42
            printAndSaveProp(uns,unso,onecomp,take,mmult) # print properties for the component
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
43
        unso.save()
44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# 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
66
def printAndSaveProp(uns,unso,comp,take,mmult):
67 68 69 70 71
    info="""\
    ----------------------------------------------
    Component : [%s]
    ----------------------------------------------
    """
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
72
    print (info % (comp))
73 74 75
    # return a 1D numpy data array with mass
    ok,mass=uns.getArrayF(comp,"mass")
    if ok:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
76
        print ("mass =",mass.size,mass)
77
        mass=keepEveryTake(mass,1,take)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
78
        if mmult==1:
(no author)'s avatar
(no author) committed
79
            mass=mass*take # multiply mass by take
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
80
        print ("mass =",mass.size,mass)
81 82 83 84 85
        unso.setArrayF(comp,"mass",mass) # save mass
        
    # return a 1D numpy data array with pos
    ok,pos=uns.getArrayF(comp,"pos")
    if ok:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
86
        print ("pos =",pos)
87 88 89 90 91 92
        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:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
93
        print ("vel =",vel)
94 95 96
        vel=keepEveryTake(vel,3,take)
        unso.setArrayF(comp,"vel",vel) # save vel

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

   # return a 1D numpy data array with rho
105 106
    ok,rho=uns.getArrayF(comp,"rho")
    if ok:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
107
        print ("rho =",rho)
108 109 110 111 112 113
        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:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
114
        print ("temp =",temp)
115 116 117 118 119 120
        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:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
121
        print ("u =",u.size,u.dtype,u)
122
        u=keepEveryTake(u,1,take)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
123
        print ("u =",u.size,u.dtype,u)
124 125 126 127 128
        unso.setArrayF(comp,"u",u) # save internal energy

    # return a 1D numpy data array with hsml
    ok,hsml=uns.getArrayF(comp,"hsml")
    if ok:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
129
        print ("hsml =",hsml)
130 131 132 133 134 135
        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:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
136
        print ("age =",age)
137 138 139 140 141 142
        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:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
143
        print ("metal =",metal)
144 145 146
        metal=keepEveryTake(metal,1,take)
        unso.setArrayF(comp,"metal",metal) # save mettalicity

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


155 156 157 158

    # return a 1D numy data array with id
    ok,indexes=uns.getArrayI(comp,"id")
    if ok:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
159
        print ("indexes =", indexes)
160 161 162 163 164 165 166 167 168 169 170 171 172
        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
173
    mmult=1
174 175

    try:
(no author)'s avatar
(no author) committed
176
        opts,args=getopt.getopt(argv,"hi:o:k:c:b:m:",["in=","out=","take=","comp=","bits=","mmult="])
177 178

    except getopt.GetoptError:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
179
        print ("\nUnknown parameters, please check ....\n\n")
180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195
        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
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
196
        elif opt in ("-m", "--mmult"):
(no author)'s avatar
(no author) committed
197
	            mmult = int(arg)
198 199 200


    if (infile != '' and out != '' and take !=0):
(no author)'s avatar
(no author) committed
201
        compute(infile,out,comp,take,bits,mmult)
202
    else:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
203
        print ("\nYou must provide input, output files and number of particles to take.....\n\n")
204 205 206 207 208 209 210 211 212 213 214 215
        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
216
    Syntaxe : %s  -i <inputfile> -o <outfile> -c <components> -k <taKe> -b <bits> -m <mmult>
217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232
    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
233 234
                       bits="" (every array, default)    
        mmult      : multiply mass by take (default 1 - yes,  else  0 -  no)
235 236

    """
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
237
    print (help % (prog,prog))
238 239 240 241 242
        
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# main
if __name__ == "__main__":
    main(sys.argv[0:])