Source

yt.geometry_tests / test_art_octre.py

from yt.mods import *
from yt.utilities.amr_utils import RAMSESOctreeContainer
import yt.geometry.selection_routines as sel

NX = 128
LE = na.array([0.0, 0.0, 0.0], dtype='float64')
RE = na.array([1.0, 1.0, 1.0], dtype='float64')

art_oct = RAMSESOctreeContainer(NX, LE, RE)
art_oct.allocate_domains(1)

for level in range(MAX_LEVEL+1):
    # read level info into POS
    # 1 is the CPU, -1 is not used, 1 is the domain ID
    art_oct.add(1, level, -1, pos, 1)

# Now let's test a selector
class Empty(object):
    pass

sp_obj = Empty()
sp_obj.center = [0.5, 0.5, 0.5]
sp_obj.radius = 0.1
sp = sel.SphereSelector(sp_obj)

mask = na.zeros(sp.nocts, dtype="bool")

cells = art_oct.count_cells(sp, mask)
nc = cells.sum()
print nc, sp.nocts
# 1 is domain_id
icoords = art_oct.icoords(1, mask, nc)
ires = art_oct.ires(1, mask, nc)

print icoords
print ires