Source

IRATE-format / irate / enbidbinary.py

Full commit
#!/bin/python
from __future__ import division

import h5py
from numpy import array,empty,zeros,fromstring
import struct,os,sys


def read_enbid_binary(fname):
    """Reads a binary file outputted by EnBiD with ICFormat = 1 -- that is,
		with a Gadget format header and data in binary format, written as float
		data, in the same order as the Gadget file that EnBiD was run on.
		Returns the header, then the data for each group as separate arrays.
	
	:param fname:  The binary file to read.
	"""
    
    print "\nOpening "+fname
    f = open(fname,'rb')

    #First read the header.
    header_size = struct.unpack('<I',f.read(4))[0]

    #number of particles of each type in this file
    nfile = struct.unpack('<6i',f.read(24)) #Number of particles in this file

    masstable = struct.unpack('<6d',f.read(48))  #masses of the particle groups
        
    a = struct.unpack('<d',f.read(8))[0]        #expansion factor
    z = struct.unpack('<d',f.read(8))[0]        #redshift

    flag1 = struct.unpack('<i',f.read(4))[0] #unused
    flag1 = struct.unpack('<i',f.read(4))[0] #unused

    ntot = struct.unpack('<6i',f.read(24))      #total number of particles in the simulation (= nfile if numfiles == 1)
        
    flag3 = struct.unpack('<i',f.read(4))[0]  #unused
    numfiles = struct.unpack('<i',f.read(4))[0]   #number of files in each snapshot
    boxsize = struct.unpack('<d',f.read(8))[0] #Size of the box, if periodic
    omega0 = struct.unpack('<d',f.read(8))[0]  #matter density at z = 0
    omegaL = struct.unpack('<d',f.read(8))[0]  #vacuum energy density at z = 0
    h = struct.unpack('<d',f.read(8))[0] #hubble parameter in units of 100 km/s/Mpc
    
    flag4 = struct.unpack('<i',f.read(4))[0] #unused
    ndim = struct.unpack('<i',f.read(4))[0] #number of dimensions used for density calculation
    flagdensity = struct.unpack('<i',f.read(4))[0]  #indicates that this is an enbid output file
    
    if flagdensity != 1:
        print "Sorry, but this doesn't appear to be an Enbid file."
        sys.exit(1337)

    #That's it for the header, so let's assemble it into a list then skip to the end:
    header = [nfile,masstable,a,z,ntot,numfiles,boxsize,omega0,omegaL,h,ndim]
    
    f.seek(264)     #Seek to the end of the header, plus the size blocks
    
    [ngas,nhalo,ndisk,nbulge,nstar,nbndry] = nfile
    
    rhosize = struct.unpack('<I',f.read(4))[0]
    
    gasrho = fromstring(f.read(4*ngas),dtype='f')
    halorho = fromstring(f.read(4*nhalo),dtype='f')
    diskrho = fromstring(f.read(4*ndisk),dtype='f')
    bulgerho = fromstring(f.read(4*nbulge),dtype='f')
    starrho = fromstring(f.read(4*nstar),dtype='f')
    bndryrho = fromstring(f.read(4*nbndry),dtype='f')
    
    if struct.unpack('<I',f.read(4))[0] != rhosize:
        print "Size at end of block doesn't match that of beginning.  Exiting."
        sys.exit(1337)
        
    f.close()
    
    print "Read data from EnBiD binary file."
        
    return [header,gasrho,halorho,diskrho,bulgerho,starrho,bndryrho]
    
    
def add_enbid(outname,header,gasrho,halorho,diskrho,bulgerho,starrho,bndryrho,
    gastree,gasname,halotree,haloname,disktree,diskname,
    bulgetree,bulgename,startree,starname,bndrytree,bndryname):
	"""Writes the already read data from EnBiD to an existing IRATE file
		under the Analysis group.
	:param outname:  The IRATE file to save the data to.  Must exist already
		and must already have a /Analysis group (and not have a /Analysis/EnBiD
		group
	:param header:  The EnBiD header, in the order that it's provided (which
		is the same as it's provided in Gadget files.
	:param ___rho:  The density arrays for each group of files to be saved.
	:param ___tree:  The group to save the data for each respective particle
		type under.  Must be either dark, star, or gas.  NOTE:  CURRENTLY NOT
		IMPLIMENTED (so it can be anything).  ALL DATA IS SAVED UNDER:
		/Analysis/(Groupname), but I should someday save them under
		/Analysis/(Grouptree)/(Groupname) for consistency sake.
	:param ___name:  The name to give the datasets for each respective particle
		type.
	"""
    
    print "Creating new datasets in existing IRATE file."
    
    f = h5py.File(outname,'a')
    
    enbid = f['Analysis'].create_group('EnBiD')
     
    eheader = enbid.create_group('EnBiDHeader')
    eheader.attrs['NumPart_ThisFile'] = header[0]
    eheader.attrs['MassTable'] = header[1]
    eheader.attrs['Time'] = header[2]
    eheader.attrs['Redshift'] = header[3]
    eheader.attrs['NumPart_Total'] = header[4]
    eheader.attrs['NumFilesPerSnapshot'] = header[5]
    eheader.attrs['BoxSize'] = header[6]
    eheader.attrs['Omega0'] = header[7]
    eheader.attrs['OmegaLambda'] = header[8]
    eheader.attrs['HubbleParam'] = header[9]
    eheader.attrs['NumDimensions'] = header[10] #Number of dimensions used in EnBiD density calculation
    
    if header[0][0] > 0:    #Don't bother even trying anything unless there are gas particles
        enbid.create_dataset(gasname,data=gasrho)            
    if header[0][1] > 0:
        enbid.create_dataset(haloname,data=halorho)
    if header[0][2] > 0:
        enbid.create_dataset(diskname,data=diskrho)
    if header[0][3] > 0:
        enbid.create_dataset(bulgename,data=bulgerho)            
    if header[0][4] > 0:
        enbid.create_dataset(starname,data=starrho)
    if header[0][5] > 0:
        enbid.create_dataset(bndryname,data=bndryrho)
    f.close()