Test script for allocating EnzoE Blocks using the new EnzoEOctreeContainer and the Check-00900 dataset.

new_script.py

import yt
import h5py
import glob
import numpy as np
from collections import defaultdict

from yt.geometry.oct_container import EnzoEOctreeContainer

nroot = 16
octree = EnzoEOctreeContainer([16, 16, 16], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0], 16)

domains = defaultdict(dict)
domain_count = []
for i, fn in enumerate(sorted(glob.glob("../Check-000900/block*.block_list"))):
    print(fn)
    levels = domains[i] = defaultdict(list)
    for line in open(fn):
        s, l = line.split()
        if int(l) < 0: continue
        xs, ys, zs = s[1:].split("_")
        if int(l) == 0:
            levels[0].append([int(_, 2) for _ in (xs, ys, zs)] + [0, 0, 0])
            continue
        rxpos, xpos = xs.split(":")
        rypos, ypos = ys.split(":")
        rzpos, zpos = zs.split(":")
        levels[int(l)].append([int(_, 2) for _ in (rxpos, rypos, rzpos, xpos, ypos, zpos)])

    domain_count.append(sum(len(_) for _ in levels.values()))

print(domain_count)
octree.allocate_domains([_ + 1000 for _ in domain_count], 16**3)
for i in sorted(domains):
    for level in sorted(domains[i]):
        pos = np.array(domains[i][level], dtype="uint64")
        if pos.size == 0: continue
        root_pos = pos[:, :3]
        ipos = pos[:, -3:]
        print(i, level, root_pos.shape, ipos.shape)
        print(octree.add(i, level, ipos, root_pos))