# -*- coding: utf-8 -*-
from sympy import Derivative
from sympy.core.function import UndefinedFunction, AppliedUndef
from sympy.core.symbol import Symbol
from sympy.interactive.printing import init_printing
from sympy.printing.conventions import split_super_sub
from sympy.printing.latex import LatexPrinter, translate
from sympy.printing.pretty.pretty import PrettyPrinter
from sympy.printing.pretty.pretty_symbology import center_accent
from sympy.printing.str import StrPrinter
__all__ = ['vprint', 'vsstrrepr', 'vsprint', 'vpprint', 'vlatex',
'init_vprinting']
class VectorStrPrinter(StrPrinter):
"""String Printer for vector expressions. """
def _print_Derivative(self, e):
from sympy.physics.vector.functions import dynamicsymbols
t = dynamicsymbols._t
if (bool(sum([i == t for i in e.variables])) &
isinstance(type(e.args[0]), UndefinedFunction)):
ol = str(e.args[0].func)
for i, v in enumerate(e.variables):
ol += dynamicsymbols._str
return ol
else:
return StrPrinter().doprint(e)
def _print_Function(self, e):
from sympy.physics.vector.functions import dynamicsymbols
t = dynamicsymbols._t
if isinstance(type(e), UndefinedFunction):
return StrPrinter().doprint(e).replace("(%s)" % t, '')
return e.func.__name__ + "(%s)" % self.stringify(e.args, ", ")
class VectorStrReprPrinter(VectorStrPrinter):
"""String repr printer for vector expressions."""
def _print_str(self, s):
return repr(s)
class VectorLatexPrinter(LatexPrinter):
"""Latex Printer for vector expressions. """
def _print_Function(self, expr, exp=None):
from sympy.physics.vector.functions import dynamicsymbols
func = expr.func.__name__
t = dynamicsymbols._t
if hasattr(self, '_print_' + func) and \
not isinstance(type(expr), UndefinedFunction):
return getattr(self, '_print_' + func)(expr, exp)
elif isinstance(type(expr), UndefinedFunction) and (expr.args == (t,)):
name, supers, subs = split_super_sub(func)
name = translate(name)
supers = [translate(sup) for sup in supers]
subs = [translate(sub) for sub in subs]
if len(supers) != 0:
supers = r"^{%s}" % "".join(supers)
else:
supers = r""
if len(subs) != 0:
subs = r"_{%s}" % "".join(subs)
else:
subs = r""
if exp:
supers += r"^{%s}" % self._print(exp)
return r"%s" % (name + supers + subs)
else:
args = [str(self._print(arg)) for arg in expr.args]
# How inverse trig functions should be displayed, formats are:
# abbreviated: asin, full: arcsin, power: sin^-1
inv_trig_style = self._settings['inv_trig_style']
# If we are dealing with a power-style inverse trig function
inv_trig_power_case = False
# If it is applicable to fold the argument brackets
can_fold_brackets = self._settings['fold_func_brackets'] and \
len(args) == 1 and \
not self._needs_function_brackets(expr.args[0])
inv_trig_table = ["asin", "acos", "atan", "acot"]
# If the function is an inverse trig function, handle the style
if func in inv_trig_table:
if inv_trig_style == "abbreviated":
pass
elif inv_trig_style == "full":
func = "arc" + func[1:]
elif inv_trig_style == "power":
func = func[1:]
inv_trig_power_case = True
# Can never fold brackets if we're raised to a power
if exp is not None:
can_fold_brackets = False
if inv_trig_power_case:
name = r"\operatorname{%s}^{-1}" % func
elif exp is not None:
name = r"\operatorname{%s}^{%s}" % (func, exp)
else:
name = r"\operatorname{%s}" % func
if can_fold_brackets:
name += r"%s"
else:
name += r"\left(%s\right)"
if inv_trig_power_case and exp is not None:
name += r"^{%s}" % exp
return name % ",".join(args)
def _print_Derivative(self, der_expr):
from sympy.physics.vector.functions import dynamicsymbols
# make sure it is in the right form
der_expr = der_expr.doit()
if not isinstance(der_expr, Derivative):
return r"\left(%s\right)" % self.doprint(der_expr)
# check if expr is a dynamicsymbol
t = dynamicsymbols._t
expr = der_expr.expr
red = expr.atoms(AppliedUndef)
syms = der_expr.variables
test1 = not all([True for i in red if i.free_symbols == {t}])
test2 = not all([(t == i) for i in syms])
if test1 or test2:
return LatexPrinter().doprint(der_expr)
# done checking
dots = len(syms)
base = self._print_Function(expr)
base_split = base.split('_', 1)
base = base_split[0]
if dots == 1:
base = r"\dot{%s}" % base
elif dots == 2:
base = r"\ddot{%s}" % base
elif dots == 3:
base = r"\dddot{%s}" % base
elif dots == 4:
base = r"\ddddot{%s}" % base
else: # Fallback to standard printing
return LatexPrinter().doprint(der_expr)
if len(base_split) is not 1:
base += '_' + base_split[1]
return base
class VectorPrettyPrinter(PrettyPrinter):
"""Pretty Printer for vectorialexpressions. """
def _print_Derivative(self, deriv):
from sympy.physics.vector.functions import dynamicsymbols
# XXX use U('PARTIAL DIFFERENTIAL') here ?
t = dynamicsymbols._t
dot_i = 0
syms = list(reversed(deriv.variables))
while len(syms) > 0:
if syms[-1] == t:
syms.pop()
dot_i += 1
else:
return super(VectorPrettyPrinter, self)._print_Derivative(deriv)
if not (isinstance(type(deriv.expr), UndefinedFunction)
and (deriv.expr.args == (t,))):
return super(VectorPrettyPrinter, self)._print_Derivative(deriv)
else:
pform = self._print_Function(deriv.expr)
# the following condition would happen with some sort of non-standard
# dynamic symbol I guess, so we'll just print the SymPy way
if len(pform.picture) > 1:
return super(VectorPrettyPrinter, self)._print_Derivative(deriv)
# There are only special symbols up to fourth-order derivatives
if dot_i >= 5:
return super(VectorPrettyPrinter, self)._print_Derivative(deriv)
# Deal with special symbols
dots = {0 : u"",
1 : u"\N{COMBINING DOT ABOVE}",
2 : u"\N{COMBINING DIAERESIS}",
3 : u"\N{COMBINING THREE DOTS ABOVE}",
4 : u"\N{COMBINING FOUR DOTS ABOVE}"}
d = pform.__dict__
#if unicode is false then calculate number of apostrophes needed and add to output
if not self._use_unicode:
apostrophes = ""
for i in range(0, dot_i):
apostrophes += "'"
d['picture'][0] += apostrophes + "(t)"
else:
d['picture'] = [center_accent(d['picture'][0], dots[dot_i])]
d['unicode'] = center_accent(d['unicode'], dots[dot_i])
return pform
def _print_Function(self, e):
from sympy.physics.vector.functions import dynamicsymbols
t = dynamicsymbols._t
# XXX works only for applied functions
func = e.func
args = e.args
func_name = func.__name__
pform = self._print_Symbol(Symbol(func_name))
# If this function is an Undefined function of t, it is probably a
# dynamic symbol, so we'll skip the (t). The rest of the code is
# identical to the normal PrettyPrinter code
if not (isinstance(func, UndefinedFunction) and (args == (t,))):
return super(VectorPrettyPrinter, self)._print_Function(e)
return pform
[docs]def vprint(expr, **settings):
r"""Function for printing of expressions generated in the
sympy.physics vector package.
Extends SymPy's StrPrinter, takes the same setting accepted by SymPy's
`sstr()`, and is equivalent to `print(sstr(foo))`.
Parameters
==========
expr : valid SymPy object
SymPy expression to print.
settings : args
Same as the settings accepted by SymPy's sstr().
Examples
========
>>> from sympy.physics.vector import vprint, dynamicsymbols
>>> u1 = dynamicsymbols('u1')
>>> print(u1)
u1(t)
>>> vprint(u1)
u1
"""
outstr = vsprint(expr, **settings)
from sympy.core.compatibility import builtins
if (outstr != 'None'):
builtins._ = outstr
print(outstr)
def vsstrrepr(expr, **settings):
"""Function for displaying expression representation's with vector
printing enabled.
Parameters
==========
expr : valid SymPy object
SymPy expression to print.
settings : args
Same as the settings accepted by SymPy's sstrrepr().
"""
p = VectorStrReprPrinter(settings)
return p.doprint(expr)
def vsprint(expr, **settings):
r"""Function for displaying expressions generated in the
sympy.physics vector package.
Returns the output of vprint() as a string.
Parameters
==========
expr : valid SymPy object
SymPy expression to print
settings : args
Same as the settings accepted by SymPy's sstr().
Examples
========
>>> from sympy.physics.vector import vsprint, dynamicsymbols
>>> u1, u2 = dynamicsymbols('u1 u2')
>>> u2d = dynamicsymbols('u2', level=1)
>>> print("%s = %s" % (u1, u2 + u2d))
u1(t) = u2(t) + Derivative(u2(t), t)
>>> print("%s = %s" % (vsprint(u1), vsprint(u2 + u2d)))
u1 = u2 + u2'
"""
string_printer = VectorStrPrinter(settings)
return string_printer.doprint(expr)
[docs]def vpprint(expr, **settings):
r"""Function for pretty printing of expressions generated in the
sympy.physics vector package.
Mainly used for expressions not inside a vector; the output of running
scripts and generating equations of motion. Takes the same options as
SymPy's pretty_print(); see that function for more information.
Parameters
==========
expr : valid SymPy object
SymPy expression to pretty print
settings : args
Same as those accepted by SymPy's pretty_print.
"""
pp = VectorPrettyPrinter(settings)
# Note that this is copied from sympy.printing.pretty.pretty_print:
# XXX: this is an ugly hack, but at least it works
use_unicode = pp._settings['use_unicode']
from sympy.printing.pretty.pretty_symbology import pretty_use_unicode
uflag = pretty_use_unicode(use_unicode)
try:
return pp.doprint(expr)
finally:
pretty_use_unicode(uflag)
[docs]def vlatex(expr, **settings):
r"""Function for printing latex representation of sympy.physics.vector
objects.
For latex representation of Vectors, Dyadics, and dynamicsymbols. Takes the
same options as SymPy's latex(); see that function for more information;
Parameters
==========
expr : valid SymPy object
SymPy expression to represent in LaTeX form
settings : args
Same as latex()
Examples
========
>>> from sympy.physics.vector import vlatex, ReferenceFrame, dynamicsymbols
>>> N = ReferenceFrame('N')
>>> q1, q2 = dynamicsymbols('q1 q2')
>>> q1d, q2d = dynamicsymbols('q1 q2', 1)
>>> q1dd, q2dd = dynamicsymbols('q1 q2', 2)
>>> vlatex(N.x + N.y)
'\\mathbf{\\hat{n}_x} + \\mathbf{\\hat{n}_y}'
>>> vlatex(q1 + q2)
'q_{1} + q_{2}'
>>> vlatex(q1d)
'\\dot{q}_{1}'
>>> vlatex(q1 * q2d)
'q_{1} \\dot{q}_{2}'
>>> vlatex(q1dd * q1 / q1d)
'\\frac{q_{1} \\ddot{q}_{1}}{\\dot{q}_{1}}'
"""
latex_printer = VectorLatexPrinter(settings)
return latex_printer.doprint(expr)
def init_vprinting(**kwargs):
"""Initializes time derivative printing for all SymPy objects, i.e. any
functions of time will be displayed in a more compact notation. The main
benefit of this is for printing of time derivatives; instead of
displaying as ``Derivative(f(t),t)``, it will display ``f'``. This is
only actually needed for when derivatives are present and are not in a
physics.vector.Vector or physics.vector.Dyadic object. This function is a
light wrapper to `sympy.interactive.init_printing`. Any keyword
arguments for it are valid here.
{0}
Examples
========
>>> from sympy import Function, symbols
>>> from sympy.physics.vector import init_vprinting
>>> t, x = symbols('t, x')
>>> omega = Function('omega')
>>> omega(x).diff()
Derivative(omega(x), x)
>>> omega(t).diff()
Derivative(omega(t), t)
Now use the string printer:
>>> init_vprinting(pretty_print=False)
>>> omega(x).diff()
Derivative(omega(x), x)
>>> omega(t).diff()
omega'
"""
kwargs['str_printer'] = vsstrrepr
kwargs['pretty_printer'] = vpprint
kwargs['latex_printer'] = vlatex
init_printing(**kwargs)
params = init_printing.__doc__.split('Examples\n ========')[0]
init_vprinting.__doc__ = init_vprinting.__doc__.format(params)
