from yt.config import ytcfg; ytcfg["yt","suppressStreamLogging"] = "True"
from yt.mods import *
from pyne.xs.channels import sigma_t
from pyne.material import Material, from_atom_frac
from pyne.mesh import Mesh, IMeshTag, MetadataTag, ComputedTag

import yt

yt.__file__

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)
coords = [c5, c5, [0.0, 1.0]]

m = Mesh(structured_coords=coords, structured=True)
fuel = from_atom_frac({'U235': 0.045, 'U238': 0.955, 'O16': 2.0}, density=10.7)
cool = from_atom_frac({'H1': 2.0, 'O16': 1.0}, density=1.0)
for i in range(len(m)):
    m.mats[i] = cool
m.mats[len(m)/2] = fuel
m.mats[len(m)/4] = fuel
for i, c in enumerate(c5[:-1]):
    m.mats[i*len(c5)] = fuel

m.density 

m.density[42]

print m.density[::100]  # slice

print m.density[m.density[:] >= 10]  # mask

print m.density[[10, 0, 11, 100]]  # fancy index is fancy

m.tags.keys()

pf = PyneMoabHex8StaticOutput(m)
s = SlicePlot(pf, 'z', 'density', origin='native')
s.display()

m.density2 = ComputedTag(lambda mesh, i: mesh.density[i]**2)

pf = PyneMoabHex8StaticOutput(m)
s = SlicePlot(pf, 'z', 'density2', origin='native')
s.display()

m.sigma_t = ComputedTag(lambda mesh, i: sigma_t(mesh.mats[i], group_struct=[10.0, 1e-6], phi_g=[1.0])[0])

pf = PyneMoabHex8StaticOutput(m)
s = SlicePlot(pf, 'z', 'sigma_t', origin='native')
s.display()