Carry units through a whole analysis

Carry units through a whole analysis#

Three libraries that know nothing about each other, and the units survive all of it. That is what the unxts.* packages are for, and it is much easier to watch than to describe.

So let’s watch. We will take a dataset whose units are only labels, make them real, compute a derived quantity from them, and plot the result — passing the data through xarray, unxt and matplotlib without writing one conversion factor or axis label by hand.

You need unxt, unxts.interop.xarray, unxts.interop.matplotlib, xarray and matplotlib installed. pip install "unxt[interop-xarray,interop-mpl]" gets all of it.

Set up#

>>> import matplotlib
>>> matplotlib.use("Agg")  # draw to memory, so no window opens

>>> import matplotlib.pyplot as plt
>>> import numpy as np
>>> import xarray as xr
>>> from astropy import constants as const

>>> import unxt as u
>>> import unxts.interop.xarray       # registers the .unxt accessor
>>> import unxts.interop.matplotlib   # registers the plotting converter

Those two imports are the entire integration. Each registers itself with the library it bridges, and then gets out of the way — see unxts.interop.xarray and unxts.interop.matplotlib.

Load data whose units are only labels#

Here is a rotation curve: how fast things orbit the centre of a galaxy, as a function of how far out they are. Recorded the way scientific data usually is, with the units written in the attributes:

>>> ds = xr.Dataset(
...     {
...         "radius": ("i", np.array([2.0, 4.0, 6.0, 8.0, 10.0]),
...                    {"units": "kpc"}),
...         "v_circ": ("i", np.array([120.0, 180.0, 210.0, 220.0, 218.0]),
...                    {"units": "km/s"}),
...     }
... )

>>> ds["v_circ"].attrs
{'units': 'km/s'}

That 'km/s' is a string sitting beside the numbers. Nothing enforces it — the values will happily be doubled into something that is no longer km/s and keep the label anyway:

>>> ds["v_circ"].values * 2
array([240., 360., 420., 440., 436.])

Make the units real#

quantify() reads each units attribute and replaces the plain arrays with unxt quantities:

>>> qds = ds.unxt.quantify()

>>> r = qds["radius"].data
>>> r
Quantity(Array([ 2.,  4.,  6.,  8., 10.], dtype=float32), unit='kpc')

>>> v = qds["v_circ"].data
>>> v
Quantity(Array([120., 180., 210., 220., 218.], dtype=float32), unit='km / s')

The unit is part of the data now, not a note attached to it.

Compute something new#

The mass enclosed within radius \(r\) for a circular orbit is \(M = v^2 r / G\). We need the gravitational constant, in whatever units it comes in:

>>> G = u.Q(const.G.value, "m3 / (kg s2)")
>>> G
Quantity(Array(6.6743e-11, dtype=float32), unit='m3 / (kg s2)')

Now just write the formula. Kiloparsecs, kilometres per second and SI metres all in one expression:

>>> M = v**2 * r / G

Look at what came out:

>>> M.unit
Unit("km2 kg kpc / m3")

That is an ugly unit, and it is exactly right — unxt did the algebra without tidying up after itself. Ask what it is, though, and the answer is clean:

>>> u.dimension_of(M)
PhysicalType('mass')

A mass. Notice we have verified the formula is dimensionally correct before converting anything, using the dimensions rather than the units. Now ask for it in units an astronomer reads:

>>> M_sun = M.uconvert("solMass")
>>> M_sun[3]
Quantity(Array(9.00273e+10, dtype=float32), unit='solMass')

About 9 × 10¹⁰ solar masses inside 8 kpc — which is roughly the right answer for the Milky Way inside the Sun’s orbit. We never wrote down a single conversion factor between kiloparsecs, kilometres, metres, kilograms and solar masses.

Plot it#

Hand the quantities straight to matplotlib:

>>> fig, ax = plt.subplots()
>>> _ = ax.plot(r, v)

and look at the axes:

>>> print(ax.get_xlabel())
$\mathrm{kpc}$

>>> print(ax.get_ylabel())
$\mathrm{km\,s^{-1}}$

We did not call set_xlabel. The converter read the unit off each quantity and labelled the axis with it, typeset for the figure.

Plot the mass we derived, and the label follows the same route:

>>> fig2, ax2 = plt.subplots()
>>> _ = ax2.plot(r, M_sun)
>>> print(ax2.get_ylabel())
$\mathrm{M_{\odot}}$

>>> plt.close("all")

The solar-mass symbol, from a quantity we computed rather than declared.

Put it back#

Store the derived column alongside the originals and hand the dataset back to the world it came from:

>>> qds["m_enc"] = ("i", M_sun)
>>> out = qds.unxt.dequantify()

>>> out["m_enc"].attrs
{'units': 'solMass'}

dequantify() moved the unit back into the attributes, ready to be written to NetCDF. The label on the way out was derived, not typed.

What we built#

A complete analysis — load, compute, plot, save — in which the units were attached to the data at the start and carried themselves through three libraries to the end. The only place a unit name was written by hand was the one place it was a genuine choice: asking for the answer in solar masses.

The rest of the ecosystem#

unxt is the core; each unxts.* package bridges it to something else. Two of them appeared above.

Package

For

Docs

unxts.interop.xarray

labelled arrays and datasets

docs

unxts.interop.matplotlib

plotting quantities

docs

unxts.interop.gala

unit systems from gala

docs

unxts.parametric

the physical dimension in the type, checked at runtime

docs

unxts.linalg

matrices whose elements carry different units

docs

unxts.hypothesis

property-based testing over quantities

docs

unxts.api

the abstract API, for making your own types speak unxt

docs

Each has its own tutorial; any of them will drop into a pipeline like this one the same way these two did.

Where to go next#