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

# # Interoperability with SymPy
# 
# ```{versionadded} 0.10.0
# 
# ```
# 
# Starting with version 0.10, symbolic expressions in heyoka.py can be converted to/from [SymPy](https://www.sympy.org/en/index.html) expressions. This feature can be very useful, as it unlocks the possibility of using the wide arsenal of SymPy algorithms on heyoka.py's expressions. Possible use cases include:
# 
# - the automatic simplification of heyoka.py's expressions,
# - the application of symbolic algorithms not available in heyoka.py (e.g., symbolic integration),
# - LaTeX pretty-printing.
# 
# In this tutorial, we will show how SymPy interoperability can be used for expression simplification purposes, which ultimately leads to a measurable performance increase in the numerical integration of a system of ODEs.
# 
# For this example, we will consider the relativistic dynamics of Mercury's orbit around the Sun, which, in the first post-Newtonian approximation, can be described by the following Hamiltonian:
# 
# $$
# \mathcal{H}_\mathrm{1PN} \left(v_x, v_y, v_z, x, y, z \right) = \frac{1}{2}v^2-\frac{\mu}{r} + \varepsilon \left(\frac{\mu^2}{2r^2} -\frac{1}{8}v^4-\frac{3}{2}\frac{\mu v^2}{r} \right).
# $$
# 
# where $\left(v_x, v_y, v_z, x, y, z \right)$ is Mercury's Cartesian state vector with respect to the Sun, $\varepsilon = 1/c^2$, $v=\sqrt{v_x^2+v_y^2+v_z^2}$, $r=\sqrt{x^2+y^2+z^2}$ and $\mu=GM_\odot$ is the gravitational parameter of the system. See the [tutorial about Mercury's relativistic precession](<./mercury_precession.ipynb>) for more details about this dynamical system.
# 
# Let us begin by defining the Hamiltonian as a symbolic expression in heyoka.py's expression system:

# In[1]:


import heyoka as hy

# Create the symbolic variables.
vx, vy, vz, x, y, z = hy.make_vars("v_x", "v_y", "v_z", "x", "y", "z")

# mu and eps constants.
mu = 0.01720209895 * 0.01720209895 * 365 * 365
eps = 2.5037803127808595e-10

# Auxiliary quantities.
v2 = vx * vx + vy * vy + vz * vz
r2 = x * x + y * y + z * z
r = hy.sqrt(r2)

# Initial conditions corresponding (roughly)
# to Mercury's current orbit.
ic = (
    -10.461611630114545,
    6.667253667139184,
    0.818635813947965,
    0.16614243942411336,
    0.2568228239702581,
    0.0315338776710321,
)

# Define the Hamiltonian.
Ham = (
    1.0 / 2 * v2
    - mu / r
    + eps * (mu**2 / (2 * r2) - 1 / 8.0 * v2 * v2 - 3.0 / 2.0 * mu * v2 / r)
)


# As usual, we are using Solar masses, astronomical units and years as units of measure.
# 
# Let us now convert the heyoka.py Hamiltonian to a SymPy expression via the ``to_sympy()`` function and pretty-print it to screen:

# In[2]:


hy.to_sympy(Ham)


# The SymPy expression can be converted back to heyoka.py using the ``from_sympy()`` function:

# In[3]:


hy.from_sympy(hy.to_sympy(Ham))


# In order to formulate Hamilton's equations for this system we need, as usual, to differentiate the Hamiltonian with respect to the state variables. Let us do it with heyoka.py's ``diff()`` function:

# In[4]:


# Formulate the equations of motion
# using heyoka.py's expression system.
ta = hy.taylor_adaptive(
    # Hamilton's equations.
    [
        (vx, -hy.diff(Ham, x)),
        (vy, -hy.diff(Ham, y)),
        (vz, -hy.diff(Ham, z)),
        (x, hy.diff(Ham, vx)),
        (y, hy.diff(Ham, vy)),
        (z, hy.diff(Ham, vz)),
    ],
    # Initial conditions.
    ic,
)


