Source

yt.geometry_tests / art_header.py

import numpy as na
import struct

def read_header(parameter_filename):
    header_struct = [
        ('>i','pad byte'),
        ('>256s','jname'),
        ('>i','pad byte'),
        
        ('>i','pad byte'),
        ('>i','istep'),
        ('>d','t'),
        ('>d','dt'),
        ('>f','aexpn'),
        ('>f','ainit'),
        ('>i','pad byte'),
        
        ('>i','pad byte'),
        ('>f','boxh'),
        ('>f','Om0'),
        ('>f','Oml0'),
        ('>f','Omb0'),
        ('>f','hubble'),
        ('>i','pad byte'),
        
        ('>i','pad byte'),
        ('>i','nextras'),
        ('>i','pad byte'),

        ('>i','pad byte'),
        ('>f','extra1'),
        ('>f','extra2'),
        ('>i','pad byte'),

        ('>i','pad byte'),
        ('>256s','lextra'),
        ('>256s','lextra'),
        ('>i','pad byte'),
        
        ('>i', 'pad byte'),
        ('>i', 'min_level'),
        ('>i', 'max_level'),
        ('>i', 'pad byte'),
        ]
    
    f = open(parameter_filename, "rb")
    header_vals = {}
    for format, name in header_struct:
        size = struct.calcsize(format)
        output = struct.unpack(format, f.read(size))[0]
        header_vals[name] = output
    for to_skip in ['tl','dtl','tlold','dtlold','iSO']:
        _skip_record(f)

    (ncell,) = struct.unpack('>l', _read_record(f))
    # Try to figure out the root grid dimensions
    est = int(na.rint(ncell**(1.0/3.0)))
    # Note here: this is the number of *cells* on the root grid.
    # This is not the same as the number of Octs.
    domain_dimensions = na.ones(3, dtype='int64')*est 

    root_grid_mask_offset = f.tell()
    #_skip_record(f) # iOctCh
    root_cells = domain_dimensions.prod()
    root_iOctCh = _read_frecord(f,'>i')[:root_cells]
    root_iOctCh = root_iOctCh.reshape(domain_dimensions,order='F')
    root_grid_offset = f.tell()
    _skip_record(f) # hvar
    _skip_record(f) # var

    iOctFree, nOct = struct.unpack('>ii', _read_record(f))
    child_grid_offset = f.tell()

    f.close()
    return root_grid_mask_offset,root_grid_offset,child_grid_offset,\
           header_vals['min_level'],header_vals['max_level']


def _skip_record(f):
    s = struct.unpack('>i', f.read(struct.calcsize('>i')))
    f.seek(s[0], 1)
    s = struct.unpack('>i', f.read(struct.calcsize('>i')))

def _read_frecord(f,fmt):
    s1 = struct.unpack('>i', f.read(struct.calcsize('>i')))[0]
    count = s1/na.dtype(fmt).itemsize
    ss = na.fromfile(f,fmt,count=count)
    s2 = struct.unpack('>i', f.read(struct.calcsize('>i')))[0]
    assert s1==s2
    return ss

def _read_record(f,fmt=None):
    s = struct.unpack('>i', f.read(struct.calcsize('>i')))[0]
    ss = f.read(s)
    s = struct.unpack('>i', f.read(struct.calcsize('>i')))
    if fmt is not None:
        return struct.unpack(ss,fmt)
    return ss

def _read_record_size(f):
    pos = f.tell()
    s = struct.unpack('>i', f.read(struct.calcsize('>i')))
    f.seek(pos)
    return s[0]

def _count_art_octs(f, offset, 
                   MinLev, MaxLevelNow):
    level_oct_offsets= [0,]
    level_child_offsets= [0,]
    f.seek(offset)
    nchild,ntot=8,0
    Level = na.zeros(MaxLevelNow+1 - MinLev, dtype='i')
    iNOLL = na.zeros(MaxLevelNow+1 - MinLev, dtype='i')
    iHOLL = na.zeros(MaxLevelNow+1 - MinLev, dtype='i')
    for Lev in xrange(MinLev + 1, MaxLevelNow+1):
        level_oct_offsets.append(f.tell())

        #Get the info for this level, skip the rest
        #print "Reading oct tree data for level", Lev
        #print 'offset:',f.tell()
        Level[Lev], iNOLL[Lev], iHOLL[Lev] = struct.unpack(
           '>iii', _read_record(f))
        #print 'Level %i : '%Lev, iNOLL
        #print 'offset after level record:',f.tell()
        iOct = iHOLL[Lev] - 1
        nLevel = iNOLL[Lev]
        nLevCells = nLevel * nchild
        ntot = ntot + nLevel

        #Skip all the oct hierarchy data
        ns = _read_record_size(f)
        size = struct.calcsize('>i') + ns + struct.calcsize('>i')
        f.seek(f.tell()+size * nLevel)

        level_child_offsets.append(f.tell())
        #Skip the child vars data
        ns = _read_record_size(f)
        size = struct.calcsize('>i') + ns + struct.calcsize('>i')
        f.seek(f.tell()+size * nLevel*nchild)

        #find nhydrovars
        nhydrovars = 8+2
    f.seek(offset)
    return nhydrovars, iNOLL, level_oct_offsets, level_child_offsets

def _read_art_level_info(f, level_oct_offsets,level,root_level=15):
    pos = f.tell()
    f.seek(level_oct_offsets[level])
    #Get the info for this level, skip the rest
    junk, nLevel, iOct = struct.unpack(
       '>iii', _read_record(f))
    
    #fortran indices start at 1
    
    #Skip all the oct hierarchy data
    le     = na.zeros((nLevel,3),dtype='int64')
    fl     = na.ones((nLevel,6),dtype='int64')
    iocts  = na.zeros(nLevel+1,dtype='int64')
    idxa,idxb = 0,0
    chunk = long(1e6) #this is ~111MB for 15 dimensional 64 bit arrays
    left = nLevel
    while left > 0 :
        this_chunk = min(chunk,left)
        idxb=idxa+this_chunk
        data = na.fromfile(f,dtype='>i',count=this_chunk*15)
        data=data.reshape(this_chunk,15)
        left-=this_chunk
        le[idxa:idxb,:] = data[:,1:4]
        fl[idxa:idxb,1] = na.arange(idxa,idxb)
        #pad byte is last, LL2, then ioct right before it
        iocts[idxa:idxb] = data[:,-3] 
        idxa=idxa+this_chunk
    del data
    
    #ioct always represents the index of the next variable
    #not the current, so shift forward one index
    #the last index isn't used
    ioctso = iocts.copy()
    iocts[1:]=iocts[:-1] #shift
    iocts = iocts[:nLevel] #chop off the last index
    iocts[0]=iOct #starting value

    #now correct iocts for fortran indices start @ 1
    iocts = iocts-1

    assert na.unique(iocts).shape[0] == nLevel
    
    #ioct tries to access arrays much larger than le & fl
    #just make sure they appear in the right order, skipping
    #the empty space in between
    idx = na.argsort(iocts)
    
    #now rearrange le & fl in order of the ioct
    le = le[idx]
    fl = fl[idx]

    #left edges are expressed as if they were on 
    #level 15, so no matter what level max(le)=2**15 
    #correct to the yt convention
    #le = le/2**(root_level-1-level)-1

    #try without the -1
    le = le/2**(root_level-2-level)-1

    #now read the hvars and vars arrays
    #we are looking for iOctCh
    #we record if iOctCh is >0, in which it is subdivided
    iOctCh  = na.zeros((nLevel+1,8),dtype='bool')
    
    f.seek(pos)
    return le,fl,nLevel