Python pytools.indices_in_shape函数代码示例

    def map_numpy_array(self, expr):
        if not self.visit(expr):
            return

        from pytools import indices_in_shape
        for i in indices_in_shape(expr.shape):
            self.rec(expr[i])

def make_common_subexpression(field, prefix=None):
    from pytools.obj_array import log_shape
    from hedge.tools import is_zero
    from pymbolic.primitives import CommonSubexpression

    ls = log_shape(field)
    if ls != ():
        from pytools import indices_in_shape
        result = numpy.zeros(ls, dtype=object)

        for i in indices_in_shape(ls):
            if prefix is not None:
                component_prefix = prefix+"_".join(str(i_i) for i_i in i)
            else:
                component_prefix = None

            if is_zero(field[i]):
                result[i] = 0
            else:
                result[i] = CommonSubexpression(field[i], component_prefix)

        return result
    else:
        if is_zero(field):
            return 0
        else:
            return CommonSubexpression(field, prefix)

def make_common_subexpression(field, prefix=None):
    try:
        from pytools.obj_array import log_shape
    except ImportError:
        have_obj_array = False
    else:
        have_obj_array = True

    if have_obj_array:
        ls = log_shape(field)

    if have_obj_array and ls != ():
        from pytools import indices_in_shape
        result = numpy.zeros(ls, dtype=object)

        for i in indices_in_shape(ls):
            if prefix is not None:
                component_prefix = prefix+"_".join(str(i_i) for i_i in i)
            else:
                component_prefix = None

            if is_constant(field[i]):
                result[i] = field[i]
            else:
                result[i] = CommonSubexpression(field[i], component_prefix)

        return result
    else:
        if is_constant(field):
            return field
        else:
            return CommonSubexpression(field, prefix)

def with_object_array_or_scalar_n_args(f, *args):
    oarray_arg_indices = []
    for i, arg in enumerate(args):
        if is_obj_array(arg):
            oarray_arg_indices.append(i)

    if not oarray_arg_indices:
        return f(*args)

    leading_oa_index = oarray_arg_indices[0]

    ls = log_shape(args[leading_oa_index])
    if ls != ():
        from pytools import indices_in_shape
        result = np.zeros(ls, dtype=object)

        new_args = list(args)
        for i in indices_in_shape(ls):
            for arg_i in oarray_arg_indices:
                new_args[arg_i] = args[arg_i][i]

            result[i] = f(*new_args)
        return result
    else:
        return f(*args)

    def __call__(self, a):
        from pytools import indices_in_shape

        # assumes nonempty, which is reasonable
        return max(
                abs(self.scalar_kernel(a[i]))
                for i in indices_in_shape(a.shape))

 def map_numpy_array(self, expr, *args, **kwargs):
     import numpy
     result = numpy.empty(expr.shape, dtype=object)
     from pytools import indices_in_shape
     for i in indices_in_shape(expr.shape):
         result[i] = self.rec(expr[i], *args, **kwargs)
     return result

def to_obj_array(ary):
    ls = log_shape(ary)
    result = np.empty(ls, dtype=object)

    from pytools import indices_in_shape
    for i in indices_in_shape(ls):
        result[i] = ary[i]

    return result

    def apply_mask(field):
        from hedge.tools import log_shape
        ls = log_shape(field)
        result = discr.volume_empty(ls)
        from pytools import indices_in_shape
        for i in indices_in_shape(ls):
            result[i] = mask * field[i]

        return result

def grad_D(kernel, arg, dim):
    from pytools.obj_array import log_shape
    arg_shape = log_shape(arg)
    result = np.zeros(arg_shape+(dim,), dtype=object)
    from pytools import indices_in_shape
    for i in indices_in_shape(arg_shape):
        for j in range(dim):
            result[i+(j,)] = IntGdMixed(kernel, arg[i], j)
    return result

def ptwise_mul(a, b):
    from pytools.obj_array import log_shape
    a_log_shape = log_shape(a)
    b_log_shape = log_shape(b)

    from pytools import indices_in_shape

    if a_log_shape == ():
        result = np.empty(b_log_shape, dtype=object)
        for i in indices_in_shape(b_log_shape):
            result[i] = a*b[i]
    elif b_log_shape == ():
        result = np.empty(a_log_shape, dtype=object)
        for i in indices_in_shape(a_log_shape):
            result[i] = a[i]*b
    else:
        raise ValueError("ptwise_mul can't handle two non-scalars")

    return result

def make_sym_array(name, shape):
    vfld = Variable(name)
    if shape == ():
        return vfld

    import numpy as np
    result = np.zeros(shape, dtype=object)
    from pytools import indices_in_shape
    for i in indices_in_shape(shape):
        result[i] = vfld[i]

    return result

def ptwise_dot(logdims1, logdims2, a1, a2):
    a1_log_shape = a1.shape[:logdims1]
    a2_log_shape = a1.shape[:logdims2]

    assert a1_log_shape[-1] == a2_log_shape[0]
    len_k = a2_log_shape[0]

    result = np.empty(a1_log_shape[:-1]+a2_log_shape[1:], dtype=object)

    from pytools import indices_in_shape
    for a1_i in indices_in_shape(a1_log_shape[:-1]):
        for a2_i in indices_in_shape(a2_log_shape[1:]):
            result[a1_i+a2_i] = sum(
                    a1[a1_i+(k,)] * a2[(k,)+a2_i]
                    for k in xrange(len_k)
                    )

    if result.shape == ():
        return result[()]
    else:
        return result

    def __call__(self, discr, t, fields, x, make_empty):
        result = self.make_func(discr)(
                t=numpy.float64(t), x=x, fields=fields)

        # make sure we return no scalars in the result
        from pytools.obj_array import log_shape, is_obj_array
        if is_obj_array(result):
            from pytools import indices_in_shape
            from hedge.optemplate.tools import is_scalar
            for i in indices_in_shape(log_shape(result)):
                if is_scalar(result[i]):
                    result[i] = make_empty().fill(result[i])

        return result

def count_dofs(vec):
    try:
        dtype = vec.dtype
        size = vec.size
        shape = vec.shape
    except AttributeError:
        from warnings import warn
        warn("could not count dofs of vector")
        return 0

    if dtype == object:
        from pytools import indices_in_shape
        return sum(count_dofs(vec[i])
                for i in indices_in_shape(vec.shape))
    else:
        return size

def with_object_array_or_scalar(f, field, obj_array_only=False):
    if obj_array_only:
        if is_obj_array(field):
            ls = field.shape
        else:
            ls = ()
    else:
        ls = log_shape(field)
    if ls != ():
        from pytools import indices_in_shape
        result = np.zeros(ls, dtype=object)
        for i in indices_in_shape(ls):
            result[i] = f(field[i])
        return result
    else:
        return f(field)