# -*- 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 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) 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) 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)