n6tohdf5.py

#!/usr/local/bin/python

import numpy as np
import struct
import h5py
import sys
import argparse
import matplotlib.pyplot as plt


def read_nbody6(filename, filenamehr = None, filenametidal = None):
    """Read in the nbody6 output file(s).
    
       Keyword arguments:
       filename -- the OUT3 file generated by nbody6
       filenamehr -- the file containing HR data (usually fort.83) (default None)
       filenametidal -- the file containing tidal escapers (usually OUT33) (default None)
    
       Takes one or two files. We at least need the binary dump OUT3.
       If fort.83 is also passed in, then temperature and luminosity 
       are also read in. These can be problematic; if a star is in a 
       close encounter it isn't written to fort.83. At the moment this
       code deals with that situation by using the last recorded temp
       and luminosity for that star.
    """
    # open up OUT3
    f = open(filename,'r')
    
    if filenamehr == None:
        HR = False
    else:
        HR = True
        
    if filenametidal == None:
        tidal = False
    else:
        tidal = True
        
    # open up fort.83 if present
    if HR:
        fhr = open(filenamehr, 'r')
    
    # open up OUT33 if present
    if tidal:
        ft = open(filenametidal, 'r')
    
    # start looping over time and reading the nbody6 data. 
    # itime tracks the number of timesteps outputted.
    itime = -1
    
    # runs on some computers seem to generate extra bytes after the record length.
    # 'extra' and its associated reads below take care of that.
    extra = 0 
    
    byte = f.read(4) # begin header block
    if tidal:
        byte2 = ft.read(4)    
    
    while byte != "":
        blocksize1 = struct.unpack('i',byte)[0] # header block size
        if extra == 1:
            f.read(4)
        if tidal:
            blocksize3 = struct.unpack('i',byte2)[0] # header block size
            if extra == 1:
                ft.read(4)   
                
        itime = itime + 1
        # read the header for this time
  
        ntot = struct.unpack('i',f.read(4))[0] #ntot
        if ntot < 10 and itime == 0:
            # check for a nonsensical ntot value, 
            # which is symptomatic of bizarre extra record padding that occurs on some
            # systems. if needed, include extra reads to deal with these.
            extra = 1
            ntot = struct.unpack('i',f.read(4))[0]
            print 'extra! '+str(extra)
            
        print 'ntot ',ntot
        
        if itime == 0:
            ntotorig = ntot #sometimes there are immediate escapers, so include a 
                                #small buffer here in calculating the max number of stars.
        if tidal and ntot < ntotorig:
            # this triggers tidal reading once we have some escapers
            tidalstart = True
        else: 
            tidalstart = False                           
                                
        print struct.unpack('i',f.read(4))[0] #model
        print struct.unpack('i',f.read(4))[0] #nrun
        nk = struct.unpack('i',f.read(4))[0] #nk

        if tidalstart:
            blocksize3 = struct.unpack('i',byte2)[0] # header block size
            if extra == 1:
                ft.read(4)        
            ntidal = struct.unpack('i',ft.read(4))[0] #number of tidal members
            nkt = struct.unpack('i',ft.read(4))[0] #nk
            print 'ntidal ',ntidal
        else:
            ntidal = 0
        
        if itime == 0:
            ntotorig = ntot + ntidal

        blocksize2 = struct.unpack('i',f.read(4))[0] #end header block size
        if extra == 1:
            f.read(4)
        if tidalstart:
            blocksize4 = struct.unpack('i',ft.read(4))[0] # end header block size
            if extra == 1:
                ft.read(4)         
            
        # check for consistency
        if blocksize1 != blocksize2:
            print 'header trouble! t = '+str(itime)+' '+str(blocksize1)+' '+str(blocksize2)
            sys.exit(1)
        if tidalstart:
            if blocksize3 != blocksize4:
                print 'header trouble tidal! t = '+str(itime)+' '+str(blocksize3)+' '+str(blocksize4)
                sys.exit(2)
                
        # now read the star data
        blocksize1 = struct.unpack('i',f.read(4))[0] #begin data block size
        if extra == 1:
            f.read(4)
        if tidalstart:
            blocksize3 = struct.unpack('i',ft.read(4))[0] # begin data block size
            if extra == 1:
                ft.read(4) 
                        
        alist = []
        for i in range(nk):
            alist.append(struct.unpack('f',f.read(4))[0]) #Sverre's 'as'
        alist.append(ntot - alist[1]) # add number of single stars at the end

        if tidalstart:
            alistt = []
            for i in range(nkt):
                alistt.append(struct.unpack('f',ft.read(4))[0])

        if HR:
            stardata = np.zeros((ntot + ntidal, 10))
        else:
            stardata = np.zeros((ntot + ntidal, 8))
   
        datalist=[] # masses
        for i in range(ntot):
            datalist.append(struct.unpack('f',f.read(4))[0])
        for i in range(ntidal):
            datalist.append(struct.unpack('f',ft.read(4))[0])            
        stardata[:,0] = datalist
       # print stardata[:,0]
         
        datalistx=[] # positions
        datalisty=[]
        datalistz=[]
        for i in range(ntot):           
            datalistx.append(struct.unpack('f',f.read(4))[0])
            datalisty.append(struct.unpack('f',f.read(4))[0])
            datalistz.append(struct.unpack('f',f.read(4))[0]) 
        for i in range(ntidal): # ntidal = 0 if we're not reading tidal data right now
            datalistx.append(struct.unpack('f',ft.read(4))[0])
            datalisty.append(struct.unpack('f',ft.read(4))[0])
            datalistz.append(struct.unpack('f',ft.read(4))[0])                
        stardata[:,1] = datalistx
        stardata[:,2] = datalisty
        stardata[:,3] = datalistz
        
        datalistx=[] # velocities
        datalisty=[]
        datalistz=[]
        for i in range(ntot):           
            datalistx.append(struct.unpack('f',f.read(4))[0])
            datalisty.append(struct.unpack('f',f.read(4))[0])
            datalistz.append(struct.unpack('f',f.read(4))[0])
        for i in range(ntidal):           
            datalistx.append(struct.unpack('f',ft.read(4))[0])
            datalisty.append(struct.unpack('f',ft.read(4))[0])
            datalistz.append(struct.unpack('f',ft.read(4))[0])            
        stardata[:,4] = datalistx
        stardata[:,5] = datalisty
        stardata[:,6] = datalistz    
            
        datalist=[] # names
        for i in range(ntot):
            datalist.append(struct.unpack('i',f.read(4))[0])
        for i in range(ntidal):
            datalist.append(struct.unpack('i',ft.read(4))[0])             
        stardata[:,7] = datalist
 
        blocksize2 = struct.unpack('i',f.read(4))[0] #end data block size
        if extra == 1:
            f.read(4)
        if tidalstart:
            blocksize4 = struct.unpack('i',ft.read(4))[0] # end data block size
            if extra == 1:
                ft.read(4) 
                
        # check for consistency
        if blocksize1 != blocksize2:
            print 'star data trouble! '+itime
            sys.exit(3)
        if tidalstart:
            if blocksize3 != blocksize4:
                print 'star data trouble tidal! t = '+str(itime)+' '+str(blocksize3)+' '+str(blocksize4)
                sys.exit(4)
                
                
        # put into order by name
        stardata = stardata[ stardata[:,7].argsort() ]
        # restrict to names in the range [1, nstars]   
        stardata = stardata[0 < stardata[:,7]]
        # remove binary centers of mass
        stardata = stardata[:ntot - alist[1],:]
        maxname = int(stardata[:,7].max())

        if HR:
            # now find the luminosity and Teff for these stars
            # after itime = 0 the leading line at each time is already read below.
            if itime == 0:
                line = fhr.readline()
                line = line.strip()
                line = line.split()
            
            if len(line) > 0:
                nhr = int(line[0])
            print 'lenline ',len(line), nhr
            # create an array, indexed by name, of (Lum, Teff) pairs
            HRarray = np.zeros((max(ntotorig,maxname),2)) - 1000
            for i in range(nhr+1):
                line = fhr.readline()
                line = line.strip()
                line = line.split()
                if len(line) == 2: # this is the next time
                    break
                if len(line) > 0:
                    hrname = int(line[0])
                    hrlum = float(line[4])
                    hrteff = float(line[6])
                    #use hrname-1 since we're 0 indexed, n6 is 1 indexed
                    HRarray[hrname-1][0]= hrlum 
                    HRarray[hrname-1][1] = hrteff
    
            # initialize the teffs and lums to the last recorded value
            if itime > 0:   
                for i in range(len(stardata)):
                    ind = int(stardata[i,7] - 1)
                # print ind, len(stardata), maxname, len(HRarraylast)
                    stardata[i,8:10] = HRarraylast[ind]
        
            for i in range(len(stardata)):
                ind = int(stardata[i,7] - 1)
                # check here for existence of these values!
                lum = HRarray[ind,0]
                teff = HRarray[ind,1]
                # if it doesn't exist in this snap because it's part of a close 
                # interaction, use the last recorded value
                if lum > -999 and teff > -999:
                    stardata[i,8:10] = HRarray[ind]
        
            HRarraylast = HRarray
            for i in range(len(stardata)):
                ind = int(stardata[i,7] - 1)
                HRarraylast[ind] = stardata[i,8:10]
        
            # split it up into mass, position, velocity, name, teff, lum
            splitarray = np.hsplit(stardata, [1,4,7,8,9])
        else:
            splitarray = np.hsplit(stardata, [1,4,7])
        
        # output as an hdf5 file
        outputname = 'n6snap.'+str(itime).zfill(4)+'.hdf5'
        print outputname
        h5output(outputname, alist, splitarray)   
        
        # next header blocksize for loop control
        byte = f.read(4) # begin header block
        if tidalstart:
            bytet = ft.read(4)
        
                          
    # close OUT3 and other files if present
    f.close()
    if HR:
        fhr.close()
    if tidal:
        ft.close()
      


