Uncertainty extension for IPython

Wouldn’t it be nice if we had uncertainty with a nice notation in IPython? The current method would be to use raw Python,

from uncertainties import ufloat

print(ufloat(12.34, 0.01))

Let’s use the infix library to make the notation easier. We’ll define |pm| to mean +/-.

Note: this article was written with infix 1.0. As of 1.1, you can now directly use pow_infix without creating your own.

from infix import or_infix, custom_infix, base_infix

pm = or_infix(ufloat)
print(12.34 | pm | 0.01)

Aside

Our | operator has very low precedence. So, things like this are not very useful:

2 * 12.34 | pm | 0.01

The highest order we can do with the infix library is __mul__,

pm = custom_infix("__rmul__", "__mul__")(ufloat)
2 * 12.34 * pm * 0.01

The library doesn’t have support for python’s only right-associative operator, **. We can add that easily, though:

class pow_infix(base_infix):
    def right(self, right):
        return self.func(right, self.__infix__)

    __pow__ = base_infix.left
    __rpow__ = right
pm = pow_infix(ufloat)
2 * 12.34**pm**0.01

Extending IPython

We’ve made an improvement, but this is ugly notation. Let’s patch IPython to find the notation +/- or ± and replace it with our infix operator before passing it to the underlying python kernel:

from IPython.core.inputtransformer import StatelessInputTransformer


@StatelessInputTransformer.wrap
def my_special_commands(line):
    return line.replace("+/-", "**pm**").replace("±", "**pm**")

We only needed a line transformer, and there was no need to even handle storing a state. Thankfully, we didn’t need something like the AST this time, since it’s a simple replacement only; that’s the beauty of starting with a valid python construct for infix operations.

Now, let’s use our transformer in the current IPython instance:

import IPython

ip = IPython.get_ipython()
ip.input_splitter.logical_line_transforms.append(my_special_commands())
ip.input_transformer_manager.logical_line_transforms.append(my_special_commands())
print(112 +/- 2 - 1321 +/- 2)
print(112 ± 2 - 1321 ± 2)

Loadable extension

Let’s make a IPython extension that makes all of this automatic. I’m removing the need for the infix library (since it’s only a few lines).

%%writefile uncert_ipython.py

from uncertainties import ufloat
import IPython
from IPython.core.inputtransformer import StatelessInputTransformer

class pow_infix(object):
    def __init__(self, func):
        self.func = func

    def __pow__(self, left):
        self.__infix__ = left
        return self

    def __rpow__(self, right):
        return self.func(right, self.__infix__)

pm = pow_infix(ufloat)

@StatelessInputTransformer.wrap
def my_special_commands(line):
    return line.replace('+/-','**pm**').replace('±','**pm**')

comm1 = my_special_commands()
comm2 = my_special_commands()

def load_ipython_extension(ipython):
    ipython.input_splitter.logical_line_transforms.append(comm1)
    ipython.input_transformer_manager.logical_line_transforms.append(comm2)
    ipython.user_ns['pm'] = pm

def unload_ipython_extension(ipython):
    ipython.input_splitter.logical_line_transforms.remove(comm1)
    ipython.input_transformer_manager.logical_line_transforms.remove(comm2)

After restarting the kernel, we can test this:

%load_ext uncert_ipython
1 +/- .1
%unload_ext uncert_ipython
1 +/- .1
comments powered by Disqus