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))
12.340+/-0.010

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)
12.340+/-0.010

Aside

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

2 * 12.34 | pm | 0.01
    24.68+/-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
24.68+/-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
24.68+/-0.02

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)
-1209.0+/-2.8
print(112 ± 2 - 1321 ± 2)
-1209.0+/-2.8

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)
Overwriting uncert_ipython.py

After restarting the kernel, we can test this:

%load_ext uncert_ipython
1 +/- .1
1.0+/-0.1
%unload_ext uncert_ipython
1 +/- .1
  File "", line 1
    1 +/- .1
       ^
SyntaxError: invalid syntax
comments powered by Disqus