# heyoka.py's ``diff()`` function applies elementary calculus rules in order to differentiate a symbolic expression, performing only very basic simplifications in the process. We can get an estimate of the complexity of the resulting expressions by taking a look at the size of the Taylor decomposition for this dynamical system (this is the number of elementary subexpressions in which the equations of motion have been decomposed):

# In[5]:


# NOTE: subtract 12 to avoid counting the 6 definitions of the state variables
# and the 6 definitions of their differential equations.
len(ta.decomposition) - 12


# Let us now try to simplify Hamilton's equations of motion using SymPy. We first define a small wrapper to invoke SymPy's ``simplify()`` function on a heyoka.py expression:

# In[6]:


import sympy


def simplify(ex):
    return hy.from_sympy(sympy.simplify(hy.to_sympy(ex)))


# Then, we create a new integrator with the simplified equations of motion:

# In[7]:


ta_spy = hy.taylor_adaptive(
    # Hamilton's equations.
    [
        (vx, simplify(-hy.diff(Ham, x))),
        (vy, simplify(-hy.diff(Ham, y))),
        (vz, simplify(-hy.diff(Ham, z))),
        (x, simplify(hy.diff(Ham, vx))),
        (y, simplify(hy.diff(Ham, vy))),
        (z, simplify(hy.diff(Ham, vz))),
    ],
    # Initial conditions.
    ic,
)


# Let us take a look at the size of the Taylor decomposition for the new integrator:

# In[8]:


len(ta_spy.decomposition) - 12


# In other words, thanks to SymPy's expression simplification capabilities, we reduced by $\sim 50\%$ the symbolic complexity (i.e., the number of elementary subexpressions) of our ODE system.
# 
# What is the performance impact of these simplifications? Let us find out:

# In[9]:


time ta_spy.propagate_until(10000.)


# In[10]:


time ta.propagate_until(10000.)


# The performance increase is measurable but not as large as the reduction in symbolic complexity.
# 
# As a word of warning, users should be aware that the general-purpose ``simplify()`` function from SymPy can be quite expensive from a computational point of view, especially for large expressions. SymPy provides various specialised [simplification primitives](https://docs.sympy.org/latest/tutorial/simplification.html) that can be considerably more efficient.
# 
# ## More details on the conversion process
# 
# When converting a heyoka.py expression to SymPy, the following conventions are adopted:
# 
# * variables are converted to SymPy ``Symbol``s,
# * constants are converted to SymPy numbers,
# * [runtime parameters](<./ODEs with parameters.ipynb>) are converted to SymPy ``Symbol``s conventionally called ``"par[n]"``, where ``n`` is the parameter index,
# * functions are mapped to the corresponding SymPy functions.
# 
# Similarly, when converting a SymPy expression to heyoka.py, the following conventions are adopted:
# 
# * SymPy ``Symbol``s are converted to heyoka.py variables, unless the name of the symbol is ``"par[n]"``, in which case the ``Symbol`` is converted to a runtime parameter,
# * SymPy numbers are converted to heyoka.py constants, if the conversion is lossless (otherwise an error is raised),
# * SymPy functions are mapped to heyoka.py functions.
# 
# When converting a SymPy expression to heyoka.py, it is possible to customise the conversion process by passing a substitution dictionary to the ``from_sympy()`` function. Here's an example:

# In[11]:


# Create a few SymPy symbols.
x, y, z = sympy.symbols("x y z")

# Convert the expression x*(y+z) to heyoka.py,
# but provide a custom substitution rule for
# the subexpression y+z.
hy.from_sympy(x * (y + z), {y + z: hy.expression("t")})


# Note that the same effect could be achieved by doing the substitution in SymPy *before* the conversion to heyoka.py. However, for large expressions, using the custom substitution dictionary in the ``from_sympy()`` function will perform better.
# 
# ## Handling rational numbers
# 
# heyoka.py is very conservative when converting from SymPy expressions containing rational numbers. In particular, if a rational number cannot be converted *exactly* to a floating-point value, heyoka.py will raise an error. In order to avoid this, you can use the ``sympy.N`` function to forcibly convert all rational numbers to floating-point values:

# In[12]:


hy.from_sympy(sympy.N(x / 3))