def h5output(outputfilename, alist, splitarray):
    """output the nbody6 snapshot data in hdf5 format. Plots the first snapshot
       just to make sure it's sensible looking.
    
       Keyword arguments:
       outputfilename -- the name of the hdf5 file to create
       alist -- header data from the nbody6 data. this is sverre's AS list.
       splitarray -- the data that has been parsed from the nbody6 files
    """
    
    floattype = np.float32
    
    # open up the hdf5 file!
    f5 = h5py.File(outputfilename,'w')
    
    # create groups for the file info and data
    stargrp = f5.create_group("Stars")
    headgrp = f5.create_group("Header")
    
    # populate the groups with data
    massdset = stargrp.create_dataset('Masses', data = splitarray[0], dtype=floattype)
    posdset = stargrp.create_dataset('Positions', data = splitarray[1], dtype=floattype)
    veldset = stargrp.create_dataset('Velocities', data = splitarray[2], dtype=floattype)
    namedset = stargrp.create_dataset('Names', data = splitarray[3], dtype=np.int32) 

    if len(splitarray) > 4:
        lumdset = stargrp.create_dataset('Luminosity', data = splitarray[4], dtype=floattype) 
        teffdset = stargrp.create_dataset('Teff', data = splitarray[5], dtype=floattype) 
     
        
    if alist[0] == -1:
        pos = splitarray[1]
        fig = plt.figure(figsize=(5,5))
        ax = fig.add_subplot(1, 1, 1)
        ax.plot(pos[:,0],pos[:,2], 'ro')
        ax.set_aspect(1.0)
        plt.show()   
           
    # header and meta information
    headgrp.attrs["Time"] = alist[0]
    headgrp.attrs["Binaries"] = alist[1]
    headgrp.attrs["Rbar"] = alist[2]
    headgrp.attrs["Mbar"] = alist[3]
    headgrp.attrs["Rtide"] = alist[4]
    headgrp.attrs["Tidal"] = alist[5]
    headgrp.attrs["DensityCenterX"] = alist[6]
    headgrp.attrs["DensityCenterY"] = alist[7]
    headgrp.attrs["DensityCenterZ"] = alist[8]
    headgrp.attrs["ToverTcross"] = alist[9]
    headgrp.attrs["Tscale"] = alist[10]
    headgrp.attrs["Vstar"] = alist[11]
    headgrp.attrs["Rcore"] = alist[12]
    headgrp.attrs["Ncore"] = alist[13]
    headgrp.attrs["Vcore"] = alist[14]
    headgrp.attrs["Rhom"] = alist[15]
    headgrp.attrs["Cmax"] = alist[16]
    headgrp.attrs["Rscale"] = alist[17]
    headgrp.attrs["Rsmin"] = alist[18]
    headgrp.attrs["Dmin1"] = alist[19]
    headgrp.attrs["Nstars"] = alist[20]
    headgrp.attrs["TimeMyr"] = alist[0] * alist[10]
    
    massdset.attrs["Units"] = "Nbody units"
    posdset.attrs["Units"] = "Nbody units"               
    veldset.attrs["Units"] = "Nbody units" 
    if len(splitarray) > 4:
        lumdset.attrs["Units"] = "Log10 solar luminosity"
        teffdset.attrs["Units"] = "Log10 K"
    
    # add conversion factors to the data in each group
    massdset.attrs["to_msun"] = alist[3]
    posdset.attrs["to_pc"] = alist[2]
    veldset.attrs["to_kms"] = alist[11]
    
    # close it up              
    f5.close()       
    return()
        
        
def main():
    parser = argparse.ArgumentParser(description = 'convert nbody6 data to hdf5 snapshots')
    parser.add_argument('starfile', 
        help = 'the OUT3 file to parse')
    parser.add_argument('--hr',
        help = 'the optional HR data (usually fort.83)')
    parser.add_argument('--t', 
        help = 'tidal tail member data (usually OUT33)')
    args = parser.parse_args()
    
    read_nbody6(args.starfile, args.hr, args.t)        
        
        
if __name__ == '__main__':
    main()