Source code for sympy.physics.vector.printing

# -*- coding: utf-8 -*-

from sympy import Derivative
from sympy.core.function import UndefinedFunction
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.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):
            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":
                    func = func
                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
        from sympy.core.function import AppliedUndef
        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
        if len(base_split) is not 1:
            base += '_' + base_split[1]
        return base

    def parenthesize(self, item, level, strict=False):
        item_latex = self._print(item)
        if item_latex.startswith(r"\dot") or item_latex.startswith(r"\ddot") or item_latex.startswith(r"\dddot"):
            return self._print(item)
        else:
            return LatexPrinter.parenthesize(self, item, level, strict)


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
        can_break = True
        syms = list(reversed(deriv.variables))
        x = None

        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)

        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__
        pic = d['picture'][0]
        uni = d['unicode']
        lp = len(pic) // 2 + 1
        lu = len(uni) // 2 + 1
        pic_split = [pic[:lp], pic[lp:]]
        uni_split = [uni[:lu], uni[lu:]]

        d['picture'] = [pic_split[0] + dots[dot_i] + pic_split[1]]
        d['unicode'] =  uni_split[0] + dots[dot_i] + uni_split[1]

        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)