Source

yt.geometry_tests / read_ramses_header.py

Full commit
import struct
import numpy as na
import pprint
import os
from yt.utilities._amr_utils.geometry_utils import \
    get_hilbert_points, \
    get_hilbert_indices
from yt.utilities.fortran_utils import read_attrs, read_vector, skip

hvals = {}

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')
          )

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)
for ilvl in range(hvals['nlevelmax']):
    cur = f1.tell()
    if cur == end: break
    print "HANDLING LEVEL %s, %s from end" % (ilvl, end - cur)
    vv = pointers[ilvl] = {}
    ng = hvals['numbl'][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
    #for i in range(8):
    #    vv.append(read_vector(f1, "I")) # 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

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

# Let's compare our hilbert curves
for ilvl in range(hvals['nlevelmax']):
    if ilvl not in pointers: break
    vv = pointers[ilvl]
    ind = vv['ind']
    dx = 2.0**-(ilvl+1)
    pos = vv['pos'] - dx
    hind = get_hilbert_indices(ilvl, (pos/dx - 0.5).astype("int64"))
    hpos = get_hilbert_points(ilvl, na.arange(ind.size).astype("int64"))
    hpos = hpos.astype("float64") * dx * 2
    assert(na.unique(hind).size == hind.size)
    print (hind[1:] - hind[:-1] > 0).sum() / float(hind.size)
    if ilvl == 1: break

import matplotlib;matplotlib.use("Agg")
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()

plt.clf()
ax = fig.add_subplot(111, projection='3d')
ax.plot(hpos[:,0], hpos[:,1], hpos[:,2], '-x')
ax.set_xlim(0,1)
ax.set_ylim(0,1)
ax.set_zlim(0,1)
ax.set_title("From yt")
plt.savefig("hind_yt.png")

plt.clf()
ax = fig.add_subplot(111, projection='3d')
ax.plot(pos[:,0], pos[:,1], pos[:,2], '-x')
ax.set_xlim(0,1)
ax.set_ylim(0,1)
ax.set_zlim(0,1)
ax.set_title("From RAMSES")
plt.savefig("hind_ramses.png")

print hind[0], ind[0]