Source

IRATE-format / irate / enbid.py

Full commit
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
#!/bin/python
from __future__ import division
import numpy as np

def read_and_write_enbid_binary(fname,outname,snapnum,dset_name="EnBiDDensity",
    t0_name="Gas",t1_name="Dark_Halo",t2_name="Dark_Disk",
    t3_name="Dark_Bulge",t4_name="Star",t5_name="Dark_Boundary",
    units_name="1e10 M_sun per comoving Mpc^3 h^2",units_ar=np.array([6.767991e-31,2,-1])):

    """
    Reads an EnBiD binary file, as produced when run on a Gadget snapshot,
    and saves it into an HDF5 file in the IRATE format.  Currently, no support
    for units because it's not totally clear to me what units EnBiD uses.
    
    :param fname:
        The name of the EnBiD binary file to be read
    :param outname:
        The name of the IRATE file to be created or added to
    :param int snapnum:
        The number to add to 'Snapshot' that becomes the name of
        the group that data is added to.
    :param dset_name:
        The name to be used for the dataset that stores the data from the file
    :param t#_name:
        Determines the name of the group that contains the data from #
        
    # refers to the same names as Gadget2:
    
    * 0 = gas particles
    * 1 = halo particles
    * 2 = disk particles
    * 3 = bulge particles
    * 4 = star particles
    * 5 = bndry particles
    
    """

    import h5py,struct,os
    from numpy import fromstring
    from .core import add_cosmology,check_version

    try:
        snapnum = int(snapnum)
    except ValueError:
        raise ValueError("The number used for the snapshot is not an integer--please check and provide it by hand.")
    print "\nOpening "+fname    
    f = open(fname,'rb')

    #Grab the cosmology to check that it's the same
    f.seek(4+24+48+8+8+4+4+24+4+4+8)   
    #skip headsize, nfile, masstable, a, z, flag1, flag2, ntot, flag3, numfiles, boxsize
    omegaM = 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
    hubble = struct.unpack('<d',f.read(8))[0] #hubble parameter in units of 100 km/s/Mpc
    f.seek(0)
    
    if os.path.isfile(outname):
        new = False
        print "Opening {0} to add data to it.".format(outname)
        irate = h5py.File(outname,'a')
        
        try:
            check_version(irate,update=False)
        except ValueError,e:
            print "\nWARNING:  ",e,"\n" 
        
        if 'Cosmology' not in irate.keys():
            add_cosmology(outname,omegaM=omegaM,omegaL=omegaL,h=hubble,update=True)
        else:
            add_cosmology(outname,omegaM=omegaM,omegaL=omegaL,h=hubble,update=False)
    else:
        new = True
        print "Creating new IRATE file {0}".format(outname)
        irate = h5py.File(outname,'w')
        check_version(irate,update=True)
        add_cosmology(outname,omegaM=omegaM,omegaL=omegaL,h=hubble,update=True)
        
    #Create the file structure
    try:
        snap = irate.create_group("Snapshot{0:05}".format(snapnum))
        print "Created new group for Snapshot{0:05}".format(snapnum)
    except ValueError:
        snap = irate['Snapshot{0:05}'.format(snapnum)]
        print "Adding data to existing group Snapshot{0:05}".format(snapnum)
        
    #Now I need to read the header to know what particle types I have
    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
    flag2 = 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:
        msg = "Sorry, but this doesn't appear to be an Enbid file."
        raise ValueError(msg)
        
    if "Redshift" in snap.attrs.keys():
        if snap.attrs['Redshift'] != z:
            msg = "The existing redshift for group /Snapshot{0:05} doesn't match that in the Gadget file.  Please specify the snapshot group manually.".format(snapnum)
            raise ValueError(msg)
    else:
        snap.attrs['Redshift'] = z
    
    if "ScaleFactor" in snap.attrs.keys():
        if snap.attrs['ScaleFactor'] != a:
            msg = "The existing scale factor for group /Snapshot{0:05} doesn't match that in the Gadget file.  Please specify the snapshot group manually".format(snapnum)
            raise ValueError(msg)
    else:
        snap.attrs['ScaleFactor'] = a
        
    f.seek(264)     #Seek to the end of the header, plus the size blocks
    
    [ngas,nhalo,ndisk,nbulge,nstar,nbndry] = nfile
    
    if "ParticleData" in snap.keys():
        pdata = snap['ParticleData']
    else:
        pdata = snap.create_group('ParticleData')
        
    if dset_name+'unitname' in pdata.attrs.keys():
        if pdata.attrs[dset_name+'unitname'] != units_name or (pdata.attrs[dset_name+'unitcgs'] != units_ar).all():
            msg = "Density units in existing file are different than those provided.\n"
            raise ValueError(msg)
    else:
        pdata.attrs[dset_name+'unitname'] = units_name
        pdata.attrs[dset_name+'unitcgs'] = units_ar            
        
    rhosize = struct.unpack('<I',f.read(4))[0]
    
    snapname = "Snapshot{0:05}".format(snapnum)
    
    if ngas:
        if t0_name in pdata.keys():
            try:
                dset = pdata[t0_name].create_dataset(dset_name,data=fromstring(f.read(4*ngas),dtype='f'))
            except KeyError:
                msg = "{0} already exists under {1} in {2}".format(dset_name,t0_name,snapname)
                raise KeyError(msg)
        else:
            pdata.create_group(t0_name)
            pdata[t0_name].create_dataset(dset_name,data=fromstring(f.read(4*ngas),dtype='f'))
    
    if nhalo:
        if t1_name in pdata.keys():
            try:
                pdata[t1_name].create_dataset(dset_name,data=fromstring(f.read(4*nhalo),dtype='f'))
            except KeyError:
                msg = "{0} already exists under {1} in {2}".format(dset_name,t1_name,snapname)
                raise KeyError(msg)
        else:
            pdata.create_group(t1_name)
            pdata[t1_name].create_dataset(dset_name,data=fromstring(f.read(4*nhalo),dtype='f'))
    
    if ndisk:
        if t2_name in pdata.keys():
            try:
                pdata[t2_name].create_dataset(dset_name,data=fromstring(f.read(4*ndisk),dtype='f'))
            except KeyError:
                msg = "{0} already exists under {1} in {2}".format(dset_name,t2_name,snapname)
                raise KeyError(msg)
        else:
            pdata.create_group(t2_name)
            pdata[t2_name].create_dataset(dset_name,data=fromstring(f.read(4*ndisk),dtype='f'))
    
    if nbulge:
        if t3_name in pdata.keys():
            try:
                pdata[t3_name].create_dataset(dset_name,data=fromstring(f.read(4*nbulge),dtype='f'))
            except KeyError:
                msg = "{0} already exists under {1} in {2}".format(dset_name,t3_name,snapname)
                raise KeyError(msg)
        else:
            pdata.create_group(t3_name)
            pdata[t3_name].create_dataset(dset_name,data=fromstring(f.read(4*nbulge),dtype='f'))
    
    if nstar:
        if t4_name in pdata.keys():
            try:
                pdata[t4_name].create_dataset(dset_name,data=fromstring(f.read(4*nstar),dtype='f'))
            except KeyError:
                msg = "{0} already exists under {1} in {2}".format(dset_name,t4_name,snapname)
                raise KeyError(msg)
        else:
            pdata.create_group(t4_name)
            pdata[t4_name].create_dataset(dset_name,data=fromstring(f.read(4*nstar),dtype='f'))
    
    if nbndry:
        if t5_name in pdata.keys():
            try:
                pdata[t5_name].create_dataset(dset_name,data=fromstring(f.read(4*nbndry),dtype='f'))
            except KeyError:
                msg = "{0} already exists under {1} in {2}".format(dset_name,t5_name,snapname)
                raise KeyError(msg)
        else:
            pdata.create_group(t5_name)
            pdata[t5_name].create_dataset(dset_name,data=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()
    irate.close()
    
    print "Read data from EnBiD file {0} and added it to IRATE file {1}.".format(fname,outname)
    
    

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.
    
    .. warning::
        This function hasn't been maintained and updated to the new format yet,
        so it's not yet ready for use.
    """
    from numpy import array,empty,zeros,fromstring
    import struct,os,sys
    
    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
    flag2 = 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.
    :param ___name:  
        The name to give the datasets for each respective particle
        type.
        
    .. warning::
        This function hasn't been maintained and updated to the new format yet,
        so it's not yet ready for use.
    """
    
    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
    
    trees = [gastree,halotree,disktree,bulgetree,startree,bndrytree]
    if gas in trees:
        gtree = enbid.create_group('Gas')
    if dark in trees:
        dtree = enbid.create_group('Dark')
    if star in trees:
        stree = enbid.create_group('Star')
    
    if header[0][0] > 0:    #Don't bother even trying anything unless there are gas particles
        if gastree == 'gas':
            print "Saving gas data under '/Analysis/Gas/'"+gasname
            gtree.create_dataset(gasname,data=gasrho)
        elif gastree == 'dark':
            print "Saving gas data under '/Analysis/Dark'"+gasname
            dtree.create_dataset(gasname,data=gasrho)
        elif gastree == 'star':
            print "Saving gas data under '/Analysis/Star/'"+gasname
            stree.create_dataset(gasname,data=gasrho)
        else:
            print "I shouldn't be here!  Somehow you got away with trying to place the particle data for the gas group into a nonexistant tree."
            sys.exit(1337)
    if header[0][1] > 0:
        if halotree == 'gas':
            print "Saving halo data under '/Analysis/Gas/'"+haloname
            gtree.create_dataset(haloname,data=halorho)
        elif halotree == 'dark':
            print "Saving halo data under '/Analysis/Dark'"+haloname
            dtree.create_dataset(haloname,data=halorho)
        elif halotree == 'star':
            print "Saving halo data under '/Analysis/Star/'"+haloname
            stree.create_dataset(haloname,data=halorho)
        else:
            print "I shouldn't be here!  Somehow you got away with trying to place the particle data for the halo group into a nonexistant tree."
            sys.exit(1337)
    if header[0][2] > 0:
        if disktree == 'gas':
            print "Saving disk data under '/Analysis/Gas/'"+diskname
            gtree.create_dataset(diskname,data=diskrho)
        elif disktree == 'dark':
            print "Saving disk data under '/Analysis/Dark'"+diskname
            dtree.create_dataset(diskname,data=diskrho)
        elif disktree == 'star':
            print "Saving disk data under '/Analysis/Star/'"+diskname
            stree.create_dataset(diskname,data=diskrho)
        else:
            print "I shouldn't be here!  Somehow you got away with trying to place the particle data for the disk group into a nonexistant tree."
            sys.exit(1337)
    if header[0][3] > 0:
        if bulgetree == 'gas':
            print "Saving bulge data under '/Analysis/Gas/'"+bulgename
            gtree.create_dataset(bulgename,data=bulgerho)
        elif bulgetree == 'dark':
            print "Saving bulge data under '/Analysis/Dark'"+bulgename
            dtree.create_dataset(bulgename,data=bulgerho)
        elif bulgetree == 'star':
            print "Saving bulge data under '/Analysis/Star/'"+bulgename
            stree.create_dataset(bulgename,data=bulgerho)
        else:
            print "I shouldn't be here!  Somehow you got away with trying to place the particle data for the bulge group into a nonexistant tree."
            sys.exit(1337)
    if header[0][4] > 0:
        if startree == 'gas':
            print "Saving star data under '/Analysis/Gas/'"+starname
            gtree.create_dataset(starname,data=starrho)
        elif startree == 'dark':
            print "Saving star data under '/Analysis/Dark'"+starname
            dtree.create_dataset(starname,data=starrho)
        elif startree == 'star':
            print "Saving star data under '/Analysis/Star/'"+starname
            stree.create_dataset(starname,data=starrho)
        else:
            print "I shouldn't be here!  Somehow you got away with trying to place the particle data for the star group into a nonexistant tree."
            sys.exit(1337)
    if header[0][5] > 0:
        if bndrytree == 'gas':
            print "Saving bndry data under '/Analysis/Gas/'"+bndryname
            gtree.create_dataset(bndryname,data=bndryrho)
        elif bndrytree == 'dark':
            print "Saving bndry data under '/Analysis/Dark'"+bndryname
            dtree.create_dataset(bndryname,data=bndryrho)
        elif bndrytree == 'star':
            print "Saving bndry data under '/Analysis/Star/'"+bndryname
            stree.create_dataset(bndryname,data=bndryrho)
        else:
            print "I shouldn't be here!  Somehow you got away with trying to place the particle data for the boundary group into a nonexistant tree."
            sys.exit(1337)
    f.close()