# sym.py
'''
Symbolic maths with sympy.
'''

import sympy as sp

#
# Preliminiaries.
#

# Define the symbols and expressions.
x, y = sp.symbols('x y')
expr1 = x + 2*y
expr2 = sp.sin(x) + y
print(expr1)
print(expr2)
# Expressions are immutable.

# You can call the symbols any name.
# But better choose consistent names.
x, y = sp.symbols('x z')

# Expressions don't change if we change the symbols objects.
x = 4
print(expr1)

# To perform such a change use 'subs'.
x = sp.symbols('x')
expr1.subs(x, 2)

# expand and factor.
expr3 = sp.expand(expr1*x)
sp.factor(expr3)

# Equalities.
x + 1 == 4  # This represents 'exact structural equality testing'.
# Need to do:
expr = sp.Eq(x + 1, 4)
# Symbolic comparison:
a = (x + 1)**2
b = x**2 + 2*x + 1
sp.simplify(a - b)
# Comparing numerically:
a = sp.cos(x)**2 - sp.sin(x)**2
b = sp.cos(2*x)
a.equals(b)


#
# Operations on expressions.
#

# substitution
x = sp.symbols('x')
expr = sp.cos(x) + 1
expr.subs(x, 1)
expr.subs(x, sp.sin(x))

y = sp.symbols('y')
expr = sp.sin(x) + sp.cos(y)
expr.subs([(x, 2), (y, 3)])

# Strings to sympy expressions.
str_exp = "x**2 + 3*x + 1/2"
expr = sp.sympify(str_exp)

# Evaluate numerically.
expr = expr.subs(x, 3)
expr.evalf()
sp.pi.evalf(100)
expr = x**2 + 3*x + 1
expr.evalf(subs={x: 2.4})


#
# Lambdify.
#

expr = sp.sin(x)
f = sp.lambdify(x, expr, "numpy")
t = np.linspace(0, 7, 20)
f(t)


#
# Printing.
#

# Set up printing.
sp.init_printing()
sp.Integral(sp.sqrt(1/x), x)

# Pretty printer.
sp.pprint(sp.Integral(sp.sqrt(1/x), x, use_unicode=True))

# LaTex
print(sp.latex(sp.Integral(sp.sqrt(1/x), x)))


#
# Simplifications.
#

sp.simplify(sp.cos(x)**2 + sp.sin(x)**2)

sp.factor(x**2 + 2*x + 1)

sp.expand((x + 2)**2)

# Collect common powers of a term.
x, y, z = sp.symbols('x y z')
expr = x*y + x - 3 + 2*x**2 - z*x**2 + x**3
sp.collect(expr, x)

# Simplify trigonometric expressions.
expr = sp.sin(x)**4 - 2*sp.cos(x)**2*sp.sin(x)**2 + sp.cos(x)**4
sp.trigsimp(expr)


#
# Calculus.
#

# Derivatives.
sp.diff(sp.cos(x), x)
sp.diff(sp.cos(x), x, x)
sp.diff(sp.cos(x*y), x, y)
sp.diff(sp.cos(x), x, 2)
sp.diff(sp.cos(x*y), x, 3, y, 2)

# Non-performed derivatives.
deriv = sp.Derivative(sp.cos(x), x)
# Perform derivative.
deriv.doit()

# Integrate.
sp.integrate(sp.cos(x), x)
sp.integrate(x**2, (x, 0, 1))
sp.integrate(sp.exp(-x), (x, 0, sp.oo))
sp.integrate(sp.exp(-x**2 - y**2), (x, -sp.oo, sp.oo), (y, -sp.oo, sp.oo))

# Unvalued integral.
expr = sp.Integral(sp.log(x)**2, x)
# Perform integral.
expr.doit()

# Limits.
sp.limit(sp.sin(x)/x, x, 0)
sp.limit(sp.exp(-x), x, sp.oo)

# Unvalued limit.
expr = sp.Limit(so.sin(x)/x, x, 0)
expr.doit()

# Directions.
sp.limit(abs(x)/x, x, 0, '+')
sp.limit(abs(x)/x, x, 0, '-')

# Series expansion.
expr = sp.exp(sp.sin(x))
expr.series(x, 0, 4)


#
# Solvers.
#

# Define equation.
# !!! Don't use '=' or '==' !!!
sp.Eq(x, y)
sp.solveset(sp.Eq(x**2, 1), x)
# or set difference to 0
sp.solveset(sp.Eq(x**2 - 1, 0), x)
sp.solveset(x**2 - 1, x)

# Specify the domain.
sp.solveset(x**2 - x, x, domain=sp.S.Reals)
# If not able to find solution:
sp.solveset(sp.exp(x), x)
sp.solveset(sp.cos(x) - x, x)

# solveset not for non-linear multivariate systems
# use solve instead
sp.solve([x*y - 1, x - 2], x, y)

# for a linear system of equations
sp.linsolve([x + y + z - 1, x + y + 2*z - 4], (x, y, z))


#
# PDEs
#

# Define functions.
f, g = sp.symbols('f g', cls=sp.Function)
f = sp.Function('f')
g = sp.Function('g')(x)
f(x).diff(x)
g.diff(x)

# Define the PDE.
diffeq = sp.Eq(f(x).diff(x, x) + 2*f(x).diff(x) + f(x), sp.sin(x))
sp.dsolve(diffeq, f(x))


#
# Plotting.
#

# Plot symbolic functions.
# Uses matplotlib in the background.
x = sp.symbols('x')
sp.plotting.plot(x)
sp.plotting.plot(x**2, (x, -5, 5))
sp.plotting.plot(x, x**2, x**3, (x, -5, 5))

# parametric plot
u = sp.symbols('u')
sp.plotting.plot_parametric(sp.cos(u), sp.sin(u), (u, -5, 5))

# 3d plots
x, y = sp.symbols('x y')
sp.plotting.plot3d(x*y, (x, -5, 5), (y, -5, 5))


#
# Vector Calculus.
#

from sympy.physics.vector import ReferenceFrame
from sympy.physics.vector import curl
from sympy.physics.vector import cross
from sympy.physics.vector import divergence
from sympy.physics.vector import dot
from sympy.physics.vector import gradient
from sympy.physics.vector import vprint

# Define the spatial coordinates.
ee = ReferenceFrame('e')
t = sp.Symbol('t')
# Coordinate x.
ee[0]
# Unit vector.
ee.x

# Define the fields: magnetic, current, Lorentz force, rotation Lorentz
# and the velocity field.
bx = sp.Function('bx')(ee[0], ee[1], ee[2], t)
by = sp.Function('by')(ee[0], ee[1], ee[2], t)
bz = sp.Function('bz')(ee[0], ee[1], ee[2], t)
bb = bx*ee.x + by*ee.y + bz*ee.z
jj = curl(bb, ee)
ff_l = cross(jj, bb)
rot_f_l = curl(ff_l, ee)

# Check if the magnetic field is divergence free.
div_bb = divergence(bb, ee)
print "div(bb) = ", div_bb

# Check if the rotational component of the Lorentz force is finite.
print "curl(ff_l) = ", rot_f_l