from pyne.mesh import Mesh
import numpy as np

def cantor(n):
    return [0.] + cant(0., 1., n) + [1.]

def cant(x, y, n):
    if n == 0:
        return []
    new_pts = [2.*x/3. + y/3., x/3. + 2.*y/3.]
    return cant(x, new_pts[0], n-1) + new_pts + cant(new_pts[1], y, n-1)

c5 = cantor(5)
c7 = cantor(7)
c8 = cantor(8)

points = c5

my_mesh = Mesh(structured_coords=[points]*3, structured=True)
flux_tag = my_mesh.mesh.createTag("flux", 1, float)

for ent in my_mesh.structured_iterate_hex("xyz"):
    coords = my_mesh.mesh.getVtxCoords(ent)
    flux_tag[ent] = sum(coords**2) 

my_mesh.mesh.save("c5.h5m")