#!/usr/bin/env python
# coding: utf-8

# # 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:

# In[2]:


import pymbolic.primitives as p

x = p.Variable("x")


# In[3]:


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.

# In[4]:


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


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

# In[5]:


from pymbolic.mapper import IdentityMapper


# In[6]:


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


# Now apply it:

# In[7]:


dd = DerivDoubler()

dd(u)


# In[ ]: