Change the type setting of sympy output functions that contain the Heaviside step function from Theta(t) to u(t)

Question

When running Jupyter notebooks, or co-lab notebooks, and a function is returned that includes a step function, it displays the step function as Theta(t) and not u(t). While Theta(t) is common, when teaching step function concepts to students who are new to the field, switching symbols can cause a bit of unnecessary cognitive load.

Can one change the typesetting from Theta(t) to u(t)?

I tried as many searches of title variations as possible, from the specific change of typesetting to generic, and I can't find any resources on it.

Answer 1

This is how you modify a LatexPrinter to achieve your objective.

from sympy import *
from sympy.printing.latex import LatexPrinter

class MyLatexPrinter(LatexPrinter):
    def _print_Heaviside(self, expr, exp=None):
        pargs = ', '.join(self._print(arg) for arg in expr.pargs)
        tex = r"u\left(%s\right)" % pargs
        if exp:
            tex = r"\left(%s\right)^{%s}" % (tex, exp)
        return tex

def my_latex(expr, **settings):
    return MyLatexPrinter(settings).doprint(expr)

init_printing(latex_printer=my_latex)

With this edit, you can use your expression to achieve other goals, for example plotting:

var("s, t")
expr = 400 / (s**2 + 12*s + 400)
inv_lap = inverse_laplace_transform(expr, s, t)
display(inv_lap)
plot(inv_lap, (t, 0, 2))

This works because inv_lap contains an Heaviside function, which can be evaluated.