Defining Custom Node Types

Mathematical expressions are only the first step. Most of the time, in mathematical software, the interesting aspects are special "things" that are strung together by expressions.

So it would be helpful to be able to define our own expression types:

import pymbolic.primitives as p

x = p.Variable("x")

class DerivativeOperator(p.Expression):
    def __init__(self, operand):
        self.operand = operand

    def __getinitargs__(self):
        return (self.operand,)

    mapper_method = "map_derivative_operator"

__getinitargs__ tells pymbolic what the arguments of the constructor were. This is used for printing and comparisons.

u = x/DerivativeOperator((x + 23)**0.5)
u

Quotient(Variable('x'), DerivativeOperator(Power(Sum((Variable('x'), 23)), 0.5)))

We can then also define custom mappers (let's call ours DerivDoubler) that operate on these node types:

from pymbolic.mapper import IdentityMapper

class DerivDoubler(IdentityMapper):
    def map_derivative_operator(self, expr):
        return 2*DerivativeOperator(self.rec(expr.operand))

Now apply it:

dd = DerivDoubler()

dd(u)

Quotient(Variable('x'), Product((2, DerivativeOperator(Power(Sum((Variable('x'), 23)), 0.5)))))