Python ScdMesh.iterateHex方法代码示例

# 需要导入模块: from r2s.scdmesh import ScdMesh [as 别名]
# 或者: from r2s.scdmesh.ScdMesh import iterateHex [as 别名]
    def test_large_iterator(self):

        print "building large mesh"
        big = ScdMesh(range(1,100), range(101,200), range(201,300))
        print "iterating (1)"
        for i in big.iterateHex():
            pass
        print "iterating (2)"
        for i in big.iterateHex( 'yzx' ):
            pass

# 需要导入模块: from r2s.scdmesh import ScdMesh [as 别名]
# 或者: from r2s.scdmesh.ScdMesh import iterateHex [as 别名]
def create_ww_mesh(flux_mesh, e_group_names):
    """
    This function reads a flux mesh and returns a weight window mesh tagged with 
    all zeros in every energy group. This is used for the intial MAGIC weight
    window generation, for which no preexisting weight window mesh is given.

    Parameters
    ----------
    flux_mesh : ScdMesh
        A ScdMesh tagged with fluxes in the form X_group_YYY and or
        X_group_total. Addition required tags are "particle" (1 for n, 2 for p) 
        and E_group_bounds (vector of energy upper bounds).
    e_group_names : vector of energy names
        In the form X_group_YYY or X_group_total.
        
    """

    ww_mesh = ScdMesh(flux_mesh.getDivisions('x'),\
                       flux_mesh.getDivisions('y'),\
                       flux_mesh.getDivisions('z'))

    # create ww tags
    for e_group_name in e_group_names:
            ww_tag = ww_mesh.imesh.createTag('ww_{0}'.format(e_group_name), 1, float)
            for voxel in ww_mesh.iterateHex('xyz'):
                ww_tag[voxel] = 0

    # create e_upper_bound tags
    e_upper_bounds = \
        flux_mesh.imesh.getTagHandle("E_upper_bounds")[flux_mesh.imesh.rootSet]
    if isinstance(e_upper_bounds, float):
        e_upper_bounds = [e_upper_bounds]

    e_tag = ww_mesh.imesh.createTag("E_upper_bounds", len(e_upper_bounds), float)
    e_tag[ww_mesh.imesh.rootSet] = \
        flux_mesh.imesh.getTagHandle("E_upper_bounds")[flux_mesh.imesh.rootSet]

    # create  particle tag
    particle = flux_mesh.imesh.getTagHandle("particle")[flux_mesh.imesh.rootSet]
    particle_tag = ww_mesh.imesh.createTag("particle", 1, int)
    particle_tag[ww_mesh.imesh.rootSet] = particle

    return ww_mesh

