def expression_indexed(expr, parameters): aggregate, multiindex = (expression(c, parameters) for c in expr.children) return pym.Subscript(aggregate, multiindex) extents = [int(numpy.prod(expr.aggregate.shape[i+1:])) for i in range(len(multiindex))] make_sum = lambda x, y: pym.Sum((x, y)) index = reduce(make_sum, [pym.Product((e, m)) for e, m in zip(extents, multiindex)]) return pym.Subscript(aggregate, (index,))
def map_constant(self, expr): """Convert complex values not within complex64 to a product for loopy """ if not isinstance(expr, complex): return IdentityMapper.map_constant(self, expr) if complex(self.float_type(expr.imag)) == expr.imag: return IdentityMapper.map_constant(self, expr) return expr.real + prim.Product((expr.imag, 1j))
def _expression_flexiblyindexed(expr, ctx): var = expression(expr.children[0], ctx) rank = [] for off, idxs in expr.dim2idxs: for index, stride in idxs: assert isinstance(index, gem.Index) rank_ = [off] for index, stride in idxs: rank_.append(p.Product((ctx.active_indices[index], stride))) rank.append(p.Sum(tuple(rank_))) return p.Subscript(var, tuple(rank))
def test_field_diff(proc_shape): if proc_shape != (1, 1, 1): pytest.skip("test field only on one rank") from pystella import diff y = ps.Field("y") assert diff(y, y) == 1 assert diff(y[0], y[0]) == 1 assert diff(y[0], y[1]) == 0 y = ps.DynamicField("y") assert diff(y, y) == 1 assert diff(y[0], y[0]) == 1 assert diff(y[0], y[1]) == 0 import pymbolic.primitives as pp assert diff(y**3, y, "t") == pp.Product((3, 2, y, y.d(0))) assert diff(y**3, "t", y) == pp.Product((3, y.d(0), 2, y)) for i, x in enumerate(["t", "x", "y", "z"]): assert diff(y, x) == y.d(i) assert diff(y[1, 3], x) == y.d(1, 3, i) assert diff(y[1]**2, x) == 2 * y[1] * y.d(1, i)
def _expression_product(expr, ctx): return p.Product(tuple(expression(c, ctx) for c in expr.children))
def _mult(x, y): return p.Product((x, y))
def _sub(x, y): return p.Sum((x, p.Product(((-1), y))))
def map_Mul(self, expr): return prim.Product(tuple(self.rec(arg) for arg in expr.args))
def _mult(x, y): # noqa return p.Product((x, y))
def _neg(x): # noqa return p.Product((-1), x)
def parse_postfix(self, pstate, min_precedence, left_exp): import pymbolic.primitives as primitives did_something = False next_tag = pstate.next_tag() if next_tag is _openpar and _PREC_CALL > min_precedence: pstate.advance() args, kwargs = self.parse_arglist(pstate) if kwargs: left_exp = primitives.CallWithKwargs(left_exp, args, kwargs) else: left_exp = primitives.Call(left_exp, args) did_something = True elif next_tag is _openbracket and _PREC_CALL > min_precedence: pstate.advance() pstate.expect_not_end() left_exp = primitives.Subscript(left_exp, self.parse_expression(pstate)) pstate.expect(_closebracket) pstate.advance() did_something = True elif next_tag is _if and _PREC_IF > min_precedence: from pymbolic.primitives import If then_expr = left_exp pstate.advance() pstate.expect_not_end() condition = self.parse_expression(pstate, _PREC_LOGICAL_OR) pstate.expect(_else) pstate.advance() else_expr = self.parse_expression(pstate) left_exp = If(condition, then_expr, else_expr) did_something = True elif next_tag is _dot and _PREC_CALL > min_precedence: pstate.advance() pstate.expect(_identifier) left_exp = primitives.Lookup(left_exp, pstate.next_str()) pstate.advance() did_something = True elif next_tag is _plus and _PREC_PLUS > min_precedence: pstate.advance() right_exp = self.parse_expression(pstate, _PREC_PLUS) if isinstance(left_exp, primitives.Sum): left_exp = primitives.Sum(left_exp.children + (right_exp,)) else: left_exp = primitives.Sum((left_exp, right_exp)) did_something = True elif next_tag is _minus and _PREC_PLUS > min_precedence: pstate.advance() right_exp = self.parse_expression(pstate, _PREC_PLUS) if isinstance(left_exp, primitives.Sum): left_exp = primitives.Sum(left_exp.children + ((-right_exp),)) # noqa pylint:disable=invalid-unary-operand-type else: left_exp = primitives.Sum((left_exp, -right_exp)) # noqa pylint:disable=invalid-unary-operand-type did_something = True elif next_tag is _times and _PREC_TIMES > min_precedence: pstate.advance() right_exp = self.parse_expression(pstate, _PREC_PLUS) if isinstance(left_exp, primitives.Product): left_exp = primitives.Product(left_exp.children + (right_exp,)) else: left_exp = primitives.Product((left_exp, right_exp)) did_something = True elif next_tag is _floordiv and _PREC_TIMES > min_precedence: pstate.advance() left_exp = primitives.FloorDiv( left_exp, self.parse_expression(pstate, _PREC_TIMES)) did_something = True elif next_tag is _over and _PREC_TIMES > min_precedence: pstate.advance() left_exp = primitives.Quotient( left_exp, self.parse_expression(pstate, _PREC_TIMES)) did_something = True elif next_tag is _modulo and _PREC_TIMES > min_precedence: pstate.advance() left_exp = primitives.Remainder( left_exp, self.parse_expression(pstate, _PREC_TIMES)) did_something = True elif next_tag is _power and _PREC_POWER > min_precedence: pstate.advance() left_exp = primitives.Power( left_exp, self.parse_expression(pstate, _PREC_TIMES)) did_something = True elif next_tag is _and and _PREC_LOGICAL_AND > min_precedence: pstate.advance() from pymbolic.primitives import LogicalAnd left_exp = LogicalAnd(( left_exp, self.parse_expression(pstate, _PREC_LOGICAL_AND))) did_something = True elif next_tag is _or and _PREC_LOGICAL_OR > min_precedence: pstate.advance() from pymbolic.primitives import LogicalOr left_exp = LogicalOr(( left_exp, self.parse_expression(pstate, _PREC_LOGICAL_OR))) did_something = True elif next_tag is _bitwiseor and _PREC_BITWISE_OR > min_precedence: pstate.advance() from pymbolic.primitives import BitwiseOr left_exp = BitwiseOr(( left_exp, self.parse_expression(pstate, _PREC_BITWISE_OR))) did_something = True elif next_tag is _bitwisexor and _PREC_BITWISE_XOR > min_precedence: pstate.advance() from pymbolic.primitives import BitwiseXor left_exp = BitwiseXor(( left_exp, self.parse_expression(pstate, _PREC_BITWISE_XOR))) did_something = True elif next_tag is _bitwiseand and _PREC_BITWISE_AND > min_precedence: pstate.advance() from pymbolic.primitives import BitwiseAnd left_exp = BitwiseAnd(( left_exp, self.parse_expression(pstate, _PREC_BITWISE_AND))) did_something = True elif next_tag is _rightshift and _PREC_SHIFT > min_precedence: pstate.advance() from pymbolic.primitives import RightShift left_exp = RightShift( left_exp, self.parse_expression(pstate, _PREC_SHIFT)) did_something = True elif next_tag is _leftshift and _PREC_SHIFT > min_precedence: pstate.advance() from pymbolic.primitives import LeftShift left_exp = LeftShift( left_exp, self.parse_expression(pstate, _PREC_SHIFT)) did_something = True elif next_tag in self._COMP_TABLE and _PREC_COMPARISON > min_precedence: pstate.advance() from pymbolic.primitives import Comparison left_exp = Comparison( left_exp, self._COMP_TABLE[next_tag], self.parse_expression(pstate, _PREC_COMPARISON)) did_something = True elif next_tag is _colon and _PREC_SLICE >= min_precedence: pstate.advance() expr_pstate = pstate.copy() assert not isinstance(left_exp, primitives.Slice) from pytools.lex import ParseError try: next_expr = self.parse_expression(expr_pstate, _PREC_SLICE) except ParseError: # no expression follows, too bad. left_exp = primitives.Slice((left_exp, None,)) else: left_exp = _join_to_slice(left_exp, next_expr) pstate.assign(expr_pstate) did_something = True elif next_tag is _comma and _PREC_COMMA > min_precedence: # The precedence makes the comma left-associative. pstate.advance() if pstate.is_at_end() or pstate.next_tag() is _closepar: if isinstance(left_exp, (tuple, list)) \ and not isinstance(left_exp, FinalizedContainer): # left_expr is a container with trailing commas pass else: left_exp = (left_exp,) else: new_el = self.parse_expression(pstate, _PREC_COMMA) if isinstance(left_exp, (tuple, list)) \ and not isinstance(left_exp, FinalizedContainer): left_exp = left_exp + (new_el,) else: left_exp = (left_exp, new_el) did_something = True return left_exp, did_something