def test_vtx_iterator():
# use vanilla izip as we"ll test using non-equal-length iterators
izip = itertools.izip
sm = Mesh(structured=True,
structured_coords=[range(10, 15),
range(21, 25),
range(31, 34)])
it = meshset_iterate(sm.mesh, sm.structured_set, types.MBVERTEX)
# test the default order
for (it_x,
sm_x) in itertools.izip_longest(it,
sm.structured_iterate_vertex("zyx")):
assert_equal(it_x, sm_x)
# Do the same again, but use an arbitrary kwarg to structured_iterate_vertex
# to prevent optimization from kicking in
it.reset()
for (it_x, sm_x) in itertools.izip_longest(
it, sm.structured_iterate_vertex("zyx", no_opt=True)):
assert_equal(it_x, sm_x)
it.reset()
for (it_x, sm_x) in izip(it,
sm.structured_iterate_vertex("yx", z=sm.dims[2])):
assert_equal(it_x, sm_x)
it.reset()
for (it_x, sm_x) in izip(it, sm.structured_iterate_vertex("x")):
assert_equal(it_x, sm_x)
def test_iterate_3d():
# use izip_longest in the lockstep iterations below; this will catch any
# situations where one iterator turns out to be longer than expected.
sm = Mesh(structured=True,
structured_coords=[range(10, 15),
range(21, 25),
range(31, 34)])
I = range(0, 4)
J = range(0, 3)
K = range(0, 2)
izip = itertools.izip_longest
it = meshset_iterate(sm.mesh, sm.structured_set, types.MBHEX)
# Test the zyx order, which is default; it should be equivalent
# to the standard pyne.mesh iterator
for it_x, sm_x in izip(it, sm.structured_iterate_hex()):
assert_equal(it_x, sm_x)
# testing xyz
all_indices_zyx = itertools.product(I, J, K)
# Test the xyz order, the default from original mmGridGen
for ijk_index, sm_x in izip(all_indices_zyx,
sm.structured_iterate_hex("xyz")):
assert_equal(sm.structured_get_hex(*ijk_index), sm_x)
def _tuple_sort(collection, indices):
# sorting function for order test
def t(tup):
# sort this 3-tuple according to the order of x, y, and z in
# indices
return (tup["xyz".find(indices[0])] * 100 +
tup["xyz".find(indices[1])] * 10 +
tup["xyz".find(indices[2])])
return sorted(collection, key=t)
def test_order(order, *args, **kw):
all_indices = itertools.product(*args)
for ijk_index, sm_x in izip(_tuple_sort(all_indices, order),
sm.structured_iterate_hex(order, **kw)):
assert_equal(sm.structured_get_hex(*ijk_index), sm_x)
test_order("yxz", I, J, K)
test_order("yzx", I, J, K)
test_order("xzy", I, J, K)
test_order("zxy", I, J, K)
# Specify z=[1] to iterator
test_order("xyz", I, J, [1], z=[1])
# Specify y=2 to iterator
test_order("zyx", I, [2], K, y=2)
# specify x and y both to iterator
test_order("yzx", [1, 2, 3], J[:-1], K, y=J[:-1], x=[1, 2, 3])