ENH: Add support for ipython latex printing to polynomial

Choices made, and the alternatives rejected (for no particularly strong reason):

1. Show terms in ascending order, to match their internal representation
  * alternative: descending, to match convention
2. Shows 0 terms in gray
  * alternative: omit entirely
  * alternative: show normally to aid comparison
3. Write each term as `basis(ax + b)
  * alternative: write as `basis(u) ... where u = ax + b`
  * alternative: show the normalized polynomial

In future it would perhaps make sense to expose these options to the end user

8 files changed, 156 insertions, 2 deletions

diff --git a/numpy/polynomial/_polybase.py b/numpy/polynomial/_polybase.py
index 78392d2a2..9f4d30e53 100644
--- a/numpy/polynomial/_polybase.py
+++ b/numpy/polynomial/_polybase.py
@@ -9,7 +9,7 @@ abc module from the stdlib, hence it is only available for Python >= 2.6.
 from __future__ import division, absolute_import, print_function
 
 from abc import ABCMeta, abstractmethod, abstractproperty
-from numbers import Number
+import numbers
 
 import numpy as np
 from . import polyutils as pu
@@ -82,6 +82,10 @@ class ABCPolyBase(object):
     def nickname(self):
         pass
 
+    @abstractproperty
+    def basis_name(self):
+        pass
+
     @abstractmethod
     def _add(self):
         pass
@@ -273,6 +277,89 @@ class ABCPolyBase(object):
         name = self.nickname
         return format % (name, coef)
 
+    @classmethod
+    def _repr_latex_term(cls, i, arg_str, needs_parens):
+        if cls.basis_name is None:
+            raise NotImplementedError(
+                "Subclasses must define either a basis name, or override "
+                "_repr_latex_term(i, arg_str, needs_parens)")
+        # since we always add parens, we don't care if the expression needs them
+        return "{{{basis}}}_{{{i}}}({arg_str})".format(
+            basis=cls.basis_name, i=i, arg_str=arg_str
+        )
+
+    @staticmethod
+    def _repr_latex_scalar(x):
+        # TODO: we're stuck with disabling math formatting until we handle
+        # exponents in this function
+        return r'\text{{{}}}'.format(x)
+
+    def _repr_latex_(self):
+        # get the scaled argument string to the basis functions
+        off, scale = self.mapparms()
+        if off == 0 and scale == 1:
+            term = 'x'
+            needs_parens = False
+        elif scale == 1:
+            term = '{} + x'.format(
+                self._repr_latex_scalar(off)
+            )
+            needs_parens = True
+        elif off == 0:
+            term = '{}x'.format(
+                self._repr_latex_scalar(scale)
+            )
+            needs_parens = True
+        else:
+            term = '{} + {}x'.format(
+                self._repr_latex_scalar(off),
+                self._repr_latex_scalar(scale)
+            )
+            needs_parens = True
+
+        # filter out uninteresting coefficients
+        filtered_coeffs = [
+            (i, c)
+            for i, c in enumerate(self.coef)
+            # if not (c == 0)  # handle NaN
+        ]
+
+        mute = r"\color{{LightGray}}{{{}}}".format
+
+        parts = []
+        for i, c in enumerate(self.coef):
+            # prevent duplication of + and - signs
+            if i == 0:
+                coef_str = '{}'.format(self._repr_latex_scalar(c))
+            elif not isinstance(c, numbers.Real):
+                coef_str = ' + ({})'.format(self._repr_latex_scalar(c))
+            elif not np.signbit(c):
+                coef_str = ' + {}'.format(self._repr_latex_scalar(c))
+            else:
+                coef_str = ' - {}'.format(self._repr_latex_scalar(-c))
+
+            # produce the string for the term
+            term_str = self._repr_latex_term(i, term, needs_parens)
+            if term_str == '1':
+                part = coef_str
+            else:
+                part = r'{}\,{}'.format(coef_str, term_str)
+
+            if c == 0:
+                part = mute(part)
+
+            parts.append(part)
+
+        if parts:
+            body = ''.join(parts)
+        else:
+            # in case somehow there are no coefficients at all
+            body = '0'
+
+        return r'$x \mapsto {}$'.format(body)
+
+
+
     # Pickle and copy
 
     def __getstate__(self):
@@ -338,7 +425,7 @@ class ABCPolyBase(object):
         # there is no true divide if the rhs is not a Number, although it
         # could return the first n elements of an infinite series.
         # It is hard to see where n would come from, though.
-        if not isinstance(other, Number) or isinstance(other, bool):
+        if not isinstance(other, numbers.Number) or isinstance(other, bool):
             form = "unsupported types for true division: '%s', '%s'"
             raise TypeError(form % (type(self), type(other)))
         return self.__floordiv__(other)
diff --git a/numpy/polynomial/chebyshev.py b/numpy/polynomial/chebyshev.py
index 946e0499c..f2162a894 100644
--- a/numpy/polynomial/chebyshev.py
+++ b/numpy/polynomial/chebyshev.py
@@ -2188,3 +2188,4 @@ class Chebyshev(ABCPolyBase):
     nickname = 'cheb'
     domain = np.array(chebdomain)
     window = np.array(chebdomain)
+    basis_name = 'T'
diff --git a/numpy/polynomial/hermite.py b/numpy/polynomial/hermite.py
index 75c7e6832..8c33ee863 100644
--- a/numpy/polynomial/hermite.py
+++ b/numpy/polynomial/hermite.py
@@ -1851,3 +1851,4 @@ class Hermite(ABCPolyBase):
     nickname = 'herm'
     domain = np.array(hermdomain)
     window = np.array(hermdomain)
+    basis_name = 'H'
diff --git a/numpy/polynomial/hermite_e.py b/numpy/polynomial/hermite_e.py
index 125364a11..6166c03fd 100644
--- a/numpy/polynomial/hermite_e.py
+++ b/numpy/polynomial/hermite_e.py
@@ -1848,3 +1848,4 @@ class HermiteE(ABCPolyBase):
     nickname = 'herme'
     domain = np.array(hermedomain)
     window = np.array(hermedomain)
+    basis_name = 'He'
diff --git a/numpy/polynomial/laguerre.py b/numpy/polynomial/laguerre.py
index 2b9757ab8..0e4554071 100644
--- a/numpy/polynomial/laguerre.py
+++ b/numpy/polynomial/laguerre.py
@@ -1801,3 +1801,4 @@ class Laguerre(ABCPolyBase):
     nickname = 'lag'
     domain = np.array(lagdomain)
     window = np.array(lagdomain)
+    basis_name = 'L'
diff --git a/numpy/polynomial/legendre.py b/numpy/polynomial/legendre.py
index a83c5735f..0a20707e6 100644
--- a/numpy/polynomial/legendre.py
+++ b/numpy/polynomial/legendre.py
@@ -1831,3 +1831,4 @@ class Legendre(ABCPolyBase):
     nickname = 'leg'
     domain = np.array(legdomain)
     window = np.array(legdomain)
+    basis_name = 'P'
diff --git a/numpy/polynomial/polynomial.py b/numpy/polynomial/polynomial.py
index adbf30234..7c43e658a 100644
--- a/numpy/polynomial/polynomial.py
+++ b/numpy/polynomial/polynomial.py
@@ -1643,3 +1643,15 @@ class Polynomial(ABCPolyBase):
     nickname = 'poly'
     domain = np.array(polydomain)
     window = np.array(polydomain)
+    basis_name = None
+
+    @staticmethod
+    def _repr_latex_term(i, arg_str, needs_parens):
+        if needs_parens:
+            arg_str = r'\left({}\right)'.format(arg_str)
+        if i == 0:
+            return '1'
+        elif i == 1:
+            return arg_str
+        else:
+            return '{}^{{{}}}'.format(arg_str, i)
diff --git a/numpy/polynomial/tests/test_classes.py b/numpy/polynomial/tests/test_classes.py
index 738741668..15e24f92b 100644
--- a/numpy/polynomial/tests/test_classes.py
+++ b/numpy/polynomial/tests/test_classes.py
@@ -562,6 +562,56 @@ def test_ufunc_override(Poly):
     assert_raises(TypeError, np.add, x, p)
 
 
+
+class TestLatexRepr(object):
+    """Test the latex repr used by ipython """
+
+    def as_latex(self, obj):
+        # right now we ignore the formatting of scalars in our tests, since
+        # it makes them too verbose. Ideally, the formatting of scalars will
+        # be fixed such that tests below continue to pass
+        obj._repr_latex_scalar = lambda x: str(x)
+        try:
+            return obj._repr_latex_()
+        finally:
+            del obj._repr_latex_scalar
+
+    def test_simple_polynomial(self):
+        # default input
+        p = Polynomial([1, 2, 3])
+        assert_equal(self.as_latex(p),
+            r'$x \mapsto 1.0 + 2.0\,x + 3.0\,x^{2}$')
+
+        # translated input
+        p = Polynomial([1, 2, 3], domain=[-2, 0])
+        assert_equal(self.as_latex(p),
+            r'$x \mapsto 1.0 + 2.0\,\left(1.0 + x\right) + 3.0\,\left(1.0 + x\right)^{2}$')
+
+        # scaled input
+        p = Polynomial([1, 2, 3], domain=[-0.5, 0.5])
+        assert_equal(self.as_latex(p),
+            r'$x \mapsto 1.0 + 2.0\,\left(2.0x\right) + 3.0\,\left(2.0x\right)^{2}$')
+
+        # affine input
+        p = Polynomial([1, 2, 3], domain=[-1, 0])
+        assert_equal(self.as_latex(p),
+            r'$x \mapsto 1.0 + 2.0\,\left(1.0 + 2.0x\right) + 3.0\,\left(1.0 + 2.0x\right)^{2}$')
+
+    def test_basis_func(self):
+        p = Chebyshev([1, 2, 3])
+        assert_equal(self.as_latex(p),
+            r'$x \mapsto 1.0\,{T}_{0}(x) + 2.0\,{T}_{1}(x) + 3.0\,{T}_{2}(x)$')
+        # affine input - check no surplus parens are added
+        p = Chebyshev([1, 2, 3], domain=[-1, 0])
+        assert_equal(self.as_latex(p),
+            r'$x \mapsto 1.0\,{T}_{0}(1.0 + 2.0x) + 2.0\,{T}_{1}(1.0 + 2.0x) + 3.0\,{T}_{2}(1.0 + 2.0x)$')
+
+    def test_multichar_basis_func(self):
+        p = HermiteE([1, 2, 3])
+        assert_equal(self.as_latex(p),
+            r'$x \mapsto 1.0\,{He}_{0}(x) + 2.0\,{He}_{1}(x) + 3.0\,{He}_{2}(x)$')
+
+
 #
 # Test class method that only exists for some classes
 #