How do I extend SymPy pretty printing for new structures in Jupyter notebook?

Question

Note: This is a duplicate of this question in Math Stack Exchange. I had to first pose the question in Math StackExchange because StackOverflow doesn't have MathJax. However, almost all SymPy questions are on StackOverflow. Please refer to the Math Stack Exchange version for the typesetting of the desired output. An editor on Math StackExchange suggested I cross-post it here.

In Jupyter notebook if we execute this code:

import sympy as sp
sp.init_printing()
x,y=sp.symbols('x,y')
x**2+sp.sin(y)

We will get a nice output, with no further coding, due to SymPy's pretty printing process, that looks like

Now suppose we do:

class MetricSpace:
    def __init__(self, M, d):
        self.M = M
        self.d = d
        
    def __repr__(self):
        return f"$({self.M}, {self.d})$ {self.__class__.__name__}"

Re,d=sp.symbols(r'\Re,d')

MetricSpace(Re,d)

Then the output we get is

If we do instead

from IPython.core.display import Markdown
Markdown(repr(MetricSpace(Re,d)))

then I get the desired output, which looks like

How do we code the above so that SymPy's pretty printer provides the desired output in Jupyter notebook without having to wrap it in Markdown(repr(...))?

Answer 1

If all you care about is integration with Jupyter notebook, then you want the IPython _repr_latex_ hook:

class MetricSpace:
   def _repr_latex_(self):
       # this is text-mode latex, so needs $ to enter math mode
       return f"$({self.M}, {self.d})$ {self.__class__.__name__}"

This will make MetricSpace(...) show as latex in a jupter notebook, and is completely independent of sympy.

---

Separately, sympy has a latex function. If you want to support that, you need to implement _latex:

class MetricSpace:
   def _latex(self, printer):
       # this is math-mode latex, so needs \text to enter text mode
       # if the class members support this hook, you can use `printer._print` to recurse
       return f"{printer._print(self.M)}, {printer._print(self.d)}$ \text{{{self.__class__.__name__}}}"

This will make sympy.latex(MetricSpace(...)) work.

---

Finally, there's the sympy integration that connects these two modes together, after init_printing(use_latex='mathjax') is called. This has changed between Sympy 1.6 and 1.7.

• In Sympy 1.6, an ipython latex formatter is provided for
  • subclasses of sympy.Basic (and some others)
  • lists, sets, etc of those objects
• In Sympy 1.7, an ipython latex formatter is provided for
  • subclasses of sympy.printing.defaults.Printable
  • lists, sets, etc of any object that implements _latex or subclasses Printable

Subclassing Expr is a bad idea, as that adds tens of possibly meaningless methods and operator overloads to your class. Subclassing Basic is less bad, but suffers from a lesser version of the same problem if your class is not intended to be used within sympy

Answer 2

Note: I have already accepted the answer from Izaak van Dongen above. This answer is to provide context for the question for those that may be interested.

I was reading a book on SDEs and I asked a question on the measure theory presentation of probability space. There are a large number of structural definitions involved. I wanted to use SymPy to organize them and present the type signatures in something mimicking Axiom style.

To do this pleasantly in Python 3 and SymPy I needed a few things:

• A way of dynamically enforcing function signatures.
• A way of pretty-printing complex algebraic type signatures (this question).

I started implementing the definitions. To check that they were organized correctly, I asked

• A question about extending the concept of distribution function to any ordered space, which naturally arose in formalizing the definitions for probability space. I don't think anybody got, outside of the SymPy formalization context, why this question was natural or interesting.

• A question confirming that my stack of definitions for probability space was correct (without bringing SymPy into it).

With that in hand, and the above solution for the pretty-printing, the following few definitions give the style of my solution (without quoting the whole thing which is about 85 definitions):

import sympy as sp  # I am at version 1.6.1
from typen import strict_type_hints, enforce_type_hints
from traits.api import Array, Either, Enum, Instance, Int, Str, Tuple

class Concept(sp.Expr):
    def __init__(self, name, value):
        self.name = name
        self.value = value
        
    def _latex(self, printer=None):
        return f"{self.name}:\\ \\text{{{self.__class__.__name__}}}"

class NonemptySet(Concept):
    def __init__(self, name, value):
        if value==sp.S.EmptySet:
            raise ValueError("Set must not be empty")
        super().__init__(name, value)
        
    def _latex(self, printer=None):
        return self.name

class Reals(NonemptySet):
    
    @strict_type_hints
    def __init__(self):
        self.name = sp.symbols('\\Re')
        super().__init__(self.name,sp.Reals)
        
    def _latex(self, printer=None):
        return self.name

class SampleSpace (NonemptySet):
    pass

class Algebra(NonemptySet):
    
    @strict_type_hints
    def __init__(self, 
                 name: sp.Symbol, 
                 Ω: SampleSpace, 
                 A: Either(NonemptySet, sp.Symbol)):
        self.Ω=Ω
        super().__init__(name, A)

    def _latex(self, printer=None):
        math=str(self.name).replace('$','')
        math2 = self.Ω._latex(printer)
        return f"{math}:\\ \\text{{{self.__class__.__name__} on }} ({math2})"

class 𝜎Algebra(Algebra):
    
    @strict_type_hints
    def __init__(self, name: sp.Symbol, Ω: SampleSpace, A: Algebra):
        self.Ω=Ω
        super().__init__(name, Ω, A)

class EventSpace(𝜎Algebra):
    
    @strict_type_hints
    def __init__(self, name: sp.Symbol, Ω: SampleSpace, A: 𝜎Algebra):
        super().__init__(name, Ω, A)

class 𝜎AdditiveFunction(Concept):
    
    @strict_type_hints
    def __init__(self, 
                 name: sp.core.function.UndefinedFunction, 
                 Ω: SampleSpace, 
                 A: 𝜎Algebra, 
                 f: sp.core.function.UndefinedFunction):
        self.Ω = Ω
        self.A = A
        super().__init__(name, f)

    def _latex(self, printer=None):
        math2 = self.A._latex(printer)
        return f"{self.name}: {self.A.name} \\to \\Re \\ \\text{{{self.__class__.__name__} on }} {math2}"

and so on. Any comments or suggestions on a more "SymPy-thonic" way of improving the above sketch would be greatly appreciated.

Answer 3

Here is a snippet that functions correctly. It might not do quite all the things you want yet, but hopefully it will start you off.

import sympy as sp
sp.init_printing()

class MetricSpace(sp.Expr):
    def __init__(self, M, d):
        self.M = M
        self.d = d

    def _latex(self, printer=None):
        return f"({printer.doprint(self.M)}, {printer.doprint(self.d)})\\ \\text{{{self.__class__.__name__}}}"

Re, d = sp.symbols(r'\Re,d')

MetricSpace(Re, d)

See here for more information on the Sympy printing system and how to hook into it. The key things to note is that you should subclass Expr if you want your object to come into the domain of Sympy's pretty printing, and that the latex printer calls the _latex method to get its LaTeX. It doesn't need to be wrapped in dollar signs, and you should use printer.doprint to get the LaTeX of nested expressions.

Here is a screenshot of what this produces on my notebook:

PS: I'm intrigued by a metric space with underlying set denoted by "\Re". If you mean the reals, might I suggest using sp.Reals?