^{Department of Physics and Astronomy}_{The Forbes Group
SymPy for Code Generation}

Code Generation with SymPy¶

I am often working with density functionals where equations of motion follow by varying an energy functional with respect to some densities. In this post I explore using SymPy as a tool for expressing these relationships and for generating the relevant code.

In [1]:

import sympy
sympy.init_printing(use_latex=True)
from sympy import S, var

Differentiating a non-linear expression is straightforward, the complication is dealing with gradient terms. For example, consider a Brachistochrone problem to find the curve $y(x)$ that minimizes the time for a particle to slide from $y(0) = 0$ to $y(1) = y_1$. From energy conservation starting with $v(0) = 0$ we have $v^2/2 = gy$ where $g \approx -10$m/s$^2$. We can formulate the problem in terms of the horizontal velocity

$$ v_x = \dot{x} = \frac{v}{\sqrt{1+(y')^2}} $$

where $y'(x)$ is the slope to find the total time $T$ which is to be minimized:

$$ \DeclareMathOperator{\d}{d} T = \int_0^T \d{t} = \int_0^1 \d{x}\;\diff{}{t}{x} = \int_0^1 \d{x}\;\frac{1}{v_x} = \int_0^1 \d{x}\;\sqrt{\frac{1+(y')^2}{2gy}} $$

In [28]:

x, y, z, g = sympy.symbols('x, y, z, g')
y_x = sympy.Function('y')(x)
dy_x = sympy.Derivative(y(x), x)
f = sympy.sqrt((1+dy_x**2)/(2*g*y_x))
F = sympy.Integral(f, (x, 0, 1))
F, f

Out[28]:

$$\left ( \int_{0}^{1} \frac{\sqrt{2}}{2} \sqrt{\frac{\left(\frac{d}{d x} y{\left (x \right )}\right)^{2} + 1}{g y{\left (x \right )}}}\, dx, \quad \frac{\sqrt{2}}{2} \sqrt{\frac{\left(\frac{d}{d x} y{\left (x \right )}\right)^{2} + 1}{g y{\left (x \right )}}}\right )$$

The Euler-Lagrange equations can be obtained by differentiating with respect to the variable and derivative, but sympy can't do the latter without help:

In [29]:

f.diff(dy_x)

Out[29]:

$$\frac{\sqrt{2} \sqrt{\frac{\left(\frac{d}{d x} y{\left (x \right )}\right)^{2} + 1}{g y{\left (x \right )}}}}{2 \left(\frac{d}{d x} y{\left (x \right )}\right)^{2} + 2} \frac{d}{d x} y{\left (x \right )}$$

We can do this with substitutions:

In [22]:

def ddiff(f, dy_x):
    return f.subs(dy_x, '_dy_x').diff('_dy_x').subs('_dy_x', dy_x)
ddiff(f, dy_x)

Out[22]:

$$\frac{\sqrt{2} \sqrt{\frac{1}{g y} \left(\left(\frac{d}{d x} y\right)^{2} + 1\right)} \frac{d}{d x} y}{2 \left(\frac{d}{d x} y\right)^{2} + 2}$$

In [23]:

(f.diff(y_x) - ddiff(f, dy_x).diff(x)).simplify()

Out[23]:

$$- \frac{\sqrt{2}}{4 y} \sqrt{\frac{1}{g y} \left(\left(\frac{d}{d x} y\right)^{2} + 1\right)}$$

In [18]:

[sympy.simplify(_s.subs(y(x), y)) for _s in sympy.calculus.euler.euler_equations(f.subs(y, y(x)))]

Out[18]:

$$\left [ - \frac{\sqrt{2} \sqrt{\frac{1}{g y} \left(\left(\frac{d}{d x} y\right)^{2} + 1\right)} \left(2 y \frac{d^{2}}{d x^{2}} y + \left(\frac{d}{d x} y\right)^{2} + 1\right)}{4 y \left(\left(\frac{d}{d x} y\right)^{4} + 2 \left(\frac{d}{d x} y\right)^{2} + 1\right)} = 0\right ]$$

In [7]:

Out[7]:

$$\left [ - \frac{\sqrt{2} \sqrt{\frac{\left(\frac{d}{d x} y{\left (x \right )}\right)^{2} + 1}{g y{\left (x \right )}}} \left(2 y{\left (x \right )} \frac{d^{2}}{d x^{2}} y{\left (x \right )} + \left(\frac{d}{d x} y{\left (x \right )}\right)^{2} + 1\right)}{4 \left(\left(\frac{d}{d x} y{\left (x \right )}\right)^{4} + 2 \left(\frac{d}{d x} y{\left (x \right )}\right)^{2} + 1\right) y{\left (x \right )}} = 0\right ]$$

In [26]:

from sympy.tensor import IndexedBase, Idx
from sympy import symbols
M = IndexedBase('M')
i, j = symbols('i, j', cls=Idx)
M[i]
M[i].diff(M[j])

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-26-4d1aa6970312> in <module>()
      4 i, j = symbols('i, j', cls=Idx)
      5 M[i]
----> 6 M[i].diff(M[j])

/data/apps/anaconda/envs/work/lib/python2.7/site-packages/sympy/core/expr.pyc in diff(self, *symbols, **assumptions)
   2773         new_symbols = list(map(sympify, symbols))  # e.g. x, 2, y, z
   2774         assumptions.setdefault("evaluate", True)
-> 2775         return Derivative(self, *new_symbols, **assumptions)
   2776 
   2777     ###########################################################################

/data/apps/anaconda/envs/work/lib/python2.7/site-packages/sympy/core/function.pyc in __new__(cls, expr, *variables, **assumptions)
   1016                 from sympy.utilities.misc import filldedent
   1017                 raise ValueError(filldedent('''
-> 1018                 Can\'t calculate %s-th derivative wrt %s.''' % (count, v)))
   1019 
   1020             if all_zero and not count == 0:

ValueError: 
Can't calculate 1-th derivative wrt M[j].

In [34]:

f = sympy.Function('f')
i = sympy.Symbol('i')
f(i).diff(i)

Out[34]:

$$\frac{d}{d i} f{\left (i \right )}$$

In [19]:

from sympy.physics.vector import ReferenceFrame
R = ReferenceFrame('R')

In [ ]: