Source

yt.geometry_tests / quadtree.py

Full commit
from yt.mods import *
from yt.utilities.amr_utils import QuadTree
import time
t1 = time.time()
#pf = load("DD0039/output_0039")
#pf = load("output_00045/info_00045.txt")
pf = load("output_00700/info_00700.txt")
#pf = load("output_00007/info_00007.txt")
axis = 0
xax = x_dict[axis]
yax = y_dict[axis]
dd = pf.h.all_data()
pf.domain_dimensions[:] = 3
tree = QuadTree(na.array([pf.domain_dimensions[xax],
                          pf.domain_dimensions[yax]],
                          dtype="int64"), nvals=1)

class ProjMock(dict):
    pass

level_extrema = defaultdict(lambda: (1e30, -1e30))
for i,chunk in enumerate(dd.chunks(["Density"], "io")):
    print i, chunk
    pxs = chunk.icoords[:,xax]
    pys = chunk.icoords[:,yax]
    ires = level = chunk.ires
    pvals = chunk["Density"].astype("float64").reshape((chunk.size, 1))
    opvals = pvals.copy()
    pvals /= (2**(level+1))[:,None]
    wvals = chunk["Density"].astype("float64")
    #pvals *= wvals[:,None]
    tree.add_chunk_to_tree(pxs, pys, level, pvals, wvals)
    for l in na.unique(level):
        li = (level == l)
        level_extrema[l] = (min(opvals[li].min(), level_extrema[l][0]),
                            max(opvals[li].max(), level_extrema[l][1]))
for l, (mi, ma) in sorted(level_extrema.items()):
    print "% 2i: %0.3e %0.3e" % (l, mi, ma)
t2 = time.time()

print "Took %0.3e" % (t2-t1)

coord_data, field_data, weight_data, dxs = [], [], [], []
level = 0
for level in range(pf.max_level + 1):
    npos, nvals, nwvals = tree.get_all_from_level(level)
    coord_data.append(npos)
    #nvals /= nwvals[:,None]
    field_data.append(nvals)
    weight_data.append(nwvals)
    dxs.append(1.0/(2.0**(level+1)) * na.ones(nvals.shape[0], dtype="float64"))
    if nvals.shape[0] > 0:
        print "% 2i: %0.3e %0.3e" % (
            level, field_data[-1].min(), field_data[-1].max())
    level += 1
coord_data = na.concatenate(coord_data, axis=0).transpose()
field_data = na.concatenate(field_data, axis=0).transpose()
weight_data = na.concatenate(weight_data, axis=0).transpose()
dxs = na.concatenate(dxs, axis=0).transpose()
proj = ProjMock()
proj.pf = pf
proj['px'] = (coord_data[0,:] + 0.5) * dxs
proj['py'] = (coord_data[1,:] + 0.5) * dxs
proj['pdx'] = dxs / 2.0
proj['pdy'] = dxs / 2.0
proj['weight_field'] = weight_data
fields = ["Density"]
for field, arr in zip(fields, field_data):
    proj[field] = arr
proj.axis = 0
frb = FixedResolutionBuffer(proj, (0.0, 1.0, 0.0, 1.0), (1024, 1024))
write_image(na.log10(frb["Density"]), "first_light.png")