Source

yt.geometry_tests / test_oct.py

Full commit
from yt.geometry.oct_container import \
    OctreeContainer
import struct
import numpy as na
import pprint
import os
import time
from yt.utilities._amr_utils.geometry_utils import \
    get_hilbert_points, \
    get_hilbert_indices
from yt.geometry.selection_routines import \
    SphereSelector

f1 = open("output_00007/amr_00007.out00001")
header = ( ('ncpu', 1, 'i'),
           ('ndim', 1, 'i'),
           ('nx', 3, 'i'),
           ('nlevelmax', 1, 'i'),
           ('ngridmax', 1, 'i'),
           ('nboundary', 1, 'i'),
           ('ngrid_current', 1, 'i'),
           ('boxlen', 1, 'd'),
           ('nout', 3, 'I')
          )

def read_attrs(f, attrs):
    vv = {}
    net_format = "="
    for a, n, t in attrs:
        net_format += "".join(["I"] + ([t] * n) + ["I"])
    size = struct.calcsize(net_format)
    vals = list(struct.unpack(net_format, f.read(size)))
    vv = {}
    for a, b, n in attrs:
        s1 = vals.pop(0)
        v = [vals.pop(0) for i in range(b)]
        s2 = vals.pop(0)
        if s1 != s2:
            size = struct.calcsize("=I" + "".join(b*[n]) + "I")
            print "S1 = %s ; S2 = %s ; %s %s %s = %s" % (
                    s1, s2, a, b, n, size)
            raise RuntimeError
        assert(s1 == s2)
        print s1, s2, a, b, n
        if b == 1: v = v[0]
        vv[a] = v
    return vv

def read_vector(f, d):
    fmt = "=I"
    ss = struct.unpack(fmt, f .read(struct.calcsize(fmt)))[0]
    ds = struct.calcsize("=%s" % d)
    if ss % ds != 0:
        print "fmt = '%s' ; ss = %s ; ds = %s" % (fmt, ss, ds)
        raise RuntimeError
    count = ss / ds
    #print "READING %s (%s %s)" % (count, ss, ds)
    tr = na.fromfile(f, d, count)
    vec = struct.unpack(fmt, f.read(struct.calcsize(fmt)))
    assert(vec[-1] == ss)
    global examine
    examine = tr
    return tr

def skip(f, n = 1):
    for i in range(n):
        fmt = "=I"
        ss = struct.unpack(fmt, f.read(struct.calcsize(fmt)))[0]
        #print "SKIPPING %s byte record" % (ss)
        f.seek(ss + struct.calcsize("=I"), os.SEEK_CUR)

hvals = {}
hvals.update(read_attrs(f1, header))

noutput, iout, ifout = hvals['nout']

next_set = ( ('tout', noutput, 'd'),
             ('aout', noutput, 'd'),
             ('t', 1, 'd'),
             ('dtold', hvals['nlevelmax'], 'd'),
             ('dtnew', hvals['nlevelmax'], 'd'),
             ('nstep',  2, 'i'),
             ('stat', 3, 'd'),
             ('cosm', 7, 'd'),
             ('timing', 5, 'd'),
             ('mass_sph', 1, 'd') )

hvals.update(read_attrs(f1, next_set))

tree_header = ( ('headl', hvals['nlevelmax'], 'i'),
                ('taill', hvals['nlevelmax'], 'i'),
                ('numbl', hvals['nlevelmax'], 'i'),
              )
hvals.update(read_attrs(f1, tree_header))
skip(f1)

for i in range(hvals['nboundary']):
    skip(f1, 2) # headb, tailb
    skip(f1) # grid boundary

next_set = ( ('free_mem', 5, 'i'), )
hvals.update(read_attrs(f1, next_set))

hvals['ordering'] = "".join(read_vector(f1, 'c')).strip()
skip(f1, 4) 

pointers = {}
begin = f1.tell()
f1.seek(0, os.SEEK_END)
end = f1.tell()
f1.seek(begin)

print hvals['nx']
otc = OctreeContainer(hvals['nx'],
    (0.0, 0.0, 0.0), 
    (hvals['boxlen'], hvals['boxlen'], hvals['boxlen']))

mi = 1e30
ma = -1e30

loki = raw_input("Connected?")

for ilvl in range(hvals['nlevelmax']):
    cur = f1.tell()
    if cur == end: break
    ng = hvals['numbl'][ilvl]
    print "HANDLING LEVEL %s (%s), %s from end" % (ilvl, ng, end - cur)
    vv = pointers[ilvl] = {}
    ind = na.empty(ng, dtype='int64')
    ind[:] = read_vector(f1, "I")
    vv['ind'] = ind 
    # Skip vv to next / prev
    skip(f1, 2)
    # Let's break this down.  Oliver skips 3+3+6+8+8+8
    # I think this is:
    #   center (N,3) double
    #   parent indices (N) int
    #   neighbor indices (N, 8) int
    #   child indices (N, 8) int
    #   cpu_map (N, 8) int
    pos = na.empty((ng, 3), dtype='float64')
    pos[:,0] = read_vector(f1, "d") # center_x
    pos[:,1] = read_vector(f1, "d") # center_y
    pos[:,2] = read_vector(f1, "d") # center_z
    vv['pos'] = pos
    vv['parents'] = read_vector(f1, "I") # parents
    skip(f1, 6) # neighbors
    #for i in range(6):
    #    vv.append(read_vector(f1, "I")) # neighbor
    children = na.empty((ng, 8), dtype='int64')
    for i in range(8):
        children[:,i] = read_vector(f1, "I") # children
    vv['children'] = children
    #skip(f1, 8) # cpu map
    cpu_map = na.empty((ng, 8), dtype='int64')
    for i in range(8):
        cpu_map[:,i] = read_vector(f1, "I") # cpu map
    vv['cpu_map'] = cpu_map
    rmap = na.empty((ng, 8), dtype='int64')
    for i in range(8):
      rmap[:,i] = read_vector(f1, "I") # refinement map
    vv['rmap'] = rmap
    mi = min(mi, pos.min())
    ma = max(ma, pos.max())
    print "NUMGRIDS", ng
    aa = otc.nocts
    otc.add_ramses(1, ilvl, ng, pos, vv['ind'], cpu_map)
    print otc.nocts - aa, ng

cur = f1.tell()
print "Ended %s bytes from end" % (end - cur)
#pprint.pprint(hvals)

# Let's compare our hilbert curves
print mi, ma

class SphereMock(object):
    center = [0.5, 0.5, 0.5]
    radius = 0.1

ss = SphereMock()
sobj = SphereSelector(ss)
t1 = time.time()
for i in range(10000):
    mask = sobj.select_octs(otc)
t2 = time.time()
print mask.sum(), mask.size
print "%0.3e %0.3e" % (t2-t1, (t2-t1)/10000.)