def __init__(self, ds, comm_rank=0, comm_size=1, left=None, right=None,
min_level=None, max_level=None, data_source=None):
self.ds = ds
try:
self._id_offset = ds.index.grids[0]._id_offset
except AttributeError:
self._id_offset = 0
if data_source is None:
data_source = ds.all_data()
self.data_source = data_source
if left is None:
left = np.array([-np.inf]*3)
if right is None:
right = np.array([np.inf]*3)
if min_level is None: min_level = 0
if max_level is None: max_level = ds.index.max_level
self.min_level = min_level
self.max_level = max_level
self.comm_rank = comm_rank
self.comm_size = comm_size
self.trunk = Node(None, None, None, left, right, -1, 1)
self.build()
def rebuild_tree_from_array(self, nids, pids, lids,
rids, les, res, gids, splitdims, splitposs):
del self.tree.trunk
self.tree.trunk = Node(None,
None,
None,
les[0], res[0], gids[0], nids[0])
N = nids.shape[0]
for i in range(N):
n = self.get_node(nids[i])
n.set_left_edge(les[i])
n.set_right_edge(res[i])
if lids[i] != -1 and n.left is None:
n.left = Node(n, None, None,
np.zeros(3, dtype='float64'),
np.zeros(3, dtype='float64'),
-1, lids[i])
if rids[i] != -1 and n.right is None:
n.right = Node(n, None, None,
np.zeros(3, dtype='float64'),
np.zeros(3, dtype='float64'),
-1, rids[i])
if gids[i] != -1:
n.grid = gids[i]
if splitdims[i] != -1:
n.create_split(splitdims[i], splitposs[i])
mylog.info('AMRKDTree rebuilt, Final Volume: %e' % self.tree.trunk.kd_sum_volume())
return self.tree.trunk
class Tree(object):
def __init__(self, ds, comm_rank=0, comm_size=1, left=None, right=None,
min_level=None, max_level=None, data_source=None):
self.ds = ds
try:
self._id_offset = ds.index.grids[0]._id_offset
except AttributeError:
self._id_offset = 0
if data_source is None:
data_source = ds.all_data()
self.data_source = data_source
if left is None:
left = np.array([-np.inf]*3)
if right is None:
right = np.array([np.inf]*3)
if min_level is None: min_level = 0
if max_level is None: max_level = ds.index.max_level
self.min_level = min_level
self.max_level = max_level
self.comm_rank = comm_rank
self.comm_size = comm_size
self.trunk = Node(None, None, None, left, right, -1, 1)
self.build()
def add_grids(self, grids):
gles = np.array([g.LeftEdge for g in grids])
gres = np.array([g.RightEdge for g in grids])
gids = np.array([g.id for g in grids], dtype="int64")
self.trunk.add_grids(gids.size, gles, gres, gids,
self.comm_rank, self.comm_size)
del gles, gres, gids, grids
def build(self):
lvl_range = range(self.min_level, self.max_level+1)
for lvl in lvl_range:
#grids = self.data_source.select_grids(lvl)
grids = np.array([b for b, mask in self.data_source.blocks if b.Level == lvl])
if len(grids) == 0: continue
self.add_grids(grids)
def check_tree(self):
for node in self.trunk.depth_traverse():
if node.grid == -1:
continue
grid = self.ds.index.grids[node.grid - self._id_offset]
dds = grid.dds
gle = grid.LeftEdge
nle = self.ds.arr(node.get_left_edge(), input_units="code_length")
nre = self.ds.arr(node.get_right_edge(), input_units="code_length")
li = np.rint((nle-gle)/dds).astype('int32')
ri = np.rint((nre-gle)/dds).astype('int32')
dims = (ri - li).astype('int32')
assert(np.all(grid.LeftEdge <= nle))
assert(np.all(grid.RightEdge >= nre))
assert(np.all(dims > 0))
# print grid, dims, li, ri
# Calculate the Volume
vol = self.trunk.kd_sum_volume()
mylog.debug('AMRKDTree volume = %e' % vol)
self.trunk.kd_node_check()
def sum_cells(self, all_cells=False):
cells = 0
for node in self.trunk.depth_traverse():
if node.grid == -1:
continue
if not all_cells and not node.kd_is_leaf():
continue
grid = self.ds.index.grids[node.grid - self._id_offset]
dds = grid.dds
gle = grid.LeftEdge
nle = self.ds.arr(node.get_left_edge(), input_units="code_length")
nre = self.ds.arr(node.get_right_edge(), input_units="code_length")
li = np.rint((nle-gle)/dds).astype('int32')
ri = np.rint((nre-gle)/dds).astype('int32')
dims = (ri - li).astype('int32')
cells += np.prod(dims)
return cells