# 需要导入模块: from r2s.scdmesh import ScdMesh [as 别名]
# 或者: from r2s.scdmesh.ScdMesh import iterateHex [as 别名]
class ScdMeshIterateTest(unittest.TestCase):

    def setUp(self):
        self.mesh = iMesh.Mesh()
        self.sm = ScdMesh( range(10,15), # i = 0,1,2,3
                           range(21,25), # j = 0,1,2
                           range(31,34), # k = 0,1
                           self.mesh )

        self.I = range(0,4)
        self.J = range(0,3)
        self.K = range(0,2)

    def test_bad_iterates(self):

        self.assertRaises( ScdMeshError,  self.sm.iterateHex, 'abc' )
        self.assertRaises( TypeError,  self.sm.iterateHex, 12 )
        self.assertRaises( ScdMeshError,  self.sm.iterateHex, 'xxyz' )
        self.assertRaises( ScdMeshError,  self.sm.iterateHex, 'yyx' )
        self.assertRaises( ScdMeshError,  self.sm.iterateHex, 'xyz',z=[0,1,2] )

    def test_iterate_3d(self):
        
        # use izip_longest in the lockstep iterations below; this will catch any
        # situations where one iterator turns out to be longer than expected.
        izip = itertools.izip_longest

        it = self.sm.scdset.iterate( iBase.Type.region, 
                                     iMesh.Topology.hexahedron )

        print "testing zyx"

        # Test the zyx order, which is default; it should be equivalent
        # to the standard imesh iterator
        for it_x, sm_x in izip( it, self.sm.iterateHex() ):
            self.assertEqual( it_x, sm_x )

        print "testing xyz"

        all_indices_zyx = itertools.product( self.I, self.J, self.K )
        # Test the xyz order, the default from original mmGridGen
        for ijk_index, sm_x in izip( all_indices_zyx, 
                                     self.sm.iterateHex('xyz') ):
            self.assertEqual( self.sm.getHex(*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 ):
            print 'testing',order
            all_indices = itertools.product(*args)
            for ijk_index, sm_x in izip( tuple_sort(all_indices, order),
                                         self.sm.iterateHex(order,**kw) ):
                self.assertEqual(self.sm.getHex(*ijk_index), sm_x)

        test_order( 'yxz', self.I, self.J, self.K )
        test_order( 'yzx', self.I, self.J, self.K )
        test_order( 'xzy', self.I, self.J, self.K )
        test_order( 'zxy', self.I, self.J, self.K )

        # Specify z=[1] to iterator
        test_order( 'xyz', self.I, self.J, [1], z=[1] )
        # Specify y=2 to iterator
        test_order( 'zyx', self.I, [2], self.K, y=2 )
        # specify x and y both to iterator
        test_order( 'yzx', [1,2,3],self.J[:-1], self.K, y=self.J[:-1], x=[1,2,3] )

    def test_iterate_2d(self):
        def test_order( iter1, iter2 ):
            for i1, i2 in itertools.izip_longest( iter1, iter2 ):
                self.assertEqual( i1, i2 )

        test_order( self.sm.iterateHex('yx'), self.sm.iterateHex('zyx', z=[0] ) )
        test_order( self.sm.iterateHex('yx',z=1), self.sm.iterateHex('zyx',z=[1]) )
        test_order( self.sm.iterateHex('yx',z=1), self.sm.iterateHex('yzx',z=[1]) )
        test_order( self.sm.iterateHex('zy',x=[3]), self.sm.iterateHex('zxy',x=3) )

        # Cannot iterate over multiple z's without specifing z order
        self.assertRaises( ScdMeshError, self.sm.iterateHex, 'yx', z=[0,1] )

    def test_iterate_1d(self):
        
        def test_equal( ijk_list, miter ):
            for ijk, i in itertools.izip_longest( ijk_list, miter ):
                self.assertEqual( self.sm.getHex(*ijk), i )

        test_equal( [[0,0,0],[0,0,1]], 
                    self.sm.iterateHex('z') )

        test_equal( [[0,1,1],[0,2,1]],
                    self.sm.iterateHex('y', y=[1,2], z=1) )

        test_equal( [[2,0,0],[2,1,0],[2,2,0]],
                    self.sm.iterateHex('y', x=2) ) 
#.........这里部分代码省略.........

# 需要导入模块: from r2s.scdmesh import ScdMesh [as 别名]
# 或者: from r2s.scdmesh.ScdMesh import iterateHex [as 别名]
def magic_wwinp(flux_mesh, ww_mesh='None', total_bool=False, null_value=0, tolerance=0.1):
    """This function reads in a flux mesh and a ww mesh as well as relevant paramters
       then the magic method is applied and a newly tagged flux is returned.
    """

    # find meshtal type
    tag_names = []
    for tag in flux_mesh.imesh.getAllTags(flux_mesh.getHex(0,0,0)):
        tag_names.append(tag.name)

    if 'n_group_001' in tag_names or 'n_group_total' in tag_names:
       particle = 'n'
    elif 'p_group_001' in tag_names or 'p_group_total' in tag_names:
        particle = 'p'
    else:
        print >>sys.stderr, 'Tag X_group_YYY or X_group_total not found'
        sys.exit(1)

    # find number of e_groups
    num_e_groups = find_num_e_groups(flux_mesh, particle)

    if total_bool == False:
        e_groups = ['{0:03d}'.format(x) for x in range(1, num_e_groups + 1)]
        print "\tGenerating WW for {0} energy groups".format(num_e_groups)
    else:
        e_groups = ['total']
        print "\tGenerating WW for Total energy group"

    # find the max flux value for each e_group, store in vector
    max_fluxes = find_max_fluxes(flux_mesh, particle, e_groups, total_bool)

    if ww_mesh == 'None':
        print "\tNo WW mesh file supplied; generating one based on meshtal"
        ww_bool = False # mesh file NOT preexisting
        # create a mesh with the same dimensions as flux_mesh
        ww_mesh = ScdMesh(flux_mesh.getDivisions('x'),\
                          flux_mesh.getDivisions('y'),\
                          flux_mesh.getDivisions('z'))
        # create a tag for each energy group
        for e_group in e_groups:
            group_name = "ww_{0}_group_{1}".format(particle, e_group)
            ww_mesh.imesh.createTag(group_name, 1, float)   

        # create energy bounds
        tag_e_groups = ww_mesh.imesh.createTag("e_groups", len(e_groups), float)

        if e_groups != ['total']:
            tag_e_groups[ww_mesh.imesh.rootSet] = \
                flux_mesh.imesh.getTagHandle("e_groups")[flux_mesh.imesh.rootSet]
        else:
            tag_e_groups[ww_mesh.imesh.rootSet] = 1E36 # usual MCNP value           


    else:
        ww_bool = True # mesh file preexisting
        # make sure the supplied meshes have the same dimenstions
        ww_mesh = ScdMesh.fromFile(ww_mesh)
        try:
            for i in ('x', 'y', 'z'):
                flux_mesh.getDivisions(i) == ww_mesh.getDivisions(i)

        except:
            print >>sys.stderr, 'Mismatched dimensions on WWINP and flux meshes'
            sys.exit(1)

    print "\tSupplied meshes confirmed to have same dimensions"
    
    # iterate through all voxels          
    flux_voxels = flux_mesh.iterateHex('xyz')
    ww_voxels = ww_mesh.iterateHex('xyz')

    for (flux_voxel, ww_voxel) in zip(flux_voxels, ww_voxels):
        for i, e_group in enumerate(e_groups):
            flux = flux_mesh.imesh.getTagHandle(\
                '{0}_group_{1}'.format(particle, e_group))[flux_voxel]
            error = flux_mesh.imesh.getTagHandle(\
                 '{0}_group_{1}_error'.format(particle, e_group))[flux_voxel]
            if ((ww_bool == False and error != 0.0) \
            or (0.0 < error and error < tolerance)):
                if ww_bool == True:
                    if ww_mesh.imesh.getTagHandle('ww_{0}_group_{1}'\
                    .format(particle, e_group))[ww_voxel] != -1:       
                        ww_mesh.imesh.getTagHandle('ww_{0}_group_{1}'\
                        .format(particle, e_group))[ww_voxel]\
                        = flux/(2*max_fluxes[i]) # apply magic method

                else:
                    ww_mesh.imesh.getTagHandle(\
                        'ww_{0}_group_{1}'.format(particle, e_group))[ww_voxel]\
                         = flux/(2*max_fluxes[i]) # apply magic method

            elif ww_bool == False and error == 0.0 :
                ww_mesh.imesh.getTagHandle(\
                    'ww_{0}_group_{1}'.format(particle, e_group))[ww_voxel]\
                     = null_value

    return ww_mesh, e_groups